推理效率狂飙60倍：DiDi-Instruct让扩散大模型16步超越千步GPT

近日，来自普渡大学、德克萨斯大学、新加坡国立大学、摩根士丹利机器学习研究、小红书 hi-lab 的研究者联合提出了一种对离散扩散大语言模型的后训练方法 —— Discrete Diffusion Divergence Instruct (DiDi-Instruct)。经过 DiDi-Instruct 后训练的扩散大语言模型可以以 60 倍的加速超越传统的 GPT 模型和扩散大语言模型。

DiDi-Instruct 蒸馏得到的 “学生” 模型与教师模型、GPT-2 的文本生成效率对比。

DiDi-Instruct 提出了一种独创的概率分布匹配的后训练策略，可以将原本需要 500 步以上的昂贵的扩散语言 “教师”（diffusion Large Language Model, dLLM）模型，蒸馏成一个仅需 8-16 步生成整个文本段落的 “学生” 模型。在 OpenWebText 标准数据集上，DiDi-Instruct 语言模型既实现了超过 64 倍以上的推理加速，又在性能上同时显著超越了被蒸馏的教师扩散语言模型（dLLM，1024 步生成）和自回归的 GPT2 模型（1024 步生成）。DiDi-Instruct 算法同时提升了大语言模型的推理效率和推理效果。为极端高效的大语言模型落地提供了新的方案。

研究背景 | 大语言模型生成的 “速度极限” 是多少？

近年来，以自回归（ARMs）范式为核心的大语言模型（如 ChatGPT，DeepSeek 等模型）取得了巨大成功。然而，自回归模型逐词串行生成的固有瓶颈，使其在长文本生成时面临难以逾越的延迟 “天花板”，即使强大的并行计算硬件也无计可施。作为一种新兴的替代范式，扩散语言模型（后文将用 dLLM 指代）应运而生。dLLM 将文本生成重塑为一个从完全噪声（或掩码）序列中迭代去噪、恢复出完整文本的过程。这一模式天然支持并行化语言段落生成，相较于自回归模型生成速度更快。然而尽管如此，现有最好的 dLLM 在同等模型尺寸下为了达到与 GPT-2 相当的性能，仍然需要多达上百次模型迭代。这个困境不禁让人疑惑：是否存在模型在极端少的迭代次数下（如 8-16 次迭代）下能显著超越 1024 次迭代的 GPT 模型？

破局者 | DiDi-Instruct：分布匹配训练实现语言模型极致加速

在上述研究背景下，本篇文章提出了 DiDi-Instruct。简而言之，DiDi-Instruct 是一个 dLLM 的后训练算法。一个 dLLM 通过 DiDi-Instruct 算法训练蒸馏之后，可以将原本的 1024 次推理次数压缩至 8 到 16 步，同时可以显著提升的 dLLM 的建模效果。

DiDi-Instruct 的理论来源于连续扩散模型中的一个经典单步蒸馏算法：Diff-Instruct。从理论上看，DiDi-Instruct 训练算法的核心思想是最小化一个少采样步数的 “学生” 模型与多采样步数的 “教师” dLLM 模型在整个离散 Token 去噪轨迹上分布的积分 KL 散度（Integral Kullback-Leibler Divergence）。该目标把不同时间的 KL 以权重积分汇总，避免只对齐末端样本而训练不稳的问题，从而让学生以一种全局、全过程匹配的方式，高效 “学习” 教师的精髓。一旦积分 KL 散度被优化至收敛（接近 0 值），少步生成的 “学生” 模型便在概率意义上吸收了 "教师 dLLM" 的知识。

图片

DiDi-Instruct 流程示意：学生模型（Student）与教师模型（Teacher）从全掩码序列重建 “干净文本”，并同时进行加噪处理。随后，判别器（Discriminator）对两者输出进行区分并给出奖励分数，用作学生模型的更新信号，使其在后续生成中逼近教师分布。经过反复迭代，Student 能以更少步数获得接近 Teacher 的生成质量。

然而，想要直接优化积分 KL 散度面临诸多例如离散文本不可微分等理论困难。针对这些挑战，DiDi-Instruct 提出了一套系统性的解决方案，其关键创新包括：

基于策略梯度的分布匹配目标：DiDi-Instruct 巧妙地将蒸馏目标重构为一种策略梯度（Policy Gradient）的数学形式，然后通过引入一个奖励函数来指导学生模型的更新，优雅地绕过了在离散空间中求导的难题。

通过对抗学习动态塑造奖励函数：为了获得上述奖励信号，DiDi-Instruct 引入了一个辅助的判别器网络（discriminator）。该网络通过对抗训练，学习区分 “学生” 和 “教师” 在任意中间步骤生成的噪声样本，其输出的对数密度比（log-density ratio）便构成了指导学生优化的精确奖励信号。

稳定训练与高质量推理的关键技术：DiDi-Instruct 还引入多项关键设计对该方法进行系统性优化，以稳定训练、缓解熵坍塌、提升推理质量。

分组奖励归一化（Grouped Reward Normalization）：借鉴深度求索（DeepSeek）提出的组相对策略优化（GRPO），DiDi-Instruct 在每个小批量（mini-batch）内对奖励进行标准化。该操作显著降低了训练梯度的方差，有效提升了训练的稳定性。

分步式中间状态匹配（Intermediate-state Matching）：通过分解梯度信息，DiDi-Instruct 使学生模型在训练中接触到不同噪声水平的中间状态。这个机制有效缓解了困扰许多后训练算法的模型熵坍塌问题（mode collapse），保证了学生模型真正学习到生成复杂，多样性的内容。

奖励驱动的祖先采样（Reward-guided Ancestral Sampling）：在推理阶段，利用训练好的判别器获得奖励信号，对生成过程进行 “梯度引导 + 多候选重排序”，进一步提升了最终生成文本的质量。

图片