誘導ヘッド(Induction Heads)とin-context learning — LLMが文脈から学ぶ仕組みに迫る機構的解釈入門

大規模言語モデルの不思議な能力のひとつに in-context learning(ICL、文脈内学習) があります。プロンプトに「英語→フランス語」の例を数個並べるだけで、パラメータを1つも更新していないのに、続きの翻訳ができてしまう。学習は訓練時に終わっているはずなのに、モデルは文脈の中で新しいパターンを拾って使いこなします。重みが固定されたただの行列計算の塊が、どうやって「その場で学ぶ」のでしょうか。

この謎に対する最有力の手がかりが、Anthropic の Olsson らが2022年に提示した 誘導ヘッド(induction head) です。誘導ヘッドとは、ざっくり言えば

「前に [A][B] という並びがあった。いままた [A] が来た。なら次は [B] だろう」

を実行する attention ヘッドのことです。ハリー・ポッターの「Mr. Dur」まで読んだモデルが「sley」を即答できるのは、前に出てきた「Dursley」をこの機構がコピーしてくるからです。単純すぎる機構に見えますが、Olsson らは「これこそが ICL の主要な源泉ではないか」という仮説を、複数の証拠を積み上げて論じました。

誘導ヘッドを理解すると、次のような場面で役に立ちます。

  • LLMの挙動の解像度が上がる: few-shot プロンプトが効く理由、繰り返しの多い文書でモデルが強い理由を、内部機構のレベルで説明できる
  • 機構的解釈(mechanistic interpretability)への入口: 「モデルの中の回路を特定する」という研究分野の、最も成功した具体例を知ることができる

本記事の内容

  • 誘導ヘッドの定義 — prefix matching と copying という2つの性質
  • なぜ1層では作れないのか — previous-token head と K-composition による最小回路の分解
  • 学習ゼロの手作り attention 回路で誘導を決定的に再現(的中率100%)
  • 2層 attention-only Transformer の訓練実験 — 繰り返し系列の「2周目」だけ損失が急落し、1層モデルは失敗する
  • 訓練中の相転移 — 誘導ヘッドの形成と ICL 能力が同時に立ち上がる

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が深まります。

誘導ヘッドの概念図: 前回のAの直後に注意を向け、Bをコピーする

上の図が誘導ヘッドの動作です。いま「山」というトークンを処理しているとき、①文脈をさかのぼって前回「山」が出た場所の直後(=「川」の位置)に注意を向け(prefix matching)、②注意先のトークン「川」の出力確率を引き上げます(copying)。この2つの性質を併せ持つヘッドが誘導ヘッドです。

原論文の定義図も見ておきましょう。

誘導ヘッドの定義: prefix matchingとcopying(原論文より)

出典: Olsson et al., “In-context Learning and Induction Heads”, Transformer Circuits Thread / arXiv:2209.11895, 2022

ランダムなトークン列が繰り返される状況で、2周目の「node」を処理中のヘッドが、1周目の「node」の直後の「struction」へ attention を張り(赤矢印)、そのトークンの logit を引き上げる——図の注釈がそのまま prefix matching と copying の定義になっています。重要なのは、この機構が「node struction という並びを訓練データで暗記した」のではなく、その場の文脈だけから任意の [A][B] ペアに働くことです。だからこそ「文脈内学習」の部品になれるのです。

では、この機構は attention でどう実装されるのでしょうか。実は、1層では原理的に作れません

なぜ2層必要なのか — 最小回路の分解

誘導ヘッドの仕事を分解すると、「自分のトークン $t_i$ と同じトークンの直後の位置を探す」検索です。attention の検索は Query と Key の内積で行われるので、位置 $j$ の Key に「位置 $j-1$ のトークンが何だったか」という情報が入っていなければ、この検索は成立しません。

ところが1層目の attention に入力される表現は「自分のトークン+自分の位置」だけです。Key に「直前のトークン」を含める手段がない——これが1層モデルの原理的な限界です(1層でできるのは、現在のトークンから直接「よく続くトークン」を引く skip-trigram 的な統計まで、というのが Elhage らの “A Mathematical Framework for Transformer Circuits”(2021)の分析です)。

2層あれば話が変わります。

