近日，MIT CSAIL 的一個研究團隊（一作為 MIT 在讀博士陳博遠）成功地將全序列擴散模型與下一 token 模型的強大能力統(tǒng)合到了一起，提出了一種訓(xùn)練和采樣范式：Diffusion Forcing（DF）。

• 論文標(biāo)題：Diffusion Forcing:Next-token Prediction Meets Full-Sequence Diffusion
• 論文地址：https://arxiv.org/pdf/2407.01392
• 項目網(wǎng)站：https://boyuan.space/diffusion-forcing
• 代碼地址：https://github.com/buoyancy99/diffusion-forcing?

如下所示，擴散強制在一致性和穩(wěn)定性方面都明顯勝過全序列擴散和教師強制這兩種方法。

在該框架中，每個 token 都關(guān)聯(lián)了一個隨機的、獨立的噪聲水平，并且可使用一種共享的下一 token 預(yù)測模型或下幾 token 預(yù)測模型根據(jù)任意的、獨立的、每 token 的方案對 token 進行去噪。

該方法的研究靈感來自這一觀察：對 token 加噪聲的過程就是一種形式的部分掩碼過程 —— 零噪聲就意味著未對 token 加掩碼，而完整噪聲則是完全掩蔽 token。因此，DF 可強迫模型學(xué)習(xí)去除任何可變有噪聲 token 集合的掩碼（圖 2）。

與此同時，通過將預(yù)測方法參數(shù)化為多個下一 token 預(yù)測模型的組合，該系統(tǒng)可以靈活地生成不同長度的序列，并以組合方式泛化到新的軌跡（圖 1）。

該團隊將用于序列生成的 DF 實現(xiàn)成了因果擴散強制（Causal Diffusion Forcing/CDF），其中未來 token 通過一個因果架構(gòu)依賴于過去 token。他們訓(xùn)練該模型一次性去噪序列的所有 token（其中每個 token 都有獨立的噪聲水平）。

在采樣期間，CDF 會將一個高斯噪聲幀序列逐漸地去噪成潔凈的樣本，其中不同幀在每個去噪步驟可能會有不同的噪聲水平。類似于下一 token 預(yù)測模型，CDF 可以生成長度可變的序列；不同于下一 token 預(yù)測，CDF 的表現(xiàn)非常穩(wěn)定 —— 不管是預(yù)測接下來的一個 token，還是未來的數(shù)千 token，甚至是連續(xù) token。

此外，類似于全序列擴散，它也可接收引導(dǎo)，從而實現(xiàn)高獎勵生成。通過協(xié)同利用因果關(guān)系、靈活的范圍和可變噪聲調(diào)度，CDF 能實現(xiàn)一項新功能：蒙特卡洛樹引導(dǎo)（MCTG）。相比于非因果全序列擴散模型，MCTG 能極大提升高獎勵生成的采樣率。圖 1 給出了這些能力的概況。

Diffusion Forcing（擴散強制）

1、將加噪過程視為部分掩碼

首先，我們可以將任意 token 集合（不管是否為序列）視為一個通過 t 索引的有序集合。那么，使用教師強制（teacher forcing）訓(xùn)練下一 token 預(yù)測便可被解釋成掩蔽掉時間 t 的每個 token x_t 基于過去 x_{1:t?1} 預(yù)測它們。

對于序列，可將這種操作描述成：沿時間軸執(zhí)行掩碼。我們可以將全序列前向擴散（即逐漸向數(shù)據(jù)

添加噪聲的過程）看作一種部分掩碼（partial masking），這可被稱為「沿噪聲軸執(zhí)行掩碼）。

事實上，在 K 步加噪之后，

（大概）就是白噪聲了，不再有任何有關(guān)原數(shù)據(jù)的信息。

如圖 2 所示，該團隊建立了一個統(tǒng)一視角來看待沿這兩個軸的掩碼。

2、擴散強制：不同 token 的噪聲水平不同

擴散強制（DF）框架可用于訓(xùn)練和采樣任意序列長度的有噪聲 token 

，其中的關(guān)鍵在于每個 token 的噪聲水平 k_t 會隨時間步驟而變化。

這篇論文關(guān)注的重點是時間序列數(shù)據(jù)，因此他們通過一種因果架構(gòu)實例化了 DF，并由此得到了因果擴散強制（CDF）。簡單來說，這是使用基礎(chǔ)循環(huán)神經(jīng)網(wǎng)絡(luò)（RNN）獲得的一種最小實現(xiàn)。

