在這個(gè) AI 技術(shù)與應(yīng)用大爆發(fā)的時(shí)代，我們最不缺的就是「熱詞」，從自回歸到擴(kuò)散模型，從注意力機(jī)制到狀態(tài)空間模型，從思維鏈到推理模型…… 有時(shí)候，其中一些熱詞會(huì)聚攏一處，為 AI 世界創(chuàng)造出新的可能性。

今天我們要介紹的這項(xiàng)研究便是如此，集齊了長(zhǎng)上下文、狀態(tài)空間模型（SSM）、擴(kuò)散模型、世界模型等「熱詞」，創(chuàng)造了一種全新的「視頻世界模型」。該研究來(lái)自斯坦福大學(xué)、普林斯頓大學(xué)和 Adobe Research，在社交網(wǎng)絡(luò)上引起了不少關(guān)注。

• 論文標(biāo)題：Long-Context State-Space Video World Models
• 論文地址：https://arxiv.org/pdf/2505.20171

要了解這項(xiàng)研究的貢獻(xiàn)，首先需要先界定一下相關(guān)概念。在這篇論文中，世界模型（world model）是指用于預(yù)測(cè)世界狀態(tài)如何隨動(dòng)作而演變的因果生成式模型。其可實(shí)現(xiàn)對(duì)復(fù)雜環(huán)境的交互式模擬。

而視頻擴(kuò)散模型已成為一種頗具前景的世界建模方法。不過，早期的視頻擴(kuò)散模型僅限于生成固定長(zhǎng)度的視頻，因此不適用于交互式應(yīng)用，而近期的架構(gòu)已可通過自回歸式的滑動(dòng)窗口預(yù)測(cè)實(shí)現(xiàn)無(wú)限長(zhǎng)度的視頻生成。這為一種新的范式鋪平了道路：基于交互式控制信號(hào)，視頻擴(kuò)散模型可以通過連續(xù)生成視頻幀而實(shí)現(xiàn)對(duì)視覺世界的交互式模擬。

然而，由于注意力機(jī)制的上下文長(zhǎng)度有限，現(xiàn)有視頻世界模型的時(shí)間記憶非常有限。這一限制使它們難以模擬具有長(zhǎng)期一致性的世界。

例如，當(dāng)使用現(xiàn)有視頻世界模型模擬游戲時(shí)，玩家只需向右看然后再次向左看，整個(gè)環(huán)境就可能完全改變（見圖 1）。

原因很容易理解：模型的注意力窗口中已經(jīng)沒有包含原始環(huán)境的幀了。

雖然理論上可以通過更長(zhǎng)的上下文窗口來(lái)擴(kuò)展記憶，但這種方法有兩大問題：

• 訓(xùn)練的計(jì)算成本會(huì)與上下文長(zhǎng)度呈二次方增長(zhǎng)，使其成本過高；
• 每幀推理時(shí)間隨上下文長(zhǎng)度線性增長(zhǎng)，導(dǎo)致生成速度越來(lái)越慢，這對(duì)于需要實(shí)時(shí)、無(wú)限長(zhǎng)度生成的應(yīng)用（例如游戲）來(lái)說(shuō)，根本沒法用。

那么，今天我們要介紹的這篇論文有何創(chuàng)新之處呢？

簡(jiǎn)單來(lái)說(shuō)，他們使用了狀態(tài)空間模型（SSM）來(lái)實(shí)現(xiàn)長(zhǎng)期記憶，其中關(guān)鍵在于 Mamba 的逐塊掃描（block-wise scan）方案 —— 能在保留時(shí)間因果關(guān)系的同時(shí)，實(shí)現(xiàn)時(shí)間記憶與空間一致性的最佳平衡。另外，該團(tuán)隊(duì)還對(duì)該方案進(jìn)行了補(bǔ)充：在相鄰幀之間設(shè)置了密集的局部注意力機(jī)制，從而能以最小的計(jì)算開銷實(shí)現(xiàn)高保真度的生成。

該團(tuán)隊(duì)介紹說(shuō)：「不同于以往針對(duì)非因果視覺任務(wù)改進(jìn) SSM 的方法，我們的方法有根本上的差異：我們專門使用了 SSM 來(lái)處理因果時(shí)間動(dòng)態(tài)并追蹤世界狀態(tài)，充分利用了其在序列建模方面的固有優(yōu)勢(shì)?！?/span>

