💬 NLP

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

This paper introduces OPDLM, a data-efficient framework that transforms autoregressive language models into diffusion language models via on-policy distillation, effectively eliminating train-inference mismatch and preserving prior knowledge while reducing training token requirements by up to 7,000x.

Original authors: Xingyu Su, Jacob Helwig, Shubham Parashar, Atharv Chagi, Lakshmi Jotsna, Degui Zhi, James Caverlee, Dileep Kalathil, Shuiwang Ji

Published 2026-06-08

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Xingyu Su, Jacob Helwig, Shubham Parashar, Atharv Chagi, Lakshmi Jotsna, Degui Zhi, James Caverlee, Dileep Kalathil, Shuiwang Ji

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a brilliant, world-class chef (the Autoregressive Model, or ARLM) who has spent years perfecting their craft by cooking one dish at a time, step-by-step, always looking at what they just put in the pot to decide the next ingredient. This chef is incredibly fast and knows a lot, but they only know how to cook in a straight line.

Now, imagine you want this chef to learn a new, revolutionary style of cooking called Diffusion (the Diffusion Language Model, or DLM). In this new style, instead of cooking step-by-step, the chef starts with a bowl of random, unrecognizable ingredients (a "masked" sequence) and gradually refines them into a perfect meal all at once. This new style is great because it can guess several ingredients at the same time, potentially cooking much faster.

The Problem: The "Translation" Gap
The paper explains that previous attempts to turn the step-by-step chef into the "all-at-once" chef had two big problems:

The Memory Loss: When you force the chef to stop cooking step-by-step and start guessing the whole dish at once, they tend to forget the thousands of recipes they learned during their years of training. It's like telling a master pianist to suddenly play with their eyes closed and in a different key; they might lose their touch.
The Practice vs. Performance Mismatch: In the old method, the chef practiced by randomly throwing ingredients into a bowl and trying to fix them (random masking). But in the real world (inference), the chef actually refines the dish based on how confident they are in their guesses. The practice didn't match the performance, so the chef was always a bit rusty when it counted.

The Solution: OPDLM (The "On-Policy" Mentor)
The authors introduce a new method called OPDLM (On-Policy Diffusion Language Model). Think of this as a smart training camp that solves both problems using a technique called On-Policy Distillation.

Here is how it works, using a simple analogy:

The Student and the Teacher: The "Student" is the new Diffusion chef. The "Teacher" is the original, frozen step-by-step chef (who is kept in a glass case, unchanged).
The "On-Policy" Twist: Instead of practicing on random, messy bowls of ingredients (the old way), the Student chef is allowed to cook a full meal using their own new, "all-at-once" style. They generate their own path from a messy bowl to a finished dish.
The Distillation: Once the Student finishes a meal, the frozen Teacher chef looks at that exact same meal and says, "If I had cooked this specific dish, here is exactly what I would have said for every missing ingredient."
The Learning: The Student compares their guess with the Teacher's answer and learns from it.

Why This is a Game-Changer
Because the Student practices on meals they actually generate (matching the real performance) and gets corrected by the original expert, they don't forget their old knowledge, and they learn much faster.

The paper claims this method is incredibly efficient:

Data Savings: It requires 15 to 7,000 times fewer training examples than previous methods. If other methods needed to read a library of a billion books to learn, this method might only need a few dozen.
No "Pre-training" Cost: You don't need to train a new Diffusion chef from scratch (which is expensive and slow). You just take an existing expert, give them this specific training, and they transform.
Keeping the Magic: The new model keeps the "thinking" and "multilingual" skills of the original chef, even though it wasn't explicitly taught those things during the transformation. It's like the chef naturally remembers how to speak French or solve complex riddles because the training preserved their original brain.

The Result
The paper shows that this new "OPDLM" chef performs just as well as the best existing chefs on math, coding, and general knowledge tasks, but they got there with a tiny fraction of the effort and data. It turns the difficult process of changing a model's "brain" into a simple, efficient post-training upgrade.

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

1. Problem Statement

The paper addresses the challenges inherent in transforming pretrained Autoregressive Language Models (ARLMs) into Diffusion Language Models (DLMs). While recent work has attempted to convert ARLMs to DLMs by replacing causal attention with bidirectional attention and retraining with a diffusion objective, the authors identify two critical "distribution shifts" that limit training efficiency and performance:

Knowledge-Retention Mismatch: Transitioning from a next-token prediction objective (ARLM) to a diffusion objective (DLM) can weaken or discard the substantial knowledge acquired during the original ARLM pretraining.
Train-Inference State Mismatch: Standard DLMs are trained on sequences corrupted by uniformly random masking (forward process). However, during inference, the model generates sequences via reverse unmasking trajectories guided by confidence-based decoding heuristics. The states encountered during inference differ significantly from the random states seen during training, leading to suboptimal learning.

