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

In plain words

When a small model learns to reason, a bigger "teacher" model can tell it, at every word it writes, which choice looks better. That signal is useful, but the teacher is not always right, and nobody knew at which words, for which questions, or with which teacher the hints actually helped.

This paper builds a way to check that without training anything. The student writes many answers to a question. At each point where the answers split into different next words, the authors ask which words actually led to correct answers. Then they compare the direction the teacher pushes with the direction that would really help, and measure how well the two agree.

The main surprise is that the hints help most on answers the student got wrong, and mostly add noise on answers it already got right. The best kind of hint also depends on the student's size and on the task, so there is no single recipe that wins everywhere.

TL;DR

• The paper defines an "ideal" per-token gradient from estimated success probabilities and scores any teacher configuration by its cosine similarity to that gradient §2.2 §2.4.
• Targeted resampling makes the estimate affordable for long reasoning traces §2.5.
• Teacher guidance aligns better with success on failing rollouts than on correct ones §4.1 fig. 3.
• The best teacher changes with student size and task, so the paper argues for per-task diagnostics rather than a fixed recipe §4.2 §4.4.

Problem

On-policy distillation gives a dense, per-token training signal from a teacher, in place of the sparse outcome reward used by RL methods such as GRPO §1. It is unclear which teacher to use, and, for self-distillation, which context to show the student-as-teacher §1. Answering these questions usually needs full training runs, and aggregate metrics hide what happens at individual tokens §1. The paper sets out to measure, per token, per question and per teacher configuration, whether the teacher's guidance points toward correct answers §1.

Method

The method has three parts: a reference ideal gradient, the distillation gradients it is compared with, and a rollout procedure that estimates the reference at scale.

Generation tree and success estimates. The student samples G rollouts per question. Rollouts that share a prefix form a tree. At each node u, the empirical success probability \hat{P}_{\mathrm{succ}}^k is the share of rollouts that reach a correct answer after taking token k at u §2.1.

Teacher advantage. For a node u with a set \mathcal{S}_u of sufficiently visited tokens, the advantage is the teacher's expected success under its own probabilities minus the student's, with both renormalised over \mathcal{S}_u §2.1 eq. 2.1. A positive value means the teacher favours tokens that succeed more often §2.1.

Ideal reference gradient. The ideal local objective is the student's expected success at the node eq. 2.2. Differentiating with respect to a logit z_j with the softmax Jacobian gives P_\theta^j \, (P_{\mathrm{succ}}^j - \bar{P}_{\mathrm{succ}}), where \bar{P}_{\mathrm{succ}} is the student's expected success at u §2.2 eq. 2.3. The paper computes this gradient directly from the empirical estimates rather than from per-rollout advantages, which it reports as more accurate and lower in noise §2.2.

Dr. GRPO connection. The paper argues that Dr. GRPO (GRPO without the 1/|o_i| length normalisation) yields, in expectation, a gradient proportional to the ideal one §2.2 eq. 2.4 eq. 2.5. The derivation drops the KL penalty by taking \beta = 0 and uses batch-level constants app. B.2.

Distillation gradients. Each method is written as a per-node gradient of one common form eq. 2.11:

• GKD minimises \mathrm{KL}(\pi_\theta \| \pi_{\mathrm{te}}) at each node eq. 2.6. Its gradient is P_\theta^j(\ell_j - \bar{\ell}) with \ell_k = \log P_\theta^k - \log P_{\mathrm{te}}^k eq. 2.7 eq. 2.8.
• The single-sample estimator of the Thinking-Lab paper uses only the sampled token, with a weight of \log P_{\mathrm{te}}^{r_i} - \log P_\theta^{r_i} - 1. In expectation it gives the GKD gradient with the opposite sign eq. 2.9 app. B.4.
• MiniLLM uses a REINFORCE-style form whose per-node factor depends on the full downstream trajectory eq. 2.10 app. B.5.

Gradient alignment score. At node u, the score is the cosine between the ideal gradient and the distillation gradient, both restricted to \mathcal{S}_u §2.4 eq. 2.12. It lies in [-1, +1]: positive means the distillation gradient pushes toward success, near zero means it is orthogonal, and negative means it pushes toward failure §2.4. The experiments score GKD only §3 app. E.5.

Targeted rollouts. Sparse visits are the main obstacle. The method launches targeted rollouts from a prefix to an under-visited node and completes them to the end of the response. Each rollout updates every node along its path, so the total budget grows more slowly than sequence length §2.5.

