Entropy-Aware On-Policy Distillation of Language Models

Dennis Wei; Kimin Lee; Nathalie Baracaldo; Swanand Ravindra Kadhe; Taywon Min; Woogyeol Jin; Yi Zhou; Yongjin Yang

arxiv: 2603.07079 · v2 · pith:66MAX7UCnew · submitted 2026-03-07 · 💻 cs.LG · cs.CL

Entropy-Aware On-Policy Distillation of Language Models

Woogyeol Jin , Taywon Min , Yongjin Yang , Swanand Ravindra Kadhe , Yi Zhou , Dennis Wei , Nathalie Baracaldo , Kimin Lee This is my paper

Pith reviewed 2026-05-25 07:02 UTC · model grok-4.3

classification 💻 cs.LG cs.CL

keywords on-policy distillationentropy-aware distillationreverse KL divergenceforward KL divergencelanguage model distillationmath reasoning benchmarksgeneration diversityknowledge transfer

0 comments

The pith

Augmenting reverse KL with forward KL on high-entropy tokens improves math reasoning accuracy while preserving generation diversity in on-policy distillation.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

Standard on-policy distillation uses reverse KL divergence to make the student match the teacher's confident predictions, but this mode-seeking behavior reduces output variety and produces unstable signals when the teacher itself is uncertain. The paper introduces an entropy-aware variant that adds forward KL divergence precisely on tokens where teacher entropy is high, allowing the student to cover multiple plausible answers in uncertain cases while keeping tight imitation elsewhere. This hybrid loss is applied along the student's own trajectories without changing the on-policy efficiency. Across six math benchmarks the approach produces Pass@8 gains that increase with model size, from 1.37 points on the 0.6B base to 5.05 points on the 4B base, relative to plain reverse-KL distillation. The result matters because it shows that explicitly accounting for teacher uncertainty can improve both alignment and diversity at once.

Core claim

The central claim is that the mode-seeking property of reverse KL reduces generation diversity and yields unstable signals when teacher entropy is high, and that augmenting the objective with forward KL on those high-entropy tokens captures the full range of plausible outputs while retaining precise imitation on low-entropy tokens, thereby maintaining sustained token-level entropy and lowering forward KL on high-entropy tokens.

What carries the argument

The entropy-aware objective that augments standard reverse KL with forward KL when teacher entropy on the student's on-policy trajectories exceeds an implicit threshold.

If this is right

Generation diversity measured by token-level entropy is sustained rather than collapsed.
Student-teacher alignment improves specifically on high-entropy tokens as shown by lower forward KL.
Pass@8 accuracy rises by +1.37, +2.39 and +5.05 points on the 0.6B, 1.7B and 4B Qwen3 base models across six math benchmarks.
Knowledge transfer remains efficient because the method stays fully on-policy.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

The same entropy-triggered switch could be tested on non-math tasks such as code generation where teacher uncertainty also varies by token.
If the gains scale with model size, the technique may become more valuable precisely when distilling the largest teachers.
The method implicitly assumes entropy is a sufficient statistic for uncertainty; alternative uncertainty measures such as mutual information could be substituted in follow-up work.

Load-bearing premise

Teacher entropy can be computed reliably from the student's trajectories and the simple reverse-plus-forward augmentation will not introduce instabilities or require new hyperparameters that explain the observed gains.

What would settle it

A controlled run in which the entropy-aware switch is applied but forward KL on high-entropy tokens does not decrease or token-level entropy does not remain higher than the plain reverse-KL baseline.

read the original abstract

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.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Desk Editor's Note private letter to a colleague

The entropy-triggered mix of forward and reverse KL is a targeted fix for a real weakness in on-policy distillation, but the gains rest on an unablated rule whose details are missing from the abstract.

read the letter

