世界モデル/計画2026/8/17arXiv:2608.16287

SCALE: 正しい幾何学におけるJEPA計画のための状態校正潜在埋め込み

SCALE: State-Calibrated Latent Embeddings for JEPA Planning in the Right Geometry

エンドツーエンドの世界モデルLeWMの潜在表現に、タスク関連状態空間との距離整合性を持たせることで、計画性能を向上させる手法SCALEを提案した。

詳しい要約

1. どんなもの？

SCALE (State-CAlibrated Latent Embeddings) は、JEPA (Joint-Embedding Predictive Architecture) に基づく世界モデルの計画性能を向上させるための手法。エンドツーエンドで学習される LeWorldModel (LeWM) の潜在表現に対して、タスク関連の状態空間における距離と潜在空間の距離を相関させる正則化を加えることで、DINO-WM のような事前学習済み特徴量が持つ「主要成分に状態情報が多く含まれる」という幾何学的特性を誘導する。これにより、ユークリッド距離に基づく計画コストが状態情報をより強く反映するようになり、計画精度が向上する。

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

先行研究では、非崩壊表現を得るために、DINO-WM のように事前学習済み特徴量を継承する方法と、LeWM のようにエンドツーエンドで anti-collapse 正則化 (SIGReg) を用いて学習する方法があった。これらの手法はタスクによって相補的な強みを持つが、LeWM の表現は DINO-WM に比べて主要成分に状態情報が少なく、計画コストが高分散方向に支配されやすいという問題があった。SCALE は、LeWM のエンコーダを置き換えずに、その表現に DINO-WM の好ましい幾何学的特性を誘導する点が新しい。

3. 技術・手法の肝は？

SCALE の核心は、学習時にサンプリングされたペアの潜在表現間の距離と、標準化されたタスク関連状態空間での距離との相関を最大化する軽量な正則化項を追加すること。これにより、潜在空間の距離が状態空間の距離と整合するようになり、高分散方向に状態情報が多く含まれるようになる。エンコーダは LeWM のものをそのまま使用し、計画時には追加の計算コストは発生しない。

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

5つのタスク、3つの計画ソルバー、5つの計算予算にわたって評価し、SCALE が LeWM と比較してすべてのタスク・ソルバー平均を改善することを示した。また、潜在表現から状態を回帰する制御実験 (latent-to-state regression) を行い、SCALE と同等のデコーダビリティを持つが、潜在空間と状態空間の距離整合はほぼ変化せず、計画改善も一貫しないことを確認した。これにより、SCALE の効果が単なる状態情報の存在ではなく、幾何学的特性の変化によるものであることを示した。

5. 議論はある？

要旨からは、SCALE の正則化がタスク関連状態の定義に依存する可能性や、状態空間の標準化方法の影響、また、DINO-WM と比較した場合の性能差の詳細は不明。さらに、SCALE が LeWM の表現に与える影響が、計画以外の下流タスクにどのように影響するかは議論の余地がある。

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

要旨で参照されている DINO-WM と LeWorldModel (LeWM) の論文が次に読むべき論文として挙げられる。また、関連する JEPA の基礎論文や、表現の崩壊を防ぐ正則化手法 (SIGReg) に関する研究も有用である。

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

著者: Jiaming Hu, Yan Zheng, Tian Wang

分類: cs.LG

原文アブストラクト

Joint-embedding predictive world models plan by scoring predicted terminal embeddings against a goal embedding using a cost defined on the representation itself. Two prominent strategies for obtaining non-collapsed representations are to inherit a pretrained feature space, as in DINO-WM, and to learn an embedding end to end with anti-collapse regularization, as in LeWorldModel (LeWM) with SIGReg. These strategies show complementary strengths across tasks. Although task-relevant state is decodable from the full embeddings of both models, DINO-WM's leading principal components usually retain substantially more state information than LeWM's. Because Euclidean planning costs are dominated by high-variance directions, this difference affects how strongly state can influence candidate selection. We propose SCALE (State-CAlibrated Latent Embeddings) to give the end-to-end LeWM representation the favorable geometric property observed in DINO-WM. SCALE induces this property by correlating sampled pairwise latent distances with distances in a standardized task-relevant state space, without replacing LeWM's learned encoder. Across five tasks, three planning solvers, and five compute budgets, SCALE improves every task--solver average over LeWM. A latent-to-state regression control matches or exceeds SCALE's full-embedding decodability yet leaves latent--state distance alignment essentially unchanged and yields less consistent planning gains. SCALE adds a single lightweight training-time regularizer and no planning-time overhead. These results show that planning depends not only on whether task-relevant information is present, but also on whether it shapes the geometry consumed by the planner.

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