On-Policy Transition Distillation

Runyang You; Yongqi Li; Tiezheng Yu; Wenjie Li

cs.CLOct 20, 2026

On-Policy Transition Distillation

Authors: Runyang You, Yongqi Li, Tiezheng Yu, Wenjie Li

Organizations: The Hong Kong Polytechnic University · Huawei Technologies Ltd.

Abstract

On-policy distillation (OPD) consolidates expert capabilities into a single student by matching teacher and student final-layer distributions on student-generated rollouts. We investigate whether depthwise prediction transitions, the changes in predictions between layers, can provide additional task-relevant supervision. Our diagnostic study finds that stronger teacher--student agreement in intermediate transitions is associated with higher answer correctness within questions, with correct responses exhibiting lower mean transition KL across all measured intermediate layer pairs. These observations motivate On-Policy Transition Distillation (OPTD), which augments final-layer distribution matching by aligning teacher and student layer-to-layer log-probability increments on student-generated rollouts. Evaluations across three backbones and eight benchmarks spanning mathematics, coding, and instruction following show that OPTD achieves the best or tied-best overall performance in both dual- and three-domain distillation. In the three-domain setting, OPTD recovers $\approx 90\%$ of the overall teacher--student capability gap, compared with $\approx 75\%$ for OPD.

Figures & tables

Figure 1: Teacher agreement and answer correctness. (a) Correctness by within-question agreement quartile; shading shows 95% confidence intervals. (b) Relative transition-KL gaps across layer pairs; positive values indicate closer teacher agreement on correct responses. Diagnostic details are in App. C.6 .

	Math							Coding
Method	AIME24		AIME25		AIMO		Avg.	LiveCodeBench		CodeContests		CodeElo			Avg.	Overall
	Avg	Pass	Avg	Pass	Avg	Pass	Avg.	Avg	Pass	Avg	Pass	Avg	Pass	Elo	Avg.
DeepSeek-R1-Distill-Qwen-1.5B
OPD	47.29	80.00	30.00	53.33	74.40	91.57	50.56	22.98	37.63	9.58	25.45	8.55	19.36	843.26	13.70	32.13
LastOPD	43.75	76.67	31.46	56.67	75.90	93.98	50.37	22.40	38.35	10.57	26.06	9.87	21.57	916.11	14.28	32.33
OPRD	36.88	73.33	27.50	53.33	71.46	91.57	45.28	20.03	34.77	9.58	22.42	8.55	19.36	843.26	12.72	29.00

Table 1: Dual-domain distillation results on mathematics and coding benchmarks across model backbones. All scores are percentages except CodeElo reports the estimated rating. Highest and second-highest values in each column are bold and underlined, respectively.

	Math					Coding						Instruction Following
Method	AIME24		AIME25		Avg.	LiveCodeBench		CodeContests		CodeElo	Avg.	IFEval		IFBench		Avg.	Overall
	Avg	Pass	Avg	Pass	Avg.	Avg	Pass	Avg	Pass	Elo	Avg.	prompt	instr.	prompt	instr.	Avg.
Student	30.63	70.00	21.04	33.33	25.84	14.43	27.96	5.15	21.82	618.23	9.79	41.77	53.59	13.84	15.70	31.23	29.11
Teachers	48.96	76.67	34.79	56.67	41.88	21.86	36.56	10.11	23.03	943	15.99	69.67	77.90	16.00	17.74	45.33	40.83
OPD	43.54	73.33	30.63	50.00	37.09	22.45	36.56	10.57	26.06	946.11	16.51	61.54	69.74	13.84	16.42	40.39	37.89
recovery rate	70.4%	49.9%	69.7%	71.4%	70.1%	107.9%	100.0%	109.3%	350.0%	100.9%	108.4%	70.9%	66.4%	0.0%	35.3%	65.0%	74.9%

