🤖 machine learning

Sub-JEPA: Subspace Gaussian Regularization for Stable End-to-End World Models

本論文は、全環境空間ではなく複数のランダム部分空間にガウス正則化を適用することで、世界モデル構築における Joint-Embedding Predictive Architectures (JEPAs) の安定性と柔軟性を向上させる Sub-JEPA という手法を提案し、これにより優れたバイアス - バランスのトレードオフを実現し、最先端の LeWorldModel を凌駕する性能を達成するものである。

原著者： Kai Zhao, Dongliang Nie, Yuchen Lin, Zhehan Luo, Yixiao Gu, Deng-Ping Fan, Dan Zeng

公開日 2026-05-12

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Kai Zhao, Dongliang Nie, Yuchen Lin, Zhehan Luo, Yixiao Gu, Deng-Ping Fan, Dan Zeng

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

以下は、論文「Sub-JEPA」を平易な言葉と日常的な比喩を用いて説明したものです。

全体像：ロボットに未来を想像させること

ロボットに迷路を navigated させると想像してください。壁のすべてのレンガや床のすべてのほこり粒（これらはデータが多すぎます）をロボットに見せるのではなく、ロボットが「心の地図」を学ぶようにしたいのです。この地図は、自分がどこにいて、どこへ向かっているのかを要約した簡略化されたものです。

AI の世界では、この心の地図をワールドモデルと呼びます。特定の種類のモデルであるJEPAは、現在の地図と行動（例えば「左に曲がる」など）に基づいて、次の心の地図がどのように見えるかを予測するため、効率的であることで人気があります。

問題：「厳しすぎる」教師

この論文は、これらのモデルが現在どのように訓練されているかについて、重大な問題点を指摘しています。

ロボット心の地図を、テストを受ける学生だと想像してください。学生がカンニングしたり、あきらめたり（すべての質問に同じ答えを書くという「崩壊」と呼ばれる問題）しないようにするため、教師は次のようなルールを課します：「あなたの答えは、点の完全な円のように、完全に均等に広がっていなければなりません。」

これが以前の手法であるLeWorldModel (LeWM) が行っていたことです。それは、ロボットの心の地図を、巨大な高次元空間内で完璧で均一なデータの雲のように強制しました。

欠点：
著者らは、このルールが多くの現実世界のタスクには厳しすぎると主張しています。

比喩： 綱渡りをする人を想像してください。彼らの動きは複雑ですが、主に綱という一本の細い線に沿って起こります。もし綱渡りをする人に、綱の周りで完璧な 3 次元の球体を動くよう強制すれば、彼らの自然な動きを制限することになります。彼らは使う必要のない方向に動くよう強制されるのです。
技術的な用語で言えば、この論文は、現実世界のタスクはしばしば「高次元空間（巨大で空っぽの部屋）」の中に存在する「低次元多様体（単純な経路）」上で生きていると述べています。データを部屋全体に均等に満たすよう強制することは、性能を損なう悪いバイアスを生み出します。

解決策：Sub-JEPA（「部分空間」アプローチ）

著者らはSub-JEPAを提案します。ロボットの心の地図を巨大な部屋全体で完璧な球体にする代わりに、その部屋を多くのより小さなランダムな「部分部屋（部分空間）」に分割します。

仕組み：

データをスライスする： ロボットの心の地図を取り、多くの異なるランダムな 2 次元または 3 次元のビューにスライスすると想像してください（彫刻をさまざまな角度から見るようなものです）。
スライスをチェックする： これらの小さなスライスのそれぞれにおいて、モデルがデータがきれいで均一なベルカーブ（ガウス分布）のように見えるかを確認します。
結果： モデルは大きな部屋の中で完璧な球体である必要はありません。これらのより小さくランダムなスライスにおいて「きれいで均一」に見えるだけでよいのです。

利点：
これは、綱渡りをする人に次のように伝えるようなものです：「あなたは 3 次元の部屋全体を満たす必要はありません。横から見たときにバランスが取れて見え、前から見たときにもバランスが取れて見えるようにすればよいのです。」

