[論文レビュー] Thinking Like Transformers

Gail Garfinkel Weiss, Yoav Goldberg+1 more|arXiv (Cornell University)|Jun 13, 2021

Advanced Neural Network Applications被引用数 11

ひとこと要約

本稿では、注意機構とフィードフォワード演算を1等級の原始的演算に抽象化することで、Transformerエンコーダーの計算的挙動をモデル化する記号的プログラミング言語RASP（Restricted Access Sequence Processing Language）を紹介する。この言語により、Transformerの能力に関する形式的推論が可能となり、RASPプログラムを用いて、ソート、ヒストグラム計算、Dyck言語認識といったタスクを解くために必要な最小の層数やヘッド数を予測できる。また、訓練済みのTransformerはRASPから導かれた注意パターンと一致するよう学習し、99%以上の精度を達成している。

ABSTRACT

What is the computational model behind a Transformer? Where recurrent neural networks have direct parallels in finite state machines, allowing clear discussion and thought around architecture variants or trained models, Transformers have no such familiar parallel. In this paper we aim to change that, proposing a computational model for the transformer-encoder in the form of a programming language. We map the basic components of a transformer-encoder -- attention and feed-forward computation -- into simple primitives, around which we form a programming language: the Restricted Access Sequence Processing Language (RASP). We show how RASP can be used to program solutions to tasks that could conceivably be learned by a Transformer, and how a Transformer can be trained to mimic a RASP solution. In particular, we provide RASP programs for histograms, sorting, and Dyck-languages. We further use our model to relate their difficulty in terms of the number of required layers and attention heads: analyzing a RASP program implies a maximum number of heads and layers necessary to encode a task in a transformer. Finally, we see how insights gained from our abstraction might be used to explain phenomena seen in recent works.

研究の動機と目的

• RNNの有限オートマトンに類似した、Transformerエンコーダーの形式的計算モデルを提供すること。
• ドメイン特化言語を用いてタスクを記号的プログラムとして表現することで、Transformerの能力に関する高レベルの推論を可能にすること。
• 特定のシーケンス処理タスクを実装するためにTransformerで必要な最小の層数およびアテンションヘッド数を予測すること。
• 訓練済みのTransformerがRASPから導かれた解法と一致するアテンションパターンを学習できることを検証し、記号的推論とニューラルネットワークの挙動を結びつけること。
• RASPの抽象化を用いて、Transformerの変種における実験的現象や、効率的アーキテクチャの限界を説明すること。

提案手法

• 選択、フィルタリング、集約のための原始的演算を備えた最小限のプログラミング言語RASPを設計し、シーケンス上での注意とフィードフォワードの制約を捉える。
• RASPプログラムを、Transformerにおける必要なアテンションパターンとレイヤー構造を定義する同等のアテンションおよびMLP演算にコンパイルする。
• 出力とアテンションパターンの両方の監督を用いてTransformerを訓練：出力には標準の交差エントロピー損失、アテンションパターンにはMSE損失を適用し、RASP解法と一致させる。
• RASPプログラムを用いて、特定のタスクに必要な層数およびヘッド数の理論的下界を導出する。
• RASPで予測された最小サイズのTransformerを訓練し、ヘッド数や層数を減らした際の精度低下を測定することで、RASPプログラムの実行可能性を分析する。
• Dyck-k言語認識のための先行研究の構成を再表現・改善し、RASPを用いてすべてのkについてDyck-kを完全に認識可能であることを証明する。

実験結果

リサーチクエスチョン

• RQ1 RNNの有限オートマトンに類似した方法で、Transformerエンコーダーの計算的挙動をモデル化できる記号的プログラミング言語を構築できるか？
• RQ2 記号的プログラムによって予測されるように、特定のシーケンス処理タスクを実装するためにTransformerで必要な最小の層数およびアテンションヘッド数は何か？
• RQ3 訓練済みのTransformerは、特定のタスクに対してRASPプログラムから導かれたアテンションパターンを正確に学習できるか？
• RQ4 モデルサイズ（層数またはヘッド数）を減らした場合の性能低下はどのように影響を受けるか？RASPはこれらの失敗点を予測できるか？
• RQ5 RASPを用いて、効率的アーキテクチャにおける性能低下などのTransformerの変種における実験的挙動を説明または予測できるか？

主な発見

• 二重ヒストグラム、ソート、最頻出トークンの特定といったタスクのRASPプログラムは、訓練済みのTransformerが99%以上のテスト精度で再現できるアテンションパターンにコンパイル可能であった。
• 出力とアテンションパターンの両方の監督を用いて訓練したTransformerは、二重ヒストグラムタスクで99.9%、ソートタスクで99.8%の精度を達成し、RASPで指定されたアテンションパターンと密接に一致した。
• RASPが予測した最小サイズを下回るモデルサイズにした場合、精度が著しく低下した——例として、1層少ない場合の逆順処理タスクで41.2%の精度に低下——RASPのアーキテクチャ的限界の予測力が裏付けられた。
• RASPの抽象化により、先行研究よりもより簡潔かつ正確なDyck-k言語認識の構成が可能となり、すべてのkについてTransformerがDyck-kを完全に認識可能であることが証明された。
• ソートタスクでは、1層1ヘッドのTransformerが一様なアテンションパターンによりほぼ完璧な精度を達成したため、バケットソートの一種を効果的に実装していると考えられる。
• 本研究により、注意機構に制限を設えた「効率的」なTransformerは、複雑な情報フローを必要とするタスクで失敗する可能性があることが判明し、RASP分析の予測と一致した。