💬 NLP

How Do Transformers Learn to Associate Tokens: Gradient Leading Terms Bring Mechanistic Interpretability

本論文は、トレーニング勾配ダイナミクスに基づいて、ビグラム、トークン交換可能性、およびコンテキストマッピングの合成として初期段階の重みを導出することで、トランスフォーマが意味的関連性をどのように学習するかを説明するメカニズム的解釈を提供する。

原著者： Shawn Im, Changdae Oh, Zhen Fang, Sharon Li

公開日 2026-05-14

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Shawn Im, Changdae Oh, Zhen Fang, Sharon Li

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

現代の AI チャットボットの頭脳であるトランスフォーマーを、巨大で無地のノートブックだと想像してください。トレーニングを開始すると、そのページは空っぽです。あなたが尋ねているこの論文は、シンプルながら深遠な問いを投げかけています：このノートブックは、数百万の文を読むだけで「鳥」と「飛んだ」といった単語をどうやって結びつけるのでしょうか？

完成した複雑な AI を眺めるのではなく、著者たちは学習の最初の数瞬間に焦点を当てました。彼らは数学的な「拡大鏡」を用いて、AI が記憶を更新し始める際に取る最初の数ステップを詳しく観察しました。

以下に、彼らの発見をシンプルなアナロジーを用いて解説します。

1. 「最初のステップ」のショートカット

AI のトレーニングは、文の次の単語を予測することに似ています。通常、AI が記憶を更新する背後にある数学は、信じられないほど複雑で厄介です。

著者たちは、学習のごく初期段階では、AI が複雑な微積分を行う必要はないことに気づきました。それは、データ内にある最も強く、明白なシグナルに従うだけで十分なのです。彼らはこれを「主要項」と呼びます。山登りを始めるハイカーを想像してみてください。出発点では、道は明確で直線的です。高さを稼いでからでないと、道は曲がりくねり複雑になります。著者たちは、驚くほど長い間、AI の「記憶の重み」はこの単純で直線的な初期の経路に非常に近い状態にとどまることを証明しました。

2. 3 つの構築ブロック

この論文は、AI の記憶がランダムなカオスではないことを明らかにしています。むしろ、それは 3 つの特定の種類の「レゴブロック」（著者たちは基底関数と呼んでいます）を積み重ねて作られています。これらのブロックは、AI が読むテキストの統計データから直接導き出されます。

「次の単語」ブロック（2 語連接マップ）:
- アナロジー: 「The」の後に「cat」や「dog」のような名詞がほぼ必ず続くことを、子供が学ぶと想像してください。
- 役割: このブロックは単に記録します。「単語 A が見えたら、次に最も可能性が高い単語 B は何か？」これは、シンプルで即座のつながりを捉えます。
「交換可能」ブロック（交換可能性マップ）:
- アナロジー: 類語辞典を想像してください。文法を崩さずに「car」を「truck」に置き換えられるなら、それらは「交換可能」です。
- 役割: このブロックは、特定の単語が同じ役割を果たすことを学びます。「red」が通常「car」を修飾し、「fast」もまた「car」を修飾する場合、このブロックは「red」と「fast」が同じ単語と付き合っているため、関連しているかもしれないと AI に気づかせます。これは、単なる隣接関係ではなく、単語をその機能によってグループ化します。
「文脈」ブロック（文脈マップ）:
- アナロジー: 「魚」についての物語を読んでいると想像してください。「pond」という単語が「fish」のすぐ隣にないとしても、物語の前半で「water」や「lake」、「boat」に触れられているかもしれません。
- 役割: このブロックは、文のより奥深くを振り返ります。ある単語が特定の文脈（例えば自然についての物語）に現れる場合、たとえ遠く離れていても、その同じ文脈に属する他の単語と関連する可能性が高いことを学びます。

3. それらがどう連携するか

この論文は、AI の内部の「重み」（思考の仕方を決定する数値）が、これら 3 つのブロックの単なる組み合わせであることを示しています。

