Transformer回路のための数学的フレームワーク (2021)

Transformer回路のための数学的フレームワーク

目次

• モデルの単純化

• 高レベルアーキテクチャ

• 仮想重みと残差ストリームを通信チャネルとして

• アテンションヘッドは独立かつ加算的

• アテンションヘッドの情報移動として

• パス展開トリック

• アテンションヘッドの項をクエリ・キーと出力・値の回路に分割

• スキップトリグラムとしての解釈

• OV/QK行列の要約

• 1層モデルを「完全に理解」できているか？

• 3種類の合成

• ロジットのパス展開

• アテンションスコアのパス展開 QK回路

• 2層モデルの分析

• 誘導ヘッド

• 項の重要度分析

• 仮想アテンションヘッド

Transformer言語モデルは、GPT-3、LaMDA、Codex、Meena、Gopherなどのシステムで実用化が進んでいる新興技術です。しかし、これらのモデルがスケールアップするにつれて、そのオープンエンド性と高いキャパシティは、予期せぬ、時には有害な挙動の範囲を増大させます。大規模モデルがトレーニングされてから数年経っても、作成者もユーザーも、以前は知らなかったモデルの能力（問題のある挙動を含む）を発見することが日常的にあります。

これらの問題に対処する一つの道筋は、メカニズム的解釈可能性（mechanistic interpretability）であり、これはTransformerが行う詳細な計算を、プログラマーが複雑なバイナリを人間が読めるソースコードにリバースエンジニアリングしようとするのと同様に、リバースエンジニアリングしようとする試みです。これが可能になれば、現在の安全性の問題の説明、新たな問題の特定、そしてまだ構築されていない強力な将来のモデルの安全性問題の予測さえも、より体系的なアプローチで提供できる可能性があります。以前のプロジェクトであるDistill Circuitsスレッドは、ビジョンモデルのリバースエンジニアリングを試みましたが、Transformerや言語モデルに関しては、これまで同等のプロジェクトはありませんでした。

本稿では、Transformerのリバースエンジニアリングに向けた、最初の、非常に予備的な一歩を試みます。現代の言語モデルの信じられないほどの複雑さとサイズを考慮すると、最も単純なモデルから始めてそこから積み上げていくのが最も有益であることがわかりました。私たちの目標は、より大きく複雑なモデルに後で適用できる、単純なアルゴリズムパターン、モチーフ、またはフレームワークを発見することです。具体的には、本稿では、注意ブロックのみを持つ、2層以下のTransformerを研究します。これは、96層を持ち、注意ブロックとMLPブロックを交互に配置するGPT-3のような大規模な現代のTransformerとは対照的です。

Transformerの動作を、新しいが数学的に等価な方法で概念化することにより、これらの小規模モデルを理解し、それらが内部でどのように動作するかについてかなりの理解を得ることができることがわかりました。特に注目すべきは、「誘導ヘッド」と名付けた特定の注意ヘッドが、これらの小規模モデルにおけるインコンテキスト学習を説明できること、そしてこれらのヘッドは少なくとも2つの注意層を持つモデルでのみ発達することを発見しました。また、特定のデータ上で動作するこれらのヘッドの例もいくつか見ていきます。

この最初の論文では、私たちの洞察をより大きなモデルに適用しようとはしませんが、今後の論文では、Transformerを理解するための私たちの数学的フレームワークと誘導ヘッドの概念が、はるかに大きく現実的なモデルに対しても少なくとも部分的に関連し続けることを示します。ただし、そのようなモデルを完全にリバースエンジニアリングできるようになるには、まだ程遠い状況です。

---

結果の要約

リバースエンジニアリングの結果

Transformerのリバースエンジニアリングという課題を探求するために、いくつかのトイ（toy）の、注意機構のみのモデルをリバースエンジニアリングしました。その過程で、以下のことがわかりました。

