![]()
近日,来自普渡大学、德克萨斯大学、新加坡国立大学、摩根士丹利机器学习研究、小红书 hi-lab 的研究者联合提出了一种对离散扩散大语言模型的后训练方法 ——DiscreteDiffusion DivergenceInstruct(DiDi-Instruct)。经过 DiDi-Instruct 后训练的扩散大语言模型可以以 60 倍的加速超越传统的 GPT 模型和扩散大语言模型。
![]()
DiDi-Instruct 提出了一种独创的概率分布匹配的后训练策略,可以将原本需要 500 步以上的昂贵的扩散语言 “教师”(diffusion Large Language Model, dLLM)模型,蒸馏成一个仅需 8-16 步生成整个文本段落的 “学生” 模型。在 OpenWebText 标准数据集上,DiDi-Instruct 语言模型既实现了超过 64 倍以上的推理加速,又在性能上同时显著超越了被蒸馏的教师扩散语言模型(dLLM,1024 步生成)和自回归的 GPT2 模型(1024 步生成)。DiDi-Instruct 算法同时提升了大语言模型的推理效率和推理效果。为极端高效的大语言模型落地提供了新的方案。
![]()
论文标题:Ultra-Fast Language Generation via Discrete Diffusion Divergence Instruct论文链接:www.arxiv.org/abs/2509.25035代码仓库:github.com/haoyangzheng-ai/didi-instruct项目地址:haoyangzheng.github.io/research/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):在推理阶段,利用训练好的判别器获得奖励信号,对生成过程进行 “梯度引导 + 多候选重排序”,进一步提升了最终生成文本的质量。
![]()
DiDi-Instruct 后训练算法。
![]()
奖励驱动的祖先采样算法。
科学实验 | 效率与性能的双重飞跃
研究团队在公开的 OpenWebText 数据集上进行了详尽的实验,结果出人出人意料:经过 DiDi-Instruct 后训练的语言模型在效率和效果上得到了双重提升。
1.性能与质量新标杆:DiDi-Instruct 在生成质量和效率上均达到了新的 SOTA 水平。该工作系统性地将 DiDi-Instruct 与 GPT-2、MDLM、DUO、SDTT 等多个基准模型进行了比较。结果显示,在 OpenWebText 数据集上,DiDi-Instruct 在 8 到 128 步的所有函数评估次数(NFEs)设置下,其困惑度(Perplexity)指标全面且持续地优于所有基准模型。一个尤为亮眼的成果是,仅需 16 步函数评估,DiDi-Instruct 生成的文本质量 Perplexity(PPL)就已经超越了需要 1024 步才能完成生成的教师模型,相比最强的基线模型提升超过 30%。同时,这些性能增益是在几乎没有熵损失(约 1%)的情况下实现的,充分保证了生成内容的多样性。
![]()
DiDi-Instruct 蒸馏所得学生模型与基准模型在不同函数评估次数(NFEs)下的文本生成困惑度(PPL)对比。
2.训练效率大幅提升:DiDi-Instruct 不仅生成质量高,其训练(蒸馏)过程也极为高效。出人意料的时候,整个蒸馏框架的训练仅需在单张 NVIDIA H100 GPU 上运行约 1 小时即可完成。相比之下,其他同类蒸馏方法(基线模型)通常需要超过倍以上的训练时间。这意味着 DiDi-Instruct 将训练效率提升了超过 20 倍,极大地降低了开发者迭代和部署高性能生成模型的门槛。
3.跨领域通用性验证:研究团队在报告中指出,DiDi-Instruct 的蒸馏框架是为离散扩散模型设计的,并不局限于语言模型。为了验证这一点,团队将其成功应用于一个完全不同的领域:无条件蛋白质序列生成。他们使用一个预训练的蛋白质语言扩散模型(DPLM)作为教师模型进行蒸馏。结果表明,蒸馏后的学生模型保留了教师模型生成可变长度序列的能力,同时大幅降低了推理成本。更重要的是,学生模型在极少步数下即可生成结构合理的高置信度蛋白质结构。这一跨领域实验有力地证实了 DiDi-Instruct 作为通用离散序列生成加速框架的巨大潜力。
![]()
由 DiDi-Instruct 蒸馏得到的学生模型生成的 高置信度蛋白质序列(pLDDT > 70)。
4.深入消融实验,探究各组件的核心贡献:为了科学地验证每个创新组件的必要性和贡献,研究团队还进行了详尽的 “逐项累加”(cumulative)和 “逐一剔除”(leave-one-out)的消融研究。这些实验揭示了模型性能的关键驱动因素:
中间状态匹配是框架稳定的基石:实验表明,虽然单独加入该模块对性能提升有限,但在完整的模型中一旦移除,模型性能会灾难性下降(PPL > 30,000),证明了其在复杂优化环境下的关键稳定作用。时间步耦合能高效提升蒸馏性能:该技术将 8 步生成下的困惑度从 600 + 骤降至 100 左右,凸显了对齐奖励信号与分数函数中间状态的重要性。而在目标函数中增加权重信息则能进一步提升模型训练效果。正则化项扮的 “双重角色”:在极少步数下(如 8 NFEs),它能有效稳定训练,防止离散误差导致训练目标偏离。然而在更多步数(≥ 16 NFEs)的采样中,移除正则化反而能取得更好的结果,这表明此时过强的约束会限制模型的表达能力。引导式推理的作用解读:在少步数(如 8 NFEs)时,它能显著降低困惑度(困惑度相对改善约 30%),提升文本生成质量。而在多步数下,它对困惑度影响甚微,但能显著提升生成样本的多样性(熵从 5.00 提升至 5.15),这与奖励驱动的祖先采样设计的先 “梯度引导” 后 “多候选重排序” 的混合策略设计完美契合。
![]()
“逐项累加” 消融实验结果见表 1,“逐一剔除” 消融实验结果见表 2。
技术展望 | 开启高效生成模型新范式
DiDi-Instruct 的提出,不仅是离散扩散模型加速技术的一次技术突破,也为广泛的大语言模型的极限加速,对齐和强化学习提供了新的思路。它首次成功地将分布匹配蒸馏思想应用于基于掩码的离散扩散模型,并建立了一套集 “分布匹配目标、稳定训练、高效推理” 于一体的完整框架。这项工作展示了通过系统性的算法与框架设计,可以克服现阶段大语言模型在生成效率上的瓶颈,使其成为下一代 AI 内容生成中(多模态生成、代码生成、生物序列设计等领域)极具竞争力的选项。我们非常期待将 DiDi-Instruct 应用于最前沿的超大规模的扩散语言模型的效果。
团队简介
本论文第一作者郑昊阳,目前于美国普渡大学攻读博士学位,导师为林光老师。林光是普渡大学的 Moses Cobb Stevens 教授兼理学院副院长。论文的两位通讯作者罗维俭和邓伟分别是小红书 hi-lab 的多模态研究员和纽约摩根士丹利的机器学习研究员。





京公网安备 11011402013531号