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

cs.LGMar 7, 2026

Entropy-Aware On-Policy Distillation of Language Models

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

Organizations: KAIST AI, Seoul, South Korea · University of Toronto, Ontario, Canada · Vector Institute, Ontario, Canada · IBM Research, San Jose, CA, USA

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.

Figures & tables

Figure 1: Top-10 change rate for Scenario A (blue), where the teacher distribution has low entropy, and Scenario B (red), where the teacher distribution has high entropy across 3 seeds. For Scenario B, a student optimized with reverse KL fails to capture the teacher’s distribution, as evidenced by persistently high and fluctuating Top-10 change rates.

Teacher Entropy	Temp.	Top-1 Change Count
(A) Low	0.3	$7.3\pm 1.6$
(B) High	1.0	$84.0\pm 16.7$

Table 1: Top-1 Change Count for low teacher entropy scenario and a high teacher entropy scenario. We observe that when the teacher entropy is high, the top-1 index frequently changes.

\displaystyle\left(\mathbf{q},\mathbf{x},\{Q_{t}\}_{t=1}^{\mid\mathbf{x}\mid}\right)\rightarrow\;\mathcal{R}

Algorithm 1 Entropy-Aware On-Policy Distillation

Benchmark	Metric	Qwen3-0.6B-Base (MATH)				Qwen3-1.7B-Base (MATH)				Qwen3-4B-Base (DAPO-Math-14k)
Benchmark	Metric	KD	GRPO	OPD	EOPD (ours)	KD	GRPO	OPD	EOPD (ours)	KD	GRPO	OPD	EOPD (ours)
MATH500	Avg@8	47.80	51.83	50.09	52.02	63.80	68.83	67.76	68.73	74.73	80.47	78.81	80.20
MATH500	Pass@8	69.60	74.40	73.20	76.00	84.20	84.60	84.80	87.60	92.20	91.00	90.80	93.00
AMC23	Avg@8	23.43	25.25	24.69	27.81	38.11	40.62	39.06	41.88	48.13	58.44	57.33	60.94
AMC23	Pass@8	52.50	55.00	57.50	55.00	70.00	70.00	70.00	75.00	85.00	75.00	80.00	87.50
Minerva	Avg@8	18.61	20.08	19.98	20.82	28.14	29.46	29.83	30.15	34.65	40.81	40.08	39.71

Table 2: Main results (accuracy %) on six mathematical reasoning benchmarks. EOPD demonstrates consistent improvements across benchmarks and model scales. Parentheses indicate training data. Bold indicates best performance and underline indicates second-best.

Panel of Figure 2 — Figure 2: Pass@ $k$ performance comparison between OPD and EOPD on the AIME and AMC benchmarks with the Qwen3-1.7B-Base student. EOPD achieves higher Pass@ $k$ compared to OPD, with the performance gap becoming more pronounced as $k$ increases.

Benchmark	Metric	KD	GRPO	OPD	EOPD
GPQA-Diamond	Avg@8	27.01	27.86	30.08	31.50
GPQA-Diamond	Pass@8	75.20	77.83	77.78	81.31
MMLU-Pro	Pass@1	37.54	41.46	42.26	43.20
AlpacaEval 2.0	LC-WR	19.30	16.86	22.86	23.83
AlpacaEval 2.0	WR	23.63	20.55	29.92	29.54

Table 3: Out-of-domain benchmark results on Qwen3-1.7B-Base, covering general reasoning and instruction following. Our method achieves higher average accuracy and pass@8 on general reasoning, and is competitive or superior in win rate ( WR ) and length-controlled win rate ( LC-WR ). Best and second-best results are shown in bold and underline , respectively.

Figure 3: Token-level entropy histograms comparing the Qwen3-8B teacher with Qwen3-1.7B-Base trained using OPD and EOPD on the AIME 2024 and 2025 benchmarks. While both methods exhibit similar distributions to the teacher in the mid-entropy range, EOPD preserves more probability mass in the high-entropy region, staying closer to the teacher than OPD.