• ゼロ層Transformerは、バイグラム（bigram）統計をモデル化します。バイグラムテーブルには、重みから直接アクセスできます。
• 1層の注意機構のみのTransformerは、バイグラムモデルと「スキップトリグラム」（「A… B C」のような形式のシーケンス）モデルのアンサンブルです。バイグラムおよびスキップトリグラムテーブルには、モデルを実行せずに重みから直接アクセスできます。これらのスキップトリグラムは驚くほど表現力豊かです。これには、非常に単純なインコンテキスト学習の一種の実装が含まれます。
• 2層の注意機構のみのTransformerは、注意ヘッドの合成を使用して、より複雑なアルゴリズムを実装できます。これらの合成アルゴリズムも、重みから直接検出できます。特に、2層モデルは注意ヘッドの合成を使用して「誘導ヘッド」を作成します。これは非常に一般的なインコンテキスト学習アルゴリズムです。誘導ヘッドについては、今後の論文でさらに詳しく説明します。
• 1層および2層の注意機構のみのTransformerは、インコンテキスト学習を実行するために非常に異なるアルゴリズムを使用します。2層の注意ヘッドは、質的に洗練された推論時間アルゴリズム（特に、誘導ヘッドと呼ぶ特別な種類の注意ヘッド）を使用してインコンテキスト学習を実行します。これは、より大きなモデルに関連する重要な遷移点となります。

コンセプト的なテイクアウェイ

Transformerアーキテクチャの多くの微妙な詳細は、InceptionV1 Circuitsとはかなり異なる方法でリバースエンジニアリングに取り組む必要があることがわかりました。これらのポイントは以下のセクションで詳しく説明しますが、ここでは簡単に要約します。また、適切なセクションに到達した際に、導入する多くの用語を拡張します。（明確にするために、これらのポイントのいずれかが必ずしも新しいものであると主張するつもりはありません。多くは、他の論文に暗黙的または明示的に存在します。）

• アテンションヘッドは独立した操作として理解でき、それぞれが出力結果を残差ストリームに加算します。アテンションヘッドは、計算効率のために、しばしば代替的な「連結して乗算する」形式で説明されますが、これは数学的に同等です。
• 注意機構のみのモデルは、トークンからロジットへの変化を解釈可能なエンドツーエンド関数として記述できます。これらの関数はモデル内の「パス」に対応し、アテンションパターンを固定すると線形になります。
• Transformerには膨大な量の線形構造があります。行列の連鎖を分解して乗算するだけで、多くのことを学ぶことができます。
• アテンションヘッドは、2つのほぼ独立した計算を持つと理解できます。1つはアテンションパターンを計算するQK（クエリ・キー）回路、もう1つは、アテンションされた場合に各トークンが出力にどのように影響するかを計算するOV（出力・値）回路です。
• キー、クエリ、値ベクトルは、低ランク行列 W_Q^T W_K および W_O W_V の計算における中間結果と見なすことができます。それらを参照せずにTransformerを記述すると役立つ場合があります。
• アテンションヘッドの合成は、Transformerの表現力を大幅に向上させます。アテンションヘッドが合成できる方法は3つあり、それぞれキー、クエリ、値に対応します。キーとクエリの合成は、値の合成とは大きく異なります。
• Transformerのすべてのコンポーネント（トークン埋め込み、アテンションヘッド、MLP層、アンエンベディング）は、残差ストリームの異なるサブスペースを読み書きすることで互いに通信します。残差ストリームベクトルを直接分析するのではなく、モデル内のパスに対応する、残差ストリームで発生するさまざまな通信ストリームを分離して分析すると役立つ場合があります。

---

Transformerの概要

Transformerをリバースエンジニアリングしようとする前に、Transformerの高レベルな構造を簡単にレビューし、私たちがどのように考えているかを説明すると役立ちます。

多くの場合、Transformerを等価だが非標準的な方法で再構成することが有益であることがわかりました。メカニズム的解釈可能性は、モデルを人間が理解できる部分に分解することを必要とします。重要な最初のステップは、モデルについて推論するのに最も簡単な表現を見つけることです。現代のディープラーニングでは、正当な理由から、計算効率に多くの重点が置かれており、モデルの数学的記述は、モデルを実行するための効率的なコードの書き方に関する決定をしばしば反映しています。しかし、同じ計算を表す多くの等価な方法がある場合、最も人間が理解しやすい表現と最も計算効率の良い表現は異なる可能性が高いです。

Transformerをレビューすることで、用語を統一することもできます。また、その過程でいくつかの記法を導入しますが、これらの記法は多くのセクションで使用されるため、読者が簡潔に参照できるように、すべての記法の詳細な説明を記法付録で提供します。

モデルの単純化

この論文のアイデアを最もクリーンな形で示すために、「トイTransformer」にいくつかの単純化を加えて焦点を当てます。

