Less is More: Early Stopping Rollout for On-Policy Distillation

Unknown

Conceptio › Archive › arXiv CS

arXiv CSopen access

Less is More: Early Stopping Rollout for On-Policy Distillation

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

machine learning, deep learning, neural networks

Less is More: Early Stopping Rollout for On-Policy Distillation Zhou Ziheng1, # , Jiaqi Li2 , Huacong Tang1 , Ying Nian Wu1 , Demetri Terzopoulos1 1

University of California, Los Angeles 2 Beijing Institute of General Artificial Intelligence # [email protected]

arXiv:2605.27028v1 [cs.LG] 26 May 2026

Abstract

However, we point out that the late-position token reward is ill-posed: at the first few tokens, the teacher’s score is conditioned only on the t=1 prompt πteacher (τstudent | xprompt ) - indeed what we expect the teacher to score on. However, at a later position m, it becomes conditioned on the student’s own previously generated tokens too: t=m t=1:m ). This condiπteacher (τstudent | xprompt , τstudent tioning context is off-policy to the teacher model, drifting away from teacher’s model distribution. Recent works in the LLM alignment field find that LLMs may revert to pre-training behaviors when they see contexts not covered by their posttraining (Anthropic, 2025; Tice et al., 2026; Kutasov et al., 2026). Therefore, the teacher my no longer continue to correct the student tokens to solve the answer but merely continues the autocompletion. We confirm and measure this decay by running a preliminary experiment by having the teacher to continue from an early-stopped student rollout. As shown in Figure 1, the teacher’s performance decays toward the student’s quickly after 100 tokens, and reaches the student baseline level within only 300 tokens. Motivated by this finding, we propose Early Stopping Rollout(ESR): restrict the student rollout to its first N tokens and compute the distillation loss only on this early window. The change is a single line in any on-policy distillation loop. Despite its simplicity, ESR consistently outperforms full-rollout OPD across tasks (math, code, function calling), training regimes (LoRA, full finetuning(FFT)), model scales (students 1.5B–32B, teachers 1.7B–72B), and model families (Qwen2.5, Qwen3, Gemma 2, Gemma 3), while reducing wallclock cost by up to 24× and peak training memory by up to 4×. Moreover, although normally the teacher is expected to be the upper bound of the distillation, we observe that ESR-trained students can often exceed the teacher. Moreover, importantly, ESR remains stable

On-policy distillation has recently emerged as a promising alternative to standard sequencelevel imitation, training a student by scoring its own rollouts with a teacher model. However, we observe “Off-policy Teacher Decay” problem in this paradigm: for the later tokens, with student’s earlier trajectory as context that is off-policy to the teacher, the teacher’s ability to produce a corrective score would decay, and may fall back to token-completion behavior learned in the pre-training stage. We empirically verify this problem, and we propose a simple method Early Stopping Rollout (ESR) to fix it: simply restricting the rollout generation to the first N response tokens. We show that ESR both surpasses the full rollout OPD performance across model size, family, tasks and traning regime, and exhibit much higher GPU efficiency and training stability, especially under cross model family scenarios. We further investigate the mechanism behind this surprising performance and discovered “Cascading Alignment” and “Sub-mode Commitment” effect of ESR that may explain why it works effectively and even sometimes exceeding the teacher model performance. Besides, we show that this position-based token selection strategy cannot be fully explainable by KL divergence and entropy signals.

Introduction

On-policy distillation (OPD) has emerged as a dominant paradigm for model distillation in industrial practice. The student generates its own rollouts τ , which are then scored by the teacher: at each token, the teacher’s probability πteacher (τstudent | xprompt ) serves as the soft target for the student (Agarwal et al., 2024; Gu et al., 2024). Viewed through an RL lens, OPD can be understood as using the teacher as a dense, token-level reward model that judges the student’s own behavior on a given prompt (Thinking Machines, 2025). 1

across model generations (eg. Qwen 2.5 to Qwen 3) and families (eg. Gemma to Qwen)(Table 1). We find that OPD brings little gain for same-family same-generation pairs, possibly due to that the teacher and student often share upstream data or were themselves co-distilled. The gain is much salient only when cross generation or cross family, but full-rollout OPD becomes very unstable in these settings and frequently collapses. Therefore, the stability and effectiveness of ESR is very valuable. To better understand the surprising effectiveness of ESR, we conduct a series of ablation to investigate the potential reasons: 1) Firstly, we identified an important mechanism that we named as Cascading Alignment after training on the early window, KL divergence on the untrained late tokens also drops by 30–40%. Therefore, we find that with ESR, the KL divergence does not have to see late positions to repair them. 2) Secondly, we discovered the Sub-mode Commitment behavior of ESR that may explain why ESR sometimes even exceeds the teacher: the ESR-trained student commits to a sub-mode of supported teacher modes instead of chasing the dominant mode. This sub-mode, however, may be better than the dominant mode sometimes. This finding indicates a potential path of superceding the teacher model in distillation that worth future investigations. 3) Lastly, we ablate over it relevance to KL and entropy signals and show that position is an independent factor from KL and entropy. Our contributions are summarized below:

Off-Policy Teacher Decay in OPD

