[論文レビュー] Tracr: Compiled Transformers as a Laboratory for Interpretability

David Lindner, János Kramár+4 more|arXiv (Cornell University)|Jan 12, 2023

Explainable Artificial Intelligence (XAI)被引用数 6

ひとこと要約

Tracr は人間が読みやすい RASP プログラムを、構造が明確で解釈可能な標準的なデコーダー専用トランスフォーマーモデルにコンパイルする。これにより、解釈可能性に関する制御された実験が可能になる。Tracr は、真の計算メカニズムが明確なモデルを生成するため、スーパーポジションや勾配ベースの帰属分析、因果トレーシングといった解釈可能性手法の検証が可能になる。

ABSTRACT

We show how to "compile" human-readable programs into standard decoder-only transformer models. Our compiler, Tracr, generates models with known structure. This structure can be used to design experiments. For example, we use it to study "superposition" in transformers that execute multi-step algorithms. Additionally, the known structure of Tracr-compiled models can serve as ground-truth for evaluating interpretability methods. Commonly, because the "programs" learned by transformers are unknown it is unclear whether an interpretation succeeded. We demonstrate our approach by implementing and examining programs including computing token frequencies, sorting, and parenthesis checking. We provide an open-source implementation of Tracr at https://github.com/google-deepmind/tracr.

研究の動機と目的

• 解釈可能性研究における真の説明の欠如に応えるために、人間が読みやすい計算メカニズムが明確なトランスフォーマーモデルを構築すること。
• 真の計算が事前に分かっているため、再現可能で構造的な環境を提供し、解釈可能性ツールのテストを可能にすること。
• 複数ステップのアルゴリズムを実行するモデルにおける、スーパーポジションなどのニューラルネットワーク現象の制御された調査を可能にすること。
• コンパイル済みでトレース可能なモデルを通じて、トランスフォーマーのメカニズムを理解する教育的・実験的プラットフォームを提供すること。
• 真の動作が分かっているモデルを用いて、分類器プローブ、勾配ベースの帰属分析、因果トレーシングといった解釈可能性手法の評価を支援すること。

提案手法

• Tracr は、トランスフォーマー計算のためのドメイン固有言語である RASP を、構造が明確な標準的なデコーダー専用トランスフォーマー構造にコンパイルする。
• コンパイラは、各レイヤーの関数（例：インジケータの計算、トークンの選択、値の集約）が高レベルのプログラム操作に直接対応するモデルを生成する。
• 残差ストリームは中心的なメモリ機構として使用され、アテンション層と MLP 層は、選択や集約といった特定の計算ステップを明示的に実装するように設計されている。
• この手法はトランスフォーマー回路の視点を活用しており、クエリ・キーと値出力の行列をパrameter化することで、モジュラーで解釈可能な設計を実現している。
• モデルは勾配降下法で学習され、内部表現の分析を通じて、周波数カウント、ソーティング、括弧の整合性チェックといったアルゴリズムの正しく実装されているかを検証する。
• スーパーポジションの研究では、学習済みモデルを圧縮し、特徴が低次元部分空間にどのように表現されたり、不必要な特徴がどのように削除されるかを観察する。

実験結果

リサーチクエスチョン

• RQ1 RASP プログラムからコンパイルされたトランスフォーマーモデルは、解釈可能性手法の評価に用いる真のベンチマークとして利用可能か？
• RQ2 圧縮されたトランスフォーマーモデルは、スーパーポジションにおいて特徴をどのように表現するか？理論予測通りに、重要でない特徴が削除されるか？
• RQ3 勾配ベースの帰属分析や因果トレーシングといった解釈可能性ツールは、Tracr が生成するモデルの既知の計算構造をどの程度正しく回復できるか？
• RQ4 ソーティングや括弧の整合性チェックといった複数ステップのアルゴリズムは、解釈可能なアテンションおよび MLP パatters で信頼性高くコンパイル可能か？
• RQ5 ランク1の述語制約がある場合、注目ベースのセレクタで複雑な論理演算（例：OR、AND）を実装する際の制限は何か？

主な発見

• Tracr は、トークンの周波数カウント、ソーティング、括弧の整合性チェックといったアルゴリズムを、内部メカニズムが明確なトランスフォーマーモデルに正しくコンパイルした。
• コンパイル済みモデルは明確で解釈可能なアテンションおよび MLP パatters を示しており、例えば1つのアテンションヘッドが選択を、別のヘッドが集約を実行しており、意図したプログラム論理と一致している。
• 勾配降下法による圧縮によって、Tracr モデルは不要な特徴を削除し、重要度の低い特徴をスーパーポジションに表現するようになることが確認され、Elhage 他 (2022b) の予測を裏付けた。
• モデルは解釈可能性ツールの検証が可能である。例えば、勾配ベースの帰属分析や因果トレーシングは、既知の真のメカニズムと比較可能な基準を提供する。
• OR 演算子のようなブール演算子は、特定の条件（例：共有の s-ops）が満たされている場合にのみ実装可能であり、一般にはランク1の述語制約によって制限を受ける。
• 本手法は、複雑なアルゴリズムが予測可能でモジュラーなコンponents を持つトランスフォーマーにコンパイル可能であることを示しており、研究および教育的応用を支援する。