WorldKV: Efficient World Memory with World Retrieval and Compression

技術サマリー：WorldKV

問題定義
自己回帰型ビデオ拡散モデルは、リアルタイムかつ行動条件付きの世界生成を可能にしましたが、以前に観た視点への再訪問で一貫したコンテンツを維持する持続的な世界を構築することは、依然として未解決の課題です。完全な KV キャッシュアテンションはこの一貫性を維持しますが、ロールアウト長に伴うメモリフットプリントとアテンションコストの線形増加を招き、リアルタイム制約を破綻させ、メモリ不足（OOM）エラーを引き起こします。一方、スライディングウィンドウ推論は古いコンテキストを破棄することでスループットを回復させますが、長期的な一貫性を犠牲にし、シーンを再訪問した際にモデルが新しいコンテンツを幻視したり、ドリフトしたりする原因となります。最近の観察では、これらのモデルにおける KV キャッシュが、世界記憶の創発的な形態として機能することが示唆されていますが、履歴全体へのアテンションにかかる完全な計算コストなしに、この記憶に効率的にアクセスする方法は未解決のままです。

手法：WorldKV
著者は、自己回帰型ビデオ世界モデルにおける効率的な長期記憶を可能にするトレーニングフリーのフレームワーク「WorldKV」を提案します。これは既存のスライディングウィンドウ推論アーキテクチャの上に動作し、「World Retrieval（世界検索）」と「World Compression（世界圧縮）」という 2 つの相補的なコンポーネントを導入します。

1. World Retrieval（世界検索）:

   • メカニズム: 排除された KV キャッシュチャンクを破棄する代わりに、WorldKV はそれらを GPU または CPU メモリに保存します。モデルがシーンを再訪問する際、カメラまたは行動の対応関係に基づいてシーンに関連するチャンクを選択的に検索し、ネイティブなアテンションウィンドウに再度挿入します。
   • 統合: 検索されたチャンクは、再エンコードやアーキテクチャの変更なしにネイティブに挿入されます。このフレームワークは検索アルゴリズムに依存せず、カメラ/行動ベースの類似性やアテンションベースの戦略をサポートします。
   • アテンションウィンドウ構造: アクティブなウィンドウは 4 つの領域に分割されます。シンク KV キャッシュ（初期フレーム）、検索された KV キャッシュ（保存された履歴から）、最近の KV キャッシュ（直前のフレーム）、および現在のノイズ除去チャンクです。

2. World Compression（世界圧縮）:

   • 動機: 全ての排除されたキャッシュを保存することはメモリ集約的です。著者は、時間的に隣接するフレームが実質的な視覚コンテンツを共有し、ほぼ重複する KV キャッシュを生成することを観察しました。
   • メカニズム: 各保存されたチャンク内（例：3 フレーム）で、最初のフレームを「アンカー」として指定します。システムは、非アンカーフレームのキーとアンカーフレームのキーとの間のキー - キーコサイン類似度を計算します。
   • プルーニング: アンカーと高い類似度を持つトークンは冗長として除去され、低い類似度のトークン（新たに露出した領域や動的コンテンツを表すもの）は保持されます。
   • 効率性: このプロセスにより、各チャンクは元のサイズの約半分（例：アンカーと他の 2 フレームから 25% のトークンを保持）に圧縮され、固定されたメモリ予算内で約 2 倍の履歴を格納可能にします。

主要な貢献

• World Retrieval: カメラ/行動またはアテンションベースのシグナルを使用して、排除された KV キャッシュチャンクを保存し、選択的に検索するフレームワーク。アーキテクチャの変更なしにモデルが長期コンテキストにアクセスすることを可能にします。
• World Compression: キー類似度に基づくプルーニングメカニズムにより、チャンクあたりの保存量を約 50% 削減。同じメモリ制約下で履歴カバレッジを 2 倍にし、再訪問時の忠実度を維持または向上させます。
• トレーニングフリーの効率性: このアプローチは微調整、蒸留、専用メモリモジュールを必要とせず、凍結されたバックボーン上で推論時に完全に動作します。

実験結果
このフレームワークは、異なるスケールの 2 つの自己回帰型ビデオ世界モデルで評価されました。Matrix-Game-2.0（13 億パラメータ、短いシーケンスでトレーニング済み）とLingBot-World-Fast（140 億パラメータ、長いビデオでトレーニング済み）です。

• スループット対忠実度: WorldKV は、フル KV アテンションのメモリ忠実度と同等かそれ以上の性能を発揮しつつ、スループットを約2 倍維持します。LingBot-World-Fast においては、ネイティブモードであるフル KV と同等の性能を達成しながら、はるかに高い FPS を実現します。Matrix-Game-2.0 においては、スライディングウィンドウおよびフル KV ベースラインの両方を上回ります。特に、フル KV は分布外に蓄積されたキャッシュへのアテンションによる累積エラーにより性能が低下しますが、WorldKV の選択的検索はこれを緩和します。
• SOTA との比較: WorldKV は、メモリ固有のトレーニングを行わずに、LPIPS、PSNR、FID などの指標において、メモリトレーニング済みベースライン（WorldPlay、Yume-1.5）を上回ります。
• アブレーション研究:
  • チャンク内圧縮: 中程度の圧縮（アンカー＋他のトークンの 50% を保持）が最適な性能をもたらします。アンカーのみを保持すると忠実度が低下し、すべてのトークンを保持すると限界効用逓減が見られます。
  • チャンク間カバレッジ: 固定予算に収まるようより多くのチャンクを圧縮する（例：6 チャンク→3 チャンク予算）ことは、フル解像度でより少ないチャンクを検索する場合と比較して忠実度を向上させ、広範な歴史的カバレッジの価値を浮き彫りにします。

意義と主張
本論文は、WorldKV が、明示的に長いシーケンスでトレーニングされていないモデルであっても、KV キャッシュ自体が根本的かつ創発的な世界記憶の形態であることを実証していると主張しています。検索と圧縮を通じてこの既存のリソースを効率的に管理することで、WorldKV はリアルタイム推論制約を遵守しつつ、一貫したシーンの再訪問を維持する持続的な世界モデルを実現します。著者は、これが 3D 再構築のレイテンシや専用メモリモジュールのトレーニングオーバーヘッドなしに達成されていることを強調しています。

限界
著者は、WorldKV が推論時の管理手法であるため、視覚的忠実度は基盤となる事前学習モデルの生成品質によって制限されることを指摘しています。極めて長いロールアウトでは、依然としてエラーの蓄積が発生する可能性があります。さらに、CPU オフローディングは VRAM 使用量をさらに削減できますが、現在のホスト - デバイス間の転送レイテンシにより、その特定の構成ではリアルタイム生成が不可能であり、これは今後の課題として残されています。