擴(kuò)散方法打入語言模型領(lǐng)域！

最近，康奈爾博士生Subham Sahoo，在X介紹了擴(kuò)散大語言模型的最新工作。

這項(xiàng)研究引發(fā)了AI研究領(lǐng)域的思考。

英偉達(dá)研究院杰出研究科學(xué)家Pavlo Molchanov說：「擴(kuò)散大語言模型正在崛起！」

谷歌研究院學(xué)生研究員、康奈爾大學(xué)博士生Yash Akhauri更是指出：「自回歸危在旦夕」。

這項(xiàng)新鮮出爐的研究，提出了突破性的方法：Esoteric Language Models（Eso-LMs）。

論文鏈接：https://arxiv.org/abs/2506.01928

項(xiàng)目鏈接：https://s-sahoo.com/Eso-LMs/

這是首個(gè)在保持并行生成的同時(shí)，引入KV緩存機(jī)制的方法。

推理速度相比標(biāo)準(zhǔn)MDM提升了 65 倍，相比支持KV緩存的半自回歸基線模型快3–4倍。

這是一種新的語言建模框架，融合了自回歸（AR）和離散擴(kuò)散模型（MDM）兩種范式，性能超越了之前的混合方法BD3-LMs。

研究者還發(fā)現(xiàn)，BD3-LMs 在低采樣步數(shù)下性能下降，而新方法在低計(jì)算量（NFE）場景下與離散擴(kuò)散模型相當(dāng)，在高計(jì)算量場景下與自回歸模型相當(dāng)。

這次的結(jié)果為離散擴(kuò)散模型建立了新的困惑度（perplexity）最優(yōu)水平，縮小了與自回歸模型的差距。

另外值得一提的是，除了共同一作Zhihan Yang外，還有多位華人作者，其中包括知名華人學(xué)者邢波（Eric Xing）。

語言也能擴(kuò)散

這并非擴(kuò)散方法首次「入侵」文本生成領(lǐng)域。

甚至達(dá)到商用級(jí)別的擴(kuò)散語言模型，都不止一個(gè)。

斯坦福、UCLA和康奈爾的三位教授聯(lián)合創(chuàng)立了Inception Labs，推出了全球首個(gè)商用級(jí)別的擴(kuò)散語言模型

擴(kuò)散語言模型最大特點(diǎn)就是快：推理速度可達(dá)ChatGPT的6倍！

IBM甚至認(rèn)為擴(kuò)散模型就是下一代AI，GPT這類自回歸范式受到有力挑戰(zhàn)。

不過，三位教授具體如何實(shí)現(xiàn)這一突破，目前尚屬商業(yè)機(jī)密，外界難以得知。

而在AI巨頭中，谷歌是第一家嘗試擴(kuò)散語言模型——

在I/O大會(huì)上，它放出了實(shí)驗(yàn)版語言模型Gemini Diffusion：推理每秒可達(dá)1400多token。

而這次的新論文，作者Arash Vahdat是英偉達(dá)研究院的科研總監(jiān)（Research Director），領(lǐng)導(dǎo)基礎(chǔ)生成式人工智能（GenAIR）團(tuán)隊(duì)。

莫非英偉達(dá)也要押注擴(kuò)散語言模型？

擴(kuò)散模型：后來者居上？

眾所周知，掩蔽擴(kuò)散模型（Masked Diffusion Models，MDMs）是自回歸（AR）語言模型的有力替代方案——

但它們有兩個(gè)致命短板：

速度慢：沒有KV緩存 = 實(shí)際上比AR慢得多；

質(zhì)量差：在復(fù)雜任務(wù)中表現(xiàn)不佳，似然度低于AR。

塊擴(kuò)散（Block Diffusion）模型BD3-LM，在每個(gè)區(qū)塊內(nèi)執(zhí)行擴(kuò)散過程，以先前區(qū)塊為條件，實(shí)現(xiàn)分塊生成token序列。

它融合了自回歸模型與擴(kuò)散模型的優(yōu)勢：在支持可變長度生成的同時(shí)，利用KV緩存和并行采樣提升推理效率，從而克服兩種傳統(tǒng)方法的局限性——

既能實(shí)現(xiàn)更高質(zhì)量的生成效果，又能保持高效推理特性。

但是，BD3-LM的速度與質(zhì)量仍需權(quán)衡：

低采樣步數(shù)下出現(xiàn)模式崩塌，導(dǎo)致樣本質(zhì)量差；

而且只支持部分緩存，塊內(nèi)鍵值緩存仍缺失。

針對現(xiàn)有方法在速度與質(zhì)量之間的權(quán)衡，研究者提出了一種全新的混合范式：Eso-LM。

