PaperlayerView exact revision on arXiv

Source-linked AI summary

Unmasking On-Policy Distillation: Where It Helps, Where It Hurts, and Why

Mohammadreza Armandpour, Fatih Ilhan, David Harrison, Ajay Jaiswal, Duc N. M Hoang, Fartash Faghri, Yizhe Zhang, Minsik Cho, Mehrdad Farajtabar

arXiv:2605.10889v1cs.LGcs.AI

TL;DR

On-policy distillation lacks principled guidance about which teacher or context to use and whether the optimal signal varies by question or token. The paper introduces a training-free per-token diagnostic based on ideal-gradient estimation and finds that guidance aligns better on failing rollouts, while context effectiveness depends on student capacity and task, with no universal recipe.

  • Problem

    Practitioners lack principled guidance for choosing teachers and contexts across questions and tokens, while aggregate training metrics obscure token-level effects.

  • Method

    The paper estimates an ideal per-node success-improving gradient and compares distillation gradients with it using scalable offline gradient-alignment diagnostics.

  • Results

    Distillation aligns better on failing rollouts, while effective contexts vary with student capacity and task, yielding no universal configuration.

  • Takeaways & Limitations

    Distillation should be diagnosed per task, token, and student rather than applied through a fixed teacher or context pipeline.

  • Takeaways & Limitations

    The diagnostic evaluates only sufficiently sampled tokens and requires substantial per-question compute, making it primarily an offline evaluation tool.

Abstract

from arXiv · show

On-policy distillation offers dense, per-token supervision for training reasoning models; however, it remains unclear under which conditions this signal is beneficial and under which it is detrimental. Which teacher model should be used, and in the case of self-distillation, which specific context should serve as the supervisory signal? Does the optimal choice vary from one token to the next? At present, addressing these questions typically requires costly training runs whose aggregate performance metrics obscure the dynamics at the level of individual tokens. We introduce a training-free diagnostic framework that operates at the highest resolution: per token, per question, and per teacher. We derive an ideal per-node gradient defined as the parameter update that maximally increases the student's probability of success. We then develop a scalable targeted-rollout algorithm to estimate this gradient efficiently, even for long chains of intermediate thoughts. The gradient alignment score, defined as the cosine similarity between this ideal gradient and any given distillation gradient, quantifies the extent to which a particular configuration approximates the ideal signal. Across a range of self-distillation settings and external teacher models, we observe that distillation guidance exhibits substantially higher alignment with the ideal on incorrect rollouts than on correct ones, where the student already performs well and the teacher's signal tends to become noisy. Furthermore, we find that the optimal distillation context depends jointly on the student model's capacity and the target task, and that no single universally effective configuration emerges. These findings motivate the use of per-task, per-token diagnostic analyses for distillation.

1 Introduction

The paper introduces a training-free, per-token diagnostic framework that links teacher guidance to downstream success, addressing decisions that aggregate training metrics cannot resolve. Across tasks and teacher configurations, alignment depends on rollout correctness, student capacity, context, and question difficulty, with no universal recipe.

  • Motivation: Practitioners lack principled guidance on teacher choice and context design because expensive training runs obscure token-level effects.The unresolved choices include larger external teachers versus self-distillation, full traces versus summaries, and whether the answer varies by question or token.
  • Method: The framework evaluates teacher guidance per token, question, and teacher configuration by comparing it with behaviors that lead to correct answers.It derives an ideal per-token gradient from empirical success probabilities and estimates it scalably for long reasoning chains using targeted rollouts.
  • Key findings: Distillation guidance is more reliable on incorrect rollouts, while signals on already-correct paths are noisy and weakly aligned with the ideal.This pattern holds across all tested settings and metrics.
  • Key findings: Context design interacts with student capacity: summarized solutions nearly double alignment for 1.7B, slightly hurt 0.6B, and 32B teachers help 0.6B on simple tasks but fail on hard math.External teachers outperform self-distillation only for the larger student.
  • Key findings: No universal recipe exists because the best teacher or context shifts with task difficulty, while divergence is only a weak alignment filter with |ρ| < 0.05.Contrastive examples can hurt on simple reasoning but help on hard math, and a wrong demonstration achieves the highest alignment on AIME 2025 hard math problems.

2 Methodology