The paper's core move is to add forward KL only on high-entropy tokens from the teacher while keeping reverse KL elsewhere, all inside an on-policy loop. That produces the reported Pass@8 lifts of roughly 1.4 to 5 points on the six math benchmarks for the Qwen3 base models, plus better sustained token entropy. The idea is not a routine extension; it directly tackles the mode-seeking behavior that hurts diversity when the teacher is uncertain. That part is useful and worth noting for anyone running distillation on reasoning tasks. The experiments also track both accuracy and a diversity metric, which is better than many distillation papers that report only the former. The numbers scale with model size in the right direction, which adds some credibility. The main gap is exactly what the stress-test flags: the abstract gives no rule for the entropy threshold, no mixture weight, no per-token versus sequence decision, and no ablations showing that these choices were not tuned separately from the baselines. Without those, it is hard to know whether the gains come from the entropy awareness itself or from extra hyperparameter effort. There are also no error bars or implementation specifics on how teacher entropy is estimated on student trajectories. This is the sort of paper that belongs in a distillation or LLM training reading group. Practitioners who already run on-policy loops would get immediate value from trying the switch if the full paper supplies the missing rule and checks. It deserves a serious referee because the underlying problem is real, the empirical direction is positive, and the method is simple enough to test. Referees will just need to press on the hyperparameter question before the claim can be taken as settled.

Referee Report

4 major / 2 minor

Summary. The paper proposes Entropy-Aware On-Policy Distillation, which augments the standard reverse KL objective used in on-policy distillation with forward KL divergence on high-entropy tokens from the teacher. The method is motivated by the claim that reverse KL alone reduces diversity and yields unstable signals under high teacher entropy. Experiments on six math reasoning benchmarks report Pass@8 gains of +1.37, +2.39, and +5.05 for Qwen3-0.6B-Base, Qwen3-1.7B-Base, and Qwen3-4B-Base relative to baseline on-policy distillation, while claiming to maintain generation diversity via sustained token-level entropy and improved alignment on high-entropy tokens.

Significance. If the central empirical claim holds after clarification of implementation details, the work would offer a targeted, efficiency-preserving modification to on-policy distillation that addresses a known limitation of reverse KL in uncertain regions of the teacher distribution. This could be relevant for improving robustness in knowledge transfer for reasoning tasks without requiring off-policy sampling or additional model components.

major comments (4)

[Method section] Method section: The precise decision rule (entropy threshold value, per-token vs. sequence-level application, and mixing weight for the forward KL term) is not specified or derived. This is load-bearing for the central claim, as the reported Pass@8 gains cannot be attributed to the augmentation idea rather than unablated hyperparameter choices.
[Experiments section] Experiments section: No error bars, standard deviations across runs, or statistical significance tests are reported for the Pass@8 accuracies on the six benchmarks. This prevents assessment of whether the +1.37/+2.39/+5.05 gains are reliable or could arise from training variance.
[Experiments section] Experiments section: No ablation is provided on the entropy threshold, the mixing weight, or the on-policy restriction of the forward KL term. Without these, it is impossible to isolate whether the gains stem from the proposed entropy-aware mechanism or from other implementation factors.
[Abstract and Experiments] Abstract and Experiments: The claim that teacher entropy is computed reliably on the student's on-policy trajectories and that the augmentation does not introduce new instabilities is stated but not verified experimentally (e.g., via training curves or divergence measurements before/after the switch).

minor comments (2)

[Abstract] The abstract mentions 'sustained token-level entropy' and 'lower forward KL on high-entropy tokens' as supporting metrics but does not report their numerical values or how they were measured.
[Method section] Notation for the combined objective (reverse KL + forward KL) should be formalized with an explicit equation showing the condition and weighting, even if the exact threshold is left as a hyperparameter.

Simulated Author's Rebuttal

4 responses · 0 unresolved

We thank the referee for the constructive feedback. We address each major comment below and will revise the manuscript to improve clarity and add requested analyses where possible.

read point-by-point responses

Referee: [Method section] The precise decision rule (entropy threshold value, per-token vs. sequence-level application, and mixing weight for the forward KL term) is not specified or derived. This is load-bearing for the central claim, as the reported Pass@8 gains cannot be attributed to the augmentation idea rather than unablated hyperparameter choices.

Authors: We agree that these details are critical for reproducibility and attribution. The revised manuscript will explicitly document the entropy threshold (chosen via validation-set statistics), confirm per-token application, and state the mixing weight for the forward KL term, along with a short justification for the selected values. revision: yes
Referee: [Experiments section] No error bars, standard deviations across runs, or statistical significance tests are reported for the Pass@8 accuracies on the six benchmarks. This prevents assessment of whether the +1.37/+2.39/+5.05 gains are reliable or could arise from training variance.