この論文のほとんどの部分で、非常に実質的な変更を加えます。MLP層を持たない「注意機構のみ」のTransformerに焦点を当てます。これはTransformerアーキテクチャの非常に劇的な単純化です。その理由の一部は、注意ヘッドを持つ回路がDistill Circuitsの作業では直面しなかった新しい課題を提示しており、それらを分離して考慮することで、それらの問題に対する特にエレガントな扱いが可能になることです。しかし、MLP層の理解に関しては、これまで成功がはるかに少なかったということもあります。注意機構とMLP層の両方を持つ通常のTransformerでは、研究できる注意ヘッドによって主に仲介される多くの回路があり、それらの一部は非常に重要であるように見えますが、MLP部分はより困難でした。これは私たちの仕事の大きな弱点であり、将来的に対処することを目指しています。それにもかかわらず、後のセクションではMLP層を持つTransformerについてもいくつか議論します。

また、より表面的であると考えられ、主に明瞭さと単純さのために行われたいくつかの変更も行います。バイアスは考慮しませんが、バイアスを持つモデルは、それらを重みに折り込み、常に1である次元を作成することで、常にバイアスなしでシミュレートできます。さらに、注意機構のみのTransformerにおけるバイアスは、主にロジットに対するバイアスとして機能します。レイヤー正規化も無視します。明示的に考慮するとかなりの複雑さが加わりますが、変数スケーリングをモジュロとして、レイヤー正規化は隣接する重みにマージできます。また、実装上のいくつかの煩わしさ（annoyances）をモジュロとして、レイヤー正規化はバッチ正規化（これは完全に隣接するパラメータに折り込むことができます）に置き換えることができると予想しています。

高レベルアーキテクチャ

Transformer言語モデルにはいくつかのバリエーションがあります。GPT-3のような自己回帰型のデコーダーオンリーTransformer言語モデルに焦点を当てます。（オリジナルのTransformer論文は、翻訳をサポートするために特別なエンコーダー・デコーダー構造を持っていましたが、多くの現代の言語モデルにはこれらは含まれていません。）

Transformerは、トークン埋め込みから始まり、一連の「残差ブロック」、そして最後にトークンアンエンベディングが続きます。各残差ブロックは、アテンション層とMLP層で構成されます。アテンション層とMLP層はどちらも、残差ストリームから入力を「読み込み」（線形射影を実行することにより）、結果を「書き込み」（線形射影を加算することにより）ます。各アテンション層は複数のヘッドで構成され、これらは並列に動作します。

［画像: Transformerの高レベルアーキテクチャの図］

仮想重みと残差ストリームを通信チャネルとして

Transformerの高レベルアーキテクチャの主な特徴の1つは、各層がその結果を「残差ストリーム」と呼ぶものに加算することです。残差ストリームを持つモデルの構築は、Schmidhuberグループの初期の作業、例えばハイウェイネットワークやLSTMに遡ります。これらは、より最近の残差ネットワークアーキテクチャで大きな現代的な成功を収めています。Transformerでは、残差ストリームベクトルはしばしば「埋め込み」と呼ばれます。私たちは、残差ストリームという用語を好みます。なぜなら、それが残差的な性質（重要だと信じている）を強調するだけでなく、残差ストリームが現在のトークン以外のトークンに専用のサブスペースをしばしば割り当て、埋め込みという用語が示唆する直感を破るからです。残差ストリームは、単にすべての前の層の出力と元の埋め込みの合計です。私たちは一般的に、残差ストリームを通信チャネルと見なします。なぜなら、それはそれ自体では何も処理せず、すべての層がそれを介して通信するからです。

［画像: 残差ストリームの図］

残差ストリームは、非常に線形な構造を持っています。注意すべきは、完全に線形な残差ストリームは、ニューラルネットワークアーキテクチャの中でも非常に珍しいということです。広く使用されている最も類似したアーキテクチャであるResNetでさえ、残差ストリームに非線形活性化関数があるか、残差ストリームにアクセスするたびに適用されます！各層は、線形変換を実行して残差ストリームから情報を「読み込み」（これは各層の先頭にあるレイヤー正規化を無視していますが、定数スケーリングをモジュロとして、レイヤー正規化は定数アフィン変換であり、線形変換に折り込むことができます。レイヤー正規化の扱いについては付録の議論を参照してください）、そして「書き込み」の前に別の線形変換を実行して、その出力を残差ストリームに加算します。残差ストリームのこの線形かつ加算的な構造は、多くの重要な意味を持っています。1つの基本的な結果は、残差ストリームには「特権基底（privileged basis）」がないということです。モデルの動作を変えることなく、それと相互作用するすべての行列を回転させることで、それを回転させることができます。

