JustOPD: On-Policy Distillation from Autoregressive to Diffusion Language Models

Zikun Zhang; Jiayuan Sheng; David D. Yao; Wenpin Tang

cs.LGOct 12, 2026

JustOPD: On-Policy Distillation from Autoregressive to Diffusion Language Models

Authors: Zikun Zhang, Jiayuan Sheng, David D. Yao, Wenpin Tang

Organizations: Department of Industrial Engineering and Operations Research, Columbia University, New York, NY 10027.

Abstract

On-policy distillation (OPD) is a post-training approach for large language models (LLMs) that enables a student model to learn a teacher model's distribution from sequences sampled from the student's own policy. Despite its promising capability, OPD for diffusion language models (DLMs) remains underexplored and existing studies primarily focuses on DLM-to-DLM distillation. The limited availability of strong teacher DLMs and the flexible generation order of DLMs impose additional technical challenges as well. However, it remains unclear how to effectively distill knowledge from a strong autoregressive language model (ARM) into a weaker DLM, given the mismatch between the left-to-right decoding of ARMs and the flexible decoding order of DLMs. Inspired by Ni et al. (2026), we propose JustOPD, an ARM-to-DLM on-policy distillation method that uses autoregressive rollouts during training to align the student's generation process with the ARM teacher's left-to-right decoding. This alignment enables effective distillation from ARM teachers while preserving the DLM's flexible generation-order capability at inference time. Experiments using Qwen3-14B as the teacher and Dream-7B as the student demonstrate that JustOPD consistently improves mathematical reasoning performance over the base DLM across multiple divergence choices.

Figures & tables

Divergence	$f(u)$	$\tilde{f}(u)$	$A_{f}(u)$
Forward KL	$u\log u$	$u\log u-u+1$	$u$
Reverse KL	$-\log u$	$-\log u+u-1$	$\log u-1$
$\alpha$ -divergence, $\alpha\in(0,1)$	$\dfrac{u^{\alpha}-\alpha u-(1-\alpha)}{\alpha(\alpha-1)}$	unchanged	$\dfrac{u^{\alpha}-1}{\alpha}$
Jensen-Shannon	$\frac{1}{2}\big(u\log u-(1+u)\log\frac{1+u}{2}\big)$	unchanged	$\frac{1}{2}\log\frac{1+u}{2}$

Table 1: Typical

f

-divergences and their centralized versions

\tilde{f}

and per-token advantages

A_{f}

\mathcal{L}(\theta)=\frac{1}{|{\cal B}_{\text{mini}}|}\sum_{({\mathbf{q}},{\mathbf{x}})\in{\cal B}_{\text{mini}}}\frac{1}{|{\cal H}^{\mathtt{te}}|}\sum_{t\in{\cal H}^{\mathtt{te}}}\min\left(\rho_{t}^{\theta}\cdot A_{t},\ \mathrm{clip}(\rho_{t}^{\theta},1-\epsilon,1+\epsilon)\cdot A_{t}\right)

Algorithm 1 JustOPD for DLMs: Sampled-Token Distillation

Figure 1: The

\alpha

-divergence advantage interpolates between the logarithmic growth of reverse KL and the linear growth of forward KL. The Jensen-Shannon advantage also grows logarithmically.

	GSM8K		MATH500
Accuracy	Dream-7B	Qwen3-14B	Dream-7B	Qwen3-14B
$L=256$	77.03	51.71	31.80	14.40
$L=512$	78.77	93.93	48.20	48.20
$L=1024$	56.25	94.84	36.80	70.80

Table 2: Performance of Dream-7B and Qwen3-14B with different generation lengths

L

	GSM8K		MATH500
Accuracy	Dream-7B	Qwen3-14B-FT	Dream-7B	Qwen3-14B-FT
$L=256$	77.03	93.48	31.80	53.20

Table 3: Performance of Dream-7B and fine-tuned Qwen3-14B with a short sequence length of