論文鏈接：https://arxiv.org/abs/2503.09573

這次研究人員結(jié)合掩蔽擴(kuò)散和自回歸，提出了新的語言建模范式：Esoteric Language Models (Eso-LMs)。

新范式兼顧了速度與質(zhì)量，超越了BD3-LM。

正如圖1所示，Eso-LM包含擴(kuò)散和順序兩個(gè)階段：

在擴(kuò)散階段（Diffusion Phase），Eso-LM每一步去噪一個(gè)或多個(gè)可能不相鄰的掩蔽token (圖1中底部字母「M」)。

在順序階段（Squential Phase），Eso-LM從左到右逐個(gè)去噪剩余的掩蔽token。

與BD3-LM不同，Eso-LM (B)允許在兩個(gè)階段使用統(tǒng)一的KV緩存，藍(lán)色邊框框住了正在構(gòu)建KV緩存的transformer單元；當(dāng)單元的KV緩存構(gòu)建完成時(shí)，該單元變成藍(lán)色。

下方的序列顯示了transformer中token的自然順序。

圖1: 使用Eso-LM (B) 高效生成示例序列。

這招「KV緩存」原本是自回歸模型加速推理的「殺手锏」。

但Eso-LM利用創(chuàng)新的混合訓(xùn)練方法，將KV緩存引入了擴(kuò)散模型。

具體來說：

• 混合訓(xùn)練：Eso-LM在訓(xùn)練時(shí)一半數(shù)據(jù)采用AR風(fēng)格（干凈的上下文預(yù)測下一個(gè)單詞），另一半采用擴(kuò)散風(fēng)格（打亂輸入，部分掩碼，逐步去噪）。
• 推理優(yōu)化：在生成過程中，Eso-LM只對部分單詞（掩碼和干凈單詞）進(jìn)行前向計(jì)算，并緩存干凈單詞的KV對，大幅減少計(jì)算量。

爆改Transformer

靈活切換注意力

自回歸模型（AR）需要因果注意力和逐個(gè)token解碼，而掩碼去噪模型（MDM）依賴雙向注意力。

要想同時(shí)支持順序（AR）和并行（MDM）生成模式，并使用共享的Transformer架構(gòu)，必須解決它們之間的架構(gòu)不匹配問題。

研究者引入了注意力偏置矩陣A，調(diào)整標(biāo)準(zhǔn)的自注意力機(jī)制：

其中Q，K，V分別表示自注意力機(jī)制中query、key和value矩陣。

偏置矩陣A控制注意力流：當(dāng)Ai,j=0時(shí)，表示「允許」從tokeni注意到j(luò)；當(dāng)Ai,j=?∞時(shí)，表示「阻止」這種注意力。

這種機(jī)制只要一個(gè)transformer，就能根據(jù)需要模擬因果（單向）和雙向注意力行為。

基于統(tǒng)一的注意力機(jī)制，研究者提出了兩個(gè)變體：Eso-LM(A)和Eso-LM(B)。

Eso-LM(A)通過稀疏化注意力并在每一步擴(kuò)散過程中僅將去噪transformer應(yīng)用于部分遮蔽token，從而降低計(jì)算量。

Eso-LM(B)進(jìn)一步擴(kuò)展了這個(gè)想法，不僅對遮蔽token應(yīng)用因果mask，還對干凈token應(yīng)用，從而實(shí)現(xiàn)更高效的KV緩存（KV-caching）——代價(jià)是困惑度略有下降。

擴(kuò)散階段

在擴(kuò)散階段，標(biāo)準(zhǔn)的采樣方法會(huì)浪費(fèi)大量FLOPs。

為了提高效率，研究者對標(biāo)準(zhǔn)采樣和訓(xùn)練過程提出了兩個(gè)關(guān)鍵改進(jìn)。

在采樣過程中，預(yù)先計(jì)算擴(kuò)散去噪計(jì)劃SMDM=(S1,…,S1/T)，其中S_t是在擴(kuò)散步驟t去噪的遮蔽token的索引集合。

而且不再處理整個(gè)序列，而只對子序列{zt?∣?∈C(zt)∪St}進(jìn)行前向傳播——即，干凈的token和計(jì)劃去噪的token——

這在處理長序列時(shí)顯著降低了計(jì)算量。

這種方法支持在擴(kuò)散過程中進(jìn)行高效的KV緩存。

關(guān)鍵思想借用了已有的方法AO-ARM（見下文）：遮蔽token可以按任何順序揭示。

論文鏈接：https://openreview.net/forum?id=sMyXP8Tanm