對(duì)視頻擴(kuò)散模型和狀態(tài)空間模型的基礎(chǔ)數(shù)學(xué)描述請(qǐng)參看原論文，下面將更詳細(xì)地介紹這項(xiàng)研究的創(chuàng)新。

新方法詳解

模型架構(gòu)

由于這個(gè)模型會(huì)以自回歸的方式（一次一幀）生成視頻幀，因此時(shí)間維度（幀序列）必須位于掃描順序的末尾。這種「空間主 / 時(shí)間次」的排序可確保模型在移動(dòng)到下一幀之前處理完當(dāng)前幀內(nèi)的所有空間信息，從而保留因果約束并防止模型訪問未來(lái)幀的信息。

然而，以空間為主的掃描順序會(huì)使得捕捉長(zhǎng)期時(shí)間依賴性變得困難，因?yàn)樵谡蛊降?token 序列中，時(shí)間上相鄰的 token 彼此之間會(huì)變得相當(dāng)遙遠(yuǎn)。

為了解決這一限制，該團(tuán)隊(duì)提出了一種平衡時(shí)間記憶和空間一致性的方法，即對(duì)時(shí)空 token 進(jìn)行逐塊重新排序（block-wise reordering）。

逐塊 SSM 掃描。如圖 3（右下）所示，新提出的方法會(huì)將原始 token 序列沿空間維度分解為大小為 (b_h, b_w, T) 的塊，其中 b_h 和 b_w 是與層相關(guān)的塊高度 / 寬度，T 是數(shù)據(jù)的時(shí)間維度。

需要注意，這里并不會(huì)對(duì)所有 token 序列進(jìn)行一次掃描，而是對(duì)每個(gè) token 塊進(jìn)行單獨(dú)的掃描。通過控制 b_h 和 b_w 的值，可以在時(shí)間相關(guān)性和空間一致性之間取得平衡。現(xiàn)在，時(shí)間上相鄰的 token 以 b_h × b_w token 分隔，而不是像傳統(tǒng)的以空間為主的掃描中那樣以 H × W token 分隔，其中 H、W 表示每幀的高度 / 寬度。

然而，較小的塊會(huì)導(dǎo)致空間一致性更差，因?yàn)楠?dú)立的掃描會(huì)阻止不同塊中的 token 交互。因此，塊大小的選擇代表了一種在一致性長(zhǎng)期記憶和短期空間一致性之間進(jìn)行權(quán)衡的有效方法。通過在不同的層中采用不同的 b_h 和 b_w 值，該模型可充分利用大塊和小塊的優(yōu)勢(shì)。

由于固定維度的 SSM 狀態(tài)的表征能力有限，因此 SSM 在處理視覺生成等高復(fù)雜度任務(wù)時(shí)可能會(huì)遇到困難。新提出的逐塊掃描方法可通過有效地增加每層的 SSM 狀態(tài)的維度來(lái)緩解這一限制，因?yàn)槊總€(gè)塊都被分配了一個(gè)單獨(dú)的狀態(tài)。

幀局部注意力機(jī)制。研究已經(jīng)證明，Mamba 等線性注意力機(jī)制的變體在與聯(lián)想回憶相關(guān)的任務(wù)中表現(xiàn)不佳。在視頻生成中，Mamba 無(wú)法檢索精確的局部信息，導(dǎo)致幀間質(zhì)量不佳，并會(huì)喪失短期時(shí)間一致性。

之前有研究表明，將局部注意力機(jī)制與 SSM 相結(jié)合的混合架構(gòu)可以提升語(yǔ)言建模的效果。在新提出的模型中，會(huì)在每次 Mamba 掃描后引入一個(gè)逐幀局部注意力模塊，如圖 3 所示。在訓(xùn)練過程中，應(yīng)用逐塊因果注意力機(jī)制，其中每個(gè) token 只能關(guān)注同一幀中的 token 以及一個(gè)固定大小的前幾幀窗口。注意力掩碼 M 的形式為：

其中 i 和 j 是序列中幀的索引，k 是窗口大小。

