この論文のクライマックスだ。アテンション層が二つになった瞬間、1層では原理的に不可能だったことが可能になる。前の層のヘッドの出力が、後の層のヘッドの入力になること。つまり合成(composition)だ。
合成の三種類
第2回の絵を思い出そう。ヘッドは残差ストリーム(residual stream)という共用ボードから読み、ボードに書く。層が二つになると、新しいことが可能になる。前の層のヘッドがボードに書いた結果を、後の層のヘッドが読んでいけるのだ。そして第2回で見たように、後の層のヘッドがボードから読む場所は三つある。検索語(クエリ)を作るとき、索引(キー)を作るとき、中身(バリュー)を作るとき。前の層の結果がどこへ流れ込むかによって、合成も三種類になる。
- Q-合成(Q-Composition): 前層の出力が後層のクエリに影響する。どこを見るかが前層の計算に依存するようになる。
- K-合成(K-Composition): 前層の出力が後層のキーに影響する。各位置がどれだけ注目されるかが前層の計算に依存するようになる。
- V-合成(V-Composition): 前層の出力が後層のバリューに影響する。運ばれる中身そのものが前層で加工されたものになる。
経路展開(path expansion)を2層に適用すると項が増える。直接経路(direct path)、ヘッドを一つだけ通る経路、そして前層ヘッドと後層ヘッドを続けて通る経路。特にV-合成でつながれた二つのヘッドの項は $(A^{h_2} A^{h_1}) \otimes (W_{OV}^{h_2} W_{OV}^{h_1})$ の形にまとまる。パターンはパターン同士、OVはOV同士が掛け合わされた、事実上一つの新しいヘッドだ。論文はこれを仮想アテンションヘッド(virtual attention head)と呼ぶ。なぜV-合成だけがこうまとまるのかは、言葉でも説明できる。情報の移動に情報の移動を合成しても、やはり情報の移動だ。一方、Q-合成とK-合成はアテンションパターンそのものを変える仕事なので、一つのヘッドに還元されない新しいものを作る。
図1. 2層ロジットの経路展開。直接経路の項はバイグラムを担い、個別ヘッドの項は1層のときと同じで、最後の仮想アテンションヘッドの項がV-合成に当たる。
誘導ヘッド
学習済みの2層モデルを分析すると、合成を実際に使うはっきりしたパターンが繰り返し見つかる。誘導ヘッド(induction head)だ。
やることは規則一つに要約できる。文脈で[A][B]を見たことがあり、今のトークンが再び[A]なら、次のトークンとして[B]を予測せよ。前に出たパターンをそのまま続けるのだ。探しているトークンが文脈になければ、開始トークンに注目したまま何もしない。一種の待機位置(resting position)だ。
図2. ハリー・ポッター冒頭で動く誘導ヘッド1:8のアテンションパターンとロジット効果。現在のトークンのクエリが、前回の出現の直後のトークンのキーに注目し、そのトークンのロジットを押し上げる。
仕組みは二つのヘッドの協業で、核心はK-合成だ。具体例でたどってみよう。文脈に「Harry Potter」が出てきて、今のトークンが再び「Harry」だ。次は「Potter」であるべきだ。
- 前層の**直前トークンヘッド(previous token head)**が準備作業をする。各位置で一つ前のトークンを読んで、残差ストリームに書いておくのだ。これで「Potter」の位置には「私の前はHarryだった」という付箋が付いている。
- 後層の誘導ヘッドは、索引(キー)をその付箋から作る。これがK-合成だ。だから「Potter」の位置の索引は「前がHarryだった場所」になる。今のトークン「Harry」の検索語(クエリ)はまさにそういう場所を探し、アテンションは「Potter」に突き刺さる。
- OV回路(output-value circuit)はただのコピーだ。注目した「Potter」を次トークン予測として押し上げる。
一般化すれば [A][B] … [A] → [B] だ。第2回の言葉で要約すると、QKは前層の付箋のおかげでずっと賢い検索(「前回の自分の直後の場所」)ができるようになり、OVは第4回のコピーそのままだ。新しい部品ができたのではなく、新しい組み合わせができたのである。
図3. QK回路をヘッドではなくトークン単位で展開した図。キーとクエリの強さは、各トークンがアテンションスコアをどれだけ上げるかを表す。キーが1トークン後ろにずれて計算されるため、同じトークンを探すクエリはその直後の位置に刺さる。
1層のコピーヘッドと比べると質的な違いが見える。1層のコピーは「出たトークンはまた出る」という水準だ。誘導ヘッドは「出たパターンは続く」を実装し、そのパターンが何であるかを問わない。完全にランダムなトークン列を繰り返させても、つまり学習分布から大きく外れた入力でも動く。統計を暗記したのではなく、文脈から規則を読むアルゴリズムを学んだのだ。
論文はこの理論を二つの方向から検証する。一つは重みの検証だ。このアルゴリズムが正しいなら、誘導ヘッドのOV回路はコピー行列であり、K-合成が作るQK項は同じトークンを探す一致行列でなければならない。前回の固有値の物差しで二つの軸を測ると、誘導ヘッドだけが正-正の極端な隅に集まる。もう一つはアブレーションだ。展開項を次数ごとに消してみる項の重要度分析(term importance analysis)は、性能の大部分が2層ヘッドの項にあり、仮想ヘッドの項の取り分は小さいことを示す。
図4. 2層のヘッドを、OV回路の固有値の正値性(横)とK-合成QK項の正値性(縦)で配置した平面。誘導ヘッドだけが右上、コピーと一致の両方が強い極端な隅に集まる。
合成は思ったより希薄だ
注意すべき経験的発見が一つ。彼らの2層モデルで、合成はどこにでもあるものではなく希薄だった。ほとんどのヘッドは合成をほぼ使わず、1層モデルのようにスキップトライグラム(skip-trigram)を計算する。モデル全体は、おおむね大きな1層モデルの上に誘導ヘッドを載せたもののように振る舞う。新しい能力は、層を満たす洪水ではなく、数本の細い流れとして入ってきたわけだ。
正誤表を一つ。論文公開後、この合成測定に使われた著者らのライブラリにバグが見つかり、修正版の図が文書に追加された。実際のモデルには最初に見えたよりも多くの合成があり、一部の誘導ヘッドは直前トークンヘッドだけでなく、直近数トークンを見るヘッドとも合成していた。ただし核心の結論、つまり直前トークンヘッドとのK-合成が誘導ヘッドの骨格だという事実はそのままだ。主張と訂正を同じ文書に添えておくこのスレッドの文化には、最終回で再会する。
図5. 正誤表の合成ダイアグラム。上がバグのあった元の図、下が修正版(青緑: 誘導ヘッド、赤: 直前トークンヘッド)。修正版ではより多くの合成が見えるが、K-合成という骨格はそのままだ。二つの図の線の太さは直接比較できない。
仮想アテンションヘッドも同様だ。理論上、仮想ヘッドはヘッド数の積だけ存在しうるので、大きなモデルでは膨大な表現力の源になりうる。論文が挙げる状況証拠を一つ。直前のトークンを見るヘッドはよくあるのに、二つ前を見るヘッドはあまり見つからない。二つ前の予測力は、直前トークンヘッド二つの仮想合成で得られるからかもしれない。人称や時制のように運ぶ情報がごく少ない仕事には、丸ごと席を占める正規ヘッドより仮想ヘッドの方が経済的でもある。ただしこの論文が分析した小さなモデルで明確に大きな役割を果たしていたのは誘導ヘッドの方で、仮想ヘッドの取り分は小さかった。可能性の兆しと確認された事実を区別する論文の態度は見習う価値がある。
今回の要点
- 2層からQ-、K-、V-合成が可能になり、V-合成の経路は仮想ヘッドという新しい項になる。
- 誘導ヘッドは直前トークンヘッドとのK-合成で [A][B] … [A] → [B] を実装する。
- これは任意のパターンで動く文脈内学習(in-context learning)アルゴリズムであり、のちの大規模モデル研究の中心素材になる。
- ただし観察された合成は希薄だ。2層モデルはおおむね、大きな1層モデル + 誘導ヘッドのように振る舞う。
次回は論文の締めくくりだ。このフレームワークはどこまで来て、何が残ったのか。
原文: A Mathematical Framework for Transformer Circuits のTwo-Layer Attention-Only Transformersセクション。本文の図はすべて原文からの引用。