因此，在訓(xùn)練過程中，新方法要采樣隨機(jī)順序σ～PL，并對每個(gè)σ，強(qiáng)制執(zhí)行對遮蔽token的因果注意力。

具體來說，要求遮蔽token只能對干凈token和根據(jù)順序σ排列的先前遮蔽token進(jìn)行注意力計(jì)算。

Eso-LM(A)采用了這一策略，在采樣過程中顯著減少了計(jì)算量，同時(shí)保持了性能。

而Eso-LM(B)對干凈token強(qiáng)制施加類似的因果mask，進(jìn)一步擴(kuò)展了這一思想，從而實(shí)現(xiàn)了KV緩存。

盡管在困惑度上稍微差一些，Eso-LM(B)在采樣過程中提供了顯著的加速（最多65倍）。

順序階段

自回歸模型隨后從左到右填充遮蔽token，使用順序去噪計(jì)劃，其中要求每個(gè)單元素集合按其唯一元素升序排列。

不同于標(biāo)準(zhǔn)的自回歸解碼，每個(gè)x~?同時(shí)依賴其左側(cè)上下文（完全由干凈token構(gòu)成）和右側(cè)干凈的token，從而實(shí)現(xiàn)更豐富的生成。

我們跳過對右側(cè)遮蔽token的評(píng)估，減少不必要的計(jì)算。

順序階段自然支持KV緩存。

我們將統(tǒng)一的去噪計(jì)劃表示為S=SMDM∪SAR，它將兩個(gè)采樣計(jì)劃連接起來以劃分集合[L]。

當(dāng)α0=1時(shí)，所有token都由擴(kuò)散生成，因此S=S_MDM，且S_AR=?；

當(dāng)α0=0時(shí)，所有token都由順序方式生成，因此S=S_AR，且S_MDM=?。

完整采樣算法如下。

實(shí)際例子

在擴(kuò)散階段，去噪Transformer接收zt～qt(?∣x)，其中包含待去噪的掩碼token，以及目標(biāo)序列x。

從排列分布PL中采樣一個(gè)隨機(jī)排列σ，并滿足一個(gè)自然約束：在排列σ中，zt中的干凈token必須排在掩碼token之前。

下圖展示了一個(gè)示例的注意力掩碼及其排序?qū)崿F(xiàn)，其中x=(A,B,C,D,E,F)，zt=(A,M,C,M,M,F)，排列σ=(3,1,6,4,5,2)。

在順序階段，去噪Transformer接收z0⊕x∈V2L，其中z0～q0(?∣x)包含待去噪的掩碼token，并通過比較Transformer在z0上的輸出與目標(biāo)序列x來計(jì)算損失。

在訓(xùn)練過程中需要將z0與x進(jìn)行拼接作為輸入，這是因?yàn)椴幌馎R模型那樣在輸出端使用逐步移動(dòng)（shift-by-one）。

從排列分布PL中采樣一個(gè)隨機(jī)排列σ，該排列滿足以下兩個(gè)約束：

（i）σ中z0的未掩碼token排在掩碼token前；

（ii）掩碼token在σ中保持其自然順序。

下方展示了一個(gè)示例的注意力掩碼及其排序?qū)崿F(xiàn)，

其中x=(A,B,C,D,E,F)，z0=(A,M,C,M,M,F)，σ=(3,1,6,2,4,5)。

在順序生成過程中，模型需要從左到右地對由z0～pθMDM(?)生成的遮蔽token進(jìn)行去噪。

圖2：擴(kuò)散階段訓(xùn)練中注意力偏置的比較。橙色代表0（有注意力），灰色代表?∞（無注意力）

干凈的原始序列為x=(A,B,C,D,E,F)。

經(jīng)過隨機(jī)遮蔽后，得到zt=(A,M,C,M,M,F)。

圖中整數(shù)表示位置索引，其中遮蔽token的索引集為M(zt)={2,4,5}，干凈token的索引集為C(zt)={1,3,6}。

隨機(jī)順序?yàn)棣?(3,1,6,4,5,2)～P6，其中干凈token出現(xiàn)在遮蔽token之前。

混合訓(xùn)練

設(shè)x～qdata(x)為數(shù)據(jù)分布中的樣本，pθ是由參數(shù)θ定義的模型分布。

ESO-LM將模型分布pθ分解為兩部分：自回歸模型（Autoregressive Model, AR）和掩碼擴(kuò)散模型（Masked Diffusion Model, MDM）。

具體生成過程為：首先，掩碼擴(kuò)散模型生成一個(gè)部分掩碼的序列，然后自回歸模型以從左到右的方式完成剩余的解掩碼步驟，生成條件分布。

