PaperlayerView exact revision on arXiv

Source-linked AI summary

Memorization Dynamics in Knowledge Distillation for Language Models

Jaydeep Borkar, Karan Chadha, Niloofar Mireshghallah, Yuchen Zhang, Irina-Elena Veliche, Archi Mitra, David A. Smith, Zheng Xu, Diego Garcia-Olano

arXiv:2601.15394v2cs.CL

TL;DR

Training-data memorization in knowledge distillation is less understood than in standard pre-training and fine-tuning. The paper studies this across three model families and three datasets, finding reduced student memorization, predictable high-risk examples, and greater teacher-specific inheritance under hard distillation.

  • Problem

    Memorization dynamics in knowledge distillation remain poorly understood despite extensive study of memorization in standard pre-training and fine-tuning.

  • Method

    The study compares distilled students, teachers, and same-size standard fine-tuning baselines across Pythia, OLMo-2, and Qwen-3 models and three datasets.

  • Results

    78% of teacher generalization is recovered while only 2% of teacher memorization is inherited, and hard distillation inherits 2.7× more teacher-specific examples than soft distillation.

  • Takeaways & Limitations

    Distillation improves utility while significantly reducing training-data memorization relative to standard fine-tuning, with easy-to-memorize examples accounting for most student memorization.

  • Takeaways & Limitations

    The analysis assumes teacher and baseline models were trained on data similar to the distillation data.

Abstract

from arXiv · show

Knowledge Distillation (KD) is increasingly adopted to transfer capabilities from large language models to smaller ones, offering significant improvements in efficiency and utility while often surpassing standard fine-tuning. Beyond performance, KD is also explored as a privacy-preserving mechanism to mitigate the risk of training data leakage. While training data memorization has been extensively studied in standard pre-training and fine-tuning settings, its dynamics in a knowledge distillation setup remain poorly understood. In this work, we study memorization across the KD pipeline using three large language model (LLM) families (Pythia, OLMo-2, Qwen-3) and three datasets (FineWeb, Wikitext, Nemotron-CC-v2). We find: (1) distilled models memorize significantly less training data than standard fine-tuning (reducing memorization by more than 50%); (2) some examples are inherently easier to memorize and account for a large fraction of memorization during distillation (over ~95%); (3) student memorization is predictable prior to distillation using features based on zlib entropy, KL divergence, and perplexity; and (4) while soft and hard distillation have similar overall memorization rates, hard distillation poses a greater risk: it inherits $2.7\times$ more teacher-specific examples than soft distillation. Overall, we demonstrate that distillation can provide both improved generalization and reduced memorization risks compared to standard fine-tuning.

1 Introduction

Knowledge distillation is studied as a potentially efficient and privacy-relevant alternative to standard fine-tuning, but its memorization dynamics remain poorly understood. Across multiple model families and datasets, the paper finds reduced memorization, predictable example-level risk, and different inheritance patterns for soft versus hard distillation.

  • Knowledge distillation transfers knowledge from larger teacher models to smaller student models for utility and efficiency benefits.
  • Training-data memorization in knowledge distillation remains poorly understood despite extensive study in pre-training and fine-tuning.
  • The study examines memorization across Pythia, OLMo-2, and Qwen-3 using FineWeb, WikiText, and Nemotron-CC-v2.
  • 78% of the teacher’s generalization is recovered while only 2% of its memorization is inherited under logit-level distillation.
  • Certain inherently easy-to-memorize examples are consistently memorized across model sizes and preferentially memorized by distilled students.
  • Student memorization is predictable before distillation from computable features, including sequence-level entropy and log probability analyses.
  • Hard distillation inherits 2.7× more teacher-specific memorization of difficult examples than soft distillation, despite similar overall overlap.

2 Background and Experimental Setup

The paper defines its teacher–student–baseline setup, trains models with either distillation or standard cross-entropy, and evaluates discoverable memorization on training sequences. The setup also examines temperature effects and compares soft and hard distillation.

  • The teacher is the larger model guiding training, the student is the smaller model mimicking it, and the baseline is independently fine-tuned with cross-entropy.
  • The study extends beyond the primary setup to OLMo-2 and Qwen-3 models and Nemotron-CC-v2 and WikiText-103 datasets.
  • Increasing distillation temperature reduces memorization in the student model.
  • The primary setup uses 1M FineWeb examples of length 256, with a Pythia 12B teacher and Pythia 1.4B student.
  • The student is trained with forward KL divergence, while the same-size baseline is trained independently using standard cross-entropy under comparable computational budgets.
  • Temperature is a distillation parameter set to 2.0 in the primary experiments and scales the token probabilities used in training.
  • An example is counted as memorized when greedy generation from its 50-token prefix exactly matches its 50-token ground-truth suffix.
  • The evaluation is applied to 1M examples from the training dataset, and the paper compares soft with sequence-level hard distillation.

