Overcoming Scaling Limits in On-Policy Self-Distillation for LLM Reasoning

AI summary

Overview

Research area: Machine learning / large language model reasoning, specifically on-policy self-distillation (OPSD) for mathematical reasoning, privileged-information imitation, and training efficiency.

Technical level: Advanced. The paper builds on forward-KL distillation, conditional mutual information, and imitation-gap theory, though its core empirical claim is stated in plain terms.

Scope: A diagnostic study of where along a student's own trajectory OPSD supervision is applied, plus a new method (OASIS) that supervises only outcome-verified on-policy trajectories and replaces written reference solutions with same-problem model rollouts.

What This Paper Is About

On-policy self-distillation trains a student model to match a teacher's token-level distribution along trajectories the student itself generated, where the teacher sees extra privileged information (usually a written gold solution) that the student will not have at test time. This paper asks whether it matters where along the student's trajectory that supervision lands, and shows that the teacher's advantage is concentrated on failed trajectories, where the signal is least imitable and where the student cannot reproduce it. The authors then propose OASIS, which keeps the OPSD loss unchanged but applies it to verified rollouts and uses a different model-generated attempt as the teacher's context, requiring only final-answer labels.

Key Contributions

1. Characterization of the irreducible imitation gap in OPSD. A formal result (Proposition 1) shows that under forward-KL distillation, a student that cannot see the privileged context converges to the context-averaged teacher distribution, with a residual loss equal to the conditional mutual information, and that the optimal student is at least as entropic as the average teacher.

2. Separation of two roles that standard OPSD couples. The paper distinguishes the trajectory scaffold (which prefixes receive supervision) from the teacher context (the target distribution at those prefixes), and introduces two diagnostic protocols to study them separately: a signal-replay test on recorded scaffolds and a truncated-continuation recovery benchmark with cut fractions f ∈ {0.1, …, 0.9}.

3. Empirical finding that scaffold correctness matters more than context correctness. Across Qwen3-1.7B and Qwen3-4B, the gold-context recovery advantage is much larger on prefixes from non-verified trajectories, while teacher targets on outcome-verified trajectories are relatively insensitive to alternative contexts.

4. OASIS (On-policy Alignment via Scaffold-Isolated Supervision). A method that samples K rollouts per problem, verifies final answers, trains along the shortest verified rollout, and conditions the teacher on a distinct same-problem rollout. It requires only final-answer labels rather than written solutions.

Main Findings

• Teacher advantage is concentrated on failed trajectories. On failed prefixes at 1.7B, the gold-context teacher recovers a correct answer in 77% of continuations versus 37% for the student (Δ = +0.40), while empty and unrelated references give 37% and 32%. On verified prefixes the gap is small (0.99 versus 0.95, Δ = +0.04 at 1.7B and +0.02 at 4B).

• The advantage shrinks with scale. From 1.7B to 4B, the student's own recovery on failed prefixes rises from 0.37 to 0.50 and the teacher's advantage falls by more than half (Δ = +0.40 to +0.17). Reference content accounts for 29.6% of the teacher–student KL gap at 1.7B but 17.3% at 4B.

• Most OPSD supervision lands on unverified trajectories. At 1.7B and 4B, only 28.3% and 30.6% of training rollouts produce a correct answer, while 62.0% and 64.4% reach the token cap without an answer; 71.7% and 69.4% of rollouts end without a verified answer. Over 100 training steps, the verifying fraction falls from 0.333 to 0.229 at 1.7B and from 0.394 to 0.218 at 4B.

• Teacher behavior differs by trajectory outcome. On failed states the teacher has higher entropy (0.461 versus 0.297), assigns lower probability to the token the student sampled (0.775 versus 0.834), and produces a lower KL signal (0.181 versus 0.274) than on verified states.

• Context sensitivity on verified scaffolds is low. Measured in total-variation distance from the gold-context teacher: the student's own failed attempt 0.065, unrelated 0.056, empty 0.051, template-only 0.056, and a line-shuffled gold reference 0.025.

