PaperlayerView exact revision on arXiv

Source-linked AI summary

Entropy-Aware On-Policy Distillation of Language Models

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

arXiv:2603.07079v3cs.LGcs.CL

TL;DR

Reverse KL-based on-policy distillation can collapse diversity and destabilize learning when teacher uncertainty is high. EOPD selectively applies forward KL to high-entropy teacher tokens while retaining reverse KL elsewhere, preserving diversity and improving results across six mathematical reasoning benchmarks.

  • Problem

    Reverse KL can reduce generation diversity and destabilize learning at token positions where the teacher distribution has high entropy.

  • Method

    EOPD selectively applies reverse KL in low-entropy regions and forward KL in high-entropy regions during on-policy distillation.

  • Results

    EOPD yields consistent gains over standard on-policy distillation across six mathematical reasoning benchmarks.

  • Takeaways & Limitations

    Explicitly modeling teacher uncertainty is essential for stable, diverse, and effective knowledge transfer.

Abstract

from arXiv · show

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.

1. Introduction

On-policy distillation avoids the distribution mismatch of off-policy training but reverse KL can reduce diversity and destabilize learning under teacher uncertainty. EOPD addresses this by switching between reverse and forward KL, improving diversity and reasoning performance.

  • Off-policy distillation trains on teacher data, creating a mismatch with sequences generated by the student at inference time.
  • On-policy distillation instead corrects samples generated by the student, typically using reverse KL and integrating with reinforcement-learning pipelines.
  • Reverse KL reduces student diversity and produces unstable learning signals when teacher distributions have high entropy.
  • Standard on-policy distillation retains 6.8% of high-entropy tokens, compared with 18.5% in the teacher.
  • EOPD applies reverse KL in low-entropy regions and forward KL in high-entropy regions to preserve diversity without naive forward-KL overhead.
  • Pass@8 improves by +1.37, +2.39, and +5.05 for Qwen3-0.6B-Base, Qwen3-1.7B-Base, and Qwen3-4B-Base, respectively, across six math benchmarks.

2. Preliminaries

The preliminaries define KL divergence and contrast forward and reverse KL through their reference distributions and optimization behavior. They then describe on-policy distillation, its reverse-KL objective, and PPO-style stabilization.

  • 2.1. KL-Based Divergences: KL divergence is a non-symmetric measure of how well Q approximates reference distribution P, so reversing its arguments changes optimization behavior.
  • 2.1. KL-Based Divergences: For language models, token distributions are conditioned on context, with student πθ(· | c) and teacher πte(· | c) defining the distillation distributions.
  • 2.1. KL-Based Divergences: Forward KL penalizes the student for assigning low probability to teacher-likely tokens and encourages mode-covering behavior.
  • 2.1. KL-Based Divergences: Reverse KL averages over student samples, penalizes teacher-unlikely generated tokens, and ignores teacher modes the student does not visit.
  • 2.2. On-Policy Distillation: On-policy distillation trains the student against teacher token probabilities on the student's own generated sequences, mitigating compounding errors from off-policy imitation.
  • 2.2. On-Policy Distillation: The OPD objective uses student-generated sequences and token contexts to optimize an on-policy reverse-KL loss through policy-gradient-style updates.
  • 2.3. OPD with Clipped-Reverse KL: PPO-style importance sampling and clipping stabilize OPD by correcting behavior-policy sampling and limiting overly large policy updates.

3. Diversity Degradation and Instability in On-Policy Distillation

Standard on-policy distillation loses diversity and becomes unstable when teacher uncertainty is high. Controlled experiments show reverse KL converges under low entropy but fails to stabilize top predictions under high entropy.

  • High-entropy tokens encode multiple plausible reasoning paths and meaningful uncertainty rather than mere noise.
  • Token-Level Entropy Analysis: Reverse KL favors dominant modes, reducing exploration and hindering transfer of the teacher’s uncertainty.
  • Token-Level Entropy Analysis: 6.8% versus 18.5%: the on-policy-distilled student retains fewer high-entropy tokens than the teacher.The comparison uses an entropy threshold of ≥1.0 on AIME24 and AIME25 prompts.
  • Instability of Reverse KL-based Reward: Under low teacher entropy, top-10 change rates decrease steadily and top-1 changes are rare.
  • Instability of Reverse KL-based Reward: Under high teacher entropy, top-1 indices change frequently and top-10 predictions fail to converge.The high-entropy scenario shows persistent, fluctuating top-10 change rates across three seeds.

4. Entropy-Aware On-Policy Distillation

