数据高效的自回归到扩散语言模型：通过在线策略蒸馏实现

背景

近年来，扩散语言模型（Diffusion Language Models, DLMs）作为一种新兴的文本生成范式，因其能够并行生成文本而受到广泛关注。然而，从头开始预训练一个强大的 DLM 需要巨大的计算成本和海量数据。现有的研究尝试通过将预训练的自回归语言模型（Autoregressive Language Models, ARLMs）中的因果注意力机制替换为双向注意力机制，并将其目标函数转换为 DLM 目标，来实现 ARLM 向 DLM 的转化。

尽管这种“转换”策略避免了从头预训练的高昂成本，但 prior work（先前工作）存在两个主要的分布偏移（distribution shifts）问题：

1. 知识丢失：从“下一个 token 预测”目标转向 DLM 目标的过程中，ARLM 在训练期间习得的大量知识可能会被丢弃。
2. 训练-推理不匹配（Train-Inference Mismatch）：标准的 DLM 训练损失定义在随机掩码序列上，而推理阶段产生的轨迹是由基于置信度的解码（confidence-based decoding）生成的。这种训练与推理数据分布的不一致限制了模型的实际性能。

为了解决上述挑战，研究人员提出了一种新的方法，旨在通过更高效的蒸馏策略，将 ARLM 转化为高性能的 DLM。

核心内容

本文提出了一种名为 On-Policy Diffusion Language Model (OPDLM) 的新框架，其核心在于使用 On-Policy Distillation (OPD，在线策略蒸馏) 来实现从 ARLM 到 DLM 的转变。

1. 方法论：OPDLM 与 OPD

OPDLM 通过“自 OPD”（self-OPD）的方式进行训练。这一过程涉及两个关键角色：

• 学生模型（Student）：这是一个具有双向注意力的 ARLM。它负责生成自己的生成轨迹（trajectories）。
• 教师模型（Teacher）：这是原始的、冻结权重的 ARLM。它通过在学生模型生成的轨迹上提供目标 logits（对数几率），来蒸馏其知识。

2. 解决分布偏移

OPDLM 的设计直接针对了先前方法中的两个痛点：

• 消除训练-推理不匹配：通过以“在线策略”（on-policy）的方式直接训练，即学生模型生成的轨迹直接用于训练，OPDLM 消除了标准 DLM 中训练损失定义与推理轨迹之间的分布差异。
• 增强知识保留：通过从原始冻结的 ARLM 进行蒸馏，模型能够保留 ARLM 在预训练阶段习得的丰富知识，避免了因目标函数改变而导致的信息丢失。

3. 效率与性能

实证结果显示，OPDLM 在保持强大性能的同时，极大地提高了数据效率。具体而言，OPDLM 所需的训练 token 数量比传统从头预训练的方法减少了 15 倍到 7,000 倍。

这一结果证明了 DLM 的转换可以被视为一种 ARLM 的**后训练（post-training）**形式，而非独立的预训练任务。这种方法避免了 DLM 预训练的禁止性成本（prohibitive cost），为高效构建扩散语言模型提供了一条可行路径。

关键要点

• 创新框架：提出了 OPDLM（On-Policy Diffusion Language Model），利用在线策略蒸馏（OPD）将自回归模型转化为扩散语言模型。
• 双重优势：
  • 知识保留：通过从原始冻结 ARLM 蒸馏，保留了预训练知识。
  • 消除不匹配：通过在线策略训练，解决了训练数据（随机掩码）与推理数据（置信度解码轨迹）之间的分布偏移问题。
• 极高的数据效率：相比从头预训练，OPDLM 所需的训练 token 减少了 15x 至 7,000x，显著降低了计算资源需求。
• 范式转变：将 DLM 的构建从“从头预训练”重新定位为“ARLM 的后训练”过程，降低了扩散语言模型的采用门槛。
• 广泛适用性：在多种任务中均表现出强大的性能，证明了该方法的有效性和通用性。

意义与影响

这项研究对自然语言处理领域，特别是生成式 AI 的发展具有深远意义：

1. 降低扩散模型门槛：扩散模型在图像生成中已证明其优势，但在文本生成中因训练成本高而难以普及。OPDLM 证明了可以通过高效蒸馏现有强大的 ARLM（如 Llama、Mistral 等）来快速获得高性能的扩散语言模型，无需昂贵的从头预训练。
2. 优化推理效率：扩散模型的核心优势在于并行生成。OPDLM 解决了训练-推理不匹配问题，意味着实际部署时，模型能更好地发挥并行生成的速度优势，同时保持生成质量。
3. 推动混合架构研究：该方法展示了结合自回归模型（强大的知识基础）和扩散模型（并行生成能力）的潜力，为未来混合架构的语言模型设计提供了新思路。
4. 资源友好型 AI：对于计算资源有限的研究机构或企业，OPDLM 提供了一种经济高效的途径，利用现有的开源或商业 ARLM 快速构建具备扩散特性的语言模型，促进了技术的民主化。

总之，OPDLM 不仅是一个技术改进，更是一种范式上的转变，它使得扩散语言模型从“理论上的高效”走向“实际可用的高效”，为下一代文本生成模型的发展奠定了重要基础。