誘導ヘッドの最小回路: previous-token headとK-composition

  • 第1層の previous-token head が、各位置 $j$ の残差ストリームに「直前のトークン $t_{j-1}$ の情報」を書き込む
  • 第2層の induction head が、Query に自分のトークン $t_i$、Key に第1層が書き込んだ「直前トークン情報」を使う。すると内積が大きくなるのは「直前が $t_i$ だった位置」=前回の $t_i$ の直後だけになる

第2層の Key が第1層の出力を使ってスコアを作る——この合成を K-composition と呼びます。誘導ヘッドは単独のヘッドではなく、層をまたいだ回路(circuit)として実現されるのです。

この分解が正しいなら、学習に頼らなくても、重みを手で設計するだけで誘導が動くはずです。やってみましょう。

実験1: 学習ゼロの手作り回路で誘導を動かす

埋め込みを「トークンのone-hot|位置のone-hot|直前トークン用の空きスロット」と区画し、上の分解をそのまま numpy で書きます。

import numpy as np

def softmax_rows(S):
    S = S - S.max(axis=-1, keepdims=True)
    e = np.exp(S)
    return e / e.sum(axis=-1, keepdims=True)

def handcrafted_induction(tokens, V, beta=30.0):
    """学習なしの2層attention誘導回路。
    埋め込み: [トークンone-hot | 位置one-hot | 直前トークンone-hot(第1層が書く)]"""
    L = len(tokens)
    X = np.zeros((L, V + L + V))
    X[np.arange(L), tokens] = 1.0                # トークン one-hot
    X[np.arange(L), V + np.arange(L)] = 1.0      # 位置 one-hot
    mask = np.tril(np.ones((L, L)))              # 因果マスク

    # 第1層: previous-token head (位置だけで i -> i-1 に注意)
    Q1 = X[:, V:V + L]
    K1 = np.zeros((L, L)); K1[1:, :L - 1] = np.eye(L - 1)
    K1 = X[:, V:V + L] @ K1.T                    # 位置jのkey = onehot(位置j+1)
    A1 = softmax_rows(np.where(mask > 0, beta * (Q1 @ K1.T), -1e9))
    prev_tok = A1 @ X[:, :V]                     # 直前トークンを取ってくる
    prev_tok[0] = 0.0                            # 位置0に「直前」は無い
    X[:, V + L:] = prev_tok                      # 残差ストリームに書き込む

    # 第2層: induction head (Query=自トークン, Key=直前トークン情報)
    S2 = beta * (X[:, :V] @ X[:, V + L:].T)      # K-composition
    S2 = np.where(mask > 0, S2, -1e9)
    np.fill_diagonal(S2, -1e9)
    A2 = softmax_rows(S2)
    out = A2 @ X[:, :V]                          # 注意先のトークンをコピー
    return A1, A2, out.argmax(axis=1)            # 次トークン予測

rng = np.random.default_rng(0)
V, T = 12, 10
pat = rng.permutation(V)[:T]
tokens = np.concatenate([pat, pat])              # 同じランダム列を2回
A1, A2, pred = handcrafted_induction(tokens, V)
L = len(tokens)
print(f"2周目の次トークン的中率 = {(pred[T:L-1] == tokens[T+1:L]).mean():.3f}")
print(f"1周目の次トークン的中率 = {(pred[1:T-1] == tokens[2:T]).mean():.3f}")
# => 2周目の次トークン的中率 = 1.000
# => 1周目の次トークン的中率 = 0.000

手作り誘導回路のattentionパターン: 第1層はサブ対角、第2層は誘導の縞

2周目の次トークン的中率は100%です。左のヒートマップが第1層で、教科書どおりの「1つ下の対角」(各位置が直前を見る)。右が第2層で、2周目(シアンの線から下)の各行が「前回の自分の直後」をピンポイントで指す縞が現れています。1周目の的中率が0%なのも重要です——ランダム列の初出部分は原理的に予測不能で、誘導ヘッドは「文脈に前例があるときだけ」働く機構だと分かります。

回路の設計図が正しいことは確認できました。次の問いは「この回路は勾配降下で自然に生まれるのか」です。

実験2: 2層attention-onlyモデルを訓練する

ランダムなトークン列(語彙30)を2回繰り返した系列で、attention-only(FFNなし)の自己回帰 Transformer を訓練します。繰り返しの周期 $T$ はバッチごとに12〜20でランダムに変えます。こうすると「常に $T$ 個前を見る」という固定オフセットの位置ショートカットでは解けなくなり、本物の誘導(内容による検索)だけが解になります。

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