The method estimates token-level success probabilities from a generation tree, derives an ideal success-improving gradient, and compares distillation gradients against it through cosine alignment. A scalable rollout procedure makes these per-node, per-teacher diagnostics practical for long sequences and many teacher configurations.

  • Generation-tree analysis: Generation trees estimate each token’s probability of ultimately reaching a correct answer from sampled student trajectories.Nodes represent token positions, while observed continuations and outcomes provide empirical success probabilities for sufficiently visited tokens.
  • Ideal gradient: The ideal local gradient increases logits for tokens whose success probability exceeds the student’s current expected success and decreases them otherwise.This gradient is the reference direction that maximally improves the student’s chance of success at each node.
  • Distillation gradients: Distillation methods produce different per-node gradients from the same teacher, while MiniLLM additionally couples each node to downstream reward-to-go.The shared gradient structure enables comparison across GKD, its single-sample estimator, MiniLLM, and the ideal reference.
  • Gradient alignment: Gradient alignment is the cosine similarity between an ideal gradient and a distillation gradient restricted to sufficiently sampled tokens.The restriction is necessary because success probabilities are reliable only for tokens observed often enough; positive, near-zero, and negative alignment distinguish helpful, irrelevant, and harmful guidance.
  • Scalable computation: The scalable procedure avoids exhaustive branching rollouts and reuses one rollout tree across eight teacher configurations, reducing total compute by ∼7×.Naïve evaluation would require thousands of rollouts for every possible token at every node in long sequences with 150K-token vocabularies.

3 Experimental Setup

The experiments compare two Qwen3 student scales across eight self-distillation and external-teacher configurations on BoolQ and MMLU, using per-node gradient-alignment and related statistics. AIME 2025 case studies extend the analysis to substantially longer reasoning traces.

  • Student models: The study evaluates Qwen3-0.6B and Qwen3-1.7B as student models.
  • Teacher configurations: Eight teacher configurations span self-distillation contexts and larger external Qwen3 teachers.Self-distillation includes correct, wrong, summarized, and Qwen3-32B-generated contexts; external teachers are Qwen3-4B, Qwen3-8B, and Qwen3-14B.
  • Datasets: Evaluation uses BoolQ for True/False reading comprehension with short reasoning chains and MMLU for multiple-choice knowledge with medium-length chains.
  • Datasets: AIME 2025 case studies examine reasoning traces of approximately 5K–30K tokens.
  • Metrics: At eligible branching nodes, the analysis computes gradient alignment, teacher advantage, and success-rate statistics aggregated by path, question, and teacher.Eligibility requires at least two children with at least 20 visits each and a nonzero success-rate range; teacher means include 95% confidence intervals across questions.

4 Results

Results show that distillation alignment is weak on average but highly variable across tokens, with stronger guidance on incorrect than correct paths. The best teacher and context depend on model capacity and task difficulty, while no single feature reliably predicts useful alignment.

  • Overall alignment: Mean cosine alignment is only +0.027 for 0.6B and +0.026 for 1.7B on BoolQ, despite per-token standard deviations of approximately 0.83–0.91.The analysis covers approximately 88K BoolQ and 95K MMLU decision points across the two model scales.
  • Path correctness: Incorrect paths consistently have higher gradient alignment than correct paths because teacher guidance opposes failing trajectories, whereas successful students receive weaker, noisier reward-aligned signals.This pattern is reported as the most consistent finding across settings.
  • Teacher choice: For 0.6B, correct-only self-distillation teachers achieve 2–3× higher alignment than external teachers, while Qwen3-8B leads for the 1.7B student on BoolQ.Teacher rankings therefore reverse across model scales, with external teachers becoming competitive for larger students.
  • Demonstration composition: Including wrong demonstrations consistently hurts BoolQ and MMLU, with 1C1W variants underperforming 1C variants; Self-1C for 0.6B uniquely has ∆≈0 between path types.The results suggest negative examples add noise on these shorter-reasoning tasks.
  • Context choice: Context effects are capacity- and task-dependent: summaries nearly double 1.7B BoolQ alignment from 0.028 to 0.050, whereas 0.6B benefits from verbose traces.Self-1C (32B) works well for 0.6B but poorly for 1.7B and AIME, consistent with differences in context comprehensibility.
  • Token-level predictors: Teacher–student divergence correlates positively with alignment, but all within-path feature correlations remain weak at |ρ| ≈0.02–0.04, so no feature reliably predicts useful signal.The useful signal tends to occur where the teacher disagrees with the student.
  • Longer reasoning traces: On four AIME 2025 questions, incorrect paths again align better, but Self-1C1W becomes best on harder questions and raw demonstrations outperform summaries.For example, Q0 shows incorrect alignment of +0.097 versus correct alignment of −0.011, with ∆= −0.108.

