技術的サマリー：UWM-JEPA

問題定義

結合埋め込み予測アーキテクチャ（JEPAs）は、生観測を再構成するのではなく、潜在空間において将来の状態を予測し、表現学習と生成的デコーディングを効果的に分離する。しかし、部分的に観測可能な環境において、標準的なベクトル値の潜在表現は、隠れた継続性に対する「信念」を保持するために必要な内部構造を欠いている。現在の観測履歴が隠れた状態を一意に決定できない場合、未来は表現空間内の単一の点ではなく、互換性のある隠れた状態にわたる分布となる。

指数移動平均（EMA）ターゲットに対して訓練された標準的なベクトル予測器は損失を最小化できるが、損失関数だけでは、どの不確実性を保持すべきか、どの不確実性を平均化すべきか、あるいは反事実的な行動が信念状態をどのようにシフトさせるべきかを指定しない。その結果、ベクトル潜在 JEPAs は、「ブラインドロールアウト」（新しい観測なしに将来の状態を想像すること）中に、行動に敏感なダイナミクスを維持したり、タスクに関連する情報を保持したりすることに失敗することが多い。

手法

著者らは、潜在状態を決定論的な点ではなく、明示的な信念状態として扱う JEPA フレームワークの改変であるUWM-JEPA（Unitary World Model JEPA）を導入する。

1. 信念状態潜在

ベクトルではなく、オンラインエンコーダ $E_\theta$ は、観測されたコンテキストを双部分潜在空間 $\mathcal{H}_T = \mathcal{H}_S \otimes \mathcal{H}_E$（システム $\otimes$ 環境）上の密度行列 $\rho_t$ にマッピングする。

• 次元: システム因子（$d_{sys}=8$）は損失とプローブによって読み取られ、環境因子（$d_{env}=2$）は隠れたモードとシステム - 環境相関を運ぶレジスターとして機能する。
• 制約: 密度行列はエルミート、半正定値、単位トレースであるように制約される。これらは構築によって強制される。すなわち、初期状態は最大混合状態であり、観測更新は正の測定演算子を適用し、予測ステップはユニタリ共役を使用する。

2. ユニタリ予測器

予測器はユニタリ共役を介して潜在状態を進める：
$$\hat{\rho}_{t+k|t} = U_k \rho_t (U_k)^\dagger$$
ここで、$U = \exp(-iH\Delta t)$ であり、$H$ はエルミート生成子である。

• 行動条件付け: 生成子は行動 $a$ に対して線形である：$H(a) = H_0 + a H_1$。これにより、ユニタリ軌道を行動によって誘導できる。
• 訓練目的: モデルは標準的な JEPA スケフォールドを使用する。停止勾配を伴う EMA ターゲットエンコーダ $E_\xi$ と、システム因子への部分トレースおよび投影ヘッドの後で評価される潜在空間マッチング損失である。
• 反事実的ターゲット: 予測器が行動チャネルを使用することを学習することを保証するため、訓練ターゲットは、教師強制された観測軌道ではなく、反事実的な行動シーケンス $\tilde{a}$ 下でのシミュレーターロールアウトによって生成される。

3. 理論的保証

核心的な理論的主張は結合状態の忘却しないことである。ユニタリ共役は等長写像であるため、システム - 環境密度行列のスペクトル（固有値）、純度、フォン・ノイマンエントロピーを、ブラインドロールアウト中に正確に保存する。

• 限界: この保証は完全な結合状態に適用される。部分トレースによって得られる縮小システム状態はユニタリではないため、情報がシステム - 環境相関に移行し、結合状態が非散逸的であっても、線形プローブからアクセスできなくなる可能性がある。

主要な結果

1. 行動感応性には反事実的ターゲットが必要

• 教師強制の失敗: 教師強制されたターゲット（観測された未来に一致する）で訓練された場合、行動チャネルは不活性化する。行動依存項と基底ダイナミクス項の比率は $\|H_1\|/\|H_0\| \approx 0.03$ まで崩壊する。ターゲットはすでに行動の影響を含んでいるため、予測器は行動に関係なくターゲットに一致することを学習する。
• 反事実的成功: ターゲットをサンプリングされた反事実的行動下でのシミュレーターロールアウトに置き換えると、行動感応性が回復し、$\|H_1\|/\|H_0\| = 1.00 \pm 0.15$ となる。
• 摂動制御: 反事実的ターゲットで訓練された UWM-JEPA（UWM-JEPA-CF）は、行動が摂動（反転、シャッフル、否定）されるにつれて予測精度が単調に低下するのに対し、パラメータ整合された LSTM-JEPA-CF はすべての条件下で多数派クラスの精度（0.53）に崩壊する。

2. ブラインドロールアウトと短期保持

• 隠れた速度指標: 隠れた速度の符号を予測するために 5 ステップ先のシミュレーションを必要とするタスク（ターゲット観測がマスクされている）において、UWM-JEPA-CF は0.77 の精度を達成する。整合された LSTM-JEPA-CF は失敗し、すべての例に対して多数派クラスを予測する。
• 情報保持（$\Delta R^2$）: ブラインドロールアウト中の教師プローブに対する情報損失を測定する：
  • 短期間（$k=1, 3$）において、UWM-JEPA は $R^2$ が 10 ポイント未満減少する（それぞれ 0.023 および 0.069）。
  • LSTM-JEPA ベースラインは著しく多く失う（それぞれ 0.413 および 0.678）。
  • 訓練された時間範囲（$k=5$）において、LSTM は再整列するが、UWM は中間ステップで優れた性能を維持する。
  • 長期間（$k>10$）において、両モデルとも劣化し、長期的保持が未解決の課題であることを示している。

3. コンテキスト表現は駆動要因ではない

• エンコーダ同等性: フロートエンコーダ表現に対する保持された線形プローブは、UWM-JEPA と LSTM-JEPA が同一に機能することを示す（$R^2 = 0.901 \pm 0.032$ 対 $0.894 \pm 0.021$、$p=0.70$）。
• 結論: 性能の格差は、UWM-JEPA が優れたコンテキストエンコーダを持っていることによるものではない。分離は、ロールアウト中に信念状態構造を維持する予測器の能力に完全に起因する。

意義と主張

本論文は、潜在状態を隠れた継続性に対する信念（密度行列）として明示的に扱い、予測器が結合状態スペクトルを散逸させることなくこの信念を進化させるように制約された、最初の JEPA 系ワールドモデルとして UWM-JEPA を主張する。

• 限定的な範囲: 著者らは明示的に、これは汎用プランナーでも、普遍的に強力なエンコーダでも、長期的計画のための競争力のある解決策でもないとしている。「忘却しない」という性質は、結合状態に対してのみ正確であり、必ずしも下流タスクで使用される縮小状態に対しては成り立たない。
• 中核的貢献: この研究は、部分的な観測性下で JEPA ワールドモデルが想像するためには、凍結されたコンテキストエンコーディング能力だけでなく、潜在幾何学と予測器ダイナミクスが重要であることを実証している。
• 実践的含意: この研究は、反事実的ターゲットに対する訓練が、行動に敏感な予測器を学習するために不可欠であることを浮き彫りにしており、この知見は特定のユニタリパラメータ化を超えて拡張される。

結果は、構造化された部分的に観測可能なオシレーターおよび CartPole 風の環境で検証されている。著者らは、現在の手法は反事実的訓練のためにシミュレーターアクセスを必要とし、まだ構造化されていない視覚予測や隠れマルコフ連鎖には拡張されていないことを認めている。