Authors: We acknowledge the absence of variance estimates. Given the computational expense of full multi-seed training, we will add a limitations paragraph discussing observed variance from smaller-scale runs and note that the reported gains should be interpreted with this caveat. Additional seeds will be included for the 0.6B model if resources allow. revision: partial
Referee: [Experiments section] No ablation is provided on the entropy threshold, the mixing weight, or the on-policy restriction of the forward KL term. Without these, it is impossible to isolate whether the gains stem from the proposed entropy-aware mechanism or from other implementation factors.

Authors: We will add ablations on the entropy threshold and mixing weight to the experiments section or appendix. The on-policy restriction is central to the efficiency claim; we will strengthen the text with a brief theoretical argument for why off-policy forward KL would alter the setting, while noting that a full off-policy comparison lies outside the current scope. revision: partial
Referee: [Abstract and Experiments] The claim that teacher entropy is computed reliably on the student's on-policy trajectories and that the augmentation does not introduce new instabilities is stated but not verified experimentally (e.g., via training curves or divergence measurements before/after the switch).

Authors: We will include training curves and per-epoch measurements of token-level entropy and forward KL on high-entropy tokens in the revised appendix to empirically verify stable computation on on-policy trajectories and absence of new instabilities. revision: yes

Circularity Check

0 steps flagged

No significant circularity; empirical method with independent experimental validation

full rationale

The paper describes an algorithmic change to on-policy distillation (augmenting reverse KL with forward KL conditioned on teacher entropy) and reports empirical gains on math benchmarks. No equations, derivations, or self-citations are presented that reduce the claimed improvements to fitted inputs, self-definitions, or prior author results by construction. The central claims rest on experimental comparisons rather than a closed mathematical chain, making the work self-contained against external benchmarks.

Axiom & Free-Parameter Ledger

0 free parameters · 0 axioms · 0 invented entities

Only the abstract is available; no explicit free parameters, axioms, or invented entities are stated. The method implicitly relies on an entropy threshold or mixing coefficient whose value is not reported.

pith-pipeline@v0.9.0 · 5784 in / 1180 out tokens · 23118 ms · 2026-05-25T07:02:47.390082+00:00 · methodology

discussion (0)

Lean theorems connected to this paper

Citations machine-checked in the Pith Canon. Every link opens the source theorem in the public Lean library.

IndisputableMonolith/Cost/FunctionalEquation.lean washburn_uniqueness_aczel unclear

?

unclear
Relation between the paper passage and the cited Recognition theorem.

augmenting the standard reverse KL objective with forward KL when teacher entropy is high... hyperparameter τ controls this transition
IndisputableMonolith/Foundation/BranchSelection.lean branch_selection unclear

?

unclear
Relation between the paper passage and the cited Recognition theorem.

LEOPD_t(θ;ct)=LOPD_t(θ;ct)+I[Hte_t>τ]LFKL_t(θ;ct)

What do these tags mean?

matches: The paper's claim is directly supported by a theorem in the formal canon.
supports: The theorem supports part of the paper's argument, but the paper may add assumptions or extra steps.
extends: The paper goes beyond the formal theorem; the theorem is a base layer rather than the whole result.
uses: The paper appears to rely on the theorem as machinery.
contradicts: The paper's claim conflicts with a theorem or certificate in the canon.
unclear: Pith found a possible connection, but the passage is too broad, indirect, or ambiguous to say the theorem truly supports the claim.

Forward citations

Cited by 27 Pith papers

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score.

Visual-Advantage On-Policy Distillation for Vision-Language Models
cs.CV 2026-05 unverdicted novelty 7.0

VA-OPD improves VLM performance over standard on-policy distillation by reweighting rollouts and separating KL terms according to token-level visual advantage on math and visual benchmarks.
Decoupling KL and Trajectories: A Unified Perspective for SFT, DAgger, Offline RL, and OPD in LLM Distillation
cs.LG 2026-05 unverdicted novelty 7.0

