Divergence controls entropy in distillation

Nicolas Zucchet; Scott W. Linderman

cs.LGOct 2, 2026

Divergence controls entropy in distillation

Authors: Nicolas Zucchet, Scott W. Linderman

Organizations: Stanford University

Abstract

Distillation has become a core primitive of large language model training, but its properties are not yet well understood. We take an entropic perspective, studying how the entropy of the student depends on the data and the divergence that define the distillation objective. We prove that forward KL inflates the entropy of the student above that of the teacher. Since cross-entropy training is a special case, this yields an identity that we verify quantitatively in pretraining and supervised finetuning. Other divergences come with no such guarantee: reverse KL deflates entropy until the gap between student and teacher gets too large, and interpolating between the two changes entropy smoothly early in training but abruptly at convergence. The lower entropy of on-policy distillation comes from token-level reverse KL, not from on-policy sampling. The divergence therefore acts as an implicit entropy regularizer, whose role is clearest in self-distillation: as conditioning on privileged information deflates entropy, the divergence hyperparameters that work best are those that compensate for it.

Figures & tables

Figure 1: Decoupling divergence from sampling in distillation. A prompt

x

is completed by the sampling distribution

p

, and the student

\pi_{\theta}

is then compared to the teacher

\pi^{*}

on each token of the completion through a divergence

\ell

between their next-token distributions. Writing the objective at the token level keeps the two choices independent, whereas a sequence-level divergence ties them together. Details in Section 2 .

Figure 2: A toy language model reveals phase transitions in how the divergence affects entropy. Toy model (a). A linear softmax student maps a fixed embedding

e_{i}

per context

i

to next-token probabilities and is trained towards a teacher distribution per context with forward KL, reverse KL, or the generalized Jensen–Shannon (JS) divergence whose weight

\lambda

interpolates between them (Section 3.1 ). Its shared weights

W

prevent it from reproducing all teacher distributions when there are too many contexts. Simulations (b–e). At convergence, forward KL (blue) always inflates entropy relative to the teacher and reverse KL (red) tends to deflate it, unless the teacher is too hard to model because there are too many contexts (b) or because it is too diffuse (c), in which case reverse KL inflates entropy too. Along the generalized JS family (colors in between), entropy smoothly decreases with the weight

\lambda

early in training (e), whereas at convergence it drops abruptly (d).

Figure 3: Restricting the divergence to the top-

k

tokens shifts entropy depending on how the support and the tail are handled. Same setup as in Figure 2 with

512

contexts, this time varying

k

out of a vocabulary of

64

, when using the student (left) or the teacher (right) to select the top-

k

tokens and when using the renormalized or tail bucket cropping strategy. Both choices substantially change how the restriction affects entropy. Forward KL here; other divergences in Figure 8 .

Figure 4: Cross-entropy tracks entropy in pretraining and supervised finetuning. Average entropy of the student

\bar{H}(\pi_{\theta})

against its average cross-entropy loss on the data, both computed from the student alone, for OLMo 2 at different sizes (marker size), before (hollow) and after (filled) supervised finetuning (SFT), on the pretraining (left, green) and finetuning (right, orange) mixtures; the grey line is entropy equals cross-entropy. The identity holds on the data each model was last trained on, as predicted, and also, without guarantee, on the finetuning data after pretraining, whereas supervised finetuning breaks it on the pretraining data.

Figure 5: The lower entropy of on-policy distillation is due to token-level reverse KL, on-policy sampling nudges it up, and teacher size barely matters. Average entropy of the Qwen3-0.6B-Base student on its own MATH500 completions over distillation from Qwen3-Instruct teachers of increasing size in non-thinking mode (left to right). Color is the divergence, with the full color for on-policy and a lighter shade for off-policy, and each curve is a single seed.

Figure 6: Self-distillation tends to deflate entropy, the choice of divergence must compensate for it to avoid entropy collapse. Self-distillation experiments on the Chemistry dataset following Hübotter et al. (2026) , varying the weight of the generalized Jensen–Shannon (JS) divergence (left), the frequency at which the teacher is updated (center), and how many tokens are taken into account (right); defaults are

\lambda=0.5

, a teacher that is an exponential moving average (EMA) of the student with decay

0.95

, and the top-

100

tokens of the student with tail bucket. Each curve is a single seed. Conditioning on privileged information lowers the entropy of the teacher ( Nicolicioiu et al., 2026 ) , and as the teacher derives from the student, this can compound into entropy collapse, e.g., with reverse KL (left) or a live teacher (center). Divergences that inflate entropy, a slowly moving teacher, or a tail bucket (right) counterbalance it.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: The entropy gap phase transition moves with the divergence. Average entropy gap $\bar{H}(\pi_{\theta})-\bar{H}(\pi^{*})$ at convergence over the plane of teacher perplexity and number of contexts in the setup of Section 3 , at five weights $\lambda$ of the generalized Jensen–Shannon family, with the forward KL on the left and the reverse KL on the right. Each cell is its own configuration, averaged over $5$ seeds, and color indicates the entropy gap. The forward KL always inflates entropy and increasing $\lambda$ progressively moves the diagonal boundary, that is the task needs to be harder for the student to inflate entropy.