這一混合生成過程的邊緣似然表示為：

雖然上述求和難以直接計(jì)算，但可以通過引入后驗(yàn)分布q(z0∣x)來對真實(shí)似然進(jìn)行變分下界估計(jì)。

由于建模的是掩碼序列，可以選擇一個(gè)簡單的掩碼分布q，具體定義如下：，即以概率1?α0獨(dú)立掩碼每個(gè)token，其中α0∈[0,1]。

由此推導(dǎo)得到變分下界：

在原文附錄中，研究者分析了KL項(xiàng)并給出負(fù)證據(jù)下界（NELBO）：

當(dāng)α?=1時(shí)，后驗(yàn)采樣z?=x，所有token均由MDM（掩碼擴(kuò)散模型）生成，此時(shí)上式負(fù)證據(jù)下界中的AR損失為零，NELBO（負(fù)證據(jù)下界）退化為純MDM損失。

反之，當(dāng)α?=0時(shí)，所有token均被掩碼，MDM損失消失，NELBO退化為純AR（自回歸）損失。

因此，ESO-LM通過超參數(shù)α?的調(diào)控，實(shí)現(xiàn)了自回歸（AR）與掩碼擴(kuò)散（MDM）兩種生成范式的平滑插值。

這能夠在兩種風(fēng)格之間流暢切換，實(shí)現(xiàn)以下方面的完美平衡：本通順度、生成質(zhì)量和推理速度。

速度與質(zhì)量的完美平衡

Eso-LM模型在兩個(gè)標(biāo)準(zhǔn)語言建模基準(zhǔn)上進(jìn)行了評(píng)估：十億詞數(shù)據(jù)集（LM1B）和OpenWebText（OWT）。

所有模型均采用提出的基于擴(kuò)散Transformer的架構(gòu)，并引入旋轉(zhuǎn)位置編碼。

實(shí)驗(yàn)結(jié)果表明，在LM1B和OWT基準(zhǔn)測試中，Eso-LM模型實(shí)現(xiàn)了擴(kuò)散模型的最優(yōu)困惑度表現(xiàn)，同時(shí)在掩碼擴(kuò)散模型(MDM)與自回歸模型(AR)的困惑度區(qū)間實(shí)現(xiàn)了更精細(xì)的插值調(diào)控(見表1和表2)。

具體而言：

1. 性能突破：在LM1B上，Eso-LM將擴(kuò)散模型的困惑度記錄從18.7顯著降低至16.3，相對提升達(dá)13%；
2. 動(dòng)態(tài)調(diào)控：通過調(diào)節(jié)擴(kuò)散步數(shù)(T=10至T=1000)，模型可平滑過渡生成質(zhì)量與速度，相鄰步長困惑度差異保持在0.8以內(nèi)；
3. 長程優(yōu)勢：在OpenWebText(OWT)長文本評(píng)估中，1024上下文窗口下的困惑度從21.5優(yōu)化至19.1，驗(yàn)證了模型對長距離依賴的有效建模；
4. 評(píng)估嚴(yán)謹(jǐn)：采用序列打包技術(shù)使LM1B評(píng)估更具挑戰(zhàn)性(基準(zhǔn)困惑度提升2.1)，但模型仍保持12-15%的相對性能優(yōu)勢。

當(dāng)生成長度為8192的序列，并使用最大數(shù)量的函數(shù)評(píng)估（NFEs=8192）時(shí)，Eso-LM模型的推理速度最多比MDLM快65倍，比BD3-LMs快3～4倍。

對在OWT數(shù)據(jù)集上訓(xùn)練的模型，研究者使用生成困惑度（Generative Perplexity，Gen. PPL）來評(píng)估所生成樣本的質(zhì)量。

Gen. PPL越低，表示生成質(zhì)量越高。

為比較采樣效率，研究者還記錄了每種方法生成一個(gè)樣本（即batch size=1）所需的采樣時(shí)間中位數(shù)（單位為秒，基于5次試驗(yàn)）。

Eso-LM模型在采樣速度–質(zhì)量的帕累托前沿（Pareto frontier）上達(dá)到了新的SOTA（最先進(jìn)水平），重新定義了生成模型的可能性：

• 在高速采樣條件下實(shí)現(xiàn)與MDLM相當(dāng)?shù)?/span>困惑度；
• 在需要時(shí)，可達(dá)到與自回歸模型（AR）相同的困惑度水平；
• 在采樣步驟較少時(shí)不會(huì)出現(xiàn)模式崩潰（mode collapse）——這是Block Diffusion 模型所無法做到的。