技術的サマリー：LeJEPA はいつ世界モデルを学習するか？

問題定義

本論文は、自己教師あり学習（SSL）における根本的なギャップに取り組んでいる。すなわち、Joint-Embedding Predictive Architectures（JEPAs）が視覚、動画、計画のための表現学習において経験的な成功を収めているにもかかわらず、これらの表現が「世界モデル」を構成するという理論的保証はないという点である。具体的には、学習された表現が、これらの要因を混同するのではなく、世界の潜在的な構造（位置、速度、色など）を忠実に復元する条件は未だ不明である。

真の自由度を撹乱する表現（例えば、物体の位置と色を混同する）は、狭義のタスクでは良好に機能する可能性があるが、世界が変化する際や構成的な汎化に用いられる際には失敗する。核心的な問いは以下の通りである：学習された表現は、いつ世界の潜在変数を線形的に復元するか？ 著者らはこの性質を線形識別可能性と定義し、信頼性の高い線形プロービングと最適な潜在空間計画のための必要条件としている。

手法と理論的枠組み

世界モデル

著者らは、「世界」を観測されない潜在変数 $z \in \mathbb{R}^n$ を持つシステムとして形式化する。観測データ $x$ は、未知の、おそらく高度に非線形な混合関数 $g$ によって生成され、$x = g(z)$ となる。学習目的には、独立した加法ノイズ遷移を持つ定常過程によって生成された正のペア $(z, z')$ が含まれる。

本論文は、以下の特性を持つ特定の世界のクラスに焦点を当てている：

1. 独立性: 潜在変数とその遷移は、次元間で独立である。
2. 定常性: 潜在変数の周辺分布は、ビュー間で一定に保たれる。
3. 加法ノイズ: 遷移は $z' = m(z) + \eta$ に従い、ここで $\eta$ は独立なノイズである。

重要なのは、著者らがこのクラスをガウス世界（$z \sim \mathcal{N}(0, I_n)$）に特化している点である。これらの仮定の下、遷移ダイナミクスは Ornstein–Uhlenbeck（OU）過程に対応する：$z' = \rho z + \sqrt{1-\rho^2}\eta$。

学習器：LeJEPA

学習アーキテクチャはLeJEPA（LeCun の JEPA）であり、エンコーダ $f$ を訓練して埋め込み $h(z) = f(g(z))$ を生成する。これは、2 つの項からなる損失関数を最小化することで行われる：

1. アライメント損失: 正のペアの埋め込み間の距離を最小化する、$E[\|h(z') - h(z)\|^2]$。これは表現が低速で予測可能な特徴を捉えるように促す。
2. SIGReg（スケッチされた等方性ガウス正則化）: 埋め込み分布を明示的に等方性ガウス分布 $h(z) \sim \mathcal{N}(0, I_n)$ として正則化し、表現の崩壊を防ぐ。

最適化問題は以下の通りである：
$$\min_h E[\|h(z') - h(z)\|^2] \quad \text{s.t.} \quad h(z) \sim \mathcal{N}(0, I_n)$$

理論的アプローチ

著者らは、遷移演算子 $T$（$(Th)(z) = E[h(z')|z]$ で定義される）のスペクトル解析を採用している。

• エルミート展開: ガウス潜在変数の場合、遷移演算子の固有関数はエルミート多項式である。固有値は、$d$ を多項式の次数として、$\rho^d$ として幾何級数的に減衰する。
• スペクトルペナルティ: アライメント損失は、遷移下でより速く相関を失うため、線形成分よりも高次（非線形）の成分をより重く罰する。
• シュトゥルム・リウヴィル理論: 一般的な（非ガウシアンの）分布の場合、固有関数はシュトゥルム・リウヴィル理論によって特徴づけられる。最初の非定数固有関数は常に単調であるが、線形性は分布がガウスである場合にのみ保証される。

主要な貢献

1. 順方向：ガウス世界における線形識別可能性

定理 1 は、潜在変数がガウスであり、遷移が OU 過程である場合、ガウス性制約付きの LeJEPA 目的関数の任意の最小化器は、真の潜在変数の線形直交変換でなければならないことを証明する。すなわち、ある直交行列 $Q$ に対して $h(z) = Qz$ となる。

• メカニズム: この証明はスペクトル分解に依存しており、任意の非線形性は正のペア間の相関を厳密に減少させる。アライメント損失はこの相関を最大化しようとするため、大域的最適解は表現を純粋に線形（1 次のエルミート多項式）に強制する。
• 結果: 学習された表現は、大域的な回転まで潜在変数を復元し、遷移ダイナミクスは潜在空間内で正確に復元される。

2. 逆方向：ガウス分布の一意性

定理 2 は、ガウス分布が、この世界クラスにおいて線形識別可能性が成立する唯一の潜在分布であることを確立する。

• メカニズム: シュトゥルム・リウヴィル理論を用いて、遷移演算子の最初の固有関数がアフィン（線形）であることを要求すると、潜在分布のスコア関数が線形でなければならないことを示す。この条件はガウス分布を一意に特徴づける。
• 含意: 潜在変数が非ガウスである場合、LeJEPA は線形識別可能性を保証できない。学習された表現は必然的に潜在変数の非線形変換となる。

3. 近似識別可能性

定理 3 は、目的関数が近似してのみ満たされる場合（すなわち、不完全なアライメントと不完全なホワイトニング）の復元誤差に対する定量的な境界を提供する。

• 誤差は、アライメントギャップ $\delta$ とホワイトニング誤差 $\epsilon$ に対して滑らかにスケーリングする。
• この境界は、アライメントギャップが復元誤差の支配的な要因であることを強調しており、一方、近似ホワイトニングは比較的「容易に」達成可能であることを示している。

4. 最適潜在計画

定理 4 は、線形識別可能性を制御理論と結びつける。エンコーダが線形識別可能（$h(z) = Qz$）である場合、直交変換に対して不変なコスト（例えば、目標までのユークリッド距離など）を持つ任意の最適制御問題は、学習された潜在空間内で最適に解けることを証明する。潜在空間における最適行動と価値関数は、コスト関数の修正を必要とせずに、真の世界に直接マッピングされる。

実験結果

著者らは、いくつかの実験設定において理論を検証した：

1. 2 次元シミュレーション: 様々な非線形混合関数（スパイラル、せん断、RealNVP）を使用し、潜在変数がガウスである場合、LeJEPA が回転まで潜在変数を正常に復元することを確認し、定理 1 を裏付けた。
2. 高次元スケーリング: $N=1024$ 次元までの実験により、バッチ統計量正則化器（SIGReg、VICReg）がほぼ完全な線形識別可能性（$R^2 > 0.999$）を維持するのに対し、対照的手法（InfoNCE）はカーネル幅の不一致によりスケールとともに劣化することが示された。
3. 分布アブレーション: 潜在分布の形状を（重尾から一様まで）掃引した結果、線形復元はガウス（$\alpha=2$）で鋭くピークに達し、非ガウス分布では劣化することが示され、定理 2 が検証された。
4. ピクセルベースのロボティクス制御: DeepMind Control Suite の Reacher を用い、合成ガウス遷移（OU）と実 RL 軌跡での訓練を比較した。
   • OU 条件: 高い線形識別可能性（$R^2 \approx 0.95$）を達成。
   • 軌跡条件: 実世界のポリシーは非ガウスで異方性のある周辺分布を誘起し、識別可能性が著しく低下した（$R^2 < 0.5$）。
5. 計画検証: Reacher タスクにおいて、ガウス訓練エンコーダの学習された潜在空間内での直線補間をデコードすると、関節空間ではほぼ直線的な経路（最適計画）となる。一方、軌跡訓練エンコーダは曲がった最適でない経路を生み出し、定理 4 が確認された。

意義と主張

本論文は、経験的に成功したレシピを数学的保証に変える、JEPAs に対する最初の識別可能性結果を提供すると主張している。

• 基礎的保証: 本作業は、LeJEPA が特定の明確な条件（ガウス潜在変数、定常加法ノイズ）の下で、世界の構造（潜在変数とダイナミクス）を証明的に復元することを確立している。
• ガウス性の役割: 本論文は古典的な ICA の物語を逆転させる。線形 ICA では、ガウス性はソース分離が失敗する分布である。しかし、この非線形で自己教師ありな設定では、ガウス性が線形識別可能性を可能にする唯一の分布である。
• 実践的含意: この理論は、データをガウス領域に保つ等方ランダムウォークを近似する探索戦略が、有用な世界モデルの学習に不可欠である理由を説明する。一方、データを低エントロピー領域に崩壊させる目標指向ポリシーは、識別可能性を破る可能性がある。
• 計画の基盤: 線形識別可能性が回転不変なコストによる最適計画を可能にすることを証明することで、本論文は、真の物理座標へマッピングし直すことなく、制御のために学習された潜在空間を使用する理論的根拠を提供する。

著者らは限界を指摘している。すなわち、理論はエンコーダ次元が潜在次元と一致し、集団レベルの最適性に依存していることである。しかし、彼らは保証が有限サンプルと最適化誤差に対して滑らかに劣化することを示し、また Lean 4 定理証明器を用いて主要な定理を形式的に検証している。