PaperlayerView exact revision on arXiv

Source-linked AI summary

Does On-Policy Distillation Need to Complete Every Rollout?

Zhengxin Zhang, Maosen Tang, Ning Wang, Changhao Li, Jingren Zhou, Bo Dai, Sainyam Galhotra, Alex Townsend, Claire Cardie

arXiv:2610.21337v1cs.LGcs.AI

TL;DR

Complete rollouts and teacher scoring for every token make OPD computationally expensive. STOPD randomly truncates rollouts and applies blockwise inverse-probability weighting. It reduces generated tokens by 25–48% and achieves 1.3–2.2× training speedups with accuracy comparable to Vanilla OPD.

  • Problem

    Complete student rollouts and teacher scoring for every token make on-policy distillation computationally expensive.

    Source: §1 · p. 2

  • Method

    STOPD randomly truncates student rollouts and uses blockwise inverse-probability weighting to correct sampling bias while controlling token weights.

    Source: §1 · p. 2

  • Results

    25–48% fewer generated tokens and 1.3–2.2× training speedups, with accuracy comparable to Vanilla OPD across mathematics, code, and instruction following.

    Source: §1 · p. 2, §7 · p. 14

  • Takeaways & Limitations

    STOPD shows that OPD can learn effectively from partial rollouts while preserving Vanilla OPD’s gradient sum in expectation at lower training cost.

    Source: §7 · p. 14

    When all rollouts start together, a long surviving rollout can keep decoding rounds nearly unchanged, limiting generation gains when fixed per-round costs dominate.

    Source: §3.3 · p. 8

Abstract

from arXiv · show

On-policy distillation (OPD) effectively transfers capabilities from a teacher to a student through feedback on student-generated rollouts. However, generating complete student rollouts and computing teacher feedback for every token make OPD computationally expensive. We introduce Stochastic Truncated On-Policy Distillation (STOPD), which accelerates OPD by randomly truncating student rollouts and correcting the resulting sampling bias with blockwise inverse-probability weighting (IPW). At each training step, we sample a stopping position for each rollout before generation and let the student generate until that position or natural termination. We obtain teacher feedback on every generated token and use blockwise IPW to compensate for the lower probability of retaining later tokens while controlling the weights. We prove that this weighting preserves Vanilla OPD's gradient sum in expectation and derive one-step update-error bounds for SGD, Adam, and Muon. We evaluate STOPD across multiple teacher-student model pairs, task domains, and optimizers. STOPD reduces generated tokens by 25-48% while maintaining accuracy comparable to Vanilla OPD, achieving 1.3-2.2x training speedups.

1. Introduction

Complete rollouts and token-level teacher scoring make OPD costly, motivating partial generation without sacrificing effective updates. STOPD uses random truncation and blockwise inverse-probability weighting, preserving expected gradient sums while reducing training cost.

1. Introduction

  • Complete generation, teacher scoring of every token, and student forward and backward passes become more costly as rollouts lengthen and models scale.Sparse supervision reduces student-update time but still requires complete generation; a fixed prefix cutoff omits later-token gradients and biases the gradient-sum estimate.Source: §1 · p. 2
  • STOPD randomly truncates each student rollout and uses blockwise inverse-probability weighting to correct unequal token inclusion while limiting extreme late-token weights.Each rollout stops at a sampled block boundary or natural termination, and every generated token receives teacher feedback.Source: §1 · p. 2
  • Blockwise weighting preserves the full-generation gradient and token-count sums in expectation; the paper also bounds one-step update error for SGD, Adam, and Muon.Token savings can yield greater time savings when long-context computation or cache-limited concurrency improves generation throughput.Source: §1 · p. 2
  • 25–48% fewer generated tokens and 1.3–2.2× training speedups, with accuracy comparable to Vanilla OPD, across mathematics, code, and instruction following.The evaluation spans students from 1.7B to 8B and teachers from 4B to 30B.Source: §1 · p. 2

2. Related Work

Prior OPD work develops objectives, feedback allocation, and multi-teacher variants, while sparse supervision selects loss positions without shortening generation. STOPD instead combines random truncation with inverse-probability weighting to preserve expected gradient sums.

