AI自我糾錯，Diffusion超越自回歸，質量提升55%，已達理論證據下界

Diffusion模型，學會了自我糾正！無需強化學習等後訓練，擴散在語言建模中實現了自我糾錯，達到了計算效率匹配的最優性能，找到了證據下界（ELBO）的理論上的閉式解，在實驗中將樣本質量最高提升了55%。

如果大語言模型（LLMs）能夠發現並糾正自己的錯誤，那豈不是很好？

而且，如果能夠直接從預訓練中實現這一點，而無需任何監督微調（SFT）或強化學習（RL），那會怎樣呢？

最新提出的離散擴散模型，稱為GIDD，它能夠做到這一點。

在語言建模中，GIDD實現了計算效率匹配的最優性能！

請注意：不是基於自回歸的LLM，是擴散語言模型。

來自蘇黎世聯邦理工學院ETH Zurich等組織的研究團隊，推廣了掩碼擴散（masked diffusion），並推導出一系列廣義插值離散擴散模型（general interpolating discrete diffusion，GIDD）的理論基礎。

GIDD不僅更靈活，而且在理論上得到了證據下界（evidence lower bound，ELBO）的閉式解。

實驗結果表明：GIDD樣本質量（生成困惑度，PPL）最高可提升55%。

GIDD核心主要有3大特徵：

1、推廣性強：GIDD適用於多種擴散過程，不局限於掩碼擴散。

2、混合擴散訓練：訓練了一個結合掩碼（masking）和均勻噪聲（uniform noise）的混合擴散模型。

3、雙重能力：不僅可以填補空缺（填充被掩蓋的token），還可以評估已填充token的正確性，並在必要時用更合理的 token替換錯誤的部分。

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

項目地址：https://github.com/dvruette/gidd/

預測下一個token，雖然成果顯著，但存在固有的局限性，例如無法修改已經生成的token。

這促使研究者探索替代方法，如離散擴散。

然而，由於簡單性和有效性，掩碼擴散（masked diffusion）成為流行選擇，但重新引入了無法修改token的這種局限性。

基於對擴散模型添加噪聲的重要性認識，新研究旨在探索離散擴散模型的設計空間，並嘗試不同的擴散過程。

廣義插值離散擴散（GIDD）是新的離散擴散方法，將掩碼擴散推廣到任意插值噪聲過程。

在擴散過程中，GIDD可以在任何時刻靈活地選擇添加不種類型的噪聲。

有趣的是，任何邊際分佈符合上述方程的擴散過程，都可以推導出證據下界（ELBO）。

想要掩碼擴散？沒問題！

只需將 πt設為 one_hot([MASK])，然後在GIDD的ELBO上進行訓練即可。

不幸的是，掩碼擴散無法進行自我糾正。就像自回歸模型一樣，一旦token被確定，就無法再更改。

因此，如果模型在某個步驟出錯，就沒有辦法進行修正。

為瞭解決這個問題，從BERT中汲取了靈感：如果除了掩碼token外，隨機用其他token替換一部分token，會怎麼樣？

這樣，模型不僅需要學會「填空」，還要學會識別並修正錯誤的token。

由於GIDD的ELBO具有高度的靈活性，只需選擇合適的 πt來捕捉所需的動態特性。

設定 πt使得均勻噪聲的比例隨著噪聲水平的變化而上升或下降，並在 t=0.5時達到峰值。

真地有效嗎？

第一次實驗結果有些讓人失望，因為使用均勻噪聲訓練的模型，其困惑度比僅使用掩碼的模型要差。

但考慮到這些模型除了需要填補缺失的token外，還要修正錯誤的token，這並不令人感到意外。

那樣本質量又如何呢？

下圖發現發現生成性PPL（使用Gemma-2-9b測量）有了顯著的改進，特別是對於使用均勻噪聲訓練的模型來說更是如此，尤其是在推理預算緊張的情況下！

這是否意味著模型在糾正自己的錯誤呢？

