今回から、AnthropicのA Mathematical Framework for Transformer Circuits(2021)をセクション単位でたどっていく。この論文はTransformer Circuitsスレッド全体の土台だ。後に出るToy Models of Superpositionも、誘導ヘッド(induction head)の論文も、ここで定義された言葉の上に立っている。
論文が掲げる動機は安全性だ。言語モデルは、訓練が終わって何年も経ってから、作った本人すら知らなかった能力や問題行動が見つかることがある。コンパイルされたバイナリをソースコードにリバースエンジニアリングするように、モデルの内部計算を読み取れるなら、現在の安全性の問題を説明し、新しい問題を見つけ、さらにはまだ作られていない強力なモデルの問題を先取りする体系的なアプローチが可能になる、というわけだ。
戦略: モデル生物
生物学が大腸菌とショウジョウバエから始まったように、この論文は最小のトランスフォーマーから始める。MLPを取り除いたアテンション専用モデル、それも0層、1層、2層のもの。目標は性能ではなく完全な理解だ。小さなモデルを最後までリバースエンジニアリングできれば、そこで得た概念が大きなモデルを見るときの言語になる。
この戦略は実際に機能した。2層のトイモデルで発見された誘導ヘッドは、のちに大規模モデルの文脈内学習(in-context learning)を説明する最有力候補になる。
主要結果の地図
論文が要約する結果を、先に一枚にまとめておく。それぞれは今後の回で一つずつ扱う。
残差ストリーム(residual stream)は通信チャネルである。 トランスフォーマーの各層は残差ストリームから読み、残差ストリームに足す。ストリーム自体は何の計算もしない線形の通路だ。おかげで、離れた層同士も仮想重み(virtual weights)という積を通じて直接つながっているかのように分析できる。
アテンションヘッドは独立で加算的である。 一つの層の複数のヘッドは絡み合っているように見えるが、数学的には、各自が計算して残差ストリームに足す独立した項に分解される。
ヘッドは二つの回路に分かれる。 どこに注目するかを決めるQK回路(query-key circuit)と、注目した場所から何を運ぶかを決めるOV回路(output-value circuit)。この二つは互いに独立に分析できる。アテンションとは情報の移動である。
層が増えるたびに質的に新しい能力が生まれる。 0層モデルはバイグラム統計しか学べない。1層モデルはスキップトライグラム(skip-trigram)まで学ぶ。2層モデルからヘッドの合成(composition)が可能になり、ここで誘導ヘッドが登場する。誘導ヘッドは、文脈の中で以前見たパターンを繰り返すメカニズム、つまり文脈内学習の原型だ。
道具は経路展開(path expansion)である。 モデル全体を、入力から出力に至る経路の和として書き直すと、各経路が一つの解釈可能な項になる。この論文の数学の大部分は、この展開を正当化して活用する作業だ。
なぜこの論文から読むのか
機械論的解釈可能性の多くの論文は「何を発見したか」を語る。この論文はその前に「何で見るか」を決める。残差ストリーム、QK/OV、合成、誘導ヘッド。これらは今やこの分野の標準語彙だ。語彙を先に身につけておけば、後の論文がずっと速く読める。
次回は、この論文のやり方でトランスフォーマーを描き直す。
原文: A Mathematical Framework for Transformer Circuits (Elhage, Nanda, Olsson他, Anthropic, 2021)。