2. Related Work

  • OPD variants refine distribution-matching objectives and teacher-feedback allocation, while MOPD and OPSD extend supervision to specialist teachers and privileged self-teacher contexts [Ma et al., 2026; Zhao et al., 2026].Lightning OPD precomputes and reuses teacher feedback during offline training [Wu et al., 2026].Source: §2 · pp. 2–3
  • Sparse-supervision methods select token positions for distillation, but selecting loss positions alone does not shorten rollout generation [Xu et al., 2026; Wang et al., 2026; Huang et al., 2025; Sheng et al., 2026].ESR truncates at a fixed prefix, leaving later positions with zero supervision probability and an uncompensated gradient contribution [Zhou et al., 2026].Source: §2 · p. 3
  • STOPD differs by randomly truncating rollouts and applying blockwise inverse-probability weighting, preserving the expected full-generation local gradient sum while controlling extreme weights [Beatson & Adams, 2019].The method characterizes the trade-off between generation cost and estimator variance.Source: §2 · p. 3

3. Stochastic Truncated On-Policy Distillation

STOPD samples blockwise rollout cutoffs before generation, weights retained tokens by inverse inclusion probability, and analyzes expected gradient preservation, update error, token savings, and conditions for generation-time gains.

3.1. Training Procedure

  • Randomly sampled stopping positions are rounded up to block ends, and the student generates to that cutoff or natural termination before teacher-scoring retained tokens with inverse-probability weights.Blockwise truncation limits extreme weights; larger blocks trade additional generation for less extreme weighting.Source: §3.1 · pp. 3–4
  • After Vanilla OPD warmup, STOPD estimates the complete token count from weighted batches and uses the preceding batch’s estimate to normalize the current weighted loss.The procedure uses at least one warmup step and initializes the estimate from the last warmup batch’s exact token count.Source: §3.1 · p. 4

3.2. Gradient Preservation and Update Error

  • Blockwise inverse-probability weighting restores each position’s gradient and token-count contribution in expectation, under conditions including budget-independent next-token distributions and unchanged retained-token losses.Stopping positions must also be independently sampled across rollouts and independently of full-rollout and teacher-feedback randomness.Source: §3.2 · p. 5
  • Batch averaging reduces stochastic gradient discrepancy at rate B^-1/2, and the analysis bounds its one-step parameter-update effect for SGD, Adam, and Muon.The update comparison uses the same optimizer, parameters, and prior optimizer state.Source: §3.2 · pp. 6–7

3.3. Token and Time Savings

  • For block width K much smaller than rollout length T, expected token savings are approximately T^2/(2L), or about 25% at T=L/2 and nearly 50% at T=L.When rollouts greatly exceed prompt length, the corresponding expected KV-area savings are about 33% at half the maximum length and 67% at the maximum.Source: §3.3 · p. 7
  • In 1.7B mathematics experiments, 34% fewer tokens corresponded to a 2.0× generation speedup, exceeding the 1.5× ratio implied by token savings alone.Generation-time gains depend on decoding costs and scheduling; simultaneous rollout starts can leave decoding-round counts nearly unchanged when a long rollout survives.Source: §3.3 · p. 8

4. Experimental Setup

Experiments compare STOPD and Vanilla OPD under shared model, data, training, evaluation, and timing settings across mathematics, code, and instruction following.

4. Experimental Setup

  • The main experiment distills Qwen3-4B-Instruct-2507 into Qwen3-1.7B-Base; mathematics scaling runs use a Qwen3-30B-A3B teacher with 4B and 8B students.All models run in non-thinking mode.Source: §4 · p. 8
  • Mathematics uses Nemotron-Cascade-RL-Math, code uses DeepCoder-Preview, and instruction following uses Nemotron-RL instruction-following data.Main runs use AdamW for 250 steps, batch size 256, three seeds, and STOPD block width 128 after 25 Vanilla warmup steps.Source: §4 · p. 8
  • Evaluation averages sampled-response scores across domain-specific benchmarks, and timing measures wall-clock time per training step including generation, teacher scoring, student updates, and overhead.Main results are averaged over three training seeds.Source: §4 · p. 8

5. Experiments

