Predictive Objectives Discard Exogenous Control-Relevant Features: A Controlled Mechanistic Study

技術要約：予測目的関数による外生的かつ制御関連性の高い特徴量の破棄

問題提起

本研究は、Joint-embedding Predictive Architectures (JEPA) や同様の自己予測目的関数における特定の失敗モードを調査している。これらのモデルは、将来の潜在状態を予測することで表現を学習するが、**制御への関連性（control-relevance）**ではなく、時間的な予測可能性（temporal predictability）に対して最適化を行う。本論文は、この最適化バイアスにより、モデルが外生的（exogenous）（エージェントによって制御不能）でありながら、制御に関連する（control-relevant）（最適な行動や報酬を決定づける）特徴量を、たとえそれが容易にエンコード可能であっても破棄してしまうことを提唱している。

ここにある核心的な対立は、以下の2つの理論的枠組みの間にある：

1. 自己予測（Self-Prediction）： スケーラビリティとノイズとなるテクスチャに対する堅牢性を達成するために、「予測不可能な詳細」を破棄する。
2. 双模倣理論（Bisimulation Theory）： 最適な制御のための表現は、予測可能性に関わらず、報酬に関連する未来によって状態を区別する特徴を保持しなければならないと主張する。

本論文は、特徴量の分類における特定の象限、すなわち「制御不能（外生的）だが関連性が高い」特徴量に焦点を当てている。先行研究では、制御可能な特徴量や、無関係な外生的ディストラクター（妨害要素）については広く研究されてきたが、「制御不能かつ関連あり」というケースは、独立した失敗モードとして分離・測定されてこなかった。

方法論

著者らは、特徴の制御可能性（controllability）と関連性（relevance）を独立して変化させる2 × 2 実験デザインを用いた、制御されたメカニスティックな研究を採用している。

• 環境： 主要な実験では、表面形態（解像度、テクスチャ）は異なるが、基礎となるロジックは同一である合成環境（QuadrantEnv、SwitchColorEnv、GridWorldHiddenRuleEnv）を使用している。
  • 特徴のダイナミクス： 特徴は以下の2つの軸によって定義される：
    • 制御可能性： エージェントの行動がその特徴の次状態に影響を与えるかどうか。
    • 関連性： その特徴が報酬を決定するかどうか。
  • 予測可能性のノブ（The Predictability Knob）： パラメータ $p_{repeat}$ が、外生的特徴の時間的な予測可能性を制御する。$p_{repeat} = 0.5$ のとき、特徴は i.i.d. ベルヌーイ分布からの抽出（最大に予測不可能）となり、$p_{repeat} = 1.0$ のときは一定となる。極めて重要な点として、制御関連性は全範囲にわたって固定されており、予測可能性と関連性をデカップリング（分離）している。
• 比較される目的関数： 本研究では、教師あり参照モデルを含む6つの目的関数を評価している：
  1. 再構成（Reconstruction）： ピクセル再構成（ベースライン）。
  2. JEPA： 潜在的な自己予測（ベースライン）。
  3. 行動条件付きJEPA（Action-Conditioned JEPA）： 行動を条件として将来の潜在状態を予測する。
  4. 制御可能性ベースのJEPA（Controllability-Based JEPA）： 逆ダイナミクス・ヘッドを用いて行動を予測する。
  5. 逆ダイナミクス（ランダムポリシー）（Inverse Dynamics (Random Policy)）： ランダムなポリシーの下での行動を予測する。
  6. 報酬接地型JEPA（Reward-Grounded JEPA）： 報酬ラベルを組み込んだ提案手法のバリアント。
  7. 教師あり参照（Supervised Reference）： 正解の（グラウンドトゥルース）特徴ラベルで訓練されたもの（上限値）。
• 指標：
  • 線形プローブ精度（Linear Probe Accuracy）： 固定された潜在表現から特徴をデコードできるかどうかを測定する。
  • InfoNCE相互情報量（InfoNCE Mutual Information）： $I(Z; c)$ を推定し、共有された決定境界なしに情報の保持を確認する。
  • 有効ランク（Effective Rank）： 表現が崩壊（低ランク化）したのではなく、選択的に特徴を落としたものであることを確認する。

主な結果

実験の結果、報酬フリーの予測目的関数と、報酬に基づいた手法との間で、外生的な制御関連特徴量（セル4）に関する顕著な乖離があることが明らかになった。

1. 報酬フリーの予測目的関数の失敗：

   • 評価されたすべての報酬フリー目的関数（JEPA、行動条件付きJEPA、制御可能性ベースのJEPA、およびランダムポリシー下の逆ダイナミクス）は、外生的な制御関連特徴量を保持することに失敗した。
   • 性能： 線形プローブ精度はチャンスレベル付近（$\approx 0.51$）であり、相互情報量はゼロに近い（$\approx 0$ nats）。
   • 選択性： これは表現の崩壊ではない。有効ランクは依然として高い（QuadrantEnv で 38–42）ままであり、表現は豊かであるが、この特定の機能を選択的に破棄したことを示している。
   • 容量への依存性： 潜在次元を16から1024に増やしても、JEPAにおいて特徴は回復しない。これは、失敗がアーキテクチャの容量不足ではなく、目的関数レベルの構造的問題であることを裏付けている。

2. 報酬接地の成功：

   • 報酬接地型JEPAのバリアントは、外生的な制御関連特徴量を選択的に保持し、線形プローブ精度 1.00、および相互情報量 $\approx 0.693$ nats（1ビットの特徴における理論的最大値）を達成した。
   • ラベル効率： 解決策は非常にサンプル効率が高い。遷移のわずか 2% に報酬ラベルを付与するだけで、特徴を回復（精度 $\approx 0.77$）できる。
   • 選択性： すべての特徴を保持する再構成とは異なり、報酬接地型バリアントは、関連するセル（制御可能・関連あり、および外生・関連あり）のみを保持し、無関係なセルを破棄している。

3. 双模倣（Bisimulation）との比較：

   • 外生的な制御関連特徴量に関する理論的な双模倣距離は $1.0$ である（状態は区別されるべきである）。
   • 教師あり参照および再構成のベースラインは、$\approx 1.998$ の潜在クラス分離を実現している。
   • 一方、JEPAは $\approx 0.105$ の分離しか実現しておらず、JEPAの潜在空間の幾何学が、双模倣理論が予測する報酬関連の区別を符号化できていないことを裏付けている。

意義と主張

本論文は、自然なタスクにおける世界モデルのベンチマークを行うことや、一般的な表現学習においてすべての予測目的関数に欠陥があると主張することが目的ではないと明記している。その意義は以下にある：

• 失敗モードの孤立化： 「外生的かつ制御関連」という失敗モードを初めて制御下で分離・測定し、予測可能性と制御関連性が直交し得ることを示したこと。
• メカニスティックな洞察： 純粋な予測目的関数は時間的な予測可能性を最適化するため、その情報が予測不可能である場合、制御に不可欠な情報の体系的な喪失を招く可能性があることを示したこと。
• 選択的な救済策： 最小限の報酬接地（遷移の2%）によって、予測的目的関数の無関係なノイズに対する堅牢性を損なうことなく、失われた情報を選択的に回復できることを示したこと。
• 理論的ギャップ： 双模倣理論が要求するクラス分離と、JEPAの潜在空間が実際に実現している分離との間のギャップを定量化したこと。

著者らは、予測的目的関数は設計された目的（ノイズとなるテクスチャの無視など）には効果的であるが、構造的に、時間的な予測可能性が低い制御関連情報を破棄するリスクを抱えていると結論付けている。このリスクは、最小限の報酬接地によって軽減可能である。