これにより、安全網（ロボットが退屈で反復的な答えに崩壊するのを防ぐ）を維持します。
しかし、ロボットに人工的で過度に厳格なルールと戦うのではなく、実際のタスクの経路に沿って自然に動く自由を与えます。

実験が示したもの

研究者たちは、この手法を 4 つの異なるビデオゲームのような制御タスク（ブロックを押し出すロボットアームや、部屋を飛ぶドローンなど）でテストしました。

より良い性能： Sub-JEPA は一貫して古い手法（LeWM）を凌駕しました。ロボットはより良く計画し、移動することを学びました。
「ランク」との関連： 彼らはロボットの心の地図がどれほど「複雑」かを測定しました。新しい手法が地図をより単純化（実際のタスクに適合するよりコンパクトな形）することを許したとき、ロボットが賢くなることがわかりました。古い手法は、地図を不必要に複雑にするよう強制していました。
滑らかな経路： ロボットの心の旅を可視化したとき、新しい手法はより直線的で滑らかな線を生み出しました。古い手法では、ロボット的思考が時間とともに揺れ動き、漂流していました。
長期的な安定性： 現実世界を見ずに未来の動きの長いシーケンスを想像するよう求められたとき、新しい手法はより長く正確さを保ちました。古い手法は幻覚を見始め、コースから外れ始めました。

まとめ

Sub-JEPAは、複雑な問題に対するシンプルな解決策です。巨大で抽象的な空間において AI モデルを「完全に均一」にするよう強制することは、しばしば過度に制限的であることに気づいています。代わりに、その空間のより小さくランダムなスライスの中で秩序を検査することで、AI がより自然で効率的かつ安定した世界の理解を学ぶことを可能にします。

これは、魚を完璧な立方体の中で泳がせること（古い手法）と、どの角度から見てもバランスが取れていれば、水の中で自然に泳がせること（新しい手法）の違いです。

技術的サマリー：Sub-JEPA

問題定義

Joint-Embedding Predictive Architectures（JEPAs）は、生ピクセルを再構成することなく将来の潜在表現を予測することで、世界モデルを学習するための枠組みを提供します。しかし、これらのモデルをエンドツーエンドで訓練する際には、重要なバイアス - バランスのトレードオフに直面します。十分な制約がない場合、エンコーダはすべての入力を同一の埋め込みへ写像する自明な解に収束します（エンコーダの崩壊）；逆に、過度な制約は学習された表現の豊かさを抑制します。

最近の研究であるLeWorldModel（LeWM）は、潜在埋め込みに対して等方性ガウス事前分布を適用することでこの問題を解決し、安定したエンドツーエンド訓練を可能にしました。しかし、著者らはこのアプローチに根本的な限界があることを指摘しています：自然な制御タスクは、高次元の環境空間内において低次元多様体上に存在するのが一般的です。高次元の環境空間内で直接グローバルな等方性ガウス事前分布を強制することは、基礎となるダイナミクスの内在的幾何学と一致しない過度に強いバイアスを導入し、低内在次元性を持つタスクのパフォーマンスを阻害する可能性があります。

手法：Sub-JEPA

著者らは、Sub-JEPA（部分空間ガウス正則化）を提案します。これは、ガウス制約を環境空間から複数の低次元部分空間へ移動させることで、バイアス - バラントのフロンティア上で有利な動作点を目指す手法です。

中核メカニズム

直交部分空間射影： $D$ 次元の潜在ベクトル $z$ 全体を正則化するのではなく、 $z$ を $K$ 個の凍結された行直交ランダム射影行列 $\{P_k\}$ onto 射影します。各射影は、潜在ベクトルを次元 $d_s$ （ここで $d_s \approx D/K$ ）の低次元部分空間へ写像します。
$z^{(k)} = P_k z$
多部分空間ガウス正則化：ガウス制約は、各 $K$ $K$ 個の部分空間内で独立して適用されます。
- 各部分空間 $k$ 内では、射影された埋め込みがさらに $M$ 個のランダムな単位ベクトル onto 射影されます。
- これらの1次元射影に対してEpps–Pulley 正規性統計量を計算し、ガウス性を強制します。
- 総正則化損失 $L_{reg}$ は、すべての部分空間およびランダム方向におけるこれらの統計量の平均です。
訓練目的：モデルは、結合された損失を用いてエンドツーエンドで訓練されます。
$L_{total} = L_{pred}(\hat{z}_{t+1}, z_{t+1}) + \lambda L_{reg}(Z)$
ここで、 $L_{pred}$ は潜在状態の予測誤差（MSE）であり、 $\lambda$ は正則化重みです。