權(quán)重為 θ 的 RNN 維護著獲悉過去 token 影響的隱藏狀態(tài) z_t，其會通過一個循環(huán)層根據(jù)動態(tài)

而演化。當(dāng)獲得輸入噪聲觀察

時，就以馬爾可夫方式更新該隱藏狀態(tài)。

當(dāng) k_t=0 時，這就是貝葉斯過濾中的后驗更新；而當(dāng) k_t= K（純噪聲、無信息）時，這就等價于建模貝葉斯過濾中的「后驗分布」p_θ(z_t | z_{t?1})。

給定隱藏狀態(tài) z_t，觀察模型 p_θ(x_t^0 | z_t) 的目標(biāo)是預(yù)測 x_t；這個單元的輸入 - 輸出行為與標(biāo)準(zhǔn)的條件擴散模型一樣：以條件變量 z_{t?1} 和有噪聲 token 為輸入，預(yù)測無噪聲的 x_t=x_t^0，并由此間接地通過仿射重新參數(shù)化預(yù)測噪聲 ε^{k_t}。因此，我們就可以直接使用經(jīng)典的擴散目標(biāo)來訓(xùn)練（因果）擴散強制。根據(jù)噪聲預(yù)測結(jié)果 ε_θ，可以對上述單元進行參數(shù)化。然后，通過最小化以下?lián)p失來找到參數(shù) θ：

算法 1 給出了偽代碼。重點在于，該損失捕獲了貝葉斯過濾和條件擴散的關(guān)鍵元素。該團隊也進一步重新推斷了用于擴散強制的擴散模型訓(xùn)練中的常用技術(shù)，詳見原論文的附錄部分。他們也得出了一個非正式的定理。

定理 3.1（非正式）。擴散強制訓(xùn)練流程（算法 1）是在期望對數(shù)似然

上優(yōu)化證據(jù)下限（ELBO）的重新加權(quán)，其中期望值會在噪聲水平上平均，而

是根據(jù)前向過程加噪。此外，在適當(dāng)條件下，優(yōu)化 (3.1) 式還可以同時最大化所有噪聲水平序列的似然下限。

擴散強制采樣和所得到的能力

算法 2 描述了采樣過程，其定義是：在二維的 M × T 網(wǎng)格 K ∈ [K]^{M×T} 上指定噪聲調(diào)度；其中列對應(yīng)于時間步驟 t，m 索引的行則決定了噪聲水平。

為了生成長度為 T 的整個序列，先將 token x_{1:T} 初始化為白噪聲，對應(yīng)于噪聲水平 k = K。然后沿著網(wǎng)格逐行向下迭代，并從左到右逐列去噪，直到噪聲水平達到 K。到最后一行 m = 0 時，token 的噪聲已清理干凈，即噪聲水平為 K_{0,t} ≡ 0。

這個采樣范式會帶來如下新能力：

• 讓自回歸生成變得穩(wěn)定
• 保持未來的不確定
• 長期引導(dǎo)能力

將擴散強制用于靈活的序列決策

擴散強制的新能力也帶來了新的可能性。該團隊基于此為序列決策（SDM）設(shè)計了一種全新框架，并且將其成功應(yīng)用到了機器人和自主智能體領(lǐng)域。

首先，定義一個馬爾可夫決策過程，該過程具有動態(tài) p (s_{t+1}|s_t, a_t)、觀察 p (o_t|s_t) 和獎勵 p (r_t|s_t, a_t)。這里的目標(biāo)是訓(xùn)練一個策略 π(a_t|o_{1:t})，使得軌跡 

的預(yù)期累積獎勵最大化。這里分配 token x_t = [a_t, r_t, o_{t+1}]。一條軌跡就是一個序列 x_{1:T}，其長度可能是可變的；訓(xùn)練方式則如算法 1 所示。

在執(zhí)行過程的每一步 t，都有一個隱藏狀態(tài) z_{t-1} 總結(jié)過去的無噪聲 token x_{1:t-1}?；谶@個隱藏狀態(tài)，根據(jù)算法 2 采樣一個規(guī)劃

，其中

包含預(yù)測的動作、獎勵和觀察。H 是一個前向觀察窗口，類似于模型預(yù)測控制中的未來預(yù)測。在采用了規(guī)劃的動作之后，環(huán)境會得到一個獎勵和下一個觀察，從而得到下一個 token。其中隱藏狀態(tài)可以根據(jù)后驗 p_θ(z_t|z_{t?1}, x_t, 0) 獲得更新。