Figure 4: Average policy entropy during training for the Qwen3-1.7B-Base student trained with OPD + Entropy Bonus, OPD + Advantage Shaping, and EOPD. Advantage Shaping converges to a lower entropy regime, while Entropy Bonus maintains entropy levels comparable to EOPD.

Figure 5: Average forward KL divergence measured during training at token positions where the teacher distribution exhibits high entropy (entropy

\geq 0.8

) for the Qwen3-1.7B-Base student. Compared to OPD + Entropy Bonus and OPD + Advantage Shaping, EOPD maintains lower forward KL values throughout training, indicating closer alignment with the teacher distribution in regions of high uncertainty.

Benchmark	Metric	Entropy	Advantage	EOPD
Benchmark	Metric	Bonus	Shaping	EOPD
MATH500	Avg@8	66.87	67.90	68.73
MATH500	Pass@8	86.00	85.20	87.60
AMC23	Avg@8	39.69	36.56	41.88
AMC23	Pass@8	75.00	75.00	75.00
AIME24	Avg@8	9.58	8.75	10.42

Table 4: Comparison with entropy-driven exploration baselines for the Qwen3-1.7B-Base student. EOPD achieves higher Avg@8 and Pass@8 compared to OPD + Entropy Bonus and OPD + Advantage Shaping. Bold indicates best performance and underline indicates second-best.

Benchmark	Metric	GKD	GKD	GKD	EOPD
Benchmark	Metric	( $\lambda=0.5$ )	( $\lambda=0.8$ )	( $\lambda=1.0$ )	EOPD
MATH500	Avg@8	64.53	65.28	67.20	68.73
MATH500	Pass@8	84.40	84.20	85.00	87.60
AMC23	Avg@8	35.62	37.50	38.44	41.88
AMC23	Pass@8	67.50	70.00	67.50	75.00
AIME24	Avg@8	5.83	6.67	9.17	10.42

Table 5: Comparison with GKD baselines for the Qwen3-1.7B-Base student. We evaluate GKD with different

\lambda

values following the original setup. EOPD consistently outperforms all GKD variants across reasoning benchmarks. Bold indicates best performance and underline indicates second-best.

Benchmark	Metric	OPD-FKL	OPD	EOPD
MATH500	Avg@8	67.20	67.76	68.73
MATH500	Pass@8	85.00	84.80	87.60
AMC23	Avg@8	38.44	39.06	41.88
AMC23	Pass@8	67.50	70.00	75.00
AIME24	Avg@8	9.17	8.33	10.42
AIME24	Pass@8	23.33	20.00	23.33

Table 6: Comparison of different KL objectives for the Qwen3-1.7B-Base student. EOPD outperforms both OPD and OPD-FKL across most benchmarks, while OPD generally outperforms OPD-FKL. Bold indicates the best performance and underline indicates second-best.

Benchmark	Metric	Full	Random	EOPD
MATH500	Avg@8	67.58	67.33	68.73
MATH500	Pass@8	86.20	84.80	87.60
AMC23	Avg@8	39.39	36.88	41.88
AMC23	Pass@8	75.00	67.50	75.00
AIME24	Avg@8	8.75	7.92	10.42
AIME24	Pass@8	23.33	20.00	23.33

Table 7: Comparison with different forward KL placement strategies for the Qwen3-1.7B-Base student. EOPD shows competitive performance against full and random forward KL, while random forward KL underperforms the other two approaches. Bold indicates the best performance and underline indicates second-best.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	Value
Learning rate	1e-5
LR scheduler type	cosine
Optimizer	AdamW
CE loss weight	0.5
KL loss weight	0.5
Training Batch size	128

Appendix

Table 8: Hyperparameters used for Off-policy distillation.