仮想重み

残差ストリームが線形であることの特に有用な結果は、任意の2つの層（他の多くの層によって隔てられている層でさえ）を直接接続する暗黙の「仮想重み」を、残差ストリームを介したそれらの相互作用を掛け合わせることによって考えることができるということです。これらの仮想重みは、ある層の出力重みと別の層の入力重み（つまり、W_I^2 W_O^1）の積であり、後の層が前の層によって書き込まれた情報をどの程度読み込むかを記述します。注意層の場合、入力重みはW_Q、W_K、W_Vの3種類があることに注意してください。単純さと一般性のために、ここでは層は単に入力重みと出力重みを持つと考えることにします。

［画像: 仮想重みの図］

サブスペースと残差ストリーム帯域幅

残差ストリームは高次元ベクトル空間です。小規模なモデルでは数百次元かもしれませんが、大規模なモデルでは数万次元に達することもあります。これにより、層は情報を異なるサブスペースに格納することで、異なる層に異なる情報を送信できます。これは、個々のヘッドが比較的小さなサブスペース（しばしば64または128次元）で動作し、完全に分離したサブスペースに書き込んで相互作用しないことが非常に容易であるため、特にアテンションヘッドの場合に重要です。

一度追加された情報は、別の層がそれを積極的に削除しない限り、サブスペース内に永続します。この観点から、残差ストリームの次元は「メモリ」または「帯域幅」のようなものになります。元のトークン埋め込みとアンエンベディングは、主に次元の比較的小さな部分集合と相互作用します。トークン埋め込みとアンエンベディングのPCA分析を実行しました。d_modelが大きいモデルでは、スペクトルは急速に減衰し、埋め込み/アンエンベディングは全体的な次元の比較的小さな部分集合に集中していました。それらが同じサブスペースまたは異なるサブスペースを占めているかどうかを把握するために、正規化された埋め込み行列とアンエンベディング行列を連結し、PCAを適用しました。この共同PCAプロセスは、「混合」次元と一方のみで使用される次元の両方の組み合わせを示しました。一方のみで使用される次元の存在は、それらが同じサブスペースを使用する範囲の上限と見なすことができます。これにより、ほとんどの次元が他の層が情報を格納するために「自由」に残ります。残差ストリーム帯域幅の需要は非常に高いと予想されます！一般的に、「計算次元」（ニューロンやアテンションヘッドの結果次元など）は、情報移動のための次元よりも残差ストリームの方がはるかに多くなります。単一のMLP層でさえ、通常、残差ストリームの次元の4倍のニューロンがあります。したがって、例えば、50層Transformerの25層目では、残差ストリームは、それより前のニューロンの100倍の次元を持ち、それより後のニューロンの100倍のニューロンと通信しようとしており、重ね合わせで通信しています！このようなテンソルを「ボトルネックアクティベーション」と呼び、解釈が異常に困難であると予想されます。（これは、残差ストリームを直接研究するのではなく、残差ストリームを介したさまざまな通信ストリームを仮想重みの観点から分離しようとする主な理由の1つです。）

おそらくこの残差ストリーム帯域幅の高い需要のため、一部のMLPニューロンとアテンションヘッドは、一種の「メモリ管理」の役割を果たし、他の層によって設定された残差ストリーム次元を、情報を読み込んで負のバージョンを書き出すことでクリアしている兆候が見られます。一部のMLPニューロンは、入力と出力の重みの間に非常に負のコサイン類似度を持ち、これは残差ストリームからの情報の削除を示唆している可能性があります。同様に、一部のアテンションヘッドは、OV行列に大きな負の固有値を持ち、主に現在のトークンにアテンションを向け、情報を削除するメカニズムとして機能する可能性があります。これらは情報の「メモリ管理」削除の一般的なメカニズムであるかもしれませんが、条件付きで情報を削除するメカニズムでもあり、特定のケースでのみ動作することに注意する価値があります。

［画像: 残差ストリーム帯域幅の図］

アテンションヘッドは独立かつ加算的

