推理效率狂飆60倍：DiDi-Instruct讓擴散大模型16步超越千步GPT

近日，來自普渡大學、德克薩斯大學、新加坡國立大學、摩根士丹利機器學習研究、小紅書 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 提出了一套系統性的解決方案，其關鍵創新包括：

1. 基于策略梯度的分布匹配目標：DiDi-Instruct 巧妙地將蒸餾目標重構為一種策略梯度（Policy Gradient）的數學形式，然后通過引入一個獎勵函數來指導學生模型的更新，優雅地繞過了在離散空間中求導的難題。
2. 通過對抗學習動態塑造獎勵函數：為了獲得上述獎勵信號，DiDi-Instruct 引入了一個輔助的判別器網絡（discriminator）。該網絡通過對抗訓練，學習區分 “學生” 和 “教師” 在任意中間步驟生成的噪聲樣本，其輸出的對數密度比（log-density ratio）便構成了指導學生優化的精確獎勵信號。
3. 穩定訓練與高質量推理的關鍵技術：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 的多模態研究員和紐約摩根士丹利的機器學習研究員。