📊 statistics

Assign and Add: A Mechanistic Study of Compositional Arithmetic

本論文は、変数割り当てとモジュラー加算という制御された設定において、小型のトランスフォーマーがいかにして構成的汎化を達成するかを調査し、モデルが直接入力と間接入力の両方に対して内部メカニズムを再利用すること、および未知の組み合わせへと自然に汎化するために明確な学習フェーズを経て進展することを明らかにしている。

原著者： Brady Exoo, Alberto Bietti, John Sous

公開日 2026-06-01

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

CC BY 4.0

原著者： Brady Exoo, Alberto Bietti, John Sous

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

あなたは、非常に賢いけれど、非常に字義通りにしか理解できないロボットに数学を教えていると想像してください。あなたは次のようなパズルを解かせたいと考えています。「もし『A』が5で、『B』が10なら、『A + B』は何ですか？」

厄解なのは、ロボットは「A」と「B」という文字がまさにその順番で使われるのを一度も見たことがないということです。ロボットは、数字同士が足し算されているもの（例：5 + 10）や、他の文字が異なる方法で使われているものは見たことがあります。研究者が問いかけた大きな疑問は、ロボットはいかにして、これら2つの異なるスキル（文字が数字を表していることを理解することと、実際に計算を行うこと）を組み合わせ、未学習の新しいパズルを解く方法を学ぶのか？ということです。

論文「Assign and Add」が発見した内容を、分かりやすく説明します。

1. セットアップ：2ステップのダンス

研究者たちは、小さなシンプルな「脳」（ニューラルネットワーク）を構築し、特定のタスクを与えました。彼らはロボットに次のような文章を入力しました。

c=17, b=42, b+19=?

ロボットは、b が実際には 42 であり、つまり計算は 42 + 19 であることを理解し、答えを出す必要があります。

ロボットがどのように学習したかを確認するために、学習を異なる「フレーバー（種類）」のパズルに分割しました。

簡単な方法： 数字のみ（例：17 + 19 = ?）。
中間の方法： 1つの文字と1つの数字（例：b + 19 = ?）。
難しい方法： 2つの文字（例：b + c = ?）。

2. 学習の3つのフェーズ

ロボットはすべてを一度に学んだわけではありません。学習は、段階的に主題をマスターしていく学生のように、3つの明確なフェーズを経て進みました。

フェーズ1：数学の達人。 まず、ロボットは実際の足し算を学びました。それは数字を足し合わせることに非常に長けていきました。この段階では、17 + 19 は得意でしたが、文字については全く混乱していました。これは、変数とは何かを知らない計算機のようです。
フェーズ2：翻訳者。 次に、ロボットは「代入」のスキルを学びました。c=17 を見て、c が 17 を意味することを覚え、文字を数字に置き換える方法を理解しました。
フェーズ3：指揮者。 最後に、ロボットはこれら2つのスキルを組み合わせる方法を学びました。「あ、『翻訳者』のスキルを使って数字を見つけ、『数学の達人』のスキルを使って足し算ができるんだ」と気づいたのです。

驚きの事実： ロボットは、翻訳を完全に習得する前に、数学を学んでしまいました。翻訳に苦戦している間でも、足し算は完璧にできていたのです！

3. ロボットの「脳」の仕組み（メカニズム）

研究者たちは単に最終スコアを見たのではなく、ロボットがどのように考えているのかを知るために、その「脳」の内部を観察しました。そこで、2つの特定のツールが連携して動いていることを見つけました。

ツールA：「前方のトークン」ヘッド（翻訳者）。
現在の項の直前にあるアイテムを掴み取るロボットアームを想像してください。ロボットの第1層において、それは文章の中で以前に割り当てられた文字に対応する数字を見つけるために、このアームを使用することを学びました。もし b を見たら、b が以前にどのように割り当てられたかを遡って探し出します。
ツールB：「フーリエ」数学モジュール（計算機）。
第2層において、ロボットは波やパターンを用いた特別な数学的トリック（フーリエ変換に関連するもの）を使用して、足し算を行います。これは、ロボットがモジュラー算術（時計のように循環する計算）を行うための既知の方法です。