設計上の選択

凍結された射影：射影行列は直交するように初期化され、凍結されたままにされます。これにより、正則化器が変化する潜在分布に適応することを防ぎ、制約が安定かつ独立であることを保証します。
直交性：直交射影を使用することで、部分空間が潜在表現の冗長性のないバランスの取れた視点を提供し、幾何学的構造を保持することが保証されます。

主要な貢献

バイアスの特定：本論文は、LeWM におけるグローバルな等方性ガウス事前分布が、低内在次元性を持つタスクに対して過度なバイアスを導入し、タスクの基礎となる幾何学と不一致を生じさせることを特定しました。
Sub-JEPA フレームワーク：著者らは、複数のランダムな低次元の行直交射影ビューに対してガウス正則化を適用することでグローバル事前分布を緩和する手法を提案します。これにより、崩壊防止効果を維持しつつ、潜在空間がダイナミクスの内在次元性に適応することを可能にします。
実証的検証：4 つの連続制御環境における広範な実験により、Sub-JEPA が一貫して LeWM を上回ることを示しました。性能向上は、学習された埋め込みの実効ランクの低下と相関しており、この手法が内在次元性をより適切に尊重していることを確認しました。

実験結果

本手法は、4 つの連続制御ベンチマークで評価されました：Two-Room（低内在次元）、Reacher、PushT、およびOGB-Cube（視覚的に豊かな 3D）。

計画性能：Sub-JEPA はすべての環境で LeWM を一貫して上回りました。改善は、LeWM の全空間ガウス事前分布がタスクの低次元性と最も不一致であったTwo-Roomで最も顕著でした。Sub-JEPA は、事前学習された視覚バックボーンに依存することなく、複数のタスクにおいてPLDM（多目的ベースライン）およびDINO-WM（凍結エンコーダベースライン）をも凌駕しました。
実効ランク分析：LeWM から Sub-JEPA への実効ランクの低下と、計画成功率の増加との間に直接的な相関が観察されました。ランク圧縮が大きいタスク（Two-Room、OGB-Cube）ほど、性能向上が大きかったことが示されました。
物理的解読可能性：PushT環境におけるプロービング実験により、Sub-JEPA の表現は、エージェントの位置やブロックの位置などの物理的状態情報を、LeWM と同等かそれ以上に保持し、非線形解読性が向上していることが示されました。
潜在軌道の質：
- 時間的整合性：可視化により、Sub-JEPA は LeWM の整理されていない進化と比較して、より時間的に整合性の取れた潜在軌道を生成することが示されました。
- 経路の直線化：Sub-JEPA は潜在速度ベクトルにおいてより高い「直線性」を達成し、計画に不可欠な、より滑らかで線形的に発展するダイナミクスを示しました。
- 長視野安定性：オープンループロールアウトにおいて、Sub-JEPA はシーンの幾何学を維持し、LeWM に比べて長視野にわたる空間的ドリフトが少なかったことが示されました。

意義と主張

本論文は、Sub-JEPAが、JEP A ベースの世界モデルにおけるバイアス - バラントのトレードオフに対するシンプルかつ効果的な解決策を提供すると主張しています。ガウス正則化を部分空間へ移動させることで、本手法は以下の点を実現します：

崩壊防止の恩恵を失うことなく、グローバルな等方性事前分布を緩和します。
特に状態空間が低次元であるタスクにおいて、基礎となるダイナミクスの内在次元性とより良く一致します。
構造的制約を単純化し、柔軟性を高めつつ安定性を犠牲にしないことを示すことで、将来の JEP A ベースの世界モデル研究のための強力なベースラインとなります。

著者らは、このアプローチが予測損失と正則化器のみを使用する LeWM の訓練レシピの単純性を維持しつつ、優れた計画性能と表現の質をもたらす、より柔軟な構造的事前分布を導入していることを強調しています。