• Depth windows grow exponentially, for example tokens 1–50, 51–150, 151–350 §2.5. Each window gets a fixed budget of tokens chosen by GKD gradient magnitude and by the student–teacher probability difference §2.5 app. D.
• Tokens count as important when P_\theta^k > \tau or P_{\mathrm{te}}^k > \tau, with \tau = 0.02 §2.5. Each important token is enriched to N_{\min} = 100 visits §2.5.
• A node is kept only when its children have at least N_{\mathrm{sig}} visits: 20 in general, 40 for AIME §2.5 app. A.1.
• The generation tree and \hat{P}_{\mathrm{succ}} do not depend on the teacher. One enriched tree is shared across all teachers, which the paper says cuts compute by about 7× §2.5.

Experiments. Students are Qwen3-0.6B and Qwen3-1.7B. There are eight teacher configurations: five self-distillation variants (the student with one correct solution, a summarised correct one, one correct plus one wrong, both summarised, and a Qwen3-32B-generated correct solution) and three external teachers (Qwen3-4B, 8B, 14B) §3. Datasets are BoolQ and MMLU, with AIME 2025 case studies §3 §4.4.

Pseudocode. The three quantities below follow the paper's equations.

# Ideal gradient on student logits at node u  [§2.2] [eq. 2.3]
# Input:  student_probs (|S_u|,), success (|S_u|,)  Output: (|S_u|,)

def ideal_gradient(student_probs, success):
    # Student's expected success at this node  [§2.2]
    mean_success = sum(student_probs * success)
    # Tokens that succeed more often than average get pushed up  [§2.2]
    return student_probs * (success - mean_success)

# GKD update direction at node u, negated loss gradient  [§2.3] [eq. 2.8]
# Input:  student_probs (|S_u|,), teacher_probs (|S_u|,)  Output: (|S_u|,)

def gkd_direction(student_probs, teacher_probs):
    # Per-token log-ratio of student to teacher  [§2.3] [eq. 2.6]
    log_ratio = log(student_probs) - log(teacher_probs)
    # Expected log-ratio under the student, which equals the KL  [§2.3]
    mean_log_ratio = sum(student_probs * log_ratio)
    # Sign flipped, as the paper writes it  [§2.3]
    return -student_probs * (log_ratio - mean_log_ratio)

# Gradient alignment score at node u  [§2.4] [eq. 2.12]
# Input:  two vectors over S_u  Output: scalar in [-1, 1]

def alignment_score(ideal, distill):
    # Cosine similarity, both vectors restricted to S_u  [§2.4]
    dot = sum(ideal * distill)
    return dot / (norm(ideal) * norm(distill))

Hyperparameters

The paper trains no model with this method. Values below are for the diagnostic runs. Anything not in this table is not stated in the paper.

Results

Author-reported. Setup: BoolQ and MMLU, Qwen3-0.6B and 1.7B, eight teachers. Values are mean gradient alignment (cosine) with 95% CIs across questions tab. 1. The overall mean is about +0.027 (0.6B) and +0.026 (1.7B) on BoolQ, with a per-token standard deviation of roughly 0.83–0.91 §4.

Source: tab. 1 §4.2. Bold in the paper marks the top teacher per column: Self-1C (0.6B BoolQ), Self-1C (32B) (0.6B MMLU), Self-Sum-1C and Qwen3-8B (1.7B BoolQ, which is 0.050 vs 0.053), and Self-Sum-1C (1.7B MMLU).

Correct vs incorrect rollouts. Alignment is higher on incorrect paths §4.1 fig. 3 app. E.6. Weighted cosine is the cosine weighted by the success-rate range at each node.

Source: fig. 3 app. E.6.

Teacher-advantage. The paper also reports a teacher-advantage column in app. E.6. It is larger on incorrect paths in all four settings, with Δ between −0.004 and −0.013; the per-path means lie between −0.002 and 0.011.

Oracle selective distillation. Keeping only tokens with positive alignment raises the mean signal 10–15× while keeping about half the tokens app. E.3. For 0.6B BoolQ, full GKD gives 0.007; selecting at threshold 0 gives 0.070 with 51.6% of tokens; threshold 0.3 gives 0.076 with 46.6% app. E.3. The paper calls this an oracle, since the true alignment is unknown at training time.

Teacher-ranking robustness. On BoolQ, the fraction of positive tokens and the weighted cosine give the same 0.6B ordering. For 1.7B the weighted cosine puts Self-Sum-1C (0.088) above Qwen3-8B (0.072) app. E.4 §4.2.

AIME 2025 case studies. Four questions, two per student §4.4 app. A.1.

Source: app. A.1. On Q0 the best teacher gives +0.097 on incorrect rollouts and −0.011 on correct ones (Δ = −0.108) §4.4 app. A.1. Per the paper, the incorrect-over-correct pattern replicates on all four questions §4.4.

Correlations within a path. Divergence measures (KL, JS, L2) correlate positively with alignment, and cosine similarity of the two distributions correlates negatively. All magnitudes are small, with |\rho| \lesssim 0.04 on BoolQ and MMLU §4.3 app. E.7.

