技術的概要：トランスフォーマーにおける幾何学的事実想起

問題提起

大規模言語モデル（LLM）は、膨大な事実的関連付けの集合を記憶し、想起する能力を有している。この能力を理解する上での核心的な問いは、トランスフォーマーアーキテクチャがこれらの関連付けをどのように記憶するかである。既存の理論的枠組みは、これを連想記憶の問題としてモデル化することが多く、内部の重み行列（具体的には MLP またはアテンション機構）が入力と出力の埋め込みのペアを、ランダムでほぼ直交するベクトルとして記憶すると考える。この見方によれば、$N$個の事実を記憶するために必要なパラメータ数は、$N$に対して線形にスケーリングする。

しかし、経験的観察は異なるメカニズムを示唆している：トランスフォーマーは幾何学的記憶を利用する可能性があり、学習された埋め込みがデータの関係構造を直接符号化する。本論文は、トランスフォーマーが、蛮力的な連想記憶に依存するのではなく、埋め込み空間を幾何学的に構造化することで、事実的関連付けをより効率的に記憶できるかどうかを調査する。著者らは、$N$個の主体、$R$個の関係、および共有された属性の集合を含む制御された設定に焦点を当て、モデルが主体 - 関係ペアに関連する属性を予測するか、または多段クエリ（例：「$x$の妻の母親は誰か？」）においてこれらの関係を構成する必要がある状況を対象とする。

手法

本論文は、合成データおよび実世界データ上での理論的構築と経験的検証を組み合わせる。

理論的枠組み

著者らは、埋め込みを学習可能なパラメータとして扱う場合の、事実想起タスクを解くトランスフォーマーの容量を分析する。彼らは以下の 2 つの領域を区別する：

1. 単段想起：主体と関係が与えられた場合の属性の予測。
2. 多段想起：中間出力なし（直接予測）で $k$ 個の関係 sequentially に構成するか、または思考連鎖（CoT）を用いて構成するか。

理論的分析は、学習可能な埋め込みを持つトランスフォーマーに対する明示的な解の構築に依存する。これらの構築の鍵となるのは、主体埋め込み内での線形重ね合わせの使用である。事実を異なる直交ベクトルとして記憶する代わりに、主体埋め込みは、その $R$ 個の可能な属性すべての表現を含むブロックベクトルとして構築される。MLP は、ルックアップテーブルとして機能するのではなく、ReLU ゲーティングを介して関連する属性ブロックを抽出する関係条件付きセレクターとして機能する。

著者らはまた、$R$-正則辺彩色グラフ上の数え上げ論法を用いて、$k$段想起に必要なパラメータ数と埋め込み次元の情報理論的下界を導出する。

経験的設定

勾配降下法がこれらの理論的構造を発見するかどうかを検証するために、著者らは合成データ上で広範な実験を行う：

• データセット：共有属性プールにマッピングされる、$R \in \{2, \dots, 16\}$個の関係を持つ $N=4096$個の主体。
• アーキテクチャ：一様または学習されたアテンションと 2 層 MLP を持つ単層トランスフォーマー。
• 対照実験：学習可能な埋め込みを持つモデルと、入力エンティティ埋め込みを初期化時にランダムに固定（先行研究の固定埋め込み領域を模倣）する対照モデルを比較する実験。
• 分析：学習後の分析には、重ね合わせ構造を確認するための線形読み出し、MLP の選択性をテストするための因果介入、および学習されたメカニズムの一般性をテストするための「固定・交換」転移実験が含まれる。
• 実世界プロービング：著者らは事前学習済み LLM（例：Qwen、Llama、Phi）をプローブし、主体の隠れ状態が理論的予測に類似した線形関係符号化を示すかどうかを確認する。

主要な貢献

1. 共有属性による対数記憶

本論文は、MLP が小さく関係条件付きである場合、単層トランスフォーマーが埋め込み次元 $d = O(R \log N)$ を用いて $N \times R$個の事実を記憶できることを証明する。

