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

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

arXiv · PDF · Open in the Atlas

Code that ran

We lifted 2 functions out of this paper's own repositories and ran 2 of them in a sandbox. "Ran" means the function executed on a synthesized input and returned a value. It is not a reproduction of the paper's results.

Repositories linked to this paper

• SaFo-Lab/PW-OPSD

Some links come from the archived Papers with Code dataset (CC BY-SA 4.0): attribution and licence.

Abstract

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 studentvisited 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 studenttemplate continuation recovers the correct answer. On Qwen3-4B, we find that an oriented within-sequence position score is the strongest tested predictor of teachertoken reliability, reaching an area-under-ROC-curve (AUROC) of 0.83 with a 95% cluster-bootstrap interval of [0.66, 0.95]; 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 trajectorystructured and can be utilized without additional teacher computation. The code is available at https://github.com/SaFo-Lab/PW-OPSD

For agents

The same record, over MCP at https://syntology.ai/mcp:

get_harvested_code_for_paper("2605.21606")
get_code_for_paper("2605.21606")
have("2605.21606")

Connect an agent — have() is free.