Fast and Effective On-policy Distillation from Reasoning Prefixes

TL;DR

This paper introduces on-policy prefix distillation, a simple modification of on-policy distillation (OPD) that supervises only the initial prefix of each student trajectory and gradually increases the trained prefix length via a linear schedule. By focusing training signals on early tokens where learning is most impactful for reasoning, prefix distillation achieves comparable performance to full OPD while reducing training FLOP by a factor of 2x–47x on AI-for-Math and out-of-domain benchmarks. The approach preserves the on-policy nature, uses a lightweight prefix strategy, and integrates a minimal special-token design to align teacher-student behavior. Empirical results across MATH, AIME, GPQA, and MMLU-Pro show that scheduled prefix OPD maintains strong accuracy with lower compute, offering practical advantages for scalable reasoning tasks while highlighting the importance of early-token supervision and careful tail handling. Limitations include potential tail-safety calibration concerns and the need for further exploration of adaptive schedulers and broader task domains.

Abstract

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.

Figures (8)

• Figure 1: On-policy prefix distillation with prefix scheduling. During training, gradients are backpropagated only through the first $L$ tokens of each student rollout, and sampling is terminated after $L$. Prefix scheduling increases $L$ by a fixed amount at each training step. At test time, the student generates autoregressively until producing the final answer.
• Figure 2: Distribution of reverse-KL loss along the output position during full on-policy distillation over OpenThoughts3. Top: Student model is Qwen3-1.7B-Base. Bottom: Student model is Qwen3-8B-Base.
• Figure 3: Mean Acc.@16 on AIME-24 when an untrained student (Qwen3-8B-Base and Qwen3-1.7B-Base) continues from varying lengths of a teacher (Qwen3-8B)-generated prefix.
• Figure 4: Comparison of training compute vs Mean Acc.@16 on AIME-24. Top: Student is Qwen3-1.7B-Base. Bottom: Student is Qwen3-8B-Base.
• Figure 5: Distribution of reverse-KL loss along the output position during on-policy prefix distillation with prefix=1024. Top: Student model is Qwen3-1.7B-Base. Bottom: Student model is Qwen3-8B-Base.