為了弄清楚這一點，提出了一種自我糾正算法，通過一次修復一個token（使用模型）來改進已經生成的樣本，直到收斂到一個穩定點。

這樣可以不斷提高樣本質量（以生成PPL衡量），甚至超越僅僅通過增加去噪預算所能達到的水平，提供了一種簡單的方法來擴展測試時計算資源！

值得注意的是，僅在去掩碼任務上訓練的模型不具備這種能力。

在理論和實踐上，新研究的主要貢獻包括以下兩方面：

• 理論方面：將掩碼擴散（masked diffusion）框架擴展為「廣義插值離散擴散」過程。GIDD能更靈活地選擇噪聲添加方式，同時仍然能夠得到累積狀態轉移和擴散證據下界（ELBO）的閉式解。
• 實踐方面：基於理論分析，在實踐中取得了改進。在計算資源相當的情況下，優化訓練目標達到了當前最優的性能。同時，通過引入均勻噪聲，提高了樣本質量，並實現了自校正能力（見圖1和表1）。

表1.使用20%均勻噪聲訓練的GIDD+BASE模型進行自我糾正的例子（綠色替換紅色）。

GIDD模型能夠在未明確訓練的情況下糾正語法錯誤、改進詞彙選擇，甚至提升事實準確性。

自回歸模型缺陷

從觀察中提取世界的結構是智能的基本機制。

生物體可以自然而然地做到這一點，而機器在這方面的能力直到最近才取得重大突破。

近年來，深度生成模型取得了顯著提升。

顧名思義，生成模型生成新的、逼真的樣本，其中「逼真」通常指的是樣本在某個參考分佈下具有較高的概率。

對於某些數據分佈（例如自然圖像或自然語言），單個樣本所包含的信息量可能非常龐大。

為了降低生成模型的計算負擔，一種常見策略是將生成單個樣本的任務拆分為多個推理步驟。

每個步驟本身相對簡單，但當所有步驟組合在一起時，能夠恢復完整的分佈。

在自然語言處理領域，最典型的方法是自回歸建模（autoregressive modeling）。

在此類方法中，生成一個句子（或序列）的過程被拆解為逐個生成單詞（或token），並使用已生成的單詞作為上下文，來預測下一個單詞。

自回歸模型存在一些天生的缺陷：（1）計算成本高；（2）長程依賴與連貫性（long-term dependencies and coherence）問題。

為瞭解決這些問題，研究者們付出了大量努力。

例如，通過強化學習（RL）進行後訓練，讓模型在多個自回歸步驟中學會序列推理，從而提升連貫性。

去噪擴散模型，提出了不同的生成任務分解方式，可以解決這兩個限制。

擴散模型的調用次數與樣本大小無關。

在圖像生成上，擴展模型取得了成功，其中添加的高斯噪聲。

然而，比如自然語言上，擴散方法並不如在圖像上如意。

廣義插值擴散（GIDD）

掩碼擴散（masked diffusion）技術雖廣泛應用，但仍存在根本性局限。

主要問題源於其底層Markov鏈設計：一旦token被填充，便無法修改，可能導致錯誤累積或token不兼容，且缺乏修正機制，直接影響結果質量。

另一個局限是，僅掩碼token貢獻損失信號，未掩碼token保持無噪聲狀態，減少了有效批大小。

一種有效解決方案是借鑒BERT，將掩碼機制與均勻噪聲結合。這種方法不僅能解決上述問題，還帶來額外優勢：

通過這種方式，模型學會識別「正確」與「錯誤」token後，可能發展出自我校正\自我糾錯能力。

然而，在特定理想擴散路徑上訓練擴散模型時，仍面臨技術挑戰，需進一步研究。

標準的訓練目標——擴散證據下界（diffusion ELBO）—— 需要已知Markov狀態轉移才能推導出來。

但構造具有特定性質的Markov鏈通常是複雜的逆問題，並不容易直接求解。

