Induction Headsとは — 文脈内学習を担う誘導ヘッドの仕組み

induction headsとは何か、なぜ調べる価値があるのか

「In-context Learning and Induction Heads」は2022年3月にAnthropicが公開した論文で、前作のTransformer Circuits論文で発見した誘導ヘッド(induction head)という回路が、大規模言語モデル全般の文脈内学習(in-context learning)を担う主な仕組みではないかという仮説を検証したものです。

誘導ヘッドは、attentionヘッドが2層に分かれたTransformerで発生する回路です。1つ目のヘッド(「1つ前のトークンヘッド」)が各トークンに1つ前のトークンの情報をコピーし、2つ目のヘッド(誘導ヘッド本体)がその情報を使って「現在のトークンAが以前どこに出現したか」を探し、そのすぐ次にあったトークンBに注意を向けてコピーします。結果として、…[A][B]…[A]という系列が[B]で完結する確率が上がります。名前の由来は帰納推論(inductive reasoning)で、文脈中で一度AがBに続いた実績があれば、再びAが出たときもBが続きやすいという推論を、この回路がそのまま体現しているためです。

前作の論文では2層のattentionのみモデルでこの仕組みを数式で完全に説明できましたが、実際にターゲットとする最先端の言語モデルは数百層・数十億〜数兆パラメータを持ち、MLP層も含みます。そこまで複雑なモデルを数式だけで解析するのは困難です。本論文はアプローチを変え、学習過程を観察・摂動しながら間接的な証拠を積み重ねる手法を取りました。神経科学者が神経の発達過程や損傷例の観察、動物実験での摂動を通じて脳機能を理解するのに近いやり方です。

文脈内学習をどう定義し、どう測るか

文脈内学習には大きく2つの捉え方があります。1つはFew-shot学習に注目する立場で、タスクの実例を複数プロンプトに含めて、モデルがそのタスクをこなせるかを見ます。もう1つは、文脈中のトークン位置ごとの損失(loss)を観察し、文脈が長くなるほど予測がどれだけ改善するかを見る立場です。

本論文は後者の定義を採用しています。特定のタスク設計に依存せず、より一般的にモデルの「文脈から学習する能力」全体を測れるためです。具体的には次の指標を使います。

文脈内学習スコア = 文脈中500トークン目の損失 − 文脈中50トークン目の平均損失(データセット全体で平均)

50トークン目は言語や文書の種類など基本的な性質が確立された直後、500トークン目は512トークンの文脈のほぼ終端にあたります。この指標が大きいほど、文脈が長くなるにつれて予測が改善している、つまり文脈内学習が強く働いていることを示します。

6つの根拠 — 誘導ヘッド仮説をどう裏づけたか

論文は34個のTransformer(訓練過程を通じて)、5万件を超えるattentionヘッドのアブレーション実験のデータをもとに、6つの補完的な論拠を積み重ねています。

これらを総合すると、誘導ヘッドが最先端のTransformerモデルにおける文脈内学習の大部分を担っている可能性がある、という状況証拠が積み上がります。著者ら自身、これは決定的な証明ではなく今後の研究で強化されるべき初期段階の証拠だと位置づけています。それでも、こうした間接的な証拠を早い段階から報告する慣習を解釈可能性研究の分野に根づかせたい、という狙いも論文には書かれています。

相転移という現象と安全性への含意

論文の証拠の中心にあるのが「相転移(phase change)」です。2層以上を持つあらゆる規模の言語モデルで、訓練の初期に必ず起きる現象で、訓練損失のグラフに小さな山(bump)として現れます。この相転移の最中に、文脈内学習能力(文脈の早い位置と遅い位置のトークンの損失差で測定)の大部分が獲得されると同時に、抽象的でファジーなパターン補完を実行できる誘導ヘッドがモデル内に形成されます。

論文はこの2つの現象が本当に因果関係にあるかを確かめるため、Transformerのアーキテクチャを意図的に変更して相転移の発生タイミングをずらす実験も行っています。すると誘導ヘッドの形成時期と文脈内学習の向上時期が、そろって同じだけ動いたといいます(論拠②)。

