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

1 Introduction

• Prior work leaves unclear what new knowledge explains OPD gains and whether a teacher’s gain can be predicted before training.Earlier analyses attribute improved distillation to knowledge acquired through additional teacher post-training but diagnose OPD post hoc.Source: §1 · p. 1
• OPD gains are unaffected by whether the teacher solves training problems, suggesting the transferred behavior appears on both correct and incorrect trajectories.The authors therefore investigate teacher behavior at high-entropy fork tokens.Source: §1 · p. 1
• OPD teaches both where to keep reasoning paths open and which candidate tokens to consider, while entropy shape and ESC measure positional uncertainty and its alignment.Entropy shape is the sequence of next-token entropies along a trajectory; Overlap@k tracks candidate-token overlap.Source: §1 · pp. 1–2, Figure 1 · p. 2
• Counterfactual teachers show that a successful teacher’s entropy shape improves the student while preserving token preferences, and ESC can rise independently of Overlap@k.Training runs that improve show ESC rising with performance; failing runs start with high ESC and quickly plateau.Source: §1 · p. 2
• Pre-OPD ESC correctly ranks and selects the best teacher for both students, outperforming predictors based on teacher performance, size, and Overlap@k.The paper frames prediction as selecting the teacher expected to maximize a student’s downstream gain.Source: §1 · p. 2

2 Related work

• OPD research includes dense supervision, continual learning, and self-distillation, where a student acts as its own teacher using privileged information (Hübotter et al., 2026; Shenfeld et al., 2026; Zhao et al., 2026).The cited work reports on-policy dense supervision as more efficient than outcome-reward RL and promising for continual learning.Source: §2 · pp. 2–3
• Prior OPD analyses track progressive alignment of student and teacher top-k candidates, while this work identifies where to branch as a distinct learning signal (Li et al., 2026).MiniLLM optimizes reverse KL on student-generated rollouts, and GKD minimizes token-level divergences on mixed on- and off-policy data (Agarwal et al., 2024).Source: §2 · pp. 2–3
• Teacher–student capacity gaps can diminish or reverse distillation benefits, prompting teacher selection and teacher-construction approaches in off-policy settings (Cho and Hariharan, 2019; Mirzadeh et al., 2020).The passage notes that OPD shows a similar stronger-teacher failure, whose cause remains underexplored.Source: §2 · p. 3

3.1 Notation

• Notation distinguishes the student’s learnable policy from the frozen teacher and defines their next-token distributions over a shared vocabulary.Responses are indexed by token position, with y<t denoting the prefix preceding step t.Source: §3.1 · p. 3

3.2 On-policy distillation

• OPD minimizes token-level KL divergence on trajectories sampled from the student; sampled-token OPD evaluates only the student-sampled token as an unbiased estimator.The paper uses sampled-token OPD in its experiments and describes it as the most lightweight variant.Source: §3.2 · p. 3

3.3 Dynamic metrics

• Overlap@k measures shared top-k candidate tokens, while entropy gap measures mean-entropy difference but discards positional information.Overlap@k is computed as the average proportion of shared tokens; the paper uses k = 16.Source: §3.3 · p. 4

4 What does on-policy distillation learn?

• The paper investigates what OPD transfers beyond the knowledge gap between stronger and weaker teachers, testing whether it depends on correct supervision.Its analysis proceeds from correctness-controlled comparisons to entropy shape and then counterfactual teachers.Source: §4 · p. 4, §4.1 · p. 4

4.1 OPD does not rely on correct supervision

• Post-trained teachers remain effective on prompts they fail to solve, while same-recipe teachers remain ineffective even on prompts they solve.The both-correct and both-wrong subsets yield comparable outcomes, suggesting the teacher gap is not explained by training-data knowledge or supervision correctness.Source: §4.1 · p. 5

4.2 OPD learns where and how to branch

• Entropy shape—the sequence of next-token entropies—captures where reasoning paths remain open, complementing token overlap, which captures which candidates are considered.ESC measures the correlation between student and teacher entropy sequences; it is invariant to shifts and positive scaling, isolating shape alignment rather than absolute entropy levels.Source: §4.2 · p. 5
• In successful runs, ESC rises alongside performance, whereas it starts high and quickly plateaus in failing runs; entropy gap does not distinguish outcomes.These observations are correlational; the counterfactual-teacher experiments test whether entropy shape drives gains.Source: §4.2 · pp. 5–6