L=256

Acc. $(L=256)$	Dream-7B	$\alpha=0$	$\alpha=0.2$	$\alpha=0.5$	$\alpha=0.8$	$\alpha=1$	JSD
GSM8K	77.03	80.97	81.80	80.36	80.67	79.68	80.52
MATH500	31.80	38.80	36.60	38.40	39.60	38.40	39.40

Table 4: Performance of Dream-7B before and after OPD with Qwen3-14B-FT on GSM8K and MATH500 on the short sequence length of

L=256

across different

\alpha

-divergences and Jensen-Shannon divergence (JSD).

\alpha=0

is the reverse KL and

\alpha=1

is the forward KL.

Acc.	Dream-7B	$\alpha=0$	$\alpha=0.5$	$\alpha=1$	JSD
GSM8K ( $L=512$ )	78.77	84.45	83.24	79.23	81.96
MATH500 ( $L=1024$ )	36.80	43.80	46.20	42.80	48.00

Table 5: Performance of Dream-7B before and after OPD on GSM8K and MATH500 with Qwen3-14B across different

\alpha

-divergences and JSD. The generation length

L

is specified for each benchmark.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

\mathcal{L}(\theta)=\frac{1}{|{\cal B}_{\text{mini}}|}\sum_{({\mathbf{q}},{\mathbf{x}})\in{\cal B}_{\text{mini}}}\frac{1}{|{\cal H}^{\mathtt{te}}|}\sum_{t\in{\cal H}^{\mathtt{te}}}D_{f}^{\dagger}\big(\pi_{\mathtt{te}}(\cdot\mid{\mathbf{q}},{\mathbf{x}}^{<t})\,\|\,\pi_{\theta}(\cdot\mid{\mathbf{q}},{\mathbf{x}}^{<t})\big)

Appendix

Algorithm 2 JustOPD for DLMs: Full/TopK-Vocabulary Distillation

\mathcal{L}(\theta)=-\frac{1}{|{\cal B}_{\text{mini}}|}\sum_{({\mathbf{q}},{\mathbf{x}})\in{\cal B}_{\text{mini}}}\frac{1}{|{\cal T}|}\sum_{t\in{\cal T}}\frac{1}{|Q_{t}|}\sum_{y^{t}\in Q_{t}}\log\pi_{\theta}(y^{t}\mid{\mathbf{q}},{\mathbf{x}}^{<t})

Appendix

Algorithm 3 JustOPD for DLMs: Black-Box Teacher Distillation

	$f$	$\tilde{f}$
Signal $D_{f,t}$	$\Theta(\gamma^{2})$	$\Theta(\gamma^{2})$
Single-sample variance $\sigma_{t}^{2}$	$\Theta(\gamma^{2})$	$O(\gamma^{4})$
Single-sample standard deviation $\sigma_{t}$	$\Theta(\gamma)$	$O(\gamma^{2})$
Noise-to-signal ratio $\sigma_{t}/D_{f,t}$	$\Theta(1/\gamma)$	$O(1)$

Appendix

Table 6: Comparison of signal and noise scaling for

f

and

\tilde{f}

\mathcal{L}(\phi)=-\frac{1}{|\mathcal{R}|L}\sum_{(\tilde{{\mathbf{q}}},{\mathbf{x}}_{i},A_{i})\in\mathcal{R}}A_{i}\cdot\sum_{t=1}^{|{\mathbf{x}}_{i}|}\log\pi_{\phi}(x_{i}^{t}\mid\tilde{{\mathbf{q}}},{\mathbf{x}}_{i}^{<t}).

Appendix

Algorithm 4 Length-Aware Fine-Tuning of Qwen3-14B

	GSM8K		MATH500		AIME 24
Gen. Length	Acc.		Acc.		Pass@16		Avg@16
	Stu.	Tr.	Stu.	Tr.	Stu.	Tr.	Stu.	Tr.