研究人員將插值擴散（interpolating diffusion）擴展到任意（隨時間變化的）插值策略，避免了單獨求解特定的掩碼與均勻噪聲組合的逆問題，同時提升了模型設計的靈活性，

具體來說，提出了廣義插值離散擴散（GIDD），是一類具有邊際前向轉移（marginal forward transitions）的擴散模型，其形式如下：

其中，π_t是一個隨時間平滑變化的概率分佈。

值得注意的是，當π_t=m時，GIDD便退化為掩碼擴散的特例。

可以證明，在適當選擇αt和πt的情況下，確實存在一個Markov鏈可以產生這些邊際分佈，並推導其條件轉移關係以及訓練所需的ELBO公式。

前向過程

GIDD旨在提供最大程度的靈活性，使得在任意時間點都可以對數據添加不同類型的噪聲。

其核心由兩個部分組成：

1）混合率αt：定義了隨時間變化的信噪比（signal-to-noise ratio）。

2）混合分佈πt：決定了數據在某一時刻被噪聲化後的目標分佈。

研究人員將這兩個函數的組合稱為擴散過程的「混合調度」（mixing schedule）。

定義 3.1（混合速率）：設（累積）混合速率αt和βt（其中βt=1−αt）為時間可微且遞減的函數αt:[0,1]→[0,1]，滿足初始條件α0=1（表示無混合狀態）和最終條件α1=0（表示完全混合狀態）。

這一設定決定了信噪比（SNR），即 SNR=αtβt。隨著t的增加，αt減小，表明信號成分逐漸減少，而噪聲成分逐漸增加。

定義 3.2（混合分佈）：設混合分佈πt是一個依賴於時間的概率向量，即時間可微函數 πt:[0,1]→Δ∣V∣−1，這裏 Δ∣V∣−1表示 ∣V∣維單純形。

混合分佈πt描述了在任意給定時間點 tt添加到數據中的噪聲類型。因此，π1特別地代表了擴散過程的先驗分佈，它刻畫了在時間 t=1時的數據噪聲特性。

在此過程中，研究人員已經成功構建了一個馬爾可夫鏈，其邊緣分佈按邊際前向轉移公式所述。

為了後續推導ELBO（證據下界），還需要定義相應的連續時間馬爾可夫鏈（CTMC）的轉移速率，具體如下。

最終，得到了GIDD的CTMC前向速率。

反向過程

擴散模型的標準分佈pθ(zs∣zt)由以下公式給出：

其中，qt(zt∣xθ)的簡寫形式為：

這裏，xθ(Zt,t)是一個神經網絡，用於預測在噪聲序列Zt條件下的x的分佈。

此外，ELBO（證據下界）的推導涉及連續時間馬爾可夫鏈（CTMC）的反向速率

GIDD的證據下界

為了訓練GIDD模型，需要一種可微分的方法來估計其似然函數（likelihood）。

證據下界（ELBO）正是用於此目的：通過最大化ELBO，實際上也是在最大化模型的（最壞情況下的）似然函數。

在計算ELBO時，需要用到GIDD的前向速率（forward rate）和反向速率（backward rate），這些已在前文推導完成。

然後，基於Campbell等人提出的ELBO形式進行一定修改，並將GIDD的前向速率和反向速率代入，經過化簡後得到了定理3.7。

深入分析GIDD的ELBO（證據下界），可以發現它實際上是在同時優化兩個任務：

1、匹配模型與前向分佈的邊際概率

2、最小化rθ(zt,x)以匹配邊際概率

有趣的是，這兩個優化任務具有相同的全局最優解。

這意味著，如果模型能夠完美優化ELBO，它就能同時滿足這兩個目標。這一特性對於理解ELBO的全局最小值及其優化過程具有重要意義。

采樣

給定一個采樣時間表0≈t0

具體來說，從一個全掩碼token的序列開始，即所有ztT都設為掩碼tokenm。

然後，按照以下方式迭代采樣，其中i=T,…,1：

自校正步驟（Self-Correction Step）

