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

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

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

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

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

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

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

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

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

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

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

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

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

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

該文章最初發表於 Synced。