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

Abstract: We study the transformation of autoregressive models (ARLMs) into diffusion LLMs (DLMs). Rather than pretraining from scratch, prior work replaces the causal attention in ARLMs with bidirectional attention and then trains the resulting model using a DLM objective. However, these approaches incur two distribution shifts. First, transitioning from a next-token prediction objective to a DLM objective can discard knowledge acquired by the ARLM during training. Second, standard DLMs suffer from a train-inference mismatch, as the training loss is defined on randomly masked sequences rather than the trajectories encountered at inference produced by confidence-based decoding. To address both challenges, we introduce an On-Policy Diffusion LLM (OPDLM) in which On-Policy Distillation (OPD) is employed for ARLM-to-DLM transformation. Specifically, OPDLM is trained via self-OPD, where the student, an ARLM with bidirectional attention, generates its own trajectories, and the teacher, the original frozen ARLM, distills its knowledge by providing target logits on these trajectories. By training directly in an on-policy manner, OPDLM eliminates the train-inference mismatch in DLMs, while distillation from the original model enhances knowledge retention from the ARLM. Empirical results demonstrate that OPDLM requires 15x to 7,000x fewer training tokens with strong performance across a wide variety of tasks. OPDLM avoids the prohibitive cost of DLM pretraining and positions DLM transformation as a form of ARLM post-training.

• The paper presents on-policy distillation to convert ARLMs into DLMs, effectively bridging the train-inference state mismatch.
• It employs reverse-sampled student rollouts with curriculum learning, reducing token requirements by 15× to 7,000× compared to prior methods.
• OPDLM achieves competitive performance on benchmarks like AIME-24 while preserving structural priors and zero-shot capabilities from ARLMs.

Data-Efficient ARLM-to-Diffusion Conversion via On-Policy Distillation

Motivation and Problem Formulation

Diffusion LLMs (DLMs) have gained prominence due to their ability to sample and decode multiple tokens per denoising step and their improved convergence and data learning over ARLMs. However, masked DLMs require trillions of tokens for pretraining, and the conversion process from ARLMs (causal Transformer LMs) to DLMs with bidirectional attention remains expensive and typically discards much acquired ARLM knowledge. Moreover, there is a fundamental mismatch in DLM training: models are trained on random masked states, while inference follows reverse unmasking trajectories guided by confidence-based heuristics, generating a domain shift that manifests as exposure bias and suboptimal performance.

This paper introduces On-Policy Diffusion LLMs (OPDLMs), leveraging On-Policy Distillation for ARLM-to-DLM conversion. OPDLMs operate by initializing a DLM student from a pretrained ARLM checkpoint, letting it generate its own reverse diffusion trajectories, and distilling knowledge from the original ARLM on these inference-time states. This framework provides token-level distribution matching and effectively eliminates the train-inference mismatch characteristic of prior DLM approaches.

OPDLM Training Framework

OPDLM reformulates block diffusion modeling objectives to operate on reverse-sampled trajectories, employing a curriculum to stabilize early rollouts. For each inference-time state in the reverse trajectory, masked positions are supervised via KL divergence against the original ARLM predictive distribution, conditioned on causal prefixes of the terminal sequence. The loss is computed for each block and masked token, weighted uniformly over trajectory time.

Figure 1: OPDLM training framework and distillation loss computation with ARLM supervision on reverse trajectory states.

Curriculum learning is used during training, restricting initial rollouts to short sequences and gradually increasing the rollout length as the student DLM improves, thus avoiding training instability from incoherent outputs.

Empirical Results: Training Efficiency and Performance

OPDLM achieves dramatic improvements in training efficiency, requiring 15× to 7,000× fewer tokens compared to prior AR-to-DLM conversion baselines. OPDLM-8B, for instance, attains competitive results on the AIME-24 benchmark using only $0.066$B tokens and $4.2 \times 10^{18}$ FLOPs.

