去年初，OpenAI 的視頻生成模型 Sora 帶火了擴散模型。

如今，擴散模型被廣泛用于生成圖像和視頻，并在生成文本或生物序列等離散數(shù)據(jù)方面變得越來越有效。從技術(shù)上講，與自回歸模型相比，擴散模型具有加速生成和提高模型輸出可控性的潛力。 

目前，離散擴散模型目前面臨至少三個限制。首先，在聊天系統(tǒng)等應(yīng)用中，模型必須生成任意長度的輸出序列（例如對用戶問題的回答）。但是，大多數(shù)最新的擴散架構(gòu)僅能生成固定長度的向量。其次，離散擴散模型在生成過程中使用雙向上下文，因此無法使用 KV 緩存重用以前的計算，這會降低推理效率。第三，以困惑度等標準指標衡量的離散擴散模型，質(zhì)量落后于自回歸方法，進一步限制了其適用性。

本文中，來自 Cornell Tech、斯坦福大學、Cohere 的研究者提出通過塊離散去噪擴散語言模型（Block Discrete Denoising Diffusion Language Models，BD3-LMs）來解決以上限制，該模型在擴散和自回歸模型之間進行插值。

具體來講，塊擴散模型（也是半自回歸模型）定義了離散隨機變量塊的自回歸概率分布，而給定先前塊的條件概率由離散去噪擴散模型指定。

• 論文標題：Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models
• 論文地址：https://arxiv.org/pdf/2503.09573
• 項目主頁：https://m-arriola.com/bd3lms/

下圖為 Block Diffusion 與自回歸、擴散模型的生成效果對比：

研究者表示，開發(fā)有效的 BD3-LM 面臨以下兩個挑戰(zhàn)：一是使用神經(jīng)網(wǎng)絡(luò)的一次標準前向傳遞無法有效地計算塊擴散模型的訓練目標，需要開發(fā)專門的算法。二是擴散目標梯度的高方差阻礙了訓練，導致 BD3-LM 即使在塊大小為 1 的情況下（當兩個模型等效時）也表現(xiàn)不佳。

因此，研究者推導出梯度方差的估計量，并證明它是自回歸和擴散之間困惑度差距的關(guān)鍵因素。然后，他們提出了自定義噪聲過程，以實現(xiàn)最小化梯度方差并進一步縮小困惑度差距。

實驗部分，研究者在多個語言建模基準上評估了 BD3-LM，并證明它們能夠生成任意長度的序列，包括超出其訓練上下文的長度。此外，BD3-LM 在離散擴散模型中實現(xiàn)了新的 SOTA 困惑度。與對嵌入進行高斯擴散的替代半自回歸方法相比，本文離散方法實現(xiàn)了易于處理的似然估計，并在少一個數(shù)量級生成步驟的情況下，生成的樣本在困惑度方面得到了改進。

論文一作 Marianne Arriola 發(fā)推稱，擴散語言模型在并行文本生成領(lǐng)域正在崛起，但與自回歸模型相比，它們存在質(zhì)量、固定長度限制和缺乏 KV 緩存等問題。本文 Block Diffusion 將自回歸和擴散模型結(jié)合了起來，實現(xiàn)了兩全其美。

BD3-LMs 模型概覽

研究者結(jié)合建模范式，從自回歸模型中獲得更好的似然估計和靈活的長度生成，并從擴散模型中獲得了快速的并行生成效果。

塊擴散似然

研究者提出了一個建模框架，該框架對 token 塊進行自回歸建模，并在每個塊內(nèi)執(zhí)行擴散操作。他們對長度為 L′ 的 B 個塊進行似然分解，如下所示：

每個 pθ(x^b|x^<b) 都使用包含 L′個 token 的塊上的離散擴散 ELBO 進行建模，并通過優(yōu)化以下似然邊界來獲得原則性學習目標 L_BD (x,θ)：

研究者使用簡單的離散擴散參數(shù)化對每個塊的似然進行建模，最終目標是對交叉熵項進行加權(quán)總和：

高效的訓練與采樣算法

簡單來說，研究者想要通過在一個 loop 中應(yīng)用B 次來計算 logits。不過，他們只需要兩次前向傳遞。第一次傳遞分別預計算完整序列 x 的鍵和值 K^1:B、V^1:B，在第二次前向傳遞中使用同時計算所有塊的去噪預測。