EOPD adapts the on-policy distillation objective to the teacher’s token-level entropy. It uses reverse KL for confident predictions and forward KL for uncertain regions, preserving diversity while retaining efficiency.

  • EOPD combines reverse and forward KL because they provide complementary efficiency and distribution-covering properties.Reverse KL targets confident predictions, while forward KL transfers uncertainty and global structure.
  • EOPD selectively applies forward KL according to the teacher’s token-level uncertainty.
  • In low-entropy regions, EOPD uses reverse KL for efficiency and fast convergence.
  • In high-entropy regions, EOPD uses forward KL to preserve probability mass across plausible continuations and prevent mode collapse.
  • EOPD approximates forward KL over the teacher’s top-k tokens instead of sampling, limiting tail coverage and computational cost.The paper reports k = 16 as a practical balance between cumulative probability mass and memory cost.
  • EOPD adapts to local uncertainty to balance training stability with diversity preservation across confident and ambiguous regions.

5. Experiments

Experiments evaluate EOPD across mathematical, out-of-domain, entropy-transfer, and objective-placement settings. EOPD improves reasoning performance while better preserving and matching teacher uncertainty, especially in high-entropy regions.

  • Main Results: EOPD improves Avg@8 by +1.16 and Pass@8 by +1.37 for Qwen3-0.6B-Base across six mathematical reasoning benchmarks.
  • Main Results: EOPD improves Avg@8 by +0.99 and Pass@8 by +2.39 for Qwen3-1.7B-Base, and by +1.80 and +5.05 for Qwen3-4B-Base.
  • Main Results: EOPD achieves higher Pass@k than OPD on AIME and AMC, with larger gaps as the number of sampled rollouts increases.Pass@k measures the probability of obtaining at least one correct solution among k sampled rollouts.
  • Out-of-domain Evaluation: On out-of-domain benchmarks, OPD and EOPD outperform KD and GRPO, while EOPD surpasses OPD except on AlpacaEval 2.0 win rate.
  • Token-Level Entropy Analysis: EOPD retains substantially more probability mass at entropy ≥1.0 than OPD and remains closer to the teacher, while mid-entropy distributions are similar.
  • Entropy-Driven Baselines: EOPD achieves lower forward KL than entropy bonus and advantage shaping at teacher-entropy threshold τ = 0.8, indicating closer alignment in uncertain regions.Entropy Bonus maintains entropy comparable to EOPD, so entropy preservation alone does not explain EOPD’s advantage.
  • Objective Comparisons: EOPD consistently outperforms tuned GKD and benefits from selectively applying forward KL in high-uncertainty regions rather than placing it randomly.

6. Related Work

Related work covers knowledge distillation, on-policy distillation, reinforcement-learning approaches, and divergence-adaptation strategies. EOPD specifically targets reverse-KL instability and diversity degradation in high-entropy on-policy regions.

  • Knowledge Distillation: Knowledge distillation transfers larger-model output distributions through token-level matching, teacher-generated sequences, attention scores, or alternative divergence objectives.
  • Adaptive Divergences: Prior work adapts divergence objectives using entropy gaps, head-tail mismatch, or vocabulary-level discrepancy measures.
  • EOPD Positioning: EOPD differs by selectively applying forward KL to high-uncertainty positions within on-policy distillation, addressing reverse-KL instability and diversity degradation.
  • Reasoning Models: Language-model reasoning research includes prompting, test-time scaling, distillation, verifiable-reward reinforcement learning, and process-level supervision.

7. Conclusion

The paper identifies reverse KL’s failure to preserve diversity and distributional structure under teacher uncertainty and introduces EOPD to address it. Across six mathematical reasoning benchmarks, EOPD improves knowledge transfer while retaining on-policy efficiency.

  • EOPD selectively applies forward KL to high-entropy teacher tokens while adapting the on-policy distillation objective.This complements reverse KL’s mode-seeking behavior with forward KL to better preserve uncertain distributions.
  • EOPD better preserves the teacher-like distribution while retaining on-policy training efficiency.
  • EOPD yields consistent gains over standard on-policy distillation across six mathematical reasoning benchmarks.
  • Explicitly modeling teacher uncertainty is essential for stable, diverse, and effective knowledge transfer.

Impact Statement

The work targets the transfer of reasoning capabilities to smaller, more efficient language models by addressing reverse KL’s loss of diversity and distributional structure in high-uncertainty regions.

  • The framework adaptively applies KL objectives based on the teacher’s token-level uncertainty.
  • Reverse-KL distillation can fail to preserve distributional structure and diversity in high-uncertainty regions critical for reasoning.