We first identify a failure mode of on-policy distillation (OPD), which we call Off-policy Teacher Decay. In OPD, the teacher T provides token-level supervision by scoring the student S’s rollout at S ), and the loss is each position t, i.e., πT (· | x, y<t typically averaged uniformly across positions. This procedure implicitly assumes that, after conditioning on the student’s partial trajectory, the teacher can still provide a useful corrective signal. HowS , which ever, as t increases, the student prefix y<t is off-policy to the teacher, may move increasingly far away from the teacher’s own high-probability reasoning regions. The teacher may then no longer operating from its natural reasoning state; instead, it could fall back to the behavior that completes next tokens from this off-policy state induced by the student (Anthropic, 2025; Tice et al., 2026; Kutasov et al., 2026).. We propose that this drifting issue can be measured by the teacher’s recoverability gap after conditioning on a student-generated prefix: S ∆decay (t) = AT (x) − AT (x | y<t ),

where AT (x) denotes the teacher’s accuracy when S ) solving from the original prompt, and AT (x | y<t denotes its accuracy when continuing from a lengtht student-generated prefix. A larger ∆decay (t) indicates that the teacher is less able to recover from the student-induced prefix, and therefore its lateposition token distribution is less likely to represent a reliable corrective target. To empirically verify this decay, we feed the teacher a k-token student-generated prefix on MATH-500 and then let it continue autoregressively. The teacher’s avg@4 accuracy decays from its unconditional baseline of 65.30% to 62.70% at N =100, and further to 51.75% at N =300, approaching the student-baseline performance (Figure 1). This suggests that late-position teacher scores are not independent assessments of the original problem; they increasingly reflect how the teacher continues a trajectory that the student has already committed to. Uniformly weighting all token positions in OPD gives undue emphasis over those regions where the teacher signal is no longer corrective.

1. Method: Early Stopping Rollout(ESR) outperforms full-rollout on-policy distillation. A one-line change—restricting the rollout length to the first N response tokens—beats full rollout OPD distillation across tasks, model families, scales, and training, while being dramatically more efficient and stable to train, particularly for cross-family scenarios. 2. Deep dive: Investigation of why it works with systematic experiments. We show with experiments that: 1) ESR mitigates the Offpolicy Teacher Decay from full-rollout OPD. 2) The Cascading Alignment effect enables ESR to work for late-position tokens without training on them. 3) The Sub-mode Commitment behavior of ESR enables it to even sometimes exceed the teacher.

Method: Early Stopping Rollout (ESR)

Let πs denote the student and πt the teacher. In standard on-policy reverse-KL distillation, the student 2

Figure 1: (Left) Off-policy Teacher Decay. The teacher loses accuracy quickly as the student-generated rollout grows over a few hundred tokens. MATH-500, avg@4 (n=4, t=0.7). Teacher = Qwen3-1.7B; student = Qwen2.5-Math-1.5B. After ∼300 student tokens the teacher has effectively been dragged down to student-baseline performance. (Right) Rollout length N sweep on MATH-500. LoRA, Qwen2.5-Math-1.5B → Qwen3-1.7B; best avg@4 across training steps. OPD and the undistilled baseline are shown as horizontal references. Performance saturates for N ∈ [50, 200] and all beat OPD.

generates a response y = (y1 , . . . , yT ) conditioned on prompt x, and the loss is Lfull = Ey∼πs (·|x)

T hX

many experiments, due to resource constraints, the main experiments use LoRA (Hu et al., 2022) (r=32, α=64. But we conduct full finetune ablations to confirm its validity. Each training step processes a batch of 16 problems with 1 rollout per problem (nsamples =1, batch size 16). We train for 200 steps on all tasks, and saving checkpoints every 50 steps. Training data are drawn from from NuminaMath (LI et al., 2024), CodeUltraFeedback (Weyssow et al., 2024), and glaivefunction-calling-v2 (Glaive AI, 2023). Our method uses N =100 unless otherwise specified. For pairs whose student and teacher use different tokenizers (all cross-generation and cross-family pairs in our setup), we decode the student rollout to text and re-encode it under the teacher’s tokenizer to obtain teacher token-level log-probabilities; the reverseKL loss is then computed on tokens that are tokenaligned across the two vocabularies via a greedy text-span match. Evaluation. MATH-500 (Hendrycks et al., 2021; Lightman et al., 2023) with n=4 samples at temperature 0.7 (reporting avg@4), HumanEval (Chen et al., 2021; Liu et al., 2023) at temperature 0.0 (reporting pass@1), BFCL (Patil et al., 2025) reporting full accuracy: correct function name and arguments.

KL πs (· | x, y<t )

t=1

i ∥ πt (· | x, y<t ) . (1) ESR (position cutoff N , with N ≪ T in practice) truncates the student rollout to its first N tokens, and the loss is computed over exactly those tokens: LESR (N ) = Ey∼πs , |y|≤N

|y| X

KL πs (· | x, y<t )

t=1

∥ πt (· | x, y<t ) . (2) If the student emits EOS before position N , the rollout terminates naturally. Everything else— generation temperature, LoRA target modules, optimizer, scorer—is unchanged from the standard on-policy KD loop.

Main Experiments

4.1

Setup

Models. We evaluate across three regimes: samefamily same-generation (e.g. Qwen2.5→Qwen2.5, Qwen3→Qwen3), same-family cross-generation (e.g. Qwen2.5→Qwen3, Gemma-2→Gemma-3), and cross-family (Gemma→Qwen). Student sizes range from 1.5B to 32B and teacher sizes from 1.7B to 72B. Training. We employed reverse KL divergence loss with learning rate 5×10−5 ), and generate sequences with temperature 0.7. Since we have

4.2

Overall performance

ESR beats OPD across model families, generations, sizes. Across every cell of Table 1, ESR matches or beats OPD’s best score, and surpasses the teacher reference in many of them. For same family same generation setting, we test three sizes 3