Figure 8: The top- $k$ approximation across the divergence family. Average student entropy over training, for the two choices the approximation makes: which model defines the support, i.e. whether we take the student’s own top- $k$ or the teacher’s, and what happens to the mass left outside the support, either renormalizing both sides over the window or keeping it in one extra tail bucket. Rows are the generalized Jensen–Shannon weight, from the forward KL at $\lambda=0$ to the reverse KL at $\lambda=1$ . Curves are averaged over $5$ seeds $\pm\mathrm{std}$ , and the dashed line is the teacher entropy.

Figure 9: At the end of pretraining, a model’s mean entropy equals its mean cross-entropy. Similar to Figure 4 (left), for different model families and at different stages of pretraining. (left) OLMo 2 at 1B, 7B, 13B and 32B, at the end of stage 1 of pretraining (light) and at the end of pretraining (dark). (right) the eight Pythia models from 70M to 12B on the Pile ( Gao et al., 2020 ) .

Figure 10: Teacher-to-student distillation from non-thinking teachers. Average entropy of the student on frozen teacher completions (top) and on its own completions (middle, as in Figure 5 ), and MATH500 accuracy (bottom), against distillation step, for teachers of increasing size (left to right). Color is the divergence, with the full color for on-policy and a lighter shade for off-policy. Each curve is a single seed.

Setup
Trainer	general online logit distillation trainer of TRL 1.9.2 ( von Werra et al., 2020 )
Reference implementation	Zhao et al. (2026a)
Generation engine	vLLM ( Kwon et al., 2023 ) , colocated
Dataset	DeepScaleR ( Luo et al., 2025 )
Student model	Qwen3-0.6B-Base ( Yang et al., 2025 )
Teacher models	Qwen3-1.7B, Qwen3-4B, Qwen3-8B, thinking disabled

Appendix

Table 1: Hyperparameters of the on- vs. off-policy distillation experiments. Values not motivated in the text follow Zhao et al. (2026a) , with three teachers rather than their two and one seed rather than their three.

Figure 11: Accuracy of the student model as a function of its entropy, in the self-distillation experiment of Figure 6 . Each dot is a different checkpoint, the larger the dot, the later in training it is. JS stands for Jensen–Shannon and EMA for exponential moving average.

Figure 12: Self-distillation on GSM8K behaves as on the Chemistry dataset. Entropy of the student (top) and accuracy, averaged over

16

samples per question (bottom), against training step, varying the weight of the generalized Jensen–Shannon (JS) divergence (left), the teacher (center), and how many tokens are taken into account and whether a tail bucket is used for the remaining tokens or not (right); defaults are

\lambda=0.5

, an exponential moving average (EMA) teacher with decay

0.95

, and top-

100

with tail bucket. Lines are means over three seeds and bands their range, and the dashed line is the accuracy of the base model.

Figure 13: The more a configuration deflates entropy, the smaller the learning rate it tolerates. Entropy at step

50

(top) and accuracy over the last five evaluations (bottom) against learning rate, varying the divergence weight (left) and the teacher (right), on GSM8K. Markers are means over three seeds and bars their range. JS stands for Jensen–Shannon and EMA for exponential moving average.

Figure 14: Entropy settles within

500

steps, and

\lambda=0

is the only weight that ends up inflating it. Entropy of the student over

500

GSM8K steps, one seed per generalized Jensen–Shannon (JS) divergence weight. Forward KL converges at a high entropy level, every other

\lambda

deflates below the initial entropy, and all configurations have flattened by the end of training, in the ordering the shorter runs already show. The gap between

\lambda=0

and

\lambda=0.25

is the largest and is consistent with a sharp phase transition in that interval, but the five weights are too coarsely spaced to confirm it.

	Chemistry	GSM8K
Setup
Trainer	self-distillation trainer of TRL 1.9.2 ( von Werra et al., 2020 )
Reference implementation	Hübotter et al. (2026)
Sharding, generation engine	DeepSpeed ZeRO-3, vLLM ( Kwon et al., 2023 )
Dataset	SciKnowEval chemistry	GSM8K
Student model	Olmo-3-7B-Instruct	Qwen3-0.6B

Appendix

Table 2: Default hyperparameters of the self-distillation experiments , on the Chemistry dataset (Figure 6 ) and on GSM8K (Figures 12 to 14 ). A value spanning both columns is shared by the two setups. Defaults follow the released implementation of Hübotter et al. (2026) , and each ablation moves a single field away from them. The learning rate and its warmup are the exception: Hübotter et al. (2026) report

10^{-6}

without warmup for the setup we replicate and

10^{-5}

with warmup for their other ones, and we use the latter, which reproduces the accuracy they report.

Kullback-Leibler Divergence

See all 27 papers

Mar 7, 2026

Entropy-Aware On-Policy Distillation of Language Models

