PaperlayerView exact revision on arXiv

Source-linked AI summary

When Are Teacher Tokens Reliable? Position-Weighted On-Policy Self-Distillation for Reasoning

Xiaogeng Liu, Xinyan Wang, Yingzi Ma, Yechao Zhang, Chaowei Xiao

arXiv:2605.21606v1cs.LGcs.AI

TL;DR

The paper asks whether locally ambiguous teacher targets are genuinely unreliable in on-policy reasoning distillation. It introduces a branch-viability diagnostic and position-weighted OPSD, finding stronger positional reliability prediction than local uncertainty and consistent benchmark gains without additional teacher computation.

  • Problem

    Uniform OPSD weighting treats privileged teacher targets as equally useful, while local ambiguity can represent either unreliable branches or benign solution diversity.

  • Method

    The paper tests teacher-token reliability by forcing teacher alternatives through student-template continuations, then applies an increasing position weight in PW-OPSD while retaining the OPSD training components.

  • Results

    PW-OPSD improves Avg@12 over OPSD by +1.0 pp on AIME 2024 and +1.1 pp on AIME 2025, with consistent aggregate improvements on DeepSeek-R1-Distill-Llama-8B and Olmo-3-7B-Think.

  • Takeaways & Limitations

    Teacher-token reliability in reasoning distillation is trajectory-structured and can be used through position weighting without additional teacher computation.

Abstract

from arXiv · show

On-policy self-distillation (OPSD) trains a student on its own rollouts using a privileged teacher, but its standard objective weights all generated tokens equally, implicitly treating the privileged teacher target as equally reliable at every student-visited prefix. Existing entropy-based OPD methods relax this uniformity by modulating token-level supervision with teacher entropy, but high teacher entropy in reasoning has an ambiguous reliability meaning: it can reflect either non-viable uncertainty or benign solution diversity. To identify this phenomenon, we introduce a branch-viability diagnostic. Specifically, we record next-token alternatives from the privileged-answer teacher prompt, force each alternative after the student prompt plus its on-policy spine prefix, and test whether the resulting student-template continuation recovers the correct answer. On Qwen3-4B, we find that an oriented within-sequence position score is the strongest tested predictor of teacher-token reliability, reaching an area-under-ROC-curve (AUROC) of 0.83; local uncertainty scores are at most 0.57. Motivated by this trajectory-level structure, we propose Position-Weighted On-Policy Self-Distillation (PW-OPSD), which applies an increasing position weight while keeping the same student rollout, privileged teacher pass, and clipped forward-KL target as OPSD. In our comprehensive evaluations with different random seeds, the diagnostic-derived PW-OPSD improves AIME 2024 and AIME 2025 Avg@12 by +1.0 and +1.1 points, and a generalization evaluation on two larger-scale models from different families, DeepSeek-R1-Distill-Llama-8B and Olmo-3-7B-Think, also demonstrates consistent aggregate Avg@12 improvements. These results show that teacher-token reliability in reasoning distillation is trajectory-structured and can be utilized without additional teacher computation.

1 Introduction

OPSD trains students on their own visited prefixes with privileged teacher targets, but uniform token weighting overlooks whether those targets are reliable. A branch-viability diagnostic reveals a positional reliability pattern, motivating PW-OPSD, which improves difficult reasoning benchmarks without extra teacher computation.

  • On-policy self-distillation: OPSD supervises student-visited prefixes with token-level targets from a privileged teacher, avoiding trajectory mismatch while injecting privileged information.The student samples from the ordinary problem prompt, while the privileged model provides targets along that same rollout.
  • Reliability problem: Uniform forward-KL supervision treats the privileged teacher target as equally useful at every student-visited prefix.This motivates estimating reliability rather than relying solely on uniform token weighting.
  • Branch-viability diagnostic: The branch-viability diagnostic forces teacher-proposed alternatives into student-template contexts and tests whether continuations recover the correct answer.Teacher-template continuation is avoided because privileged information could make almost any forced branch recover successfully.
  • Diagnostic finding: 0.83 AUROC was achieved by oriented position scoring, while tested local uncertainty diagnostics reached residual AUROC at most 0.57.Early high-ambiguity branch points were more often unreliable, whereas later positions were more reliable under forced continuation.
  • PW-OPSD: PW-OPSD applies an increasing sigmoid position weight while retaining the OPSD rollout, privileged teacher pass, and clipped forward-KL surrogate.The method changes only the outer reliability-weighted aggregation and adds no extra teacher pass or auxiliary verifier.
  • Evaluation: +1.0 point on AIME 2024 Avg@12 and +1.1 points on AIME 2025 Avg@12 were obtained by diagnostic-derived PW-OPSD with the same teacher cost.The method matched OPSD on MATH-500 Avg@12 at 95.34 versus 95.33, while the schedule sweep gave HMMT 2025 Avg@12 a +1.48-point gain.