該框架既可以作為策略，也可以作為規(guī)劃器，其優(yōu)勢包括：

• 具有靈活的規(guī)劃范圍
• 可實現(xiàn)靈活的獎勵引導(dǎo)
• 能實現(xiàn)蒙特卡洛樹引導(dǎo)（MCTG），從而實現(xiàn)未來不確定性

實驗

該團隊評估了擴散強制作為生成序列模型的優(yōu)勢，其中涉及視頻和時間序列預(yù)測、規(guī)劃和模仿學(xué)習(xí)等多種應(yīng)用。

視頻預(yù)測：一致且穩(wěn)定的序列生成和無限展開

針對視頻生成式建模任務(wù)，他們基于 Minecraft 游戲視頻和 DMLab 導(dǎo)航為因果擴散強制訓(xùn)練了一個卷積 RNN 實現(xiàn)。

圖 3 展示了擴散強制與基準(zhǔn)的定性結(jié)果。

可以看到，擴散強制能穩(wěn)定地展開，甚至能超過其訓(xùn)練范圍；而教師強制和全序列擴散基準(zhǔn)會很快發(fā)散。

擴散規(guī)劃：MCTG、因果不確定性、靈活的范圍控制

擴散強制的能力能為決策帶來獨有的好處。該團隊使用一種標(biāo)準(zhǔn)的離線強化學(xué)習(xí)框架 D4RL 評估了新提出的決策框架。

表 1 給出了定性和定量的評估結(jié)果。可以看到，擴散強制在全部 6 個環(huán)境中都優(yōu)于 Diffuser 和所有基準(zhǔn)。

可控的序列組合生成

該團隊發(fā)現(xiàn)，僅需修改采樣方案，就可以靈活地組合訓(xùn)練時間觀察到的序列的子序列。

他們使用一個 2D 軌跡數(shù)據(jù)集進行了實驗：在一個方形平面上，所有軌跡都是始于一角并最終到達對角，形成一種十字形。

如上圖 1 所示，當(dāng)不需要組合行為時，可讓 DF 保持完整記憶，復(fù)制十字形的分布。當(dāng)需要組合時，可讓模型使用 MPC 無記憶地生成更短的規(guī)劃，從而實現(xiàn)對這個十字形的子軌跡的縫合，得到 V 形軌跡。

機器人：長范圍模仿學(xué)習(xí)和穩(wěn)健的視覺運動控制

擴散強制也為真實機器人的視覺運動控制帶來了新的機會。

模仿學(xué)習(xí)是一種常用的機器人操控技術(shù)，即學(xué)習(xí)專家演示的觀察到動作的映射。但是，缺乏記憶往往會讓模仿學(xué)習(xí)難以完成長范圍的任務(wù)。DF 不僅能緩解這個短板，還能讓模仿學(xué)習(xí)更穩(wěn)健。

使用記憶進行模仿學(xué)習(xí)。通過遙控 Franka 機器人，該團隊收集了一個視頻和動作數(shù)據(jù)集。如圖 4 所示，任務(wù)就是利用第三個位置交換蘋果和橘子的位置。水果的初始位置是隨機的，因此可能的目標(biāo)狀態(tài)有兩個。

此外，當(dāng)?shù)谌齻€位置有一個水果時，就無法通過當(dāng)前觀察推斷出所需結(jié)果 —— 策略必須記住初始配置才能決定移動哪個水果。不同于常用的行為克隆方法，DF 可以自然地將記憶整合進自己的隱藏狀態(tài)中。結(jié)果發(fā)現(xiàn)，DF 能實現(xiàn) 80% 的成功率，而擴散策略（當(dāng)前最佳的無記憶模仿學(xué)習(xí)算法）卻失敗了。

此外，DF 還能更穩(wěn)健地應(yīng)對噪聲并助益機器人預(yù)訓(xùn)練。

時間序列預(yù)測：擴散強制是一種優(yōu)秀的通用序列模型

對于多變量時間序列預(yù)測任務(wù)，該團隊的研究表明 DF 足以與之前的擴散模型和基于 Transformer 的模型媲美。

更多技術(shù)細節(jié)和實驗結(jié)果請參閱原論文。

本文轉(zhuǎn)自機器之心 ，作者：機器之心

原文鏈接:??https://mp.weixin.qq.com/s/kz4RvqdK6nGtA11y5nq5xQ??