Dataset distillation for memorized data: Soft labels can leak held-out teacher knowledge

What came before

Classic knowledge distillation (Hinton et al. 2015) and label-worth-2024 both treat soft labels as carriers of "structured" distributional information (inter-class similarity, dark knowledge) that helps a student generalize; theoretical accounts (Menon et al.'s statistical perspective on distillation, Boix-Adsera's theory of model distillation) likewise assume the teacher reflects meaningful structure in the data. Separately, the memorization literature (Zhang et al.'s random- label fitting result, Nanda et al.'s grokking/modular-addition setup) shows neural networks readily memorize facts unrelated to any generalizable structure. No prior work had asked whether soft-label distillation transfers a teacher's memorized, non-generalizing knowledge about specific training examples, as opposed to only generalizable structure.

The problem

If a teacher used to relabel a distilled dataset has memorized specific facts about its training data (as large modern classifiers do), it is unknown whether a student trained only on that teacher's soft labels for a held-in fraction of examples can recover non-trivial information about the teacher's held-out memorized examples — a privacy/leakage question that is invisible in standard accuracy benchmarks, which cannot distinguish "the student learned real structure" from "the student learned what the teacher happened to memorize."

The idea

Isolate memorization from generalization by constructing settings where generalization is a priori impossible (finite i.i.d. random inputs with i.i.d. random labels), so that a teacher which fits the data must be purely memorizing it; then show a student trained on the teacher's soft labels for only a fraction of that data achieves non-trivial, sometimes perfect, accuracy on the held-out remainder it never observed. Softmax temperature acts as a dial between two student solutions: at low temperature the student recovers only class identity (soft labels approach one-hot, no leakage); at high temperature the student can recover the teacher's function almost exactly, including on the memorized held-out data.

How it works

Four settings, each partitioning a finite teacher training set $\mathcal{D}^{\text{T}}_\star$ into a student-visible fraction $\rho$ and a held-out remainder, with the student trained by cross-entropy against the teacher's temperature-$\tau$ softmax soft labels $\hat{\mathbf{y}}=\sigma_\tau(f^\star(\mathbf{x}))$ only on the visible fraction: (1) a single-layer transformer on modular addition ($p{=}113$, the grokking/delayed-generalization setup of Nanda et al.), with teachers checkpointed at different points to obtain memorizing-only, partially- and fully-generalizing teachers; (2) multinomial logistic regression on i.i.d. Gaussian inputs with i.i.d. random labels, where three sample-complexity thresholds $\alpha=n/d$ are derived in closed form in the high-dimensional proportional limit — teacher memorization capacity $\alpha^{\text{T}}_{\text{label}}\le 2$ (Cover's theorem), student identifiability $\alpha^{\text{S}}_{\text{id}}(\rho)=1/\rho$, and student memorization capacity $\alpha^{\text{S}}_{\text{label}}(\rho)$; (3) matched-capacity single-hidden-layer ReLU MLPs (multi-class, up to $c{=}100$), where two distinct leaking solutions are identified — one that memorizes the soft-label vectors themselves, another that functionally matches the teacher — with a sharp phase transition between them at the identifiability threshold; (4) a pretrained GPT-2 with a linear classification head fine-tuned to memorize 6000 random three-number token sequences mapped to random classes, to test whether the phenomenon persists in a large pretrained architecture. In the loop: one pretrained/fitted teacher per setting (frozen once fit); no image synthesis or bi-level optimization — the "distilled" artifact is purely the teacher's per-sample soft-label vector.

Evidence

Modular addition, single-layer transformer, $p{=}113$, $\tau{=}10$ (Fig. 2): for a non-generalizing (memorizing-only) teacher, the student's accuracy on the held-out memorized set $\mathcal{D}^{\text{S}}_{\text{test}}$ reaches 1.0 as $\rho$ grows, while its accuracy on true task-generalization data $\mathcal{D}_{\text{val}}$ stays near the (low) teacher's own validation accuracy — the leaked information is specific to what the teacher memorized, not the underlying modular-addition structure. Binary logistic regression, $d{=}1600$, $\rho{=}0.8$ (Fig. 3): held-out test accuracy on memorized data grows monotonically with $\alpha{=}n/d$ and reaches $\ge$99% at $\alpha^{\text{S}}_{\text{id}}(\rho{=}0.8)\simeq1.26$ — a fifth of the "memorized" data is recovered essentially perfectly from soft labels alone, just above the identifiability threshold. GPT-2 on random sequences (Sec. 5.3, Fig. 7): teacher fits 6000 memorized sequence-class pairs (1000 possible classes, random-guess baseline $\approx$0.1%); students trained on soft labels for $\rho=\{0.2,0.5,0.8\}$ of that set reach held-out test accuracy $\{21.3\%,52.4\%,65.2\%\}$ while training accuracy is $\approx$99.5% in all three cases — at $\rho{=}0.8$ the student recovers >60% accuracy on sequences it never saw, from a teacher whose only "knowledge" is memorization by construction. Multi-class scaling (Fig. 4B, $\rho{=}0.8$, $\tau{=}10$): the four relevant capacity/ leakage thresholds scale differently with class count $c$ ($\alpha^{\text{S}}_{\text{id}}\sim1/c$; $\alpha^{\text{T}}_{\text{label}}\sim\alpha^{\text{S}}_{\text{label}}\sim1/\log c$; a class-shuffled control threshold $\alpha^{\text{S-shuffle}}_{\text{label}}\sim1/\sqrt{c}$).

Limitations

Explicitly restricted to synthetic, structured-or-explicitly-memorized data (modular addition, i.i.d. random Gaussian/label data, random token sequences), not natural data distributions; uses simple models (logistic regression, one-hidden-layer ReLU MLPs, small transformers) chosen for theoretical tractability, plus a single GPT-2 experiment whose leakage thresholds are not shown to generalize to deeper architectures; does not explore regularization or optimizer choices as factors in capacity; teachers are constructed to either purely memorize or purely generalize, not the realistic mixture large models exhibit; no privacy defense or mitigation is proposed or evaluated, left as future work ("whether it can be hidden for privacy reasons").