近日，來自普渡大學、德克薩斯大學、新加坡國立大學、摩根士丹利機器學習研究、小紅書 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）：在推理階段，利用訓練好的判別器獲得獎勵信號，對生成過程進行 “梯度引導 + 多候選重排序”，進一步提升了最終生成文本的質量。

圖片