Figure 2: OPDLM establishes a new pareto frontier for AIME-24, demonstrating vastly improved training efficiency versus established DLM baselines.

Across general knowledge, math, and code generation tasks, OPDLM achieves accuracy comparable to competitive baselines trained on significantly larger corpora, with stronger relative gains on rigorous tasks (e.g., GPQA-Diamond and AIME) than on saturated datasets. Importantly, OPDLM preserves structural priors and zero-shot capabilities from its ARLM teacher, evidenced by its extended reasoning and multilingual proficiency despite the absence of such data during conversion.

Ablation: On-Policy vs. Off-Policy Distillation

The primary driver of OPDLM performance is on-policy data: distillation on states sampled from the student’s generation process (reverse trajectories) consistently outperforms off-policy random masking objectives. When the training state is switched to on-policy but random corruption is used, performance remains comparable, indicating that the crucial factor is using student rollouts for distillation.

Multi-Token Decoding and Throughput

OPDLM enables flexible control over inference throughput via two mechanisms: the decoding confidence threshold $\gamma$ and training block size. Lowering $\gamma$ or increasing block size enables parallel decoding of multiple tokens per denoising step, with a modest trade-off in task accuracy.

Figure 3: MATH-500 accuracy and average tokens/denoising step for varying decoding thresholds and block sizes, illustrating the efficiency-throughput-accuracy tradeoff.

Figure 4: Block size effects at different thresholds across diverse benchmarks; larger blocks increase throughput but reduce accuracy.

Figure 5: Decoding threshold effects at various block sizes across benchmarks; lower confidence threshold increases tokens/step at variable accuracy costs.

Figure 6: Training curves for average token throughput per denoising step; higher block sizes achieve greater parallelism.

Specialist DLMs and Task-Specific Post-Training

The OPDLM framework offers direct conversion of ARLM task specialists to DLMs without requiring intermediate DLM teacher pretraining or reinforcement learning with reward signals. For instance, in mathematical reasoning domains, OPDLM-MATH with chain-of-thought distillation achieves strong performance (e.g., AIME-24 accuracy up to 50% at 8B scale) in the absence of explicit reasoning supervision during conversion.

Implications and Future Directions

Practical Implications

• Compute and Data Efficiency: OPDLM provides a highly practical route for deploying DLMs, making them accessible even with limited resources, and substantially reducing environmental costs associated with large-scale pretraining.
• Robustness and Generalization: Preservation of ARLM priors and zero-shot capabilities directly addresses the train-inference mismatch, enabling robust application across heterogeneous task domains and inference procedures.

Theoretical Implications

• Knowledge Retention via Distribution Matching: The alignment of the student’s denoising prediction with the ARLM teacher’s distribution on inference-time rollouts represents a new paradigm for knowledge transfer, potentially applicable to other modalities beyond text.
• Exposure Bias Mitigation in Diffusion Modeling: On-policy distillation for diffusion models eliminates exposure bias analogously to recent advances in autoregressive models, suggesting broader applications in sequence modeling and generative tasks.

Future Developments

• Chain-of-Thought Distillation at Scale: Extending OPDLM to distill reasoning behaviors at scale may yield DLMs with advanced cognitive abilities for complex reasoning benchmarks.
• Cross-Family and Cross-Scale Distillation: Systematic exploration of teacher-student scale mismatch and transfer across model families may further optimize knowledge distillation.
• Data Curation: Given OPDLM’s data efficiency, incremental improvements in training corpus quality are likely to yield disproportionate gains in downstream performance.

Conclusion

OPDLM transforms ARLM-to-DLM conversion into an efficient post-training procedure, resolving both knowledge retention and train-inference state mismatch through on-policy distillation from frozen ARLMs. The resulting DLMs attain competitive task performance across general and specialist domains with orders-of-magnitude reductions in compute and data requirements. This approach positions DLMs as practical successors to autoregressive models, facilitating flexible, high-throughput sequence generation and opening new avenues for efficient model adaptation and deployment (2606.06712).