Adobe 研究團隊利用狀態空間模型解鎖視頻世界模型中的長期記憶

視頻世界模型能夠基於動作預測未來幀，在人工智能領域具有巨大潛力，使智能體能夠在動態環境中進行規劃和推理。近期，視頻擴散模型在生成逼真未來序列方面表現出色，但一個關鍵瓶頸仍然存在：長期記憶。由於傳統注意力層在處理長序列時計算成本過高，現有模型難以記住遠距離過去的事件和狀態，限制了其在需要持續理解場景的複雜任務中的表現。

一篇由斯坦福大學、普林斯頓大學和 Adobe 研究團隊合著的新論文《長上下文狀態空間視頻世界模型》提出了一種創新解決方案。他們引入了一種新型架構，利用狀態空間模型（SSM）在不犧牲計算效率的前提下擴展時間記憶。

核心問題在於注意力機制相對於序列長度的二次計算複雜度。隨着視頻上下文的增長，注意力層的資源消耗急劇增加，使得長期記憶在實際應用中變得不可行。這意味着模型在經過一定幀數後實際上會“忘記”早期事件，從而影響其在需要長距離連貫性或長時間推理的任務中的性能。

作者的關鍵洞察是利用狀態空間模型在因果序列建模方面的固有優勢。與之前嘗試將 SSM 改造用於非因果視覺任務不同，這項工作充分利用了其在高效處理序列方面的優勢。

提出的長上下文狀態空間視頻世界模型（LSSVWM）包含了幾個關鍵設計選擇：

塊狀 SSM 掃描方案：這是設計的核心。他們並非用單個 SSM 掃描處理整個視頻序列，而是採用塊狀方案。這策略性地犧牲了一定空間一致性（在塊內），以顯著擴展時間記憶。通過將長序列分解為可管理的塊，模型可以在塊之間保持壓縮的“狀態”，從而有效延長記憶範圍。

密集局部注意力：為了補償塊狀 SSM 掃描可能帶來的空間連貫性損失，模型引入了密集局部注意力。這確保塊內和塊間的連續幀保持強關聯，保留了逼真視頻生成所需的細粒度細節和一致性。這種全局（SSM）和局部（注意力）的雙重處理方法使得模型既能實現長期記憶，又能保持局部保真度。

論文還介紹了兩種關鍵訓練策略以進一步提升長上下文性能：

擴散強制：該技術鼓勵模型基於輸入前綴生成幀，從而迫使其學習在更長持續時間內保持一致性。通過有時不採樣前綴並使所有標記保持噪聲，訓練相當於擴散強制，這被強調為前綴長度為零的長上下文訓練的特例。這推動模型即使在最小初始上下文下也能生成連貫序列。

幀局部注意力：為了加快訓練和採樣速度，作者實現了“幀局部注意力”機制。它利用 FlexAttention 實現了相比完全因果掩碼的顯著加速。通過將幀分組為塊（例如，塊大小為5，幀窗口大小為10），塊內的幀保持雙向性，同時也能關注前一個塊中的幀。這使得在優化計算負載的同時保持了有效的感受野。

研究人員在具有挑戰性的數據集上評估了 LSSVWM，包括 Memory Maze 和 Minecraft，這些數據集專門設計用於通過空間檢索和推理任務測試長期記憶能力。

實驗結果表明，該方法在保持長期記憶方面顯著超越了基線。定性結果（如補充圖 S1、S2、S3 所示）顯示，與僅依賴因果注意力或沒有幀局部注意力的 Mamba2 相比，LSSVWM 能夠在更長時間內生成更連貫且準確的序列。例如，在迷宮數據集的推理任務中，該模型在長時間範圍內保持了更好的一致性和準確性。同樣，在檢索任務中，LSSVWM 在回憶和利用遙遠過去幀的信息方面表現出改進的能力。至關重要的一點是，這些改進是在保持實際推理速度的同時實現的，使得模型適用於交互式應用。

論文《長上下文狀態空間視頻世界模型》可在 arXiv 上獲取。

該文章最初發表於 Synced。