Training-Free On-Policy Distillation

Bangji Yang; Jiajun Fan; Zhaoyu Zhu; Hongbo Ma; Ge Liu

cs.LGOct 15, 2026

Training-Free On-Policy Distillation

Authors: Bangji Yang, Jiajun Fan, Zhaoyu Zhu, Hongbo Ma, Ge Liu

Organizations: University of Illinois at Urbana-Champaign · Tsinghua University

Abstract

On-policy distillation consolidates specialist models into one student at the price of rollouts, teacher passes and gradient steps for every teacher. We show that the consolidation does not need training: with forward passes only, the capabilities of separately tuned specialists can be aggregated into one model that is, on average, more accurate than either specialist alone. When specialists share a base model, distillation has an exact layer-wise solution: each linear layer of the fused model must reproduce every specialist's output on the inputs that specialist receives while reading its own answers, a least-squares problem whose solution needs only the second moments of those inputs. One forward pass over a few hundred rollouts provides them. We call this training-free on-policy distillation. What the rollouts contain matters more than the solver: answers truncated in repetition loops are few among a mathematics specialist's rollouts but carry a large share of their tokens, and removing them lifts the fused model on AIME25 from task arithmetic's level to the specialist's at matched precision. On Olmo-3-7B, a few GPU-hours of forward computation per specialist yield a generalist that keeps most of each specialist's gain on each of seven mathematics and code benchmarks.

Figures & tables

Figure 1: A generalist with the highest mean accuracy, without training (Olmo 3 7B). (a) Mean accuracy on the five mathematics benchmarks against mean accuracy on the two code benchmarks; dashed lines mark what each specialist reaches on its own domain, and TF-OPD, computed from the two with forward passes only, lies next to the corner they define. (b) Every benchmark: the base model and the specialists on the upper line of each row, TF-OPD below, and at the right the score of that domain’s specialist and of TF-OPD. (c) Mean accuracy over the seven benchmarks: TF-OPD is more accurate on average than either specialist it was built from.

Route	No training	On-policy inputs	Matches teachers
Multi-domain RL	$\times$	$\checkmark$	$\times$
On-policy distillation	$\times$	$\checkmark$	$\checkmark$
Task arithmetic, TIES, DARE	$\checkmark$	$\times$	$\times$
RegMean	$\checkmark$	$\times$	$\checkmark$
[][]TF-OPD (ours)	$\checkmark$	$\checkmark$	$\checkmark$

Table 1: Routes from several specialists to one model. TF-OPD needs no training, gathers its statistics on text the models generate, and is fit to match the specialists’ outputs.

Figure 2: Overview of TF-OPD. Two specialists trained from one base model each answer training problems of their own domain (1); answers that end in a repetition loop or fail grading are dropped (2); each specialist reads its kept answers once while hooks record the second moments

G_{d}

of every layer’s inputs (3); and each layer of the fused model is obtained by one linear solve (4). Steps 1 to 3 are forward passes and step 4 is linear algebra, so no gradient is computed at any point. Step 2 decides the hard problems (Section 3.4 ).

	Mathematics					Code
Model	MATH-500	OlympiadBench	AIME24	AIME25	HMMT25	HumanEval+	MBPP+	Mean
Base model	63.4	33.4	15.8	10.0	8.3	45.7	44.2	31.6
Math specialist	91.8	68.7	50.0	43.3	30.8	71.3	62.4	59.8
Code specialist	85.6	62.5	26.7	30.0	15.0	86.0	69.3	53.6
[][]TF-OPD (ours)	90.8	69.4	50.0	38.3	27.5	84.1	67.5	61.1

Table 2: One generalist, ahead of both specialists on average (Olmo 3 7B). Accuracy (%). AIME24, AIME25 and HMMT25 average four sampled answers per problem; MATH-500, OlympiadBench, HumanEval+ and MBPP+ use one. Released checkpoints are evaluated in bfloat16 and the fused model in float16. The mean is taken over the seven benchmarks before rounding. The best value of each column is in bold.

Model	MATH-500	HumanEval+	MBPP+
No statistics
Task arithmetic	92.0	81.7	69.3
Low-rank sum	91.6	82.9	68.8
TF-OPD, specialist rollouts
128 problems, 16k tokens	92.0	84.8	68.5
256 problems, 32k tokens	93.2	87.2	70.1

