Technical Summary: Unmasking On-Policy Distillation

Problem Statement

On-policy distillation (OPD) has emerged as a core post-training technique for reasoning models, offering dense, per-token supervision to complement sparse-reward reinforcement learning (RL). In this paradigm, a "teacher" model (either a larger external model or the student itself with privileged context) guides the student's token generation. However, practitioners lack principled guidance on critical design choices:

• Teacher Selection: Should the teacher be an external larger model or the student itself (self-distillation)?
• Context Design: What specific context (e.g., full solution traces, summarized solutions, or examples of wrong answers) should serve as the supervisory signal?
• Granularity: Does the optimal configuration vary by task, student capacity, or even individual tokens?

Current evaluation relies on aggregate performance metrics from expensive training runs, which obscure the dynamics at the level of individual tokens. It remains unclear under which conditions the teacher's signal is beneficial, neutral, or detrimental.

Methodology

The authors introduce a training-free diagnostic framework that operates at the highest resolution: per token, per question, and per teacher configuration. The framework consists of three main components:

1. The Ideal Per-Node Gradient

To evaluate teacher guidance, the authors derive an "ideal" gradient that maximally increases the student's probability of reaching a correct answer.

• Generation Tree: They decompose the generation process into a tree where nodes represent token positions and edges represent next-token choices.
• Empirical Success Probability ($\hat{P}^k_{succ}$): By sampling rollouts from the student, they estimate the probability of reaching a correct answer given the selection of token $k$ at node $u$.
• Ideal Objective: The ideal local objective is to maximize the expected success rate: $L_{ideal}(u) = \sum_k P^k_\theta P^k_{succ}$.
• Ideal Gradient: The gradient of this objective with respect to the student's logits is derived as $g^{ideal}_j = P^j_\theta (P^j_{succ} - \bar{P}_{succ})$. This direction pushes the student toward tokens that lead to success more often than average.
• Connection to Dr. GRPO: The authors show that Dr. GRPO (a variant of GRPO without length normalization) recovers this ideal gradient in expectation, validating it as an unbiased estimator of the reward-based signal.

2. Distillation Gradient Analysis

The framework analyzes the gradients produced by various distillation algorithms (GKD, single-sample estimators, MiniLLM). It demonstrates that these methods share a common local structure: $\frac{\partial L}{\partial z_j} = P^j_\theta (f_j - \bar{f})$, where $f_j$ is derived from the teacher's distribution (e.g., log-ratios) or trajectory-level rewards.

3. Gradient Alignment Score

To quantify how well a specific teacher configuration approximates the ideal signal, the authors define the Gradient Alignment Score:
$$\text{Align}(u) = \cos(g^{ideal}_u, g^D_u)$$
This cosine similarity between the ideal gradient and the distillation gradient at each token measures whether the teacher's guidance pushes toward success (positive alignment), failure (negative alignment), or is orthogonal/noisy (near-zero alignment).

4. Scalable Estimation via Targeted Rollouts

Computing $\hat{P}^k_{succ}$ for all tokens in long reasoning chains is computationally prohibitive. The authors propose a targeted-rollout algorithm:

• Exponential Depth Windows: The generation sequence is partitioned into windows of increasing size.
• Prioritization: Within each window, rollouts are allocated to tokens where the teacher disagrees most strongly with the student (high GKD gradient magnitude or probability difference).
• Efficiency: This approach scales with a user-chosen budget rather than sequence length, making it feasible for traces ranging from ~200 tokens (BoolQ) to ~30,000 tokens (AIME).

Key Results

The framework was applied to Qwen3-0.6B and Qwen3-1.7B models across BoolQ, MMLU, and AIME 2025 datasets with 8 different teacher configurations.

1. Distillation Helps More on Incorrect Paths

The most consistent finding is that distillation guidance exhibits significantly higher alignment on incorrect rollouts than on correct ones.

• When the student is on a failing path, the teacher reliably pushes toward tokens that lead to success.
• When the student is already on a correct path, the teacher's signal becomes noisy and weakly aligned with the ideal, as the student is already performing well.
• This pattern holds across all settings, model scales, and datasets.

2. Context Design and Student Capacity Interaction

The optimal distillation context is not universal; it depends jointly on the student's capacity and the task difficulty:

• Small Models (0.6B): Benefit most from self-distillation with correct-only demonstrations in their own reasoning style. They struggle to parse the reasoning patterns of much larger external teachers. Summarized contexts can hurt performance if the student needs full step-by-step traces.
• Larger Models (1.7B): Can effectively leverage external teachers (e.g., Qwen3-8B) and summarized contexts, which nearly double alignment compared to raw traces.
• Wrong Demonstrations: Including wrong examples (Self-1C1W) generally hurts performance on short-reasoning tasks (BoolQ, MMLU) by introducing noise. However, on hard math problems (AIME), wrong demonstrations provide useful contrastive signals, becoming the best teacher configuration for difficult questions.

3. No Universal Recipe

There is no single teacher or context configuration that maximizes alignment across all tasks.

• Task Dependency: External teachers outperform self-distillation for the 1.7B model on BoolQ but not on MMLU.
• Difficulty Dependency: On hard math, including a wrong demonstration improves alignment, whereas it degrades it on simpler tasks.
• Predictive Limits: While divergence between student and teacher distributions (KL, JS) correlates positively with alignment, the correlation is weak ($|\rho| < 0.05$). Divergence can serve as a cheap necessary-condition filter but is not a reliable predictor of whether a specific disagreement is helpful.

Significance and Claims

The paper claims to provide the first rigorous, training-free methodology to assess the quality of teacher guidance at the token level. Its primary contributions are:

1. Diagnostic Framework: A scalable pipeline to compute gradient alignment offline, revealing that teacher signals are a heterogeneous mixture of beneficial, neutral, and detrimental contributions.
2. Mechanistic Insight: The finding that distillation is most effective on failing rollouts suggests that current training pipelines, which treat all tokens equally, may be wasting gradient budget on tokens where the student is already correct.
3. Contextual Nuance: The demonstration that "comprehensibility" is a key factor—signals are only useful if the student can parse them—explains why small models fail to learn from much larger external teachers and why context design (summarization vs. raw traces) must be tailored to model capacity.
4. Future Directions: The results motivate per-task, per-token diagnostic analyses and suggest future training algorithms that could weight distillation loss based on alignment (e.g., emphasizing failing trajectories) or use divergence as a gating mechanism.

The authors conclude that while on-policy distillation is a powerful tool, its effectiveness is highly conditional, and the "optimal" configuration shifts dynamically based on the student's state, the task difficulty, and the specific token being generated.