Hyperparameter	OPD, EOPD	GRPO
Learning rate	3e-6	3e-6
LR scheduler type	cosine	cosine
Optimizer	AdamW	AdamW
Training batch size	128	128
Mini batch size	32	32
Samples per prompt	1	8

Appendix

Table 9: Hyperparameters used for On-policy distillation and GRPO.

Panel of Figure 6 — Figure 6: Teacher–student distributions under low and high-entropy scenarios. When the teacher distribution has low entropy, the student model accurately converges to the teacher. In contrast, when the teacher distribution has high entropy, optimization with reverse KL based reward leads the student to align with a few dominant indices of the teacher distribution, while the global structure is not fully approximated.

Panel of Figure 7 — Figure 7: Pass@ $k$ performance comparison between OPD and EOPD on the AIME and AMC benchmarks with the Qwen3-0.6B-Base student. EOPD achieves higher Pass@ $k$ compared to OPD.

Panel of Figure 8 — Figure 8: Pass@ $k$ performance comparison between OPD and EOPD on the AIME and AMC benchmarks with the Qwen3-4B-Base student. EOPD achieves higher Pass@ $k$ compared to OPD, with the performance gap becoming more pronounced as $k$ increases.

Figure 9: Ratio of tokens with high teacher entropy in rollouts produced by the student model during EOPD training.

Figure 10: Tradeoff between cumulative probability mass and memory usage under top-

k

truncation. Memory increases linearly with

k

, while the cumulative probability mass quickly saturates.

Benchmark	Metric	KD	GRPO	OPD	EOPD
MATH500	Avg@8	42.60	45.01	43.60	44.90
MATH500	Pass@8	65.80	66.00	64.80	67.20
AMC23	Avg@8	18.43	20.31	19.06	22.18
AMC23	Pass@8	52.50	50.00	47.50	52.50
AIME24	Avg@8	1.66	2.91	2.08	2.91
AIME24	Pass@8	13.33	10.00	10.00	13.33

Appendix

Table 10: Comparison with other baselines on mathematical reasoning benchmarks using a Llama series model. EOPD achieves competitive performance across benchmarks. Bold indicates the best performance and underline indicates second-best.

Figure 11: Average forward KL divergence measured during training at token positions where the teacher distribution exhibits high entropy (entropy

\geq 0.8

) for the Llama-3.2-3B-Instruct student.

Benchmark	Metric	BD-KD	AKL	ToDi	EOPD
MATH500	Avg@8	67.45	66.43	67.43	68.73
MATH500	Pass@8	86.80	85.20	86.80	87.60
AMC23	Avg@8	38.75	37.81	38.75	41.88
AMC23	Pass@8	72.50	70.00	75.00	75.00
AIME24	Avg@8	9.17	7.50	8.33	10.42
AIME24	Pass@8	20.00	16.67	20.00	23.33

Appendix

Table 11: Comparison with adaptive KL-based baselines. Bold indicates the best performance and underline indicates the second-best performance.

$\tau$	MATH500		AMC23
$\tau$	Avg@8	Pass@8	Avg@8	Pass@8
0.6	68.24	86.80	39.69	75.00
0.8	68.73	87.60	41.88	75.00
1.0	67.58	84.00	41.53	72.50
1.2	64.82	84.80	39.69	75.00
1.4	67.61	83.80	38.72	72.50

Appendix

Table 12: Performance variation with respect to the entropy threshold

\tau

for the Qwen3-1.7B-Base student. EOPD is not highly sensitive to

\tau

, while we observe an overall trend where Pass@8 performance decreases as

\tau

increases.

Benchmark	Metric	$\boldsymbol{\alpha=0.6}$	$\boldsymbol{\alpha=1.0}$	$\boldsymbol{\alpha=1.2}$	$\boldsymbol{\alpha=1.5}$
MATH500	Avg@8	67.37	68.73	69.55	67.50
MATH500	Pass@8	86.20	87.60	87.40	86.00
AMC23	Avg@8	38.12	41.88	41.56	37.81
AMC23	Pass@8	72.50	75.00	75.00	67.50
AIME24	Avg@8	8.75	10.42	10.83	9.17
AIME24	Pass@8	16.67	23.33	23.33	23.33

