アテンション層が一つできると、モデルは初めて文脈を見る。今回は1層のアテンション専用モデルを最後まで開いて、このモデルにできることの全部を確認する。
図1. 1層のアテンション専用トランスフォーマー。埋め込まれたトークンに各ヘッドの出力が足され、すぐに逆埋め込みされてロジットになる。
経路展開
このモデルの出力は、二種類の道の和だ。一つは前回見た直接経路(direct path)、つまり内蔵の自動補完。もう一つはアテンションヘッドを通るヘッド経路たちだ。数式で書くと、正確にこう展開される。
$$T = \underbrace{W_U W_E}_{\text{直接経路}} + \sum_h \underbrace{A^h \otimes (W_U W_{OV}^h W_E)}_{\text{ヘッド } h \text{ の経路}}$$数式の $\otimes$ は二つの次元の結合を表す表記だ。左側(パターン)はどの位置から取ってくるかを、右側(行列)は取ってきた内容がどう変わるかを受け持つ。第2回のQK/OV分解が、数式にそのまま刻まれているわけだ。
図2. 層ごとに掛けられていた項を展開すると、直接経路とヘッド経路の和になる。アテンションパターンは、トークンのペアを $W_{QK}$ の両側に掛けたうえで、自己回帰マスク付きソフトマックスを取って計算される。
各ヘッド経路の正体は、第2回で見た二枚の表だ。どこを見るかを決めるQK表と、注目されたトークンが出力に何を押し上げるかを決めるOV表。どちらもトークンからトークンへの表なので、モデルを動かさず、重みを読むだけで解釈できる。
図3. ソーストークンからロジットまで続くOV回路(金色、$W_U W_O W_V W_E$)と、ソース・目的地トークンのペアにアテンションスコアを与えるQK回路(赤紫、$W_E^\top W_Q^\top W_K W_E$)。それぞれ一つの行列にまとまる。
スキップトライグラム
表を実際に読むと、どんな規則が出てくるのか。具体例で見よう。文脈の前方に「keep」があり、今のトークンが「in」だ。
第一段階、QK表を引く。「in」のクエリは前方の「keep」に強く惹かれるよう学習されている。
(keep, in) 大きい <- 「in」のときは「keep」を見ろ
第二段階、OV表を引く。「keep」が注目されると「mind」の点数が上がるよう学習されている。
(keep -> mind) 大きい <- 「keep」を見たら「mind」を上げろ
二つの表がつながって、規則が一つ完成する。「keep … in → mind」。前方どこかの[ソース]、今の[目的地]、押し上げる[出力]。ソースと目的地が離れていてもよい3トークンのパターンなので、論文はこれをスキップトライグラム(skip-trigram)と呼ぶ。学習済みヘッドからはこの種の規則が大量に見つかり、これが1層モデルの能力のすべてだ。自動補完の表の上に、スキップトライグラムの表を重ねたもの。
図4. 展開されたQK/OV行列の大きな項目。ソーストークンごとに、注目する目的地トークンと確率が上がる出力トークンを読めば、そのままスキップトライグラムの表になる。4行目はLaTeXを、5行目はHTMLエスケープを学んだ痕跡だ。
すべてのヘッドが内容を見て注目するわけではない。主に位置を根拠に、たとえば直前のトークンの近くに注目するヘッドも見つかる。今は脇役に見えるが、次回、直前トークンヘッド(previous token head)という名前で誘導ヘッド(induction head)のパートナーになる。
コピーヘッドと固有値
見つかったスキップトライグラムの最も多い形はコピーだ。「Tom … T → Tom」のように、出たトークンをまた押し出す規則。「文書にすでに出た単語はまた出やすい」という統計の実装であり、論文はこれを原始的な文脈内学習(in-context learning)と呼ぶ。役割分担はもうおなじみだ。OV表は注目されたトークン自身の点数を上げるようセットされ、QK表はそのトークンが次に来そうな場所からだけ後ろを振り返る。コピーするが、自動補完が許す場所にだけコピーするのだ。
面白い特殊例もある。トークナイザーは普通、空白を単語の頭にくっつけて切るので、珍しい単語が空白なしで現れると二つのトークンに割れる(" Ralph"は1トークン、“Ralph"は"R”+“alph”)。一部のヘッドはこのケースを部分的に専門とし、断片トークン"R"を見ると前方の" Ralph"を探して"alph"を予測する。次回見る誘導ヘッドを、1層モデルが真似ている非常に特殊な事例というわけだ。
図5. 空白なしで割れた単語のコピーを扱う項目。完全な単語(" Ralph")をソースに、断片トークン(“R”)を目的地にして、残りの断片(“alph”)を出力する。
コピーヘッドかどうかは、どうすれば一目で分かるのか。直観は表の対角線だ。「Tomを見たらTomを上げる」という規則は、OV表の対角線のマスが大きいという意味である。ただし 5万 x 5万 の表の対角線を一つずつ見る代わりに、論文は固有値という要約統計一つでこれを測る。固有値とは、行列が特定方向のベクトルを伸ばす倍率のことだ。$Mv = \lambda v$ で $\lambda$ が正なら、その方向 $v$ は反転せずに拡大されて出てくる。だからOV表の固有値が正に偏っているというのは、トークン(の組み合わせ)を入れるとまさにそのトークンたちの点数が大きくなる、つまり表が全体としてコピーだという意味だ。この物差しで測ると、分析した12ヘッドのモデルでは10個のヘッドがはっきりしたコピー傾向を示した。巨大に見えるこの行列は実はランク $d_{head}$(64か128)の低ランク行列なので、計算も可能だ。ただし但し書きが付く。コピー行列は必ず正の固有値を持つが逆は成り立たないので、正の固有値は「平均的にコピーしている」という強い証拠であって、決定的な証明ではない。この要約統計は後の論文でも使われ続ける。
図6. 固有値のうち正の割合でヘッドを並べたヒストグラム。12個中10個が右端、つまりコピー側の極端に集まっている。
分解構造が生むバグ
同じヘッドが「keep … in → mind」と「keep … at → bay」を両方学んだとしよう。ヘッドの規則は常にQK表とOV表、二枚の積だ。QK表では (keep, in) と (keep, at) のマスが大きく、OV表では (keep → mind) と (keep → bay) のマスが大きい。ところが積の構造には、「inのときはmindだけ、atのときはbayだけ」と組み合わせを縛るマスがない。だから交差した組み合わせ「keep … in → bay」にも確率が漏れる。三変数の関数 $f(a, b, c)$ を $f_1(a, b) \cdot f_2(a, c)$ の形でしか書けないようなもので、三つのトークンが一緒に作る相互作用を自由に表現できないのだ。モデルの誤りがランダムなノイズではなく構造の論理的帰結だということ。リバースエンジニアリングが与える予測力の好例だ。
図7. 限られた表現力が、外から見ると奇妙なバグを生む。正しいスキップトライグラムを学んだヘッドは、ソースと出力が交差した組み合わせ(ハイライト)の確率も一緒に上げてしまう。
完全に理解したのか
論文はこの問いに慎重に答える。ある意味では、そうだ。経路展開(path expansion)によってすべてのパラメータが文脈の中で意味を持つようになり、アルゴリズム的な謎は残っていない。しかしこの理解は、巨大な線形回帰の重みを理解している、あるいは大きなデータベースへの問い合わせ方を知っている、と言うときの理解だ。語彙が5万なら、展開されたOV行列一つに約25億の項目がある。要約する方法をさらに作らない限り、人間が頭に収められる量ではない。論文の言い回しを借りれば、1層モデルが圧縮された中国語の部屋だったことを突き止めたら、今度は巨大なカードの山が残った、という状況だ。
技術的な留保もいくつか付く。相関した変数の上の線形モデルの常として、二つのヘッドが互いの役割を代替できるため、重みゼロが「やらない」という意味とは限らない。また、異なるクエリのQKスコアを比較する標準的な正規化方法もない。契約書を読める形にすることには成功したが、まだ読み切ってはいない。それが論文の自己評価だ。
今回の要点
- 1層モデルは、直接経路(自動補完)とヘッド経路(スキップトライグラム)の和として完全に展開される。
- スキップトライグラムはQK表のマス一つとOV表のマス一つがつながった規則で、代表格はコピー。コピーはOV表の正の固有値に要約される(強い証拠だが決定的証明ではない)。
- 規則が二枚の表の積だという構造そのものが、交差した組み合わせが漏れるバグを予測させる。
- アルゴリズム的な謎は消えたが、25億の項目を要約なしに頭へ収めることはできない。残っているのは要約の問題だ。
次回がこの論文のクライマックスだ。層が二つになるとヘッドがヘッドと合成され、誘導ヘッドが登場する。
原文: A Mathematical Framework for Transformer Circuits のOne-Layer Attention-Only Transformersセクション。本文の図はすべて原文からの引用。