SPOT: Outcome-Aware On-Policy Distillation for Reasoning Models
Lead
On-policy distillation (OPD) is a practical way to train smaller reasoning models: the student generates its own trajectories, and a stronger teacher provides supervision along those paths. But standard OPD also has a blind spot. When training is driven by reverse KL, the student may assign too little probability to alternative continuations that are still plausible—and sometimes more useful for reaching a correct final answer.
The paper SPOT: Sparse Probing and Outcome Calibration for On-Policy Distillation proposes a more selective and outcome-aware approach. Its basic message is simple: do not probe every token position, and do not treat the teacher’s local probability as the only signal. Instead, spend limited compute on the positions where extra information is most valuable, then use downstream verification to decide which candidate continuations deserve more weight.
Key ideas
- Why standard OPD can under-cover alternatives: Dense teacher supervision is useful, but reverse-KL-style training may concentrate probability mass too narrowly around the teacher’s preferred next tokens.
- Why entropy alone is insufficient: Teacher entropy can indicate uncertainty, but it does not reveal whether uncertainty is concentrated among a few strong candidates or spread across a long tail. It also says little about whether the student already represents those candidates well.
- A three-stage procedure: SPOT is organized around acquisition, exploration, and exploitation, jointly answering where to probe and what to distill.
- Sparse probing during acquisition: The method builds a position-level score from normalized teacher entropy, the probability mass captured by a small top-k candidate set, and student-teacher mismatch. This score allocates a limited probing budget to the most informative positions.
- Outcome-calibrated targets: During exploration, teacher-proposed candidates are tested through student continuations, which are then scored by a verifier. During exploitation, these scores are converted into a closed-form KL-regularized target that favors candidates with better downstream outcomes while staying anchored to the teacher distribution.
Why it matters
SPOT shifts distillation from local imitation toward result-aware imitation. In reasoning tasks, a next-token probability is not always a reliable proxy for eventual success. A continuation that looks less dominant locally may lead to a stronger solution path later. By introducing sparse future-looking checks, SPOT tries to improve both solution quality and coverage without turning every training step into an expensive search procedure.
According to the reported results, across three student scales and multiple reasoning benchmarks, SPOT improves macro Avg@8/Pass@8 over OPD by 0.47–1.48/4.55–5.28 points, and over EOPD by 0.29–0.68/2.49–3.19 points. These gains suggest that calibrating targets with downstream outcomes can be more effective than relying on teacher entropy or local probabilities alone.
For practitioners, the important takeaway is not just the performance gain, but the training philosophy: better distillation may require asking which divergences are worth investigating and whether a candidate actually helps the final answer. If verifier cost can be managed, outcome-calibrated OPD could become a useful tool for compressing reasoning capability into smaller language models.
Source: Hugging Face Daily Papers
Comments
Checking sign-in status...
Loading comments...