Source-linked AI summary
Fast and Effective On-policy Distillation from Reasoning Prefixes
Dongxu Zhang, Zhichao Yang, Sepehr Janghorbani, Jun Han, Andrew Ressler, Qian Qian, Gregory D. Lyng, Sanjit Singh Batra, Robert E. Tillman
TL;DR
On-policy distillation can improve learning from student-induced states but requires costly on-the-fly sampling, especially for long-form reasoning. This paper distills only prefixes of truncated student rollouts, finding that prefix OPD closely matches full OPD while reducing training FLOP by 2×–47×.
Problem
On-policy distillation requires expensive on-the-fly student sampling, while full rollouts can be prohibitively costly for long-form reasoning.
Method
The method truncates student rollouts, applies token-level distillation only to their prefixes, and progressively increases the trained prefix length during training.
Results
Prefix OPD closely matches full OPD on AI-for-Math and out-of-domain benchmarks while reducing training FLOP by 2×–47×.
Takeaways & Limitations
Scheduled prefix OPD matches full OPD performance across settings while using less than 50% of the FLOP.
Takeaways & Limitations
The approach is evaluated primarily on long-form reasoning tasks, may require warmup for weaker students, and requires teacher and student to share a vocabulary.
Abstract
from arXiv · showhide
On-policy distillation (OPD), which samples trajectories from the student model and supervises them with a teacher at the token level, avoids relying solely on verifiable terminal rewards and can yield better generalization than off-policy distillation. However, OPD requires expensive on-the-fly sampling of the student policy during training, which substantially increases training cost, especially for long responses. Our initial analysis shows that, during OPD, training signals are often concentrated in the prefix of each output, and that even a short teacher-generated prefix can significantly help the student produce the correct answer. Motivated by these observations, we propose a simple yet effective modification of OPD: we apply the distillation objective only to prefixes of student-generated outputs and terminate each sampling early during distillation. Experiments on a suite of AI-for-Math and out-of-domain benchmarks show that on-policy prefix distillation matches the performance of full OPD while reducing training FLOP by 2x-47x.
1 Introduction
On-policy distillation addresses limitations of teacher-only training and reward-based reinforcement learning by supervising student-generated trajectories token by token, but long rollouts make each update expensive. The paper proposes prefix distillation, which truncates training rollouts, supervises only their prefixes, and matches full OPD while reducing compute.
- Knowledge distillation compresses large language models into smaller students for deployment where resources are limited.
- Sequence-level distillation trains students on fixed teacher-generated trajectories, so students rarely learn from their own induced states.This exposure bias can allow errors to accumulate at test time.
- On-policy learning samples trajectories from the current student, reducing train–test mismatch, but many reinforcement-learning methods require expensive full rollouts for terminal rewards.
- OPD provides dense token-level teacher feedback without task-specific reward design, yet on-the-fly student rollouts and teacher scoring become prohibitively expensive for long reasoning trajectories.
- The paper asks whether prefix-only supervision can preserve OPD’s benefits while reducing training cost, motivated by evidence that early-token learning signals are often sufficient.
- On-policy prefix distillation truncates student rollouts, applies distillation only to the prefix, and progressively increases the trained prefix length during training.The method’s reported results match full OPD while reducing training FLOP by 2×–47×.
2 Preliminary Study
On-policy distillation trains on student-generated trajectories while matching a teacher token by token, avoiding off-policy distribution mismatch but incurring substantial sampling and scoring costs. Preliminary analyses find that training signal concentrates in early tokens and that teacher prefixes can improve untrained students’ continuation accuracy.
- On-policy distillation: On-policy distillation minimizes token-level divergence between student-generated outputs and a fixed teacher.Its expectation over the student’s state distribution avoids the distributional mismatch associated with off-policy imitation.
- On-policy distillation: Reverse-KL loss is estimated from sampled student trajectories rather than all possible responses.The implementation sums token-level losses over sequence length and avoids storing full-vocabulary logits for every trajectory token.
- Prefix signal: Training signal concentrates in earlier generated tokens across training steps and student model sizes.The authors conjecture that early positions capture high-level planning, whereas later positions more often execute an existing plan.
- Teacher-prefix study: Teacher prefixes improve untrained students’ AIME-24 continuation performance, with sharp gains after the first few thousand tokens and slower increases thereafter.The study varies teacher-prefix lengths of 2048, 4096, 6144, and 8192 tokens while preventing the final answer from appearing in the prefix.
3 Methodology
Prefix on-policy distillation limits training rollouts to a chosen prefix, back-propagates through the truncated student trajectory, and removes the cap at evaluation. A linear schedule can progressively increase the trained prefix length.
- Prefix distillation: Prefix distillation early-stops on-policy generation at Ltrain and back-propagates through the truncated student trajectory.The on-policy property is preserved for the first Ltrain tokens without modifying the loss or sampling procedure.
- Prefix distillation: The training cap is removed at evaluation so final performance is not artificially constrained.The remaining evaluation budget allows up to 16,384 tokens.
- Special tokens: The training prompt enforces <think> at its end to align student and teacher behavior around special-token generation.This avoids repeated attempts to switch tokens that can produce unreadable content.
- Prefix scheduling: Prefix scheduling increases Ltrain after each gradient update to expose the student to longer-horizon feedback.The reported schedule starts at 1 token and increments by 256 tokens.
4 Experiments
Experiments evaluate prefix on-policy distillation across in-domain and out-of-domain math benchmarks, comparing accuracy and training cost across student sizes and masking strategies. Prefix scheduling matches full OPD performance with substantially less computation, while very short prefixes can hurt smaller students’ out-of-domain accuracy.
- Main Results: Prefix OPD maintains competitive in-domain performance while reducing training cost, using less than 10% of full-budget FLOP for Qwen3-8B-Base.It outperforms limited-budget baselines on AIME-24 and AIME-25 with half the FLOP and is only slightly worse than 300-step SeqKD on MATH500.
- Main Results: For Qwen3-1.7B-Base, very short prefixes can reduce out-of-domain accuracy below the pretrained student, whereas prefix scheduling remains competitive.For Qwen3-8B-Base, increasing the prefix length moves performance closer to full OPD.
- Scheduling: Scheduled prefix OPD matches full OPD’s upper-bound performance across settings while using less than 50% of the FLOP.The linear schedule progressively increases the supervised prefix length during training.
- Ablation Study: The prefix window performs best among the tested masking windows, supporting “earlier is more important” as a heuristic.The ablation compares prefix, interior sliding-window, and tail supervision using equal-length windows.
- Case Study: The case study shows that a 128-token teacher prefix helps the base student avoid early reasoning errors and reach the correct answer, while prefix OPD enables this without an explicit teacher prefix.The example progresses from an incorrect direct solution to the correct result m + n = 33 after teacher-prefix guidance or prefix OPD.
5 Discussion
Prefix length creates a compute–accuracy trade-off: short prefixes deliver strong early gains, while longer prefixes can perform better with sufficient training budget. Prefix learning can also affect tokens beyond the supervised region, with effects depending on student size and training setting.
- Cost and Performance Trade-off: Shorter prefixes often provide larger accuracy gains early, whereas longer prefixes can achieve higher accuracy under larger training budgets.The authors recommend prefix scheduling rather than full-length distillation from the beginning and report that this can often save at least half the training cost.
- Effect on the Tail: Prefix distillation reduces losses outside the supervised prefix, with the strongest reduction near the prefix and diminishing effects toward the tail.This result is reported for Qwen3-8B-Base with a prefix of 1024 tokens.
- Effect on the Tail: Student size changes the tail effect: prefix learning benefits sufficiently large students but can increase tail loss in Qwen3-1.7B-Base.The paper attributes the small-model behavior to competition between prefix and tail loss reduction rather than shared benefit.
- Generalized Prefix Distillation: On-policy plus reverse KL is the strongest and most stable short-prefix setting, while off-policy reverse KL yields only mild gains and off-policy forward KL can hurt performance.On-policy forward KL remains effective but performs worse than on-policy reverse KL with the same prefix budget; similar findings hold for Qwen3-8B-Base.
- Generalized Prefix Distillation: Prefix truncation can amplify distribution mismatch in off-policy training because supervised states may not be reliably visited by the student during rollout.With supervision limited to the prefix, the student must generate the tail from potentially mismatched states.
- Generalized Prefix Distillation: Short-prefix distillation also reduces the memory footprint of full-distribution objectives for long-form reasoning.The authors report that increasing rollout length beyond 512 can cause out-of-memory errors in their Qwen3 GKD experiments.
6 Related Work
The work builds on on-policy learning and token-selection research that emphasizes policy alignment and the unequal usefulness of reasoning tokens. It focuses distillation on informative early tokens to retain most full-OPD benefits while reducing compute on long rationales.
- On-policy Distillation: On-policy learning collects training data from the policy being optimized, reducing distributional mismatch between training and deployment.This alignment is described as often leading to more robust behavior.
- Token Selection during LLM Training: Prior token-selection studies suggest that reasoning tokens are not equally useful and that later chain-of-thought tokens can become noisy or redundant.The paper distinguishes its task from these studies while sharing their focus on informative early tokens.
- Token Selection during LLM Training: The proposed distillation follows this early-token emphasis, yielding substantial compute savings while retaining most benefits of full OPD on long rationales.This connection is presented as an instinct shared with token-selection work rather than as the same training task.
7 Conclusion
The paper introduces on-policy prefix distillation, which truncates student rollouts and applies token-level teacher feedback only to their prefixes. Across AI4Math benchmarks, it improves over the base model, remains competitive with full OPD, and reduces training cost.
- Method: On-policy prefix distillation truncates student rollouts and applies token-level teacher feedback only to each sample’s prefix.It is presented as a simple modification of on-policy distillation.
- Results: Across AI4Math benchmarks, the approach achieves substantial gains over the base model and remains competitive with full on-policy distillation.The conclusion reports this performance pattern without narrowing it to a single benchmark.
- Takeaway: The method significantly reduces training cost while retaining performance close to full on-policy distillation.This summarizes the paper’s stated efficiency–performance conclusion.
8 Limitations
The paper identifies scope and deployment boundaries for prefix distillation, including reasoning-focused evaluation, requirements on student capability and shared vocabulary, and risks from under-training later-response behaviors.
- The evaluation covers reasoning tasks with long chain-of-thought outputs, leaving applicability to summarization and story generation unknown.
- Prefix-only training assumes the base model can follow instructions and produce valid output formats; weaker students may need SFT or full OPD warmup.
- The method requires teacher and student to share a vocabulary so token-level losses are well-defined.
- Applying gradients only to early rollout tokens may under-train later behaviors such as safety refusals, calibration, and self-correction.The paper recommends progressively increasing prefix length and evaluating safety and calibration before deployment.
Qwen3-8B-Base after Prefix On-policy Distillation (Correct)
Solving the logarithmic system by substituting a=log_2 x, b=log_2 y, and c=log_2 z yields the requested value and final answer 33.
- Setup: The three ratio equations become a linear system in a=log_2 x, b=log_2 y, and c=log_2 z.
- Solving the system: Solving the system gives a=-7/24, b=-3/8, and c=-5/12.
- Computing the target: The target logarithm equals 4a+3b+2c=-25/8, so its absolute value is 25/8.
- Final answer: 33 is the final value of m+n when m/n=25/8.
B FLOP Estimation Details
The paper estimates training compute with a consistent FLOP-accounting scheme that combines student training, student sampling and scoring, and teacher scoring, while treating FLOP as a relative proxy.
- Ftrain(M,L) estimates one student training update for model M on sequence length L, including forward, backward, and optimizer computation.
- Forward-only cost is approximated as Ffwd(M,L)≈1/3 Ftrain(M,L), applied uniformly across methods.
- For OPD, per-rollout compute is estimated as FOPD(L)≈Ftrain(S,L)+2Ffwd(S,L)+Ffwd(T,L), with prefix OPD capped at Ltrain.
- With batch size B and K rollouts per prompt, step cost is Fstep≈B·K·FOPD(L), summed over optimization steps.
- FLOP is a consistent comparison proxy rather than an exact hardware-cost measurement because implementation effects such as caching and communication are omitted.
C Prefix SeqKD with Qwen3-8B-Base
The paper tests whether prefix distillation transfers to off-policy SeqKD with Qwen3-8B-Base, comparing AIME-24 and MATH500 results with prefix OPD.
- The off-policy SeqKD setup caps maximum context length at 2048 and plots prefix-2048 OPD for comparison.
- The reported finding is that prefix distillation effectiveness does not transfer well to off-policy SeqKD, with large learning rates quickly causing degradation.
D Effect of Learning Rate
The learning-rate experiments find that reducing the learning rate does not significantly improve prefix on-policy distillation’s out-of-domain generalization. The study uses AIME-24 for checkpoint selection and development, then evaluates the selected checkpoints on the remaining benchmarks, while leaving broader hyperparameter analysis for future work.
- Learning-rate sweep: 5e-5 is the optimal learning rate with batch size 512 in the reported prefix OPD experiment setting.Performance differs significantly across learning rates in the sweep.
- Checkpoint selection: Checkpoints with similar AIME-24 accuracy are selected across learning rates to control in-domain performance when evaluating out-of-domain generalization.The selected checkpoints use learning rates of 5e-5, 2e-5, and 5e-6 at the reported training steps.
- Generalization result: Reducing the learning rate does not significantly improve prefix on-policy distillation’s generalization capability.This finding is based on the results reported in Table 4.
- Scope and limitations: The analysis does not cover a broader grid of learning rate, scheduling, step size, batch size, and prefix-length settings.The authors leave this extensive analysis for future work; learning rate 2e-4 was excluded because it dramatically hurt model performance.