近日，來(lái)自普渡大學(xué)、德克薩斯大學(xué)、新加坡國(guó)立大學(xué)、摩根士丹利機(jī)器學(xué)習(xí)研究、小紅書(shū) hi-lab 的研究者聯(lián)合提出了一種對(duì)離散擴(kuò)散大語(yǔ)言模型的后訓(xùn)練方法 —— Discrete Diffusion Divergence Instruct (DiDi-Instruct)。經(jīng)過(guò) DiDi-Instruct 后訓(xùn)練的擴(kuò)散大語(yǔ)言模型可以以 60 倍的加速超越傳統(tǒng)的 GPT 模型和擴(kuò)散大語(yǔ)言模型。

DiDi-Instruct 蒸餾得到的 “學(xué)生” 模型與教師模型、GPT-2 的文本生成效率對(duì)比。

DiDi-Instruct 提出了一種獨(dú)創(chuàng)的概率分布匹配的后訓(xùn)練策略，可以將原本需要 500 步以上的昂貴的擴(kuò)散語(yǔ)言 “教師”（diffusion Large Language Model, dLLM）模型，蒸餾成一個(gè)僅需 8-16 步生成整個(gè)文本段落的 “學(xué)生” 模型。在 OpenWebText 標(biāo)準(zhǔn)數(shù)據(jù)集上，DiDi-Instruct 語(yǔ)言模型既實(shí)現(xiàn)了超過(guò) 64 倍以上的推理加速，又在性能上同時(shí)顯著超越了被蒸餾的教師擴(kuò)散語(yǔ)言模型（dLLM，1024 步生成）和自回歸的 GPT2 模型（1024 步生成）。DiDi-Instruct 算法同時(shí)提升了大語(yǔ)言模型的推理效率和推理效果。為極端高效的大語(yǔ)言模型落地提供了新的方案。


研究背景  |  大語(yǔ)言模型生成的 “速度極限” 是多少？

近年來(lái)，以自回歸（ARMs）范式為核心的大語(yǔ)言模型（如 ChatGPT，DeepSeek 等模型）取得了巨大成功。然而，自回歸模型逐詞串行生成的固有瓶頸，使其在長(zhǎng)文本生成時(shí)面臨難以逾越的延遲 “天花板”，即使強(qiáng)大的并行計(jì)算硬件也無(wú)計(jì)可施 。作為一種新興的替代范式，擴(kuò)散語(yǔ)言模型（后文將用 dLLM 指代）應(yīng)運(yùn)而生。dLLM 將文本生成重塑為一個(gè)從完全噪聲（或掩碼）序列中迭代去噪、恢復(fù)出完整文本的過(guò)程 。這一模式天然支持并行化語(yǔ)言段落生成，相較于自回歸模型生成速度更快。然而盡管如此，現(xiàn)有最好的 dLLM 在同等模型尺寸下為了達(dá)到與 GPT-2 相當(dāng)?shù)男阅埽匀恍枰噙_(dá)上百次模型迭代。這個(gè)困境不禁讓人疑惑：是否存在模型在極端少的迭代次數(shù)下（如 8-16 次迭代）下能顯著超越 1024 次迭代的 GPT 模型？

破局者  |  DiDi-Instruct：分布匹配訓(xùn)練實(shí)現(xiàn)語(yǔ)言模型極致加速

在上述研究背景下，本篇文章提出了 DiDi-Instruct。簡(jiǎn)而言之，DiDi-Instruct 是一個(gè) dLLM 的后訓(xùn)練算法。一個(gè) dLLM 通過(guò) DiDi-Instruct 算法訓(xùn)練蒸餾之后，可以將原本的 1024 次推理次數(shù)壓縮至 8 到 16 步，同時(shí)可以顯著提升的 dLLM 的建模效果。

DiDi-Instruct 的理論來(lái)源于連續(xù)擴(kuò)散模型中的一個(gè)經(jīng)典單步蒸餾算法：Diff-Instruct。從理論上看，DiDi-Instruct 訓(xùn)練算法的核心思想是最小化一個(gè)少采樣步數(shù)的 “學(xué)生” 模型與多采樣步數(shù)的 “教師” dLLM 模型在整個(gè)離散 Token 去噪軌跡上分布的積分 KL 散度（Integral Kullback-Leibler Divergence）。該目標(biāo)把不同時(shí)間的 KL 以權(quán)重積分匯總，避免只對(duì)齊末端樣本而訓(xùn)練不穩(wěn)的問(wèn)題，從而讓學(xué)生以一種全局、全過(guò)程匹配的方式，高效 “學(xué)習(xí)” 教師的精髓。一旦積分 KL 散度被優(yōu)化至收斂（接近 0 值），少步生成的 “學(xué)生” 模型便在概率意義上吸收了 "教師 dLLM" 的知識(shí)。

