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

Geometric Factual Recall in Transformers

May 12, 2026

著者: Shauli Ravfogel, Gilad Yehudai, Joan Bruna, Alberto Bietti

cs.AI

要旨

トランスフォーマー言語モデルは、事実の関連付けをどのように記憶しているのだろうか。一般的な見解では、内部の重み行列を埋め込みのペアに対する連想記憶とみなし、パラメータ数が事実の数に比例して増加するとされている。我々は、学習された埋め込みが関係構造を直接符号化し、MLPが質的に異なる役割を果たすという、代替的な幾何学的記憶の理論的・実証的説明を展開する。単層トランスフォーマーがサブジェクトから共有属性集合へのランダムな全単射を記憶しなければならない制御された設定において、対数的な埋め込み次元で十分であることを証明する。すなわち、サブジェクト埋め込みは、関連する属性ベクトルの線形重ね合わせを符号化し、小さなMLPは連想キー・バリューマッピングとしてではなく、ReLUゲーティングを介して関連属性を抽出する関係条件付きセレクターとして機能する。我々はこれらの結果をマルチホップ設定、すなわち「xの妻の母親は誰か？」のような関係クエリの連鎖に拡張し、チェーン・オブ・ソートを用いる場合と用いない場合の構成を提供する。これらの構成は証明可能な容量・深さのトレードオフを示し、情報理論的な下界と一致する。経験的に、勾配降下法は予測された構造を正確に持つ解を発見する。訓練後、サブジェクト埋め込みが適切に再初期化されると、MLPはまったく新しい全単射にゼロショット転移し、特定の事実集合を記憶するのではなく、汎用的な選択機構を学習したことが明らかになる。

English

How do transformer language models memorize factual associations? A common view casts internal weight matrices as associative memories over pairs of embeddings, requiring parameter counts that scale linearly with the number of facts. We develop a theoretical and empirical account of an alternative, geometric form of memorization in which learned embeddings encode relational structure directly, and the MLP plays a qualitatively different role. In a controlled setting where a single-layer transformer must memorize random bijections from subjects to a shared attribute set, we prove that a logarithmic embedding dimension suffices: subject embeddings encode linear superpositions of their associated attribute vectors, and a small MLP acts as a relation-conditioned selector that extracts the relevant attribute via ReLU gating, and not as an associative key-value mapping. We extend these results to the multi-hop setting -- chains of relational queries such as ``Who is the mother of the wife of x?'' -- providing constructions with and without chain-of-thought that exhibit a provable capacity-depth tradeoff, complemented by a matching information-theoretic lower bound. Empirically, gradient descent discovers solutions with precisely the predicted structure. Once trained, the MLP transfers zero-shot to entirely new bijections when subject embeddings are appropriately re-initialized, revealing that it has learned a generic selection mechanism rather than memorized any particular set of facts.