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.10889v1 [cs.LG] 11 May 2026
Apple ∗ Equal contribution
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. Correspondence: Mohammadreza Armandpour, Mehrdad Farajtabar {marmandpour, farajtabar}@apple.com Date: May 12, 2026
1
Introduction
On-policy distillation has rapidly become a core post-training technique for reasoning models: Qwen3 (Yang et al., 2025), MiMo (Xiao et al., 2026), and GLM-5 (Zeng et al., 2026) all adopt it in their pipelines, and multiple concurrent works (Hübotter et al., 2026; Zhao et al., 2026; Ye et al., 2026; Shenfeld et al., 2026) demonstrate strong gains from self-distillation variants, establishing it as a practical and compute-efficient complement to sparse-reward RL. The idea is simple: guide the student at every token using a teacher’s distribution (Agarwal et al., 2024; Thinking Lab, 2025). In teacher distillation, a larger model provides supervision (Hinton et al., 2015). In self-distillation, the student serves as its own teacher with extra context (such as a correct solution) unavailable at test time. Both complement the sparse binary reward of RL methods like GRPO (Shao et al., 2024; DeepSeek-AI, 2025) with a dense gradient at every token. Yet practitioners face a series of decisions with no principled guidance: Should the teacher be a larger external model, or the student itself with access to a correct solution? Should the context include a full solution trace or a concise summary? Does the answer depend on the question? On the token? Today, these questions require expensive training runs whose aggregate metrics hide what happens at the level of individual tokens. Our objective was to develop a more rigorous methodology: a framework capable of assessing, at the finest feasible level of granularity (per token, per question, per teacher configuration), the extent to which the teacher’s guidance is aligned with the behaviors that yield correct answers. Figure 1 demonstrates that, even 1
Problem: A bookshelf has 3 shelves. Each shelf holds 4 books. If 5 books are removed, how many remain? Answer: 7 Student token Teacher suggestion student answer:
3 shelves
four
···
therefore
notation (stylistic)
4
books
···
×4=12.
word choice (stylistic)
each
···
so,
total is
operator choice (reasoning-critical)
3
+4=7.
=7 ✓
12−5:
=7 ✓
error recovery (post-hoc)
Rem. 5:
=2 ×
Figure 1 Not all teacher guidance points toward success. Generation tree for a bookshelf problem. The teacher’s
distribution disagrees with the student at four branching points (orange), but not all disagreements are equal: some reflect stylistic preferences (“four” vs. “4”, “therefore” vs. “so,”) rather than reasoning corrections. Standard distillation treats all four signals equally, mixing noise from irrelevant preferences with updates that actually matter.
within an individual reasoning trajectory, the teacher’s points of disagreement comprise a heterogeneous mixture of beneficial, neutral, and detrimental contributions, which cannot be reliably differentiated without explicitly linking each token to its downstream effects. To evaluate teacher guidance quality at each token, we derive an ideal per-token gradient from empirical success probabilities: the direction that maximally improves the student’s chance of reaching a correct answer. We show that Dr. GRPO (Liu, 2025) recovers this gradient in expectation, making it an unbiased estimator of the ideal (Section 2.2). We further show that major distillation objectives (GKD (Agarwal et al., 2024), the single-sample estimator of Thinking Lab (2025), MiniLLM (Gu et al., 2023)) produce gradients with the same local structure: for reward-based methods the signal comes from success probability, for distillation methods it comes from the teacher’s distribution. To estimate the ideal gradient scalably even for long reasoning chains, we design a targeted-rollout algorithm with exponential depth windows whose compute scales with a user-chosen budget rather than sequence length. The gradient alignment score (cosine similarity between the ideal and the distillation gradient at each token) then evaluates how well any teacher configuration approximates the ideal, offline (Section 2.4). Applying this framework to Qwen3-0.6B and Qwen3-1.7B across 8 teacher configurations on BoolQ and MMLU, we find that: Key findings.
• Distillation guidance is more reliable on incorrect rollouts. When the student is already on a correct path, the teacher’s signal becomes noisy and weakly aligned with the ideal; on failing rollouts, the teacher reliably pushes toward success. This holds across all settings and metrics. • Context design and student capacity interact strongly. In self-distillation, the form of context shown to the student-as-teacher matters: a summarized solution nearly doubles alignment for 1.7B compared to the raw trace, but slightly hurts 0.6B (which needs full step-by-step reasoning). A 32B-generated solution helps 0.6B on simple tasks but fails on hard math where the reasoning style becomes foreign. External teachers outperform self-distillation only for the larger student. We hypothesize that comprehensibility is the underlying factor: the gradient signal is only useful if the student can parse what it is given. • No universal recipe exists. Among self-distillation variants, contrastive examples (correct + wrong) hurt on simple reasoning but help on hard math. Comparing self-distillation to external teachers, external teachers outperform for 1.7B on BoolQ but not on MMLU. Which teacher or context achieves the highest alignment shifts with question difficulty, motivating per-task diagnostics rather than fixed pipelines. • Divergence predicts alignment, but weakly. Within-path correlations show that divergence between student and teacher prediction distributions (KL, JS, L2 ) is positively associated with alignment while their similarity (cosine of probability vectors) is negatively associated, consistently across all settings. Magnitudes are small (|ρ| < 0.05), indicating divergence as a cheap necessary-condition filter but not a reliable predictor. We further test these patterns on AIME 2025 as case studies (Section 4.4); the finding that incorrect rollouts exhibit higher alignment replicates, but the best self-distillation context changes: including a wrong demonstration, which hurts on short-reasoning tasks, produces the highest alignment on hard math problems.
2
2
Methodology
2.1
Not all teacher guidance is useful
At each token position, the teacher’s distribution may differ from the student’s for many reasons: it may prefer a stylistic variant, it may encourage the student along a productive reasoning path, or it may redirect the computation entirely toward a different continuation. The core difficulty is that one cannot distinguish these cases from the teacher’s probability alone. A token where the teacher and student assign substantially different probabilities could reflect any of these, and only some improve the student’s chance of reaching a correct answer (cf. Figure 1). To tell them apart, we need to connect the teacher’s token-level signal to downstream outcomes. We do this by decomposing the generation process into a generation tree: given G trajectories sampled from the student πθ on a prompt q, each trajectory shares prefixes with others, forming a tree where each node u corresponds to a token position and each edge corresponds to a next-token choice. At each node u, we observe which next tokens were chosen across rollouts and which of those rollouts ultimately reached a correct k : the probability of reaching a answer. This gives us an empirical estimate of the success probability P̂succ correct answer after choosing token k at node u. With this quantity in hand, we can ask precisely: does the teacher push probability mass toward high-Psucc tokens, or away from them? Independent of any training algorithm, a teacher is good at node u if it places more mass on success-leading tokens than the student does. We define the teacher advantage: X X k k k Adv(u) = Pte P̂succ − Pθk P̂succ , (2.1) k∈Su
k∈Su
k where Su is the set of tokens with sufficient visit counts at node u, Pte is the teacher’s probability of token k, k Pθ is the student’s probability of token k, and probabilities are renormalized over Su . A positive advantage means the teacher “knows better” at this node; a negative advantage means following it would hurt.
But a good teacher is not sufficient: you also need an algorithm that translates the teacher’s knowledge into a useful gradient, and different algorithms (GKD (Agarwal et al., 2024), single-sample estimators (Thinking Lab, 2025), MiniLLM (Gu et al., 2023)) use the teacher differently, producing very different gradients from the same teacher. To evaluate any (teacher, algorithm) pair, we need an ideal reference: the gradient that would optimally improve the student’s success probability at each node.
2.2
The ideal reference gradient
At each node u, the ideal local objective is to maximize the student’s probability of reaching a correct answer from this point: X k Lideal (u) = Pθk Psucc . (2.2) k
This is the expected success rate under the student’s current distribution at node u. The gradient of this objective with respect to the student’s logit zj at this node is obtained via the softmax Jacobian ∂Pθk /∂zj = Pθk (δkj − Pθj ), where δkj is the Kronecker delta (1 if k = j, 0 otherwise): X ∂Lideal k = Psucc · Pθk (δkj − Pθj ) ∂zj k X j k = Pθj Psucc − Pθj Pθk Psucc = Pθj
k j Psucc − P̄succ ,
(2.3)
k where P̄succ = k Pθk Psucc is the student’s current expected success at this node. This gradient increases logit zj when token j leads to success more often than average, and decreases it otherwise. This is our reference: the direction in which the student’s logits should move to maximally improve its chance of success at this node.
P
3
A natural question is whether any existing training objective already computes this ideal gradient. Dr. GRPO (Liu, 2025) is a variant of GRPO that removes the pertrajectory length normalization 1/|oi |. The full GRPO objective includes an importance ratio πθ /πθold , a KL penalty, and division by trajectory length. After marginalizing the importance ratio, dropping Pthe KL penalty k (small β), and removing length normalization, the expected objective at node u reduces to k Pθk · Psucc up to constants independent of θ (see Appendix B for the full derivation). The empirical per-sample gradient at node u is: 1 X ∂ L̂DrGRPO = Ai δri ,j − Pθj , (2.4) ∂zj Nu i: o ∋u u Dr. GRPO recovers this gradient in expectation.
i
where Ai = (Ri − R̄)/std(R) is the normalized advantage, ri is the token chosen by rollout i, and Nu is the number of rollouts through u. In expectation, this is proportional to the ideal gradient (Equation 2.3): j E ∇zj LDrGRPO u ∝ Pθj (Psucc − P̄succ ). (2.5) This connection motivates using Equation equation 2.3 as our oracle reference: it is what reward-based training would converge toward at each node, given sufficient rollouts. Standard GRPO’s 1/|oi | factor couples the advantage to trajectory length, preventing this clean per-node decomposition. k estimates at each node, not from In practice, we compute the ideal gradient directly from empirical P̂succ per-sample advantage terms. This is more accurate than the finite-sample Dr. GRPO estimator and avoids the noise of individual trajectory rewards.
2.3
Distillation gradients
We now derive the gradient that each distillation algorithm produces at node u. GKD (Generalized Knowledge Distillation).
to teacher at each node:
GKD (Agarwal et al., 2024) minimizes the forward KL from student
LGKD (u) = KL(πθ ∥πte ) =
X
k Pθk log Pθk − log Pte
(2.6)
k k , the gradient is: Defining the per-token log-ratio ℓk = log Pθk − log Pte
X ∂ ∂LGKD Pθk ℓk = ∂zj ∂zj k X X ∂ℓk = ℓk · Pθk (δkj − Pθj ) + Pθk · ∂zj k
The second sum contributes
(2.7)
k
j j j k k Pθ · (δkj − Pθ ) = Pθ − Pθ = 0 (the softmax Jacobian sums to zero), so:
P
∂LGKD = Pθj (ℓj − ℓ̄) ∂zj
(2.8)
P where ℓ̄ = k Pθk ℓk = KL(πθ ∥πte ). Since we minimize this KL, the distillation gradient (with sign flip) is −Pθj (ℓj − ℓ̄), which pushes logits toward tokens where the teacher assigns relatively higher probability. Thinking Lab (2025) propose an importance-weighted estimator requiring only the sampled token. For rollout i choosing token ri at node u, the per-sample gradient is: Single-sample GKD estimator.
ri ∇zj ℓi = (log Pte − log Pθri − 1) · (δri ,j − Pθj )
(2.9)
In expectation this recovers −Pθj (ℓj − ℓ̄), the GKD gradient with opposite sign (the −1 baseline vanishes; see Appendix B).
4
MiniLLM.
MiniLLM (Gu et al., 2023) uses a REINFORCE-style gradient with trajectory-level reward-to-go: X o o ∇zj ℓi = −(δri ,j − Pθj ) · log Ptet′ − log Pθ t′ − 1 (2.10) t′ ≥t
This couples the gradient at node u to all downstream nodes. The local gradient still takes the form Pθj (fj − f¯) in expectation, but fk is now trajectory-dependent rather than purely local (see Appendix B). Summary.
All four methods produce per-node gradients of the form: ∂L = Pθj (fj − f¯), ∂zj
f¯ =
X
Pθk fk
(2.11)
k
k k with fk = Psucc for Dr. GRPO, fk = ±(log Pθk − log Pte ) for GKD (and its single-sample estimator), and a trajectory-dependent reward-to-go for MiniLLM. Because they share this structure, we can compare their directions via cosine similarity. A consequence of the shared Pθj factor is that the gradient magnitude for any token is gated by the student’s current probability: even if the teacher identifies a high-success token, the update is small when Pθj is small. The teacher can amplify tokens the student already partially believes in, but has limited ability to inject entirely new continuations in a single step.
2.4
The gradient alignment score
We define the gradient alignment score at node u as the cosine similarity between the ideal gradient (Equation 2.3) and the distillation gradient (e.g., GKD): P ideal D gj j∈S gj ideal D (2.12) Align(u) = cos gu , gu = ideal u ∥g ∥Su ∥gD ∥Su k where guideal is the ideal gradient computed directly from empirical P̂succ and guD is the distillation gradient vector, both restricted to Su (the set of tokens with sufficient visit counts at node u). The restriction is k is only reliably estimated for tokens that have been sampled enough times. necessary because Psucc
The score ranges from −1 to +1: • Align(u) > 0: the distillation gradient pushes toward tokens that lead to success. The teacher + algorithm combination is helpful at this node. • Align(u) ≈ 0: the distillation gradient is orthogonal to the reward signal. The teacher’s guidance is neither helpful nor harmful; it wastes gradient budget on irrelevant directions (e.g., stylistic preferences). • Align(u) < 0: the distillation gradient pushes toward tokens that lead to failure. The teacher + algorithm combination is actively harmful at this node. This score answers the question posed in Section 2.1: it distinguishes reasoning-critical disagreements (positive or negative alignment) from stylistic ones (near-zero alignment) at each token position, without requiring any training.
2.5
Computing the score at scale
k The alignment score (Equation 2.12) requires reliable estimates of Psucc at each branching node (Figure 2 summarizes the three-step computation). Naïvely, this would require thousands of rollouts through every possible next token at every node, clearly infeasible for sequences of hundreds of tokens with vocabularies of 150K.
The core challenge is sparsity: given G initial rollouts, most tokens at most nodes receive zero visits, and deep branching points may have only 1–2 rollouts passing through them. To address this, we generate targeted rollouts: given a node u and a token k that needs more visits, we construct a prefix (prompt + path to u + 5
○ 2 Teacher forward pass
○ 1 Estimate P̂succ and gideal <think>
,
Okay
so · · ·
○ 3 Compute Align(u) ∇θ L at u
= 15 teacher πte
gideal α
so
but
, wait
···
···
···
✓ ✓ ✓· · · ✓ ×
✓ ×··· ✓ × ×
× × ✓· · · × ×
P̂succ = 0.75
P̂succ = 0.40
P̂succ = 0.11
k log Pte
k ℓk = log Pθk − log Pte
GKD: full vocab Single-sample: sampled token only
j j gideal ≈ Pθj P̂succ − P̄ˆsucc
k P̂succ = Suk Nuk
gdistill
Align(u) = cos α > 0: KD aligned with ideal < 0: in conflict
gjKD = Pθj ℓj − ℓ̄
k Figure 2 Computing the gradient alignment score at a branching node u. (1) Student rollouts yield empirical P̂succ
per branch and the ideal gradient. (2) A teacher forward pass produces the distillation gradient. (3) Their cosine similarity measures whether the teacher pushes toward success (> 0) or against it (< 0).
token k) and sample completions from the student to the end of the response. Each targeted rollout enriches not only the target node but all ancestors and descendants along its path, so a single rollout launched at depth d simultaneously improves Psucc estimates at every node it passes through. This cascading effect means the total budget required grows sublinearly with sequence length. Rather than allocating rollouts uniformly across the sequence, we partition the generation into exponentially growing depth windows (e.g., tokens 1–50, 51–150, 151–350, . . . ). Within each window, we allocate a fixed budget of k tokens ranked by GKD gradient magnitude and r tokens ranked by student–teacher probability difference, prioritizing tokens where the teacher disagrees most strongly. Early windows are small and densely sampled (where branching is frequent); later windows are larger and more sparsely sampled (where reasoning chains have committed to a direction). This mirrors the natural structure of generation trees: branching diversity decreases with depth as trajectories converge. Exponential depth windows.
k > τ (τ = 0.02), i.e., those that contribute We target tokens where Pθk > τ or Pte meaningfully to the gradient, and enrich each to Nmin = 100 visits. Nodes with ≥ Nsig = 20 visits per child are retained for the alignment computation; for longer traces (AIME) where estimates are noisier, we use Nsig = 40. The total compute scales with the user-chosen budget (number of windows × per-window budget) rather than with sequence length, making the method applicable to traces ranging from ∼200 tokens (BoolQ) to ∼30K tokens (AIME) without modification. In practice, each question requires ∼45K–200K targeted rollouts depending on trace length.
Budget and scalability.
k Teacher-independent tree sharing. A key efficiency insight is that the generation tree and Psucc estimates are teacher-independent: they depend only on the student’s rollouts and outcomes. We share a single enriched tree across all 8 teacher configurations: rollout generation is done once (Phase 1), and each teacher requires only one additional forward pass to compute its gradient and alignment score (Phase 2). This reduces total compute by ∼7× compared to independent runs. Details on rollout prioritization are in Appendix D.
3
Experimental Setup
Student models.
and Qwen3-1.7B.
We evaluate two student scales from the Qwen3 family (Yang et al., 2025): Qwen3-0.6B
Teacher configurations.
For each student, we evaluate 8 teacher configurations spanning two families:
Self-distillation (teacher = same model with enriched context): Self-1C (1 correct solution in context), Self-Sum-1C (correct solution summarized by Qwen3-32B), Self-1C1W (1 correct + 1 wrong solution), Self-Sum-1C1W (both summarized), Self-1C (32B) (correct solution generated by Qwen3-32B shown to student-as-teacher). 6
Mean cosine
Fraction positive 0.600
0.15
0.3
0.575
0.10
Value
Weighted cosine (SR)
0.2
0.550
0.05
=0.044
=0.011
=0.525
0.525
0.00 zero
0.500 chance
0.05
0.475
Incorrect (n=576)
Correct (n=576)
0.425
=0.110
0.1
=0.046
0.0 zero
0.450
0.10
=0.507
0.1
Incorrect (n=576)
Correct (n=576)
Incorrect (n=576)
Correct (n=576)
Correct Incorrect
p-value
Qwen3-0.6B, BoolQ Mean cosine .011 Weighted (SR) .046
.044 .110
7×10−8 8×10−10
Qwen3-1.7B, BoolQ Mean cosine .001 Weighted (SR) .010
.058 .093
2×10−9 7×10−9
Qwen3-0.6B, MMLU Mean cosine .009 Weighted (SR) .021
.048 .118
.0001 < 10−11
Qwen3-1.7B, MMLU Mean cosine .012 Weighted (SR) .034
.028 .098
.123 < 10−5
Figure 3 Gradient alignment on correct vs. incorrect paths. Left: distribution for Qwen3-0.6B on BoolQ; the teacher’s
gradient is more aligned on paths leading to wrong answers. Right: the pattern holds across all settings under both mean cosine and SR-weighted cosine. Full results in Appendix Table 4.
External teachers (larger models, same prompt as student): Qwen3-4B, Qwen3-8B, Qwen3-14B. Datasets. We evaluate on two benchmarks: BoolQ (Clark et al., 2019), a reading comprehension task with True/False answers and short reasoning chains; and MMLU (Hendrycks et al., 2021), a multiplechoice knowledge benchmark with medium-length chains. We additionally present case studies on AIME 2025 (Mathematical Association of America, 2024) (∼5K–30K token traces) in Section 4.4. Each question requires substantial compute: G=200 initial rollouts, 4 representative paths (2 correct, 2 incorrect), and ∼45K–200K targeted rollouts at undersampled branching points (totaling ∼72 A100-days for the full experiment suite). Each important token receives up to 100 targeted samples; nodes with ≥20 visits are considered statistically significant for computing P̂succ .
At each branching node with ≥2 children having ≥20 visits and nonzero success-rate range, we compute: gradient alignment (ideal vs. GKD cosine), teacher advantage, and success rate statistics. We aggregate per path (mean cosine along the path), per question (correct/incorrect split), and per teacher (means with 95% CIs across questions). Metrics.
4
Results
We present results across two datasets and two model scales, totaling ∼88K decision points for BoolQ (0.6B) and ∼81K for BoolQ (1.7B), with MMLU providing ∼49K (0.6B) and ∼46K (1.7B). Overall, gradient alignment is weakly positive (mean cosine +0.027 for 0.6B, +0.026 for 1.7B on BoolQ) but with enormous per-token variance (std ∼0.83–0.91; see Appendix E.1 for the full distribution).
4.1
Distillation helps more on incorrect paths
Our most consistent finding across all settings is that incorrect paths exhibit significantly higher gradient alignment than correct paths (Figure 3). On incorrect paths, the reward gradient points away from the current (failing) trajectory, and the teacher (which generally prefers tokens leading to success) pushes in the same direction. On correct paths, the student is already succeeding, so the reward gradient is weaker and the teacher’s contribution is less aligned. The effect is strongest for 1.7B on BoolQ (∆ = −0.056, p < 10−9 ); even on MMLU where the mean cosine gap is not significant (p = 0.12), the weighted cosine is highly significant (p < 10−5 ).
4.2
The best teacher depends on student capacity
A striking result emerges when comparing teacher rankings across model scales (Figure 4, Table 1). For the 0.6B student, self-distillation teachers using correct-only demonstrations (Self-1C, Self-Sum-1C, Self-1C (32B)) consistently achieve 2–3× higher alignment than external teachers, on both BoolQ and 7
Gradient alignment by teacher
Fraction of aligned decision points
Self-1C (32B)
+0.052
Self-Sum-1C
+0.047
Self-1C
+0.040
Self-Sum-1C1W
+0.037
Self-1C1W
0.529 0.533 0.520 0.524
+0.027
0.514
Qwen3-14B
+0.015
0.507
Qwen3-8B
+0.011
0.505
Qwen3-4B
+0.000
-0.02
0.00
0.02
0.04
0.501
0.06
Mean cosine similarity
0.08
0.48
chance
0.50
0.52
0.54
Fraction of aligned decision points
0.56
Figure 4 Teacher ranking by gradient alignment (Qwen3-0.6B, MMLU). Self-distillation with correct demonstrations
dominates. Additional plots for BoolQ and 1.7B in Figure 11. Table 1 Teacher ranking by mean gradient alignment (with 95% CI) across datasets. Self-distillation dominates for
0.6B; external teachers become competitive for 1.7B. Qwen3-0.6B Teacher
Qwen3-1.7B
BoolQ
MMLU
BoolQ
MMLU
Self-1C Self-1C (32B) Self-Sum-1C Self-Sum-1C1W Self-1C1W
.047 ± .021 .047 ± .022 .041 ± .025 .016 ± .025 .019 ± .018
.040 ± .028 .052 ± .025 .047 ± .032 .037 ± .031 .027 ± .037
.028 ± .024 .020 ± .034 .050 ± .030 .008 ± .029 .002 ± .027
.010 ± .024 .021 ± .030 .036 ± .025 .034 ± .033 .009 ± .022
Qwen3-14B Qwen3-8B Qwen3-4B
.016 ± .021 .018 ± .022 .018 ± .024
.015 ± .032 .011 ± .032 .000 ± .026
.036 ± .035 .053 ± .028 .040 ± .024
.017 ± .035 .014 ± .037 .017 ± .041
MMLU. But for the 1.7B student on BoolQ, external teachers, particularly Qwen3-8B, achieve the highest alignment, outperforming all self-distillation variants. We interpret these findings through the lens of context comprehensibility: self-distillation helps only when the student can understand the context it is given. A small student (0.6B) cannot effectively absorb the full distributional knowledge from a much larger external model; the teacher’s reasoning patterns are too different from its own, making the gradient signal incomprehensible. Showing the same 0.6B model a correct solution in its own reasoning style (self-distillation) provides a targeted, understandable nudge toward success tokens. For the larger 1.7B student, the capacity gap to an 8B teacher is smaller, and the genuinely different knowledge encoded in the larger model’s distribution becomes comprehensible and exploitable. An additional finding: including wrong demonstrations hurts on BoolQ and MMLU. The 1C1W variants consistently underperform their 1C counterparts, suggesting that negative examples introduce noise rather than useful contrastive signal on these tasks. Breaking down the correct-vs-incorrect gap per teacher (Appendix E.8), we find that Self-1C for 0.6B is uniquely uniform: it achieves nearly equal alignment on both correct and incorrect paths (∆ ≈ 0), while all other teachers show the typical incorrect > correct pattern. Summarizing demonstrations (Self-Sum-1C) nearly doubles alignment for the 1.7B student (0.050 vs. 0.028 on BoolQ, 0.036 vs. 0.010 on MMLU), but has mixed effects for 0.6B. The interpretation is capacity-dependent: a larger student can extract the key signal from a concise summary, while a smaller student needs the full verbose trace. Similarly, Self-1C (32B) (a Qwen3-32B-generated Summarization helps larger students.
8
solution shown to the student-as-teacher) works well for 0.6B (0.047, 0.052) but poorly for 1.7B and on AIME, as the 32B reasoning style is harder to follow on complex problems. These rankings are largely robust to metric choice (Appendix E.4), though for 1.7B the weighted cosine reverses the top ranking (Self-Sum1C: 0.088 vs. Qwen3-8B: 0.072), suggesting the external teacher’s advantage is concentrated at low-stakes tokens.
4.3
What predicts alignment within a path?
To understand where alignment is positive within a reasoning chain, we compute within-path Spearman correlations between the alignment score and single-rollout features (full results in Appendix E.7). Teacher– student divergence (KL, JS, L2 ) correlates positively with alignment, while distributional similarity correlates negatively: the useful signal lives where the teacher disagrees with the student. However, correlations are uniformly weak (|ρ| ≈ 0.02–0.04): no single feature reliably predicts whether a disagreement is helpful or harmful, since that depends on downstream success probabilities not observable from a single forward pass. The consistent sign nonetheless suggests divergence could serve as a cheap necessary-condition filter for alignment-aware training. Alignment also trends slightly positive with depth (early tokens are templatic; later tokens involve actual reasoning steps).
4.4
Case studies: mathematical reasoning (AIME 2025)
To test generalization to longer reasoning traces (∼5K–30K tokens), we analyze 4 AIME 2025 questions (2 per model) using the same 8 teacher configurations (full results in Appendix A). The core finding (incorrect paths exhibit higher alignment) replicates across all four questions (e.g., best teacher on Q0: incorrect +0.097 vs. correct −0.011, ∆ = −0.108). However, teacher-choice conclusions diverge from BoolQ/MMLU: Self-1C1W (including a wrong demonstration) is the best teacher on the two harder questions, directly contradicting the shorter-reasoning benchmarks where wrong demos consistently hurt. The interpretation is comprehensibility-dependent: on hard math, seeing a common mistake provides useful contrastive signal, whereas on simpler tasks it is merely noise. Additionally, summarized contexts lose to raw demos on hard math (the 0.6B student cannot decipher a compressed summary of a complex argument) but perform well on easier questions. These observations reinforce that no universal distillation recipe exists: the optimal teacher depends on task difficulty, student capacity, and context design.
5
Related Work
Knowledge distillation (Hinton et al., 2015) trains a student on teacher soft distributions; sequence-level variants (Kim and Rush, 2016) generate teacher outputs for training but suffer from exposure bias (Bengio et al., 2015). On-policy distillation (OPD) addresses this by supervising the student on its own rollouts: MiniLLM (Gu et al., 2023) uses reverse KL via policy gradient, GKD (Agarwal et al., 2024) interpolates between on- and off-policy data, and Yang et al. (2026) show the teacher’s log-ratio acts as an implicit dense reward. OPD is now standard in industry pipelines (Yang et al., 2025; Xiao et al., 2026; Zeng et al., 2026), and has been extended to self-distillation settings where a single model serves as its own teacher under privileged conditioning (Snell et al., 2022; Hübotter et al., 2026; Zhao et al., 2026; Ye et al., 2026; Shenfeld et al., 2026; Penaloza et al., 2026). Our work does not propose a new distillation algorithm but provides a diagnostic that measures, at each token, whether the teacher’s signal agrees with the reward objective. Despite OPD’s growing adoption, when and why it fails remains poorly understood. Two concurrent works investigate this question from complementary angles. Li et al. (2026) show that OPD requires thinking-pattern consistency and genuinely new knowledge from the teacher, with success driven by progressive alignment on high-probability overlap tokens (97–99% of mass); they also reveal that reward quality degrades with trajectory depth. Kim et al. (2026) trace self-distillation degradation to the suppression of epistemic verbalization (the model’s expression of uncertainty), showing that richer conditioning contexts suppress uncertainty tokens and harm OOD generalization when task coverage is broad. More broadly, distillation can hurt with overly capable teachers (Cho and Hariharan, 2019; Mirzadeh et al., 2020; Busbridge et al., 2025), and small models struggle to learn from strong reasoners (Li et al., 2025). Our gradient alignment score provides a mechanistic
9
explanation for these phenomena: it directly quantifies, at each token position, whether the teacher’s signal is exploitable by the student, regardless of the teacher’s aggregate performance. On the reward side, GRPO (Shao et al., 2024; DeepSeek-AI, 2025; Liu, 2025) and DAPO (Yu et al., 2025) train reasoning models with sparse outcome rewards, while process reward models (Lightman et al., 2023; Uesato et al., 2022) provide step-level feedback but require separate annotation. Our gradient decomposition unifies these perspectives by showing that reward and distillation objectives share the same local structure (Pθj (fj − f¯)), enabling direct offline comparison at token granularity without training or additional models.
6
Conclusion
We set out to answer a simple question: at each token in a reasoning chain, does the teacher’s distillation signal actually point toward correct answers? We derived an ideal per-node gradient from empirical success probabilities, showed that major distillation objectives share the same local structure, and built a scalable pipeline to compute gradient alignment offline for long sequences. Our experiments reveal that distillation helps most on failing rollouts, that context design interacts strongly with student capacity, and that no universal recipe exists. Beyond these findings, the framework serves as a general-purpose offline testbed for any token-level training algorithm. Our results point to concrete future directions: rollout-weighted distillation that emphasizes failing trajectories, multi-teacher distillation that combines complementary signals from multiple teachers for better overall alignment, adaptive context selection per domain, and divergence-based filters for alignment-aware training.
References Rishabh Agarwal, Nino Vieillard, Piotr Stanczyk, Sabela Ramos, Matthieu Geist, and Olivier Bachem. On-policy distillation of language models: Learning from self-generated mistakes. In International Conference on Learning Representations, 2024. Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. In Advances in Neural Information Processing Systems, volume 28, 2015. Dan Busbridge, Amitis Shidani, Floris Weers, Jason Ramapuram, Etai Littwin, and Russ Webb. Distillation scaling laws. arXiv preprint arXiv:2502.08606, 2025. Jang Hyun Cho and Bharath Hariharan. On the efficacy of knowledge distillation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 4794–4802, 2019. Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. In NAACL, 2019. DeepSeek-AI. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv:2501.12948, 2025. Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang. Minillm: Knowledge distillation of large language models. arXiv preprint arXiv:2306.08543, 2023. Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. In International Conference on Learning Representations, 2021. Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. arXiv:1503.02531, 2015.
Distilling the knowledge in a neural network.
arXiv preprint
Jonas Hübotter, Frederike Lübeck, Lejs Behric, Anton Baumann, Marco Bagatella, Daniel Marta, Ido Hakimi, Idan Shenfeld, Thomas Kleine Buening, Carlos Guestrin, et al. Reinforcement learning via self-distillation. arXiv preprint arXiv:2601.20802, 2026. Jeonghye Kim, Xufang Luo, Minbeom Kim, Sangmook Lee, Dohyung Kim, Jiwon Jeon, Dongsheng Li, and Yuqing Yang. Why does self-distillation (sometimes) degrade the reasoning capability of LLMs? arXiv preprint arXiv:2603.24472, 2026.
10
Yoon Kim and Alexander M Rush. Sequence-level knowledge distillation. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pages 1317–1327, 2016. Yaxuan Li, Yuxin Zuo, Bingxiang He, Jinqian Zhang, Chaojun Xiao, Cheng Qian, Tianyu Yu, Huan-ang Gao, Wenkai Yang, Zhiyuan Liu, and Ning Ding. Rethinking on-policy distillation of large language models: Phenomenology, mechanism, and recipe. arXiv preprint arXiv:2604.13016, 2026. Yuetai Li, Xiang Yue, Zhangchen Xu, Fengqing Jiang, Luyao Niu, Bill Yuchen Lin, Bhaskar Ramasubramanian, and Radha Poovendran. Small models struggle to learn from strong reasoners. In Findings of the Association for Computational Linguistics: ACL 2025, pages 25366–25394, 2025. Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let’s verify step by step. arXiv preprint arXiv:2305.20050, 2023. Zijun Liu. Understanding the grpo and dr. grpo. arXiv preprint arXiv:2503.20783, 2025. Mathematical Association of America. American invitational mathematics examination, 2024. Seyed Iman Mirzadeh, Mehrdad Farajtabar, Ang Li, Nir Levine, Akihiro Matsukawa, and Hassan Ghasemzadeh. Improved knowledge distillation via teacher assistant. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 5191–5198, 2020. Emiliano Penaloza, Dheeraj Vattikonda, Nicolas Gontier, Alexandre Lacoste, Laurent Charlin, and Massimo Caccia. Privileged information distillation for language models. arXiv preprint arXiv:2602.04942, 2026. Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Mingchuan Zhang, YK Li, Y Wu, and Daya Guo. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300, 2024. Idan Shenfeld, Mehul Damani, Jonas Hübotter, and Pulkit Agrawal. Self-distillation enables continual learning. arXiv preprint arXiv:2601.19897, 2026. Charlie Snell, Dan Klein, and Ruiqi Zhong. Learning by distilling context. arXiv preprint arXiv:2209.15189, 2022. Thinking Lab. Thinking tokens for language modeling. arXiv preprint, 2025. On-policy KD variant using importanceweighted single-token estimator. Jonathan Uesato, Nate Kushman, Ramana Kumar, Francis Song, Noah Siegel, Lisa Wang, Antonia Creswell, Geoffrey Irving, and Irina Higgins. Solving math word problems with process- and outcome-based feedback. arXiv preprint arXiv:2211.14275, 2022. Bangjun Xiao, Bingquan Xia, Bo Yang, Bofei Gao, Bowen Shen, Chen Zhang, Chenhong He, Chiheng Lou, Fuli Luo, Gang Wang, et al. Mimo-v2-flash technical report. arXiv preprint arXiv:2601.02780, 2026. An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint arXiv:2505.09388, 2025. Wenkai Yang, Weijie Liu, Ruobing Xie, Kai Yang, Saiyong Yang, and Yankai Lin. Learning beyond teacher: Generalized on-policy distillation with reward extrapolation. arXiv preprint arXiv:2602.12125, 2026. Tianzhu Ye, Li Dong, Xun Wu, Shaohan Huang, and Furu Wei. On-policy context distillation for language models. arXiv preprint arXiv:2602.12275, 2026. Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Weinan Dai, Tiantian Fan, Gaohong Liu, Lingjun Liu, et al. DAPO: An open-source LLM reinforcement learning system at scale. arXiv preprint arXiv:2503.14476, 2025. Aohan Zeng, Xin Lv, Zhenyu Hou, Zhengxiao Du, Qinkai Zheng, Bin Chen, Da Yin, Chendi Ge, Chengxing Xie, Cunxiang Wang, et al. GLM-5: From vibe coding to agentic engineering. arXiv preprint arXiv:2602.15763, 2026. Siyan Zhao, Zhihui Xie, Mengchen Liu, Jing Huang, Guan Pang, Feiyu Chen, and Aditya Grover. Self-distilled reasoner: On-policy self-distillation for large language models. arXiv preprint arXiv:2601.18734, 2026.
11
A
Full AIME 2025 Case Studies
A.1
Overview
To test whether our findings extend to mathematical reasoning (where thinking traces are substantially longer, ∼5K–30K tokens, and reasoning steps more complex), we analyze four AIME 2025 questions: two with Qwen3-0.6B (Q0, pass rate 87.5%; Q3, pass rate 9.4%) and two with Qwen3-1.7B (Q7, pass rate 34.4%; Q28, pass rate 56.3%). Each question uses the same 8 teacher configurations as BoolQ/MMLU. We use sig = 40 (requiring ≥40 visits per child for significance) given the longer chains and sparser branching. Figure 7 shows teacher alignment on two representative questions. Both external and self-distillation teachers achieve positive alignment across all four questions. Self-Sum-1C (summarized correct demonstration) is the strongest teacher on Q0 (+0.043) and Q7 (+0.055), while the raw demo variants lead on Q3 (Self-1C1W: +0.042, Self-1C: +0.040) and Q28 (Self-1C1W: +0.053). Notably, Self-1C1W (which includes a wrong demonstration) is the best teacher on the two harder questions (Q3 and Q28), contrasting sharply with BoolQ/MMLU where it consistently underperforms. Teacher rankings on AIME.
Among teachers with positive alignment, the incorrect > correct pattern holds on all four questions when focusing on high-stakes decision points. On Q0, all positive teachers show higher alignment on incorrect paths (best teacher Self-Sum-1C: incorrect +0.097 vs. correct −0.011, ∆ = −0.108). On Q3, the top-ranked teachers show the same pattern (Self-1C1W: incorrect +0.056 vs. correct +0.027, ∆ = −0.029). On Q28, 7 of 8 positive teachers show incorrect > correct (Self-1C1W: incorrect +0.097 vs. correct +0.008, ∆ = −0.089). On Q7, the pattern emerges clearly at consequential nodes (best teacher SelfSum-1C: incorrect +0.093 vs. correct +0.017, ∆ = −0.076 at high success-rate-range nodes), though the gap is smaller at low-stakes nodes where the signal is noisier. This confirms that even on long mathematical reasoning chains, the teacher’s gradient is most useful on the student’s failing rollouts. Correct vs. incorrect on AIME.
While the correct/incorrect finding transfers, teacher-choice conclusions do not, further supporting the comprehensibility hypothesis. On hard AIME math (Q3), Self-1C1W (which includes a wrong demonstration alongside the correct one) is the best teacher (+0.042), directly contradicting BoolQ/MMLU where wrong demonstrations consistently hurt. On hard math, seeing a common mistake may provide useful contrastive signal (“avoid this error”) that the student can comprehend and act on, whereas on simpler tasks the wrong solution is merely distracting noise. Meanwhile, Self-Sum-1C (summarized) loses to raw Self-1C on Q3: the 0.6B student cannot decipher a compressed summary of a hard mathematical argument and needs the full step-by-step trace. On easier/medium questions (Q0, Q7), summarized contexts perform well because the reasoning is simple enough to compress without losing comprehensibility. What differs from BoolQ/MMLU.
Additionally, teacher rankings shift with filter stringency: on Q7, external teachers lead at lenient filters but Self-Sum-1C dominates at strict filters that focus on high-stakes nodes. These observations reinforce that no universal distillation recipe exists: the optimal teacher configuration depends on the task, the question difficulty, the student’s capacity to comprehend the context, and which decision points one prioritizes.
B
Full Gradient Derivations
B.1
Softmax Jacobian
The student’s transition probability at node u is Pθk = ezk /
P
k′ e
zk′
, with Jacobian:
∂Pθk = Pθk (δkj − Pθj ) ∂zj For any objective L(u) =
(B.1)
k k Pθ fk where fk does not depend on θ:
P
X X ∂L = fk · Pθk (δkj − Pθj ) = Pθj fj − Pθj Pθk fk = Pθj (fj − f¯) ∂zj k
k
12
(B.2)
Gradient alignment by teacher
Fraction of aligned decision points
Self-Sum-1C
+0.043
Qwen3-14B
+0.022
Qwen3-4B
+0.013
Self-Sum-1C1W
+0.013
Self-1C
+0.009
Qwen3-8B
+0.007
Self-1C (32B) Self-1C1W
0.518 0.518 0.512 0.504 0.502 0.505
-0.010
0.496
-0.031
-0.04
-0.02
0.00
0.02
0.04
Mean cosine similarity
0.06
0.47
0.48
0.483 chance
0.49
0.50
0.51
0.52
Fraction of aligned decision points
0.53
0.54
Figure 5 AIME Q0, Qwen3-0.6B (pass rate 87.5%)
Gradient alignment by teacher
Fraction of aligned decision points
Self-1C1W
+0.053
Qwen3-14B
+0.040
Self-Sum-1C1W
0.524
+0.027
Self-1C
0.511
+0.023
Qwen3-4B
0.507
+0.020
Self-Sum-1C
0.512
+0.016
Self-1C (32B) Qwen3-8B
0.520
0.504
+0.000
0.494
-0.012
-0.02
0.494
0.00
0.02
0.04
Mean cosine similarity
0.06
0.47
0.48
chance
0.49
0.50
0.51
0.52
0.53
Fraction of aligned decision points
0.54
Figure 6 AIME Q28, Qwen3-1.7B (pass rate 56.3%) Figure 7 Teacher ranking by gradient alignment on two AIME 2025 questions.
Both external teachers and selfdistillation variants achieve positive alignment, with Self-Sum-1C and Self-1C1W leading.
B.2
Dr. GRPO
The full GRPO objective in the limit G → ∞ is: JGRPO = Eoi ∼πθold
! 1 X πθ (oi,t | q, oi,<t ) · Ai − β DKL (πθ ∥πref ) |oi | t πθold (oi,t | q, oi,<t )
(B.3)
We make the following simplifications to obtain a per-node decomposition: 1. Drop the KL penalty.
Since β is typically small, we set β = 0.
For trajectories passing through node u and choosing token k, the importance weight πθ (oi,t | ·)/πθold (oi,t | ·) = Pθk /Pθkold at step t cancels with the sampling measure Pθkold after marginalization, leaving a net factor of Pθk . 2. Marginalize the importance ratio.
13
3. Remove length normalization (Dr. GRPO).
tory length: J(u) ∝
Standard GRPO divides by |oi |, coupling the advantage to trajecX
Pθk · E
k
Ai |oi |
oi passes u → vk
(B.4)
Since Ai and |oi | are correlated (e.g., successful trajectories may be shorter), this expectation does not factor. Dr. GRPO removes the 1/|oi | factor, giving: X JDrGRPO (u) ∝ Pθk · E[Ai | oi passes u → vk ] (B.5) k
4. Evaluate the conditional expectation.
With Ai = (Ri − R̄)/std(R) and binary rewards Ri ∈ {0, 1}:
E[Ai | oi passes u → vk ] =
k Psucc − R̄ k = α · Psucc −C std(R)
(B.6)
where α = 1/std(R) and C = R̄/std(R) are batch-level constants (each trajectory shifts them by ±1/G, negligible for large G). 5. Remove the constant.
node u reduces to:
Since
k k Pθ · (−C) = −C is independent of θ, the gradient-relevant objective at
P
LDrGRPO (u) =
X
k Pθk · Psucc
(B.7)
k j which is the ideal objective (Equation 2.2), giving gradient Pθj (Psucc − P̄succ ).
Empirical estimator.
In practice with finite samples, the per-node gradient is estimated as: 1 X ∂ L̂DrGRPO = Ai δri ,j − Pθj ∂zj Nu i: o ∋u u
(B.8)
i
where ri is the token chosen by rollout i at node u and Nu is the number of rollouts passing through u. In our k estimates, which is equivalent framework, we instead compute the ideal gradient directly from empirical P̂succ in expectation but lower variance.
B.3
GKD
P k k The GKD loss is LGKD = KL(πθ ∥πte ) = k Pθk (log Pθk − log Pte ). Defining ℓk = log Pθk − log Pte , the gradient is: ∂LGKD = Pθj (ℓj − ℓ̄) where ℓ̄ = KL(πθ ∥πte ) (B.9) ∂zj The additional term from ∂ log Pθk /∂zj produces +Pθj (1 − 1) = 0 after applying the Jacobian sum-to-zero property.
B.4
Single-sample GKD estimator (Thinking-Lab)
The empirical estimator uses importance weighting at the sampled token only. For rollout i choosing token ri at node u, the weight is: ri wi = log Pte − log Pθri − 1 (B.10) and the per-sample gradient contribution is: ∇zj ℓi = wi · (δri ,j − Pθj )
14
(B.11)
Taking the expectation over the student’s sampling distribution: X k E[∇zj ℓ] = Pθk log Pte − log Pθk − 1 (δkj − Pθj ) k
= Pθj (−ℓj − 1) − Pθj
X
Pθk (−ℓk − 1)
k
= −Pθj (ℓj − ℓ̄)
(B.12)
P The constant −1 cancels because k Pθk (δkj − Pθj )(−1) = 0. This recovers the GKD gradient (Equation 2.8) with opposite sign, confirming that the single-sample importance-weighted estimator and the full-vocabulary KL minimization produce the same gradient direction in expectation.
B.5
MiniLLM
MiniLLM (Gu et al., 2023) uses a REINFORCE-style gradient where the reward-to-go couples each node’s gradient to all downstream nodes. At node u (step t, sampled token ri ), the per-sample gradient is: X Rt′ − 1 (B.13) ∇zj ℓi = −(δri ,j − Pθj ) · t′ ≥t o
′
o
′
where Rt′ = log Ptei,t − log Pθ i,t is the per-step reward. Unlike GKD and its single-sample variant, here fk in the unified form Pθj (fj − f¯) depends on the full future trajectory rather than being purely local, making the per-node contribution path-dependent.
C
Discussion, Limitations, and Future Work
The alignment score is restricted to the set of tokens Su with sufficient visit counts at each node; tokens that are never sampled by the student cannot be evaluated. Additionally, targeted rollout enrichment requires substantial compute per question, making the diagnostic most practical as an offline evaluation tool rather than a real-time training signal.
Limitations.
Key extensions include: (i) designing training algorithms that exploit our observations, for example up-weighting the distillation loss on incorrect rollouts where alignment is strongest, or gating the teacher signal by student-teacher divergence (which correlates positively with alignment); (ii) multiteacher distillation that combines complementary signals from multiple teachers for better overall alignment; (iii) extending to long-horizon agentic tasks, where per-question variability and context-dependent teacher choice are likely even more pronounced. Future directions.
D
Computation Details
Not all nodes are equally informative for targeted enrichment. We prioritize using two criteria: (i) GKD gradient magnitude |Pθj (ℓj − ℓ̄)|, identifying tokens where the teacher disagrees most j strongly; and (ii) probability difference |Pte − Pθj |, identifying where the two distributions diverge most. We allocate a budget per depth window (smaller, denser windows early; larger, coarser windows later) and rank candidate tokens by both criteria. Rollout prioritization.
Phase 1 (shared): generate G initial rollouts, build the generation tree, and run targeted rollouts to enrich branching points. Each teacher’s logits are computed to identify important tokens; rollouts accumulate in the shared tree so later teachers benefit from prior enrichment. Phase 2 (per-teacher): k compute Psucc once from the enriched tree (teacher-independent), then for each teacher run one forward pass k to obtain Pte and evaluate the alignment score. Multi-teacher pipeline.
15
E
Additional Figures and Analysis
E.1
Alignment score distribution
The distribution of per-node gradient alignment scores across all decision points and 8 teachers on BoolQ spans [−1, +1] with a slight positive mean (+0.027 for 0.6B, +0.026 for 1.7B), demonstrating extreme pertoken heterogeneity.
E.2
Per-path cosine oscillation
Along individual reasoning paths, the gradient alignment score oscillates token-by-token between positive and negative values, confirming that the teacher’s helpfulness is highly local: even on a single path, consecutive tokens can alternate between beneficial and harmful distillation signal.
E.3
Selective distillation (oracle analysis)
Given the high per-token heterogeneity, a natural question is: what if we could apply the teacher’s gradient only at tokens where alignment is positive? Table 2 compares the mean importance-weighted signal under selective strategies that retain only decision points with alignment above a threshold t. Table 2 Selective distillation (oracle): mean signal (SR range × cosine), fraction of tokens retained, and fraction of
paths that beat full GKD. Filtering to aligned tokens yields 10–15× signal improvement using ∼50% of tokens. Setting
Strategy
Mean signal
% tokens
% beats GKD
0.6B BoolQ
Full GKD Selective (t=0) Selective (t=0.3)
0.007 ± 0.001 0.070 ± 0.002 0.076 ± 0.002
100% 51.6% 46.6%
— 100% 100%
1.7B BoolQ
Full GKD Selective (t=0) Selective (t=0.3)
0.006 ± 0.001 0.093 ± 0.003 0.096 ± 0.003
100% 51.6% 49.7%
— 100% 100%
0.6B MMLU
Full GKD Selective (t=0) Selective (t=0.3)
0.005 ± 0.001 0.057 ± 0.002 0.062 ± 0.002
100% 51.7% 46.1%
— 100% 100%
1.7B MMLU
Full GKD Selective (t=0) Selective (t=0.3)
0.006 ± 0.001 0.077 ± 0.003 0.081 ± 0.003
100% 51.2% 47.9%
— 100% 100%
Simply filtering to tokens with positive alignment (t = 0) yields 10–15× improvement in mean signal using only ∼52% of tokens, and every path benefits. This is an oracle analysis; at training time, one does not know the true alignment. However, it establishes a strong upper bound and motivates using the divergence-based predictors identified in Section 4.3 as approximate filters for alignment-aware training.
E.4
Teacher ranking by alternative metrics
Table 3 reports teacher rankings using fraction of positively aligned tokens and success-rate-range-weighted cosine. For 0.6B, both metrics preserve the same hierarchy as mean cosine (Table 1). For 1.7B, the weighted cosine reverses the top ranking: Self-Sum-1C leads over Qwen3-8B, as discussed in Section 4.2. The weighted cosine (SR range) amplifies the differences seen in mean cosine: self-distillation teachers for 0.6B achieve weighted cosine 0.108–0.120 vs. 0.053–0.062 for external baselines (2× ratio). For 1.7B, SelfSum-1C leads on weighted cosine (0.088) followed by Qwen3-8B (0.072), showing that at high-stakes decision points the summarized self-distillation context remains competitive even when mean cosine favors the external teacher. Fraction positive shows smaller but consistent differences in the same direction.
16
Table 3 Teacher ranking by fraction positive and weighted cosine (SR range) on BoolQ.
Frac. positive
Weighted cosine (SR)
Teacher
0.6B
1.7B
0.6B
1.7B
Self-1C Self-1C (32B) Self-Sum-1C Self-Sum-1C1W Self-1C1W
.526 .528 .522 .509 .511
.515 .512 .527 .505 .501
.120 .108 .114 .065 .043
.046 .039 .088 .038 .017
Qwen3-14B Qwen3-8B Qwen3-4B
.512 .511 .512
.518 .530 .520
.060 .062 .053
.053 .072 .060
E.5
Additional teacher ranking plots
E.6
Full correct vs. incorrect breakdown
Table 4 extends the main-text results (Figure 3) with teacher advantage. The pattern holds across all metrics, with the weighted cosine showing the largest effect sizes. Table 4 Full correct vs. incorrect path alignment across all metrics.
E.7
Setting
Metric
Correct
Incorrect
∆
p-value
0.6B BoolQ
Mean cosine Weighted cosine (SR range) Teacher advantage
0.011 0.046 0.002
0.044 0.110 0.008
−0.033 −0.065 −0.006
7 × 10−8 8 × 10−10 < 10−19
1.7B BoolQ
Mean cosine Weighted cosine (SR range) Teacher advantage
0.001 0.010 −0.002
0.058 0.093 0.011
−0.056 −0.083 −0.013
2 × 10−9 7 × 10−9 < 10−33
0.6B MMLU
Mean cosine Weighted cosine (SR range) Teacher advantage
0.009 0.021 0.002
0.048 0.118 0.005
−0.039 −0.097 −0.004
0.0001 < 10−11 < 10−14
1.7B MMLU
Mean cosine Weighted cosine (SR range) Teacher advantage
0.012 0.034 0.001
0.028 0.098 0.010
−0.016 −0.063 −0.009
0.123 < 10−5 < 10−14
Within-path correlation details
The divergence→alignment pattern is consistent for 0.6B and strengthens when restricting to the best teacher, but weakens for 1.7B BoolQ regardless of teacher choice. The positive sign means alignment is higher where the teacher disagrees with the student: low-divergence tokens offer little useful signal. However, high divergence is necessary but not sufficient: many high-divergence tokens still have negative alignment, since the teacher’s confidence can point toward failure as easily as success. Depth correlates positively for 0.6B (ρ ≈ +0.04), reflecting that early tokens are templatic while later tokens involve reasoning where the teacher’s contextual advantage emerges; this effect vanishes for 1.7B. On AIME (longer traces), the divergence pattern persists but depth becomes weakly negative (ρ ≈ −0.03 on hard questions), suggesting that on complex math the teacher’s advantage does not grow with depth.
E.8
Per-teacher correct vs. incorrect breakdown
Table 6 breaks down the correct-vs-incorrect gap by teacher on BoolQ. For 0.6B, Self-1C is unique: it achieves nearly equal alignment on both correct and incorrect paths (∆ ≈ 0), providing a uniformly helpful gradient regardless of path outcome. All other teachers show the typical pattern of higher alignment on 17
Table 5 Within-path Spearman correlations between single-rollout features and the gradient alignment score. Top:
averaged over all 8 teachers. Bottom: restricted to a teacher per setting. Qwen3-0.6B
Qwen3-1.7B
Feature
BoolQ
MMLU
BoolQ
MMLU
All teachers combined Depth (normalized) KL(πθ ∥πte ) JS divergence L2 distance Cosine(πθ , πte )
+.042 +.029 +.030 +.031 −.029
+.035 +.028 +.028 +.034 −.026
−.007 +.009 +.006 +.006 −.003
+.010 +.030 +.022 +.023 −.020
Self-1C
Self-1C (32B)
Self-Sum-1C
Self-Sum-1C
+.042 +.044 +.041 +.042 −.034
+.026 +.000 −.002 +.013 −.005
+.000 +.011 +.009 +.009 −.005
+.028 +.028 +.020 +.020 −.019
With a teacher only Depth (normalized) KL(πθ ∥πte ) JS divergence L2 distance Cosine(πθ , πte )
incorrect paths; for 1.7B this holds without exception. Table 6 Per-teacher correct vs. incorrect mean cosine on BoolQ.
Qwen3-0.6B
E.9
Qwen3-1.7B
Teacher
Corr.
Incorr.
∆
Corr.
Incorr.
∆
Self-1C Self-1C (32B) Self-Sum-1C Self-1C1W Self-Sum-1C1W
0.047 0.030 0.023 0.010 −0.005
0.047 0.063 0.059 0.029 0.037
≈0 −0.034 −0.036 −0.018 −0.042
−0.006 −0.027 0.021 −0.013 −0.026
0.062 0.067 0.079 0.017 0.042
−0.068 −0.094 −0.057 −0.030 −0.068
Qwen3-14B Qwen3-8B Qwen3-4B
−0.006 −0.012 0.001
0.037 0.048 0.035
−0.043 −0.059 −0.035
0.017 0.026 0.020
0.055 0.080 0.060
−0.039 −0.055 −0.040
Teacher Context Generation and Screening Details
To investigate the impact of in-context demonstrations on performance, we design a screening pipeline that measures the impact of various forms of context on pass rates across difficulty levels. For sourced context in summarized demonstrations, we use Qwen3-32B to process demonstrations. For the rest, we use the same model as the student. All models use thinking mode enabled. We screen questions from two benchmarks: MMLU (Hendrycks et al., 2021) (500 questions, multiple-choice) and BoolQ (Clark et al., 2019) (500 questions, yes/no). For each question we sample G=32 rollouts at temperature τ =1.0 and compute the pass rate as the fraction of rollouts producing a correct answer. Before screening, for each question, we generate demonstration responses by sampling from the context source model at temperature τdemo =0.7. Each response is verified against the ground-truth answer and classified as correct or incorrect. Generation retries up to 60 trials per question to collect the required number of demonstrations. Context Generation.
We evaluate seven context configurations, organized into three families: 1. Raw demonstrations. Correct and/or incorrect responses are prepended verbatim to the prompt. Variants: 1 correct (Self-1C), 1 correct + 1 wrong (Self-1C1W), and 3 correct (Self-3C). Additionally, Self-1C (32B) uses a single correct demonstration generated by Qwen3-32B rather than the 18
Table 7 Number of questions with at least one correct demonstration available among randomly selected questions for
screening, and the intersection used for analysis. Model
Benchmark
Self
32B
Both
Qwen3-0.6B Qwen3-0.6B
MMLU BoolQ
381 473
459 482
364 458
Qwen3-1.7B Qwen3-1.7B
MMLU BoolQ
357 474
459 482
350 463
Table 8 Screening results. ∆ is the absolute improvement over baseline.
MMLU
BoolQ
Qwen3-0.6B Qwen3-1.7B Qwen3-0.6B Qwen3-1.7B Context Variant Acc Baseline Self-1C Self-1C1W Self-3C Self-Sum-1C Self-Sum-1C1W Self-1C (32B)
63.8 98.6 68.2 99.4 96.4 91.5 98.8
∆
Acc
∆
Acc
∆
— +35.3 +4.7 +36.1 +32.7 +27.8 +35.5
85.2 99.7 94.0 99.8 99.3 98.2 99.7
— +14.5 +8.8 +14.6 +14.1 +13.0 +14.5
78.5 98.5 77.6 99.3 97.4 92.5 98.4
— +19.9 −0.9 +20.7 +18.8 +14.0 +19.9
Acc
∆
88.9 — 98.3 +10.0 87.7 −0.8 99.3 +11.0 98.2 +9.5 94.5 +5.8 98.7 +10.4
student. 2. Summarized demonstrations. Qwen3-32B condenses each response to its key reasoning steps and final answer. Variants: 1 summarized correct (Self-Sum-1C) and 1 summarized correct + 1 wrong (Self-Sum-1C1W). All context is injected by prepending it to the user message before applying the chat template. Figures 12–14 show the exact prompt structure for each context mode. Context Injection.
We assign each question to a difficulty bin based on its baseline (no-context) pass rate: easy (p ≥ 0.8), medium (0.25 ≤ p < 0.8), hard (0 < p < 0.25), and extremely hard (p = 0). Extremely hard questions are excluded from the primary analysis as a zero baseline provides no signal for measuring improvement during offline gradient analysis. Difficulty Binning and Filtering.
To ensure fair comparison across context variants, we restrict the analysis to questions for which the context generation process successfully produced at least one correct demonstration under both the self-context and 32B-context conditions. Table 7 reports the number of questions retained after this filter. It is worth to emphasize that the questions where no correct response could be generated for either source are excluded, ensuring that observed differences reflect the quality of injected context rather than its availability. Several patterns emerge across all models and benchmarks. First, even a single correct demonstration (Self1C) as teacher context produces dramatic improvements for teacher accuracy. Second, including wrong demonstrations alongside correct ones (Self-1C1W) consistently hurts performance relative to correct-only variants, and sometimes even degrades below the no-context baseline. Third, the gap between self-generated and 32B-generated demonstrations (Self-1C vs. Self-1C (32B)) is small. Lastly, we also provide difficulty breakdown for accuracy changes in screening questions based on context variations in Tables 9. We observe that the accuracy improvements for medium and easy questions (for baseline) are more significant compared to the improvements for hard questions, which suggests that especially the smaller models may still tend to generate wrong responses even when correct demonstrations are provided in its context.
19
Table 9 Detailed difficulty breakdown for Qwen3-0.6B (left) and Qwen3-1.7B (right). Each cell contains accuracy (%) and number of questions (n). Qwen3-0.6B Hard Variant
Acc
Med. n
Acc
Qwen3-1.7B Easy
n
Acc
All n
Acc
Hard n
∆
Variant
Acc
Med. n
MMLU Baseline Self-1C Self-1C1W Self-3C Self-Sum-1C Self-Sum-1C1W Self-1C (32B)
12.6 97.4 41.3 98.2 90.5 80.0 96.0
50 50 50 50 50 50 50
52.9 98.1 59.4 99.4 95.3 89.0 99.1
Baseline Self-1C Self-1C1W Self-3C Self-Sum-1C Self-Sum-1C1W Self-1C (32B)
12.9 92.6 35.2 96.0 84.1 68.6 94.4
29 29 29 29 29 29 29
56.2 96.9 57.1 98.7 94.8 85.1 96.3
168 168 168 168 168 168 168
Easy n
All
Acc
n
Acc
n
∆
97.8 100.0 98.1 100.0 100.0 100.0 100.0
265 265 265 265 265 265 265
85.2 99.7 94.0 99.8 99.3 98.2 99.7
350 350 350 350 350 350 350
— +14.5 +8.8 +14.6 +14.1 +13.0 +14.5
98.3 99.9 95.3 100.0 99.8 99.2 99.7
378 378 378 378 378 378 378
88.9 98.3 87.7 99.3 98.2 94.5 98.7
460 460 460 460 460 460 460
— +10.0 −0.8 +11.0 +9.5 +5.8 +10.4
MMLU 94.6 99.6 88.0 100.0 99.7 98.3 99.5
143 143 143 143 143 143 143
63.8 98.6 68.2 99.4 96.4 91.5 98.8
361 361 361 361 361 361 361
— +35.3 +4.7 +36.1 +32.7 +27.8 +35.5
Baseline Self-1C Self-1C1W Self-3C Self-Sum-1C Self-Sum-1C1W Self-1C (32B)
12.5 99.0 71.9 98.8 96.5 85.0 100.0
15 15 15 15 15 15 15
53.1 98.8 83.0 99.4 97.1 94.3 98.5
94.6 99.7 90.7 99.9 99.7 98.1 99.7
299 299 299 299 299 299 299
78.5 98.5 77.6 99.3 97.4 92.5 98.4
458 458 458 458 458 458 458
— +19.9 −0.9 +20.7 +18.8 +14.0 +19.9
Baseline Self-1C Self-1C1W Self-3C Self-Sum-1C Self-Sum-1C1W Self-1C (32B)
12.9 74.6 31.6 91.0 70.3 55.7 85.9
16 16 16 16 16 16 16
53.2 95.0 58.2 97.5 95.3 76.9 96.1
BoolQ 130 130 130 130 130 130 130
Acc
70 70 70 70 70 70 70
BoolQ
Apple and the Apple logo are trademarks of Apple Inc., registered in the U.S. and other countries and regions.
20
66 66 66 66 66 66 66
Gradient alignment by teacher
Fraction of aligned decision points
Self-1C
+0.047
Self-1C (32B)
+0.047
Self-Sum-1C
+0.041
Self-1C1W
0.526 0.528 0.522
+0.019
0.511
Qwen3-4B
+0.018
0.512
Qwen3-8B
+0.018
0.511
Self-Sum-1C1W
+0.016
0.509
Qwen3-14B
+0.016
0.00
0.02
0.512
0.04
Mean cosine similarity
0.06
0.08
0.48
chance
0.50
0.52
0.54
0.56
Fraction of aligned decision points
Figure 8 Qwen3-0.6B, BoolQ
Gradient alignment by teacher
Fraction of aligned decision points
Qwen3-8B
+0.053
0.530
Self-Sum-1C
+0.050
0.527
Qwen3-4B
+0.040
Qwen3-14B
0.520
+0.036
0.518
Self-1C
+0.028
0.515
Self-1C (32B)
+0.020
0.512
Self-Sum-1C1W
+0.008
Self-1C1W
0.505
+0.002
-0.02
0.00
0.02
0.501
0.04
0.06
0.08
Mean cosine similarity
0.48
chance
0.50
0.52
0.54
Fraction of aligned decision points
0.56
Figure 9 Qwen3-1.7B, BoolQ
Gradient alignment by teacher Self-Sum-1C
Fraction of aligned decision points +0.036
Self-Sum-1C1W
0.518
+0.034
Self-1C (32B)
+0.021
Qwen3-14B
+0.017
Qwen3-4B
0.519 0.510 0.514
+0.017
Qwen3-8B
0.511
+0.014
Self-1C
0.510
+0.010
Self-1C1W
0.506
+0.009
-0.02
0.00
0.02
0.04
Mean cosine similarity
0.507
0.06
0.08
0.48
chance
0.50
0.52
0.54
Fraction of aligned decision points
0.56
Figure 10 Qwen3-1.7B, MMLU Figure 11 Teacher ranking by gradient alignment (additional settings). For 1.7B on BoolQ, external teachers (Qwen3-
8B) achieve the highest alignment. On MMLU, self-distillation retains an edge for both model scales.
21
SYSTEM
You are a helpful assistant. USER
Answer the following multiple choice question. The last line of your response should be of the following format: ’Answer: $LETTER’ (without quotes) where LETTER is one of ABCD. Think step by step before answering. {question} A) {choice_A} B) {choice_B} C) {choice_C} D) {choice_D}
Figure 12 Baseline prompt (no context). Shown for MMLU (multiple-choice).
SYSTEM
You are a helpful assistant. USER
{question with instruction + choices} Demonstration
This is a correct response to the question: """ {correct_response_1} """ Now answer with a response of your own, including the thinking process:
Figure 13 Self-1C prompt. A single correct demonstration is prepended. For Self-1C (32B) the demonstration is
generated by Qwen3-32B instead; the prompt format is identical.
22
SYSTEM
You are a helpful assistant. USER
{question with instruction + choices} Demonstrations
Below are example responses to the question. Some are CORRECT and some are WRONG –the wrong responses contain mistakes and should NOT be imitated. Correct: """ {correct_response_1} """ WRONG (do NOT imitate): """ {wrong_response_1} """ Now answer with a response of your own, including the thinking process:
Figure 14 Self-1C1W prompt. One correct and one wrong demonstration with warning header.
SYSTEM
You are a helpful assistant. USER
{question with instruction + choices} Summarized Demonstration
This is a correct response to the question: """ {summarized_correct_response_1} """ Now answer with a response of your own, including the thinking process:
Figure 15 Self-Sum-1C prompt. The demonstration is first condensed by Qwen3-32B.
23