動(dòng)作條件。為了在自回歸生成過程中啟用交互式控制，該團(tuán)隊(duì)的做法是將與每幀對(duì)應(yīng)的動(dòng)作作為輸入。這里，會(huì)通過一個(gè)小型多層感知器 (MLP) 處理連續(xù)動(dòng)作值（例如，攝像機(jī)位置），并添加到噪聲級(jí)別嵌入中，然后通過自適應(yīng)歸一化層將其注入到網(wǎng)絡(luò)中。對(duì)于離散動(dòng)作，這里是直接學(xué)習(xí)與每個(gè)可能動(dòng)作對(duì)應(yīng)的嵌入。

長(zhǎng)上下文訓(xùn)練

該團(tuán)隊(duì)指出，盡管新提出的架構(gòu)設(shè)計(jì)可增強(qiáng)模型維持長(zhǎng)期記憶的能力，但使用標(biāo)準(zhǔn)的擴(kuò)散訓(xùn)練方案仍舊難以學(xué)習(xí)長(zhǎng)時(shí)域依賴性。視頻數(shù)據(jù)包含大量冗余，這使得模型在大多數(shù)情況下主要依賴鄰近幀進(jìn)行去噪。因此，擴(kuò)散模型經(jīng)常陷入局部最小值，無(wú)法捕捉長(zhǎng)期依賴性。

在訓(xùn)練期間，標(biāo)準(zhǔn)的 diffusion forcing 始終會(huì)向每個(gè)幀獨(dú)立添加噪聲。在這種情況下，模型參考遠(yuǎn)處上下文幀的動(dòng)力有限，因?yàn)樗鼈兺ǔ０挠杏眯畔⑸儆诰植繋?/span>

為了鼓勵(lì)模型關(guān)注遠(yuǎn)處幀并學(xué)習(xí)長(zhǎng)期相關(guān)性，該團(tuán)隊(duì)將 diffusion forcing 與一種改進(jìn)的訓(xùn)練方案結(jié)合了起來(lái)。該方案可在訓(xùn)練期間保持幀的隨機(jī)長(zhǎng)度前綴完全干凈（無(wú)噪聲），如圖 4 所示。

當(dāng)向后續(xù)幀添加較大噪聲時(shí)，干凈的上下文幀可能比嘈雜的局部幀提供更多有用信息，從而促使模型有效地利用它們。這與 Ca2VDM 中的訓(xùn)練方案類似。

通過固定長(zhǎng)度狀態(tài)進(jìn)行高效推理

在推理過程中，再根據(jù)輸入動(dòng)作自回歸地生成新的視頻幀。新提出的混合架構(gòu)可確保恒定的速度和內(nèi)存使用率。

具體而言，該模型的每一層僅跟蹤：前 k 幀的固定長(zhǎng)度 KV 緩存，以及每個(gè)塊的 SSM 狀態(tài)。這可確保整個(gè)推理過程中內(nèi)存使用率的恒定，這不同于完全因果式的 Transformer—— 在生成過程中內(nèi)存需求會(huì)隨著存儲(chǔ)所有先前幀的 KV 緩存而線性增長(zhǎng)。

同樣，新提出的方法可保持每幀生成速度恒定，因?yàn)榫植孔⒁饬C(jī)制和逐塊 SSM 計(jì)算不會(huì)隨視頻長(zhǎng)度而變化。此特性對(duì)于視頻世界模型應(yīng)用至關(guān)重要，因?yàn)檫@些應(yīng)用通常非常需要無(wú)限期地生成視頻幀而不降低性能。

實(shí)驗(yàn)表現(xiàn)

該團(tuán)隊(duì)從訓(xùn)練和推理效率以及長(zhǎng)期記憶能力方面評(píng)估了新提出的方法。

為此，他們使用了兩個(gè)長(zhǎng)視頻數(shù)據(jù)集，并評(píng)估該模型在空間記憶任務(wù)中的表現(xiàn)，這些任務(wù)為了生成準(zhǔn)確的預(yù)測(cè)，需要回憶遠(yuǎn)距離幀的信息。有關(guān)數(shù)據(jù)集和評(píng)估方法的更詳細(xì)介紹請(qǐng)?jiān)L問原論文，下面重點(diǎn)來(lái)看實(shí)驗(yàn)結(jié)果。

首先，表 2 和表 3 給出了不同模型在 Memory Maze 上進(jìn)行空間檢索和推理的定量結(jié)果。