上記で見たように、Transformerのアテンション層を、完全に並列に動作し、それぞれの出力を残差ストリームに加算する、完全に独立した複数のアテンションヘッド h ∈ H として考えています。しかし、これはTransformer層が一般的に提示される方法ではなく、それらが等価であることは明らかではないかもしれません。

Transformerに関するオリジナルのVaswani et al.の論文では、アテンション層の出力は、結果ベクトル r^{h_1}, r^{h_2},... をスタックし、出力行列 W_O^H で乗算することによって説明されています。W_O^H を各ヘッドの等しいサイズのブロックに分割しましょう（[W_O^{h_1}, W_O^{h_2}...]）。すると、次のようになります。

W_O^H \begin{bmatrix} r^{h_1} \ r^{h_2} \ ... \end{bmatrix} = \begin{bmatrix} W_O^{h_1},~ W_O^{h_2},~ ... \end{bmatrix] \cdot \begin{bmatrix} r^{h_1} \ r^{h_2} \ ... \end{bmatrix} = \sum_i W_O^{h_i} r^{h_i}

これは、ヘッドを独立に実行し、それぞれを独自の出力行列で乗算し、それらを残差ストリームに加算することと同等であることがわかります。連結（concatenate）の定義は、より大きく、より計算効率の良い行列乗算を生成するため、しばしば好まれます。しかし、Transformerを理論的に理解するためには、それらを独立に加算されるものとして考えることを好みます。

アテンションヘッドの情報移動として

しかし、アテンションヘッドが独立して動作する場合、それらは何をするのでしょうか？アテンションヘッドの基本的な動作は、情報の移動です。それらは、あるトークンの残差ストリームから情報を読み取り、別のトークンの残差ストリームに書き込みます。このセクションから得られる主な洞察は、どのトークンから情報を移動するかは、移動される「読み取られた」情報や、どのように「書き込まれた」かとは完全に分離可能であるということです。

［画像: アテンションヘッドの情報移動の図］

これを理解するには、アテンションを非標準的な方法で記述すると役立ちます。アテンションパターンが与えられた場合、アテンションヘッドの出力を計算することは、通常3つのステップで説明されます。

1. 各トークンの値ベクトルを残差ストリームから計算します（v_i = W_V x_i）。
2. アテンションパターンに従って値ベクトルを線形に結合して「結果ベクトル」を計算します（r_i = \sum_j A_{i,j} v_j）。
3. 最後に、各トークンのヘッドの出力ベクトルを計算します（h(x)_i = W_O r_i）。上記で議論したように、出力行列による乗算は、しばしばすべてのヘッドの結果を連結したものに適用される単一の行列乗算として記述されますが、このバージョンは同等です。

これらの各ステップは行列乗算として記述できます。なぜそれらを単一のステップにまとめるのでしょうか？xを2次元行列（各トークンのベクトルで構成される）と考えると、異なる側で乗算しています。W_VとW_Oは「トークンあたりのベクトル」側に乗算し、Aは「位置」側に乗算します。テンソルは、この種の行列から行列へのマッピングを記述するための、より自然な言語を提供してくれます（テンソル積記法に慣れていない場合は、記法付録に短い紹介を含めました）。役立つ可能性のある動機付けの1つは、行列から行列への線形マップを表現したいということです：[n_context, d_model] ~→~ [n_context, d_model]。数学者はそのような線形マップを「(2,2)-テンソル」（2つの入力次元を2つの出力次元にマップする）と呼びます。したがって、テンソルは、この変換を表現するための自然な言語です。

テンソル積を使用して、アテンションの適用プロセスを次のように記述できます。

h(x) ~=~ (Id ⊗ W_O)·

各トークの結果ベクトルを射影する

(h(x)_i = W_O r_i)

~ (A ⊗ Id)~~·~~~

結果ベクトルを計算するために値ベクトルをトークン間で混合する

(r_i = \sum_j A_{i,j} v_j)

~ (Id ⊗ W_V)~~·~~~

各トークの値ベクトルを計算する

(v_i=W_V x_i)~~

x

混合積の性質を適用し、恒等式を折りたたむと、次のようになります。

h(x) ~=~ (A ⊗ W_O W_V) ~~·~~~~~~

A はトークン間で混合し、

W_OW_V は各ベクトルに独立に作用する。

x

アテンションパターンについてはどうでしょうか？通常、キー k_i = W_K x_i を計算し、クエリ q_i = W_Q x_i を計算し、次にキーとクエリベクトルのドット積 A = softmax(q^T k) からアテンションパターンを計算します。しかし、キーとクエリを参照せずに、すべてを1つのステップで行うことができます：A = softmax(x^T W_Q^T W_K x)。