Table 1: Main results: ESR dominates OPD across same-family same-generation, same-family crossgeneration, and cross-family pairs, and across scales (students 1.5B–32B, teachers 1.7B–72B) on MATH-500. “Student”/“Teacher” columns are the base models with no distillation. OPD values are the peak across training; subscripts ↓−∆ give the drop from peak to the final checkpoint, shown when >5%. ‡ values denote configurations that never reach a functional checkpoint (peak < 20% across all saved steps). ESR uses N =100, LoRA. Bold: ESR beats OPD. ⋆: ESR surpasses the teacher reference. For Gemma-2 2B → Qwen3-4B, ESR uses N =50 on this pair. avg@4 Pair (Student → Teacher)

pass@4

Student Teacher

OPD

ESR

Student Teacher

OPD

ESR

69.20 73.80 77.05

77.95 72.60 72.60

65.85 73.45 77.30 ⋆

69.20 74.30 ⋆ 78.10 ⋆

81.00 83.20 84.40

86.40 84.80 84.80

78.00 83.80 86.20 ⋆

81.20 84.00 87.40 ⋆

50.95 50.95 53.60 73.80 13.45

65.30 77.95 76.15 83.85 66.60

62.35 65.85 ⋆ 67.45↓−12.4 68.95 68.85↓−6.5 68.95 5.40‡ 75.15 22.95 27.20

72.80 72.80 75.00 83.20 28.20

77.00 86.40 83.20 88.00 74.80

75.20 80.60 80.00 15.80‡ 31.40

79.80 ⋆ 81.00 81.20 85.40 39.40

13.45

77.95

16.40↓−11.5

28.20

86.40

27.00↓−17.2

30.20

Same family, same generation Qwen3-1.7B → Qwen3-4B Qwen2.5-14B → Qwen2.5-Math-72B Qwen2.5-32B → Qwen2.5-Math-72B Same family, cross generation Qwen2.5-Math-1.5B → Qwen3-1.7B Qwen2.5-Math-1.5B → Qwen3-4B Qwen2.5-Math-7B → Qwen3-14B Qwen2.5-14B → Qwen3.5-35B-A3B Gemma-2 2B → Gemma-3 4B Cross family Gemma-2 2B → Qwen3-4B

19.90

Table 2: (Left) Performance on HumanEval with pass@1 and BFCL with full accuracy. (Right) Full Finetune (FFT) performance on MATH-500 with best across training steps. Bold: ESR beats OPD. ⋆: ESR surpasses teacher. ↓−∆ gives the drop from peak to final checkpoint (shown if >4%).

Pair

Method

Pair

BFCL

Student 31.10 2.70 Qwen2.5-Math-1.5B Teacher 39.60 54.00 → Qwen3-1.7B OPD 40.20↓−13.4 58.20 ⋆ ESR 42.10 ⋆ 61.30 ⋆ Gemma-2-2B → Gemma-3 4B

Student 23.78 73.17 Teacher 20.70 72.83 OPD 22.00↓−10.4 76.83 ⋆ ESR 28.70 ⋆ 79.00 ⋆

Student Qwen2.5-Math-1.5B Teacher → Qwen3-1.7B OPD ESR

50.95 65.30 58.20 56.20

72.80 77.00 75.40 73.80

Student Teacher OPD ESR

13.45 66.60 13.90 26.65

28.20 74.80 25.00 40.40

Gemma-2-2B → Gemma-3 4B

of model (Qwen 1.7B, 14B and 32B). Full rollout OPD sometimes even fall below its original performance (1.7B and 14B), but ESR always improves. For cross generation setting, we test Qwen 2.5 Qwen 3 or 3.5, with sizes ranging from 1.5B to 14B. We also tested Gemma 2 to 3 to ensure it works in different model series. For cross family settings, we let Gemma 2 2B to learn from Qwen3 4B. ESR consistently exceed the full rollout training, with full rollout training collapse in most of the times.

Method avg@4 pass@4

HE) and tool calling tasks(BFCL). Table 2 reports FFT on MATH-500 for the Qwen2.5→Qwen3 and Gemma-2→Gemma-3 pairs. On Qwen OPD 58.20 is slightly better than ESR 56.20 avg@4, but the gap is close. For Gemma ESR dominates OPD by 12.75% avg@4 and +15.40% pass@4. ESR is therefore the safer choice in both parameter regimes. 4.3

Stability of ESR

ESR is significantly more robust than OPD in training. In cross-generation and cross-family settings, full-rollout distillation degrades or completely collapses most of the times; ESR degrades nowhere. We denote the cells with degrading or collapsing failure mode in Table 1 and Table 2 with ↓−∆ and ‡ . However, we observe the student to

