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

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

Apr 14, 2026

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

On-policy distillation (OPD) has become a core technique in the post-training of large language models, yet its training dynamics remain poorly understood. This paper provides a systematic investigation of OPD dynamics and mechanisms. We first identify that two conditions govern whether OPD succeeds or fails: (i) the student and teacher should share compatible thinking patterns; and (ii) even with consistent thinking patterns and higher scores, the teacher must offer genuinely new capabilities beyond what the student has seen during training. We validate these findings through weak-to-strong reverse distillation, showing that same-family 1.5B and 7B teachers are distributionally indistinguishable from the student's perspective. Probing into the token-level mechanism, we show that successful OPD is characterized by progressive alignment on high-probability tokens at student-visited states, a small shared token set that concentrates most of the probability mass (97%-99%). We further propose two practical strategies to recover failing OPD: off-policy cold start and teacher-aligned prompt selection. Finally, we show that OPD's apparent free lunch of dense token-level reward comes at a cost, raising the question of whether OPD can scale to long-horizon distillation.

Yaxuan Li, Yuxin Zuo, Bingxiang He +8

Tsinghua University · ShanghaiTech University · University of Illinois Urbana-Champaign +1

Language Model Distillation On-Policy Distillation Knowledge Distillation Teacher-Student Learning

Ask PDF

May 26, 2026

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

On-policy distillation (OPD) trains a student on its own rollouts with token-level teacher supervision. Recent selective OPD methods exploit the non-uniformity of OPD signals by prioritizing high-entropy or high-disagreement tokens. We revisit this principle and ask: which token-level teacher signals are actually learnable? Using a fixed-context diagnostic that measures same-context teacher-student KL reduction, we show that raw KL disagreement is a coarse proxy for learning value. It conflates learnable disagreement, where the teacher assigns corrective mass to the student's top-K candidates, with incompatible disagreement, where the teacher places mass mostly off the student's current support. We formalize this local compatibility as token teachability and show that it better predicts fixed-context improvement than raw KL alone. Motivated by this finding, we propose Teachability-Aware OPD (TA-OPD), a lightweight token-position selection method that applies OPD loss to high-teachability positions without reward models or verifiers. Across Qwen2.5 and Qwen 3 teacher-student settings, TA-OPD often surpasses full-token OPD with only 5% retained tokens and improves over entropy- and divergence-based baselines. Our results reframe selective OPD as selecting learnable teacher signals rather than merely salient tokens.

Yuanyi Wang, Su Lu, Yanggan Gu +6

The Hong Kong Polytechnic University, PolyU · Hong Kong Polytechnic University Daya Bay Technology and Innovation Research Institute

Selective Knowledge Distillation Language Model Distillation On-Policy Distillation Teacher-Student Learning Knowledge Distillation

Ask PDF

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