• OPSD's gain over the base model collapses with scale. Gains of 3.05, 1.98, and 0.14 points at 1.7B, 4B, and 8B, while roughly 70% of its updates remain on unverified trajectories.

• OASIS improves over OPSD at every scale, and the margin grows. Average improvements of 0.59, 1.86, and 3.05 points at 1.7B, 4B, and 8B across AIME 2024, AIME 2025, and HMMT 2025. At 8B, OASIS scores 77.50 / 72.22 / 47.77 (average 65.83) against OPSD's 73.75 / 68.75 / 45.84 (average 62.78).

• Dense teacher supervision contributes beyond the verifier. GRPO uses the same verifier and labels but improves over the base model by only 0.60, 0.60, and 0.23 points, versus 3.64, 3.84, and 3.19 for OASIS. SFT on written solutions is at or below the base model at all three scales.

• OASIS keeps the student closer to the teacher during training at 1.7B. Teacher top-1 agreement stays at 0.931 while OPSD's falls from 0.932 to 0.888, and student entropy reaches 0.49 versus 0.70 for OPSD. At 4B the separation lasts about two thirds of training, after which OASIS also drifts, ending at 0.853 top-1 agreement with about half of tokens carrying a clipped entry.

• The per-entry KL clip is a second, independent problem. At κ = 0.05, clipping removes 97% of positive KL mass on verified scaffolds and 94% on failed ones, affecting 29–32% of tokens. Proposition 2 shows that a clipped entry loses the teacher's upward pull and receives a downward logit gradient.

• OASIS reaches comparable accuracy with far less labelled supervision. With a different training configuration (batch size 64 up to 200 steps for baselines), OASIS sees 6.4K problems and trains on roughly 3.2K, while OPSD and AVSD see and train on 12.8K problems each with written solutions. At 4B OASIS averages 63.70 versus 63.60 for OPSD and 62.90 for AVSD; at 8B, 65.83 versus 64.80 and 64.03.

• The trade-off is generation cost. OASIS samples eight rollouts per problem, generating 51,200 rollouts against 12,800 for a one-rollout OPSD run on the larger problem set. With K = 8 at 1.7B over 100 steps it generates 43.5M tokens versus 2.76M for OPSD, but supervises fewer (1.98M versus 2.76M) because about half of problems are masked and verified scaffolds are shorter (603 versus 859 tokens).

Methodology in Plain English

The authors first run offline diagnostics rather than training runs. In a recovery probe, they take a student rollout, cut it at various fractions of its length (f ∈ {0.1, …, 0.9}), and let the model continue from that prefix under different conditions: with no context, with the gold reference, with only the final answer, with the gold reference stripped of its answer, with an unrelated solution, and with an empty reference. They then measure how often each condition reaches the correct answer. In a signal replay test, they replay recorded student and teacher distributions along existing scaffolds to measure how far the teacher's next-token distribution moves when its context is changed.

They group trajectories by whether the final answer verifies, and compare the teacher's advantage across the two groups and across model sizes (Qwen3-1.7B, 4B, 8B). They also quantify how the teacher–student KL gap splits between the chat template, the reference block and its instructions, and the reference content itself.

Building on the finding that verified scaffolds are where the teacher is imitable, OASIS changes only two things relative to OPSD: which trajectories get supervised, and what the teacher is shown. For each problem it draws K = 8 rollouts, verifies their final answers, keeps problems with at least one verified rollout, and applies the unchanged OPSD loss along the shortest verified rollout. The teacher context is a distinct rollout of the same problem, typically the shortest unverified one; if every rollout verifies, another verified rollout is used, and the scaffold is never used as its own context. The loss, teacher, prompts, clipping constant, temperature, and optimizer are identical to OPSD.

Evaluation uses Avg@12 on AIME 2024, AIME 2025, and HMMT 2025 with Qwen3's recommended sampling settings, training on the mathematical subset of OpenThoughts, with checkpoints evaluated every 20 steps under the same selection rule for all methods. Comparisons include the base model, SFT on reference solutions, GRPO with a binary outcome reward against the same answer verifier, OPSD, AVSD, and CRISP. One table additionally reports three evaluation seeds on one Qwen3-4B checkpoint per method to show that evaluation variability is small relative to the OASIS–OPSD gap.