$L=256$	77.03	51.71	31.80	14.40	0.00	0.00	0.00	0.00
$L=512$	78.77	93.93	48.20	48.20	0.00	3.33	0.00	0.83
$L=1024$	56.25	94.84	36.80	70.80	0.00	13.33	0.00	4.58

Appendix

Table 7: Performance of Dream-7B and Qwen3-14B with different generation lengths

L

. Stu. and Tr. denote the base Dream-7B (student) and Qwen3-14B (teacher) models, respectively.

	GSM8K	MATH500	AIME 24
Model	Acc. ( $L=512$ )	Acc. ( $L=1024$ )	Pass@16 ( $L=1024$ )	Avg@16 ( $L=1024$ )
Dream-7B	78.77	36.80	0.00	0.00
+JustOPD	84.45	43.80	3.33	0.21

Appendix

Table 8: Performance of Dream-7B before and after OPD on GSM8K, MATH500, and AIME 24. The generation length

L

is specified for each benchmark.

Diffusion Language Model Distillation

Diffusion language model distillation transfers knowledge or decoding behavior to make diffusion-based text generation faster and more capable.

See all 37 papers

dOPSD: On-Policy Self-Distillation for Diffusion Language Models
Jul 5, 2026Phuong Tuan Dat, Qi Li, Xinchao WangOn-Policy Self-Distillation Language Model Distillation
National University of Singapore
d-OPD: Future-Aware On-Policy Distillation for Block Diffusion Language Models
Sep 28, 2026Ruitao Liu, Qinghao Hu, Song HanDiffusion Model Distillation Blockwise Diffusion Language Models
Tsinghua University · MIT · NVIDIA
Data-Efficient Autoregressive-to-Diffusion Language Models via On-Policy Distillation
Jun 4, 2026Xingyu Su, Jacob Helwig, Shubham Parashar +6Language Model Distillation Diffusion Language Models
Department of Computer Science and Engineering, Texas A&M University · Department of Bioinformatics and Systems Medicine,2026 University of Texas Health Science Center at Houston · Department of Electrical and Computer Engineering, Texas A&M UniversityJun 4

Jul 5, 2026J/K move · Enter open · S save

dOPSD: On-Policy Self-Distillation for Diffusion Language Models

Phuong Tuan Dat, Qi Li, Xinchao Wang

National University of Singapore

On-Policy Self-Distillation Language Model Distillation Diffusion Language Model Decoding Diffusion Language Models LLM Reasoning Diffusion Language Model Distillation

Open paper AskPDF arXiv

Diffusion large language models (dLLMs) generate text by iteratively denoising a masked sequence, offering a parallel alternative to autoregressive models, but eliciting strong reasoning through post-training remains difficult: supervised fine-tuning is off-policy and suffers from exposure bias, while reinforcement learning gives only sparse, sequence-level rewards and is hard to apply without tractable sequence likelihoods. On-policy self-distillation (OPSD) offers a promising alternative, using one model as both student and teacher to provide dense, token-level, on-policy supervision, but its effectiveness hinges on giving the teacher privileged information (PI) - typically an instance-specific ground-truth reference unavailable at inference - so the student ends up distilling a weak PI-free consensus policy that yields little improvement on dLLM reasoning. We introduce dOPSD, which instead derives the teacher's privilege directly from the student's own denoising trajectory, evaluating masked positions using later, more-decoded steps of that same trajectory rather than an external label, so the teacher's advantage emerges from the model's own decoding process; on Dream and LLaDA, dOPSD improves both in-domain math reasoning and out-of-domain code generation, outperforming supervised and on-policy baselines.

Abstract

Diffusion Language Model Distillation

Explore similar work

dOPSD: On-Policy Self-Distillation for Diffusion Language Models

d-OPD: Future-Aware On-Policy Distillation for Block Diffusion Language Models

Data-Efficient Autoregressive-to-Diffusion Language Models via On-Policy Distillation