Paper vs code

No official code.

Searched, in the order the skill sets:

1. Paper and abstract: no code link. The arXiv abstract page lists no repository, only the CC BY 4.0 license. The LaTeX source contains only a commented-out template placeholder for a repository [sample.tex line 32].
2. Project page: none linked in the paper.
3. GitHub repository search: queries for the title and "gradient alignment" with "on-policy distillation" returned one repository, owned by a third party (pleyva2004), with no license. It is not an author or an author's lab, so it was not used. Searches for org:apple returned nothing.

No discrepancies between code and paper were checked because no code exists.

Internal inconsistencies in the paper, which an implementer should know about:

• The introduction says the alignment-versus-correctness pattern "holds across all settings and metrics" §1. Table app. E.6 gives a mean-cosine p-value of 0.123 for 1.7B MMLU, and §4.1 itself calls that gap not significant.
• The introduction calls Dr. GRPO an unbiased estimator of the ideal gradient §1. Section §2.2 and app. B.2 only show proportionality, with constants and the \beta = 0 assumption.
• Within-path correlations are described as consistent across all settings §4.3. The same section then says the divergence pattern weakens for 1.7B on BoolQ.

Pitfalls

• Sign convention for the distillation gradient. The ideal gradient is the ascent direction of success. The paper writes the GKD gradient as the loss gradient P_\theta^j(\ell_j - \bar\ell) eq. 2.8 and says the distillation gradient, with a sign flip, is -P_\theta^j(\ell_j - \bar\ell) §2.3. The alignment score is only meaningful with both vectors in the same convention. The paper does not spell out which one the score uses. My reading, not stated by the paper: the negated form is the one under which "helpful means positive" §2.4 holds, because the ideal gradient is an ascent direction. Check this on a constructed case with a known sign.
• KL orientation. The paper calls the GKD objective "forward KL from student to teacher" but writes it as \mathrm{KL}(\pi_\theta \| \pi_{\mathrm{te}}) §2.3 eq. 2.6. Implement the expression as written, with the student as the first argument.
• Renormalisation is not specified for the score. The advantage is renormalised over \mathcal{S}_u §2.1. The score is restricted to \mathcal{S}_u §2.4, but the paper does not say whether the student and teacher probabilities used in the gradients are renormalised over \mathcal{S}_u or taken from the full vocabulary.
• Undefined cosine. When either gradient vector is zero, the cosine is undefined. The paper only uses nodes with a nonzero success-rate range §3, and does not say how other nodes are handled.
• Visit thresholds. Nodes need N_{\mathrm{sig}} visits per child, and tokens never sampled by the student cannot be scored app. C. A node can pass the threshold while its estimate of \hat{P}_{\mathrm{succ}} is still noisy, so report counts with any aggregate.
• Teacher order affects comparisons. The tree is shared across teachers, but rollout enrichment uses each teacher's logits, and later teachers benefit from earlier enrichment app. D. Comparing teachers therefore depends on the order in which they are processed. The paper does not report how much this matters.
• Noisy per-token signal. Alignment swings in sign from token to token along one path app. E.2, and the spread is very large §4. Aggregate per path, question or teacher, as the paper does, rather than reading single tokens.
• Small AIME sample. The AIME study uses four questions, two per student §4.4. Treat the teacher ranking there as indicative only.
• Oracle, not method. The selective-distillation gain uses true alignment app. E.3. The divergence filter in §4.3 is only a suggested proxy, and the paper says it is weak.
• Compute. Each question needs roughly 45K–200K targeted rollouts §2.5 §3, and the full suite about 72 A100-days §3. The authors present the method as an offline evaluation tool, not a training signal app. C.
• Unstated values. Student sampling temperature for the main runs, maximum response length and any training optimiser settings are not stated in the paper. The screening temperature of 1.0 app. E.9 applies only to the screening pipeline.
• Context composition. The self-distillation context is prepended to the user message before the chat template is applied app. E.9. Summaries are written by Qwen3-32B §3 app. E.9.
• Version. This card follows arXiv v2 [raw metadata].

Lineage

• Builds on: Dr. GRPO, "Understanding r1-zero-like training" (arXiv:2503.20783), for the link between the reward gradient and the per-node ideal gradient §2.2 app. B.2. GKD (Agarwal et al., ICLR 2024), the single-sample estimator in the Thinking-Lab paper "On-Policy Distillation", and MiniLLM (Gu et al., ICLR 2024) are the distillation objectives it analyses §2.3. The bibliography gives no arXiv ids for these three.
• Compared against: the paper compares eight teacher configurations within its own framework §3. It contrasts its findings with related work on on-policy distillation and self-distillation §5.
• Superseded by: not stated.