2 Related Work

Related work develops on-policy distillation and adaptive token-level objectives using alternative divergences, teacher uncertainty, teacher-student discrepancy, rewards, sampling, or consensus. This paper’s distinction is to question whether local ambiguity reliably identifies useful supervision.

  • On-policy distillation: On-policy distillation trains students on prefixes they actually visit, extending knowledge distillation toward learner-visited states and reducing off-policy mismatch concerns.The related methods include GKD, MiniLLM, and DistiLLM, which use generalized on-policy distillation, sequence-level reverse KL, or skewed KL objectives.
  • Adaptive token weighting: Adaptive distillation methods vary divergences or token controllers using teacher entropy, teacher-student ratios, clipped rewards, entropy-guided sampling, or tutor consensus.These approaches depart from uniform token supervision for different reasons and with different control signals.
  • Paper’s distinction: The paper distinguishes local ambiguity from teacher reliability, arguing that entropy, mismatch, and consensus signals can reweight tokens for the wrong reason.High local ambiguity may correspond either to unreliable branches or to benign solution diversity.

3 Method

PW-OPSD estimates teacher-target reliability with a branch-viability diagnostic and uses within-sequence position as a structural weighting signal. It preserves OPSD’s student rollout, privileged teacher pass, and clipped forward-KL target while emphasizing later tokens.

  • PW-OPSD: PW-OPSD preserves OPSD’s student rollout, privileged teacher pass, and per-vocabulary clipped forward-KL surrogate while changing outer token aggregation.The method applies position-dependent reliability weighting and per-sequence averaging.
  • Branch-viability diagnostic: The branch-viability diagnostic forces teacher-proposed alternatives after the student prompt and rollout spine, then tests whether student-template continuations recover the correct answer.Candidates are selected from high teacher-side ambiguity and classified by whether forced continuations remain viable.
  • Branch-viability diagnostic: 8 real-uncertain and 271 diversity candidates across 61 problems are compared using within-problem residualized AUROC with cluster-bootstrap confidence intervals.The evaluation is restricted to correct-spine, high-truncated-entropy candidate positions.
  • Position signal: AUROC 0.83 is achieved by the oriented early-position score, while tested local uncertainty diagnostics reach at most 0.57.The score is 1 − e_r, so larger values correspond to earlier tokens; raw position is negatively associated with teacher unreliability.
  • Position signal: Early high-ambiguity positions are often teacher-unreliable, whereas later positions more often reflect viable diversity.This finding motivates lower weights for early tokens and increasing reliability weights across the sequence.
  • Reliability formulation: The formal reliability target is the posterior probability that matching a privileged teacher target helps preserve or recover the correct solution.PW-OPSD uses normalized position as a low-cost structural proxy for this posterior.

4 Experiments

Experiments compare diagnostic-derived position weighting with OPSD and adaptive baselines on Qwen3-4B, then test schedule sensitivity, reduction choices, and cross-model transfer. PW-OPSD improves selected reasoning benchmarks while retaining near-saturated MATH-500 performance and positive aggregate gains across three model families.

  • Main Qwen3-4B results: PW-OPSD Moderate improves Avg@12 over OPSD by +1.0 pp on AIME 2024 and +1.1 pp on AIME 2025.The comparison uses the same maximum generation length of 38,912 tokens and three evaluation seeds per cell.
  • Main Qwen3-4B results: On MATH-500, OPSD, EOPD, and PW-OPSD Moderate are tied within 0.01 pp Avg@12, while PW-OPSD Aggressive gives the highest MATH-500 Avg@12 and Maj@12 in the table.
  • Main Qwen3-4B results: EOPD trails OPSD by 1.6 pp and 1.0 pp Avg@12 on AIME 2024 and AIME 2025, while REOPOLD trails by 1.2 pp and 4.5 pp.
  • Schedule sensitivity: Aggressive PW-OPSD raises HMMT 2025 Avg@12 to 45.37 pp, a +1.48 pp gain over OPSD, while its Maj@12 remains +1.11 pp above OPSD.
  • Schedule sensitivity: The four schedules vary early-token discount strength, with lower wmin, larger τ, and smaller s producing stronger, later, and sharper down-weighting respectively.
  • Schedule sensitivity: Across the four-benchmark average, every schedule improves Avg@12 over OPSD’s 70.27 by +0.30 to +0.65 pp, with only 0.35 pp separating schedules.
  • Reduction-positioning ablation: The joint position-weighting and per-sequence-mean configuration is the only factorial variant matching the AIME 2024 lead of +1.0 pp Avg@12 over OPSD.
  • Cross-model transfer: With no per-model schedule retuning, Moderate yields positive aggregate Avg@12 gains of +0.39 pp on Qwen3-4B, +0.35 pp on DeepSeek-R1-Distill-Llama-8B, and +0.50 pp on Olmo-3-7B-Think.