5 Related Work

Prior work established on-policy distillation as a response to exposure bias, while identifying conditions under which teacher guidance or self-distillation can fail. This work adds token-level gradient alignment to explain these phenomena and connects distillation with reward-based objectives through a shared local structure.

  • On-Policy Distillation: On-policy distillation supervises student-generated rollouts, with MiniLLM using reverse-KL policy gradients and GKD interpolating between on- and off-policy data.Sequence-level distillation generates teacher outputs but suffers from exposure bias; OPD addresses this by training on the student’s own rollouts.
  • Failure Modes: Concurrent studies attribute OPD failures to inconsistent thinking patterns, insufficiently novel teacher knowledge, and suppressed epistemic verbalization in self-distillation.Richer conditioning contexts can suppress uncertainty tokens, while broader task coverage can harm out-of-distribution generalization.
  • Token-Level Diagnosis: Gradient alignment explains distillation failures by quantifying at each token position whether the teacher’s signal is exploitable by the student, independent of aggregate teacher performance.This provides a mechanistic account of why overly capable teachers can hurt and why small models may struggle with strong reasoners.
  • Reward-Based Training: Reward and distillation objectives share the local structure P_jθ(f_j − f̄), enabling direct offline comparison at token granularity without training or additional models.This connects sparse outcome rewards, step-level process feedback, and distillation within a common gradient framework.

6 Conclusion · A Full AIME 2025 Case Studies

The paper derives an ideal per-node gradient and a scalable offline pipeline for testing whether token-level distillation signals point toward correct answers. Experiments show that distillation helps most on failing rollouts, context design depends on student capacity, and no universal recipe exists.

  • 6 Conclusion: The study asks whether teacher distillation signals point toward correct answers at each token in a reasoning chain.This frames the paper’s central diagnostic question.
  • 6 Conclusion: The authors derive an ideal per-node gradient from empirical success probabilities.The gradient is defined at the token or node level of a reasoning chain.
  • 6 Conclusion: Major distillation objectives share the same local structure, enabling a common gradient-alignment analysis.The shared structure supports comparing different distillation objectives within the same framework.
  • 6 Conclusion: A scalable pipeline computes gradient alignment offline for long sequences.The pipeline avoids requiring immediate training runs for token-level diagnostic analysis.
  • 6 Conclusion: Distillation helps most on failing rollouts, while context design interacts strongly with student capacity.These findings indicate that distillation effectiveness varies with rollout outcome and model capacity.
  • 6 Conclusion: No universal recipe exists for distillation.The conclusion therefore emphasizes configuration-dependent behavior rather than a single best setup.

A.1 Overview … B.3 GKD

Across long AIME reasoning chains, teacher alignment is strongest on incorrect rollouts, while the best teacher varies with question difficulty, student capacity, filtering, and prioritized decision points. The appendices derive per-node gradients from softmax, Dr. GRPO, and GKD formulations.

  • A.1 Overview: Self-1C1W is the best teacher on harder AIME questions Q3 and Q28, unlike BoolQ/MMLU where wrong demonstrations consistently hurt.On Q3, Self-1C1W reaches +0.042 alignment; the paper attributes its advantage to useful contrastive information from common mistakes.
  • A.1 Overview: Across all four AIME questions, teacher alignment is higher on incorrect than correct rollouts at consequential decision points.Examples include Q0: incorrect +0.097 versus correct −0.011, and Q7: incorrect +0.093 versus correct +0.017.
  • A.1 Overview: Teacher rankings shift with filter stringency, and no universal distillation recipe emerges across task, difficulty, student capacity, and prioritized decision points.External teachers lead under lenient Q7 filters, whereas Self-Sum-1C dominates strict filters focused on high-stakes nodes.
  • B.1 Softmax Jacobian: The softmax-Jacobian appendix formulates the student’s transition probability and its Jacobian, establishing the local probability-gradient structure.The derivation uses the Jacobian’s sum-to-zero property in subsequent expressions.
  • B Full Gradient Derivations: The gradient derivation decomposes the training objective per node by simplifying the full GRPO objective and removing KL, importance-ratio, length-normalization, and constant terms.The resulting gradient-relevant objective is identified with the ideal objective, while finite-sample gradients can be estimated from rollouts passing through each node.
  • B.2 Dr. GRPO: Dr. GRPO yields a lower-variance ideal-gradient estimator by computing gradients directly from empirical success-probability estimates.The estimator is equivalent in expectation to the rollout-based per-node gradient.
  • B.3 GKD: The GKD derivation invokes the Jacobian sum-to-zero property to show that the centered term θ(1 −1) contributes zero.This identity supports simplifying the gradient expression in the GKD appendix.