この定式化は数学的に同等ですが、実際にはこの方法でアテンションを実装すること（つまり、W_O W_V および W_Q^T W_K で乗算すること）は、非常に非効率的になることに注意してください！

アテンションヘッドに関する観察

アテンションヘッドをこの形式で書き直すことの主な利点は、以前は観察が難しかった多くの構造を表面化させることです。

• アテンションヘッドは、あるトークンから別のトークンへ情報を移動させます。

• これの帰結として、残差ストリームベクトル空間（しばしば「文脈化された単語埋め込み」として解釈される）は、現在のトークンに直接関係するのではなく、他のトークンからコピーされた情報に対応する線形サブスペースを持つことが一般的です。

• アテンションヘッドは、実際には2つの線形操作、AとW_OW_Vを適用しており、これらは異なる次元で動作し、独立に作用します。

• A は、どのトークンの情報が移動元および移動先となるかを決定します。

• W_O W_V は、どの情報がソーストークンから読み取られ、どのように宛先トークンに書き込まれるかを決定します。W_{OV}=W_O W_V が、アテンションヘッドが情報を移動する際に、残差ストリームのどのサブスペースを読み取り、書き込むかを決定すると言った場合、それはどういう意味でしょうか？特異値分解 USV = W_{OV} を考慮すると役立ちます。d_head < d_model であるため、W_{OV} は低ランクであり、S の対角成分のサブセットのみがゼロではありません。右特異ベクトル V は、アテンションされている残差ストリームのどのサブスペースが「読み込まれる」（値ベクトルとして何らかの方法で格納される）かを記述し、左特異ベクトル U は、宛先残差ストリームのどのサブスペースに書き込まれるかを記述します。

• A はこの方程式の唯一の非線形部分です（softmaxから計算されるため）。これは、アテンションパターンを固定した場合、アテンションヘッドは線形操作を実行することを意味します。これはまた、A を固定しない場合、アテンションヘッドは、トークンごとの線形操作が一定であるため、ある意味で「半線形」であることを意味します。

• W_Q と W_K は常に一緒に動作します。それらは決して独立ではありません。同様に、W_O と W_V も常に一緒に動作します。

• それらは別々の行列としてパラメータ化されていますが、W_O W_V と W_Q^T W_K は常に個別の低ランク行列と見なすことができます。

• キー、クエリ、値ベクトルは、ある意味で表面的です。それらは、これらの低ランク行列を計算する際の中間的な副産物です。異なるベクトルを作成するために、低ランク行列の各因数を容易に再パラメータ化できますが、機能は同じままです。

• W_O W_V と W_Q W_K は常に一緒に動作するため、これらの結合行列を表す変数、W_{OV} = W_O W_V および W_{QK} = W_Q^T W_K を定義することを好みます。

• アテンションヘッドの積は、アテンションヘッド自体と非常によく似た振る舞いをします。分配法則により、(A^{h_2}⊗ W_{OV}^{h_2}) · (A^{h_1}⊗ W_{OV}^{h_1}) = (A^{h_2}A^{h_1})⊗(W_{OV}^{h_2}W_{OV}^{h_1}) となります。この積の結果は、2つのヘッドの合成であるアテンションパターン A^{h_2}A^{h_1} と出力値行列 W_{OV}^{h_2}W_{OV}^{h_1} を持つ、機能的にアテンションヘッドに等価なものと見なすことができます。これらを「仮想アテンションヘッド」と呼び、後で詳しく説明します。

---

ゼロ層Transformer

このセクションと同様の内容をカバーするビデオをご覧ください：0層理論

より複雑なモデルに進む前に、「ゼロ層」Transformerを簡単に検討することは役立ちます。このようなモデルは、トークンを取得し、それを埋め込み、次に次のトークンを予測するロジットを生成するためにアンエンベッドします。

T ~=~ W_U W_E

モデルは他のトークンから情報を移動できないため、現在のトークンから次のトークンを予測するだけです。これは、W_U W_E の最適な動作が、バイグラムの対数尤度を近似することであることを意味します。これは、Levy & Goldberg, 2014 による、多くの初期の単語埋め込みが対数尤度行列の行列分解と見なせるという観察と並行しています。

