技術的概要：MapFormers

問題定義

現在のトランスフォーマーベースの AI システムは、分布内タスクにおいて印象的な性能を示す一方で、分布外（OOD）のシナリオに直面すると著しい脆さを露呈する。これらは、人間や動物が容易に処理する新しい状況への汎化に失敗することが多い。著者らは、この格差の原因を、生物学的な脳が認知地図（環境の抽象的で不変な構造を、その具体的で可変な内容から分離する内部モデル）を構築することを可能にする構造的バイアスの欠如にあると帰結している。

認知地図をモデル化しようとする既存の試み、例えば RNN 内で動作依存行列を用いるトールマン・アイヒェンバウム機械（TEM）などは、2 つの決定的な限界に苦しんでいる：

1. スケーラビリティの欠如：位置を更新するために逐次的な行列乗算を必要とし、並列処理を阻害する。
2. 事前指定された動作：動作の集合（例えば、基本方向）がデータから学習されるのではなく、ハードコードされる必要があることが多い。

さらに、標準的な位置符号化（RoPE など）は静的であるか、未知の環境における経路統合に必要な動作の動的な結合をサポートしないヒューリスティックに依存している。

手法

本論文は、教師なしで次のトークン予測を通じて認知地図を学習し、経路統合を実行するように設計されたトランスフォーマーアーキテクチャのクラスであるMapFormersを導入する。中核的な革新は、リー群論に由来する入力依存行列を用いて位置符号化を更新することにある。

理論的枠組み

• リー群形式：著者らは、動作をリー群の要素（特に回転のための $SO(2)$ などのコンパクト群）としてモデル化する。すべての可能な動作に対して個別の行列を学習するのではなく、モデルはリー代数生成子を学習する。
• 指数写像による経路統合：動作は、リー代数における無限小生成子として表現される。経路統合（動作の系列の結合）は、代数内でこれらの生成子を合計（線形蓄積）し、その後、リー群における有限変換を得るために単一の指数写像を適用することによって実行される。
  • 可換な地図（例えば、2 次元グリッドナビゲーション）の場合、これにより経路統合が累積和を経て並列に計算され、単一の行列指数関数が適用されることを可能にする。
  • このアプローチは、帰納的バイアスを動作の数から、代数の構造的複雑さに軽減する。

アーキテクチャのバリエーション

本論文は、絶対的位置符号化と相対的位置符号化を、異なるメモリモデルと統合する MapFormers の 2 つの変種を提案する：

1. MapEM（エピソード記憶）：

   • 認知地図を表すために絶対的位置埋め込みを使用する。
   • 内容（観測）と構造（位置）を、分離されたニューラル集団に解きほぐす。
   • 内容と位置の結合に対してアテンションを計算する（$Q_G = \text{vec}(Q^T \cdot P)$）。これは、位置符号化が内容とは独立して更新されるアテンションマスクとして機能する。
   • 特定の状態の堅牢な想起を必要とするタスクに適している。

2. MapWM（ワーキングメモリ）：

   • 相対的位置符号化（RoPE に類似するが動的）を使用する。
   • 入力依存の回転行列（$R_{\theta PI}$）を介してニューラル活動（キーとクエリ）を異なる部分空間へ回転させることで、構造を暗黙的に表現する。
   • 回転角 $\theta$ は、入力トークンの低ランク射影から導出され、モデルがどのトークンが動作（位置を更新）として機能し、どのトークンが観測（位置を変化させない）として機能するかを学習することを可能にする。

入力依存メカニズム

両方のアーキテクチャは、各トークンに対する「統合期間」($\Delta_t$) を決定するために 2 段階の射影を採用する：

1. 内部射影：入力トークンを低次元の動作表現にマッピングする。
2. 外部射影：この表現をブロックごとの統合期間に射影する。
   最終的な回転角は $\theta_t = \omega \odot \Delta_t$ であり、ここで $\omega$ は学習可能な角速度ベクトルである。経路統合は、これらの角度の累積和（$\theta_{0 \to t} = \text{cumsum}(\theta_t)$）として計算される。

