Memorization Dynamics in Knowledge Distillation for Language Models

TL;DR

This work investigates memorization dynamics in knowledge distillation (KD) for language models, addressing privacy concerns associated with training data leakage. Through empirical analysis across three LLM families (Pythia, OLMo-2, Qwen-3) and three datasets (FineWeb, WikiText, Nemotron-CC-v2), the authors show that KD substantially reduces memorization while improving generalization compared with standard fine-tuning. They demonstrate that memorization is highly predictable from features such as $zlib$ entropy, perplexity, and KL divergence, and identify a core set of easy-to-memorize examples that dominate the student’s memorization. The study also contrasts soft (logit-level) and hard (sequence-level) distillation, finding similar overall memorization rates but higher risks with hard distillation due to increased inheritance of teacher-specific memorized content. Finally, they show that pre-identifying and removing memorized examples before distillation can dramatically lower memorization in the student, highlighting practical privacy-preserving benefits of KD and offering guidance for deploying KD in privacy-sensitive settings.

Abstract

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.

Figures (18)

• Figure 1: Experimental framework. (Left) Training setup: The Teacher and Baseline models are independently fine-tuned from Pythia 12B and Pythia 1.4B base models, respectively, on dataset $\mathcal{D}$ using cross-entropy. The Student is initialized from the Pythia 1.4B base model and distilled on the same dataset $\mathcal{D}$ to match the Teacher's logit distribution by minimizing KL divergence. (Right) Discoverable memorization evaluation: We prompt models with 50-token prefixes from training examples. An example is classified as memorized if the model's greedy generation of the subsequent 50 tokens exactly matches the ground truth suffix.
• Figure 2: Effect of Temperature on Memorization. We find that increasing the temperature during distillation reduces memorization in the Student model.
• Figure 3: Consistency heatmap. We compare the consistency of memorization by the count from three independent runs for the baseline (rows) and student (columns). The cell values represent the number of examples. The strong density in the (three, three) cell confirms that naturally "easy" examples are consistently memorized by both models. Conversely, the high count in (three, zero) highlights 494 examples that are consistently memorized by the baseline but successfully suppressed (memorized in zero runs) by the student.
• Figure 4: Overlap of memorized examples. The majority (80%) of examples memorized by Pythia 1.4B baseline are also memorized by the Pythia 12B teacher. We term these consistently memorized examples as easy-to-memorize (enclosed by a bold outline). The distilled student primarily memorizes a subset of these specific examples. To account for variance in training dynamics, we train three student and three baseline models with different random seeds and report the union of memorized examples across these runs.
• Figure 5: Intrinsic properties of easy-to-memorize examples. We plot zlib entropy versus baseline perplexity for the easy-to-memorize examples (red) compared to a random subset of 25,000 other examples (grey). They form a distinct cluster with significantly lower entropy and perplexity.