Transformer Circuits数学的フレームワークとは — 誘導ヘッド発見の道筋

Transformer Circuits論文とは何を目指した研究か

大規模言語モデルは実世界での利用が急速に広がる一方、モデルの規模が大きくなるほど予期しない挙動の余地も増えます。訓練から何年も経ってから、開発者やユーザーが問題のある挙動を発見するケースも珍しくありません。

この課題への一つのアプローチが「メカニスティック解釈可能性(mechanistic interpretability)」です。プログラマーが複雑なバイナリを人間可読なソースコードへ逆コンパイルするように、Transformerが行う計算を詳細に逆解析しようという試みです。画像認識モデルに対してはDistillのCircuitsスレッドが同様の取り組みを行っていましたが、Transformerや言語モデルに対する比較可能な研究はそれまで存在しませんでした。2021年12月にAnthropicが公開した「A Mathematical Framework for Transformer Circuits」は、この空白を埋めるための最初の一歩です。

なぜ「2層以下・attentionのみ」から始めたのか

現代の言語モデル(たとえばGPT-3)は96層を持ち、attentionブロックとMLPブロックを交互に積み重ねた構造をしています。この論文が対象にしたのはその対極、attentionブロックのみで構成された2層以下のTransformerです。

複雑で巨大なモデルをいきなり解析するのではなく、最も単純なモデルから出発し、そこで見つけたアルゴリズム的なパターンや枠組みを、より大きく複雑なモデルへ段階的に適用していくという戦略です。MLP層についてはこの時点でまだ理解が進んでおらず、著者らも「大きな弱点」と認めたうえで、まずattentionのみのモデルに範囲を絞っています。

モデルの見方を組み替える — 残差ストリームという通信路

論文の核心は、Transformerの計算を数学的に等価な、しかしより解釈しやすい形に書き換えることにあります。中心になる概念が「残差ストリーム(residual stream)」です。

Transformerの各層(attention層・MLP層)は、残差ストリームから情報を線形変換で「読み込み」、計算結果を線形変換で残差ストリームへ「書き込み」ます。残差ストリームはそれ自体では何も処理せず、すべての層がこの通信路を介して情報をやり取りしているにすぎません。この構造が完全に線形であることから、層をまたいだ「仮想的な重み(virtual weights)」を考えることができます。ある層の出力重みと別の層の入力重みを直接掛け合わせることで、間に他の層がいくつ挟まっていても、2つの層の間の情報伝達を直接記述できるのです。

もう一つの重要な整理が「attentionヘッドは独立かつ加算的である」という視点です。通常はヘッドの出力を連結してから1回の行列積で処理する実装(計算効率重視)で説明されますが、数学的には各ヘッドが独立に動作し、その出力をそれぞれ残差ストリームに足し込んでいるのと同じです。解釈可能性の観点では、この「独立に加算する」見方のほうが扱いやすいと著者らは指摘しています。

残差ストリームが高次元のベクトル空間であることも見逃せません。小さなモデルでは数百次元、大きなモデルでは数万次元に達します。各層は異なる部分空間に情報を書き込むことで、他の層と干渉せずに情報をやり取りできます。いったん書き込まれた情報は、別の層が能動的に削除しない限り残り続けるため、残差ストリームの次元は一種の「メモリー」や「帯域」のように機能します。実際には計算に使われるニューロンの数のほうが残差ストリームの次元よりはるかに多いため、この帯域は常に取り合いの状態にあると著者らは指摘しています。一部のMLPニューロンやattentionヘッドが、他の層が書き込んだ情報を打ち消す形で残差ストリームを掃除する「メモリー管理」のような役割を担っている可能性も報告されています。

QK回路とOV回路への分解

attentionヘッドの計算はさらに2つの独立した部分に分解できます。

この2つは概念的に独立しており、QK回路が「どこを見るか」、OV回路が「見た情報をどう使うか」をそれぞれ担当します。この分解の土台になっているのが「Path Expansion Trick」という論文の中核的な手法です。attentionのみのTransformerは複数の層を掛け合わせた積の形で書けますが、この積を展開すると、入力トークンから出力までのすべての経路が個々の項として並びます。各項は独立に読める1つの「回路」になり、それらを足し合わせたものがモデル全体の挙動です。QK回路とOV回路への分解も、この展開で現れる各attentionヘッドの項をさらに読み解いた結果にほかなりません。

層数を増やすと何が起きるか

論文はモデルを層数ごとに段階的に解析しています。

0層Transformerはバイグラム統計をモデル化するだけで、その対応表は重みから直接読み取れます。1層のattentionのみTransformerは、バイグラムと「スキップトライグラム」(A…B Cのような形の系列)モデルのアンサンブルとして機能し、この表も重みから直接アクセスできます。スキップトライグラムは見た目以上に表現力があり、ごく単純な文脈内学習(in-context learning)まで実装してしまいます。