Existing conversion methods often treat this as a continuation of pretraining, failing to explicitly address these mismatches or requiring prohibitively large datasets (trillions of tokens) for pretraining DLMs from scratch.

2. Methodology: On-Policy Diffusion Language Models (OPDLM)

The authors propose OPDLM, a framework that converts ARLMs to DLMs using On-Policy Distillation (OPD). This approach reframes the conversion as a post-training problem rather than a pretraining one.

Core Mechanism

OPDLM employs a "self-OPD" strategy where the student model generates its own training data, and the original ARLM acts as a frozen teacher.

Student Model: An ARLM checkpoint modified to use block-wise causal (bidirectional) attention, initialized from the original ARLM weights.
Teacher Model: The original, frozen ARLM.
Training Trajectory (On-Policy): Instead of sampling random masks from a clean dataset, the student DLM samples its own reverse decoding trajectories ( $\tau \sim \Gamma^{rev}_{\theta}$ ). These trajectories represent the actual states the model encounters during inference.
Distillation Objective:
- From a sampled reverse trajectory, a partially denoised state $x_t$ is selected.
- For each masked token position $i$ in a block, the student's predictive distribution $p_{\theta}$ is trained to match the teacher's distribution.
- Crucially, the teacher distribution is conditioned on the unmasked causal prefix of the terminal sequence $x_0$ (the final generated output of that trajectory). Specifically, the target is $p_{ARLM}(\cdot | x_{<b}^0, x_{<i}^0)$ .
- The loss is a Kullback-Leibler (KL) divergence between the teacher's token-level distribution and the student's distribution on the masked positions.

Curriculum Learning

To address instability immediately after conversion (where the student may generate low-quality trajectories due to the structural shift from causal to block-wise attention), the authors introduce a Rollout-Length Curriculum. The length of the on-policy generation is progressively increased from a minimum ( $L_{min}$ ) to a maximum ( $L_{max}$ ) over warmup steps, ensuring the student learns to match the teacher on coherent contexts before tackling longer sequences.

3. Key Contributions

Formulation of ARLM-to-DLM as Post-Training: The paper demonstrates that DLMs can be obtained without the computationally intensive stage of pretraining from scratch, positioning the conversion as an efficient post-training procedure.
Resolution of Distribution Shifts: By training on the student's own reverse trajectories (on-policy), OPDLM eliminates the train-inference state mismatch. By using the original ARLM as a teacher, it ensures knowledge retention.
Extreme Data Efficiency: The method achieves competitive performance using 15× to 7,000× fewer training tokens compared to established DLM baselines (e.g., LLaDA, Dream, SDAR).
General and Task-Specific Applicability: The framework is shown to produce both general-purpose DLMs and task-specialized models (e.g., for math) without requiring a separate DLM pretraining stage.

4. Experimental Results

The authors evaluated OPDLM (4B and 8B scales) against baselines like LLaDA, Dream, SDAR, and Fast-dLLM-v2 across general knowledge, mathematics, and coding benchmarks.

Performance vs. Efficiency: OPDLM-8B achieved strong results on benchmarks like AIME-24 (14.7% accuracy) and GPQA-Diamond (36.1%) while requiring only 0.066 billion training tokens and 4.2 × 10¹⁸ FLOPs. In contrast, baselines required hundreds of billions of tokens (e.g., LLaDA-8B used 1,500B tokens).
Ablation Studies:
- On-Policy vs. Off-Policy: Replacing the on-policy trajectory sampling with standard random masking (Off-Policy) resulted in consistent performance drops, confirming that training on the model's own inference states is the primary driver of success.
- Knowledge Retention: OPDLM retained zero-shot capabilities not explicitly targeted during conversion, including extended thinking (reasoning within <thought> tags) and multilingual generation, despite being trained on a curated English corpus.
Inference Efficiency: The models support multi-token decoding. By adjusting the decoding confidence threshold and block size, the authors demonstrated a trade-off where throughput can exceed 2–3 tokens per step with moderate accuracy costs.
Task-Specific Adaptation: When trained strictly on math data (MATH dataset), OPDLM-MATH outperformed standard DLM baselines and matched specialized RL-trained models (TraDo) without requiring a pre-trained DLM or reinforcement learning with verifiable rewards.

5. Significance and Claims

The paper claims that OPDLM establishes a new Pareto frontier for the efficiency of DLM training. By leveraging on-policy distillation, the authors argue that the conversion of ARLMs to DLMs can be achieved with minimal computational resources, making diffusion-based text generation accessible without the "prohibitive cost" of pretraining from scratch.

The authors emphasize that this approach allows the community to adapt future ARLMs into diffusion models efficiently, preserving the rich knowledge of the original models while gaining the parallel decoding benefits of diffusion. The work is presented as a practical pathway to bridge the gap between the maturity of autoregressive models and the potential advantages of diffusion models.