B.4 Single-sample GKD estimator (Thinking-Lab) · B.5 MiniLLM

B.4 presents a single-sample GKD estimator that importance-weights the sampled token and matches the full-vocabulary KL gradient direction in expectation. B.5 contrasts MiniLLM’s REINFORCE-style, reward-to-go gradient, whose per-node contribution depends on the future trajectory.

  • B.4 Single-sample GKD estimator (Thinking-Lab): B.4 uses importance weighting at the sampled token for each rollout’s node-level estimator.For rollout i choosing token r_i at node u, the empirical estimator applies the weight only to that sampled token.
  • B.4 Single-sample GKD estimator (Thinking-Lab): B.4 defines a per-sample gradient contribution for the sampled rollout and node.The passage introduces the per-sample contribution following the importance-weight definition.
  • B.4 Single-sample GKD estimator (Thinking-Lab): B.4 obtains the estimator’s expected gradient by averaging over the student’s sampling distribution.The expectation is explicitly taken with respect to the student’s sampling distribution.
  • B.4 Single-sample GKD estimator (Thinking-Lab): B.4 recovers the GKD gradient with opposite sign, showing equivalence in expected gradient direction to full-vocabulary KL minimization.The passage states that the single-sample importance-weighted estimator and full-vocabulary KL minimization produce the same direction in expectation.
  • B.5 MiniLLM: B.5 MiniLLM uses a REINFORCE-style gradient in which reward-to-go couples each node’s gradient to all downstream nodes.At node u, corresponding to step t and sampled token r_i, the gradient is defined using this reward-to-go structure.
  • B.5 MiniLLM: B.5 assigns a per-step reward to the MiniLLM formulation, distinguishing it from GKD and its single-sample variant.The cited passage identifies o_i,t′ as the per-step reward and contrasts the resulting form with GKD.
  • B.5 MiniLLM: B.5 makes each per-node contribution depend on the full future trajectory rather than remaining purely local.Consequently, the contribution is path-dependent.

C Discussion, Limitations, and Future Work … E.2 Per-path cosine oscillation

The diagnostic reveals highly heterogeneous, locally oscillating distillation quality, while its practical use is limited by sampling coverage and substantial targeted-rollout compute. Future work therefore targets adaptive training, teacher combination, and longer-horizon evaluation.

  • C Discussion, Limitations, and Future Work: The alignment score applies only to tokens with sufficient visit counts, leaving tokens never sampled by the student unevaluable.This restricts diagnostic coverage to observed student behavior.
  • C Discussion, Limitations, and Future Work: Targeted rollout enrichment requires substantial compute per question, making the diagnostic more practical offline than as a real-time training signal.The computational cost is a stated limitation of deployment during training.
  • C Discussion, Limitations, and Future Work: Future training could up-weight distillation loss on incorrect rollouts, where alignment is strongest, or gate teacher signals using student-teacher divergence.Other proposed extensions include multi-teacher distillation and long-horizon agentic tasks.
  • D Computation Details: Rollout prioritization ranks tokens by GKD gradient magnitude and probability difference, identifying strong teacher disagreement and distribution divergence.The method allocates budgets by depth window, using smaller dense windows early and larger coarser windows later.
  • D Computation Details: The multi-teacher pipeline shares initial rollouts, tree construction, and targeted enrichment before computing teacher-specific logits and alignment scores.Each teacher benefits from prior enrichment in the shared tree.
  • E.1 Alignment score distribution: Per-node alignment scores span −1 to +1 across 8 BoolQ teachers, with slight positive means of +0.027 for 0.6B and +0.026 for 1.7B.This spread demonstrates extreme per-token heterogeneity.
  • E.2 Per-path cosine oscillation: Alignment oscillates between positive and negative values along individual reasoning paths, so consecutive tokens can receive beneficial or harmful distillation signals.The teacher’s helpfulness is highly local, even within a single path.

E.3 Selective distillation (oracle analysis)

