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 & tables
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.
In-Domain
Out-domain
Training cost
MATH500
AIME-24
AIME-25
GPQA
MMLU-Pro
FLOP
Hours
Qwen3-1.7B-Base
20.1
0.2
0.0
11.2
9.1
No post training
w. limited budgets
OPD - 10 steps
43.4
5.4
3.5
18.3
30.7
8.2×1018
12.2
SeqKD - 300 steps
60.5
4.6
6.9
20.2
28.5
1.3×1019
3.7
OPD (prefix 256)
50.0
4.8
6.3
1.0
1.5
9.4×1017
2.4
Table 1: Mean accuracies and costs with Qwen3-1.7B-Base as the student.
In-Domain
Out-domain
Training cost
MATH500
AIME-24
AIME-25
GPQA
MMLU-Pro
FLOP
Hours
Qwen3-8B-Base
39.0
5.2
2.1
17.6
26.5
No post training
w. limited budgets
OPD - 10 steps
77.6
23.3
24.2
47.5
62.7
2.7×1019
23.3
SeqKD - 300 steps
88.3
31.9
25.4
38.5
52.6
4.4×1019
11.3
OPD (prefix 256)
82.0
30.8
26.3
35.5
58.9
3.4×1018
3.2
Table 2: Mean accuracies and costs with Qwen3-8B-Base as the student.
Method
Step 10
Step 20
Step 30
OPD (full)
23.3
40.8
40.0
OPD on prefix 1024
26.3
32.5
34.8
OPD on [1024,2048)
14.0
16.0
16.0
OPD on [2048,3072)
18.8
16.9
15.8
OPD on [3072,4096)
9.4
19.2
22.5
OPD on tail 1024
17.9
23.8
23.8
Table 3: Ablation of masking strategies with Qwen3-8B-Base as student and Qwen3-8B as the teacher. We report Mean Acc.@16 on AIME-24.
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.
Figure 6: Performances of prefix generalized knowledge distillation. The student is Qwen3-1.7B-Base; Questions and off-policy signals are sampled from Openthoughts3. λ=1 uses on-policy data, and 0 uses off-policy data, while β=1 uses reverse KL and 0 uses forward KL.
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 7: Performances of off-policy prefix distillation. The student is Qwen3-8B-Base; Off-policy prefix signals are sampled from Openthoughts3. Top : AIME-24 Bottom : MATH500
Figure 8: The effect of learning rate on AIME-24 during the training of on-policy prefix distillation on OpenThoughts3. The student is Qwen3-8B-Base and the teacher is Qwen3-8B.
AIME-24
MMLU-Pro
GPQA
(control)
step 0
5.2%
26.5%
17.6%
lr=5e-5, step 10
26.3%
55.0%
27.9%
lr=2e-5, step 10
25.0%
49.3%
26.6%
lr=5e-6, step 50
25.2%
47.8%
26.6%
Appendix
Table 4: Effect of learning rate on generalization capability. We report validation performance after 60 steps of OPD (Prefix=1024) on OpenThoughts3, using Qwen3-8B-Base as the student and Qwen3-8B as the teacher.
Prefix length
Checkpoint for testing
256
Step 20
512
Step 60
1024
Step 60
2048
Step 40
Prefix scheduling
Step 60
Full
Step 50
Appendix
Table 5: Checkpoints for testing with Qwen3-1.7B-Base as the student.
Prefix length
Checkpoint for testing
256
Step 60
512
Step 50
1024
Step 40
2048
Step 60
Prefix scheduling
Step 60
Full
Step 60
Appendix
Table 6: Checkpoints for testing with Qwen3-8B-Base as the student.
Problem ID
Number of Solutions
1
2
2
6
3
5
4
2
5
7
6
7
Appendix
Table 7: Number of solutions for each question in AIME-2024 from AoPS forum https://artofproblemsolving.com/wiki/index.php?title=2024_AIME_I
Language model distillation transfers a larger model’s capabilities to smaller or specialized models using outputs, probabilities, or internal representations.
On-policy distillation (OPD) improves reasoning models by applying dense teacher supervision on student-sampled trajectories. However, scaling OPD to long-horizon mathematical reasoning exposes a reliability and efficiency problem: standard OPD assigns every sampled candidate the same long rollout budget, even though some trajectories may quickly become weakly aligned with the teacher and provide less useful supervision. Prior analyses suggest that successful OPD depends on local teacher-student compatibility, which can be measured by top-k overlap on student-visited prefixes. When this overlap is low, continuing to generate or train on long suffixes may waste computation and introduce noisy learning signal. To address this, we introduce Prefix-Guided On-Policy Distillation (PG-OPD), a simple rollout-allocation framework that uses fixed-length prefixes to estimate trajectory value before expensive long-horizon generation. PG-OPD first decodes every sampled candidate to the same prefix length, computes teacher-student top-k overlap within an early probe window of that prefix, and selectively continues high-overlap candidates to a fixed long length. Low-overlap candidates stop at the fixed prefix, avoiding unnecessary suffix generation. Across diverse teacher-student combinations on AMC, AIME, and HMMT benchmarks, PG-OPD improves average accuracy by up to 4.80 points while reducing training time by up to 2.46x. These results suggest that prefix-level compatibility provides a practical signal for directing OPD computation toward trajectories that remain learnable from the teacher.
On-policy distillation transfers reasoning capabilities by training a student model on its own generated trajectories using token-level feedback from a teacher. However, we identify a critical bottleneck, Supervision Fidelity Decay (SFD): as student-generated prefixes lengthen, the teacher's next-token distribution becomes less confident and less discriminative. Consequently, the teacher-dependent corrective signal in reverse-KL distillation weakens, causing student drift to compound across long reasoning chains. To mitigate SFD, we introduce Lookahead Group Reward (\ours{}). Building on the insight that next-step teacher confidence reflects the discriminative strength of future reverse-KL supervision, \ours{} evaluates the student's top-K candidate tokens by the teacher confidence they induce at the subsequent step and assigns a group-normalized reward. To maintain computational efficiency, we further design an entropy-triggered tree-attention mechanism. Across six math and code benchmarks, \ours{} improves mean@8 by 2.57 points over OPD for a 7B student, with gains increasing in longer-generation and reaching +4.92 points on AIME-26 at 39k tokens.
Yanjiang Liu, Jie Lou, Xinyan Guan +7
University of Chinese Academy of Sciences · 2Chinese Information Processing Laboratory Institute of Software, Chinese Academy of Sciences 2University of Chinese Academy of Sciences, Beijing, China · 3Xiaohongshu
On-policy distillation (OPD) grounds token-level supervision in the student's own trajectory, yet suffers from prefix failure: once the student commits to a wrong reasoning direction, all subsequent generation builds on this deviation, producing misdirected continuations that elicit unreliable supervision and waste compute. We identify a teacher-student continuation asymmetry on failed prefixes, where the teacher tends to redirect while the student continues along the original direction, and convert it into a label-free handoff trigger in Relay On-Policy Distillation (Relay-OPD). During training, Relay-OPD constructs relay trajectories by letting the teacher briefly take over at detected trigger points to produce a teacher leg, after which the student resumes and is optimized on the resulting trajectory. A limited relay budget concentrates intervention on critical early positions while limiting departure from the student policy. With a Qwen3-4B-Instruct-2507 teacher and Qwen3-0.6B/1.7B-Non-Thinking students on eight mathematical reasoning benchmarks, Relay-OPD achieves the best or second-best results on every benchmark, outperforming standard OPD by +5.73% and the strongest baseline FastOPD by +1.49% on average for 1.7B, with consistent gains at 0.6B. Training trajectory length is reduced by over 50%.
Haolei Xu, Xiaowen Xu, Haiwen Hong +5
1Zhejiang University · 2Yuvion Team, Alibaba Group