Adobe Research、状態空間モデルを用いてビデオ世界モデルの長期記憶を実現
Adobe Research、スタンフォード大学、プリンストン大学の研究チームは、状態空間モデル（SSM）と密な局所注意を組み合わせた新しいアーキテクチャを提案し、ビデオ生成における長期記憶の課題を解決した。ブロック単位のSSMスキャン、拡散強制、フレーム局所注意などの戦略により、Memory MazeやMinecraftデータセットで優れた性能を達成し、計算効率を維持しながらインタラクティブな応用を可能にする。
ビデオ世界モデルは、行動に基づいて未来のフレームを予測するものであり、エージェントが動的環境で計画や推論を行うための可能性を秘めている。近年、ビデオ拡散モデルはリアルな未来シーケンスの生成において印象的な能力を示しているが、長期記憶の維持という大きなボトルネックが残っている。従来の注意機構はシーケンス長に対して二次的な計算複雑性を持つため、長いシーケンスを処理する際に計算コストが爆発し、モデルは遠い過去のイベントや状態を記憶できず、持続的なシーン理解を必要とする複雑なタスクの性能が制限される。
スタンフォード大学、プリンストン大学、Adobe Researchによる新しい論文「Long-Context State-Space Video World Models」は、この課題に対する革新的な解決策を提案している。彼らは、状態空間モデル（SSM）を活用して計算効率を犠牲にすることなく時間的記憶を拡張する新しいアーキテクチャを導入した。
核心は、注意機構のシーケンス長に対する二次計算複雑性にある。ビデオコンテキストが長くなるにつれて、注意層に必要なリソースが爆発的に増加し、長期記憶が現実的なアプリケーションで非現実的になる。これは、一定のフレーム数を超えるとモデルが以前のイベントを事実上「忘れる」ことを意味し、長距離の一貫性や長期にわたる推論を必要とするタスクの性能を損なう。
著者らは、因果シーケンスモデリングにおける状態空間モデルの固有の強みを活用することを着想した。以前の試みが非因果的なビジョンタスクにSSMを改造したのとは異なり、この研究はシーケンスを効率的に処理する利点を完全に活用している。
提案されたLong-Context State-Space Video World Model（LSSVWM）は、以下の重要な設計選択を含む：
ブロック単位のSSMスキャンスキーム：これは設計の中心である。ビデオシーケンス全体を単一のSSMスキャンで処理する代わりに、ブロック単位のスキームを採用する。これにより、ブロック内の空間的一貫性をある程度犠牲にする代わりに、時間的記憶を大幅に拡張する。長いシーケンスを管理可能なブロックに分割することで、ブロック間で圧縮された「状態」を維持し、モデルの記憶範囲を効果的に延長する。
密な局所注意：ブロック単位のSSMスキャンによって生じる可能性のある空間的コヒーレンスの損失を補うために、モデルは密な局所注意を組み込む。これにより、ブロック内およびブロック間の連続フレームが強い関係を維持し、リアルなビデオ生成に必要な細かい詳細と一貫性を保持する。このグローバル（SSM）とローカル（注意）の二重処理により、長期記憶と局所的な忠実度の両方を達成する。
論文では、長期コンテキスト性能をさらに向上させるための2つの主要な訓練戦略も紹介している：
拡散強制：この技術は、モデルに入力プレフィックスに条件付けられたフレームを生成させることで、長期間にわたって一貫性を維持することを学習させる。プレフィックスをサンプリングせずにすべてのトークンをノイズのままにすることもあり、これはプレフィックス長がゼロの長期コンテキスト訓練の特殊ケースとして拡散強制と等価になる。これにより、最小限の初期コンテキストからでも一貫したシーケンスを生成するようモデルが促進される。
フレーム局所注意：訓練とサンプリングを高速化するために、著者らは「フレーム局所注意」メカニズムを実装した。これはFlexAttentionを利用して、完全な因果マスクと比較して大幅な高速化を達成する。フレームをチャンクにグループ化し（例：チャンクサイズ5、フレームウィンドウサイズ10）、チャンク内のフレームは双方向性を維持しながら、前のチャンクのフレームにも注意を向ける。これにより、計算負荷を最適化しつつ、効果的な受容野を可能にする。
研究者らは、LSSVWMをMemory MazeやMinecraftなどの挑戦的なデータセットで評価した。これらのデータセットは、空間的検索や推論タスクを通じて長期記憶能力をテストするために特別に設計されている。
実験では、このアプローチが長期記憶の保持においてベースラインを大幅に上回ることが示された。定性結果（補足図S1、S2、S3など）は、因果注意のみやフレーム局所注意のないMamba2と比較して、LSSVWMが長期にわたってより一貫性のある正確なシーケンスを生成できることを示している。例えば、迷路データセットの推論タスクでは、モデルは長期にわたって一貫性と正確性を維持する。同様に、検索タスクでは、LSSVWMは遠い過去のフレームからの情報を想起し利用する能力が向上する。重要なことに、これらの改善は実用的な推論速度を維持しながら達成されており、モデルをインタラクティブなアプリケーションに適したものにしている。
論文「Long-Context State-Space Video World Models」はarXivで入手可能。
この記事は最初にSyncedに掲載されました。