技術的サマリー：世界モデルのための識別可能トークン対応

1. 問題定義

トランスフォーマーベースの世界モデルは、過去の状態と行動の系列をトークンストリームとして扱い、未来の状態を予測することで、視覚強化学習（RL）において高い性能を示してきました。しかし、これらのモデルは長期的なロールアウト中に時間的不整合に頻繁に見舞われます。一般的な失敗モードには、オブジェクトの重複、消滅、および異なるオブジェクトへの変容が含まれます。

著者らは、これらの誤差の主要な原因を特定しました：既存のアプローチは、次のフレームの予測を純粋なトークン生成問題として扱っています。それらは、時間を超えたトークン間の対応関係を明示的にモデル化することに失敗しています。現実的な視覚環境では、連続するフレームのトークンは、持続し移動する同じ基盤となる実体に対応することがよくあります。アイデンティティを追跡することなくすべてのトークンをゼロから再生成することを強制されると、モデルは持続的な構造を維持することに苦慮し、結果として時間経過とともに脆弱なアイデンティティ保持と誤差の累積を招きます。

2. 手法：識別可能トークン対応（ITC）

本論文は、潜在トークン対応変数を持つ構造化された確率推論問題として次のフレームの予測を定式化するトランスフォーマー世界モデル、**識別可能トークン対応（ITC）**を提案します。

中核メカニズム

ITC は、すべての次のフレームのトークンを新たに生成するのではなく、各次のフレームのトークンに対して、以下のいずれかを行うかを決定します：

1. 前のフレームのトークンからコピーする（アイデンティティを保持）。
2. 新しいトークンとして生成する（変化や新規コンテンツを処理）。

この割り当ては、**最適輸送（OT）**問題としてモデル化されます。

技術的アーキテクチャ

1. トークン化と埋め込み：

   • 視覚的観測は、最近傍パッチルックアップトークナイザーを使用して離散トークンに変換されます。
   • トークン埋め込みは、**3 次元回転位置エンコーディング（RoPE）**によって拡張されます。標準的な 1 次元 RoPE と異なり、この定式化は時間、垂直、水平座標を独立してエンコードし、モデルが空間と時間の両方にわたる局所的な関係構造を捉えることを可能にします。

2. 最適輸送ソルバー：

   • モデルは、ソース頂点が前の状態トークン（$u$）とトランスフォーマーによって予測された候補の次の状態トークン（$\tilde{p}$）を表し、デスティネーション頂点が最終的な次の状態トークン（$\hat{u}$）を表すグラフを構築します。
   • 前のトークンと予測トークンの間の関係をスコアリングするアフィニティ行列が計算されます。この行列には以下が含まれます：
     • $A^{(prev)}$：前のトークンと予測との間のアフィニティをスコアリングし、物理的な妥当性を強制するために空間距離によってペナルティを課します（例：オブジェクトが恣意的に遠くへテレポートすることはできない）。
     • $A^{(gen)}$：前のトークンを再利用する代わりに、新しい生成トークン（ワイルドカード）を受け入れることの「ボーナス」をスコアリングし、一定のペナルティ subject にします。
   • ソルバー（Sinkhorn アルゴリズムを使用）は、ソースからデスティネーションへの質量移動のコストを最小化する輸送計画を計算します。

3. 二値化とデコーディング：

   • 連続的な輸送計画は、貪欲な二値化手順（アルゴリズム 2）を使用して二値割り当て行列に変換されます。これにより、厳密な 1 対 1 対応が保証されます。
   • デコーディング規則： 次のフレームの各位置について、輸送計画が前のトークンを割り当てた場合、そのトークンがコピーされます。計画が「生成」トークンを割り当てた場合、モデルはトランスフォーマーの予測分布からサンプリングします。

訓練ループ

世界モデルと方策は同時に訓練されます。世界モデルは「想像」（合成軌道）を生成し、これらは実際の環境データとともに、近接方策最適化（PPO）を介して方策の更新に使用されます。

3. 主要な貢献

• トークン対応の定式化： 本論文は、トークン対応を最適輸送問題として明示的にモデル化することにより、次のフレームの予測を再定式化し、「再利用」する持続的実体と「生成」する新規コンテンツのタスクを分離します。
• 最適輸送の統合： Sinkhorn 経由で微分可能でありながら（二値化経由で）離散的なメカニズムを導入し、前のトークンのコピーと新規生成の選択を行い、幻覚を効果的に削減します。
• 3 次元位置エンコーディング： 3D RoPE の実装により、モデルは時間的および空間的差異を区別できるようになり、これはビデオ系列における移動オブジェクトの追跡に不可欠な要素です。

4. 実験結果

著者らは、ITC を 4 つの挑戦的なベンチマーク、Craftax-classic、Craftax、MinAtar、Atari 100Kで評価しました。

• Craftax-classic： ITC はリターンで72.5%、スコアで**35.6%**を達成し、67.4% のリターンと 27.9% のスコアを達成した以前の最先端（Dedieu et al., 2025）を大幅に上回りました。また、より迅速な収束も示しました。
• Craftax： より複雑なこの環境において、ITC はリターンで 7.09%、スコアで 2.40% を達成し、Simulus や Dedieu et al.を含むベースラインを上回りました。
• MinAtar： ITC は、4 つのゲーム（Asterix、Breakout、Freeway、SpaceInvaders）すべてにおいて、モデルベース RL の以前の最先端を上回りました。
• Atari 100K： 26 個のゲームのスイートに適用されたところ、ITC は四分位平均（IQM）で 1.092、最適性ギャップで 0.376 の新しい最先端性能を達成し、トークンベースの世界モデルである Simulus を上回りました。
• 予測精度： 10,000 回の遷移のデータセットにおいて、ITC はベースラインと比較して全体的な予測精度を 3.39% 向上させ、移動する生き物を含む遷移では 3.45% の顕著な改善を示しました。
• 定性的改善： 視覚分析により、DreamerV3 や Dedieu et al.などのベースラインモデルで見られるオブジェクトの重複や消滅などの一般的なアーティファクトが ITC によって排除されることが確認されました。
• 効率性： この手法は過大なオーバーヘッドを導入せず、エンドツーエンドの訓練時間をわずか 2.8% 増加させるのみでした。

5. 意義と主張

本論文は、時間を超えたトークン間の対応関係を明示的にモデル化することにより、ITC が長期的な視野にわたってオブジェクトのアイデンティティを維持できないという世界モデルにおける広範な問題を解決すると主張しています。著者らは、このアプローチによりモデルが「前のフレームと強い対応関係を持つトークンを選択的に再利用」できるようになり、すべてのトークンをトランスフォーマーに再生成させることだけに依存するのではなく、フレーム間の類似性を活用できると述べています。

その意義は、グリッドベースのオープンワールドから標準的なアタリゲームまで多様な環境において最先端の性能を達成しつつ、計算効率を維持している点にあります。著者らは、ITC が重複や消滅のような時間的不整合を効果的に防ぎ、高品質な想像と改善された方策性能につながると結論付けています。