Table 3: On the easier benchmarks every rule looks alike. Accuracy (%), one sampled answer per problem. A row marked + adds one change to the row above it.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Step	Unit	Cost
Rollouts	64 problems	1 GPU-hour
Statistics	2.6M tokens	40 min
Solve	all layers	25 min
Total	two specialists	$\approx$ 10 GPU-hours
AIME25 evaluation	one model	2.5–3 GPU-hours

Appendix

Table 4: Cost of TF-OPD on Olmo 3 7B, on 48 GB GPUs. Every step is forward computation or linear algebra, and none computes a gradient. The total is derived from the rows above for two specialists with 256 training problems each; evaluating one model on AIME25 is given for scale.

Table 8

	KL per	Top-1	Residual
Model	token	(%)	(%)
TF-OPD, all answers	0.0024	98.6	0.9 $\to$ 3.6
Task arithmetic	0.0247	95.3	1.2 $\to$ 10.4
Code specialist	0.0595	92.8	1.7 $\to$ 18.6

Appendix

Table 8: Teacher-forced agreement. Agreement with the mathematics specialist under teacher forcing along its own correct solutions to AIME25 and MATH-500 problems: KL divergence per token, top-1 agreement, and relative deviation of the residual stream at the first and at the last layer.

Matrix	Share (%)
Attn. query	13.0
Attn. key	15.0
Attn. value	8.5
Attn. output	5.0
MLP gate	11.5
MLP up	10.7

Appendix

Table 9: Unreproduced share of a teacher’s update (Equation ( 4 )). Median over layers, in percent. Left: by matrix type, for the all-answers statistics. Right: by teacher, for three sets of statistics.

Figure 6: Problems that TF-OPD solves and the specialist of the domain does not. One square per sampled answer: four per model on HMMT25, and one for TF-OPD against four seeds of the coding specialist on the code benchmarks. Numbers are the final answers given on HMMT25 problem 21, whose reference answer is 63.

On-Policy Distillation

On-policy distillation trains a model on its own outputs using teacher guidance, aligning supervision with the states it encounters at deployment.

See all 507 papers

OPRD: On-Policy Representation Distillation
Jun 4, 2026Shenzhi Yang, Guangcheng Zhu, Bowen Song +8Cross-Architecture Knowledge Distillation Language Model Distillation
Zhejiang University · Ant Group
On-Policy Transition Distillation
Oct 20, 2026Runyang You, Yongqi Li, Tiezheng Yu +1Language Model Distillation On-Policy Distillation
The Hong Kong Polytechnic University · Huawei Technologies Ltd.
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

Jun 4, 2026J/K move · Enter open · S save

OPRD: On-Policy Representation Distillation

Shenzhi Yang, Guangcheng Zhu, Bowen Song, Haobo Wang+7

Zhejiang University · Ant Group

Cross-Architecture Knowledge Distillation Language Model Distillation On-Policy Distillation Teacher-Student Learning Knowledge Distillation

Open paper AskPDF arXiv

On-policy distillation (OPD) supervises the student exclusively in the output space by matching next-token distributions. This paradigm suffers from two limitations: (i) a high-variance gradient estimator whose signal-to-noise ratio collapses as the student approaches the teacher, and (ii) an LM-head information bottleneck that discards the teacher's intermediate hidden states. We propose On-Policy Representation Distillation (OPRD), the first method to lift on-policy distillation into the hidden-state space. OPRD aligns student and teacher representations across selected layers on the same on-policy rollouts, providing dense, deterministic, per-layer supervision while bypassing the LM head entirely. Theoretically, OPRD provides a deterministic per-sample gradient, removing the token-level estimation variance that plagues OPD, and exposes structural information that any output-space objective necessarily discards. Empirically, OPRD closes the student-teacher gap on competition mathematics benchmarks (AIME 2024, AIME 2025, and AIMO), where every output-space baseline plateaus below the teacher, while training 1.44x faster and using up to 54% less memory. We further extend OPRD to the cross-architecture setting via OPRD-Bridge. By exploiting the observation that heterogeneous models share a low-rank representational structure, we construct a frozen projector pair that aligns representations across arbitrary depth and width mismatches, shifting the alignment from the output space (which depends on a shared vocabulary) to the representation space. We validate OPRD-Bridge on both cross-architecture (Qwen3-4B -> Qwen3-1.7B-Base) and cross-tokenizer (Phi-4-mini-reasoning -> Qwen3-1.7B-Base) settings, demonstrating successful knowledge transfer even when the vocabulary-based alignment channel is unavailable. Code: https://github.com/ShenzhiYang2000/OPRD.

Abstract

On-Policy Distillation

Explore similar work

OPRD: On-Policy Representation Distillation

On-Policy Transition Distillation

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