5 Conclusion, Limitations, and Future Work

The paper concludes that teacher-token reliability is positionally structured and can motivate position-weighted OPSD. It also acknowledges that PW-OPSD’s gains remain modest and proposes more sophisticated position-conditioned objectives.

  • Conclusion: The branch-viability diagnostic finds positional reliability structure, with AUROC 0.83 for oriented position versus at most 0.57 for local uncertainty.
  • Conclusion: PW-OPSD improves Avg@12 over OPSD on AIME 2024 by +1.0 pp and AIME 2025 by +1.1 pp, with additional aggregate gains on two larger models.
  • Limitations and future work: PW-OPSD’s gains over OPSD remain modest, leaving room for objectives beyond pure position weighting.
  • Limitations and future work: The proposed future direction is position-conditioned mixing of forward and reverse KL.

A Branch-viability protocol details

The branch-viability protocol uses one checkpoint with separate privileged-teacher and ordinary-student templates, labels forced alternatives by continuation viability, and evaluates residualized predictors with problem-clustered bootstrap.

  • Setup: The diagnostic applies teacher and student prompt templates to the same Qwen3-4B checkpoint, with the teacher receiving the privileged ground-truth answer.
  • Problem sampling: Phase A samples 84 problems, Phase B retains 62 correct-spine problems, and the binary-labeled pool contains 61 usable problems after excluding one gray-only problem.
  • Scoring and uncertainty: Raw uncertainty scores are mean-centered within each problem before AUROC computation, and cluster bootstrap resamples problems rather than candidates.
  • Forced continuations: For each high-truncated-entropy alternative, the protocol forces it after the student prompt and spine prefix, then continues generation under the student template.
  • Candidate labeling: Candidates are labeled diversity when at least two children have viability ≥Vhigh=0.75 and real-uncertain when every child has viability < Vlow=0.40 with mean viability below Vlow.

B MC-dropout diagnostic implementation

The diagnostic implementation combines branch-mixture analysis with MC-dropout uncertainty measurements, while PW-OPSD uses a position-based reliability proxy rather than these local uncertainty scores. The position schedule preserves the clipped forward-KL framework and increases supervision toward later sequence positions.

  • MC-dropout diagnostic: MC dropout is diagnostic-only: it computes Hfull, MI, and log κ̂ scores but does not construct training targets or token weights.Training uses the ordinary single-pass privileged-teacher target.
  • MC-dropout diagnostic: Inference-time MC dropout applies p = 0.1 to the last L = 4 transformer layers across M = 5 samples.Forward hooks inject dropout because the transformer blocks do not expose nn.Dropout submodules.
  • Numerical stability: The κ̂ estimator can become negative when sample variance exceeds categorical variance, so raw estimates are retained and log-scored with ϵ = 10^-6.The negative values are attributed to a small-M artifact rather than a training objective.
  • Branch-mixture interpretation: Branch-mixture identities show that conditional mutual-information terms quantify branch ambiguity while remaining independent of the student distribution.Therefore, the identity alone does not establish a different marginal forward-KL gradient objective.
  • Position-weighted objective: PW-OPSD retains the OPSD forward-KL inner loss and uses an increasing within-sequence position schedule as its reliability proxy.Its Moderate schedule is (wmin, τ, s) = (0.25, 0.30, 0.10), with a nonzero early-token floor.

G PW-OPSD training pseudocode

PW-OPSD samples student rollouts, scores visited prefixes under student and privileged-teacher contexts, and applies position-weighted clipped forward-KL losses. The sampled rollout remains fixed during each update, and losses are reduced within sequences before averaging across sequences.

  • Training step: PW-OPSD computes unreduced clipped forward-KL terms, sums over vocabulary, applies position weights, and averages over valid tokens and sequences.This reduction prevents longer rollouts from receiving more gradient mass solely because they contain more tokens.
  • Training step: The training step samples student rollouts and scores each valid visited prefix using privileged-teacher and ordinary-student contexts.The valid-token count ends at EOS or truncation.
  • Position weighting: Each token receives ri,t = (t − 0.5)/Li and wi,t = wmin + (1 − wmin)σ((ri,t − τ)/s).The normalized position and sigmoid schedule implement increasing within-sequence weighting.

