Distilling the Shape of Uncertainty: What On-Policy Distillation Learns and How to Predict It

Shiping Yang; Chuanqi Tang; Hossein Hajimirsadeghi; Layla El Asri; Mohammad Amin Shabani; Angel X. Chang

cs.CLOct 28, 2026

Distilling the Shape of Uncertainty: What On-Policy Distillation Learns and How to Predict It

Authors: Shiping Yang, Chuanqi Tang, Hossein Hajimirsadeghi, Layla El Asri, Mohammad Amin Shabani, Angel X. Chang

Organizations: Simon Fraser University · Borealis AI

Abstract

On-policy distillation (OPD) is increasingly adopted in the post-training of large language models. However, a stronger teacher may fail to improve the student. Prior work attributes this to a lack of new knowledge, but what that knowledge is remains unclear. We show that OPD learns not only how to branch, i.e., which candidate tokens to consider, but also where to branch, i.e., where to keep reasoning paths open. The existing token-overlap metric captures only the former. We introduce the entropy shape, the sequence of next-token entropies, to characterize the latter, and propose the Entropy Shape Correlation (ESC) to measure student-teacher agreement in entropy shape. Tracking ESC throughout training and intervening on a teacher's entropy while preserving its token preferences, we empirically show that where to branch causally drives OPD gains, independently of how to branch. Finally, we formulate OPD outcome prediction as a teacher selection task and establish an evaluation protocol for it, under which pre-OPD ESC selects the best teacher for both students and outperforms existing predictors. Code is available at https://github.com/maybenotime/opd-esc.

Figures & tables

Figure 1: Illustration of where and how to branch on a student rollout (student: R1-Distill-1.5B; teacher: JustRL-1.5B). Left : where to branch . Shaded positions mark mismatches, where the teacher tends to branch (high entropy) while the student does not (low entropy). Right : how to branch . At the first mismatch, the student follows the current path with "point", whereas the teacher spreads probability over "let" and "suppose", typical fork tokens that open alternative reasoning paths. Bold tokens appear among the top-16 candidates of both models. ESC and Overlap@k measure student-teacher agreement on where and how to branch, respectively.

Teacher	Training set	AIME24	AIME25	AMC	BeyondAIME	HMMT25	HMMT26	Average
Student: DeepSeek-R1-Distill-Qwen-1.5B (average before OPD: 25.5)
JustRL-1.5B	both-wrong	39.1	23.3	74.7	13.2	20.8	18.2	31.6 ( $+6.1$ )
JustRL-1.5B	both-correct	33.3	28.3	70.8	12.8	15.0	18.9	29.9 ( $+4.4$ )
R1-Distill-7B	both-wrong	27.5	25.0	60.2	8.8	15.8	15.9	25.5 ( $+0.0$ )
R1-Distill-7B	both-correct	25.8	20.0	62.0	8.8	17.5	15.2	24.9 ( $-0.6$ )
Student: Qwen3-1.7B (non-thinking) (average before OPD: 13.0)

Table 1: OPD results on correctness-controlled subsets. Each student is distilled from two teachers on the both-wrong and both-correct subsets (7k prompts each). Average is the mean avg@4 over the six benchmarks, followed by its absolute change relative to the student’s performance before OPD.

Figure 2: Training dynamics of successful and failing OPD runs. Top: R1-Distill-1.5B student. Successful runs are shown in color and failing runs in grey. The first three columns show ESC, Overlap@16, and

\Delta H

over training (10-step moving average). The last column shows the average performance on the six benchmarks, with successful runs evaluated every 50 steps and failing runs evaluated only at the beginning and end of training; dashed lines indicate teacher performance. Bottom: the same layout for the Qwen3-1.7B student.

Figure 3: OPD gains with original and counterfactual teachers. We study the effect of transplanting the entropy shape between a successful teacher

A

and a failing teacher

B

for two student models: Qwen3-1.7B (left) and R1-Distill-1.5B (right). We report the student’s average gain over the base model on the six benchmarks after 100 OPD steps for the two teacher models (

A

B

) and the transplanted models.

\widetilde{q}^{B\leftarrow A}

transplants

A

’s entropy shape onto

B

;

\widetilde{q}^{A\leftarrow B}

reverses the direction. Transplanting successful teacher’s entropy shape (

B\leftarrow A

) increases the performance gain, while transplanting failing teacher’s entropy shape (

A\leftarrow B

) lowers the performance gain. For R1-Distill-1.5B,

B

is the student itself and thus supplies no learning signal on its own.

	All ( $n{=}10$ )		$r_{s}$ per student		Regret per student