可以看到，新提出的模型在檢索和推理這兩個(gè)任務(wù)的所有指標(biāo)上都是最優(yōu)的。

如圖 5 和圖 6 所示，對(duì)于這兩項(xiàng)任務(wù)，其他次二次模型的幀預(yù)測(cè)在一段時(shí)間后會(huì)偏離 ground truth，而新方法在整個(gè)軌跡范圍內(nèi)都能保持準(zhǔn)確的預(yù)測(cè)。

圖 7 進(jìn)一步分析了每種方法在檢索任務(wù)上的性能，展示了隨著生成幀和檢索幀之間距離的增加，檢索準(zhǔn)確率的變化。

因果 Transformer 在其訓(xùn)練上下文中表現(xiàn)良好，但超過其最大訓(xùn)練長(zhǎng)度后會(huì)迅速下降。其他線性復(fù)雜度方法（例如 Mamba 和 Mamba2 + Frame Local Attn）由于狀態(tài)空間表達(dá)能力有限而表現(xiàn)不佳。

相比之下，本文的新方法在所有檢索距離上都保持了較高的準(zhǔn)確度，與在完整上下文上訓(xùn)練的因果 Transformer 相當(dāng)。

另外，該團(tuán)隊(duì)也在 TECO Minecraft 上進(jìn)行了實(shí)驗(yàn)，表 4 和圖 2 分別給出了定量和定性結(jié)果。這里參與對(duì)比的模型是 diffuion forcing transformer（DFoT）—— 一種在 diffuion forcing 機(jī)制下訓(xùn)練的雙向 Transformer，算得上是當(dāng)前自回歸長(zhǎng)視頻生成領(lǐng)域最先進(jìn)的架構(gòu)。然而，由于其模型的二次復(fù)雜度，DFoT 是在 25 幀的有限上下文長(zhǎng)度上訓(xùn)練的。

可以看到，新方法可以準(zhǔn)確預(yù)測(cè)先前探索過的區(qū)域，而上下文窗口有限的方法則無(wú)法做到這一點(diǎn)。

總體而言，新方法優(yōu)于 DFoT 和在 25 幀上下文上訓(xùn)練的因果 Transformer。

由于軌跡較短，所有模型在該數(shù)據(jù)集上的相似度都較低，其中模型僅獲得 100 幀上下文來(lái)預(yù)測(cè) 50 幀。通常而言，100 幀的上下文不足以讓智能體完全觀察環(huán)境，從而可能導(dǎo)致任務(wù)軌跡冒險(xiǎn)進(jìn)入先前未見過的區(qū)域，在這種情況下，逐幀相似度的信息量會(huì)降低。

該團(tuán)隊(duì)也研究了新方法的訓(xùn)練和推理成本。圖 8 使用三個(gè)指標(biāo)評(píng)估模型性能：每次迭代的訓(xùn)練成本（左）、生成期間的內(nèi)存利用率（中）以及推理期間的計(jì)算時(shí)間（右）。

可以看到，新提出的方法在所有指標(biāo)上都表現(xiàn)出了卓越的擴(kuò)展性：訓(xùn)練時(shí)間會(huì)隨上下文長(zhǎng)度線性擴(kuò)展，同時(shí)能在推理期間保持恒定的內(nèi)存和計(jì)算成本。為了比較推理運(yùn)行時(shí)間，該團(tuán)隊(duì)還比較了通過幀局部注意力機(jī)制加 SSM 更新進(jìn)行單次前向傳遞的運(yùn)行時(shí)間，以及對(duì)所有先前生成的幀進(jìn)行 KV 緩存的完整注意力機(jī)制的運(yùn)行時(shí)間。

更多詳情請(qǐng)參閱原論文。

順帶一提，正如 Meta 和蒙特利爾學(xué)習(xí)算法研究所研究者 Artem Zholus 在機(jī)器之心 ?? 帳號(hào)下評(píng)論的那樣，使用 SSM 來(lái)構(gòu)建世界模型的研究一兩年就已經(jīng)有了，感興趣的讀者可擴(kuò)展閱讀。

1. Mastering Memory Tasks with World Models

項(xiàng)目地址：https://recall2imagine.github.io/

2.  Facing Off World Model Backbones: RNNs, Transformers, and S4

項(xiàng)目地址：https://fdeng18.github.io/s4wm/