H Implementation conventions inherited from OPSD

The evaluation compares methods under shared implementation conventions inherited from OPSD, including padding, prompt-template differences, clipping, and fixed optimization settings. Rollout-seed behavior is a documented source of variation across reruns.

  • Data and prompt handling: Right-padded prompts and batch-max completion slicing are held constant across methods despite variable prompt lengths.The OpenThoughts-Math-30k prompts have median length 93 tokens and maximum length 826.
  • Data and prompt handling: Training and evaluation use different prompt templates, and this gap is shared across all reported methods.Training disables thinking, whereas evaluation enables thinking.
  • Loss conventions: The clipped forward-KL implementation clamps each vocabulary entry before summing over the vocabulary, matching the OPSD reference implementation.Clipping is applied element-wise to qt(j) log(qt(j)/pt(j)).
  • Loss conventions: Gradient accumulation uses the mean of two per-microbatch token-means for OPSD, while PW-OPSD uses per-sequence reduction without this issue.Each PW-OPSD microbatch contains the same number of sequences.
  • Reproducibility: Reruns with the same trainer seed can produce slightly different rollouts because the vLLM colocate-mode seed is independent of the trainer --seed flag.Evaluation instead uses explicitly seeded vLLM SamplingParams.seed values.

I Evaluation setup

The study evaluates PW-OPSD against uniform, entropy-conditioned, and policy-gradient baselines across math benchmarks and model families under a common on-policy protocol. Results use repeated sampling metrics and multiple evaluation seeds.

  • Models: The main comparison uses Qwen3-4B, while cross-model evidence adds DeepSeek-R1-Distill-Llama-8B and Olmo-3-7B-Think.The latter are larger models from different families.
  • Baselines: PW-OPSD is compared with OPSD, EOPD, and REOPOLD under a common training and evaluation protocol.OPSD is the uniform-weight reference, EOPD is entropy-conditioned, and REOPOLD is policy-gradient adaptive distillation.
  • Evaluation protocol: Evaluation uses 38,912 maximum generated tokens, N = 12 samples per problem, temperature T = 1.0, and top-p = 0.95.Top-k is disabled and enable_thinking=True.
  • Metrics: Reported metrics are Avg@12, Pass@12, and Maj@12, covering repeated-sample accuracy, search-style success, and aggregation stability.Majority voting uses math-equivalence clustering, with INVALID predictions scored incorrect when selected.
  • Benchmarks: Each method–benchmark pair uses three random evaluation seeds, with results reported as mean ± across-seed sample standard deviation.Benchmarks include MATH-500, AIME 2024, AIME 2025, and HMMT February 2025.

J Evaluation metric definitions

The evaluation uses repeated sampling to measure complementary aspects of reasoning performance, reporting Avg@12, Pass@12, and Maj@12. Answers are extracted and graded per sample, then aggregated across problems.

  • Repeated sampling reduces the risk that one completion understates a reasoning model’s chance of finding a correct solution.The evaluation reports results from 12 generated solutions per problem.
  • Pass@12 records whether any of 12 sampled solutions succeeds, while majority-based metrics aggregate multiple reasoning paths.
  • Each sample’s answer is extracted from its last boxed expression, with unparseable answers assigned the cluster key INVALID.
  • Samples are graded correct against the gold answer, and reported metrics are means of per-problem scores across the benchmark.Evaluation uses math_verify with a normalized string-equality fallback for parsing failures.
  • The accompanying code is available publicly, and the evaluation covers MATH-500, AIME 2024, AIME 2025, and HMMT February 2025 datasets.

L Reduction × positioning ablation

The ablation separates PW-OPSD’s position weighting from its per-sequence loss reduction. On Qwen3-4B AIME 2024, only their joint configuration matches the reported lead, while either change alone underperforms.

  • Reduction × positioning: PW-OPSD changes both token positioning and loss reduction: it uses position-dependent weights and averages token loss within each rollout.
  • Reduction × positioning: Table 6 evaluates uniform versus per-sequence reduction crossed with no positioning versus position-weighted supervision.The ablation uses Qwen3-4B on AIME 2024 and reports Avg@12 across three evaluation seeds.
  • Reduction × positioning: Only the joint configuration matches the AIME 2024 lead, while switching either axis alone underperforms by approximately 1.5 pp.The two axes are complementary rather than independently sufficient on AIME 2024.
Loading 2605.21606v1…