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.
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