此外，提出了一種不動點迭代方法，通過重新采樣部分token來改進生成結果，使其更符合模型的判斷。

具體而言，將完全去噪後的樣本Z_{t_0}輸入模型，並以溫度參數τ進行采樣。

然後，在所有與Zt0不同的采樣token中，

選擇模型置信度最高的一個token並確認它。

這個過程會持續進行，直到結果收斂（詳細內容見附錄C）。

自校正算法是一種不動點迭代方法，可以應用於任何已經（部分）去噪的生成樣本。

其核心思想是查詢模型以識別模型認為錯誤並應該替換的token，並且一次只替換一個token以避免重新引入衝突token。

具體的偽代碼實現參見算法1。

在實際操作中，發現收斂往往表現為在兩個或多個同樣優良狀態（就自準確性而言）之間的振盪，因此額外基於自準確性實現了提前停止機制。

實驗表明，提前停止的耐性值設為32時效果良好。

混合策略（Mixing Schedule）

雖然GIDD可以用於掩碼擴散，但最初提出這一廣義框架的動機是探索掩碼與均勻噪聲的結合。

為此，研究團隊設計了一種混合策略（mixing schedule），在保持掩碼先驗分佈的同時，允許在不同階段引入可調節比例的均勻噪聲。

下文中用p_u表示均勻噪聲的比例。

為了保證可解釋性，設定在數據和噪聲的中點（t=1/2）時，均勻噪聲token的期望比例達到最大值p_u。

基於這一目標，定義了混合速率（mixing rate）和混合分佈（mixing distribution）（定義3.1和3.2）。

其中：

， N表示詞彙表的大小，B是一個常數，選取它的值可以保證均勻噪聲的比例達到目標水平。

由此，得到邊際前向分佈（marginal forward distribution）：

為了在t=1/2處使均勻噪聲比例達到p_u，需要設定：

GIDD的ELBO計算涉及額外的常數和因子，需要推導相應的時間導數。

值得注意的是，當p_u=0.0時，GIDD退化回掩碼擴散（masked diffusion）。

在的實驗中，設定γ=1，但本節引入的超參數仍有許多其他可能的選擇。

訓練目標

在開始實驗之前，需要解決最後一個關鍵問題，這一改進將帶來顯著的性能提升。

仔細分析擴散證據下界（ELBO）的權重wt(zt,x)後，發現當t→0或t→1時，權重的變化非常極端。

考慮以下三種可能情況：（1） z_t=x（未被噪聲汙染的token）；（2）z_t=m（掩碼token）；（3）zt∉{x,m}（其他隨機噪聲token）

繪製w_t(z_t,x)隨時間的變化曲線後（見圖2），可以觀察到在極低或極高噪聲水平下，權重呈指數級增長。

這種現象可能會帶來問題：

• 當噪聲較低時，去噪任務變得過於簡單，模型幾乎不需要學習。
• 當噪聲較高時，去噪任務變得幾乎不可能，模型也無法從中獲得有效的訓練信號。
• 這兩種極端樣本的權重過高，可能會淹沒批次中的其他訓練樣本，從而影響整體訓練效果。

解決方案：權重裁剪（Weight Clamping）

為瞭解決這個問題，提出了兩種權重調整方案，以減少極端樣本的影響，並強調中等噪聲水平的樣本，因為這些樣本提供了最有價值的訓練信號。

最簡單直接的方法是對權重設置一個最大值w_{max}，即：

根據初步實驗，發現設定 wmax⁡=1效果最佳，因此這一設定將在後續實驗中被採用。

表3：GIDD（p_u = 0.0）和MDM的困惑度（PPL）非常接近，這與它們的理論等價性一致。

顯著的效果提升來自於選擇正確的權重函數，尤其是在p_u > 0的情況下。

最終的最佳設置包括動態損失權重和權重衰減，該設置也被稱為 GIDD+。

上述權重裁剪（clamping）方法主要影響掩碼token和均勻噪聲token的權重。

