Transformer 回路のための数学的枠組み

目次

• モデルの簡略化
  • 高レベルなアーキテクチャ
  • 仮想的な重みと通信チャンネルとしての残留ストリーム
  • アテンションヘッドの独立性と加算的動作
• パス展開のトリック
  • Query-Key と Output-Value 回路への分割
  • スキップ・トリグラム（Skip-trigrams）としての解釈
  • コピー動作と「バグ」の検出
• 二層 Transformer の解析
  • 三種類の結合（Composition）
  • 誘導ヘッド（Induction Heads）
  • 項の重要性分析と仮想的なアテンションヘッド
• 回路と解釈可能性への影響
  • 数学的枠組みの独自性
  • アーキテクチャ変更との関連

---

モデルの簡略化

Transformer のリバースエンジニアリング（メカニスティックな解釈可能性）を容易にするために、本研究では大規模モデルから段階的に簡略化されたモデルに焦点を当てています。目標は、複雑なモデルに適用可能な単純なアルゴリズムパターンや**動詞（Motifs）**を見つけることです。具体的には、96 レイヤーを持つ GPT-3 などに比べて、2 レイヤー以下かつアテンションブロックのみを持つ玩具モデルを分析しました。

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

Transformer は、GPT-3 や LaMDA などの大規模モデルにおいて未定義性の増大や予期せぬ振る舞いを引き起こす原因となっています。メカニスティックな解釈可能性は、複雑なバイナリを人間の読みやすいソースコードとしてリバースエンジニアリングすることを試みます。 本研究では以下のステップを行いました：

• 単純化: 可能な限り単純なモデル（アテンションブロックのみ）から始め、そこから段階的に進展。
• 発見: 2 レイヤー以下のモデルにおいて、コンテキスト内学習（In-context learning）を説明する特定のアテンションヘッド**「誘導ヘッド（Induction Heads）」**を発見。
• 拡張性: この数学的枠組みと誘導ヘッドの概念は、より大きな現実的なモデルでも少なくとも部分的に有効であることが示されています。

仮想的な重みと通信チャンネルとしての残留ストリーム

Transformer のリバースエンジニアリングにおいて以下の重要な発見を行いました：

• レイヤーゼロ: ビッグラム（bigram）統計をモデル化します。ビッグラムテーブルは重みから直接アクセス可能です。
• 一層モデル: ビッグラムモデルと「スキップ・トリグラム（skip-trigram）」モデルのアンサンブルです。
  • 例：形としては

    "A… B C"

    のような配列をモデル化します。
  • これらは重みから直接アクセス可能であり、驚くほど表現力が高く、単純なコンテキスト内学習を実装できます。
• 二層モデル: アテンションヘッドの結合（Composition）によりはるかに複雑なアルゴリズム（特に「誘導ヘッド」）を記述可能です。

概念的な要点

• 独立性と加算性: アテンションヘッドは独立した操作であり、それぞれが結果を計算して**残留ストリーム（Residual Stream）**に加算します。
• パス展開: トークンをロジットの変化に対応するエンドツーエンド関数の和として記述できます。アテンションパターンを凍結（freeze）すれば線形となります。
• QK/OV の分離: アテンションヘッドには二つの独立した計算があります。
  • QK 回路: アテンションパターン（どのトークンを注目するか）を計算。
  • OV 回路: 注目された場合の出力への影響を計算。
• 通信チャンネル: トークン埋め込み、アテンションヘッド、MLP、アンエンベディングはすべて、残留ストリームの異なる部分空間を読み書きすることで通信します。

アテンションヘッドの独立性と加算的動作

理論的理解のためには、連結して掛ける代替的な定式化よりも、独立して加算的であるように考えるのが好ましいです。

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

アテンションヘッドは情報を移動させる役割を持ちます：

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$)。

このプロセスは、行列乗算として記述できますが、テンソル表記法を用いるとより自然に表現されます： $$ h(x) = (\text{Id} \otimes W_O)\cdot (A \otimes \text{Id})\cdot ~(\text{Id} \otimes W_V)\cdot x $$

アテンションヘッドについての観察事項：

• 情報の移動: 一つのトークンから別のトークンへの情報移動を制御します。
• 部分空間の制御:
  • $A$（アテンションパターン）: どのトークンの情報が移動されるか（从何）、何に移動するか（何地）を制御。
  • $W_{OV}$: ソーストークンからどの情報を読み、どのように宛先に書き込むかを制御。