為了從 BD3-LM 中采樣，研究者以先前采樣的塊為條件，一次生成一個塊。生成塊后，他們緩存其鍵和值，類似于 AR。同時在每個塊的 T 個采樣步下，使用任何擴散采樣流程 SAMPLE中進行采樣。來從條件分布 pθ 

算法 1（塊擴散訓練）和算法 2（塊擴散采樣）分別如下圖（左）和（右）所示。

BD3-LM 訓練和采樣算法。

理解擴散模型與自回歸模型之間的似然差距

案例研究：單 Token 生成

該研究中的塊擴散參數(shù)化在期望上等同于自回歸負對數(shù)似然 (NLL)，特別是在 L′=1 的極限情況下。令人驚訝的是，當在 LM1B 數(shù)據(jù)集上訓練兩種模型時，研究發(fā)現(xiàn)塊擴散模型 (L′=1) 與自回歸模型之間存在兩點困惑度差距。研究確定擴散目標的高訓練方差是導致這一困惑度差距的原因。

在離散擴散 ELBO 下進行訓練時，存在高方差。

高方差訓練導致的擴散差距

直觀來說，如果采樣的掩碼率過低，重構(gòu) x 會變得容易，這不能提供有用的學習信號。如果掩碼全部內(nèi)容，最優(yōu)的重構(gòu)就是數(shù)據(jù)分布中每個標記的邊際概率，這很容易學習，同樣也沒有用處。

研究需要找到能夠最小化擴散目標引起的訓練方差，并進一步減少困惑度差距的噪聲調(diào)度方案。

基于數(shù)據(jù)的低方差訓練噪聲調(diào)度

為了避免導致高方差訓練的掩碼率，研究者在「裁剪的』掩碼率下來訓練 BD3-LMs。通過降低訓練方差，研究者在均勻采樣的掩碼率評估下改善了似然度。

由于最佳掩碼率可能會根據(jù)塊大小 L′的不同而變化，他們在訓練期間自適應(yīng)地學習 β,ω。在實踐中，研究者在每個驗證步驟后（經(jīng)過 5K 次梯度更新）使用網(wǎng)格搜索來優(yōu)化。

在下文中，研究者展示了針對每個塊大小優(yōu)化噪聲調(diào)度可以減少損失估計器的方差，并與其他替代調(diào)度方案相比實現(xiàn)最佳困惑度。

實驗結(jié)果

似然評估

BD3-LMs 在擴散模型中實現(xiàn)了最先進的似然水平。研究表明，通過調(diào)整塊長度 L′，BD3-LMs 可以在擴散和自回歸似然之間實現(xiàn)插值。

在 OWT 上測試針對 262B 標記訓練的模型的困惑度 (PPL; ↓)。

任意長度序列生成

許多現(xiàn)有擴散語言模型的一個主要缺點是，它們無法生成超過訓練時選擇的輸出上下文長度的完整文檔。例如，OpenWebText 包含最長達 131K tokens 的文檔，而離散擴散模型 SEDD（Lou 等人）僅限于生成 1024 tokens。研究表明，BD3-LMs 能夠通過解碼任意數(shù)量的塊來生成可變長度的文檔。

從在 OWT 上訓練的模型中抽樣 500 個文檔得出的生成長度統(tǒng)計信息。

研究者評估了 BD3-LMs 在變長序列上的生成質(zhì)量，使用相同數(shù)量的生成步驟（NFEs）比較了所有方法。他們用 GPT2-Large 模型測量生成序列的困惑度。結(jié)果表明，與之前所有的擴散方法相比，BD3-LMs 實現(xiàn)了最佳的生成困惑度。

300 個可變長度樣本的生成困惑度 （Gen. PPL;↓） 和功能評估次數(shù) （NFE;↓）。所有模型都在 OWT 上進行訓練，上下文長度為 L = 1024，并使用核采樣。

對于 MDLM，研究者使用了其分塊解碼技術(shù)（該技術(shù)不同于 BD3-LMs 中的分塊擴散訓練）處理 L=2048 的序列。研究者還與 SSD-LM（Han 等人提出）進行了比較，后者是一種替代性的分塊自回歸方法（也稱為半自回歸），它對詞嵌入執(zhí)行高斯擴散，但無法進行似然估計。該研究的離散方法使用比其他方法少一個數(shù)量級的生成步驟，產(chǎn)生了具有更好生成困惑度的樣本。

更多細節(jié)請參閱原論文。