Predictor	$r_{s}$	95% CI	R1-Distill-1.5B	Qwen3-1.7B	R1-Distill-1.5B	Qwen3-1.7B
ESC (ours)	$+0.98$	$[+0.78,\,+1.00]$	$+1.00$	$+1.00$	$0.0$	$0.0$
Overlap@16	$+0.49$	$[-0.31,\,+0.95]$	$+0.30$	$+1.00$	$10.3$	$0.0$
$\Delta H$	$+0.20$	$[-0.54,\,+0.89]$	$+0.70$	$-0.10$	$0.0$	$6.5$
Teacher performance	$-0.01$	$[-0.56,\,+0.81]$	$-0.60$	$+1.00$	$7.5$	$0.0$
Teacher size	$-0.61$	$[-0.76,\,+0.25]$	$-0.60$	$-0.70$	$10.8$	$10.2$

Table 2: Performance of OPD predictors. The All columns pool all ten student–teacher pairings, whereas the

r_{s}

per student and Regret per student columns compute the corresponding metrics separately within each student’s five-teacher pool, better reflecting the practical setting of selecting a teacher for a given student.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Item	Value
Training temperature	1.0
Top- $p$	1.0
Global batch size	64
Mini batch size	64
Rollout number	1
Teacher LogProb top- $K$	64

Appendix

Table 3: Default hyperparameters for OPD.

Teacher	Hugging Face id	$\Delta$
Student: DeepSeek-R1-Distill-Qwen-1.5B
JustRL-1.5B	hbx/JustRL-DeepSeek-1.5B	$+8.1$
Nemotron-1.5B	nvidia/Nemotron-Research-Reasoning-Qwen-1.5B	$+11.0$
R1-Distill-7B	deepseek-ai/DeepSeek-R1-Distill-Qwen-7B	$+0.2$
Skywork-OR1-7B	Skywork/Skywork-OR1-Math-7B	$+3.5$
Light-R1-7B	qihoo360/Light-R1-7B-DS	$+0.7$

Appendix

Table 4: Teachers used in the prediction task and the six-benchmark average gain (

\Delta

) of the student.

Figure 4: Per-position entropy (

9

-token moving average) of DeepSeek-R1-Distill-1.5B and its JustRL-1.5B teacher on the student’s own rollout, before (left,

\rho=0.85

) and after OPD (right,

\rho=0.99

); shading marks the mismatch.

Figure 5: Pre-OPD ESC against the six-benchmark average gain for the ten student-teacher pairings in § 5 . Pairings with lower pre-OPD ESC leave more room for entropy shape alignment and yield larger gains.

Figure 6: Training dynamics of counterfactual teachers. Each row tracks one of the counterfactual runs of Figure 3 . Left: ESC between the student and the model that supplies the entropy shape, and Overlap@16 between the student and the model that provides the token preferences. Right: six-benchmark average accuracy at saved checkpoints, with dashed lines marking the two source models. (a) R1-Distill-1.5B distilled from

\widetilde{q}^{B\leftarrow A}

, with preferences from the student itself and shape from JustRL-1.5B. (b) Qwen3-1.7B distilled from

\widetilde{q}^{A\leftarrow B}

, with preferences from Qwen3-4B-RL-Math and shape from Qwen3-4B. Where and how to branch can thus be learned either jointly or independently.

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 522 papers

TIP: Token Importance in On-Policy Distillation
Apr 15, 2026Yuanda Xu, Hejian Sang, Zhengze Zhou +3Knowledge Distillation On-Policy Distillation
Rethinking On-Policy Distillation of Large Language Models: Phenomenology, Mechanism, and Recipe
Apr 14, 2026Yaxuan Li, Yuxin Zuo, Bingxiang He +8Language Model Distillation On-Policy Distillation
Tsinghua University · ShanghaiTech University · University of Illinois Urbana-Champaign +1
Not All Disagreement Is Learnable: Token Teachability in On-Policy Distillation
May 26, 2026Yuanyi Wang, Su Lu, Yanggan Gu +6Selective Knowledge Distillation Language Model Distillation
The Hong Kong Polytechnic University, PolyU · Hong Kong Polytechnic University Daya Bay Technology and Innovation Research Institute

Apr 15, 2026J/K move · Enter open · S save

TIP: Token Importance in On-Policy Distillation

Yuanda Xu, Hejian Sang, Zhengze Zhou, Ran He+2

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

Open paper AskPDF arXiv

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.

Abstract

Language Model Distillation

Explore similar work

TIP: Token Importance in On-Policy Distillation

Rethinking On-Policy Distillation of Large Language Models: Phenomenology, Mechanism, and Recipe

Not All Disagreement Is Learnable: Token Teachability in On-Policy Distillation