• 線形性: アテンションパターン $A$ を固定すれば、アテンションヘッドは線形操作として動作します（「半線形」とも呼ばれます）。
• 結合行列:
  • $W_{OV} = W_O W_V$
  • $W_{QK} = W_Q^T W_K$
  • これらの結合は常に低ランク行列として扱われます。
• 仮想的なアテンションヘッド: 異なるレイヤーのヘッドが結合（積）することで、実質的に新しいアテンションヘッドのような振る舞いをします。

---

パス展開のトリック

Transformer をリバースエンジニアリングする上で重要なテクニックはパス展開です。モデルを構成する各項を展開し、これらがエンドツーエンドのパスに対応することを示します。これにより、ニューラルネットワークのパラメータ（重み）が意味づけることが可能になります。

Query-Key と Output-Value 回路への分割

一層のアテンションのみ Transformer を三つの項の積として表現できます： $$ \text{Id} \otimes W_U W_E + \sum_{h} A^h \otimes (W_U W_{OV}^h W_E) $$

各項をパスに対応する和に変換し、以下の回路に分けます：

1. Query-Key (QK) 回路: $W_E^T W_{QK}^h W_E$
   • 現在の「宛先」トークンがどの「ソース」トークンに注目するかを決定します（アテンションスコア）。
2. Output-Value (OV) 回路: $W_U W_{OV}^h W_E$
   • 与えられたトークンが注目された場合、出力ロジットへの影響を記述します。

OV と QK の独立性（「凍結」アテンションパターン）

モデルを実行してアテンションパターンを収集し、それを凍結させた上で再度実行することで、ロジットがトークンの線形関数であることが示せます。この思考実験により、QK 回路と OV 回路の構造が明確になります。

スキップ・トリグラム（Skip-trigrams）としての解釈

OV と QK 回路を展開すると、パラメータは単純な線形または双線形関数になります：

• QK 回路: 「宛先」トークンと「ソース」トークンのペアリングを決定。
• OV 回路: 次のトークンの出力予測への効果。
• 全体構造:

  [ソース] … [宛先][アウト]

  という形式のスキップ・トリグラムを形成します。

コピー動作（コピー/Primitive In-context learning）

一層モデルの多くのアテンションヘッドは、以下の動作を行います：

• ビッグラム統計が妥当だと判断する場所でのみ情報をコピー。
• 例：

  lambda… \lambda

  ,

  nbsp… > 

  , トークナイズの特殊性に関連するコピー。

「バグ」の検出

モデルが理想的には確率を増加させないスキップトリグラム（例：

QCanvas

の文脈での不整合）について洞察を与えます。これは解釈可能性を用いてモデルの失敗を理解する最初のステップです。

OV/QK 行列の総括

拡張された行列は巨大（数十億エントリ）ですが、以下の理由から理解可能です：

1. 低ランク性: ランクは $d_{head}$（64 または 128）のみ。
2. クラスター構造: 名前や場所に関連する明確なパターンが見られる。
3. コピー動作の検出:
   • 固有値分解: コピー行列は正の実数固有値を持つ傾向がある。
   • 対角線: トークン自身の確率増加に対応。

---

二層 Transformer の解析

深層学習では、**組合せ（Composition）**によって表現力が創出されます。一層モデルが「テーブル」として機能するのに対し、二層モデルは「プログラム」のように動作します。

三種類の結合（Composition）

残留ストリームを通信チャンネルとして考えます。アテンションヘッドが結合する際、以下の三つのオプションがあります：

• Q-結合: 以前のアテンションヘッドの影響を受けた部分空間を読み込み（$W_Q$）。
• K-結合: 以前のアテンションヘッドの影響を受けた部分空間を読み込み（$W_K$）。
• V-結合: 以前のアテンションヘッドの影響を受けた部分空間を読み込み（$W_V$）。

Q と K の結合はアテンションパターンに影響し、複雑なパターンを表現することを可能にします。一方、V-結合は単一の単位として動作し、追加の「仮想的なアテンションヘッド」を作成します。

ロジットのパス展開

モデル内の各レイヤーに対応する項を持ち、エンドツーエンドのパスに対応する項を持つ和に展開します： $$ \text{Total} = \text{Direct Path} + \sum \text{Attention Paths} $$

