透過結合狀態空間模型 (SSMs) 以實現高效的長距離依賴建模，並利用密集局部注意力來保持連貫性，同時採用擴散強制和影格局部注意力等訓練策略，Adobe Research 的研究人員成功克服了影片生成中長期記憶的長期挑戰。
影片世界模型能夠根據動作預測未來影格，在人工智慧領域具有巨大潛力，使代理能夠在動態環境中進行規劃和推理。最近的進展，特別是影片擴散模型，已展現出生成逼真未來序列的驚人能力。然而，一個顯著的瓶頸依然存在：維持長期記憶。由於使用傳統注意力層處理擴展序列的計算成本高昂，目前的模型難以記住遙遠過去的事件和狀態。這限制了它們執行需要持續理解場景的複雜任務的能力。
一篇由史丹佛大學、普林斯頓大學和 Adobe Research 的研究人員共同撰寫的新論文《長上下文狀態空間影片世界模型》(Long-Context State-Space Video World Models)，提出了一個創新的解決方案來應對這一挑戰。他們引入了一種新穎的架構，利用狀態空間模型 (SSMs) 來擴展時間記憶，同時不犧牲計算效率。
核心問題在於注意力機制相對於序列長度的二次方計算複雜度。隨著影片上下文的增長，注意力層所需的資源呈指數級增長，使得長期記憶在實際應用中變得不切實際。這意味著在一定數量的影格之後，模型會有效地「忘記」較早的事件，從而阻礙其在需要長距離連貫性或長時間推理的任務上的表現。
作者們的關鍵見解是利用狀態空間模型 (SSMs) 在因果序列建模方面的固有優勢。與之前將 SSMs 改裝用於非因果視覺任務的嘗試不同，這項工作充分利用了它們在高效處理序列方面的優勢。
所提出的長上下文狀態空間影片世界模型 (LSSVWM) 包含了幾個關鍵的設計選擇：
1. 區塊式 SSM 掃描機制： 這是其設計的核心。他們採用區塊式機制，而不是用單一的 SSM 掃描處理整個影片序列。這策略性地犧牲了一些空間一致性（在一個區塊內）以換取顯著擴展的時間記憶。透過將長序列分解為可管理的區塊，他們可以維持一個壓縮的「狀態」，將資訊傳遞到各個區塊，有效地擴展了模型的記憶範圍。
2. 密集局部注意力： 為了彌補區塊式 SSM 掃描可能導致的空間連貫性損失，該模型整合了密集局部注意力。這確保了區塊內部和跨區塊的連續影格之間保持強烈的關係，保留了生成逼真影片所需的細緻細節和一致性。這種全局（SSM）和局部（注意力）處理的雙重方法使他們能夠同時實現長期記憶和局部保真度。
該論文還引入了兩種關鍵的訓練策略，以進一步提高長上下文性能：
- 擴散強制： 這種技術鼓勵模型在輸入的前綴條件下生成影格，有效地迫使它學習在更長的時間內保持一致性。有時不採樣前綴並保持所有 token 噪聲化，訓練就等同於擴散強制，這被強調為長上下文訓練的一個特例，其中前綴長度為零。這促使模型即使從最少的初始上下文也能生成連貫的序列。
- 影格局部注意力： 為了加快訓練和採樣速度，作者們實施了「影格局部注意力」機制。這利用 FlexAttention 實現了相對於完全因果遮罩的顯著加速。透過將影格分組到區塊中（例如，5 個影格的區塊，影格視窗大小為 10），區塊內的影格保持雙向性，同時也關注前一個區塊中的影格。這允許有效的感受野，同時優化計算負載。
研究人員在具有挑戰性的資料集上評估了他們的 LSSVWM，包括 Memory Maze 和 Minecraft，這些資料集專門設計用於透過空間檢索和推理任務來測試長期記憶能力。
實驗證明，他們的方法在保留長距離記憶方面大幅超越了基準模型。定性結果，如補充圖（例如 S1、S2、S3）所示，說明 LSSVWM 能夠在更長的時間內生成比僅依賴因果注意力甚至沒有影格局部注意力的 Mamba2 模型更連貫和準確的序列。
例如，在迷宮資料集的推理任務上，他們的模型在長距離上保持了更好的一致性和準確性。同樣，在檢索任務中，LSSVWM 顯示出從遙遠過去的影格中回憶和利用資訊的能力有所提高。至關重要的是，這些改進是在保持實際推論速度的同時實現的，使得這些模型適用於互動式應用。
論文《長上下文狀態空間影片世界模型》(Long-Context State-Space Video World Models) 已發布於 arXiv。