魔法の瞬間： ロボットはこれら2つのツールを接続する方法を学びました。「翻訳者」のアームが正しい数字を掴み、それを「計算機」モジュールへと渡します。一度接続されると、ロボットは b + c を、たとえ b と c が一緒に足し合わされる場面を一度も見たことがなくても、解くことができるようになります。それは、文字を話す「翻訳者」と数字を話す「計算機」がいて、ついにその間に橋を架けたようなものです。

4. なぜこれが重要なのか

この論文は、**汎化（未知の問題を解くこと）**は魔法ではないことを示しています。それは、モデルが小さな原子的なスキル（数字を足すことや変数を見つけること）を学習し、それらをどのように組み合わせるかを理解したときに、自然に起こる現象です。

「グロッキング（Grokking）」現象： 研究者たちは「グロッキング」と呼ばれる現象に気づきました。時として、ロボットは訓練データを完璧に暗記しますが、新しいテストでは失敗します。しかし、ある時突然「カチッ」と音がしたように理解が進み、単なる答えの暗記ではなく、「ルール」を理解し始めるのです。
結果： ロボットが内部メカニズムを別々に学習し、その後でそれらを組み合わせたため、訓練中に許可されていなかった位置に文字が来るようなパズルであっても、対処することができました。

要約の比喩

ロボットをシェフだと考えてください。

まず、野菜の切り方（数学）を学びます。
次に、レシピカードを読み、「青いボウルの中にある玉ねぎを使ってください」という指示を理解します（変数の代入）。
最後に、これらを組み合わせる方法を学びます。カードを読み、玉ねぎを見つけ、そしてそれを切るのです。

この論文は、もしシェフにこれら2つのスキルを別々に教えれば、彼らが一度も要求されたことのない料理を作るために、それらを組み合わせる方法を最終的に自力で見つけ出すことを証明しています。スキルの「構成（コンポジション）」こそが、ロボットを賢くさせているのです。

技術要約：Assign and Add：構成的算術に関するメカニズム的研究

問題提起

大規模言語モデル（LLM）は、訓練中に未知であった複雑なタスクを実行するために、スキルを合成する能力を示しますが、この「構成的汎化（compositional generalization）」を可能にする具体的なメカニズムは依然として解明されていません。汎化は原子的なメカニズムの合成に起因するとされることが多いものの、異なるメカニズム、具体的には変数バインディング（誘導ヘッドによる連合想起）とモジュロ演算（フーリエ特徴量による構造化された表現）の相互作用については、制御された設定での系統的な研究が行われていません。

本論文は、以下の問いに取り組んでいます：トランスフォーマーは、変数の割り当てとモジュロ加算のメカニズムをどのように学習して構成し、訓練を通じてこれらの能力はどのように立ち現れるのか？

手法

著者らは、変数割り当てとモジュロ加算を組み合わせた合成タスクを学習する、メカニズム的に解釈可能な小規模なトランスフォーマーを用いた、制御された実験設定を採用しています。

モデルアーキテクチャ: $d_{model} = d_{head} = 128$ 、 $d_{mlp} = 512$ の2層・シングルヘッド・トランスフォーマー。回路発見を簡素化するため、バイアス、正規化層、および第1層のMLPを意図的に排除しています。GPT-2スタイルの絶対位置エンコーディングと因果的マスキングを使用しています。
タスク定義: 変数（定数またはコンテキスト内で動的に割り当てられた変数）をオペランドとするモジュロ加算（ $N=59$ ）を実行するようにモデルを訓練します（例：c=1, b=17, a=42, b+19=?）。
データ分布: 絶対的な位置への依存を防ぐため、シーケンスはランダムなパディングを伴ってトークン化されます。訓練セットは、汎化を評価するために以下のように分割されています：
- 0-variable (0-var): 純粋な定数加算。
- 1-variable (1-var) & 2-variable (2-var): 変数が割り当てられた加算。
- ホールドアウト集合: 学習から約30%の加算ペアが削除されています（例：もし 1+3 が未学習であれば、それを含むすべてのシーケンスを除外する）。さらに、特定の変数が訓練セット内の特定のポジションに制限されており、位置の汎化をテストしています。
分析手法: 本研究では、アテンションパターン、残差ストリームのベクトル、および重み行列を分析するために、TransformerLens ライブラリを用いたメカニスティック・インタープリタビリティ（機械論的解釈可能性）のツールを使用しています。また、線形プローブと理論的な勾配分析を用いて、訓練のダイナミクスを追跡しています。

主な貢献

