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

TL;DR 通过在线策略蒸馏将自回归模型转化为扩散语言模型，消除训练-推理不匹配并保留原始模型知识，token 效率提升最高 7000 倍。

问题

问题背景

扩散语言模型(DLM)在非自回归生成、灵活可控文本合成等场景中潜力显著，但大规模预训练成本极高。因此，将该领域已有的强大自回归语言模型(ARLM)高效转换为 DLM，成为一种兼顾质量与效率的研究热点。

现有方法局限

已有工作通常将 ARLM 的因果注意力替换为双向注意力，然后使用标准的 DLM 目标进行训练。这种两步操作引入了两个关键 分布偏移：

• 知识丢弃：从基于下一 token 预测的 ARLM 目标直接切换到基于掩码重建的 DLM 目标，会丢弃原模型在预训练中积累的丰富知识，导致性能退化。
• 训练-推理不匹配：标准 DLM 的训练损失是在随机掩码序列上计算，而推理时则依赖基于置信度的解码策略生成的轨迹，两者分布不一致，大幅降低实际生成质量。

为什么这个问题难且重要

该问题的核心难点在于：既要保留 ARLM 的预训练知识，又要在 DLM 的生成范式下消除训练与推理的 mismatch。若直接从头预训练 DLM，算力消耗难以承受（动辄数千 GPU 小时）；而简单的转换方法却会严重损害模型能力。因此，一种能高效保留知识、同时对齐训练与推理行为的转换方案，对降低扩散语言模型的应用门槛至关重要，也契合业界对 后训练技术 和 知识复用 的高度关注。

行业类比

这类似于在强化学习中，通过 on-policy 蒸馏 将大型策略网络的知识迁移到小型网络，同时修正 off-policy 带来的行为偏差——只不过这里是将自回归“教师”的知识注入到扩散“学生”，并保证推理时的轨迹与训练时一致。

核心洞察

• On-Policy 蒸馏同时解决了扩散语言模型的知识遗忘与训练-推理失配问题：传统 ARLM 转 DLM 方法直接切换目标函数，导致模型遗忘预训练知识，且训练时随机掩码与推理时置信度解码轨迹不一致。OPDLM 让学生模型（双向注意力 ARLM）生成自身推理轨迹，并由冻结的原始 ARLM 教师提供目标 logits，使训练目标与推理路径对齐，同时蒸馏教师知识，避免从头训练，从而在极低数据量下保持竞争力。
• 扩散语言模型可以视为自回归模型的高效后训练阶段，而非独立预训练范式：以往 DLM 需要从零预训练，计算成本高昂。OPDLM 通过微调现有 ARLM 并加入双向注意力，用 on-policy 蒸馏在少量 token（仅为原预训练的 1/15 至 1/7000）上完成转换，性能与从头训练 DLM 相当。这重新定位了 DLM 的训练策略，使其更易被工业界采纳，降低了尝试新生成范式的门槛。

方法

核心思路

OPDLM 提出一种将自回归语言模型（ARLM）高效转换为扩散语言模型（DLM）的方法，核心是用在线策略蒸馏（On-Policy Distillation, OPD） 消除两类分布偏移。

输入与模型结构

• 教师模型：冻结的预训练 ARLM（因果注意力），保留原始知识。
• 学生模型：从教师初始化，但将因果注意力替换为双向注意力，作为待训练的 DLM。

关键模块：Self-OPD 训练流程

1. 轨迹生成：学生模型以扩散方式（如使用 confidence-based decoding）自生成完整序列，而非从随机掩码序列开始。
2. 知识蒸馏：教师模型对这些在线轨迹进行前向计算，输出目标 logits。
3. 损失计算：学生模型的输出分布与教师提供的目标分布对齐（如 KL 散度），同时保留扩散损失的形式。
   这一过程闭环迭代：学生生成轨迹 → 教师提供软标签 → 学生更新，完全在推理时的数据分布上训练，消除常规 DLM 的训练-推理不匹配。

输出与效率

训练完成后，学生即为可直接用于置信度引导解码的 DLM，且无需从零预训练。相比从头训练 DLM 或直接转换方法，OPDLM 训练 token 数降低 15× 至 7000×，在多项任务上保持竞争力，将 DLM 转换重塑为一种 ARLM 后训练（post-training） 形式。

与同类方法的差异

以往工作（如 DiffusionBERT）常将因果注意力换成双向注意力后，直接在随机掩码序列上优化 DLM 目标，导致丢弃 ARLM 知识并引入训练-推理分布 gap。OPDLM 通过在线生成轨迹并蒸馏原始模型，首次同时解决这两类偏移，且 token 效率大幅领先。

实验

实验设计

作者将因果注意力 ARLM 转换为双向注意力 DLM，并通过 on-policy distillation 进行训练。学生模型以双向注意力 ARLM 初始化，生成自身的掩码与去噪轨迹；教师模型为冻结的原始 ARLM，在每条轨迹上提供目标 logits 作为蒸馏信号。实验覆盖语言建模与生成质量基准，对比三类方法：

