AI Papers Reader

Personalized digests of latest AI research

View on GitHub

A Knob for Self-Distillation Makes Reasoning Models More Stable Learners

Teaching a large language model to reason from its own attempts is attractive because it avoids labeled reasoning traces, but on-policy self-distillation (OPSD) is finicky in practice. University of Maryland researchers identify a hidden cause: the standard recipe fixes a regularization weight at one value. Their method, beta-OPSD, makes that weight explicit and uses it to build smoother training targets that improve math-reasoning benchmarks.

OPSD trains a student on trajectories it samples itself, while a privileged teacher—which can see a correct solution during training—provides token-level feedback. The student minimizes a Kullback-Leibler (KL) divergence, a measure of distribution mismatch, to that feedback. The team shows this objective is the beta=1 case of a broader KL-regularized policy-optimization family. Beta controls how strongly the student is penalized for moving away from a reference policy, such as a stop-gradient copy of itself. At beta=1, the teacher is the sole target; larger betas keep the target closer to the reference. The optimal policy is a geometric interpolation between reference and teacher, so each beta selects a point along that path.

Rather than run expensive reinforcement learning, beta-OPSD turns that closed-form solution into a distillation target. At each token, it averages the student logits—raw scores before softmax—with teacher logits, then normalizes. A linear schedule moves teacher weight from 0.5 to 0.8 over 200 steps, so early targets stay near the student and later ones approach the teacher. The loss also uses return-to-go credit assignment: early tokens receive weight for future mismatch they cause.

An ablation on Qwen3-1.7B shows the effect. With the same return-to-go estimator but a fixed teacher target, avg@12—average correctness across 12 sampled solutions—was 47.30 on AIME 2024. Replacing that target with the scheduled beta-OPSD interpolation raised it to 53.33, and also improved AIME 2025 and HMMT 2025. Researchers attribute this to a smoother target path rather than an abrupt projection onto the teacher.

In the main comparison, beta-OPSD outperformed vanilla OPSD on average across Qwen3-1.7B, 4B, and 8B, and also outperformed supervised fine-tuning and GRPO, a reinforcement-learning baseline, in these tests. At a 100-step checkpoint, for the 1.7B model, absolute gains were 9.16 percentage points on AIME 2024, 5.27 on AIME 2025, and 2.78 on HMMT 2025, raising the average from 31.02 to 36.76. The 4B and 8B models gained 1.76 and 1.66 percentage points on average, though beta-OPSD trailed vanilla OPSD on some individual benchmarks, including HMMT 2025 for 8B.

Those gains come with caveats. The 0.5-to-0.8 schedule was best overall among tested schedules, but no schedule dominated every benchmark, and the authors call the optimal schedule an open question. Experiments use competition math and LoRA (low-rank adapter) fine-tuning, so extension to other tasks or full fine-tuning is unresolved. The discounted return-to-go is an approximation to the exact sequence-level gradient. The broader contribution is conceptual: deriving a distillation target from policy optimization gives researchers a controllable dial for managing the student–teacher gap, which could make self-improving models more stable, though the evidence so far is limited to math benchmarks.