主要な貢献

1. MapFormers アーキテクチャ：リー代数生成子に基づく入力依存位置符号化を用いて構造を内容から解きほぐすことで、認知地図を学習する新規なトランスフォーマー変種。
2. 並列経路統合：可換リー群（例えば、$SO(2)^k$）を活用することで、モデルは並列経路統合を達成し、以前の RNN ベースの認知地図モデルの逐次的なボトルネックを克服する。
3. メモリモデルの統合：絶対的/相対的位置符号化とエピソード/ワーキングメモリの間の理論的リンクを提供し、両者が同じリー理論的原理から導出され得ることを示す。
4. 非可換地図へのスケーラビリティ：複数の生成子を用いて非可換な問題（例えば、家系図）に枠組みを拡張するが、これには逐次的な経路統合が必要となる。

実験結果

著者らは、MapFormers を、異なる認知能力を対象とした 3 つの形式的タスクで評価した：

1. 選択的コピー（ゲーティング）：

   • タスク：ダミートークンを無視しながらシーケンスをコピーする。
   • 結果：MapFormers（EM と WM の両方）は完全な OOD 汎化を達成したが、標準的なベースライン（RoPE、TAPE）は失敗した。これは動的なゲーティングメカニズムを学習する能力を示している。

2. グリッドナビゲーション（1 次元および 2 次元）：

   • タスク：交互の動作とランダムな観測のシーケンスにおける観測を予測し、モデルに潜在トポロジを推論し、位置を経路統合させる。
   • 結果：MapFormers は、長さや密度が変化する OOD シーケンスでほぼ完全な精度を達成したが、ベースライン（CoPE や PathAtt を含む）は汎化に失敗した。
   • ニューラル分析：モデルは、観測トークンをゼロ回転（単位変換）へ、動作トークンを特定の回転方向へ射影することを学習し、実質的に内容から構造を解きほぐしていた。

3. Dyck-2（ネストされた階層）：

   • タスク：スタックの追跡を必要とする括弧シーケンスの妥当な続きを予測する。
   • 結果：MapFormers は、単一層のみを使用してトレーニングで見たものよりも長く深いシーケンスに完全に汎化した。ベースラインは複数の層を必要としたにもかかわらず、OOD データでは依然として失敗した。
   • メカニズム：モデルは、異なる括弧タイプに対して直交次元を、開き/閉じ括弧に対して反対方向を使用することで、スタックを位置符号化に直接埋め込んだ。

4. 自然言語スケーリング：

   • タスク：OpenWebText での事前学習。
   • 結果：12 層の MapWM モデルは、RoPE ベースライン（18.79 対 19.14）に対して一貫したパープレキシティの改善と、より良い長さの外挿を示した。しかし、BLIMP 構文ベンチマークでは、MapWM は RoPE に対して有意な改善を示さず、可換な経路統合のみがすべての構文モデル化に十分ではない可能性を示唆している。

意義と主張

本論文は、MapFormers が、ヒューリスティックに依存するのではなく、経路統合を通じて潜在構造を捕捉することにより、体系的な OOD 汎化を可能にする重要な構造的バイアスを提供すると主張している。

• 堅牢性：モデルは、抽象的な関係を内容から解きほぐすことが、見かけのシーケンス長さ、深さ、環境密度への堅牢な汎化を可能にすることを示している。
• スケーラビリティ：可換地図上での並列計算の使用は、自然言語の結果が示すように、これらの原理を大規模な実世界ドメインにスケーリング可能にする。
• 生物学的妥当性：このアプローチは、「何か」と「どこか」を分解する生物学的認知地図と現代の深層学習の間の格差を埋め、環境の幾何学的基盤を理解することで新しい問題を通じて推論できる AI システムへの道筋を提供する。

著者らは、非可換な経路統合は並列化できず、現在の可換制約が言語における特定の複雑な構文構造のモデル化を制限する可能性があるという限界について謙虚に述べている。彼らは、より複雑な関係推論のためのハイブリッドアーキテクチャや非可換拡張を将来の研究で探求することを提案している。