Oracle selective distillation applies teacher gradients only at tokens with positive alignment, producing a large signal gain while retaining roughly half the tokens. Although true alignment is unavailable during training, this result provides an upper bound and motivates approximate divergence-based filters.

  • Oracle filtering: Selective distillation retains only decision points whose alignment exceeds a threshold t.The analysis asks whether teacher gradients can be applied only where alignment is positive.
  • Oracle filtering: 10–15× improvement in mean signal uses only ∼52% of tokens when filtering to positive alignment (t = 0).Every path benefits under this oracle strategy.
  • Training implications: The selective result is an oracle upper bound because true alignment is unknown at training time.The analysis therefore motivates divergence-based predictors as approximate filters for alignment-aware training.

E.4 Teacher ranking by alternative metrics … E.8 Per-teacher correct vs. incorrect breakdown

Alternative metrics largely preserve teacher rankings but can elevate summarized self-distillation for the larger student. Alignment is generally higher on incorrect paths and at high-divergence tokens, while depth effects vary by model and task.

  • E.4 Teacher ranking by alternative metrics: For 0.6B, fraction-positive and weighted-cosine metrics preserve mean-cosine teacher rankings, whereas weighted cosine reverses the 1.7B top ranking.Self-Sum-1C leads over Qwen3-8B for 1.7B under weighted cosine.
  • E.5 Additional teacher ranking plots: The supplied passages do not describe additional teacher-ranking plots.
  • E.4 Teacher ranking by alternative metrics: Weighted cosine gives 0.6B self-distillation teachers 0.108–0.120 versus 0.053–0.062 for external baselines, a 2× ratio.For 1.7B, Self-Sum-1C leads at 0.088, followed by Qwen3-8B at 0.072.
  • E.6 Full correct vs. incorrect breakdown: Teacher advantage persists across correct-versus-incorrect alignment metrics, with weighted cosine producing the largest effect sizes.
  • E.7 Within-path correlation details: The divergence-to-alignment pattern holds for 0.6B and strengthens for the best teacher, but weakens for 1.7B BoolQ regardless of teacher.Positive alignment means teacher disagreement with the student corresponds to higher alignment, while low-divergence tokens provide little useful signal.
  • E.7 Within-path correlation details: Depth correlates positively for 0.6B (ρ ≈+0.04), vanishes for 1.7B, and becomes weakly negative on hard AIME questions (ρ ≈−0.03).The results suggest contextual teacher advantage emerges later for 0.6B but does not grow with depth on complex math.
  • E.8 Per-teacher correct vs. incorrect breakdown: For 0.6B, Self-1C uniquely has nearly equal correct- and incorrect-path alignment (∆≈0), while other teachers align more highly on incorrect paths.For 1.7B, higher alignment on incorrect paths holds without exception.

E.9 Teacher Context Generation and Screening Details

The screening pipeline compares seven teacher-context configurations generated and verified across MMLU and BoolQ, while filtering questions by difficulty and demonstration availability. Results indicate that correct demonstrations help, wrong demonstrations hurt, and self- versus 32B-generated demonstrations differ little.

  • Screening and filtering: Screening uses MMLU and BoolQ, bins questions by baseline pass rate, and excludes extremely hard questions because their zero baseline cannot measure improvement.The bins are easy (p ≥0.8), medium (0.25 ≤p < 0.8), hard (0 < p < 0.25), and extremely hard (p = 0).
  • Context generation: Demonstration responses are sampled at temperature τdemo=0.7, checked against ground truth, classified by correctness, and regenerated for up to 60 trials.This process collects the required number of demonstrations for each question before screening.
  • Context configurations: The study evaluates seven context configurations spanning raw and summarized demonstrations, including correct-only, mixed correct-and-wrong, and self-versus-32B-generated variants.Raw variants include Self-1C, Self-1C1W, and Self-3C; summarized variants include Self-Sum-1C and Self-Sum-1C1W.
  • Screening and filtering: The analysis retains only questions with at least one correct demonstration from both the self-context and 32B-context sources.This filter is intended to make comparisons across context variants fair and to exclude questions lacking a correct generated response from either source.
  • Screening results: A single correct self-generated demonstration produces dramatic teacher-accuracy improvements, whereas adding a wrong demonstration consistently hurts and can fall below the no-context baseline.The reported patterns hold across models and benchmarks.
  • Screening results: The performance gap between self-generated and Qwen3-32B-generated single correct demonstrations is small.This comparison is specifically between Self-1C and Self-1C (32B).
Loading 2605.10889v1…