Induction Headsとは — 文脈内学習を担う誘導ヘッドの仕組み
誘導ヘッド(induction head)が大規模言語モデルの文脈内学習の主要な仕組みだとする仮説を、Anthropicが6つの根拠で裏づけた研究。訓練初期の「相転移」現象を軸に読み解きます。
induction headsとは何か、なぜ調べる価値があるのか
「In-context Learning and Induction Heads」は2022年3月にAnthropicが公開した論文で、前作のTransformer Circuits論文で発見した誘導ヘッド(induction head)という回路が、大規模言語モデル全般の文脈内学習(in-context learning)を担う主な仕組みではないかという仮説を検証したものです。
誘導ヘッドは、attentionヘッドが2層に分かれたTransformerで発生する回路です。2つのヘッドが次のように連携します。
誘導ヘッドが次のトークンを当てるまで
- 1
1つ前のトークンヘッドが情報を運ぶ
1層目のヘッドが、各トークンの位置に「1つ前のトークンは何だったか」という情報をコピーします。
- 2
誘導ヘッド本体が過去の出現位置を探す
2層目のヘッドが、現在のトークンAが以前どこに出現したかを探し、そのすぐ次にあったトークンBに注意を向けます。
- 3
Bをコピーして予測に使う
Bの情報が現在位置に運ばれ、
…[A][B]…[A]という系列が[B]で完結する確率が上がります。
名前の由来は帰納推論(inductive reasoning)です。文脈中で一度AがBに続いた実績があれば、再びAが出たときもBが続きやすいという推論を、この回路がそのまま体現しています。
前作の論文では2層のattentionのみモデルでこの仕組みを数式で完全に説明できましたが、実際にターゲットとする最先端の言語モデルは数百層・数十億〜数兆パラメータを持ち、MLP層も含みます。そこまで複雑なモデルを数式だけで解析するのは困難です。本論文はアプローチを変え、学習過程を観察・摂動しながら間接的な証拠を積み重ねる手法を取りました。神経科学者が神経の発達過程や損傷例の観察、動物実験での摂動を通じて脳機能を理解するのに近いやり方です。
文脈内学習をどう定義し、どう測るか
文脈内学習には大きく2つの捉え方があります。1つはFew-shot学習に注目する立場で、タスクの実例を複数プロンプトに含めて、モデルがそのタスクをこなせるかを見ます。もう1つは、文脈中のトークン位置ごとの損失(loss)を観察し、文脈が長くなるほど予測がどれだけ改善するかを見る立場です。
本論文は後者の定義を採用しています。特定のタスク設計に依存せず、より一般的にモデルの「文脈から学習する能力」全体を測れるためです。具体的には次の指標を使います。
文脈内学習スコア = 文脈中50トークン目の損失 − 文脈中500トークン目の損失(データセット全体で平均)
50トークン目は言語や文書の種類など基本的な性質が確立された直後、500トークン目は長さ512の文脈のほぼ終端にあたります。値が大きいほど、文脈が長くなるにつれて予測が改善している、つまり文脈内学習が強く働いていることを示します。
論文は引き算の向きを2通りで書いています。定義の箇所と「説明のつかない現象」の節では「500トークン目の損失から50トークン目の平均損失を引く」とあり、Argument 1の本文では「50トークン目の損失から500トークン目の損失を引く」形で数値を述べています。損失は文脈が進むほど下がるので、前者の向きだと値は負になります。以下の数値は後者の向きです。
6つの根拠 — 誘導ヘッド仮説をどう裏づけたか
論文は34個のTransformerを訓練の全過程にわたって分析し、5万件を超えるattentionヘッドのアブレーション実験のデータをもとに、6つの補完的な論拠を積み重ねています。
調べたモデルは2系統です。小規模モデルは1〜6層のTransformerで、MLPのあるものと注意機構のみのものがあります。どの層にも12個のattentionヘッドを持ち、文脈ウィンドウは8192トークンで、1万ステップ(約100億トークン)の訓練中に50ステップごとに計200回スナップショットを取っています。もう一方の実規模モデルは4層(非埋め込みパラメータ約1,300万)から40層(約130億)までです。
論拠③のアブレーションは、ヘッドを丸ごと消す「完全アブレーション」ではなく、「パターン保存アブレーション」が中心です。後者は、ヘッドの出力ベクトルへの寄与だけを後続のヘッドから取り除き、後続ヘッドのattentionパターンへの寄与は残します。attentionパターンの計算と出力ベクトルの計算を別々の回路として切り分けられるという前作の知見が、この操作を可能にしています。
| 論拠 | 内容 |
|---|---|
| ①マクロな同時発生 | 内容訓練初期に「相転移」が起き、誘導ヘッドの形成と文脈内学習の急激な向上が同時に起きる |
| ②マクロな共変化 | 内容アーキテクチャを変えて相転移の発生時期をずらすと、文脈内学習の向上時期も同じだけずれる |
| ③直接アブレーション | 内容小規模モデルで誘導ヘッドをテスト時に直接無効化すると、文脈内学習が大きく低下する |
| ④汎化例の提示 | 内容単純な系列コピーとして定義した誘導ヘッドが、翻訳やパターンマッチングなど抽象度の高い挙動にも関与している |
| ⑤機構としての妥当性 | 内容小規模モデルで誘導ヘッドの動作原理を数式で説明でき、より一般的な文脈内学習に転用されうる自然な経路がある |
| ⑥連続性 | 内容小規模モデルから大規模モデルまで、誘導ヘッドと文脈内学習に関するデータが滑らかに連続しており、同じ機構だと考えるのが最も単純な説明になる |
これらを総合すると、誘導ヘッドが最先端のTransformerモデルにおける文脈内学習の大部分を担っている可能性がある、という状況証拠が積み上がります。ただし論文は、決定的な証明ではなく初期段階の証拠だと明記しています。間接的な証拠を早い段階から報告する慣習を、解釈可能性研究に根づかせたいという狙いも書かれています。
論拠の強さはモデルが大きくなるほど弱まる
6つの論拠は、どのモデルにも同じ強さで効くわけではありません。論文は「誘導ヘッドが文脈内学習の一部を担う」と「大半を担う」の2つの主張について、モデルの種類ごとに最も強い論拠の強さを表にしています。
最も強い論拠でも、規模が上がると因果から相関へ
注意機構のみの小規模モデル
「一部を担う」も「大半を担う」も、強い因果的な証拠です。アブレーションに加え、機構を数式で逆解析できているためです。
MLPありの小規模モデル
「一部を担う」は強い因果的証拠ですが、「大半を担う」は中程度の因果的証拠にとどまります。
大規模モデル
「一部を担う」は相関的かつ機構的な中程度の証拠、「大半を担う」は相関的な中程度の証拠です。交絡の可能性が残ります。
大規模モデルでは、論拠①の同時発生のような相関が主な証拠になります。論文は交絡の可能性を認め、代替仮説を各論拠の節で検討しています。「誘導ヘッドは大規模モデルの文脈内学習の大半を担う」という文を読むときは、この強さの差も一緒に押さえておく必要があります。
相転移という現象と安全性への含意
論文の証拠の中心にあるのが「相転移(phase change)」です。2層以上のモデルでは、調べたどの規模でも訓練の初期に起きた現象で、訓練損失のグラフに小さな山(bump)として現れます。この相転移の最中に、文脈内学習能力(文脈の早い位置と遅い位置のトークンの損失差で測定)の大部分が獲得されると同時に、抽象的でファジーなパターン補完を実行できる誘導ヘッドがモデル内に形成されます。
論文はこの2つの現象が本当に因果関係にあるかを確かめるため、Transformerのアーキテクチャを意図的に変更して相転移の発生タイミングをずらす実験も行っています(論拠②)。改造点は「smeared key」と呼ばれ、各ヘッドのkeyベクトルを、現在のトークンのkeyと1つ前のトークンのkeyの内分に置き換えます。誘導ヘッドは「1つ前のトークンの情報を持つkey」を必要とするので、1層のモデルでも誘導ヘッドを表現しやすくなります。誘導ヘッドが文脈内学習の主因なら、1層モデルでも相転移が起き、2層以上では早まるはずです。結果はその予測どおりで、普通の1層モデルでは起きなかった文脈内学習が1層でも現れ、2層以上では相転移が早まりました。
相転移の前後で変わった数字
転移前のスコア
0.15未満
訓練の序盤
転移後のスコア
約0.4
2層〜130億パラメータでほぼ同じ
相転移の位置は、同じ論文の中で記述が揺れています。論拠①の節は約25億〜50億トークン、モデル詳細の節は小規模モデルで約10億〜30億トークンです。いずれも小規模モデルの訓練全体(約100億トークン)の序盤にあたります。訓練中のハイパーパラメータ変更は2つあります。重み減衰の引き下げは約50億トークン付近で、相転移の範囲より後なので無関係だと論文は述べています。学習率のウォームアップは最初の15億トークンにわたり、相転移の範囲と重なります。
転移後のスコアは訓練を続けても動きません。大きなモデルが全体として予測に優れていることは変わらず、その優位は文脈のごく初めの部分で付いてしまう、と論文は説明しています。50・500という測定位置を変えても同様でしたが、これは大規模モデルのグラフでの検証です。スコアが規模によらず一定になる点は、論文自身も説明がつかないとして「Unexplained Curiosities」の節で扱っています。
相転移という現象自体、安全性の観点からも重要だと論文は指摘します。多桁の足し算のような能力が、訓練が進むにつれて突然出現したり変化したりすることは知られており、望ましくない挙動や危険な挙動が急に現れうるという懸念につながります。実際、報酬ハッキング(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で紹介しています。
まとめ
誘導ヘッドが文脈内学習を担うという仮説は、モデルが小さいほど因果的な証拠で、大きいほど相関的な証拠で支えられています。論文自身がこの強さの差を認めており、「大半を担う」という主張を読むときは、どの規模の証拠なのかを見分ける必要があります。この仮説を検証した手法は、前作の数式による解析から、のちのClaude 3 Sonnetでの特徴抽出へと続く系譜に位置します。