新智元報道

　　編輯：KingHZ

　　【新智元導讀】擴散建模+自回歸，打通文本生成任督二脈！這一次，來自康奈爾、CMU等機構的研究者，提出了前所未有的「混合體」——Eso-LM。有人驚呼：「自回歸危險了。」

　　擴散方法打入語言模型領域！

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

　　這項研究引發(fā)了AI研究領域的思考。

　　英偉達研究院杰出研究科學家Pavlo Molchanov說：「擴散大語言模型正在崛起！」

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

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

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

　　項目鏈接：https://s-sahoo.com/Eso-LMs/

　　這是首個在保持并行生成的同時，引入KV緩存機制的方法。

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

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

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

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

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

　　語言也能擴散

　　這并非擴散方法首次「入侵」文本生成領域。

　　甚至達到商用級別的擴散語言模型，都不止一個。

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

　　擴散語言模型最大特點就是快：推理速度可達ChatGPT的6倍！

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

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

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

　　而這次的新論文，作者Arash Vahdat是英偉達研究院的科研總監(jiān)（Research Director），領導基礎生成式人工智能（GenAIR）團隊。

　　莫非英偉達也要押注擴散語言模型？

　　擴散模型：后來者居上？

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

　　但它們有兩個致命短板：

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

　　質量差：在復雜任務中表現不佳，似然度低于AR。

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

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

　　既能實現更高質量的生成效果，又能保持高效推理特性。

　　但是，BD3-LM的速度與質量仍需權衡：

　　低采樣步數下出現模式崩塌，導致樣本質量差；

　　而且只支持部分緩存，塊內鍵值緩存仍缺失。

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

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

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

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

　　正如圖1所示，Eso-LM包含擴散和順序兩個階段：

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

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

　　與BD3-LM不同，Eso-LM (B)允許在兩個階段使用統一的KV緩存，藍色邊框框住了正在構建KV緩存的transformer單元；當單元的KV緩存構建完成時，該單元變成藍色。

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

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

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

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

　　具體來說：

　　混合訓練：Eso-LM在訓練時一半數據采用AR風格（干凈的上下文預測下一個單詞），另一半采用擴散風格（打亂輸入，部分掩碼，逐步去噪）。

　　推理優(yōu)化：在生成過程中，Eso-LM只對部分單詞（掩碼和干凈單詞）進行前向計算，并緩存干凈單詞的KV對，大幅減少計算量。

　　爆改Transformer

　　靈活切換注意力

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

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

　　研究者引入了注意力偏置矩陣A，調整標準的自注意力機制：

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

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

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

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

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

　　Eso-LM(B)進一步擴展了這個想法，不僅對遮蔽token應用因果mask，還對干凈token應用，從而實現更高效的KV緩存（KV-caching）——代價是困惑度略有下降。

　　擴散階段

　　在擴散階段，標準的采樣方法會浪費大量FLOPs。

　　為了提高效率，研究者對標準采樣和訓練過程提出了兩個關鍵改進。

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

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

　　這在處理長序列時顯著降低了計算量。

　　這種方法支持在擴散過程中進行高效的KV緩存。

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

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

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

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

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

　　而Eso-LM(B)對干凈token強制施加類似的因果mask，進一步擴展了這一思想，從而實現了KV緩存。

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

　　順序階段

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

　　不同于標準的自回歸解碼，每個x~?同時依賴其左側上下文（完全由干凈token構成）和右側干凈的token，從而實現更豐富的生成。

　　我們跳過對右側遮蔽token的評估，減少不必要的計算。

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

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

　　當α0=1時，所有token都由擴散生成，因此S=S_MDM，且S_AR=?；

　　當α0=0時，所有token都由順序方式生成，因此S=S_AR，且S_MDM=?。

　　完整采樣算法如下。

　　實際例子

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

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

　　下圖展示了一個示例的注意力掩碼及其排序實現，其中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上的輸出與目標序列x來計算損失。

　　在訓練過程中需要將z0與x進行拼接作為輸入，這是因為不像AR模型那樣在輸出端使用逐步移動（shift-by-one）。

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

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

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

　　下方展示了一個示例的注意力掩碼及其排序實現，

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

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

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

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

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

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

　　隨機順序為σ=(3,1,6,4,5,2)～P6，其中干凈token出現在遮蔽token之前。

　　混合訓練

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

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

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

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

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

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

　　由此推導得到變分下界：

　　在原文附錄中，研究者分析了KL項并給出負證據下界（NELBO）：

　　當α?=1時，后驗采樣z?=x，所有token均由MDM（掩碼擴散模型）生成，此時上式負證據下界中的AR損失為零，NELBO（負證據下界）退化為純MDM損失。

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

　　因此，ESO-LM通過超參數α?的調控，實現了自回歸（AR）與掩碼擴散（MDM）兩種生成范式的平滑插值。

　　這能夠在兩種風格之間流暢切換，實現以下方面的完美平衡：本通順度、生成質量和推理速度。

　　速度與質量的完美平衡

　　Eso-LM模型在兩個標準語言建模基準上進行了評估：十億詞數據集（LM1B）和OpenWebText（OWT）。

　　所有模型均采用提出的基于擴散Transformer的架構，并引入旋轉位置編碼。

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

　　具體而言：

　　性能突破：在LM1B上，Eso-LM將擴散模型的困惑度記錄從18.7顯著降低至16.3，相對提升達13%；

　　動態(tài)調控：通過調節(jié)擴散步數(T=10至T=1000)，模型可平滑過渡生成質量與速度，相鄰步長困惑度差異保持在0.8以內；

　　長程優(yōu)勢：在OpenWebText(OWT)長文本評估中，1024上下文窗口下的困惑度從21.5優(yōu)化至19.1，驗證了模型對長距離依賴的有效建模；

　　評估嚴謹：采用序列打包技術使LM1B評估更具挑戰(zhàn)性(基準困惑度提升2.1)，但模型仍保持12-15%的相對性能優(yōu)勢。

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

　　對在OWT數據集上訓練的模型，研究者使用生成困惑度（Generative Perplexity，Gen. PPL）來評估所生成樣本的質量。

　　Gen. PPL越低，表示生成質量越高。

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

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

　　在高速采樣條件下實現與MDLM相當的困惑度；

　　在需要時，可達到與自回歸模型（AR）相同的困惑度水平；

　　在采樣步驟較少時不會出現模式崩潰（mode collapse）——這是Block Diffusion 模型所無法做到的。

　　參考資料：

　　https://x.com/ssahoo_/status/1929765494460239933

　　https://x.com/PavloMolchanov/status/1929944952848691309

　　https://x.com/ssahoo_/status/1929945984588755180