• メカニズム：主体埋め込みは、その $R$個の関連属性ベクトルの線形重ね合わせを符号化する。
• MLP の役割：MLP は重みにマッピングを記憶するのではなく、ReLU ゲーティングを介してクエリされた関係に対応する特定の属性ベクトルを抽出するセレクターとして機能する。
• 意義：これにより、連想記憶モデルにおける $N$に比例する線形から、MLP のサイズを増大させることなく、対数レベルの埋め込み次元へと必要とされる次元を削減する。

2. 多段想起における容量と深さのトレードオフ

著者らは、思考連鎖（CoT）なしの $k$段事実想起に対する根本的なトレードオフを確立する：

• CoT なし：$k$段クエリを解くには、指数関数的な埋め込み次元（$d = \tilde{O}(R^k)$）が必要か、または事実数に比例する MLP の幅（$O(N \cdot R)$）が必要である。これは、中間出力がない場合、モデルが単一の隠れ状態内に評価木全体を符号化しなければならないことを示す情報理論的下界（定理 4.2）によって証明される。
• CoT あり：モデルに中間関係ステップを出力させる（思考連鎖）ことで、このボトルネックを打破する。$d = \tilde{O}(R + k)$の埋め込み次元と $\tilde{O}(R)$の MLP 幅を持つ単層トランスフォーマーであれば、任意の $k$に対して十分である。CoT は、多段問題を単段問題の系列に実質的に還元し、埋め込み空間に完全な $k$段木を記憶する必要性を回避する。

3. 幾何学的構造の経験的検証

実験は、勾配降下法が予測された幾何学的構造を自然に発見することを確認する：

• 容量閾値：学習可能な埋め込みの場合、閾値 $d = \Theta(R \log N)$が経験的に観測されるのに対し、固定埋め込みの場合は $d = \Theta(N^R)$（または同様の多項式スケーリング）が必要であり、構造化された埋め込みの理論的優位性が確認される。
• 重ね合わせ：線形読み出しにより、主体埋め込みが属性ベクトルの線形重ね合わせとしてよく近似されることが示される。
• 汎用的選択：因果介入により、MLP が関係固有のセレクターとして機能することが確認される。重要なのは、主体埋め込みを新しいランダムな全単射を符号するように再初期化した場合でも、固定された MLPが高いゼロショット精度を達成することである。これは、MLP が特定の事実の記憶ではなく、汎用的な選択メカニズムを学習したことを示している。

意義と主張

本論文は、トランスフォーマーにおける幾何学的事実記憶に関する理論的かつ経験的な説明を提供し、支配的な連想記憶の見方に対する代替案を提示する。

• 効率性：埋め込みを関係トポロジーを直接符号化するように構造化することで、トランスフォーマーは、ランダム埋め込み連想記憶モデルと比較して、はるかに少ないパラメータ（$N$に対する対数）で事実想起を達成できる。
• 想起のメカニズム：これらの知見は、事実想起が構造化された幾何学的表現を通じて実装されており、MLP が汎用的な関係条件付きセレクターとして機能することを示唆している。これは、埋め込み空間が適切に再初期化されれば、モデルが新しい事実に対して一般化できる理由を説明する。なぜなら、選択メカニズムは特定の事実値から切り離されているからである。
• 思考連鎖の役割：この研究は、多段推論における思考連鎖の有効性に対する理論的説明を提供する。それは、CoT が計算を自己回帰ステップに分散させることで、直接 $k$段予測に内在する容量のボトルネックを回避し、浅いモデルが深い推論タスクを解くことを可能にすることを示している。
• 実世界への関連性：事前学習済み LLM における線形関係符号化の観察は、これらの幾何学的メカニズムが単なる理論的アーティファクトではなく、自然言語で訓練されたモデルに存在することを示唆しており、先行する経験的研究で観察された関係解読の線形性などの現象を説明する可能性がある。

著者らは、彼らの設定では関係を不透明なトークン、属性を任意のエンティティとして扱っているのに対し、実世界の知識は豊かな意味構造を含むことを指摘し、控えめな立場を維持している。彼らは、このような規則性を持つデータにこの幾何学的説明を拡張することが、今後の研究の方向性であると示唆している。