出力層（次の単語を推測する部分）は、主に**「次の単語」ブロック**です。これは AI が「今見たものに基づいて、次は何が来るか？」と言っている状態です。
アテンション層（何に注目するかを決定する部分）は、「交換可能」ブロックと**「文脈」ブロック**の混合です。これは AI が「『fish』という単語を見つけた。『pond』という単語を振り返るべきだ。なぜなら、これらは同じような物語でよく一緒に現れるからだ。たとえ隣接していなくても」と言っている状態です。

4. 証明：実際に機能する

著者たちは単に数式を書いただけではありません。実在の AI モデルでこれをテストしました。

彼らは子供向けの物語のデータセットで小さな AI をトレーニングしました。
AI が実際に学習した記憶と、彼らの数学的公式を比較しました。
結果: 一致度は驚くほど高く（場合によっては 99% 以上の類似性）、巨大な実世界の AI（Pythia-1.4B）を調べた際にも、同じパターンが成り立ちました。AI は実際には、これら 3 つの単純な統計的構築ブロックを用いて記憶を整理していました。

最大の示唆

この論文は、AI を言語を魔法のように理解する謎めいた「ブラックボックス」として見る必要はないと主張しています。むしろ、その核心において、AI は単に以下のようにして世界を統計的に整理しているに過ぎません。

何が何に続いてくるかを記録する。
同じ役割を果たす単語をグループ化する。
同じ背景物語を共有する単語を結びつける。

これら 3 つの単純なルールを理解することで、ついに機械が「鳥」と「飛んだ」を、「国」と「首都」をどうやって関連付けるのかをどのように学習しているのかを見ることができるようになります。それは魔法ではなく、読んだテキストを要約する非常に構造化された方法なのです。

技術的サマリー：トランスフォーマーはどのようにトークンを関連付けるかを学習するか：勾配の主導項が機械的解釈可能性をもたらす

問題提起

自己注意に基づく大規模言語モデル（LLM）は、事実知識の獲得や一貫性のあるテキスト生成において顕著な能力を示しており、これは主に「semantic associations（意味的関連）」を学習する能力に支えられている（例："bird"と"flew"の間のリンク、または"car"と"truck"の交換可能性）。近年の研究では、インダクションヘッドや線形意味関係のような特定のメカニズムが特定されているが、現実世界の自然言語データ上で注意機構に基づくトランスフォーマーを訓練する際に、これらの関連性がどのように出現するかについての、原理的かつ理論的根拠に基づいた説明は依然として欠けている。

従来の理論的研究は、合成構造化言語、簡略化されたアーキテクチャ（位置符号化や残差接続の欠如）、または非標準的な訓練レジーム（例：逐次的コンポーネント別訓練）といった、非現実的な仮定に依存することが多い。これらの実践からの逸脱は、現代の LLM に対する洞察の一般化可能性を制限している。本論文は、このギャップに焦点を当て、以下の問いを提起する：自然言語データ上で注意機構に基づく言語モデルを訓練する際、意味的関連性はどのように出現するか？

方法論

著者は、訓練ダイナミクス、特に訓練初期段階における勾配の主導項近似を分析する理論的枠組みを提案する。

モデル設定: 本研究は、位置符号化（相対的）、因果的マスク、および残差ストリームを備えた標準的な $L$ 層の注意機構トランスフォーマーに焦点を当て、自然言語コーパス上で標準的なフルバッチ勾配降下法により訓練されるモデルを対象とする。
主導項近似: トランスフォーマーが（意味的関係を含む）中核的な振る舞いを訓練の初期段階で獲得することを踏まえ、著者は勾配更新を近似することで重み行列（出力 $W_O$ 、値 $V^{(l)}$ 、クエリ・キー $W^{(l)}$ 、および位置符号化 $P^{(l)}$ ）の閉形式式を導出する。十分なステップ数 $s$ （ $s \leq \eta^{-1} \min(\dots)$ ）において、重みは勾配の主導項に近接して維持されることを示す。
基底関数分解: 理論的分析により、学習された重みは、コーパス統計から導出される 3 つの基本的な基底関数の単純な合成として表現できることが明らかになった：
- バイグラムマッピング（ $\bar{B}$ ）: 次のトークンの依存関係（ $P(e_j|e_i)$ ）を捉える。
- 交換可能性マッピング（ $\Sigma \bar{B} = \bar{B}^\top \bar{B}$ ）: 先行トークン分布の類似性に基づいたトークン間の機能的類似性（例：同義語または共有される文法役割）を捉える。
- コンテキストマッピング（ $\bar{\Phi}$ ）: トークンが出現する平均的なコンテキストを表す、より長距離の接頭辞 - 接尾辞の共起を符号化する。