The experiments compare STOPD with Vanilla OPD across domains, model scales, optimizers, block widths, and batch sizes. STOPD generally retains similar accuracy while reducing token use and training time, with larger batches also reducing truncation-induced fluctuations.

5.1. Main Results

  • STOPD remains closely matched to Vanilla OPD in accuracy across mathematics, code, and instruction following, with a mean relative difference of −0.83% ± 1.62%.The largest absolute difference across comparisons is 5.00%.Source: §5.1 · p. 9, Figure 2 · p. 9
  • 1.8×, 2.2×, and 1.3× training speedups accompany 34%, 48%, and 25% fewer tokens in mathematics, code, and instruction following, respectively.Generation is 2.0×, 2.4×, and 1.4× faster in those domains; student updates are 1.2–1.8× faster.Source: §5.1 · pp. 9–10, Table 1 · p. 9
  • STOPD and Vanilla have similar overall training-dynamics curves across the three domains, while STOPD shows larger local fluctuations from random truncation.The compared curves include entropy, distillation loss, gradient norm, and response length.Source: §5.1 · p. 10, Figure 3 · p. 10

5.2. Scaling

  • 1.8× speedups for the 4B student and 1.7× for the 8B student accompany 35% and 33% fewer tokens, respectively.With the 30B teacher, STOPD’s accuracy is closer to Vanilla at larger scales than in the main 4B-teacher/1.7B-student setting.Source: §5.2 · pp. 10–11, Figure 4 · p. 11

5.3. Optimizer

  • Under SGD and Muon, STOPD generates 41% and 37% fewer tokens and speeds up training by 2.2× and 1.9×, respectively, while accuracy remains close to Vanilla.The largest absolute relative accuracy differences are 3.2% under SGD and 3.6% under Muon.Source: §5.3 · p. 11, Figure 6 · p. 12

5.4. Block Width

  • For block widths 32, 128, and 512, STOPD generates 34.6%, 34.3%, and 33.1% fewer tokens and remains 1.8× faster; at width 2048, savings fall to 28.6% and speedup to 1.7×.Accuracy is consistent across widths; width 128 is the default because it balances weight control and speedup.Source: §5.4 · pp. 12–13, Figure 8 · p. 13

5.5. Batch Size

  • Across batch sizes 128, 256, and 512, STOPD stays closely matched to Vanilla and speeds up training by 2.0×, 2.1×, and 2.1×, respectively.The largest absolute accuracy difference falls from 3.40% at batch size 128 to 0.73% at 512.Source: §5.5 · p. 13, Figure 9 · p. 13
  • STOPD’s loss jitter falls from 12.2% at batch size 128 to 7.5% at 256 and 6.1% at 512, consistent with batch averaging reducing truncation-induced gradient fluctuations at an O_P(B^-1/2) rate.Vanilla’s loss jitter falls from 3.4% to 2.1% and 1.5% across the same batch sizes.Source: Figure 9 · p. 13, §5.5 · p. 14

6. Discussion

STOPD’s random truncation and blockwise IPW are presented as applicable beyond the main OPD setting. The discussion outlines adaptations for several distillation and training workflows.

6. Discussion

  • Offline KD stores truncated sequences and their inclusion probabilities during dataset construction, while online KD applies blockwise IPW within each batch using the current teacher.Both settings could reduce generation and student-update computation.Source: §6 · p. 14
  • For supervised fine-tuning, STOPD can randomly truncate reference responses and apply blockwise IPW to retained next-token cross-entropy losses.Source: §6 · p. 14
  • For on-policy self-distillation, the complete reference solution is retained while the student rollout is randomly truncated and its retained-token distillation losses are weighted.The authors state this could reduce computation cost.Source: §6 · p. 14
  • For multi-teacher OPD, each rollout keeps its domain-based teacher assignment, is independently truncated, and receives blockwise IPW to preserve the expected local gradient sum within its domain.This could reduce the computation needed to integrate specialist teachers’ capabilities into a shared student.Source: §6 · p. 14

7. Conclusion

STOPD combines random rollout truncation with blockwise IPW; the conclusion reports reduced token generation and training time with accuracy comparable to Vanilla OPD. The appendix analyzes expected-gradient preservation, weighting, variance, and normalization error.