相転移という現象自体、安全性の観点からも重要だと論文は指摘します。多桁の足し算のような能力が、訓練が進むにつれて突然出現したり変化したりすることは知られており、望ましくない挙動や危険な挙動が急に現れうるという懸念につながります。実際、報酬ハッキング(reward hacking)のような安全性上の問題が、こうした相転移として現れる例も報告されています。相転移を間近で観察し内部の仕組みを理解することは、解釈可能性という「ミクロ」な領域と、スケーリング則や学習ダイナミクスという「マクロ」な領域を橋渡しする糸口になるとも位置づけられています。

誘導ヘッドの厳密な定義 — prefix matchingとcopying

前作の論文では誘導ヘッドの定義がやや曖昧でした。「一緒に起きやすい挙動と機構の集まり」を観察し、それらをまとめて誘導ヘッドと呼んでいたにすぎません。本論文はMLP層を含む大規模モデルでは重みの数式的な解析や固有値解析が格段に難しくなることを踏まえ、誘導ヘッドを次の2つの性質で厳密に定義し直しました。

• Prefix matching(接頭辞の一致): そのヘッドが、現在のトークンや直近のトークンに続いていた過去のトークンへ注意を向け返す
• Copying(コピー): そのヘッドの出力が、注意を向けた先のトークンのロジットを増加させる

完全にランダムなトークンを繰り返した系列の上でこの2性質を測定することで、単に「次に来そうなトークンへ大まかに注意を向けるヒューリスティックなコピーヘッド」ではなく、本当に帰納的な仕組みで動いているヘッドだけを識別できます。もっとも単純な誘導ヘッドは直前の1トークンだけを照合しますが、複数の先行トークンにまたがってファジーに照合するタイプもよく観察されたといいます。

単純なコピーを超えた挙動の広がり

誘導ヘッドは本来「文字どおりの系列コピー」として狭く定義されていますが、論文はより抽象的な挙動でも同じヘッドが働いていることを観察しています。厳密な定義どおりの誘導ヘッドは「The cat sat on the mat. The cat …」のような文に対して「sat on the mat」という続きを後押しします。これは任意の系列をそのまま繰り返す、単純な形のfew-shot学習だと言えます。

さらに踏み込んだ観察が「アナロジー的な系列コピー」です。[A][B]…[A]→[B]のAとBが完全に同一のトークンである必要はなく、埋め込み空間上で似ている[A*][B*]…[A]→[B]という形でも、同じヘッドが機能することが確認されました。たとえばAとA*(およびBとB*)が異なる言語の同じ単語であれば、誘導ヘッドは文を単語ごとに「Aに似たもの」を探して見つけたAの次にあるBを頼りに、「B*に似たもの」つまりBで文を完成させることで、翻訳のような挙動を実現できます。論文はこれを含め、翻訳やパターンマッチングのような、より抽象的な形の文脈内学習にも同じヘッドが関与している具体例を提示しました(論拠④)。ここから、誘導ヘッドという1つの仕組みが、表面的には別々に見えるさまざまな文脈内学習のタスクを、実は共通の原理で説明している可能性が見えてきます。

この研究がその後の解釈可能性研究にどうつながったか

誘導ヘッドと文脈内学習を結びつけるこの一連の研究は、Anthropicの解釈可能性研究における重要な系譜の起点になっています。「小さなモデルで数式により厳密に理解する」(前作)から「大きなモデルで間接的な証拠を積み重ねる」(本作)へと手法を発展させたのち、Anthropicはさらにスケールアップした手法で、実運用中のClaude 3 Sonnetから数百万の解釈可能な特徴を取り出す研究に進みます。この研究とその象徴的な実験であるGolden Gate ClaudeについてはGolden Gate Claude論文の中身で扱っています。モデルの活性化を自然言語で読み解こうとする、より最近の取り組みはNatural Language Autoencodersで紹介しています。

まとめ

In-context Learning and Induction Heads論文は、誘導ヘッドという単一の回路が、規模を問わず言語モデルの文脈内学習の主要な駆動源になっているという仮説を、6つの補完的な証拠で裏づけようとした研究です。訓練初期に起きる「相転移」という現象を軸に、マクロな観察・アーキテクチャの操作・直接的なアブレーションを組み合わせた点が特徴で、その後の解釈可能性研究の方法論にも影響を与えました。