class AttnOnlyLM(nn.Module):
    """attention-onlyの自己回帰LM (FFNなし)"""
    def __init__(self, vocab=30, d=64, nheads=2, nlayers=2, maxlen=64):
        super().__init__()
        self.tok = nn.Embedding(vocab, d)
        self.pos = nn.Embedding(maxlen, d)
        self.layers = nn.ModuleList(
            [nn.MultiheadAttention(d, nheads, batch_first=True) for _ in range(nlayers)])
        self.lns = nn.ModuleList([nn.LayerNorm(d) for _ in range(nlayers)])
        self.head = nn.Linear(d, vocab, bias=False)

    def forward(self, x):
        B, L = x.shape
        h = self.tok(x) + self.pos(torch.arange(L))[None]
        mask = torch.triu(torch.ones(L, L, dtype=torch.bool), diagonal=1)
        for attn, ln in zip(self.layers, self.lns):
            hn = ln(h)
            out, _ = attn(hn, hn, hn, attn_mask=mask)
            h = h + out
        return self.head(h)

def make_batch(rng, B, V, Tlo=12, Thi=20):
    T = int(rng.integers(Tlo, Thi + 1))   # 周期を可変に(位置ショートカット封じ)
    seqs = np.stack([np.tile(rng.permutation(V)[:T], 2) for _ in range(B)])
    return torch.tensor(seqs), T

def train_lm(nlayers, seed=0, steps=1500, V=30):
    torch.manual_seed(seed)
    rng = np.random.default_rng(seed)
    model = AttnOnlyLM(vocab=V, nlayers=nlayers)
    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
    for step in range(steps):
        x, T = make_batch(rng, 64, V)
        logits = model(x)
        loss = F.cross_entropy(logits[:, :-1].reshape(-1, V), x[:, 1:].reshape(-1))
        opt.zero_grad(); loss.backward(); opt.step()
    return model.eval()

m2, m1 = train_lm(2), train_lm(1)
xe, Te = make_batch(np.random.default_rng(999), 128, 30)
for name, m in [("2層", m2), ("1層", m1)]:
    with torch.no_grad():
        lg = m(xe)
    logp = F.log_softmax(lg[:, :-1], -1)
    nll = -logp.gather(-1, xe[:, 1:, None]).squeeze(-1)
    print(f"{name}: 1周目損失={nll[:, :Te-1].mean():.3f} 2周目損失={nll[:, Te-1:].mean():.3f}")
# => 2層: 1周目損失=3.298 2周目損失=0.164
# => 1層: 1周目損失=3.131 2周目損失=2.861  (一様分布なら ln30 = 3.401)

訓練曲線: 2層attention-onlyだけが2周目を当てられるようになる

結果は理論の予言どおりです。2層モデルの「2周目損失」は 0.164 まで落ちる(ほぼ完全な予測)のに対し、1層モデルは 2.861 で、一様分布の $\ln 30 = 3.401$ からわずかに下がった程度。1周目(初出のランダム列)はどちらのモデルも当てられません——当てられるはずがない部分と、誘導で当てられる部分が、きれいに分離しています。

位置ごとの損失を見ると、この分離はさらに鮮明です。

位置別損失: 繰り返しの境界で損失が崖のように落ちる

繰り返しが始まる位置を境に、損失が 3.3 から 0.2 へ崖のように落ちます。「文脈の後半ほど予測が上手くなる」——これはまさに in-context learning の定義(文脈から学べている)を、最小の設定で観測した姿です。

では、訓練されたモデルの中身は、手作り回路と同じ構造になっているのでしょうか。

学習後のattentionパターン: 手作り回路と同じ構造が自然に生えている

なっています。第1層のあるヘッドは「直前トークン」への注意(サブ対角、平均0.16。一様なら0.04)、第2層のあるヘッドは2周目で「前回の自分の直後 $i-T+1$」への注意(prefix-matchingスコア0.16、同じく一様なら0.04)を示します。手作り回路ほど純粋ではなく注意は複数の場所に分散していますが——実際のモデルのヘッドは複数の仕事を兼務するのが普通です——previous-token head → induction head という2段構えの回路が、勾配降下から自然に発生しました。

