実装視点のトランスフォーマーの図は、リバースエンジニアリングには不向きだ。論文のTransformer Overviewセクションは、同じモデルを数学的に等価だが解釈に有利な形に描き直す。今回はその描き直しをたどる。
単純化
論文は分析対象をアテンション専用トランスフォーマーに絞る。MLP層は取り除き、レイヤー正規化とバイアスは隣接する重みに畳み込めるとして省略する。残るのはトークン埋め込み $W_E$、アテンション層、そして逆埋め込み $W_U$ だけだ。MLPを外すのはもちろん大きな単純化だ。しかしアテンションだけでも完全に理解することが目標なら、良い出発点になる。MLPの難しさには第6回で再会する。
図1. トークン埋め込みから残差ブロックを経て逆埋め込みに至る構造。アテンションヘッドとMLPは、それぞれの計算結果を残差ストリームに足す。
残差ストリーム: 何もしない通路
描き直した図の中心には残差ストリーム(residual stream)がある。各層は残差ストリームから入力を読み、計算結果をストリームに足す。ストリーム自体には非線形性がない。読みと加算だけの共有メモリ、すべての層が一緒に使う通信チャネルだ。
この視点の力は線形性から来る。ある層の出力は、以後のすべての層の入力にそのまま足されている。だから、離れた二つの構成要素の相互作用を、間にある層を飛ばして重み行列の積一つに要約できる。論文はこれを仮想重み(virtual weights)と呼ぶ。たとえば埋め込みがそのまま逆埋め込みにつながる直接経路(direct path)は、$W_U W_E$ という一つの行列だ。
図2. 左: 各層は線形写像で残差ストリームから読み、結果を足して書く。右: すべてが線形なので、重みを掛け合わせると層のペアを暗黙につなぐ仮想重みが現れる。異なる部分空間を使えば、特定の層だけに情報を送ることもできる。
線形性のもう一つの帰結として、残差ストリームには特権的な基底(privileged basis)がない。ストリームと相互作用するすべての行列を一緒に回転させても、モデルの挙動はそのままだ。座標軸の一本一本に意味を期待してはいけない空間だということだ。
ノート: 特権的な基底とは? 座標軸が意味を持つのは、活性化関数のように軸ごとに適用される演算があるときだ。たとえばMLPの活性値はニューロンごとに非線形が掛かるので、「3番ニューロン」という軸が実際の計算の単位になる。こうした空間には特権的な基底がある。一方、残差ストリームは読みも書きもすべて線形なので、空間全体を回転させても計算は変わらない。どの軸も特別ではなく、「何番目の次元は何を意味するか」という問い自体が成り立たない。意味を探すなら、軸ではなく方向、つまり複数次元の線形結合を見る必要がある。
狭すぎるチャネル: ボトルネックと重ね合わせの予告
トランスフォーマーには、すべての層が共有する、サイズの決まったメモリ(残差ストリーム)が一つしかない。そこに収めるべき計算結果は、それよりはるかに多い。だから複数の情報が同じ空間に重なって保存されるほかなく、これがトランスフォーマーの解釈を難しくする最大の理由の一つになる。この一段落がこのセクションのすべてだ。ここから一つずつほどいていこう。
残差ストリームの次元は有限だ。隠れ次元が4096のモデルなら、残差ストリームは長さ4096のベクトル一本。数字を4096個だけ書けるホワイトボード一枚が、情報の通る唯一の通路というわけだ。層はこのボードを丸ごと使うのではなく、部分空間(マスの束)を分け合って通信する。
ところが、書きたい側の数を数えると状況は穏やかでない。アテンションヘッドが数千個、MLPニューロンは数十万個。1層のMLPニューロンだけでも、普通はストリーム次元の4倍ある。それぞれが「私の結果も書きたい」と並ぶのに、ボードのマスは4096個しかない。
そこでボトルネックが生じる。50層モデルの真ん中、25層あたりの残差ストリームを見てみよう。前方25層が作った計算結果は、すべてこの一本のベクトルを経由する。前方のニューロンを合計すると約40万個、ストリーム次元の100倍だ。道路で言えば、100車線が4車線に合流し、また100車線に広がる地点である。論文はこの種のテンソルをボトルネック活性(bottleneck activations)と呼び、解釈が特に難しいだろうと見込む。
40万個の結果が、どうやって4096個の数字に収まるのか。ここで予告される概念が重ね合わせ(superposition)だ。ストリームの一つの次元が「犬」の情報だけを担うのではなく、犬に文法に感情に文脈まで、複数の意味を同時に担いうるということ。一つの音声信号に多くの周波数が混ざっているのと同じだ。ベクトルの中の 2.31 という数字一つにも、いくつもの意味が折り重なっていることがありうる。このテーマは後続論文 Toy Models of Superposition が正面から扱う。
そして、混み合った空間には掃除も要る。論文は、あるヘッドが「この情報はもう要らない」と判断してストリームから消す、メモリ管理(memory management)の役割を担う可能性にも触れている。書くヘッドと読むヘッドだけでなく、消すヘッドもありうるという話だ。
図3. 高次元の残差ストリームは部分空間に分かれる。層は重なる部分空間で相互作用し、互いに素な部分空間に書けば干渉しない。右のように、部分空間を読んで負を書き戻すことで情報を消す層もある。
ヘッドは独立で加算的
一つのアテンション層には複数のヘッドがある。実装ではヘッドの出力を連結して大きな行列を掛けるが、数学的には、各ヘッドが独立に計算して残差ストリームに足すのと等価だ。
$$\text{AttnLayer}(x) = x + \sum_{h} h(x)$$層という単位は実装の都合にすぎない。分析の自然な単位はヘッドだ。
QK回路とOV回路
アテンションヘッド一つは、実は二つの独立した問いに答える装置だ。どこを見るのか。そこから何を取ってくるのか。アテンションは普通、クエリ、キー、バリューという計算手順として学ぶので、二つは一塊に見える。論文の再解釈は、互いに独立な二つの機能がただ束ねられているだけだ、というものだ。この分解が、以後のすべての分析のテコになる。
具体例で見よう。「Tom likes apples. He …」まで読んで、次のトークンを予測するところだ。アテンションヘッドは二段階で動く。
第一段階、どこを見るか。現在のトークン「He」が検索語(クエリ)を出し、前方のトークンはそれぞれ索引(キー)を持っている。クエリを索引と照合してスコアを付け、ソフトマックスで配分したものがアテンションパターンだ。
Tom 0.93
likes 0.02
apples 0.05
この段階が決めるのは「Tomを見ろ」だけだ。Tomが人なのか、男性なのか、主語なのかは、ここにはない。どこを見るかだけを決める。
第二段階、何を取ってくるか。Tomを見ることにした。だがTomの位置の残差ストリームには情報が何層にも重なっている。人、男性、主語、名前、単数。ヘッドはこれを全部運ぶのではなく、自分の分だけを選び取る。たとえば「三人称単数」だけを抜き出し、変換して「He」の位置の残差ストリームに書き込む、という具合だ。次の予測が eat ではなく eats になるために必要な、まさにその情報である。
たとえるなら、第一段階はカーナビで、第二段階は宅配ドライバーだ。どの家に行くかを決めることと、その家でどの箱を積むかは、別の仕事だ。
さて、この二段階を行列で見よう。計算には四つの行列 $W_Q, W_K, W_V, W_O$ が使われるが、単独では意味を持たない。$W_Q$ と $W_K$ は第一段階で常に一緒に動き、$W_V$ と $W_O$ は第二段階で常に一緒に動く。だから実体は、積で束ねられた二つの行列だ。
QK回路(query-key circuit) $W_{QK} = W_Q^\top W_K$。トークンのペアごとに「このクエリはこのキーにどれだけ惹かれるか」というスコアを一つ出す表だ。第一段階の全体が、この表一つに要約される。
OV回路(output-value circuit) $W_{OV} = W_O W_V$。「このトークンが注目されたら出力にどんな影響を与えるか」を収めた表だ。第二段階の全体が、この表一つに要約される。
サイズの感覚も持っておこう。ストリームが4096次元でヘッドが128次元なら、$W_{QK}$ は 4096 x 4096 の表のように振る舞うが、実際の自由度はランク128しかない。巨大に見えても、実体は小さな低ランク行列二つだということ。第4回でこの表を実際に読むとき、この事実が計算を可能にしてくれる。
この視点から見ると、キー、クエリ、バリューのベクトルは副産物だ。二つの行列を別の形に因数分解すれば全く違う中間ベクトルができるが、モデルは同一に動作する。論文が、これらのベクトルに一切言及せずにトランスフォーマーを記述する方が有用なことが多いと述べる理由だ。
分解の効用は組み合わせに現れる。「常に直前のトークンを見る」という同じQK回路を持つ二つのヘッドでも、一方のOVが品詞情報を運び、もう一方のOVが感情情報を運ぶなら、二つは全く別の機能のヘッドだ。逆にOVが同じでも、QKが直前トークン、主語、同じ単語、文頭と変われば、やはり完全に別のヘッドになる。ヘッドの機能は二つの回路の組み合わせであり、だからヘッドの説明には二つの文が要る。「このヘッドは直前のトークンを見る」はQKについての説明で、「このヘッドは文法情報を伝える」はOVについての説明だ。
分離は分析のテコでもある。どこに注目するかが決まってしまえば、届く効果はソーストークンだけの関数なので、二つの表を別々に読める。第4回でスキップトライグラムを表として読み取り、そのバグまで予測できるのは、すべてこの分離のおかげだ。そしてこの分解に慣れれば、第5回の誘導ヘッドも簡単になる。QKが「前回の自分がいた場所」へ歩いていき、OVがその隣のトークンをコピーする、という組み合わせとして読めるからだ。
そこで論文のスローガンが出てくる。アテンションとは情報の移動である。ヘッドは、ある位置の残差ストリームから情報を読み、別の位置の残差ストリームへ運ぶ装置だ。
図4. アテンションヘッドは、あるトークンの残差ストリームから別のトークンの残差ストリームへ情報をコピーする。通常、読んだ場所とは別の部分空間に書く。
付け加えておくべき重要な事実が一つある。アテンション専用モデルで非線形性は、アテンションパターンを作るソフトマックスだけだ。だからパターン $A$ を固定されたものとして扱った瞬間、モデルの残り全部は線形になる。次回から使う経路展開(path expansion)は、正確にこの事実の上に立っている。
今回の要点
- 残差ストリームは計算しない線形の通信チャネルであり、その線形性が仮想重み分析を可能にする。
- 分析の単位は層ではなくヘッド。ヘッドは独立で加算的。
- ヘッドは、どこを見るか(QK)と何を運ぶか(OV)に分解される。
- 非線形性はソフトマックスだけ。パターンを固定すればモデルは線形で、これが経路展開の根拠になる。
- ストリームは全員で分け合う狭いメモリだ。重ねて書くこと、つまり重ね合わせがここで予告される。
次回は最小のモデル、0層トランスフォーマーだ。
原文: A Mathematical Framework for Transformer Circuits のTransformer Overviewセクション。本文の図はすべて原文からの引用。