これは、Transformer全般に関連しています。W_U W_E の形式の項は、すべてのTransformerの展開された方程式の形式で出現します。これは、トークン埋め込みが残差ストリームを直接アンエンベディングに流れる、「直接パス」に対応しており、どの層も通過しません。それが影響できる唯一のことは、バイグラムの対数尤度です。他のモデルの側面がバイグラムの対数尤度の一部を予測するため、より大きなモデルでは正確にバイグラム統計を表すわけではありませんが、一種の「残差」を表します。特に、W_U W_E 項は、しばしば「Barack」の後に「Obama」が続くという事実のような、より一般的な文法規則では説明されないバイグラム統計を表すのに役立つようです。これの興味深い帰結は、W_U がしばしば「アンエンベディング」行列と呼ばれているにもかかわらず、W_E との逆行列であると期待すべきではないということです。

---

1層注意機構のみのTransformer

このセクションと同様の内容をカバーするビデオをご覧ください：1層理論、1層結果。

1層注意機構のみのTransformerは、バイグラムモデルといくつかの「スキップトリグラム」モデル（「A… BC」というシーケンスの確率に影響を与える）のアンサンブルとして理解できると主張します。Mikolov et al.の単語埋め込みに関する古典的な論文での「スキップグラム」という用語の使用に触発されて、「A… BC」という形式のシーケンスを指す「スキップトリグラム」という用語を使用しています。直感的には、各アテンションヘッドは、現在のトークン（「B」）から前のトークン（「A」）に選択的にアテンションを向け、情報をコピーして、可能な次のトークン（「C」）の確率を調整できるためです。

このセクションの目標は、この対応関係を厳密に示し、Transformerの生の重みを解釈可能なスキップトリグラム確率調整テーブルに変換する方法を示すことです。

パス展開トリック

1層注意機構のみのTransformerは、トークン埋め込み、アテンション層（独立してアテンションヘッドを適用する）、そして最後にアンエンベディングで構成されていることを思い出してください。

［画像: 1層Transformerのアーキテクチャ図］

テンソル記法と、以前に導出したアテンションヘッドの代替表現を使用して、Transformerを3つの項の積として表現できます。

［画像: 1層Transformerの数式展開図］

私たちの主要なトリックは、単に積を展開することです。これにより、積（各項が層に対応する）が、和（各項がエンドツーエンドパスに対応する）に変換されます。

［画像: パス展開の数式図］

これらのエンドツーエンドパス項のそれぞれは、理解可能であり、独立に推論でき、モデルの挙動を生成するために加算的に結合されると主張します。

ゼロ層Transformerを見たときにも現れた、直接パス項、Id ⊗ W_U W_E もあります。これは位置間で情報を移動しないため（Id ⊗ ... が示すもの）、バイグラム統計にしか貢献できず、他の項が処理しないギャップを埋めます。

より興味深いのはアテンションヘッドの項です。

アテンションヘッドの項をクエリ・キーと出力・値の回路に分割

各アテンションヘッド h について、項 A^h ⊗ (W_U W_{OV}^h W_E) があります。ここで A^h = softmax( t^T · W_E^T W_{QK}^h W_E · t ) です。これらの項をモデルの挙動にどのようにマッピングできるでしょうか？そして、そのついでに、なぜ私たちの数式にこれらの特定の行列の積が現れるのでしょうか？

注意すべき重要な点は、これらの項が2つの分離可能な操作で構成されており、それらの本質は2つの [n_vocab, n_vocab] 行列であるということです。

• W_E^T W_{QK}^h W_E — この行列を「クエリ・キー（QK）回路」と呼びます。これは、すべてのクエリおよびキー（key）トークンのアテンションスコアを提供します。つまり、各エントリは、特定のクエリトークンが特定のキー（key）トークンにどれだけアテンションを向けたいかを記述します。
• W_U W_{OV}^h W_E — この行列を「出力・値（OV）回路」と呼びます。これは、アテンションされた場合に、特定のトークンが出力ロジットにどのように影響するかを記述します。

これらの積を直感的に理解するには、トークンから始まりトークンで終わるモデル内のパスとして考えることが役立ちます。QK回路は、クエリベクトルとキーベクトルをアテンションヘッドまでトレースし、そこでそれらがドット積されて双線形形式が作成される計算をたどることで形成されます。OV回路は、値ベクトルを計算するパスをたどり、それをロジットまで継続することで作成されます。

