以下は、提示された論文「HCLSM: Hierarchical Causal Latent State Machines for Object-Centric World Modeling」の技術的な要約です。

HCLSM: 物体中心の世界モデルのための階層的因果潜在状態機械

1. 問題提起

現在の動画から未来の状態を予測する「世界モデル」には、以下の 3 つの根本的な限界があります。

1. 物体の混同（Flat Latent Representations）: 既存のモデル（V-JEPA や DreamerV3 など）は、すべての物体を単一の構造化されていない潜在ベクトルにエンコードします。これにより、個々の物体（例：コップとテーブル）の状態が分離されず、因果関係や独立した運動を追跡することが困難です。
2. 時間スケールの欠如: 物理現象は連続的な軌道（ミリ秒単位）、離散的なイベント（衝突など）、長期的な戦略（分単位）など、複数の時間スケールで同時に発生します。単一の時間スケールでこれらを扱うことは非効率的です。
3. 因果構造の欠如: 「何が何を引き起こしたか」という因果関係を明示的にモデル化していないため、反事実的推論（「もしもっと強く押したらどうなるか？」）や計画立案ができません。

既存の物体中心モデル（Slot Attention など）は空間分解能を持ちますが、時間的なダイナミクスや因果推論を統合したシステムは存在しませんでした。

2. 提案手法：HCLSM

HCLSM は、**「物体中心分解」「階層的時間ダイナミクス」「因果構造学習」**の 3 つの原則を統合した新しい世界モデルアーキテクチャです。

アーキテクチャの 5 層構造

1. 知覚層（Perception）: ViT（Vision Transformer）エンコーダを用いて動画フレームをパッチ埋め込みに変換します。
2. 物体分解層（Object Decomposition）:
   • Slot Attention: 画像パッチを「スロット（潜在的な物体表現）」にグループ化します。
   • 空間ブロードキャストデコーダ（SBD）: 各スロットが独立して画像領域を再構成するように強制します。これにより、スロットが特定の物体に対応するよう学習されます。
   • GNN（グラフニューラルネットワーク）: スロット間の相互作用をエッジ重みとして学習し、因果的な関係性をモデル化します。
3. 階層的ダイナミクス層（Hierarchical Dynamics）: 3 つのレベルで時間変化を処理します。
   • Level 0（連続物理）: 選択的状態空間モデル（Selective SSM）を用いて、フレーム間の連続的な物体の運動をモデル化します。
   • Level 1（離散イベント）: イベント検出器が状態変化を検知し、その瞬間のみスパース・トランスフォーマー（Sparse Transformer）が動作します。
   • Level 2（抽象目標）: 圧縮されたトランスフォーマーを用いて、高レベルの目標や戦略を推論します。
4. 因果推論層（Causal Reasoning）: GNN のエッジ重みと NOTEARS 制約（有向非巡回グラフ制約）を用いて、物体間の因果的依存関係を学習します。
5. 継続的メモリ層（Continual Memory）: ホプフィールドネットワークや EWC（Elastic Weight Consolidation）を用いて、長期記憶を維持します。

2 段階トレーニングプロトコル

モデルが「分散表現（スロットの集合）」に依存して予測を楽にする（＝物体分解を怠る）のを防ぐため、以下の 2 段階で学習を行います。

• Stage 1（空間再構成）: 予測タスクを無効化し、SBD による再構成損失のみで学習します。これにより、各スロットが特定の空間領域（物体）を「所有」するよう強制されます。
• Stage 2（ダイナミクス予測）: 物体が分解された状態で、JEPA 風の未来予測タスクを有効化します。これにより、個別の物体の未来状態を正確に予測する能力が学習されます。

3. 主要な貢献

1. 統合アーキテクチャ: 物体スロット、3 段階の時間階層（SSM + スパース Transformer + 目標 Transformer）、GNN による因果相互作用を単一の微分可能なモデルに統合しました。
2. SBD と ViT 特徴の再構成: DINOSAUR の手法を踏襲し、ピクセルではなく凍結された ViT 特徴を再構成ターゲットとする空間ブロードキャストデコーダを採用し、実世界のロボット動画での教師なし物体分解を可能にしました。
3. 2 段階トレーニング: 空間分解を先に確立し、その後に時間ダイナミクスを学習させることで、スロットの専門化（Specialization）を成功させました。
4. 高性能実装: 選択的 SSM スキャン用のカスタム Triton カーネルを開発し、逐次的な PyTorch 実装と比較して38 倍の高速化を達成しました。
5. オープンソース化: 8,478 行の Python コード、51 のモジュール、171 の単体テストを含む完全なコードベースを公開しました。

4. 実験結果

• データセット: Open X-Embodiment データセット内の「PushT」タスク（ロボットアームで T 字型のブロックを目標へ押し込む）を使用。
• モデル規模: 68M パラメータ。
• 性能:
  • 次状態予測の MSE（平均二乗誤差）: 0.008。
  • SBD 再構成損失: 0.0075。
  • 2 段階トレーニングを採用することで、スロットが明確に空間領域を分割し、イベント境界（衝突など）を検出する能力が獲得されました。
• 高速化: NVIDIA H100 上でのカスタム Triton カーネルにより、SSM 推論がバッチ処理のボトルネックから外れ、フォワードパス時間の 5% 以下に削減されました。

5. 意義と限界

意義

HCLSM は、世界モデルが「物体」「多様な時間スケール」「因果関係」という世界の構造を反映すべきであるという前提に立ち、これらを初めて統合したアーキテクチャです。特に、2 段階トレーニングにより、物体分解と時間予測のトレードオフを解決した点は画期的です。また、実ロボットデータでの成功は、このアプローチが実世界タスクに応用可能であることを示しています。

限界と今後の課題

• スロット数: 現在のところ、32 個のスロットがすべて有効なまま（死なない）であり、3 つの物体のシーンに対して過剰なスロット数が割り当てられています。
• 因果推論の精度: 明示的な因果グラフ学習（DAG）はスパース化正則化の影響でエッジがゼロになりやすく、GNN の重みによる間接的な相互作用構造に依存しています。
• スケーラビリティ: 大規模モデル（262M パラメータ以上）やマルチ GPU 環境での学習において、bfloat16 精度での数値的不安定性（NaN 発生）が課題となっています。

結論

HCLSM は、構造化された世界モデルを実現するための堅実な基盤を提供します。物体中心の表現、階層的な時間処理、因果推論を組み合わせることで、従来のフラットな潜在空間モデルを超えた、計画や反事実的推論が可能なエージェントへの道筋を示しました。コードと評価スイートは GitHub で公開されています。