Why This Matters

Impact on research. The paper reframes a known training recipe by asking where supervision is applied rather than only what the teacher sees, and supplies a theoretical account (the conditional-mutual-information residual) for why imitation of a privileged teacher cannot be fully closed. It also documents a separate pathology in the per-entry KL clipping used by the implemented objective, showing that clipped entries reverse the sign of their logit gradient. For the community, this suggests that scaling behavior of self-distillation depends on whether teacher supervision is transferable to the student's observable context, and that outcome verification can substitute for written solutions in selecting scaffolds.

Real-world applications (as suggested by the paper's scope and its stated next step):

• Training mathematical reasoning models from datasets that supply only checkable final answers rather than worked solutions.
• Code generation with unit tests, which the authors name explicitly as a natural extension because outcome verification is inexpensive and written solutions are absent.
• Domains with automatic answer checkers, such as short-answer science or formal-answer tasks, where labels exist but rationales do not.
• Reducing annotation cost in distillation pipelines by replacing reference-solution procurement with additional sampling, which the paper notes parallelizes well and needs no labeled data.

Industry relevance. The trade-off the paper quantifies is directly operational: additional rollouts are cheap and parallelizable, whereas human-written solutions are expensive and often unavailable. OASIS matches or exceeds baselines on average at 4B and 8B using roughly a quarter of the supervised problems and no written solutions, which matters for teams that have verifiers but not rationales. The paper also notes the cost side plainly: generation becomes the binding constraint in some settings, and the K = 4 ablation halves that cost.

Future Directions

• Extending the criterion beyond competition mathematics. The authors propose settings with inexpensive outcome verification but no written solutions, such as code with unit tests, as the natural next step.
• Reducing generation cost. Generations dominate the compute budget (43.5M tokens versus 2.76M at 1.7B with K = 8); the K = 4 ablation halves Rollouts, but cheaper selection schemes remain open.
• Addressing the 4B drift. OASIS's agreement with the teacher eventually drifts at 4B, with about half of tokens carrying a clipped entry; the authors flag clipping as one possible explanation and leave it unresolved.
• Broadening model families and repetition. The study covers one model family (Qwen3), competition mathematics only, and single runs at 4B and 8B, and the abstract-level claim at 1.7B is described as within the range sampling variation can produce on 30 problems.

Target Audience

Researchers and engineers working on knowledge distillation, reasoning-model post-training, and privileged-information imitation who want to understand when self-distillation supervision transfers to the student and when it does not. It is most useful to readers already comfortable with KL-divergence objectives and on-policy training, and to practitioners who have answer verifiers but no written solutions and need to decide how to allocate a generation budget.

Authors’ abstract

On-policy self-distillation (OPSD) trains a student to match a privileged teacher distribution along its own sampled trajectory. Standard OPSD applies this supervision to unverified student rollouts while conditioning the teacher on privileged context, typically a reference solution. We separate these roles in a factorial analysis and find that scaffold correctness has a stronger effect on downstream accuracy than context correctness. Unverified scaffolds create an imitation gap because the teacher can use information unavailable to the student. This gap shrinks with model scale, yet OPSD continues to supervise mostly unverified trajectories. In contrast, verified scaffolds remain effective even when the teacher is conditioned on the student's own unsuccessful rollout. Based on this finding, we introduce OASIS, which retains the OPSD objective but supervises mostly verified by label on-policy trajectories and replaces written solutions with unverified model-generated attempts as the teacher context. OASIS therefore requires only final-answer labels. Across Qwen3-1.7B, 4B, and 8B on AIME 2024, AIME 2025, and HMMT 2025, OASIS improves over the base model by 3.2--3.8 points on average, while OPSD's gain falls from 3.05 points at 1.7B to 0.14 at 8B. At 8B, OASIS improves over OPSD by 3.05 points, showing that verified on-policy scaffolds preserve the effectiveness of self-distillation as models scale.