On-policy distillation (OPD) supplies dense token-level supervision by aligning a student to the teacher's next-token distribution on its own generated rollouts. Existing works mostly implement top-k OPD, applying reverse Kullback-Leibler (\RKL) divergence between the teacher's and the student's renormalized distributions over teacher top-k support. However, the teacher scores student-generated rollouts that are inherently off-policy for itself, leaving the reliability of its next-token guidance unclear. Restricting the divergence to the teacher top-k support also discards the remaining mass outside that support. These gaps leave two questions implicit: 1) which tokens to align, and 2) how to align those tokens. To answer these questions, we propose Overlap-gated On-Policy Distillation (O2PD). As for 1), we quantitatively analyze the reliability of teacher supervision and find that alignment on the student's trajectory prefix is more reliable than alignment on the full context. Therefore, our O2PD is applied only to the 20% prefix tokens of each student rollout. As for 2), we theoretically characterize two problems of top-k OPD: the Teacher-support Uncovering and the Dominant-Mode Mismatch. Based on these analyses, we propose two support strategies: the Reset Support strategy attaches a reset atom to every token, pooling mass outside the teacher top-k, and the Support Convergence strategy converges the support to the already aligned intersection. Our O2PD employs the designed overlap-gated α-divergence to interpolate between mode-covering and mode-seeking for per-token alignment, where the overlap is used as a gate to switch the support and adaptively modulate α.
Figures & tables
Figure 1: Framework of our O 2 PD: (a) Restricting alignment to the prefix, where teacher supervision is reliable. (b) Overlap-gated α -divergence with Reset Support and Support Convergence strategies.
Set
Role
T=TopK(p)
teacher support
S=TopK(q)
student support
I=T∩S
already aligned atoms
T∖S
teacher modes still to cover
S∖T
student-only modes to suppress
Table 1: Supports at each token.
Figure 2: (a) Cross-prefix pass@1 versus prefix ratio ℓ on MATH500. (b) Mean top- k overlap by relative-position bin when LT is applied only to the first 20% of each student trajectory. (c) Top- k OPD loss LT versus relative position at initialization and after full-trajectory training.
Figure 3: (a) Mean overlap by initial-overlap bin before and after training. (b) Destination of tokens with initial ρ<0.3 . (c) Signed gap q−p on the shared top- 1 .
Method
AIME24
AIME25
AIME26
HMMT25
Minerva
Olympiad
MATH500
Len
Avg
Qwen3-1.7B (student) ← Qwen3-4B (teacher)
Qwen3-1.7B
8.33
12.08
5.83
6.25
20.59
36.74
73.8
3049
23.37
Qwen3-4B
23.33
31.67
21.67
15.00
31.99
45.48
85.8
4127
36.42
OPD
13.75
14.17
7.08
5.83
23.53
36.00
74.0
5033
24.91
ESR
13.75
14.17
7.08
8.33
21.32
38.37
74.2
4480
25.32
EOPD
14.17
12.50
11.67
5.00
23.90
38.37
74.6
5110
25.74
Table 2: Performance (%) for three OPD student–teacher pairs. Len is the mean generated length in tokens. Bold marks the best distilled method in each column within a pair; underline marks the second best. Δ vs. OPD is O 2 PD minus OPD.
Figure 4: (a) Method-component ablation studies (%). Each bar removes one component of O 2 PD. (b) Prefix versus suffix performance. Avg is the mean of Avg@8 on AIME24, AIME25, AIME26, and HMMT25.
Figure 5: Mean top- k overlap trend.
Figure 6: Actor entropy during training at the four support-switch percentiles on the Qwen3-1.7B student with the NTRL teacher.
AIME24
AIME25
HMMT25
Percentile
Avg@8
pass@8
Avg@8
pass@8
Avg@8
pass@8
Avg
0.25
51.67
66.67
42.92
63.33
23.75
36.67
39.45
0.5
54.17
73.33
45.83
70.00
25.83
43.33
41.94
0.75
52.50
70.00
45.42
66.67
25.42
40.00
41.11
1.0
53.75
73.33
43.33
66.67
24.58
40.00
40.55
Table 3: Support-switch percentile on the Qwen3-1.7B (student) ← Qwen3-4B-NTRL (teacher).
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 7: Prefix tokens at k=16 before and after O 2 PD. (a) Mean overlap by initial-overlap bin. (b) Destination of tokens with initial ρ<0.3 . (c) Signed q−p on the shared top- 1 .
Student
Teacher
Qwen3-1.7B
Qwen3-4B
DSR1-1.5B
JustRL-1.5B
Qwen3-1.7B
Qwen3-4B-NTRL
Appendix
Table 4: Student–teacher pairs used for O 2 PD.
Setting
Value
Hardware
16 × Ascend 910C
Data
DAPO-MATH-17K, unfiltered
Optimizer
AdamW, β=(0.9,0.999) , weight decay 0.01
Learning rate
2×10−6 , cosine decay, min ratio 0.1
Warmup / grad clip
10 steps / 1.0
Train / mini-batch
64 / 16
Appendix
Table 5: Shared training recipe for O 2 PD.
Setting
Value
Top- k size k
16
Prefix fraction ℓ
0.2
Reset atom
on
α
mapped from prefix overlap min and max
Support switch
median of prefix overlap: T∪{r} below, I∪{r} at or above
Language model distillation transfers a larger model’s capabilities to smaller or specialized models using outputs, probabilities, or internal representations.
On-policy distillation (OPD) accelerates post-training by providing dense token-level supervision from a frozen teacher on the student's own rollouts. Vanilla OPD applies this supervision uniformly across prompts, without checking whether the teacher is reliable for each prompt. Because reverse KL is mode-seeking, a confidently wrong teacher can induce a strong yet misleading update. Distributional proxies, such as entropy or teacher-student likelihood agreement, measure uncertainty or agreement but do not directly verify outcome correctness. We introduce Teacher-Gated On-Policy Distillation (TGOPD), built on the principle that teacher reliability should be verified at the prompt level before dense supervision is admitted. TGOPD estimates reliability from a small set of verifier-scored teacher probes and routes each prompt exclusively to dense OPD when the reliability check passes or to verifier-grounded GRPO otherwise. Across 4B and 35B students in mathematics, code, and instruction following, TGOPD outperforms Vanilla OPD in all six single-domain settings and achieves higher seven-benchmark averages at both scales under multi-domain training. By using otherwise-idle teacher capacity for reliability estimation, TGOPD also reduces teacher-side compute waste in asynchronous OPD, increasing teacher-node GPU utilization from 9.8% to 78.9% in the measured 4B single-domain run.
On-policy distillation (OPD) is increasingly used in LLM post-training because it can leverage a teacher model to provide dense supervision on student rollouts. The standard implementation, however, usually reduces distribution matching to a sampled-token log-ratio, which can make the learning signal fragile on long rollouts whose prefixes drift away from the teacher's typical support. We revisit this formulation from both theoretical and implementation perspectives. Theoretically, token-level OPD is biased relative to sequence-level reverse-KL minimization, but admits a substantially tighter worst-case variance bound; a controlled synthetic study further shows that stronger future-reward coupling increases gradient variance and destabilizes training. Empirically, we identify three failure modes of sampled-token OPD: imbalanced token-level supervision, unreliable teacher guidance on student-generated prefixes, and tokenizer or special-token mismatch. These findings motivate teacher top-K local support matching, a truncated reverse-KL objective that compares teacher and student distributions over a teacher-supported token set at each prefix, together with top-p rollout sampling and special-token masking. Across single-task reasoning and multi-task benchmarks spanning agentic and reasoning settings, this objective improves optimization stability and yields a +19.8% performance gain over standard sampled-token OPD baselines, providing a practical recipe for more stable on-policy distillation.
Yuqian Fu, Haohuan Huang, Kaiwen Jiang +4
State Key Laboratory of Multimodal Artificial Intelligence Systems, CASIA · School of Artificial Intelligence, UCAS · Fudan University +1
On-policy distillation (OPD) trains a student model to match a teacher's next-token distributions on student-generated prefixes. Extrapolative OPD further extends the teacher's log-probability shift relative to a reference policy to construct stronger learning targets. However, existing methods couple extrapolation direction and distance, allowing the magnitude of the teacher's shift to dictate how far the target moves. We propose D2-OPD, decoupling these two factors: a per-prefix hard KL constraint bounds the target's distance from the teacher, while an adaptive combination of the teacher's post-training shift and the initial student--teacher gap determines the extrapolation direction. We optimize the student toward this target using forward KL, providing learning signals even for low-probability tokens. We evaluate three settings: where standard OPD performs well, where it struggles with base-model initialization, and where no teacher reference is available. On mathematical reasoning benchmarks, D2-OPD attains the highest score in all three settings, outperforming the strongest baseline by up to 9.45 points. Ablations validate the contributions of each component.