ワールドモデル2026/8/18arXiv:2608.17542

ガウス分布は不要：JEPAワールドモデルのための対照逆動学

No Gaussian Required: Contrastive Inverse Dynamics for JEPA World Models

JEPAワールドモデルの表現崩壊を防ぐ新しい正則化手法を提案。逆動学ヘッドを訓練時のみ追加し、行動識別タスクで崩壊を防ぐ。複雑なタスクで従来法を大きく上回る性能を達成。

詳しい要約

1. どんなもの？

本論文は、JEPA (Joint-Embedding Predictive Architecture) に基づく世界モデルの学習において、表現の崩壊 (collapse) を防ぐ新しい手法を提案する。従来のJEPAは将来の埋め込みを予測するが、定数エンコーダという自明な解が存在するため、SIGRegのような正則化で潜在分布を等方的ガウス分布に強制していた。提案手法AC-MTMは、逆動力学ヘッドを訓練時のみ追加し、Action-NCE損失を用いて、各潜在遷移がその行動をバッチ内の他の行動から識別できるようにする。これにより、崩壊したエンコーダではこの識別タスクを解けないことを利用して、分布に依存しない反崩壊信号を提供する。テスト時には逆動力学ヘッドを破棄し、エンコーディング、前方予測、プランニング、計算量は従来のLeWMと同一になる。

2. 先行研究と比べてどこがすごい？

先行研究のJEPA (LeWM) はSIGRegにより潜在分布を等方的ガウス分布に強制するが、これは環境とは無関係に表現の形状を規定する。AC-MTMは、このような事前に定義された幾何学的制約を課さず、遷移データ自体から反崩壊圧力を得る点が新しい。また、ターゲットネットワーク、ストップグラディエント、事前学習済みエンコーダ、再構成目的を一切必要としない。OGBench Visual Sceneタスクでは、SIGRegの58.0%±2.0%に対し、AC-MTMは80.0%±2.0%の成功率を達成し、20〜24ポイントの改善を示した。

3. 技術・手法の肝は？

手法の核心は、LeWMの前方潜在予測目的を維持しつつ、訓練時のみ逆動力学ヘッドを追加し、Action-NCE損失で訓練することである。Action-NCEは、各潜在遷移がその行動をバッチ内の他の行動と区別する対比学習タスクであり、崩壊したエンコーダではこのタスクが失敗することを証明する。逆動力学ヘッドは訓練後に破棄され、テスト時の推論パイプラインは従来と同一になる。また、この手法が成立するための行動空間と可観測性の仮定を特徴付けている。

4. どうやって有効だと検証した？

4つの標準的なピクセル制御タスクにおいて、一致したプランニングプロトコルでスクラッチから安定して訓練でき、平均でSIGRegと同等の性能を示した。より難しい多物体OGBench Visual Sceneタスクでは、AC-MTMは80.0±2.0%の成功率を達成し、SIGRegの58.0±2.0%を上回り、各訓練シードで20〜24ポイントの改善が見られた。また、50エピソードのランダムポリシー実行によるベースライン推定は52%であった。

5. 議論はある？

要旨からは、AC-MTMがSIGRegよりも優れている理由として、事前に定義された幾何学的制約がボトルネックになる可能性が示唆されている。しかし、標準的なピクセル制御タスクでは平均で同等であり、タスクの複雑さによって優位性が異なる可能性がある。また、提案手法が成立するための行動空間と可観測性の仮定についての詳細な議論は要旨には含まれておらず、不明な点も残る。

6. 次に読むべき論文は？

要旨で参照されている研究として、LeCun (2022)、Assran et al. (2023)、Bardes et al. (2022; 2024) が挙げられる。これらはJEPAや反崩壊機構に関する基礎的な論文であり、次に読むべきである。また、関連する対比学習や逆動力学モデルの研究も有用であろう。

※ AIが要旨から生成した要約です。正確性は原文をご確認ください。

著者: Jack Boylan, Chris Hokamp

分類: cs.LG, cs.AI

原文アブストラクト

Joint-Embedding Predictive Architectures (JEPAs) learn world models by predicting future embeddings, but the objective admits a trivial solution of a constant encoder, so every practical system adds an anti-collapse mechanism (LeCun, 2022; Assran et al., 2023; Bardes et al., 2022; 2024). LeWorldModel (LeWM) prevents collapse with SIGReg, a regularizer that forces the latent distribution to match an isotropic Gaussian: the representation is stabilized by prescribing what it must look like, independently of the environment it models. We argue that the anti-collapse pressure can instead come from the transition data itself. Action-Contrastive Masked Transition Modeling (AC-MTM) keeps LeWM's forward latent-prediction objective and adds a training-only inverse-dynamics head trained with Action-NCE: each latent transition must identify the action that produced it among the other actions in the batch, a discrimination task that a collapsed encoder provably fails. The inverse branch is discarded after training, leaving test-time encoding, forward prediction, planning, and compute identical to LeWM. On four standard pixel-control tasks under a matched planning protocol, AC-MTM trains stably from scratch and matches SIGReg on average. On the harder multi-object OGBench Visual Scene task, results are consistent with the prescribed geometry becoming a bottleneck: AC-MTM reaches 80.0$\pm$2.0% success versus 58.0$\pm$2.0% for SIGReg, improving by 20-24 points in each training seed. A single 50-episode random-policy run gives a 52% baseline estimate. Contrastive inverse dynamics thus provides a distribution-free anti-collapse signal that requires no target network, stop-gradient, pretrained encoder, or reconstruction objective, and we characterize the action-space and observability assumptions under which it holds. We make our code available at https://github.com/jackboyla/action-contrastive-jepa

arXivで原文を読む →arXiv概要ページ