二層モデルでは、以下の重要な発見があります：

• 仮想的なアテンションヘッド: V-結合により形成されます（後述）。
• QK 回路の変化: Q-結合と K-結合の影響を受け、より表現力のある第二層のアテンションパターンが生成されます。

誘導ヘッド（Induction Heads）

小さな二層モデルにおいて、組合せは主に**「誘導ヘッド」**の作成のために使用されます。

誘導ヘッドの機能

• 動作: 現在のトークンの以前の例を検索し、見つければコピーして次のトークンを予測します。
  • 一層モデル:

    [b] … [a] → [b]

    (位置依存)
  • 二層モデル:

    [a][b] … [a] → [b]

    (内容依存)
• 利点: 分布シフトに対して耐性が高く、学習された統計に依存しないため、より確実な予測が可能です。

誘導ヘッドはどのように機能するか

• K-結合の役割: キーベクトルを一つトークン分前方シフトします（$\text{Id} \otimes A^{h_{-1}} \otimes W$）。
  • これにより、現在のクエリが「類似した」キー（シフトされた位置）を見つけることで、次のトークンを発見できます。
• メカニズム: 「同一マッチング」QK 回路行列 ($\text{Id} \otimes A^{h_{-1}} \otimes W$) と「コピー」OV 回路行列を持つことで機能します。

項の重要性分析

• 仮想的なアテンションヘッド: V-組成（二次の項）はこのモデルでは比較的少ない限界効果を持ちます。重要なのは、直接パス（ビッグラム）と個別のアテンションヘッド項です。
• K-結合の重要性: 誘導ヘッドは K-結合を通じて機能するため、これが重要な役割を果たします。

仮想的なアテンションヘッド

V-結合の結果として形成されます： $$ h_2 \circ h_1: (A^{h_2}A^{h_1}) \otimes (W_{OV}^{h_2}W_{OV}^{h_1}) $$

• 特徴: 独自のアテンションパターンと OV 行列を持ち、実質的に新しいヘッドのように動作します。
• 重要性: より深いモデルでは高次の仮想ヘッド（$h_3 \circ h_2 \circ h_1$ など）が可能ですが、二層モデルでは V-結合の影響は限定的です。

---

回路と解釈可能性への影響

本研究の成果は、大規模な Transformer の一部を理解するための足場となります。

数学的枠組みの独自性

本研究の貢献は以下の点にあります：

• InceptionV1 Circuitsからの発展：線形構造を持つ残留ストリームを通信チャンネルとして扱います。
• 双線形形式: アテンションヘッドは通常、低ランクな結合行列 ($W_{QK}, W_{OV}$) を持つ点でユニークです。
• パス展開: モデルを構成する項を「パス」として解釈し、各パスの役割（直接パス、スキップ・トリグラム、誘導ヘッド）を特定します。

アテンションを説明としての批判

従来のアテンション重み（$A_{ij}$）を「どのトークンが重要か」の説明として単純に解釈することは誤謬を含みます：

• Naive な解釈の問題: 高次の項（QK/OV の相互作用）を無視しており、実際のモデル動作（誘導ヘッドなど）を説明できない。
• 例: 誘導ヘッドのアテンションパターン自体は情報豊かですが、その挙動は K-結合によって完全に決定され、Naive な解釈では誤解される可能性があります。

アーキテクチャ変更との関連

近年のアーキテクチャ改良に対する洞察：

• Primer (Convolutional Attention): 深さ方向畳み込みをすることで、K-結合なしに誘導ヘッドのような構造を実現可能です。
• Talking Heads: $W_{OV} = \sum \alpha_i W_O^i W_V^i$ のように複数のヘッドの行列が共有される場合、コピー動作などの構造は自然に説明されます。

まとめ：我々はどこに置かれるのか？

• 達成: 一層・二層のアテンションのみ Transformer を完全に理解し、「誘導ヘッド」などの具体アルゴリズムを発見しました。
• 限界: 大規模モデルの約2/3 が MLP レイヤーであり、これらを考慮しないと完全な理解は得られません。また、MLP とアテンションが相互作用するため、より深い解析が必要です。
• 展望: この枠組みを用いて、大規模言語モデルにおけるコンテキスト内学習の中央ドライバー（誘導ヘッド）や、他の回路を解明していくことが期待されます。