• 从头预训练的标准 DLM（算力基线）
• 先前的 ARLM→DLM 转换方法（如替换注意力+微调）
• 所提 self-OPD 框架

关键发现

1. 训练 token 效率跃升：OPDLM 仅需基线 DLM 预训练 token 量的 1/15 至 1/7000，便达到可比甚至更优的困惑度与生成质量。
2. 知识保留：通过蒸馏，OPDLM 几乎完全避免 ARLM 知识的遗忘，在零样本下游任务上较标准 DLM 转换有显著提升。
3. 消除训练-推理不匹配：在线策略训练完全对齐推理时的置信度解码，生成轨迹更连贯，暴露偏差显著降低。

与基线对比的深度解读

• vs. 从头预训练 DLM：OPDLM 将 DLM 训练从昂贵的大规模预训练转变为 ARLM 的 post-training 扩展，大幅降低算力门槛，同时性能持平或更好，证明 DLM 能力可以从 ARLM 高效继承。
• vs. 前人转换方法：前人仅通过替换注意力并微调，引入分布偏移，训练损失在随机掩码上计算，导致推理时性能下降。OPDLM 通过 on-policy 轨迹蒸馏 直接对齐推理分布，避免了知识损失与不匹配，是更本质的解决方案。

实验证实，OPDLM 让 DLM 的实用化不再依赖重头预训练，而是作为 ARLM 的轻量后训练步骤，为扩散语言模型的大规模部署开辟了高效路径。

行业影响

落地场景

该方法可将已有的自回归语言模型（ARLM）高效转化为扩散语言模型（DLM），适用于对文本生成质量、多样性和可控性有高要求的场景：

• 内容平台：自动生成广告文案、社交媒体贴文，允许通过扩散模型的迭代式生成灵活控制语气、风格与关键词，提升内容生产效率。
• 对话系统与客服：将微调后的 ARLM 蒸馏为 DLM，可支持更安全、更低延迟的响应生成，并在推理时通过置信度导向解码减少有害输出。
• 代码辅助工具：在代码补全中，DLM 的迭代编辑能力可帮助用户逐步细化函数体或修复 bug，且借助蒸馏可复用已有的强大代码 ARLM。

商业价值

• 大幅降低训练成本：仅需原预训练数据量的 1/15 至 1/7000，便可获得性能相近的扩散模型，直接为企业节省数百万美元的算力支出。
• 加速模型迭代：对于已有 ARLM 部署的企业，将该转换作为一种 post-training 步骤，无需重新预训练即可快速获得扩散版本，显著缩短新产品上线周期。
• 提升用户体验：DLM 在生成任务中天然支持灵活的条件控制与细粒度修订，可提供更高质量的个性化内容，增强用户粘性与付费意愿。

与现有产品/工作流的接口

OPD 蒸馏可作为 现有 LLM 推理平台的附加模块集成：

1. 模型转换阶段：对已部署的 ARLM（如 LLaMA、GPT 系列）进行结构微调（将 causal attention 替换为 bidirectional attention），并利用冻结的教师模型提供目标 logits 进行 on-policy 蒸馏，产出 DLM。
2. 推理服务：将产物 DLM 加载到标准推理引擎（如 vLLM、TensorRT-LLM），需适配扩散模型的解码算法（如 confidence-based decoding），但与现有 API 网关、负载均衡等基础设施完全兼容。
3. 监控与评测：可直接复用现有的文本质量评估管线（如 ROUGE、人类偏好模型），并对蒸馏过程的 token 效率进行持续监测。

具体落地用例

• 电商推荐文案生成：电商平台需要在不同促销场景下生成大量风格迥异的商品描述。先基于大规模 ARLM 微调得到通用文案模型，再通过 OPD 蒸馏为扩散模型，即可在推理时通过调整噪声尺度控制文案的创意性，同时保留原 ARLM 的商品知识，无需为每种风格单独训练模型。
• 医疗报告摘要：在医疗信息系统中，需将长篇病例总结为结构化摘要，且对医学术语准确性要求极高。利用已有的医学领域 ARLM 作为教师，通过 OPD 生成一个扩散版本的摘要模型，可在不泄露原始训练数据的情况下，以少量 token 完成转化，并支持医生在生成后对摘要进行交互式编辑，提升临床体验。

局限

• OPDLM 的蒸馏过程依赖一个冻结的原始自回归教师模型，这意味着当原始模型规模很大时，推理教师模型仍会引入可观的额外计算和存储开销，不适用于极致资源受限或需要完全从零训练的场景。
• 目前实验主要在中等规模模型和标准基准上进行验证，缺乏在更大规模（如数十亿参数）模型以及更复杂的生成任务（如多轮对话、长文本生成）上的评估，因此扩展到生产级大模型的可行性与收益尚不明确。
• OPDLM 训练时的 on-policy 轨迹生成需学生模型完成自回归采样，这增加了训练过程的复杂度与时间开销，同时蒸馏效果可能受教师模型自身质量上限的约束，若教师模型本身在扩散后任务上表现不佳，蒸馏很难带来额外增益。