圖片
DiDi-Instruct 流程示意：學(xué)生模型（Student）與教師模型（Teacher）從全掩碼序列重建 “干凈文本”，并同時(shí)進(jìn)行加噪處理。隨后，判別器（Discriminator）對(duì)兩者輸出進(jìn)行區(qū)分并給出獎(jiǎng)勵(lì)分?jǐn)?shù)，用作學(xué)生模型的更新信號(hào)，使其在后續(xù)生成中逼近教師分布。經(jīng)過(guò)反復(fù)迭代，Student 能以更少步數(shù)獲得接近 Teacher 的生成質(zhì)量。

然而，想要直接優(yōu)化積分 KL 散度面臨諸多例如離散文本不可微分等理論困難。針對(duì)這些挑戰(zhàn)，DiDi-Instruct 提出了一套系統(tǒng)性的解決方案，其關(guān)鍵創(chuàng)新包括：

基于策略梯度的分布匹配目標(biāo)：DiDi-Instruct 巧妙地將蒸餾目標(biāo)重構(gòu)為一種策略梯度（Policy Gradient）的數(shù)學(xué)形式，然后通過(guò)引入一個(gè)獎(jiǎng)勵(lì)函數(shù)來(lái)指導(dǎo)學(xué)生模型的更新，優(yōu)雅地繞過(guò)了在離散空間中求導(dǎo)的難題。
通過(guò)對(duì)抗學(xué)習(xí)動(dòng)態(tài)塑造獎(jiǎng)勵(lì)函數(shù)：為了獲得上述獎(jiǎng)勵(lì)信號(hào)，DiDi-Instruct 引入了一個(gè)輔助的判別器網(wǎng)絡(luò)（discriminator）。該網(wǎng)絡(luò)通過(guò)對(duì)抗訓(xùn)練，學(xué)習(xí)區(qū)分 “學(xué)生” 和 “教師” 在任意中間步驟生成的噪聲樣本，其輸出的對(duì)數(shù)密度比（log-density ratio）便構(gòu)成了指導(dǎo)學(xué)生優(yōu)化的精確獎(jiǎng)勵(lì)信號(hào)。
穩(wěn)定訓(xùn)練與高質(zhì)量推理的關(guān)鍵技術(shù)：DiDi-Instruct 還引入多項(xiàng)關(guān)鍵設(shè)計(jì)對(duì)該方法進(jìn)行系統(tǒng)性?xún)?yōu)化，以穩(wěn)定訓(xùn)練、緩解熵坍塌、提升推理質(zhì)量。
分組獎(jiǎng)勵(lì)歸一化（Grouped Reward Normalization）：借鑒深度求索（DeepSeek）提出的組相對(duì)策略?xún)?yōu)化（GRPO），DiDi-Instruct 在每個(gè)小批量（mini-batch）內(nèi)對(duì)獎(jiǎng)勵(lì)進(jìn)行標(biāo)準(zhǔn)化。該操作顯著降低了訓(xùn)練梯度的方差，有效提升了訓(xùn)練的穩(wěn)定性。
分步式中間狀態(tài)匹配（Intermediate-state Matching）：通過(guò)分解梯度信息，DiDi-Instruct 使學(xué)生模型在訓(xùn)練中接觸到不同噪聲水平的中間狀態(tài)。這個(gè)機(jī)制有效緩解了困擾許多后訓(xùn)練算法的模型熵坍塌問(wèn)題（mode collapse），保證了學(xué)生模型真正學(xué)習(xí)到生成復(fù)雜，多樣性的內(nèi)容。
獎(jiǎng)勵(lì)驅(qū)動(dòng)的祖先采樣（Reward-guided Ancestral Sampling）：在推理階段，利用訓(xùn)練好的判別器獲得獎(jiǎng)勵(lì)信號(hào)，對(duì)生成過(guò)程進(jìn)行 “梯度引導(dǎo) + 多候選重排序”，進(jìn)一步提升了最終生成文本的質(zhì)量。

圖片