構成的算術のための扱いやすいモデル: 著者らは、変数割り当てがモジュロ加算とどのように構成されるかをメカニスティックに解釈することを可能にする、簡略化されたモデルを提示しています。
分離された訓練ダイナミクス: 分析により、モジュロ加算は変数割り当てよりも前、かつ独立して学習されることが明らかになりました。モデルは、変数バインディングが完全に確立される前であっても、0-varタスクにおいてほぼ完璧な精度を達成します。
汎化のメカニズム的説明: 本論文は、汎化が以下の2つの特定の回路の合成から生じることを特定しています：
- 割り当てられた変数を取得し、その値をルーティングする、第1層の修正された誘導ヘッド。
- 取得された値に対して算術演算を実行する、第2層のモジュロ加算回路（特にOutput-Value行列と第2層のMLP）。
理論的枠組み: 著者らは、勾配ダイナミクスがいかにして変数バインディング回路とモジュロ加算回路の合成を促進するかを示す理論的分析を提供しており、第2層のQuery-Key行列に対する負の勾配が、自然に連合メモリブロックの形成を導くことを示しています。

結果

訓練ダイナミクスとフェーズ

訓練プロセスは、3つの明確なフェーズを示します：

モジュロ加算の学習: モデルはまず、変数を含まないシーケンス（0-var）に対する加算を学習し、変数バインディングが完全に確立される前に高い精度を達成します。
変数割り当ての学習: 次に、モデルは変数割り当てに必要な構造を学習します。精度における急激な「スパイク」（ステップ5,300–5,400付近）は、変数割り当てルーティング回路の出現と一致しており、そこでは第1層のアテンションが定数を割り当てられた変数へと正しくルーティングしています。
精緻化フェーズ: 第2のフェーズ（ステップ14,500–16,000）では、「困難な」シーケンス（例：特定のホールドアウトされた位置にある変数）に対するルーティングの精緻化が行われ、特定の変数位置へのアテンションの失敗を修正します。

汎化能力

位置の汎化: モデルは、新しい変数の順序（例：14+b を見た後に b+14 を予測する）に対して汎化します。
分布外（OOD）汎化: モデルは、訓練中にいかなる形式でも一度も見られなかった保持された加算ペア（例：1+3）を正常に予測し、精選されたテストセットにおいて99.4%の精度を達成しました。
データ要件: 未知の定数ペアへの汎化には、全分布の少なくとも約25%（ $f \approx 0.25$ ）の訓練が必要です。2変数シーケンスへの汎化には、訓練データにおける2-varシーケンスの相対頻度として約20%（ $r \approx 0.2$ ）が必要です。

メカニズム的知見

共有表現: 変数割り当てられた加算と純粋な定数加算の残差ストリームベクトルは、第2層のMLPの前でほぼ同一（コサイン類似度 $\approx 1$ ）であり、共通の計算メカニズムを示しています。
フーリエ特徴量: 第2層のMLPニューロンの事前活性化は、理論的なフーリエ特徴量の形式（ $\cos(\frac{2\pi k}{N}n + \alpha_k) + \cos(\frac{2\pi k}{N}m + \beta_k)$ ）と一致しており、モジュロ加算のための確立された周期メカニズムの使用を裏付けています。
誘導ヘッドの修正: 第1層のアテンションは、定数に対する「前のトークン・ヘッド」として機能し、変数の直前のトークンにアテンションを向けます。その後、第2層のQuery-Key行列は変数埋め込みに対して強い恒等パターンを形成し、これにより = トークンが関連する定数に直接アテンションを向けることが可能になります。

重要性

本論文は、トランスフォーマーにおける構成的汎化が、内部メカニズムの構成性の自然な結果となり得ることを主張しています。異なるメカニズム（誘導ヘッドとフーリエベースの算術）が独立して学習され、その後アテンションルーティングを介して合成されることを示すことで、本研究は、トランスフォーマーが複雑な推論タスクにどのように汎化するかを理解するためのメカニズム的基盤を提供しています。

著者らは、分析を容易にするために制御された小規模な設定に限定していることを述べ、研究の範囲について謙虚な姿勢を維持しています。彼らは、これらの知見がすべてのLLMの挙動を直接説明すると主張しているのではなく、トランスフォーマーにおいて汎化能力がいかにして立ち現れるかを理解するためのフレームワークを提供することを示唆しています。