Appendix

Table 13: Impact of the forward KL coefficient

\alpha

for the Qwen3-1.7B-Base student. Bold indicates the best performance and underline indicates second-best.

Benchmark	Metric	$\boldsymbol{k=1}$	$\boldsymbol{k=8}$	$\boldsymbol{k=16}$	$\boldsymbol{k=128}$
MATH500	Avg@8	67.72	67.15	68.73	65.21
MATH500	Pass@8	84.20	85.80	87.60	87.40
AMC23	Avg@8	38.19	38.25	41.88	39.42
AMC23	Pass@8	65.00	72.50	75.00	75.00
AIME24	Avg@8	9.17	10.42	10.42	8.75
AIME24	Pass@8	16.67	16.67	23.33	23.33

Appendix

Table 14: Impact of top-k selection for the Qwen3-1.7B-Base student. Bold indicates the best performance and underline indicates second-best.

Language Model Distillation

Language model distillation transfers a larger model’s capabilities to smaller or specialized models using outputs, probabilities, or internal representations.

See all 508 papers

Divergence-Routed On-Policy Distillation for Language Models
Oct 14, 2026Tiantian Peng, Guannan Lv, Hongyang Li +5On-Policy Distillation Knowledge Distillation
Peking University · Kuaishou Technology · Tsinghua University
Tail-Aware Top- $k$ On-Policy Distillation
Aug 12, 2026Huipeng Huang, Hongxin WeiLanguage Model Distillation On-Policy Distillation
Department of Statistics and Data Science, Southern University of Science and Technology
Influence-Directed Distillation: Solving the Diversity Bottleneck in Sampled-Token On-Policy Distillation
Aug 30, 2026Run Yang, Runpeng Dai, Jie Sun +5Language Model Distillation On-Policy Distillation
University of North Carolina at Chapel Hill · University of Science and Technology of China · Shanghai University of Finance and Economics

Oct 14, 2026J/K move · Enter open · S save

Divergence-Routed On-Policy Distillation for Language Models

Tiantian Peng, Guannan Lv, Hongyang Li, Qiuhe Hong+4

Peking University · Kuaishou Technology · Tsinghua University

On-Policy Distillation Knowledge Distillation Selective Knowledge Distillation Kullback-Leibler Divergence

Open paper AskPDF arXiv

On-policy distillation (OPD) transfers knowledge from a teacher model to a student along student-generated trajectories, typically by minimizing the reverse Kullback-Leibler (RKL) divergence. However, RKL can provide weak corrective signals for teacher-supported tokens to which the student assigns very low probability. This raises two questions: which positions are most consequential for the behavioral gap between the student and teacher, and what additional supervision is needed at these positions? Through cross-sampling, we find that divergence-based routing more effectively pinpoints positions contributing to the teacher-student performance gap than entropy-based routing. Further analysis reveals teacher-mode under-coverage at high Jensen-Shannon (JS) divergence positions, where standard OPD increases the large teacher-student probability gaps. Motivated by these findings, we propose Divergence-Routed On-Policy Distillation (DrOPD). DrOPD retains RKL supervision at all positions and selectively adds forward Kullback-Leibler (FKL) divergence supervision at positions with high JS divergence. Experiments across six math benchmarks show that DrOPD improves over standard OPD and entropy-based routing. DrOPD also reduces severe under-coverage cases by 37.9 percent relative to OPD, supporting targeted FKL as an effective complement to OPD.

Abstract

Language Model Distillation

Explore similar work

Divergence-Routed On-Policy Distillation for Language Models

Tail-Aware Top-kkk On-Policy Distillation

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

Tail-Aware Top- $k$ On-Policy Distillation