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

2026-06-04

AI Summary

This work addresses the challenges of knowledge loss and training-inference distribution mismatch when converting autoregressive language models (ARLMs) into diffusion language models (DLMs). The authors propose On-Policy Distillation (OPD), a novel method that constructs an On-Policy Diffusion Language Model (OPDLM), wherein the student model learns from target logits provided by the teacher ARLM along trajectories generated by the student itself during inference. This approach uniquely aligns training and inference trajectories, preserving the full knowledge of the original ARLM without requiring pretraining from scratch. By integrating bidirectional attention with diffusion modeling, OPDLM achieves strong performance using only 1/15 to 1/7,000 of the training tokens typically required by conventional DLMs, substantially reducing training costs.

Abstract

We study the transformation of autoregressive models (ARLMs) into diffusion language models (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 Language Model (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.

Problem

• distribution shift
• train-inference mismatch
• knowledge retention
• diffusion language models
• autoregressive models

Innovation

• On-Policy Distillation
• Diffusion Language Models
• Autoregressive-to-Diffusion Transformation
• Knowledge Retention
• Train-Inference Alignment