Table 2: Three-domain distillation results across mathematics, coding, and instruction following on the 1.5B backbone. Teachers denotes domain-specialized teachers evaluated on their target domains.

Panel of Table 3 — Table 3: Depth placement comparison at a fixed normalized span of $s\approx 0.214$ . Path glyphs visualize placement along the normalized depth axis $r=\ell/28$ .

Panel of Table 4 — Table 4: Span comparison at Middle and Late placements. Path glyphs visualize placement along the normalized depth axis $r=\ell/28$ .

Formulation	Weight	Centered	Math	Coding	Avg.
ReverseKL	$q^{S}$	✓	50.78	21.15	35.97
OPTD	$w_{i,t}$	×	51.01	21.68	36.35
Increment weights + raw reward	$q^{S}$	×	50.48	21.77	36.13
Predictive weights + centered reward	$w_{i,t}$	✓	43.13	0.36	21.75

Table 5: Transition update comparison on the 1.5B dual-domain setting. Weights use increment softmax

q^{S}

or predictive probabilities

w_{i,t}

; centering subtracts the

q^{S}

-weighted mean reward. Math is Math Avg@16 and Coding is LiveCodeBench Avg@4. Avg. is their equal-weight mean; bold marks the best score.

Figure 2: Teacher alignment reductions against OPD along depth. Qwen3-4B students evaluated on a shared set of saved rollout texts. Lines show reductions in layer-wise prediction reverse KL; bars show reductions in normalized-increment teacher-to-student KL between adjacent intermediate readouts. Horizontal axis is normalized depth $r=\ell/36$ ; error bars denote paired question-bootstrap 95% intervals.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Mathematics
		AIME24		AIME25		AIMO
Size	Role	Avg	Pass	Avg	Pass	Avg	Pass
1.5B	Initial student	30.63	70.00	21.04	33.33	63.33	89.16
	Math teacher	48.96	76.67	34.79	56.67	80.05	93.98
4B	Initial student	22.50	53.33	19.38	46.67	60.92	84.34
	Math teacher	53.54	80.00	41.04	63.33	79.82	95.18

Appendix

Table 6: Reference performance of initial students and domain teachers. Accuracy scores are percentages; CodeElo is an estimated rating. Dashes denote unavailable reference scores.

	IFEval				IFBench
	strict prompt	strict instr.	loose prompt	loose instr.	strict prompt	strict instr.	loose prompt	loose instr.
Initial student	39.91	51.72	43.63	55.46	12.67	13.66	15.00	17.73
IF teacher	67.51	76.16	71.82	79.64	14.67	15.70	17.33	19.77

Appendix

Table 7: Instruction-following reference accuracy (%) for the 1.5B initial student and its domain teacher. Prompt accuracy requires all constraints in a prompt to pass; instruction accuracy scores individual constraints under strict or loose checking.

Method	AIME24		AIME25		AIMO
	Avg	Pass	Avg	Pass	Avg	Pass
Student	30.63	70.00	21.04	33.33	63.33	89.16
Teachers	48.96	76.67	34.79	56.67	80.05	93.98
OPD	43.54	73.33	30.63	50.00	75.08	93.98
recovery	70.4%	49.9%	69.7%	71.4%	70.3%	100.0%
OPTD	46.67	80.00	32.50	53.33	74.32	91.57

Appendix

Table 8: Mathematics results on the 1.5B student. Recovery is the fraction of the student–teacher gap closed by that column’s score.

Method	LiveCodeBench		CodeContests		CodeElo
	Avg	Pass	Avg	Pass	Avg	Pass	Elo
Student	14.43	27.96	5.15	21.82	5.45	14.71	618.23
Teachers	21.86	36.56	10.11	23.03	9.38	22.06	943
OPD	22.45	36.56	10.57	26.06	9.99	22.79	946.11
recovery	107.9%	100.0%	109.3%	350.0%	115.5%	109.9%	100.9%
OPTD	22.13	36.92	10.72	25.45	9.59	20.83	958.12