3 Memorization During Distillation

Across model families and datasets, distillation reduces memorization while preserving or improving generalization. Memorization concentrates in inherently easy examples, whose selection is consistent within families but differs across architectures.

  • 3.1 Distilled Models Generalize Better and Memorize Less: Distilled students memorize less than same-size fine-tuning baselines across datasets and architectures.On Pythia, reductions are approximately 2.4× on FineWeb and 2.1× on Wikitext; Nemotron-CC-v2 shows 0.0012% versus 0.0091%.
  • 3.1 Distilled Models Generalize Better and Memorize Less: Distillation improves validation loss and perplexity for Pythia and OLMo-2 students compared with baselines.Table 1 reports that distilled students consistently outperform their respective baselines across Pythia, OLMo-2, and Qwen-3 on FineWeb.
  • 3.1 Distilled Models Generalize Better and Memorize Less: Only 18 of 1,955 teacher-exclusive memorized examples were inherited by the student, approximately 0.9%.The student therefore inherited generalization while rejecting most examples memorized exclusively by the teacher.
  • 3.2 Some Examples are Easier to Memorize: 96% of examples memorized by Pythia 1B persist at 1.4B, while the 12B teacher captures approximately 80% of the 1.4B baseline set.Across random seeds, the baseline also repeatedly memorizes a core set of identical examples, termed easy-to-memorize.
  • 3.2 Some Examples are Easier to Memorize: Easy-to-memorize examples have significantly lower zlib entropy and baseline perplexity than other training examples.They form a distinct cluster, and similar clustering appears across additional models and datasets.
  • 3.2 Some Examples are Easier to Memorize: 95.7% (676/706) of student memorization consists of examples shared by the teacher and baseline, although many such examples remain unmemorized by the student.The baseline consistently memorized 494 examples across three runs that the student never memorized, and the student omitted 696 examples shared by baseline and teacher.
  • 3.2.1 Do All Models Memorize the Same Easy Examples?: Pythia, OLMo-2, and Qwen-3 target low-entropy data but memorize largely non-overlapping examples.Cross-architecture perplexity comparisons show that examples easy for one model can be difficult for the others, reflecting model-specific inductive biases.

4 Identifying features associated with student memorization

The section identifies pre-distillation features that predict student memorization and shows that filtering predicted memorized examples sharply reduces residual memorization.

  • 706 examples are memorized by the student model after distillation, forming the positive class for prediction.
  • Student memorization is predicted using teacher and baseline perplexity, KL-divergence loss, and zlib entropy, or a reduced feature set without baseline information.
  • A classifier using teacher perplexity, baseline perplexity, KL-divergence, and zlib entropy achieves AUC-ROC 0.9997 ± 0.0005 and Recall 1.0000 ± 0.0000.
  • Using only teacher perplexity and zlib entropy achieves AUC 0.9998 ± 0.0005, Recall 1.0000 ± 0.0000, and Precision 0.9940 ± 0.0054.
  • Zlib entropy is the dominant feature, with lower values associated with easy-to-memorize examples and memorized examples separated from non-memorized examples.
  • Removing predicted student-memorized examples before distillation reduces memorized examples from 706 to four, a 99.4% reduction.

5 Why Does Distillation Reduce Memorization?

The analysis attributes distillation’s lower memorization to how soft targets handle difficult, uncertain examples: students remain uncertain rather than forcibly memorizing them, while retaining easy examples.

  • The mechanism is studied by comparing examples memorized by the baseline but not by the distilled student.
  • The analysis measures sequence log-probability as confidence and average Shannon entropy as intrinsic uncertainty on each example’s 50-token suffix.
  • The teacher is confident and low-entropy on memorized examples, whereas the baseline assigns high probability despite substantially higher entropy.
  • Cross-entropy forces the capacity-limited baseline to memorize sequences it models with high uncertainty.
  • KL divergence lets the student output a flatter, more uncertain distribution when it cannot match the teacher’s certainty, avoiding memorization of difficult examples.
  • Student-memorized examples cluster with the teacher’s low-entropy examples, indicating selective memorization of sequences simple enough to learn confidently.

6 Soft vs. Hard Distillation