7. Conclusion

  • Across mathematics, code, and instruction following with 1.7B–8B students, STOPD generates 25–48% fewer tokens and achieves 1.3–2.2× training speedups while maintaining accuracy comparable to Vanilla OPD.The reported gains persist with SGD or Muon in place of AdamW and across block widths.Source: §7 · p. 14

A. Gradient Preservation and Normalization

  • Random truncation preserves the expected full-generation gradient sum, and the appendix derives one-step update-error bounds for SGD, Adam, and Muon.The analysis also covers normalized gradients and batch averaging.Source: §7 · p. 14, A. Gradient Preservation and… · p. 18

A.1. Conditional preservation and estimation error

  • Under the stated conditions, full and truncated rollouts can be coupled so retained local gradients match, and inverse-probability weighting preserves the expected gradient sum and token count.The weights depend on position, rollout limit, and block width, not on complete rollout lengths.Source: §A.1 · pp. 18–19

A.2. Position weights and token coverage

  • Matching Vanilla OPD’s expected contribution for every token requires weighting retained position t by 1/p_t; positions in the first block are retained in every rollout that reaches them.Source: §A.2 · p. 21

A.3. Truncation variance and block width

  • For max-length rollouts, larger blocks reduce weighted-count variance but require more generation; with one block variance is zero, while it approaches one as the number of blocks grows.For K=32, 128, 512, and 2048, the reported relative standard deviations are 0.99, 0.98, 0.93, and 0.81, respectively.Source: §A.3 · p. 22

A.4. Error under lagged normalization

  • With lagged normalization, the current and preceding batches’ mean rollout lengths can differ, creating a scale difference in addition to cutoff error in the gradient bound.The finite-batch analysis bounds these contributions under its sampling assumptions.Source: §A.4 · pp. 22, 24

A.5. Batch size and gradient error

The appendix analyzes how batch size affects STOPD’s gradient discrepancy from Vanilla OPD. Larger batches reduce random-truncation variation and token-count sampling error, while shifts in mean rollout length can still affect normalization.

A.5. Batch size and gradient error

  • Larger batches reduce variation from random truncation and relative sampling error in the token-count estimate, but changes in mean rollout length can still introduce a scale difference.Source: §A.5 · p. 25
  • As batch size grows, the local gradient discrepancy converges to zero in probability if adjacent batches’ conditional mean rollout lengths have a ratio approaching one.When those adjacent means agree, the discrepancy is of order B^-1/2.Source: §A.5 · p. 25
  • STOPD and Vanilla OPD have the same population token-mean gradient, defined as expected gradient sum divided by expected token count.Source: §A.5 · p. 26

B. Optimizer Update Error

This section analyzes STOPD’s parameter-update error relative to Vanilla OPD, deriving optimizer-specific bounds from gradient error, including clipping and conditional mean-square guarantees. It also derives expected token savings and analyzes their effect on generation time, noting how Muon’s coefficient bounds relate small gradient errors to small update errors.

B. Optimizer Update Error

  • The conditional mean-square update error is bounded by Λ_opt^2{V/C^2 + ρ^2||g_V||^2}, separating cutoff noise from lagged-denominator mismatch; global-norm clipping is nonexpansive.The optimizer coefficient is evaluated at the clipped Vanilla reference.Source: B. Optimizer Update Error · p. 27

B.1. Stochastic gradient descent

  • For SGD, update error equals η times gradient error; with momentum, the coefficient becomes η|1−β|, and affine updates preserve the expected update at the mean gradient.The reported SGD runs use neither momentum nor clipping and set η = 5 × 10^-4 at every step.Source: §B.1 · pp. 27–28

B.2. Adam

  • Adam’s coupled dependence on each current gradient coordinate’s numerator and denominator is controlled by finite coordinatewise coefficients, which remain moderate across training.The coefficient bounds are examined across coordinates and checkpoints, with a median-by-step summary and theoretical ceiling in Figure 10.Source: §B.2 · pp. 28, 31, Figure 10 · p. 31