Appendix

Table 9: Coding results on the 1.5B student.

Method	IFEval				IFBench
	strict prompt	strict instr.	loose prompt	loose instr.	strict prompt	strict instr.	loose prompt	loose instr.
Student	39.91	51.72	43.63	55.46	12.67	13.66	15.00	17.73
Teachers	67.51	76.16	71.82	79.64	14.67	15.70	17.33	19.77
OPD	60.41	68.61	62.67	70.86	12.00	14.24	15.67	18.60
recovery	74.3%	69.1%	67.5%	63.7%	$-$ 33.5%	28.4%	28.8%	42.6%
OPTD	60.50	71.04	63.07	73.27	13.67	14.83	17.67	20.06

Appendix

Table 10: Instruction-following results on the 1.5B student.

Size	Role	Model
1.5B	Student	DeepSeek-R1-Distill-Qwen-1.5B $\nearrow$
	Math teacher	JustRL-DeepSeek-1.5B $\nearrow$
	Coding teacher	DeepCoder-1.5B-Preview $\nearrow$
	IF teacher	Nemotron-Research-Reasoning-Qwen-1.5B (ProRL v2) $\nearrow$
4B	Student	Qwen3-4B $\nearrow$
	Math teacher	Qwen3-4B-Math-RL-Specialist $\nearrow$

Appendix

Table 11: Students and domain teachers grouped by model size. The IF teacher is added only in the multi-domain extension. Link icons point to Hugging Face repositories.

Figure 3: Lower final-output KL for incorrect responses appears early and persists across training. Each curve shows the percentage difference between the incorrect- and correct-response means, computed as

\Delta\mathcal{D}_{\mathrm{final}}

in Eq. ( 9 ). Negative values indicate closer final-output agreement for incorrect responses. Questions with mixed outcomes are selected separately at each training stage.

Update branch	Maximum gradient norm	First update above 100
Increment weights + raw reward ( $q^{S}$ )	16.06	never
Predictive weights + centered reward ( $w_{i,t}$ )	4576.0	34
ReverseKL ( $q^{S}$ , centered)	3.10	never

Appendix

Table 12: Gradient-scale audit over 150 updates for the three transition-update branches shown. Among these branches, only the centered predictive branch has a recorded gradient norm above 100, reaching 4576 before clipping.

Language Model Distillation

Language model distillation transfers a larger model’s capabilities to smaller or specialized models using outputs, probabilities, or internal representations.

See all 508 papers

Verify Before You Distill: Prompt-Level Teacher Gating for On-Policy Distillation
Sep 2, 2026Zhiwei Zhang, Zechen Sun, Fei Zhao +6On-Policy Distillation RL Post-Training
AllSpark Team
OPRD: On-Policy Representation Distillation
Jun 4, 2026Shenzhi Yang, Guangcheng Zhu, Bowen Song +8Cross-Architecture Knowledge Distillation Language Model Distillation
Zhejiang University · Ant Group
Dr. OPD: Learning What to Follow for Optimal On-Policy Distillation of Large Language Models
Sep 29, 2026Zhenyu Wang, Tianze Wang, Linjun Zhang +1Language Model Distillation Bilevel Optimization
Department of Statistics, Rutgers University

Sep 2, 2026J/K move · Enter open · S save

Verify Before You Distill: Prompt-Level Teacher Gating for On-Policy Distillation

Zhiwei Zhang, Zechen Sun, Fei Zhao, Kang Peng+5

AllSpark Team

On-Policy Distillation RL Post-Training

Open paper AskPDF arXiv

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.

Abstract

Language Model Distillation

Explore similar work

Verify Before You Distill: Prompt-Level Teacher Gating for On-Policy Distillation

OPRD: On-Policy Representation Distillation

Dr. OPD: Learning What to Follow for Optimal On-Policy Distillation of Large Language Models