然而，一個更系統的方法是：在保持最大損失權重恒定的同時，仍然保留掩碼token、均勻噪聲token和無噪聲token之間的相對權重關係。

動態權重調整（Dynamic Weighting）

提出了一種動態加權函數（dynamic weighting function），其定義如下：

其中，

表示對數信噪比（log-SNR）。

該方法的相對權重關係（掩碼token/均勻噪聲token/無噪聲token=2/1/Be^{-λ_t^2}）是通過實驗經驗確定的。

需要注意的是，這種ELBO重加權方法等效於從一個非均勻分佈中采樣t。

自我糾正（Self-Correction）

目前為止，觀察到僅使用掩碼訓練的模型往往優於結合均勻噪聲的模型，但尚未討論引入均勻噪聲的核心動機：讓模型學會區分「正確」與「錯誤」token，希望它能夠具備自我糾正能力。

為了評估生成樣本的質量，採用生成困惑度（generative perplexity,PPL）這一指標。

具體而言，PPL計算的是生成樣本在更強大模型下的似然值，在的實驗中，使用Gemma 2 9B作為評估模型。

更高的似然值通常被認為對應於更高質量的樣本。

雖然PPL作為指標存在諸多局限性，但它在文獻中被廣泛採用，並且在相對比較不同模型的質量時依然具有參考價值。

除了PPL，還評估了模型的自我準確率（Self-Accuracy）。

即模型在生成過程中，對其認為「正確」的token（即在整個序列中賦予某個token最高概率）所佔的比例。

值得注意的是，在進行自我糾正之前，訓練時加入均勻噪聲的模型樣本質量已經更高。

尤其是在低計算量推理（low inference-compute）設置下，相較於僅使用掩碼的模型，其生成困惑度（generative PPL）提升尤為顯著。

例如，在32步推理時：

• GIDD+（SMALL;p_u=0.1）的PPL為387;
• 僅掩碼模型（p_u=0.0）的PPL為904;
• MDM（masked diffusion model）的PPL更高，達到1302。

這表明，訓練時加入均勻噪聲可以穩定生成過程，特別是在模型將自身生成的輸出作為輸入時，使得樣本質量更高——即使其驗證困惑度（validation PPL）略有下降。

這一發現引發了一個重要問題： 自我糾正的效果是否只是因為額外的去噪迭代次數 ？

換句話說，模型在去噪過程中可能已經在執行一定程度的自我糾正，因此自我糾正步驟帶來的提升，是否僅僅是因為額外的計算次數？

從實驗來看，雖然增加去噪步驟確實會單調提升樣本質量，但這種提升最終會趨於飽和（對於BASE模型，PPL約停留在200）。

然而，當引入自我糾正機制後，PPL可以進一步降低到100以下，這表明自我糾正帶來的改進並非僅僅是更多去噪迭代的結果，而是一種額外的、非平凡（non-trivial）的提升。

圖3|從左到右（a）不同溫度下token數變化；（b）token變化數與PPL的關係；（c）自我準確性與生成困惑度之間的相關性.

具體結果如下：

(a) 使用 GIDD+ (BASE) 模型進行自我糾正時，可以重新采樣最多10%的token，這一過程與均勻噪聲水平無關。研究發現，溫度參數τ取值在 [0.1, 0.5] 之間時效果最佳。

(b) 對於在均勻噪聲上訓練的模型來說，采樣的token數量越多，效果越好。然而，僅使用掩碼（mask-only）的模型即使重新采樣了同樣數量的token，也無法提升質量。

(c) 通過分析自我準確性（self-accuracy）與生成困惑度（generative PPL）之間的相關性，發現混合模型在評估自身生成樣本的質量方面明顯更具優勢。

下遊人物性能

在一系列基準測試中評估了模型的語言理解能力。基於混合噪聲設置p_u > 0的更高難度，預計這些模型不會超過僅使用掩碼噪聲p_u = 0的情況，實驗結果也證實了這一點。

