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

Transformer回路のための数学的フレームワーク 著者 Nelson Elhage∗†, Neel Nanda∗, Catherine Olsson∗, Tom Henighan†, Nicholas Joseph†, Ben Mann†, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, Chris Olah‡ 所属 Anthropic 公開日 2021年12月22日 ∗ コア研究貢献者; † コアインフラストラクチャ貢献者; ‡ 連絡先 colah@anthropic.com; 著者貢献声明は以下にあります。 Transformer言語モデルは、GPT-3, LaMDA, Codex, Meena, Gopherなどのシステムで実世界での利用が広がりつつある新興技術です。しかし、これらのモデルがスケールアップするにつれて、そのオープンエンド性と高い能力は、予期せぬ、時には有害な動作の範囲を増大させます。大規模モデルがトレーニングされてから数年後でも、作成者とユーザーは、以前は気づかなかったモデルの能力（問題のある動作を含む）を日常的に発見しています。これらの問題に対処する一つの方法は、メカニスティック解釈可能性であり、プログラマーが複雑なバイナリを人間が読めるソースコードに逆解析しようとするのと同様に、Transformerによって実行される詳細な計算を逆解析しようと試みることです。これが可能になれば、現在の安全性の問題の説明、新しい問題の特定、そしておそらくはまだ構築されていない強力な将来のモデルの安全性問題の予測に対して、より体系的なアプローチを提供できる可能性があります。以前のプロジェクトであるDistill Circuits threadは、ビジョンモデルの逆解析を試みましたが、これまでのところTransformerや言語モデルに対する同等のプロジェクトはありませんでした。本論文では、Transformerの逆解析に向けた、非常に予備的な初期段階を踏み出そうと試みます。現代の言語モデルの信じられないほどの複雑さとサイズを考慮すると、最も有益なのは、可能な限り単純なモデルから始めて、そこから進んでいくことだとわかりました。私たちの目標は、より大きく複雑なモデルに適用できる単純なアルゴリズムパターン、モチーフ、またはフレームワークを発見することです。具体的には、本論文では、2層以下の、アテンションブロックのみを持つTransformerを研究します。これは、96層を持ち、アテンションブロックとMLPブロックを交互に配置するGPT-3のような大規模で現代的なTransformerとは対照的です。私たちは、Transformerの操作を新しいが数学的に等価な方法で概念化することにより、これらの小さなモデルを理解し、それらが内部でどのように動作するかについて significant な理解を得ることができることを見出しました。特に注目すべきは、私たちが「誘導ヘッド」と名付けた特定のアテンションヘッドが、これらの小さなモデルにおけるインコンテキスト学習を説明できること、そしてこれらのヘッドは少なくとも2つのアテンション層を持つモデルにのみ出現することです。また、特定のデータ上で動作するこれらのヘッドの例もいくつか検討します。この最初の論文では、私たちの洞察をより大きなモデルに適用しようとはしませんが、今後の論文で、Transformerを理解するための私たちの数学的フレームワークと、誘導ヘッドの概念の両方が、はるかに大きく現実的なモデルに対しても少なくとも部分的に関連し続けることを示します。ただし、そのようなモデルを完全に逆解析できるようになるまでには、まだ非常に長い道のりがあります。 結果の要約 Transformerの逆解析結果 Transformerの逆解析の課題を探求するために、いくつかのトイ（toy）の、アテンションのみのモデルを逆解析しました。その中で、私たちは以下の発見をしました。 ゼロ層Transformerは、bigram統計をモデル化します。bigramテーブルは、重みから直接アクセスできます。 1層アテンションのみのTransformerは、bigramモデルと「skip-trigram」（「A… B C」のようなシーケンス）モデルのアンサンブルです。bigramおよびskip-trigramテーブルは、モデルを実行することなく、重みから直接アクセスできます。これらのskip-trigramは驚くほど表現力豊かであることが判明しています。これには、一種の非常に単純なインコンテキスト学習の実装が含まれます。 2層アテンションのみのTransformerは、アテンションヘッドの組み合わせを使用して、はるかに複雑なアルゴリズムを実装できます。これらの組み合わせアルゴリズムは、重みから直接検出することもできます。特に、2層モデルはアテンションヘッドの組み合わせを使用して「誘導ヘッド」を作成します。これは非常に一般的なインコンテキスト学習アルゴリズムです。 誘導ヘッドについては、今後の論文でさらに詳しく説明します。 1層および2層アテンションのみのTransformerは、インコンテキスト学習を実行するために非常に異なるアルゴリズムを使用します。2層アテンションヘッドは、質的に洗練された推論時間アルゴリズム、特に私たちが誘導ヘッドと呼ぶ特別なタイプのアテンションヘッドを使用してインコンテキスト学習を実行し、より大きなモデルに関連する重要な遷移点となります。 概念的なテイクアウェイ Transformerアーキテクチャの多くの微妙な詳細は、InceptionV1 Circuitsの作業とはかなり異なる方法で逆解析に取り組む必要があることがわかりました。これらのポイントはそれぞれ、以下のセクションで詳しく説明しますが、ここでは簡単に要約します。適切なセクションに到達したときに、ここで導入する多くの用語をさらに拡張します。（明確にするために、これらのポイントのいずれかが必ずしも新しいと主張するつもりはありません。多くは他の論文に暗黙的または明示的に存在します。） アテンションヘッドは、それぞれ結果を出力し、それが残差ストリームに追加される独立した操作として理解できます。アテンションヘッドは、計算効率のために、代替の「連結と乗算」の定式化で説明されることが多いですが、これは数学的に同等です。 アテンションのみのモデルは、トークンをロジットへの変更にマッピングする解釈可能なエンドツーエンド関数の合計として記述できます。これらの関数は、モデルを通過する「パス」に対応し、アテンションパターンを固定すると線形になります。 Transformerには膨大な量の線形構造があります。単純に和を分解し、行列の連鎖を乗算することによって、多くのことを学ぶことができます。 アテンションヘッドは、Q K（クエリ-キー）回路とO V（出力-値）回路という、ほぼ独立した2つの計算を持つと理解できます。Q K回路はアテンションパターンを計算し、O V回路は、アテンションされた場合に各トークンが出力にどのように影響するかを計算します。 キー、クエリ、値ベクトルは、低ランク行列 W_Q^T W_K および W_O W_V の計算における中間結果と考えることができます。それらを参照せずにTransformerを記述すると役立つ場合があります。 アテンションヘッドの組み合わせは、Transformerの表現力を大幅に向上させます。アテンションヘッドが組み合わせる方法は3つあり、キー、クエリ、値に対応します。キーとクエリの組み合わせは、値の組み合わせとは大きく異なります。 Transformerのすべてのコンポーネント（トークン埋め込み、アテンションヘッド、MLP層、およびアンエンベディング）は、残差ストリームの異なる部分空間を読み書きすることによって互いに通信します。残差ストリームベクトルを分析する代わりに、モデルを通過するパスに対応する、これらのさまざまな通信チャネルすべてに残差ストリームを分解することが役立つ場合があります。 Transformerの概要 Transformerを逆解析しようとする前に、Transformerのハイレベルな構造を簡単にレビューし、私たちがそれらをどのように考えているかを説明することが役立ちます。多くの場合、同等ですが非標準的な方法でTransformerを再構成することが役立つことがわかりました。メカニスティック解釈可能性は、モデルを人間が解釈可能な部分に分解することを必要とします。重要な最初のステップは、モデルについて推論するのに最も簡単な表現を見つけることです。現代のディープラーニングでは、正当な理由から、計算効率に多くの重点が置かれており、モデルの数学的記述は、モデルを実行するための効率的なコードの書き方に関する決定をしばしば反映しています。しかし、同じ計算を表す多くの同等な方法がある場合、それはおそらく