[論文レビュー] Interpreting Attention Layer Outputs with Sparse Autoencoders

Connor Kissane, Robert Krzyzanowski+3 more|arXiv (Cornell University)|Jun 25, 2024

Ferroelectric and Negative Capacitance Devices被引用数 4

ひとこと要約

この論文では、変換器のアテンション層出力にスパース自己オートエンコーダー（SAEs）を適用し、長距離コンテキスト、短距離コンテキスト、およびインダクションパターンに対応するスパースで解釈可能な特徴が抽出されることを示している。著者らは、特徴の起源を入力トークンまで遡るための再帰的直接特徴帰属法（RDFA）を導入し、アテンションメカニズムおよび回路行動の詳細な分析を可能にした。これにより、重複するインダクションヘッドの役割が解明され、間接目的語同定回路の妥当性が検証された。

ABSTRACT

Decomposing model activations into interpretable components is a key open problem in mechanistic interpretability. Sparse autoencoders (SAEs) are a popular method for decomposing the internal activations of trained transformers into sparse, interpretable features, and have been applied to MLP layers and the residual stream. In this work we train SAEs on attention layer outputs and show that also here SAEs find a sparse, interpretable decomposition. We demonstrate this on transformers from several model families and up to 2B parameters. We perform a qualitative study of the features computed by attention layers, and find multiple families: long-range context, short-range context and induction features. We qualitatively study the role of every head in GPT-2 Small, and estimate that at least 90% of the heads are polysemantic, i.e. have multiple unrelated roles. Further, we show that Sparse Autoencoders are a useful tool that enable researchers to explain model behavior in greater detail than prior work. For example, we explore the mystery of why models have so many seemingly redundant induction heads, use SAEs to motivate the hypothesis that some are long-prefix whereas others are short-prefix, and confirm this with more rigorous analysis. We use our SAEs to analyze the computation performed by the Indirect Object Identification circuit (Wang et al.), validating that the SAEs find causally meaningful intermediate variables, and deepening our understanding of the semantics of the circuit. We open-source the trained SAEs and a tool for exploring arbitrary prompts through the lens of Attention Output SAEs.

研究の動機と目的

• スパース自己オートエンコーダー（SAEs）がアテンション層出力をスパースで解釈可能な特徴に分解できることを示すこと。
• 再帰的直接特徴帰属法（RDFA）を発展・適用し、SAE特徴活性化を入力トークンおよび上流コンポONENTにまで遡って追跡すること。
• 多意味的アテンションヘッドや、明らかに重複するように見えるインダクションヘッドの機能的多様性についての未解決問題を解消すること。
• SAEが既知の回路（例：間接目的語同定回路）における因果的に意味のある中間変数を特定できることを検証すること。
• コミュニティがアテンションメカニズムを解釈するために利用可能なオープンソースのツールおよびSAEを提供すること。

提案手法

• 変換器層内の全ヘッドの連結された、線形変換前のアテンション出力（z_cat）にSAEを訓練し、高次元の活性化を再構築すること。
• 重みベースのヘッド帰属を用いて、学習されたSAE特徴を特定のアテンションヘッドに関連付けることで、多意味的問題を軽減すること。
• 再帰的直接特徴帰属法（RDFA）を実装し、SAE特徴活性化を上流コンポONENT（残留ストリーム特徴および入力トークンを含む）への寄与に再帰的に分解すること。
• 変換器の線形構造を活用して、z_catをアテンションパターン、値射影、および残留ストリーム特徴への寄与に線形分解すること。
• アテンション出力に適用したSAEと残留ストリームSAEを組み合わせ、最終的なSAE特徴から入力トークンまでの一貫した帰属を可能にすること。
• 任意のプロンプト上で特徴回路を可視化・探索できるツールを https://robertzk.github.io/circuit-explorer で公開すること。

Figure 1 : Overview. We train Sparse Autoencoders (SAEs) on $\textbf{z}_{\text{cat}}$ , the attention layer outputs pre-linear, concatenated across all heads. The SAEs extract linear directions that correspond to concepts in the model, giving us insight into what attention layers learn in practice.

実験結果

リサーチクエスチョン

• RQ1 アテンション層が学習する概念の種類は何か？ そして、それらは意味的に解釈可能な特徴に分解可能か？
• RQ2 アテンションヘッドはモデル行動にどのように寄与しているのか？ また、多意味的または重複する構成において、それぞれが果たす役割は何か？
• RQ3 アテンション出力に適用したSAEは、モデルに多数の明らかに重複するインダクションヘッドがある理由を解明するのを手伝えるか？
• RQ4 アテンション層のSAE特徴は、既知の回路（例：間接目的語同定回路）における因果的に意味のある中間変数に対応しているか？
• RQ5 再帰的帰属技術（例：RDFA）は、特定の入力トークンおよび上流コンポONENTにまで特徴の起源を遡ることができるか？

主な発見

• SAEは、アテンション層出力をスパースで解釈可能な特徴（長距離コンテキスト、短距離コンテキスト、インダクションパターンを含む）に成功して分解した。
• GPT-2 Smallの少なくとも90％のアテンションヘッドは多意味的であり、文脈に応じて複数の無関係な機能を実行している。
• RDFAは、明らかに重複するインダクションヘッドが、長接頭語または短接頭語のインダクションのいずれかに特化していることを明らかにした。これは、重要な未解決問題の一つを解消した。
• 間接目的語同定回路は、SAE特徴を介して分析することでより良く理解できるようになった。SAE特徴は因果的に意味のある中間変数を捉えている。
• 再帰的帰属によって確認された通り、モデルは「is」トークンから「isn’t」に「Dave」に関する情報をアテンション接続経由でルーティングしている。
• 著者らは、コミュニティがプロンプト全体にわたってアテンション特徴を探索できるよう、訓練済みのSAEと回路エクスプローラー・ツールを公開した。

(a)