全球首次,Transformer「混血」速度狂飆 65 倍,英偉達已下注
康奈爾等提出 Eso-LM 模型,擴散+自回歸混合,推理快 65 倍。
擴散建模+自回歸,打通文本生成任督二脈! 這一次,來自康奈爾、CMU 等機構的研究者,提出了前所未有的「混合體」——Eso-LM。 有人驚呼:「自回歸危險了。」
擴散方法打入語言模型領域!
最近,康奈爾博士生 Subham Sahoo,在 X 介紹了擴散大語言模型的最新工作。
這項研究引發了 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。
研究者還發現,BD3-LMs 在低採樣步數下性能下降,而新方法在低計算量(NFE)場景下與離散擴散模型相當,在高計算量場景下與自回歸模型相當。
這次的結果為離散擴散模型建立了新的困惑度(perplexity)最優水平,縮小了與自回歸模型的差距。
另外值得一提的是,除了共同一作 Zhihan Yang 外,還有多位華人作者,其中包括知名華人學者邢波(Eric Xing)。
語言也能擴散
這並非擴散方法首次「入侵」文本生成領域。
甚至達到商用級別的擴散語言模型,都不止一個。
斯坦福、UCLA 和康奈爾的三位教授聯合創立了 Inception Labs,推出了全球首個商用級別的擴散語言模型
擴散語言模型最大特點就是快:推理速度可達 ChatGPT 的 6 倍!
IBM 甚至認為擴散模型就是下一代 AI,GPT 這類自回歸範式受到有力挑戰。
不過,三位教授具體如何實現這一突破,目前尚屬商業機密,外界難以得知。
而在 AI 巨頭中,谷歌是第一家嘗試擴散語言模型——
在 I/O 大會上,它放出了實驗版語言模型 Gemini Diffusion:推理每秒可達 1400 多 token。
而這次的新論文,作者 Arash Vahdat 是英偉達研究院的科研總監(Research Director),領導基礎生成式人工智慧(GenAIR)團隊。
莫非英偉達也要押注擴散語言模型?
擴散模型:後來者居上?
眾所周知,掩蔽擴散模型(Masked Diffusion Models,MDMs)是自回歸(AR)語言模型的有力替代方案——
但它們有兩個致命短板:
速度慢 :沒有 KV 緩存 = 實際上比 AR 慢得多;
品質差 :在複雜任務中表現不佳,似然度低於 AR。
塊擴散(Block Diffusion)模型 BD3-LM,在每個區塊內執行擴散過程,以先前區塊為條件,實現分塊生成 token 序列。
它融合了自回歸模型與擴散模型的優勢:在支援可變長度生成的同時,利用 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 利用創新的混合訓練方法,將 KV 緩存引入了擴散模型。
具體來說:
混合訓練 :Eso-LM 在訓練時一半數據採用 AR 風格(乾淨的上下文預測下一個單詞),另一半採用擴散風格(打亂輸入,部分掩碼,逐步去噪)。
推理優化 :在生成過程中,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 的索引集合。
而且不再處理整個序列,而只對子序列{ztl∣l∈C(zt)∪St}進行前向傳播 ——即,乾淨的 token 和計劃去噪的 token——
這在處理長序列時顯著降低了計算量 。
這種方法支援在擴散過程中進行高效的 KV 緩存 。
關鍵思想借用了已有的方法 AO-ARM(見下文):遮蔽 token 可以按任何順序揭示。
論文連結:https://openreview.net/forum?id=sMyXP8Tanm
因此,在訓練過程中,新方法要採樣隨機順序 σ∼PL,並對每個σ,強制執行對遮蔽 token 的因果注意力。
具體來說, 要求遮蔽 token 只能對乾淨 token 和根據順序σ排列的先前遮蔽 token 進行注意力計算 。
Eso-LM(A)採用了這一策略,在採樣過程中顯著減少了計算量,同時保持了性能。
而 Eso-LM(B)對乾淨 token 強制施加類似的因果 mask,進一步擴展了這一思想,從而實現了 KV 緩存。
儘管在困惑度上稍微差一些,Eso-LM(B)在採樣過程中提供了顯著的加速(最多 65 倍)。
順序階段
自回歸模型隨後從左到右填充遮蔽 token,使用順序去噪計劃 ,其中要求每個單元素集合按其唯一元素升序排列。
不同於標準的自回歸解碼,每個 x~l 同時依賴其左側上下文(完全由乾淨 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 模型實現了擴散模型的最優困惑度表現,同時在掩碼擴散模型(MDM)與自回歸模型(AR)的困惑度區間實現了更精細的插值調控(見表 1 和表 2)。
具體而言:
性能突破:在 LM1B 上,Eso-LM 將擴散模型的困惑度記錄從 18.7 顯著降低至 16.3,相對提升達 13%;
動態調控:通過調節擴散步數(T=10 至 T=1000),模型可平滑過渡生成品質與速度,相鄰步長困惑度差異保持在 0.8 以內;
長程優勢:在 OpenWebText(OWT)長文本評估中,1024 上下文視窗下的困惑度從 21.5 優化至 19.1,驗證了模型對長距離依賴的有效建模;
評估嚴謹:採用序列打包技術使 LM1B 評估更具挑戰性(基準困惑度提升 2.1),但模型仍保持 12-15%的相對性能優勢。
當生成長度為 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
本文來自微信公眾號 “新智元”,作者:KingHZ,36 氪經授權發佈。