因此，研究人員將重點放在比較最佳的 SMALL GIDD+ 模型與MDM以及自回歸基線模型（即 GPT2（和重新訓練的 Llama上。

作為參考，還納入了兩個 1.1B 參數模型，一個是自回歸模型，另一個是掩碼擴散模型。

基準測試包括ARC-e和ARC-c、BoolQ、Hellaswag、PIQA、OpenBookQA以及WinoGrande。

實驗發現，平均準確率與驗證困惑度（PPL）通常有很好的相關性（見表4）。

在擴散模型中，表現最好的是僅訓練了131B token的GIDD+（p_u=0.0），超過了訓練時間兩倍的模型。

這可能是由於模型在訓練數據中的虛假模式上過擬合，雖然驗證損失仍然下降，但並未轉化為下遊任務的性能提升。

值得注意的是，最佳擴散模型GIDD+的表現優於自回歸模型 GPT2，儘管訓練數據的差異使得公平比較有些困難。實際上，最佳自回歸模型 Llama（重新訓練版本） 仍然在總體上表現最佳，但平均差距不到一個百分點。

表4：不同模型的零樣本(Zero-shot)基準準確率。在小型模型和擴散模型中分別用粗體和下劃線_標出了最佳分數。

對於使用均勻噪聲訓練的 GIDD 模型，其趨勢與驗證困惑度一致，更多的均勻噪聲通常會降低準確率。

下表5中列出了三種規模（TINY、SMALL、BASE）和所有均勻噪聲水 0.0, 0.1, 0.2的GIDD+模型的基準測試準確率。

無論均勻噪聲水平如何，模型的性能都隨著規模的增加而持續提升。然而，使用均勻噪聲訓練的模型，在性能上略微但持續落後於僅使用掩碼噪聲的模型。

然而，隨著模型規模的增加，性能持續提升，初步跡象表明，隨著規模的擴大，差距可能會縮小。

直觀地說，均勻噪聲讓訓練任務變得更難：模型不能再理所當然地認為每個未掩碼的token都是正確的，而是必須考慮上下文中的每個token，並在必要時將其替換為正確的token。

這種直觀的解釋表明，觀察到的性能差異可能是由於模型容量不足，在這種情況下，預計更大的模型受均勻噪聲的影響會更小。

為了驗證這一假設，在保持訓練時長不變的情況下擴展了參數數量，並訓練了不同規模（TINY、SMALL 和 BASE）的模型，分別在不同的均勻噪聲水平0.0, 0.1, 0.2下進行訓練。

然後，通過指數擬合繪製了計算效率前沿，反映帕累托最優的驗證 ELBO（見圖 4）。

One More Thing：只是小規模實驗

由於資源限制，實驗設置有一定的局限性：

每個噪聲水平的樣本量僅限於三種不同的計算預算，其中最大的計算預算仍然相對較小，僅為3.3*10^{20} FLOPs。

作為參考，現代大語言模型的許多標誌性能力通常需要達到 10^{22} FLOPs 左右才會開始顯現，這仍然比我們最大的計算預算高出兩個數量級。

儘管如此，確實觀察到了一個一致的趨勢，即隨著計算資源的增加，較高水平的均勻噪聲表現更好，儘管這一趨勢的幅度較小。

在僅使用掩碼噪聲的設置p_u = 0.0中，擴展指數為-0.0586，而加入均勻噪聲後，p_u = 0.1和p_u = 0.2 的擴展指數分別提高到-0.0589和-0.0621。

外推這一趨勢預測，p_u = 0.2的設置將在10^{21}FLOPs 左右超過 p_u = 0.0，這一計算預算在中到大規模的訓練中通常可以達到。

然而，必須強調的是，這次的實驗設置的局限性使得這種預測的可靠性較低。

儘管如此，觀察到的擴展行為是令人鼓舞的，值得進一步研究。

參考資料：

https://www.arxiv.org/abs/2503.04482

本文來自微信公眾號「新智元」，作者：KingHZ，36氪經授權發佈。