Decoupling prefix source from token-level KL direction in autoregressive sequence KL yields four objectives unifying SFT, DAgger, offline RL and OPD, with KL mixing and entropy-gated curriculum improving math reasonin...
Respecting Self-Uncertainty in On-Policy Self-Distillation for Efficient LLM Reasoning
cs.AI 2026-05 unverdicted novelty 7.0

EGRSD and CL-EGRSD advance the accuracy-length frontier in LLM reasoning by entropy-guided weighting of token-level distillation signals from the teacher.
Multi-Rollout On-Policy Distillation via Peer Successes and Failures
cs.LG 2026-05 unverdicted novelty 7.0

MOPD improves on-policy distillation for LLMs by using peer successes for positive patterns and failures for negative examples to create more informative teacher signals.
The Extrapolation Cliff in On-Policy Distillation of Near-Deterministic Structured Outputs
cs.LG 2026-05 unverdicted novelty 7.0

On-policy distillation has an extrapolation cliff at closed-form lambda*(p,b,c) set by teacher modal probability, warm-start mass, and clip strength, past which training shifts from format-preserving to format-collapsing.
Flow-OPD: On-Policy Distillation for Flow Matching Models
cs.CV 2026-05 conditional novelty 7.0

Flow-OPD applies on-policy distillation to flow matching models via specialized teachers, cold-start initialization, and manifold anchor regularization, lifting GenEval from 63 to 92 and OCR from 59 to 94 on Stable Di...
KL for a KL: On-Policy Distillation with Control Variate Baseline
cs.LG 2026-05 unverdicted novelty 7.0

vOPD stabilizes on-policy distillation gradients by subtracting a closed-form per-token negative reverse KL baseline as a detached control variate, preserving unbiasedness while lowering variance and matching expensiv...
Rubric-based On-policy Distillation
cs.LG 2026-05 unverdicted novelty 7.0

Rubric-based on-policy distillation allows training student models using only teacher responses by generating scoring rubrics from contrasts and using them for on-policy optimization, achieving superior performance an...
MAD-OPD: Breaking the Ceiling in On-Policy Distillation via Multi-Agent Debate
cs.CL 2026-05 unverdicted novelty 7.0

MAD-OPD recasts on-policy distillation teachers as a debating collective to supply better supervision, lifting agentic and code performance over single-teacher OPD across multiple model sizes.
TCOD: Exploring Temporal Curriculum in On-Policy Distillation for Multi-turn Autonomous Agents
cs.LG 2026-04 unverdicted novelty 7.0

TCOD stabilizes on-policy distillation for multi-turn agents via temporal curriculum on trajectory depth, improving performance up to 18 points over vanilla OPD and sometimes surpassing the teacher.
Tailoring Teaching to Aptitude: Direction-Adaptive Self-Distillation for LLM Reasoning
cs.LG 2026-05 unverdicted novelty 6.0

DASD improves math reasoning in LLMs by adaptively directing self-distillation based on per-token entropy to balance exploration and step accuracy, outperforming prior self-distillation and RLVR baselines on six benchmarks.
On-Policy Consistency Training Improves LLM Safety with Minimal Capability Degradation
cs.LG 2026-05 conditional novelty 6.0

On-Policy Consistency Training (OPCT) improves LLM safety metrics over supervised fine-tuning while largely preserving capabilities across three model families.
When Are Teacher Tokens Reliable? Position-Weighted On-Policy Self-Distillation for Reasoning
cs.LG 2026-05 unverdicted novelty 6.0

Position-Weighted On-Policy Self-Distillation (PW-OPSD) weights later tokens more heavily after a diagnostic shows position predicts teacher reliability better than entropy, yielding +1.0 and +1.1 Avg@12 gains on AIME...
Prefix Teach, Suffix Fade: Local Teachability Collapse in Strong-to-Weak On-Policy Distillation
cs.CL 2026-05 unverdicted novelty 6.0

Local teachability collapse in trajectory suffixes makes uniform dense supervision suboptimal in strong-to-weak OPD; truncating at BIC-style change points on teacher margin improves performance.
Flow-OPD: On-Policy Distillation for Flow Matching Models
cs.CV 2026-05 unverdicted novelty 6.0