主要な貢献

本論文は、以下の 3 つの主要な貢献を行う：

明示的な重みの特徴付け: 次トークン予測損失の下で現実世界のテキストコーパス上で訓練された注意機構トランスフォーマーの重みに関する、最初の明示的かつ閉形式の特徴付けを提示する。重みは、上記の 3 つの基底関数の合成であることが示される。
意味的関連性の機械的解釈: 著者は、重みで学習された特徴を、バイグラム、交換可能性、およびコンテキストマッピングの特定の合成として解釈する。これらの合成が、モデルが意味的関連性をどのように捉えるかを示す（例："fish"が"pond"とどのようにコンテキストマッピングを介して関連するか、または"red"が"truck"とどのようにバイグラムマッピングを介して関連するか）。
実証的検証: 理論的予測は、TinyStories データセットで訓練された玩具的な 3 層注意モデルと、OpenWebText で訓練された実用的な大規模 LLM（Pythia-1.4B）の両方において検証された。

結果

理論的整合性: TinyStories で訓練された 3 層トランスフォーマーにおける実験により、理論的に予測された主導項と実際に学習された重みとの間のコサイン類似度が、100 エポックの訓練後であっても極めて高い（例：注意行列と値行列で 0.99 超）ことが示された。
意味構造: 定性的分析により、基底関数が意味のある言語構造を捉えていることが確認された。例えば、バイグラムマッピング $\bar{B}$ は形容詞とそれを修飾する名詞を相関させ、コンテキストマッピング $\bar{\Phi}$ は実体とその典型的な環境（例："fish"と"pond"）を相関させる。
実在 LLM への一般化: Pythia-1.4B に適用された場合、理論的主導項特徴は、特に訓練の初期段階において、学習された埋め込みや注意相関と強い一致を示す。分析により、訓練が進むにつれて重みはより豊かな知識表現へとシフトするが、理論によって予測された基礎的な関連特徴は相当期間維持されることが明らかになった。
コンポーネント間の協調: 完全なモデル計算（式 12）の分析は、残差ストリームが平均的な次トークン予測を提供する一方で、自己注意ブロック（ $\bar{Q}$ および $\bar{\Phi}^\top \Sigma \bar{B}$ によって駆動される）が、コーパス統計に基づいて次トークンに最も指標となるトークンを選択的に焦点を当てることで、この予測を洗練させることを示している。

意義

本論文は、深層学習の訓練ダイナミクスと言語理論を結びつけることで、大規模言語モデルにおける表現学習の機械的基盤を提供すると主張している。自然主義的なテキスト分布と標準的な訓練手順に基づいて分析を根ざすことで、この研究は形式的な理論的分析と実用的な LLM 利用との間のギャップを大幅に縮小している。

著者は、その発見が各トランスフォーマーコンポーネントの機能を文脈化する厳密な方法を提供すると論じている。モデルをブラックボックスとして見るのではなく、この理論は、モデルの振る舞いが、まずこれらの基本的な関連特徴（バイグラム、交換可能性、コンテキスト）を確立することによって発展し、それがより複雑な意味関係の学習のためのアンカーとして機能することを示唆している。この研究は、異なるコンポーネントの重み行列が、共有されたコーパス統計の単純な関数に分解可能であることを示唆することで、特定のメカニズムを超えてモデル内で概念がどのように出現するかについてのより広範な仮説を可能にし、解釈可能性研究への道を開く。