BRo-JEPA: Learning Modular Arithmetic in Latent Space

技術要約: BRo-JEPA – 潜在空間における剰余演算の学習

問題提起
本論文は、ディープラーニングにおける根本的な問い、すなわち「ニューラルネットワークは真に抽象的な代数的規則を学習しているのか、それとも単に訓練パターンの記憶に過ぎないのか」を調査している。現在のモデルは既知の範囲内の数学的タスクには長けているが、未知の数値範囲や演算への外挿（エクストラポレーション）には失敗することが多い。著者らは、標準的な教師あり学習モデルや既存のJEPA（Joint Embedding Predictive Architecture）モデルは、基礎となる算術論理ではなく、統計的な相関関係に依存していると断定している。具体的には、MNISTの数字を状態として用いた剰余演算（例：$y = (x + k) \mod 10$）を課した場合、標準的なモデルは既知の演算には適合するものの、未知の演算（厳密なゼロショット汎化）には失敗する。

手法
著者らは、生のピクセル出力を予測するのではなく、コンテキストとなる数字の画像と算術演算（アクション）が与えられたときに、ターゲットとなる数字の潜在表現を予測するJEPAフレームワークを利用している。研究では以下のいくつかのアーキテクチャのバリエーションを比較している：

1. 標準的なベースライン: クロスエントロピー損失と加法的な構成アクション埋め込み（additive compositional action embeddings）を用いた教師ありモデル（MLPおよびResNet-18）。これらの埋め込みは、埋め込み空間に算術構造を強制するために、操作をベースベクトル（例：$op(k) = k \cdot v$）のスカラ倍として表現する。
2. 加法的な埋め込みを用いたJEPA: 同じ加法的な埋め込みを使用し、「構成一貫性損失（Composition Consistency Loss）」で拡張されたJEPAモデル。この補助的な損失は、合成された操作（例：$+2$）を一段階で適用した結果が、原始的な操作を逐次的に適用した結果（例：$+1$ をその後 $+1$）と同じになることを強制する。
3. BRo-JEPA (Block-Rotation JEPA): コアとなる貢献は、標準的な制約のない多層パーセプトロン（MLP）を、操作条件付きの**ブロック回転予測器（block-rotation predictor）**に置き換えた新しい予測器である。
   • メカニズム: 回転位置エンコーディング（RoPE）に着想を得たこの予測器は、潜在空間のダイナミクスを一連の2次元平面回転として扱う。潜在ベクトル $z$ と操作 $k$ に対して、予測値は $z_p = R_{block}(k \cdot \theta) \cdot z$ となる。
   • 帰納バイアス: このアーキテクチャは、モジュロ10演算の円環構造を潜在空間に明示的に課すものである。回転角 $\theta$ は固定（$2\pi/10$ に設定）することも、学習可能にすることもできる。
   • バリアント: 著者らは、すべての潜在ブロックが単一の角度を共有する単一周波数回転（Single-Frequency Rotation, SFR）と、各ブロックが独立した角度を持つ多周波数回転（Multi-Frequency Rotation, MFR）をテストしている。
   • 正則化: 「潜在崩壊（latent collapse）」（異なる入力が区別不可能な表現にマップされる現象）を防ぐため、モデルは分散・不変性・共分散正則化（VICReg）と指数移動平均（EMA）ターゲットエンコーダを採用している。

主な貢献

1. テストベッドの定義: 潜在ワールドモデルにおけるルール学習を研究するための制御されたテストベッドとして、MNIST剰余演算を導入した。
2. 失敗の分析: 加法的な埋め込みを用いた教師ありベースラインおよびJEPAモデルが、既知の演算には適合するものの、厳密なゼロショット外挿には失敗することを実証した。これらは演算の符号は捉えるが、大きさや構成的な幾何学をエンコードすることには失敗する。
3. BRo-JEPAの提案: モジュロ演算の循環構造に適合するように、操作条件付きのブロック回転予測器を備えたJEPAスタイルの潜在ダイナミクスモデルを開発した。
4. 幾何学的分析: BRo-JEPAが潜在表現を構造化された幾何学的多様体（円環多様体）へと組織化し、モデルが剰余演算のトポロジーを内部化することを可能にすることを分析した。

結果
実験は、演算 $\{-1, +1\}$ での訓練と、未知の演算 $\{\pm 2, \dots, \pm 9\}$ での評価を用いて行われた。

• ベースラインの性能: 教師ありベースラインは既知の演算に対して高い精度（~97-98%）を示したが、未知の演算に対してはほぼゼロの精度となった（ResNet-18では2.42%）。PCA可視化により、これらのモデルは未知の演算を既知の演算の表現空間の中に恣意的に詰め込んでいることが明らかになった。
• 加法的な埋め込みを用いたJEPA: これらのモデルは高い*ロールアウト（展開）*精度（$+1$ を繰り返し適用して $+k$ をシミュレートする場合）を示したが、未知の演算の直接的なゼロショット予測には失敗した（精度 ~10%）。構成一貫性損失はこれを ~99% まで向上させたが、モデルが訓練中に間接的に合成された演算にさらされていたため、「弱いゼロショット」とみなされる。
• BRo-JEPAの性能:
  • ゼロショット精度: 多周波数回転（MFR）を用いたResNetベースのBRo-JEPAは、未知の演算に対して**99.46%**のゼロショット精度を達成した。
  • ロールアウト精度: モデルはマルチステップのロールアウトにおいても**99.46%**の精度を達成し、学習された潜在ダイナミクスの安定性を確認した。
  • 潜在構造: PCA可視化により、BRo-JEPAが潜在空間を、数字のクラスが逐次的に配置された連続的な楕円/円形多様体へと強制的に組織化したことが確認された。MFRバリアントは、SFRと比較して、より分離されたクラスプロトタイプ（より低い非対角コサイン類似度）を生み出し、優れた汎化性能をもたらした。
  • 学習可能な角度: 回転角を学習可能にした場合、MFRは効果的な汎化へと迅速に収束したが、SFRは顕著な遅延と不安定性を示し、しばしば劣悪な高調波解へと収束した。

意義と主張
本論文は、潜在ワールドモデルは記号的な変換ルールを学習できるが、それはアーキテクチャの帰納バイアスが問題の構造と一致している場合に限られると主張している。結果は以下を示唆している：

• 標準的なニューラル予測器は、未知の演算に対して構成的な幾何学を保持するために必要な制約を欠いている。
• 潜在空間におけるブロック回転を通じて、モジュロ演算の循環的な性質を明示的にモデル化することで、訓練中に特定の演算にさらされることなく、強力なゼロショット汎化が可能になる。
• このアプローチは、表現学習を超えて算術以外の領域へ広がる有望な方向性を示しており、適切な変換構造を潜在遷移モデルに組み込むことが、抽象的なルールの学習において極めて重要であることを示唆している。

著者らは、自らの結果が、既知の巡回群構造を持つ高度に制御されたMNIST剰余演算の設定に特有のものであることを認め、謙虚な姿勢を保っている。彼らは、このアーキテクチャが、同様の構造的事前分布なしに、任意の記号的推論タスクや未知の変換群へと自動的に汎化できると主張しているわけではない。