A. Implementation Details

The implementation uses separate off-policy and on-policy distillation setups, with on-policy training implemented in verl alongside GRPO. EOPD adds limited computation and substantially less memory than applying forward KL at every token position.

  • Off-policy Training: Off-policy distillation samples teacher responses and trains students with cross-entropy and forward KL losses.
  • On-policy Training: On-policy distillation and GRPO use verl, with batch size B = 128, mini-batch size Bmini = 32, and one rollout per problem for distillation.GRPO instead generates eight rollouts per problem for relative trajectory comparison.
  • Computational Overhead and Efficiency Analysis: EOPD requires no additional teacher forward passes because it queries the teacher once per student-generated token position, as OPD does.
  • Computational Overhead and Efficiency Analysis: 2.16 seconds of added time per training step corresponds to approximately 4.5% of the average 47.8-second step time.The overhead comes mainly from top-k extraction, renormalization, and forward KL computation.
  • Computational Overhead and Efficiency Analysis: Selective forward KL adds 219.3 MB of memory, compared with 1617.5 MB when applied to all token positions.Selective entropy-aware routing is approximately 7.37× more memory-efficient than full forward KL application.
  • Baselines and Auxiliary Methods: Entropy Bonus adds policy entropy to the optimization objective, while Advantage Shaping modifies advantages with an entropy term.The described Advantage Shaping setup uses α = 0.1 and κ = 2.

B. Toy Experiment Visualizations

The toy visualizations show close student–teacher convergence under low teacher entropy but concentration on dominant indices under high entropy with reverse-KL-based optimization. The evaluation setup uses specified sampling and prompting protocols across math and out-of-domain benchmarks.

  • Toy Experiment Visualizations: Under low teacher entropy, the student converges closely to the teacher; under high entropy, reverse-KL optimization concentrates on a few dominant indices.The high-entropy case does not fully approximate the teacher’s global structure.
  • Evaluation Details: The math benchmark evaluation reports average@k and pass@k using zero-shot sampling with maximum sequence length 8192, temperature 1.0, and top-p 0.8.
  • Evaluation Details: MMLU-Pro uses the default 5-shot setting, while AlpacaEval follows its default protocol and reports win rate and length-controlled metrics.
  • Evaluation Details: The multiple-choice prompt instructs models to reason step by step and end with the correct answer letter.

D. Pass@k Experiments

EOPD improves Pass@k performance over OPD across Qwen3 students, with larger gains on harder benchmarks and increasing advantages as k grows. Its training analysis examines high-entropy tokens, efficient forward-KL computation, and comparisons with other baselines.

  • EOPD achieves higher Pass@k than OPD for Qwen3-0.6B-Base on AIME and AMC benchmarks.
  • EOPD achieves higher Pass@k than OPD for Qwen3-4B-Base, with the performance gap becoming more pronounced as k increases.
  • EOPD outperforms OPD on AIME24/25 and AMC23, with wider gaps on harder AIME benchmarks as k increases.The passage attributes this pattern to more effective exploration of diverse reasoning trajectories.
  • After convergence, high-entropy tokens stabilize at approximately 15% to 20%, motivating random forward KL on 20% of positions as a comparison setting.High-entropy tokens are identified using teacher entropy computed on student-generated rollouts.
  • Top-k forward-KL approximation evaluates cumulative teacher probability mass against memory usage, which increases linearly with k while probability mass quickly saturates.This analysis addresses the memory overhead of computing forward KL over the full teacher vocabulary.
  • EOPD applies forward KL only at necessary positions, reducing additional computational overhead relative to other adaptive KL baselines.

I. Additional Ablations

Additional ablations show that EOPD remains stable across entropy thresholds, forward-KL coefficients, and top-k choices, with best reported settings near α = 1.0 ∼1.2 and k = 16.

  • EOPD is not highly sensitive to the entropy threshold τ, although Pass@8 performance generally decreases as τ increases.Lower τ activates forward KL at more tokens, whereas higher τ restricts it to the most uncertain positions.
  • EOPD performs stably across forward-KL coefficient α values, with best performance achieved around α = 1.0 ∼1.2.Smaller α may under-transfer uncertainty, while excessively large α can weaken reverse KL’s dominant-mode capture.
  • EOPD remains relatively stable across top-k values, with best performance achieved at k = 16.Small k may miss sufficient probability mass, while large k includes low-probability tails and raises computational cost.
Loading 2603.07079v3…