この1層モデルの解析では、展開したQK・OV行列を実際に覗くことで、モデルの奇妙な失敗パターンも見えてきました。スキップトライグラムはOV・QK行列に分解された「factored form」で表現されるため、keep… in mind と keep… at bay の両方の確率を上げるヘッドは、原理的に keep… in bay や keep… at mind の確率まで一緒に持ち上げてしまいます。queryとkeyが分離した構造ゆえの副作用で、望ましい継続と一緒に望ましくない継続まで学習してしまう「バグ」です。些細な現象に見えますが、著者らはこれを解釈可能性によってモデルの失敗を理解できることを示した初期の実証例として重視しています。

問題は2層のattentionのみTransformerです。ここでattentionヘッドどうしの「合成(composition)」が可能になり、より複雑なアルゴリズムを実装できるようになります。合成にはquery・key・valueの3種類があり、そのうちkey合成(K-composition)とquery合成(Q-composition)がvalue合成(V-composition)とは質的に異なる働きをすることが分かりました。

誘導ヘッド(induction head)の発見

2層モデルにおける合成のもっとも重要な用途が、論文が「誘導ヘッド」と名付けた仕組みです。

誘導ヘッドは、文脈中で現在のトークンが以前どこに出現したかを探し、見つかればその「次のトークン」を見てコピーします。これにより、以前出現した系列を正確に、あるいは近似的に繰り返せるようになります。1層モデルの単純なコピーヘッドと比較すると違いが際立ちます。

誘導ヘッドはトークン間の学習済み統計(このトークンの次にこのトークンが来やすいか)に依存しないぶん、分布のずれにも強いという特徴があります。論文はこれを完全にランダムなトークン列の繰り返しに対しても機能することで実証しました。

仕組み: 誘導ヘッドの核心は、キーが1トークン分うしろにずれた状態で計算される点にあります。まず「1つ前のトークンヘッド」(K-composition)がキーベクトルを1トークン分前方にシフトします。queryは「似た」キーベクトルを探しますが、キーがずれているため、実際には次のトークンを見つけることになります。この理論の妥当性は、誘導ヘッドが「コピーするOV回路」と「一致を検出するQK回路」の両方を備えているかどうかを固有値の統計量で検証することで確認されました。実際に誘導ヘッドは、QK・OVの固有値の正負を軸にした2次元空間上で、極端な一角に集まっていたと報告されています。

仮想attentionヘッドという理論的な広がり

論文はさらに「仮想attentionヘッド」という概念も扱っています。2つのヘッドがV-compositionで合成されると、あたかも1つの独立したヘッドであるかのように振る舞う項が生まれるというものです。この2層モデルの解析では、仮想attentionヘッドの寄与はアブレーション実験によって比較的小さいことが分かりました。ただし通常のヘッド数が層数に対して線形にしか増えないのに対し、2ヘッドの合成でできる仮想ヘッドの数は2次関数的に増えるため、より大きく複雑なモデルではもっと重要になりうると考えられます。

この研究はどこへ向かったか

著者らは論文の結びで、2層以下のattentionのみモデルという限定的な設定が、より大きく一般的なTransformerの理解にどこまで役立つかを検討しています。結論は「大きなモデルの一部を理解するのには使える」というものでした。実際、大規模モデルの中にも誘導ヘッドに類似した回路が見つかっており、その基本的な構築方法(1つ前のトークンヘッドとのK-composition)はこの論文で解析した小さなモデルと同じでした。誘導ヘッドは、あらゆる規模の言語モデルにおける文脈内学習の中心的な駆動源である可能性があると著者らは述べています。

ただし限界も明確です。MLP層は標準的なTransformerのパラメータの3分の2を占め、attentionヘッドの多くがMLP層とも相互作用するため、attentionのみの解析で理解できる範囲はさらに狭くなります。この限界を踏まえ、誘導ヘッドが文脈内学習にどこまで関わっているかを大規模モデルで直接検証する研究が、続く論文「In-context Learning and Induction Heads」で行われました。詳細は誘導ヘッドと文脈内学習の関係を検証した続編論文で扱っています。

まとめ

A Mathematical Framework for Transformer Circuitsは、Transformerの計算を残差ストリーム・QK回路・OV回路という単位に分解し直すことで、ごく小さなattentionのみモデルの内部動作をほぼ完全に説明した研究です。その過程で発見された誘導ヘッドは、その後の解釈可能性研究とClaudeの内部理解の重要な出発点になりました。同じ研究系譜のもう一つの成果である、モデル内部から数百万の概念を取り出したMapping the Mind論文についてはGolden Gate Claude論文の中身で扱っています。Claudeの内部活性化を自然言語で読み解く後続研究はNatural Language Autoencodersで紹介しています。