［画像: QKおよびOV回路の図］

アテンションパターンは、ソーストークンと宛先トークンの両方の関数です。技術的には、それは開始から宛先トークンまでのすべての可能なソーストークンの関数です。なぜなら、softmaxはQK回路を介して各トークのスコアを計算し、指数化してから正規化するからです。しかし、宛先トークンがソーストークンにどれだけアテンションを向けるかを決定すると、出力への影響は、そのソーストークンのみの関数となります。つまり、複数の宛先トークンが同じソーストークンに同じ量のアテンションを向ける場合、ソーストークンは予測される出力トークンへの影響において同じになります。

OVとQKの独立性（アテンションパターン固定トリック）

OV回路とQK回路を別々に考えることは非常に役立ちます。なぜなら、それらはどちらも個別に理解可能な関数（理解している行列に作用する線形または双線形関数）だからです。

しかし、それらを独立して考えることは本当に原理的でしょうか？役立つ思考実験の1つは、モデルを2回実行することを想像することです。1回目は、各ヘッドのアテンションパターンを収集します。これはQK回路のみに依存します。より多くの層を持つモデルでは、QK回路がW_E^T W_{QK}^h W_Eよりも複雑になることがわかります。2回目は、アテンションパターンを、1回目に収集した「固定された」アテンションパターンに置き換えます。これにより、ロジットがトークンの線形関数となる関数が得られます！私たちはこれをTransformerを考える上で非常に強力な方法だと考えています。

スキップトリグラムとしての解釈

メカニズム的解釈可能性の主要な課題の1つは、ニューラルネットワークのパラメータを文脈化することによって意味のあるものにすることです（Voss et al.による重みの可視化での議論を参照）。OVとQK回路を掛け合わせることで、これを達成しました。ニューラルネットワークのパラメータは、トークン上の単純な線形または双線形関数になりました。QK回路は、現在の「宛先」トークンがどの「ソース」トークンに注意を向け、情報をコピーするかを決定します。一方、OV回路は、結果として「出力」予測にどのような影響があるかを記述します。これらを合わせると、関与する3つのトークンは [source]... [destination][out] の形式の「スキップトリグラム」を形成し、「out」が変更されます。

これは解釈が些細なことを意味するわけではないことに注意することが重要です。1つには、結果の行列は巨大です（私たちの語彙は約50,000トークンなので、単一の展開されたOV行列は約25億のエントリを持ちます）。1層注意機構のみのモデルを圧縮された中国語の部屋（Chinese room）として明らかにし、巨大なカードの山が残りました。また、相関変数に作用する一般化線形モデルの重みを理解する際に通常伴う問題や、変数間の代替可能性（fungibility）に関する問題もすべてあります。例えば、アテンションヘッドがゼロの重みを持つのは、別のアテンションヘッドが同じトークンにアテンションを向け、本来ならそのヘッドが行うはずだった役割を果たすためかもしれません。最後に、QK重みは異なるクエリベクトル間で比較できず、正規化方法についての明確な正解がないという技術的な問題があります。

それにもかかわらず、すべてのパラメータが文脈化され、理解可能になった形式のTransformerを持っています。そして、これらの微妙な点にもかかわらず、単にOVおよびQK行列からスキップトリグラムを読み取ることができます。特に、これらの行列の大きなエントリを検索すると、多くの興味深い挙動が明らかになります。

次のサブセクションでは、いくつかの興味深いスキップトリグラムと、それらがQK/OV回路にどのように埋め込まれているかのキュレーションされたツアーを提供します。しかし、最大の大きさを誇るエントリの完全で、厳選されていない例は、リンクをたどることで利用できます。

• 大きなQK/OVエントリ — 12ヘッド、d_head=64
• 大きなQK/OVエントリ — 32ヘッド、d_head=128

コピー/原始的なインコンテキスト学習

これらの行列を調べる上で最も印象的なことの1つは、1層モデルのアテンションヘッドの大部分が、その容量の大部分をコピーに費やしていることです。OV回路は、ヘッドによってアテンションされたトークンが、そのトークン、およびそれよりわずかに類似したトークンの確率を増加させるように設定されています。QK回路は、次に続くトークンとしてもっともらしいものにのみアテンションを向けます。したがって、トークンはコピーされますが、バイグラム的な統計がそれらを尤もらしく見せる場所にのみコピーされます。