Flow-OPD applies on-policy distillation to flow-matching text-to-image models, lifting GenEval from 63 to 92 and OCR accuracy from 59 to 94 while preserving fidelity.
Flow-OPD: On-Policy Distillation for Flow Matching Models
cs.CV 2026-05 unverdicted novelty 6.0

Flow-OPD applies on-policy distillation to flow matching models, achieving GenEval of 92 and OCR accuracy of 94 on Stable Diffusion 3.5 Medium while avoiding the seesaw effect of multi-reward optimization.
Flow-OPD: On-Policy Distillation for Flow Matching Models
cs.CV 2026-05 unverdicted novelty 6.0

Flow-OPD applies on-policy distillation to Flow Matching models through specialized teachers, cold-start initialization, task routing, and manifold regularization, lifting GenEval from 63 to 92 and OCR from 59 to 94 o...
SOD: Step-wise On-policy Distillation for Small Language Model Agents
cs.CL 2026-05 unverdicted novelty 6.0

SOD reweights on-policy distillation strength step-by-step using divergence to stabilize tool use in small language model agents, yielding up to 20.86% gains and 26.13% on AIME 2025 for a 0.6B model.
SimCT: Recovering Lost Supervision for Cross-Tokenizer On-Policy Distillation
cs.CL 2026-05 unverdicted novelty 6.0

SimCT enlarges the supervision space in cross-tokenizer on-policy distillation using short jointly tokenizable multi-token continuations, producing consistent gains over shared-token baselines on math and code benchmarks.
SimCT: Recovering Lost Supervision for Cross-Tokenizer On-Policy Distillation
cs.CL 2026-05 unverdicted novelty 6.0

SimCT recovers discarded teacher signal in cross-tokenizer on-policy distillation by enlarging supervision to jointly realizable multi-token continuations, yielding consistent gains on math reasoning and code generati...
UniSD: Towards a Unified Self-Distillation Framework for Large Language Models
cs.CL 2026-05 unverdicted novelty 6.0

UniSD unifies self-distillation components for autoregressive LLMs and its full integrated version improves base models by 5.4 points and baselines by 2.8 points across six benchmarks.
Uni-OPD: Unifying On-Policy Distillation with a Dual-Perspective Recipe
cs.LG 2026-05 unverdicted novelty 6.0

Uni-OPD unifies on-policy distillation across LLMs and MLLMs with dual-perspective strategies that promote student exploration and enforce order-consistent teacher supervision based on outcome rewards.
Rethinking On-Policy Distillation of Large Language Models: Phenomenology, Mechanism, and Recipe
cs.LG 2026-04 unverdicted novelty 6.0

On-policy distillation works when student and teacher models share thinking patterns and the teacher adds new capabilities, with success tied to alignment on a small set of high-probability tokens.
$\boldsymbol{f}$-OPD: Stabilizing Long-Horizon On-Policy Distillation with Freshness-Aware Control
cs.LG 2026-05 unverdicted novelty 5.0

f-OPD decomposes on-policy distillation drift into rollout and supervision components, then applies a sample-level freshness score to adaptively limit stale data influence and stabilize long-horizon agent training.
On-Policy Distillation with Best-of-N Teacher Rollout Selection
cs.CV 2026-05 unverdicted novelty 5.0

BRTS improves on-policy distillation by sampling multiple teacher rollouts and selecting the best one via a correctness-first then alignment priority rule, yielding gains on AIME and AMC math benchmarks.
On-Policy Distillation with Best-of-N Teacher Rollout Selection
cs.CV 2026-05 unverdicted novelty 5.0

BRTS improves on-policy distillation by selecting the highest-quality teacher trajectory from a small pool of samples based on correctness and alignment with the student, yielding gains on AIME and AMC math benchmarks.
UniSD: Towards a Unified Self-Distillation Framework for Large Language Models
cs.CL 2026-05 unverdicted novelty 5.0

UniSD unifies complementary self-distillation mechanisms for autoregressive LLMs and achieves up to +5.4 point gains over base models and +2.8 over baselines across six benchmarks and six models.