B.3. Muon

  • For each matrix block, a finite coefficient bounds Muon’s STOPD–Vanilla update difference pointwise and yields a conditional mean-square bound after averaging cutoffs.The derivation tracks error through the Nesterov input, norm-clamped normalization, and polynomial iterations.Source: §B.3 · pp. 32–33
  • Muon’s coefficient bounds remain moderate during training, so small gradient errors also yield small update-error bounds.The analysis examines the coefficient distributions across blocks and checkpoints, the median block, and the theoretical ceiling.Source: §B.3 · p. 34, Figure 11 · p. 35

C.1. Token savings

  • About 25% of tokens are saved for T = L/2 and nearly 50% for T = L; rollouts ending within the first block save none.When K divides L, the exact saving at maximum length is D(L) = (L − K)/2.Source: §C.1 · p. 36

C.2. Time savings

Random truncation can reduce generation and KV-cache work, with time savings depending on rollout lengths and scheduling. The section also identifies the appendix tables reporting accuracy by training step.

C.2. Time savings

  • For rollouts reaching half the maximum length, expected savings are about 25% of tokens and 33% of KV area; at maximum length, about 50% and 67%, respectively.These approximations assume rollouts are much longer than prompts and neglect discrete and block-rounding corrections for K ≪ T_i.Source: §C.2 · p. 37
  • When long-context processing dominates decoding cost, time savings can exceed the fraction of tokens saved.Random truncation skips later tokens that require processing longer cached contexts.Source: §C.2 · pp. 36, 37, 39
  • Scheduling affects speedup through the number of decoding rounds and the context processed per round.With long-context processing dominant, reducing KV area can improve speedup even when the longest retained rollout remains close to the maximum length.Source: §C.2 · pp. 37–39

D. Accuracy by Step

  • Tables 2–7 report accuracy by training step for the experiments shown in Figures 2, 4, 6, 8, 9, 12, and 13.Table 2 covers Figures 2 and 4; Tables 3–7 cover Figures 6, 8, 9, 12, and 13.Source: D. Accuracy by Step · p. 39

E. Additional Time Decomposition

In the mathematics comparison, STOPD and Vanilla continue improving after ESR’s accuracy plateaus, while ESR’s response length and truncation rate keep rising. The appendix also breaks down training-step cost across model sizes and configurations.

E. Additional Time Decomposition

F. Comparison with Early Stopping Rollout

  • At step 250, STOPD and Vanilla outperform ESR by 5.9 and 6.6 points, respectively.ESR performance saturates after approximately 50 training steps, while STOPD and Vanilla continue increasing.Source: Figure 12 · p. 42, F. Comparison with Early Stopping… · p. 42
  • ESR’s response length and truncation rate reach approximately 20k tokens and 57% at step 250, while Vanilla and STOPD peak around step 50 before declining to approximately 10k tokens and 20%.The passage attributes ESR’s divergence to fixed-prefix distillation excluding end-of-sequence supervision when responses exceed the truncation boundary.Source: F. Comparison with Early Stopping… · p. 42

G. Removing the Blockwise Inverse-Probability Weights

Setting every STOPD weight to 1 maintains accuracy close to Vanilla, but its responses can initially be longer and more frequently truncated. Blockwise weighting keeps response length closer to Vanilla over the reported training steps.

G. Removing the Blockwise Inverse-Probability Weights

  • The w_t = 1 variant maintains accuracy close to Vanilla and STOPD for both the 1.7B and 4B students.The comparison uses Qwen3-4B-Instruct-2507 and Qwen3-30B-A3B-Instruct-2507 teachers, respectively.Source: G. Removing the Blockwise… · pp. 42–43, Figure 13 · p. 43
  • At step 100, the 4B w_t = 1 variant’s responses are 43% longer and its truncation rate is 34.7%, versus 16.1% for Vanilla.By step 250, the response-length gap narrows to 3.6%.Source: G. Removing the Blockwise… · p. 43
  • Over steps 50–175, STOPD stays within 6% of Vanilla’s response length for the 4B student and within 10% for the 1.7B student.Blockwise weighting further aligns response length and truncation behavior with Vanilla compared with the w_t = 1 variant.Source: G. Removing the Blockwise… · p. 43

H. Per-Benchmark Accuracy

This section breaks average accuracy into per-benchmark curves across model and task settings, optimizers, and other training configurations. The plots show accuracy against training step, with seed variation indicated.

Per-Benchmark Accuracy

Related papers

Loading 2610.21337v1…