ESR matches or beats OPD across tasks and training regimes (LoRA vs FFT. We tests in both Qwen series and Gemma series for task and training regime generalization. Table 2 shows that ESR is also better in coding (Human Evaluation, 4

Table 3: Training efficiency. Single A6000 (48 GB), bs=16, student teacher Qwen3-1.7B. ESR uses N =100; memory values in GB. We report the average running time and GPU memory usage across student model generation, training and teacher scoring phases. Note that the real time usage can be larger if there isn’t enough GPU to hold the student and teacher models together and requires model loading and unloading. Metric

Teacher Student Student Method Total Scoring Generation Training

-step wall time

ESR OPD Speedup

1s 7s 7×

5s 180 s 36×

2s 7s 3.5×

8s 194 s 24×

ESR Peak GPU OPD memory Savings

7.3 G 14.9 G 2.0×

7.2 G 8.9 G 1.2×

9.6 G 39.5 G 4.1×

24.1 G 63.3 G 2.6×

comfortably. In our own practice, it saves a further big time overhead of model loading and unloading that we do not report here.

More Analysis on Why ESR Works

5.1

The Cascading Alignment Effect of ESR

Without training over the late-position tokens, can ESR still learns the teacher behavior comprehensively? We find “Convergence Cascade Effect” of ESR: even training on only the first N tokens with ESR, per-position KL divergence beyond [0, N ] region still drops by 30–40% (Figure 2). This shows that the student can pick up the teacher’s “global mindset” even with just the beginning tokens. Regarding to why Cascading Alignment Effect happens, one reason that we suspect is that the beginning tokens often consist of problem framing and strategic planning content. The case study in Figure 4 (Left) illustrates this concretely: on a representative MATH-500 trajectory, the first 100 tokens set up the geometry, name the unknown, and identify the key relationship (the altitude bisects the leg)—the choices that determine whether the rollout will succeed—while the last 100 tokens focus on executing the algebra that any solver can finish once the strategy is fixed. Therefore once the student picks up how to frame problems and plan the strategy, the later content naturally follows. Moreover, recently (Cloud et al., 2025) shows student models may be able to learn the teacher’s deep internal preference even with random numbers generated by the teacher, called “subliminal learning”. Therefore, the early tokens may inject a global subliminal mindset to the student rather than only altering the prefix tokens.

benefit significantly in these setting, showing more than 10 % improvement for avg accuracy many times, whereas bare improvement can be observed in the same family same generation distillation setting. Early Stopping Rolloutis not sensitive to the choice of N except for cross-family setting. A natural question naturally occurs - how to choose where to stop? Is it sensitive? We conducted a set of sweeping experiments in Figure 1 sweeps N on MATH-500 with Qwen2.5-Math-1.5B and Qwen3 1.7B, a cross-generation setting where full rollout OPD suffers from stability issue, and reveals a robust region: it reaches just as good performance starting from N =50 and remains stable to N =200. The method is not sensitive to the exact choice of N within a certain region. But we do find that for the cross-family setting (Gemma-Qwen pair), it is sensitive that it is stable with 50 tokens but not 100 tokens. Therefore, the bigger gap between the teacher and student, the more sensitive it is for choice of N . This also validates our “Off Policy Teacher Decay” diagnosis of OPD - the bigger gap between the student and teacher model, the more off-policy the student trajectory prefix is to the teacher and the bigger decay it causes. 4.4

5.2

The Sub-mode Commitment Effect

ESR exceeds the teacher in many of the main experiments1. Even the full rollout OPD model slightly exceeds the teacher in function calling experiments a few times. This shows that student has the potential exceed the teacher even in normal OPD, and our method amplifies it. Why is so? Isn’t teacher supposed to be the upper bound? We propose the reason lies in the mechanism of reverse KL KL(πs ∥ πt ), which has a modeseeking behavior: it penalizes the student for putting mass on tokens the teacher does not support, but not for concentrating mass on a single supported token. Therefore, the student has the possibility to land on a sub-mode of the teacher

Efficiency of ESR

Table 3 shows that ESR achieves a 24× wall-clock speedup and reduces peak training memory by ∼ 4×. The dominant cost in OPD is autoregressive generation (180 s/step for sequences averaging ∼1000 tokens); ESR generates only N =100 tokens (5 s/step). Note that with ESR, all the student and teacher models can be put in one A6000 GPU 5

Figure 2: (Left) Early tokens are high on student entropy, teacher entropy, and KL divergence simultaneously. (Right) The convergence cascade. Per-position KL between distilled student and teacher, before vs. after ESR training. Yellow band: positions [0, 100] that actually receive training loss. Blue band: the KL gap closed by training. Positions 100+—which see no direct training signal—drop to the same KL as the trained region, confirming that alignment on the early window cascades through the autoregressive rollout.

Table 4: Behavioral Modes Analysis. (Left) Response length distribution on -500. ESR generates far shorter outputs than OPD, teacher, and its baseline, with 10th, 50th (median), 90th percentiles, and average length shown. (Right) Alignment of student outputs with teacher distribution modes. Metrics include top-1 token probability (choice confidence), plus percentages of student top-1 matching teacher top-1 (= top-1), falling in teacher top-2–5 (∈ top 2–5), or outside teacher top-5 (∈ / top-5). ESR favors teacher non-top-1 choices and is more confident than OPD.

that is actually better. We visualize this mechanism in Figure 3, and verify it empirically below. Indeed, we verified that in comparison to the full rollout OPD, ESR can push the student more toward the non-dominant mode. We first scan the behavioral differences across distilled models. Surprisingly, ESR-trained students produce sequences 2–3× shorter than the teacher, full-rollout, and even the base student itself: ESR-100’s median length is ∼380 tokens, against ∼1,150 for the teacher and ∼1,530 for full-rollout (Table 4, left). The teacher is substantially more verbose than the student, so distilling from such a teacher generally drags the student’s length up — and indeed full-rollout produces rollouts even longer than the teacher itself. The fact that ESR learns from the same teacher yet moves in the opposite direction shows how decisive the rollout-length choice is: by removing late-position supervision, the student preserves its own succinct style while still inheriting the teacher’s reasoning strategy, leading to a more desirable outcome than simply copying the teacher.

Response length (tokens)

Metric

Student Teacher

10% 50% 90% Mean Ratio

∼400 ∼700 ∼1,190 ∼190 ∼860 ∼1,150 ∼1,530 ∼380 ∼1,500 ∼1,800 ∼1,770 ∼800 ∼990 ∼1,210 ∼1,480 ∼460 2.2× 2.6× 3.2× 1.0×

Student Top-1 prob 0.71 alignment = top-1 28.6% to teacher ∈ top 2–5 59.8% probability ∈ / top-5 11.5%

— — — —

OPD

ESR

0.77 0.79 45.7% 41.9% 44.6% 47.4% 9.7% 10.7%

dent is also more confident in its own chosen token. This exactly shows that ESR steers the student to commit to a secondary mode in the teacher’s distribution.

Furthermore, we examine quantitatively how the student’s probability output aligns with the teacher’s modes. We take the top-10% highest-KL tokens after training (n=110,894), which reveal behavioral differences most saliently, and check how often the student’s top choice agrees with the teacher’s top-1, falls in the teacher’s top 2–5, or lies outside the top-5 (Table 4, right). We find that, indeed, ESR produces a model that is more committed to the teacher’s top 2–5 choices than to the teacher’s top-1 (47.4% in top 2–5 vs 44.6% for full-rollout; 41.9% argmax agreement vs 45.7%). At the same time, ESR’s top-1 probability is higher than full-rollout’s (0.79 vs 0.77), showing the stu-

5.3

Ablation with KL and Entropy

We find that, as shown in Figure 2, early positions simultaneously have high KL divergence between student and teacher model, and high token entropy from both student and teacher models. This finding probes us to wonder if the effectiveness is intrinsically induced by the KL and entropy. To control these factors, we conduct a series of ablation experiments: pick the same amount (100) of tokens based on the highest KL divergence, highest student/teacher entropy or with them in combination 6

Figure 3: Mode-seeking on the planning region. Schematic of reverse-KL behaviour at a multi-modal planning position. The teacher supports two modes (e.g. a verbose plan and a concise correct plan); reverse KL penalises student mass outside the support but not concentration within it. A short early-window loss can collapse the student onto the better supported mode, allowing the student to exceed the teacher’s average behaviour. OPD training reverts the student toward the averaged teacher across late positions, undoing the concentration.

with ESR, regardless of position (Figure 4). If the effectiveness is indeed induced by the KL or entropy, then they should reproduce the same or better results. To our surprise, all underperform ESR, and most of them also much underperform the OPD results. Teacher or student entropy based selection can match the full sequence by them alone, but their combination falls short significantly. What’s also interesting is that KL divergence measure, the direct calculation of the loss magnitude, barely works. It only improves the baseline (50.95%) for about 3 percent. And more surprisingly, we find that the largest 100 tokens of KL occupies around 93% of the entire trajectory loss. This shows that the tokens that has larger signals are not necessarily the ones that have effective signals. Therefore, although we don’t exclude KL and entropy as potential mediator factor, we exclude them to be the sole factors that causes early tokens to be special. Position, therefore, should be considered an independent token selection dimension for the future.

teacher regions, and Agarwal et al. (2024) introduced Generalized Knowledge Distillation (GKD), which uses student-generated rollouts to obtain substantial gains over off-policy distillation on reasoning tasks. Related work has explored other divergence and sampling choices, including skew KL and adaptive off-policy schedules (Ko et al., 2024), general f -divergences (Wen et al., 2023), the modeseeking versus mean-seeking behavior of forward and reverse KL (Wu et al., 2025), and speculative knowledge distillation with interleaved teacherstudent sampling (Xu et al., 2025). Our work builds directly on the on-policy reverse-KL setting of Gu et al. (2024) and Agarwal et al. (2024), but asks a different question: holding the divergence and rollout distribution fixed, which token positions carry useful signal? Token-Level Importance in Distillation and Reasoning. A growing line of work suggests that not all tokens contribute equally to learning. In reasoning, Wang et al. (2025) found that only a small fraction of chain-of-thought tokens are high-entropy “forking tokens” that steer subsequent reasoning, while Vassoyan et al. (2025) showed that uniform KL penalties can suppress exploration on critical tokens and proposed entropy-weighted KL relaxation. Related studies also identify token-level structure in planning and credit assignment, including preplan-and-anchor behavior (Li et al., 2025) and functional importance in reasoning chains (Singh and Hakkani-Tür, 2026). In distillation specifically, several concurrent methods have explored token selection or weighting. SelecTKD uses teacher verification to mask rejected tokens (Huang et al., 2025); AdaKD adapts

Related Work

Knowledge Distillation for Language Models. Knowledge distillation (Hinton et al., 2015) transfers knowledge from a teacher to a smaller student via soft targets, and Kim and Rush (2016) extended this idea to sequence models with wordlevel and sequence-level objectives. For autoregressive LLMs, both the divergence and the data distribution are crucial. Gu et al. (2024) advocated reverse KL for generative LLM distillation, arguing that it avoids assigning mass to low-support 7

Figure 4: (Left) Case study: early tokens encode strategy; later tokens mostly execute. A representative MATH-500 trajectory. The first 100 tokens establish the geometric setup, name the unknown, and identify the key relationship (altitude bisects the leg)—these are the choices that determine whether the rollout will succeed. The last 100 tokens carry out an algebraic computation that any solver can finish once the strategy is fixed. ESR supervises the first window; the second follows for free. (Right) Token selection strategies on MATH-500. Qwen2.5-Math-1.5B → Qwen3-1.7B, LoRA, nsamples =1, best avg@4 across training steps. All selectors pick N =100 tokens from full-length rollouts; ESR is the only one that additionally truncates the rollout. Baseline (no distillation) 50.95%; OPD 62.35%. Hs : student entropy; Ht : teacher entropy; RKL: reverse KL. Problem. √ In an isosceles right triangle, the altitude to the hypotenuse has length 4 2. What is the area? First 100 tokens (strategy). “Let’s start by understanding the problem.√In an isosceles right triangle, the two legs are equal, and the hypotenuse √ is 2 times the leg. If we denote each leg by a, then hypotenuse = a 2. The altitude to the hypotenuse is half the leg, because the altitude bisects the. . . ” Last 100 tokens (execution). “. . . altitude from the right-angle vertex to√the 2 a √ = a√ = √ hypotenuse. This altitude has length h = aa·a = a2 2 . 2 2 √ a√2 √ √ a 2√ Setting this equal to 4 2: 2 = 4 2, so a 2 = 8 2, giving a = 8. The area is 21 × 8 × 8 = 32 .”

token-level temperature based on training stability (Xie et al., 2026); SE-KD disentangles selection along position, class, and sample axes and uses student-entropy filtering along the position axis (Tavor et al., 2026); and TSDKD combines entropy-based token selection with preference ranking (Kim and Baek, 2026). Our ablation in Section 5.3 shows that these scalar token-saliency criteria are insufficient: selecting the same number of tokens by top-KL, top-entropy, or combined entropy heuristics all underperform ESR, and most underperform even plain full-sequence training. This indicates that position is a load-bearing axis of supervision rather than a proxy for token saliency. It also reconciles our findings with Wang et al. (2025): in on-policy distillation, the high-entropy “forking tokens” are concentrated in the early uncontaminated window, so their conclusion is consistent with ours once position is taken into account.

Selection method

avg@4

ESR OPD

65.85 62.35

Top-RKL Top-Ht (teacher entropy) Top-Hs (student entropy) RKL·Hs Ht ·Hs (product) RKL·Ht ·Hs (triple product)

53.35 63.30 62.70 56.90 55.35 57.90

nism through which prefix tokens drive learning — none of which is addressed by this concurrent work. Li et al. (2026) likewise note as a small part of their paper that the student’s prefix can cause the teacher signal to degrade and that one may not use the full rollout. However, their focus is not on it alone, and the provided analysis and experiments are much narrower in scope than ours — they do not run the systematic cross-generation, cross-family, crossscale matrix used here, nor attribute the effect to specific mechanisms such as the convergence cascade or the reverse-KL mode-seeking behavior that lets ESR-trained students surpass their teacher.

Conclusion

We introduce Early Stopping Rollout (ESR), a minimal one-line modification to OPD that constrains the rollouts to the first N response tokens. Despite its simplicity, ESR outperforms full-rollout OPD across three core dimensions: performance, efficiency, and stability across tasks, scales, and training regimes.

Concurrent work. Zhang et al. (2026) independently report that concentrating OPD supervision on response prefixes is an effective efficiency lever; in their setting — distilling a reasoning teacher into a base model that has not yet acquired reasoning behavior — prefix OPD does not surpass full-trajectory OPD, suggesting that bootstrapping reasoning from scratch still benefits from fulltrajectory supervision. Our setting is complementary: starting from a math-SFT student that already reasons, ESR is able to achieve better performance than full rollout and may even push the student beyond the teacher; we additionally provide systematic experiments and analyses on why lateposition tokens are detrimental and on the mecha-

We also discovered a series of mechanisms that explains our method’s efficacy. We discovered the “Off Policy Teacher Decay” as the root problem our method mitigates; the “Cascading Alignment” effect that may explain why it works effectively without training on the later tokens; and “Submode Commitment” effect that explains why it even sometimes exceeds the teacher. Besides, we show that this position-based token selection strategy is an load-bearing axis beside KL divergence and entropy signals. 8

Limitations. Our experiments expects the student models to be instruction-tuned model that already learns basic thinking ability rather than pre-trained only model, which may be inferior according to the concurrent workZhang et al. (2026). Our experiments are also focus on the setting where small open-source models (<100B) are finetuned for a specific task with a limited data budget. Whether the ESR story holds at industrial scale general model capacity improvement (trillioin level model size; millions of training trajectories and above) remains unclear. It may be very likely that full rollout OPD works better in such level, although we still expect ESR to be helpful under fixed budget setting. If one samples more trajectories to cover more diverse scenarios with shorter length, it is imaginable that it may be better than full rollout trajectories with narrower diversity if the ratio is calibrated well. We also have not tested multi-modality or long-horizon tasks, which may exhibit different positional signal-quality patterns.

Anthropic. 2025. Agentic misalignment: How LLMs could be an insider threat. https: //www.anthropic.com/research/ agentic-misalignment. Anthropic research blog. Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, and 1 others. 2021. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374. Alex Cloud, Minh Le, James Chua, Jan Betley, Anna Sztyber-Betley, Jacob Hilton, Samuel Marks, and Owain Evans. 2025. Subliminal learning: Language models transmit behavioral traits via hidden signals in data. arXiv preprint arXiv:2507.14805. Glaive AI. 2023. glaive-function-callingv2. https://huggingface. co/datasets/glaiveai/ glaive-function-calling-v2. Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang. 2024. MiniLLM: Knowledge distillation of large language models. In International Conference on Learning Representations.

Ethical Considerations & Potential Risks. This work studies an algorithmic improvement to onpolicy knowledge distillation; it does not target on subjective tasks like value alignment. We see no specific ethical and risk concerns beyond beyond those generally applicable to language-model training research.

Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. 2021. Measuring mathematical problem solving with the MATH dataset. In Advances in Neural Information Processing Systems. Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. 2015. Distilling the knowledge in a neural network. In NIPS Deep Learning and Representation Learning Workshop.

Use of AI Assistants. This paper was primarily conceived, designed, and drafted by the human authors. AI assistants (including ChatGPT and Claude) were used in a supporting role for proofreading, rewriting for clarity, and assisting with code development for the simulation platform. All scientific contributions, experimental design, analysis, and intellectual direction were driven by the authors, with AI tools serving as aids for language refinement and coding assistance.

Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. 2022. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations. Haiduo Huang, Jiangcheng Song, Yadong Zhang, and Pengju Ren. 2025. SelecTKD: Selective tokenweighted knowledge distillation for LLMs. arXiv preprint arXiv:2510.24021.

Acknowledgments

Minsang Kim and Seung Jun Baek. 2026. Explain in your own words: Improving reasoning via tokenselective dual knowledge distillation. arXiv preprint arXiv:2603.13260.

Funding and competing-interest disclosures will appear here in the final version.

Yoon Kim and Alexander M. Rush. 2016. Sequencelevel knowledge distillation. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pages 1317–1327.

References Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos, Matthieu Geist, and Olivier Bachem. 2024. On-policy distillation of language models: Learning from self-generated mistakes. In International Conference on Learning Representations.

Jongwoo Ko, Sungnyun Kim, Tianyi Chen, and SeYoung Yun. 2024. DistiLLM: Towards streamlined distillation for large language models. In International Conference on Machine Learning.

Jonathan Kutasov, Adam Jermyn, Julius Steen, Minh Le, Samuel R. Bowman, Samuel Marks, Jan Leike, Amanda Askell, and Chris Olah. 2026. Teaching Claude why. https://alignment.anthropic.com/ 2026/teaching-claude-why/. Anthropic alignment-science blog post.

Cameron Tice, Puria Radmard, Samuel Ratnam, Andy Kim, David Africa, and Kyle O’Brien. 2026. Alignment pretraining: AI discourse causes self-fulfilling (mis)alignment. https://arxiv.org/abs/ 2601.10160. Preprint, arXiv:2601.10160. Jean Vassoyan, Nathanael Beau, and Roman Plaud. 2025. Ignore the KL penalty! boosting exploration on critical tokens to enhance RL fine-tuning. In Findings of the Association for Computational Linguistics: NAACL 2025.

Jia LI, Edward Beeching, Lewis Tunstall, Ben Lipkin, Roman Soletskyi, Shengyi Costa Huang, Kashif Rasul, Longhui Yu, Albert Jiang, Ziju Shen, Zihan Qin, Bin Dong, Li Zhou, Yann Fleureau, Guillaume Lample, and Stanislas Polu. 2024. NuminaMath. https://huggingface.co/AI-MO/ NuminaMath-CoT.

Shenzhi Wang, Le Yu, Chang Gao, Chujie Zheng, Shixuan Liu, Rui Lu, Kai Dang, Xionghui Chen, Jianxin Yang, Zhenru Zhang, Yuqiong Liu, An Yang, Andrew Zhao, Yang Yue, Shiji Song, Bowen Yu, Gao Huang, and Junyang Lin. 2025. Beyond the 80/20 rule: Highentropy minority tokens drive effective reinforcement learning for LLM reasoning. In Advances in Neural Information Processing Systems.

Yang Li, Zhichen Dong, Yuhan Sun, Weixun Wang, Shaopan Xiong, Yijia Luo, Jiashun Liu, Han Lu, Jiamang Wang, Wenbo Su, Bo Zheng, and Junchi Yan. 2025. Attention illuminates LLM reasoning: The preplan-and-anchor rhythm enables fine-grained policy optimization. arXiv preprint arXiv:2510.13554. Yaxuan Li, Yuxin Zuo, Bingxiang He, Jinqian Zhang, Chaojun Xiao, Cheng Qian, Tianyu Yu, Huanang Gao, Wenkai Yang, Zhiyuan Liu, and Ning Ding. 2026. Rethinking on-policy distillation of large language models: Phenomenology, mechanism, and recipe. https://arxiv.org/abs/2604. 13016. Preprint, arXiv:2604.13016. Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. 2023. Let’s verify step by step. arXiv preprint arXiv:2305.20050.

Yuqiao Wen, Zichao Li, Wenyu Du, and Lili Mou. 2023. f-divergence minimization for sequence-level knowledge distillation. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 10817– 10834. Martin Weyssow, Aton Kamanda, Xin Zhou, and Houari Sahraoui. 2024. CodeUltraFeedback: An LLM-as-a-judge dataset for aligning large language models to coding preferences. arXiv preprint arXiv:2403.09032. Taiqiang Wu, Chaofan Tao, Jiahao Wang, Runming Yang, Zhe Zhao, and Ngai Wong. 2025. Rethinking Kullback-Leibler divergence in knowledge distillation for large language models. In Proceedings of the International Conference on Computational Linguistics (COLING).

Jiawei Liu, Chunqiu Steven Xia, Yuyao Wang, and Lingming Zhang. 2023. Is your code generated by ChatGPT really correct? rigorous evaluation of large language models for code generation. In Advances in Neural Information Processing Systems.

Xurong Xie, Zhucun Xue, Jiafu Wu, Jian Li, Yabiao Wang, Xiaobin Hu, Yong Liu, and Jiangning Zhang. 2026. LLM-oriented token-adaptive knowledge distillation. In Proceedings of the AAAI Conference on Artificial Intelligence.

Shishir G. Patil, Huanzhi Mao, Fanjia Yan, Charlie Ji, Vishnu Suresh, Ion Stoica, and Joseph E. Gonzalez. 2025. The Berkeley function calling leaderboard (BFCL): From tool use to agentic evaluation of large language models. In Proceedings of the 42nd International Conference on Machine Learning, volume 267 of Proceedings of Machine Learning Research, pages 48371–48392.

Wenda Xu, Rujun Han, Zifeng Wang, Long T. Le, Dhruv Madeka, Lei Li, William Yang Wang, Rishabh Agarwal, Chen-Yu Lee, and Tomas Pfister. 2025. Speculative knowledge distillation: Bridging the teacherstudent gap through interleaved sampling. In International Conference on Learning Representations.

Janvijay Singh and Dilek Hakkani-Tür. 2026. Do LLMs encode functional importance of reasoning tokens? arXiv preprint arXiv:2601.03066.

Dongxu Zhang, Zhichao Yang, Sepehr Janghorbani, Jun Han, Andrew Ressler II, Qian Qian, Gregory D. Lyng, Sanjit Singh Batra, and Robert E. Tillman. 2026. Fast and effective on-policy distillation from reasoning prefixes. arXiv preprint arXiv:2602.15260.

Almog Tavor, Itay Ebenspanger, Neil Cnaan, and Mor Geva. 2026. Rethinking selective knowledge distillation. arXiv preprint arXiv:2602.01395. Thinking Machines. 2025. On-policy distillation. https://thinkingmachines.ai/ blog/on-policy-distillation/. Blog post.

Table 8: Complete coding results, LoRA. HumanEval (HE) pass@1.

Training Efficiency

Table 5: Training efficiency, detailed breakdown. Perstep wall-clock time on a single A6000 (48GB). Student: Qwen2.5-Math-1.5B, LoRA, bs=16. Teacher

Method

Gen (s)

Score (s)

Train (s)

Total (s)

Qwen3-1.7B

OPD Ours (N =100)

100–731 ∼5

3–10 ∼1

5–12 ∼2

∼280 ∼8

Qwen3-4B

OPD Ours (N =100)

97–733 ∼5

3–9 ∼1

6–10 ∼2

∼210 ∼8

Qwen3-8B

OPD Ours (N =100)

96–230 ∼5

7–11 ∼1

1–6 ∼2

∼170† ∼8

†

Config

s50

s100

s150

s200

s250

s300

s350

s400

ESR-50 ESR-100 ESR-150 OPD

37.8 37.2 36.6 40.2

39.0 39.0 35.4 31.7

39.6 42.1 36.6 32.3

41.5 37.8 39.0 32.9

40.2 39.0 41.5 27.4

40.9 37.8 39.6 28.0

42.1 37.8 38.4 26.8

40.9 38.4 37.2 26.8

B.3

Coding Results

B.4

Function Calling Results

Frequent OOMs; 48GB insufficient for 8B teacher + vLLM + student.

Table 9: Function calling results (BFCL), LoRA. Name accuracy / Full accuracy / Parse rate. Best full_acc in bold.

Full Experimental Results

B.1

Math Results: Per-Step Performance

Method

Table 6 presents per-step results for the primary math experiments (LoRA, n=1, 3,200 problems). Table 6: MATH-500 per-step results. LoRA, n=1, bs=16, 3,200 problems. Baseline: 50.95% avg@4. Method

Metric

Step 50

Step 100

Step 150

Step 200

ESR-50

avg@4 maj@4 pass@4

62.35 69.40 77.20

66.05 72.00 79.40

66.65 71.00 81.00

64.85 71.20 79.60

ESR-100

avg@4 maj@4 pass@4

63.75 70.00 79.80

64.45 68.40 78.40

65.15 69.60 80.20

65.85 70.80 79.80

ESR-150

avg@4 maj@4 pass@4

65.35 66.80 79.00

66.65 67.00 81.00

65.30 66.30 78.20

65.75 67.30 80.00

ESR-200

avg@4 maj@4 pass@4

66.05 71.20 81.00

64.65 68.40 79.80

65.10 70.00 80.60

65.55 71.20 80.60

B.2

Step 100

Step 150

Step 200

62.35 63.75 65.35 66.05 61.00

66.05 64.45 66.65 64.65 62.00

66.65 65.15 65.30 65.10 62.35

64.85 65.85 65.75 65.55 61.20

77.20 79.80 79.00 81.00 74.60

79.40 78.40 81.00 79.80 75.20

81.00 80.20 78.20 80.60 74.60

79.60 79.80 80.00 80.60 75.00

24.20% 75.30%

ESR-50 ESR-100 ESR-150 ESR-200 OPD

200 100 200 200 100

95.20% 86.20% 88.70% 80.80% 81.00%

57.20% 61.30% 61.50% 54.50% 58.20%

98.30% 91.30% 92.50% 90.20% 86.70%

Token Classification Methodology

2. structural: Punctuation, whitespace, formatting tokens. 3. math_number: Digits (0–9). 4. math_operator: Arithmetic operators (+, −, ×, /, =). 5. math_latex: LaTeX delimiters (\(, \[). 6. continuation: All others. Table 10: Mean KL by token category and position range.

pass@4 ESR-50 ESR-100 ESR-150 ESR-200 OPD

Parse Rate

2.70% 54.00%

1. planning: Reasoning keywords (“To”, “Let”, “First”, “Step”, “We”, “Given”, “Therefore”, “Thus”, “Since”).

avg@4 ESR-50 ESR-100 ESR-150 ESR-200 OPD

Full Acc

9.70% 75.30%

We classify each token into six categories based on string matching:

Table 7: Complete MATH-500 results, LoRA, n=1, bs=16, 3,200 problems. Best per configuration in bold. Step 50

Name Acc

— —

Math Results: Full Per-Step Trajectories (n=1, 3,200 problems, bs=16)

Config

Best Step

Baseline Teacher (Qwen3-1.7B)