4.3 Causal analysis with counterfactual teachers

• Transplanting a successful teacher’s entropy shape improves OPD gains in both model families, while replacing it with a failing teacher’s shape reduces gains.The counterfactuals preserve token preferences, and training dynamics show entropy-shape alignment can change separately from token overlap; either signal alone can drive gains.Source: §4.3 · pp. 6–7, Figure 3 · p. 7

Task formulation

• Pre-OPD ESC measures student–teacher entropy-shape agreement to predict which pairs will gain more from OPD and select a teacher before training.The prediction is framed as closing the entropy-shape gap through OPD.Source: §5 · p. 7
• The task gives a student and candidate-teacher pool, scores each pair on a small pre-training prompt probe, and selects the teacher with the highest score.The evaluation suite uses ten student–teacher pairings: five candidate teachers for each of two students.Source: §5 · pp. 7–8

Results

• ESC has perfect rank correlation for both students (rs = 1.00) and selects the best teacher for each, yielding zero regret.Table 2 evaluates predictors across ten pairings and separately within each student's five-teacher pool.Source: Table 2 · p. 8, §5 · p. 8

Baselines

• Baselines include student–teacher measures Overlap@k and entropy gap ΔH, plus teacher-only benchmark performance and parameter count.Lower Overlap@k and larger ΔH are taken to predict larger gains; higher teacher performance and size are also taken to predict larger gains.Source: §5 · p. 8

Conclusion

• Where to branch causally drives OPD gains independently of how to branch, according to the paper's conclusion.The paper measures where-to-branch agreement with ESC.Source: §6 · p. 8
• OPD transfers two aspects of reasoning behavior: where to keep paths open and which candidate tokens to consider at each fork.The paper names these aspects where to branch and how to branch.Source: §6 · p. 8
• ESC enables teacher selection before training and could also monitor early stopping or help detect whether one model was distilled from another through OPD.The latter applications are presented as possibilities.Source: §6 · pp. 8–9

Limitations

• Experiments cover only sampled-token OPD on math reasoning, with 1.5B and 1.7B students and teachers up to 8B parameters.Whether the findings extend to other OPD variants, code generation, instruction following, or larger scales remains open; teacher-selection evaluation also has only ten pairings and wide confidence intervals.Source: Limitations · p. 9

Training and Evaluation Details

• Table 3 lists the default hyperparameters for OPD.Source: Table 3 · p. 11
• Evaluation uses vLLM with temperature 0.7, top-p 0.8, top-k 20, four samples per problem, and an unweighted mean over six competitions.Source: A Training and Evaluation Details · p. 11
• Maximum generation length is 16,384 tokens for R1-Distill-1.5B in thinking mode and 8,192 for Qwen3-1.7B in non-thinking mode.Source: A Training and Evaluation Details · p. 11

• Table 4 lists ten student–teacher pairings, teacher identifiers, and each student’s realized six-benchmark average gain at step 200.Source: B Teachers in the Evaluation Suite · p. 11, Table 4 · p. 11

C Entropy Shape Before and After OPD

• Figure 4 compares student and teacher per-position entropy shapes on a student rollout before and after OPD.Source: C Entropy Shape Before and After OPD · p. 11

D Pre-OPD ESC Against the Realized Gain

• Figure 5 plots each pairing’s pre-OPD ESC against its realized OPD gain.Source: D Pre-OPD ESC Against the Realized Gain · p. 11

• Figure 4 shows per-position entropy for DeepSeek-R1-Distill-1.5B and its JustRL-1.5B teacher before and after OPD, with mismatch shading.The reported correlations are ρ = 0.85 before OPD and ρ = 0.99 after OPD.Source: Figure 4 · p. 12
• Figure 5 relates pre-OPD ESC to six-benchmark average gain; pairings with lower ESC leave more room for entropy-shape alignment and yield larger gains.Source: Figure 5 · p. 12
• Figure 6 tracks two counterfactual runs: ESC and Overlap@16 appear on the left, while six-benchmark average accuracy appears on the right.Dashed lines mark the two source models; the runs transplant entropy shape while supplying token preferences from another model or the student itself.Source: E Training Dynamics of Counterfactual… · p. 11, Figure 6 · p. 13