On-policy distillation is a promising approach for transferring knowledge between language models, where a student learns from dense token-level signals along its own trajectories. This framework typically uses reverse KL divergence, encouraging the student to match the teacher's high-confidence predictions. However, we show that the mode-seeking property of reverse KL reduces generation diversity and yields unstable learning signals when the teacher distribution has high entropy. To address this, we introduce Entropy-Aware On-Policy Distillation. Our key idea is augmenting the standard reverse KL objective with forward KL when teacher entropy is high, capturing the full range of plausible outputs while retaining precise imitation elsewhere. It balances mode-seeking precision with mode-covering robustness without sacrificing on-policy training efficiency. Experiments show that our method maintains generation diversity (sustained token-level entropy) and improves student-teacher alignment (lower forward KL on high-entropy tokens). Across six math reasoning benchmarks, this yields Pass@8 accuracy gains of +1.37 for Qwen3-0.6B-Base, +2.39 for Qwen3-1.7B-Base, and +5.05 for Qwen3-4B-Base compared to baseline on-policy distillation methods. These results demonstrate that accounting for teacher uncertainty is essential for maintaining diversity and achieving effective knowledge transfer.

Woogyeol Jin, Taywon Min, Yongjin Yang +5

KAIST AI, Seoul, South Korea · University of Toronto, Ontario, Canada · Vector Institute, Ontario, Canada +1

Language Model Distillation On-Policy Distillation Uncertainty-Aware Knowledge Distillation LLM Reasoning

Ask PDF

Apr 15, 2026

TIP: Token Importance in On-Policy Distillation

On-policy knowledge distillation (OPD) trains a student on its own rollouts under token-level supervision from a teacher. Not all token positions matter equally, but existing views of token importance are incomplete. We ask a direct question: which tokens carry the most useful learning signal in OPD? Our answer is that informative tokens come from two regions: positions with high student entropy, and positions with low student entropy plus high teacher--student divergence, where the student is overconfident and wrong. Empirically, student entropy is a strong first-order proxy: retaining

50\%

of tokens with entropy-based sampling matches or exceeds all-token training while reducing peak memory by up to

47\%

. But entropy alone misses a second important region. When we isolate low-entropy, high-divergence tokens, training on fewer than

10\%

of all tokens nearly matches full-token baselines, showing that overconfident tokens carry dense corrective signal despite being nearly invisible to entropy-only rules. We organize these findings with TIP (Token Importance in on-Policy distillation), a two-axis taxonomy over student entropy and teacher--student divergence, and give a theoretical explanation for why entropy is useful yet structurally incomplete. This view motivates type-aware token selection rules that combine uncertainty and disagreement. We validate this picture across three teacher--student pairs spanning Qwen3, Llama, and Qwen2.5 on MATH-500 and AIME 2024/2025, and on the DeepPlanning benchmark for long-horizon agentic planning, where Q3-only training on

<$$20\%

of tokens surpasses full-token OPD. Our experiments are implemented by extending the OPD repository https://github.com/HJSang/OPSD_OnPolicyDistillation, which supports memory-efficient distillation of larger models under limited GPU budgets.

Yuanda Xu, Hejian Sang, Zhengze Zhou +3

Knowledge Distillation On-Policy Distillation Selective Knowledge Distillation Uncertainty-Aware Knowledge Distillation

Ask PDF

Aug 30, 2026

Influence-Directed Distillation: Solving the Diversity Bottleneck in Sampled-Token On-Policy Distillation

Sampled-token on-policy distillation (OPD) efficiently transfers capabilities from teacher to student using student-generated tokens, requiring teacher probabilities only for sampled tokens. Yet it frequently suffers from diversity distillation failure: the student's pass@1 improves while its pass@

k

plateaus, failing to inherit the teacher's diversity. To explain this, we introduce First-Order Local Entropy Influence, a signed first-order proxy that decouples each update's entropy effect into the teacher--student log-probability gap and the student's local probability structure, and empirically links entropy contraction to negative-influence positions. Motivated by this, we propose Influence-Directed Adaptive On-Policy Distillation (IDA-OPD): rather than relying on costly full-vocabulary Forward-KL objectives, it preserves entropy-expanding updates while replacing entropy-contracting ones with divergence-adaptive advantage shrinkage, using only the teacher's sampled-token log-probability. Experiments on reasoning-oriented distillation show IDA-OPD consistently improves pass@

k

, inheriting the teacher's diversity through distillation, matches the strongest teacher-informed methods at strictly lower cost, and broadly maintains vanilla OPD's pass@1, all without full-vocabulary teacher information.

Run Yang, Runpeng Dai, Jie Sun +5

University of North Carolina at Chapel Hill · University of Science and Technology of China · Shanghai University of Finance and Economics

Language Model Distillation On-Policy Distillation Knowledge Distillation

Ask PDF

Abstract

Kullback-Leibler Divergence

Explore similar work

Entropy-Aware On-Policy Distillation of Language Models

TIP: Token Importance in On-Policy Distillation

Influence-Directed Distillation: Solving the Diversity Bottleneck in Sampled-Token On-Policy Distillation