On-policy self-distillation improves math problem solving accuracy
OPSRD: On-Policy Self-Role Distillation
Computation and LanguageArtificial IntelligenceMachine Learning
Summary
Solving math problems with AI can be tricky, especially when models make mistakes early on. The authors created a new way for AI to learn from itself by looking at multiple possible next steps, not just the one it picked. They use a version of the model pretending to be an expert to guide training without needing exact correct answers. This method helps the AI become better at solving competition-level math problems without needing special prompts when it runs. Their tests showed consistent improvements across different sizes of AI models.
What this means in practice
- •For machine learning engineers: Train language models more effectively on complex tasks using on-policy self-distillation without needing additional labeled data.
- •For ai system integrators: Deploy improved math problem-solving models that perform better without relying on specialized role-based prompts at inference time.
Authors
Weijie Ren, Yanwen Zhang, Hao Li, Zhuolin Qi, Hengyi Zhang, Naibo Wang
Abstract
Role prompting elicits specialized behavior from large language models through an expert identity, offering a lightweight way to guide reasoning on demanding tasks. However, evaluating or distilling complete role-prompted answers can miss useful next-token preferences when the sampled solution remains incorrect. Transferring these preferences also requires an objective that reaches alternatives the student rarely predicts. We introduce OPSRD, which uses a fixed expert role as privileged teaching context for on-policy self-distillation without reference solutions. A role-free student generates a trajectory, and a frozen instance of the same base model supplies role-conditioned distributions on its exact prefixes, exposing alternatives beyond the sampled continuation. Teacher-weighted forward KL targets alternatives the student underestimates, with clipping to limit individual vocabulary contributions. Supervision is restricted to the highest-entropy half of student positions, concentrating learning where predictions are uncertain. Experiments on three competition-math benchmarks with Qwen3-1.7B, 4B, and 8B show improvements over the base models without role prompts at inference. Forward KL achieves the highest macro-averaged accuracy among the three evaluated divergences at every scale. Code is available at https://github.com/zhansan114514/OPSRD.