相転移: 誘導ヘッドとICLは同時に生まれる

Olsson らの論文でもっとも印象的な発見は、誘導ヘッドが訓練の特定の時期に「相転移」的に一斉形成されることです。

誘導ヘッドは訓練中の相転移で形成される(原論文より)

出典: Olsson et al., “In-context Learning and Induction Heads”, Transformer Circuits Thread / arXiv:2209.11895, 2022

各線は1つの attention ヘッドの prefix-matching スコアです。1層モデル(左)では誘導ヘッドが一切形成されない一方、2層以上のモデル(中央・右)では、訓練のある狭い区間(黄色のハイライト)で複数のヘッドのスコアが一斉に立ち上がります。そしてこの区間は、ICL能力(文脈後半のトークンほど損失が低い度合い)が急改善する時期と正確に一致します。損失曲線にはこの時期に小さな「こぶ」さえ現れます。

私たちのおもちゃの実験でも、同じ同期が観察できます。

prefix-matchingスコアとICLスコアは同時に立ち上がる

ICLスコア(1周目損失 − 2周目損失)は訓練ステップ200〜600で0から3.1へ急上昇し、第2層ヘッドの prefix-matching スコアの立ち上がりとおおむね並走します。「誘導回路の形成」と「文脈から学ぶ能力の獲得」が同じ現象の裏表であることを示唆する動きです。

証拠の強さを区別する

ここで、Olsson らの主張の「強さ」を正しく理解しておくことが大切です。論文自身が証拠のレベルを明確に区別しています。

  • 小型のattention-onlyモデル(本記事の実験のような設定)では、誘導ヘッドがICLの主要因であることは因果的に強く支持される(ヘッドを消すとICLが消える、形成時期が一致する、など)
  • 大型の実用モデルでは、証拠は相関的です。誘導ヘッドの形成時期とICL改善の一致、誘導ヘッドの摂動がICLを損なうことなどは確認されていますが、「ICLのすべてが誘導ヘッドで説明できる」とまでは言えません

また、大型モデルの誘導ヘッドは「完全一致の [A][B]…[A]→[B]」を超えて、意味的に類似したパターンの継続([A][B]…[A]→[B] のようなソフトな照合や、翻訳のような抽象パターン)にも働くことが観察されています。few-shot プロンプトの「例示の続きを出す」動作は、この一般化された誘導の姿だと考えると腑に落ちます。

まとめ

本記事では、in-context learning の機構的な正体候補である誘導ヘッドを、定義から実験まで通しで解説しました。

  • 誘導ヘッド = prefix matching + copying: 「前回の [A] の直後」に注意を向け、そこにあった [B] をコピーする。文脈内の任意のパターンに働くため、ICLの部品になれる
  • 2層以上で初めて可能: 第1層の previous-token head が「直前トークン」を残差ストリームに書き込み、第2層が K-composition でそれを検索キーに使う。1層では検索キーに文脈を入れられない
  • 手作り回路で決定的に再現: 学習ゼロの重み設計で2周目の次トークン的中率100%。回路の設計図が正しいことの直接証明
  • 勾配降下でも自然に生まれる: 2層attention-onlyの訓練で2周目損失が3.4→0.164に急落(1層は2.861で失敗)。学習後のモデル内部に previous-token head と induction head が確認できた
  • 相転移: 誘導ヘッドの形成とICL能力の獲得は訓練中の同じ時期に同期して起きる。小型モデルでは因果的、大型モデルでは相関的な証拠——という主張の強さの区別も論文の誠実な美点

「LLMはその場で学んでいるように見える」の裏側には、文脈を連想記憶のように検索してコピーする、驚くほど具体的な回路がありました。機構的解釈という分野は、この調子でモデルの中の「プログラム」を1つずつ読み出そうとしています。

次のステップとして、以下の記事も参考にしてください。

画像なし
Attentionの数式を行列の形で完全に理解する
注意行列ヒートマップの読み方。本記事の実験図を深く理解する前提
画像なし
Modern Hopfield NetworksとAttentionの等価性
attention=連想記憶の読み出しという理論的視点。誘導ヘッドの「文脈を検索する」動作と響き合う
画像なし
Attention Sink(注意の吸い込み)とは
attentionパターンの機構的解釈のもう1つの事例。先頭トークンに注意が集まる理由