Hard distillation trains on teacher-generated sequences, while soft distillation matches full output distributions. Their overall memorization rates are similar, but hard distillation inherits more teacher-specific difficult examples.

  • Setup: Hard distillation trains the student with cross-entropy on teacher-generated output sequences instead of matching the teacher’s full probability distribution through KL divergence.
  • Setup: The hard-distilled student uses sequences generated from the teacher after prompting with the first 50 tokens and generating 206 tokens, producing sequences of length T = 256.
  • Setup: Hard and soft students are trained with the same compute budget and learning rate, and hard-distillation performance is evaluated on LAMBADA and Winogrande.
  • Performance: The hard-distilled student achieves perplexity 6.43, LAMBADA accuracy 56.65%, and Winogrande accuracy 57.46%.
  • Memorization: 0.07% is the hard-distilled student’s memorization rate, identical to the soft-distilled student and below the baseline’s 0.17%.
  • Memorization: Approximately 70% of hard-student memorized examples overlap with soft-student examples, and 90% are also memorized by the teacher and baseline.
  • Teacher-specific inheritance: Hard distillation inherits 50 teacher-specific examples, 2.7× the amount inherited by soft distillation, because it is more prone to retaining difficult teacher examples.

7 Related Work

The related work situates this study at the intersection of knowledge distillation, language-model memorization, and privacy, emphasizing that memorization remains insufficiently understood in KD.

  • Knowledge distillation transfers knowledge from large teachers to smaller students through output-distribution matching and related feature-matching approaches.
  • Prior memorization research shows that generative and fine-tuned language models can regurgitate portions of their training data.
  • Despite extensive memorization research, memorization remains poorly understood in a knowledge-distillation setup.
  • Prior KD privacy studies include membership inference and memorization inheritance, while other work examines discarded undesirable behaviors.
  • This work tracks memorization across teacher, baseline, and student models, measuring both the amount and type of memorized examples.

8 Conclusion

The paper finds that knowledge distillation improves utility while reducing training-data memorization relative to standard fine-tuning. It also identifies predictable high-risk examples and a higher inheritance risk for hard distillation.

  • Distillation improves model utility while significantly reducing training-data memorization compared to standard fine-tuning.
  • Distilled models primarily memorize easy-to-memorize examples, enabling prediction and removal of high-risk examples before distillation.
  • Soft and hard distillation memorize largely the same examples, but hard distillation has greater teacher-memory inheritance.

A.1 Easy-to-Memorize Examples on Additional Models and Datasets

Easy-to-memorize examples form a distinct low-entropy, low-perplexity cluster across additional models and datasets. This clustering indicates that the finding generalizes beyond the original setting.

  • Similar zlib-entropy and Baseline-perplexity clustering appears for Pythia on Wikitext and OLMo-2 and Qwen-3 on FineWeb.
  • Easy-to-memorize examples cluster at significantly lower zlib entropy and Baseline Perplexity than other training examples.

A.2 Robustness to the memorization definition

The reduction in memorization persists when memorization is defined approximately or evaluated with longer prefixes. These robustness checks support the main memorization finding across alternative evaluation settings.

  • 54.5% fewer sequences show approximate memorization under the BLEU-based criterion for distilled students than fine-tuning baselines.The baseline shows 1,104 approximately memorized sequences versus 502 for the student.
  • 316 examples are memorized by the student versus 630 by the fine-tuning baseline under 100-token-prefix exact matching.
  • Distillation reduces memorization under both approximate and longer-prefix evaluation settings.

A.3 Extended Generalization Results

Extended analyses show that easy-to-memorize examples, feature-based prediction, and reduced memorization also appear across models, datasets, and evaluation setups. The pre-training experiment reproduces the main memorization pattern, while hard distillation carries greater inheritance risk.

  • A.3 Extended Generalization Results: 88% of Pythia, 85% of OLMo-2, and 70% of Qwen-3 student memorization consists of easy-to-memorize examples.
  • A.3 Extended Generalization Results: The student inherits 9% of Pythia, 13% of OLMo-2, and 27% of Qwen-3 memorization directly from the teacher.
  • A.3 Extended Generalization Results: Teacher Perplexity, zlib Entropy, KL divergence, and Baseline Perplexity distinguish memorized from non-memorized student examples, with lower values associated with memorization.
  • A.3 Extended Generalization Results: Examples memorized by both baseline and student typically have lower Baseline perplexity and teacher-student KLD loss, whereas baseline-only examples have Baseline PPL below 10 and high KLD loss.
  • A.3 Extended Generalization Results: In the pre-training setup, memorization rates are 1.57% for the Teacher, 0.27% for the Baseline, and 0.06% for the Student.
  • A.3 Extended Generalization Results: Of 2,785 Baseline-memorized examples, 76% are also memorized by the Teacher, while only 361 of those are memorized by the Student.
Loading 2601.15394v2…