$D^2$-OPD: Separating Direction and Distance in Extrapolative On-Policy Distillation

Bowen Fang; Dingmin Wang; Hengzhi Pei

cs.LGNov 4, 2026

$D^2$ -OPD: Separating Direction and Distance in Extrapolative On-Policy Distillation

Authors: Bowen Fang, Dingmin Wang, Hengzhi Pei

Organizations: Columbia University · AWS AI Lab · Amazon AGI

Abstract

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 $D^2$ -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, $D^2$ -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.

Figures & tables

Figure 1: Qwen3-8B

\rightarrow

Qwen3-4B-Base distillation. Left: average score across AIME24, AIME25, and AMC23 using avg@16 (mean accuracy over 16 sampled responses per problem), with the initial student shown as a dashed reference. D 2 -OPD reaches 47.19, 19.62 points higher than standard OPD. Right: per-benchmark results. The gain is consistent across all three benchmarks. Numbers are from Table 1 .

Figure 2: Distance-aware extrapolation direction. All quantities are shown at a fixed prefix

h

, with conditioning omitted. The distances from the teacher

q^{+}

t_{\mathrm{ref}}

and

\pi_{0}

are

\sigma_{T}

and

\sigma_{\pi}

, with corresponding unit directions

u

and

v

. (a)

\theta_{T}>\theta_{\pi}

: the same orthogonal displacement

\delta

from

q^{+}

\tilde{q}

produces a larger angular deviation at the nearer reference. (b) With

\sigma_{T}<\sigma_{\pi}

, the combined direction

d_{\eta}

lies closer to

u

than to

v

, giving

\alpha_{T}<\alpha_{\pi}

Method	AIME 24	AIME 25	AMC 23	Avg	$\boldsymbol{\Delta}$
Student $\pi_{0}$	12.50	6.25	40.00	19.58	$-7.99$
Teacher $q^{+}$	75.00	64.17	94.22	77.80	$+50.23$
OPD	18.96	14.38	49.38	27.57	—
EOPD	20.00	21.67	54.22	31.96	$+4.39$
ExOPD	17.08	24.79	54.37	32.08	$+4.51$
Direct-OPD	23.75	20.62	57.97	34.11	$+6.54$

Table 1: Setting A : Qwen3-8B

\rightarrow

Qwen3-4B-Base.

Figure 3: Longer responses do not explain the gain. Average score versus mean response length for the trained methods across the three settings. Appendix E.1 includes the full response-length comparison.

Method	AIME 24	AIME 25	AMC 23	Avg	$\boldsymbol{\Delta}$
Student $\pi_{0}$	0.83	1.04	12.66	4.84	$-9.38$
Teacher $q^{+}$	24.17	15.62	64.22	34.67	$+20.45$
OPD	6.25	5.62	30.78	14.22	—
SFT	7.71	4.58	40.00	17.43	$+3.21$
SFT + OPD	14.58	10.00	37.81	20.80	$+6.58$
EOPD	4.79	9.79	38.44	17.67	$+3.45$

Table 2: Setting B : Qwen3-4B

\rightarrow

Qwen3-1.7B-Base.

Method	AIME 24	AIME 25	AMC 23	Avg	$\boldsymbol{\Delta}$
Student $\pi_{0}$	19.76	11.67	50.62	27.35	$-11.56$
Teacher $q^{+}$	73.33	62.29	99.69	78.44	$+39.53$
ExOPD	28.33	25.21	47.81	33.78	$-5.13$
EOPD	33.54	19.79	60.94	38.09	$-0.82$
OPD	30.00	25.00	61.72	38.91	—
TOP-D	29.17	30.42	70.78	43.45	$+4.54$

Table 3: Setting C : Qwen3-30B-A3B-Instruct-2507

\rightarrow

Qwen3-8B-Base. Here

t_{\mathrm{ref}}=\pi_{0}

, so the reference directions coincide.

Figure 4: The gains are not due to lower entropy. We measure full-vocabulary student entropy on each method’s on-policy prefixes. Faint lines show per-step values; solid lines show a 9-step rolling mean.

Figure 5: The

\epsilon

-constraint remains active. Fraction of token positions where the constraint binds in Setting A, with

\epsilon=1

. (a) Training at

\kappa=0.05

: per-step values (faint), a 9-step mean (solid), and the second-half median (dashed). (b) Second-half median versus

\kappa

; the filled marker denotes

\kappa=0.05

Method	Time/step (s)	Throughput (tokens/s)	Relative time
OPD	33.42	1,840	$1.00\times$
Direct-OPD	43.26	1,422	$1.29\times$
TOP-D	32.96	1,866	$0.99\times$
D 2 -OPD	38.48	1,598	$1.15\times$

Table 4: Computational cost.

Component	Variant	Avg	$\boldsymbol{\Delta}$
full		47.19	—
constraint	$-\kappa$	45.59	$-1.60$
	$-\epsilon$	43.30	$-3.89$
	$-$ both	42.73	$-4.46$
direction	$\pi_{0}$ alone	45.12	$-2.07$
direction	$t_{\mathrm{ref}}$ alone	44.27	$-2.92$

Table 5: Ablations in Setting A. Avg is the mean over AIME24, AIME25, and AMC23;

\Delta

is relative to the full method. Per-benchmark scores appear in Appendix D .

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 551 papers

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
O $^2$ PD: Overlap-Gated On-Policy Distillation on Reliable Prefix
Nov 5, 2026Shuaishuai Zu, Jihao Zhao, Pengfei Li +5Language Model Distillation On-Policy Distillation
School of Information, Renmin University of China · Huawei Technologies Ltd.
OPD+: Rethinking the Advantage Design for On-Policy Distillation
May 31, 2026Hanyang Zhao, Haoxian Chen, Han Lin +3Language Model Distillation Policy Gradient Methods
Columbia University · Amazon · Meta +1

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

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

Zhenyu Wang, Tianze Wang, Linjun Zhang, Yifan Hu

Department of Statistics, Rutgers University

Language Model Distillation Bilevel Optimization On-Policy Distillation

Open paper AskPDF arXiv

On-policy distillation (OPD) trains a student on its own generated responses using dense, token-level supervision from a stronger teacher. Vanilla OPD treats all teacher signals equally, assuming that the teacher's supervision is equally important for every token. However, teacher signals at different tokens may have very different effects on the student's performance: some correct important reasoning errors, while others have little effect on the final answer. Motivated by this observation, we introduce Dr. OPD (OPD Done Right), which defines the optimal weighted OPD to maximize the student's performance. We formulate Dr. OPD as a bilevel optimization problem in which the student learns from weighted teacher supervision, while the weights are selected to maximize the expected reward of the resulting student. To solve Dr. OPD, we develop an efficient iterative solver that updates the token weights and student policy alternatively. At each round, it updates weights in closed form and then takes one gradient step on the resulting weighted OPD objective. Under regularity conditions, we show that this weighted update achieves a higher expected reward than a vanilla OPD update. Empirically, across strong-to-weak and same-size distillation on math and code, Dr. OPD consistently outperforms all evaluated baselines. In particular, in the strong-to-weak distillation setting, Dr. OPD improves average math performance by

9.7

points over vanilla OPD, and enables the smaller student to surpass its larger teacher.

Abstract

Language Model Distillation

Explore similar work

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

O2^22PD: Overlap-Gated On-Policy Distillation on Reliable Prefix

OPD+: Rethinking the Advantage Design for On-Policy Distillation

O $^2$ PD: Overlap-Gated On-Policy Distillation on Reliable Prefix