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

📄 arXiv: 2606.06712v1 📥 PDF

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

分类: cs.CL, cs.AI

发布日期: 2026-06-04

💡 一句话要点

提出基于策略蒸馏的自回归到扩散语言模型转换方法

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 自回归模型 扩散模型 策略蒸馏 知识保留 自然语言处理 模型转化 训练效率

📋 核心要点

  1. 现有方法在自回归模型转化为扩散模型时，面临知识丢失和训练推理不匹配的挑战。
  2. 本文提出的OPDLM通过自我策略蒸馏，直接在策略上进行训练，从而解决了上述问题。
  3. 实验结果显示，OPDLM在多种任务中表现优异，训练标记需求显著降低，提升幅度达到7000倍。

📝 摘要（中文）

本研究探讨了自回归语言模型（ARLMs）向扩散语言模型（DLMs）的转化。现有方法通过将自回归模型中的因果注意力替换为双向注意力，并使用DLM目标进行训练，但面临两个分布转变的问题。为了解决这些挑战，本文提出了一种基于策略蒸馏的扩散语言模型（OPDLM），通过自我策略蒸馏训练，消除了训练与推理之间的不匹配，同时增强了知识保留。实验结果表明，OPDLM在多种任务中表现出色，所需训练标记减少了15倍至7000倍。

🔬 方法详解

问题定义：本文旨在解决自回归语言模型向扩散语言模型转化过程中的知识丢失和训练推理不匹配问题。现有方法在目标转变时可能会丢失ARLM在训练中获得的知识，且标准DLM在训练时使用随机掩码序列，导致推理时的表现不佳。

核心思路：提出的OPDLM通过自我策略蒸馏（self-OPD）进行训练，学生模型（带有双向注意力的ARLM）生成自己的轨迹，教师模型（原始的冻结ARLM）则在这些轨迹上提供目标logits，从而保留知识并消除训练推理不匹配。

技术框架：OPDLM的训练过程包括两个主要阶段：首先，学生模型生成基于自身的轨迹；其次，教师模型对这些轨迹进行知识蒸馏。整个过程在策略上进行，确保训练和推理的一致性。

关键创新：OPDLM的核心创新在于通过自我策略蒸馏实现了ARLM到DLM的高效转化，避免了传统方法中的知识丢失和训练推理不匹配问题。

关键设计：在训练过程中，采用了特定的损失函数来优化学生模型的输出，使其更接近教师模型的目标logits，同时确保双向注意力机制的有效性。

📊 实验亮点

实验结果表明，OPDLM在多种任务中表现优异，所需训练标记数量减少了15倍至7000倍，相比于传统DLM方法，显著降低了训练成本，同时保持了强大的性能表现。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理中的文本生成、对话系统和机器翻译等。通过降低训练成本，OPDLM为大规模语言模型的高效训练提供了新的思路，未来可能推动更多基于扩散模型的应用落地。

📄 摘要（原文）

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.