Technical Summary: Position-Weighted On-Policy Self-Distillation for Reasoning

Problem Statement

On-policy self-distillation (OPSD) is a practical method for distilling mathematical reasoning, where a student model is trained on its own generated rollouts using a privileged teacher (a copy of the model conditioned on ground-truth answers) to provide token-level targets. However, the standard OPSD objective applies a uniform forward-KL supervision weight to all generated tokens. This implicitly assumes that the privileged teacher's target is equally reliable at every prefix visited by the student.

Recent adaptive distillation methods attempt to relax this uniformity by modulating supervision based on local signals such as teacher entropy, teacher-student probability ratios, or consensus. The authors argue that these local ambiguity signals are insufficient proxies for teacher reliability. High teacher entropy can indicate two distinct phenomena:

1. Non-viable uncertainty: The teacher assigns probability to multiple next tokens, but only one (or none) leads to a correct solution; the others are "dead ends."
2. Benign diversity: The teacher assigns probability to multiple viable solution routes, all of which preserve correctness.

Existing methods conflate these, potentially down-weighting supervision at benign branch points or failing to down-weight supervision at non-viable branch points where the teacher derails. The core problem is identifying a signal that distinguishes non-viable uncertainty from viable diversity to weight tokens appropriately.

Methodology

1. Branch-Viability Diagnostic

To empirically investigate teacher reliability, the authors introduce a branch-viability diagnostic.

• Setup: Using a single model (Qwen3-4B), they define a "teacher" prompt (containing the ground-truth answer) and a "student" prompt (ordinary problem).
• Process:
  1. The student generates a correct "spine" rollout.
  2. The teacher identifies high-ambiguity positions (high truncated entropy) along this spine.
  3. For these positions, the teacher proposes alternative next tokens.
  4. The model is forced to continue from the student prompt + spine prefix + the forced alternative, without privileged information.
  5. Labeling: If the forced continuation fails to recover the correct answer, the branch is labeled real-uncertain (low reliability). If it recovers the answer, it is labeled diversity (benign).
• Finding: The diagnostic reveals a strong positional structure. Real-uncertain branches concentrate early in the sequence, while viable diversity branches are more common later.
• Predictive Power: An oriented within-sequence position score (early positions = low reliability) achieves an AUROC of 0.83 in distinguishing real-uncertain from diversity candidates. In contrast, standard local uncertainty metrics (predictive entropy, MC-dropout mutual information, Dirichlet precision) perform near chance (AUROC $\le$ 0.57).

2. Position-Weighted On-Policy Self-Distillation (PW-OPSD)

Motivated by the finding that early high-ambiguity positions are structurally less reliable, the authors propose PW-OPSD.

• Core Mechanism: PW-OPSD retains the standard OPSD construction (student rollout, privileged teacher pass, per-vocabulary clipped forward-KL) but modifies the outer token aggregation.
• Weighting Schedule: Instead of uniform weights, it applies an increasing sigmoid function based on the normalized position within the sequence:
  $$w_{i,t} = w_{min} + (1 - w_{min})\sigma\left(\frac{r_{i,t} - \tau}{s}\right)$$
  Where $r_{i,t}$ is the normalized position.
• Design Rationale: The schedule assigns lower weights to early tokens (where the diagnostic indicates high unreliability) and higher weights to later tokens (where reliability is higher). A floor ($w_{min}$) ensures early tokens are not completely ignored.
• Reduction: The method averages token losses within each sequence before averaging across the batch, ensuring longer rollouts do not dominate the gradient purely due to token count.
• Efficiency: It requires no additional teacher passes, auxiliary verifiers, or complex sampling; the weight is a deterministic function of position.

Key Contributions

1. Diagnostic of Reliability: The paper introduces a branch-viability diagnostic that empirically demonstrates that teacher-token reliability in reasoning is trajectory-structured (dependent on position) rather than purely local (dependent on entropy).
2. Superiority of Position over Entropy: The authors show that position is a significantly stronger predictor of teacher reliability (AUROC 0.83) than local uncertainty scores (AUROC $\le$ 0.57), challenging the assumption that entropy-based adaptive weighting is optimal.
3. PW-OPSD Algorithm: A simple, plug-and-play modification to OPSD that applies position-dependent weights to the forward-KL loss, improving reasoning performance without additional computational overhead.
4. Generalization: The method is validated across different model families (Qwen, Llama, OLMo) and scales (4B to 8B parameters) without per-model hyperparameter tuning.

Experimental Results

Main Evaluation (Qwen3-4B)

Evaluated on MATH-500, AIME 2024, AIME 2025, and HMMT 2025 with a maximum generation length of 38,912 tokens:

• AIME 2024: PW-OPSD (Moderate) improved Avg@12 by +1.0 points over standard OPSD (76.20 vs. 75.19).
• AIME 2025: PW-OPSD improved Avg@12 by +1.1 points (67.78 vs. 66.67).
• HMMT 2025: An "Aggressive" schedule variant improved Avg@12 by +1.48 points over OPSD.
• Baselines: Adaptive baselines using entropy (EOPD) or policy-gradient rewards (REOPOLD) generally underperformed OPSD on AIME benchmarks, failing to capture the reliability pattern identified by the position diagnostic.
• MATH-500: Performance remained comparable to OPSD (near-saturated), indicating the method does not degrade performance on easier tasks.

Cross-Model Transfer

The "Moderate" schedule (derived from Qwen3-4B diagnostics) was applied without retuning to:

• DeepSeek-R1-Distill-Llama-8B: Improved Avg@12 by +0.35 points.
• Olmo-3-7B-Think: Improved Avg@12 by +0.50 points.
  This suggests the positional reliability prior is robust across model families and architectures.

Ablation Studies

• Schedule Sensitivity: Various schedule configurations (Mild, Sharp, Aggressive) all outperformed OPSD on average, confirming the robustness of the position-weighting concept.
• Reduction vs. Positioning: The authors found that the improvement relies on the joint application of position-weighting and per-sequence loss reduction; applying only one axis did not yield the full performance gain.

Significance and Claims

The paper claims that teacher-token reliability in reasoning distillation is trajectory-structured. The primary significance lies in shifting the focus from local ambiguity (entropy) to structural reliability (position) for adaptive distillation.

• Modest Claims: The authors explicitly state that while PW-OPSD improves performance, the gains are "modest" (e.g., +1.0 to +1.1 points on AIME). They acknowledge that PW-OPSD is a "plug-and-play" approximation and that there is room to design more sophisticated objectives (e.g., mixing forward and reverse KL based on position) in the future.
• No Additional Cost: A key claim is that these improvements are achieved without additional teacher computation or auxiliary verification steps, making it a highly efficient modification to existing distillation pipelines.
• Theoretical Insight: The work provides a theoretical interpretation of distillation as a reliability-weighted risk, where the goal is to weight tokens by the posterior probability that the teacher's target helps preserve the correct solution. The diagnostic provides the empirical evidence that position is a strong proxy for this posterior.

The code is available at https://github.com/SaFo-Lab/PW-OPSD.