[Paper Review] Interpreting Attention Layer Outputs with Sparse Autoencoders

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

Ferroelectric and Negative Capacitance Devices4 citations

TL;DR

This paper applies sparse autoencoders (SAEs) to attention layer outputs in transformers, demonstrating that SAEs extract sparse, interpretable features corresponding to long-range context, short-range context, and induction patterns. The authors introduce recursive direct feature attribution (RDFA) to trace feature origins back to input tokens, enabling detailed analysis of attention mechanisms and circuit behavior, including resolving the role of redundant induction heads and validating the Indirect Object Identification circuit.

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.

Motivation & Objective

• To demonstrate that sparse autoencoders (SAEs) can decompose attention layer outputs into sparse, interpretable features.
• To develop and apply Recursive Direct Feature Attribution (RDFA) to trace SAE feature activations back to input tokens and upstream components.
• To resolve open questions about polysemantic attention heads and the functional diversity of seemingly redundant induction heads.
• To validate that SAEs identify causally meaningful intermediate variables in known circuits, such as the Indirect Object Identification circuit.
• To provide open-source tools and SAEs for community use in interpreting attention mechanisms.

Proposed method

• Train SAEs on concatenated, pre-linear attention outputs (z_cat) across all heads in a transformer layer to reconstruct high-dimensional activations.
• Use weight-based head attribution to associate learned SAE features with specific attention heads, mitigating polysemanticity issues.
• Implement Recursive Direct Feature Attribution (RDFA) to recursively decompose SAE feature activations into contributions from upstream components, including residual stream features and input tokens.
• Leverage the linear structure of transformers to perform linear decomposition of z_cat into contributions from attention patterns, value projections, and residual stream features.
• Combine SAEs on attention outputs with residual stream SAEs to enable end-to-end attribution from final SAE features back to input tokens.
• Release a tool at https://robertzk.github.io/circuit-explorer to visualize and explore feature circuits on arbitrary prompts.

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.

Experimental results

Research questions

• RQ1 What kinds of concepts do attention layers learn, and can they be meaningfully decomposed into interpretable features?
• RQ2 How do attention heads contribute to model behavior, and what roles do they play in polysemantic or redundant configurations?
• RQ3 Can SAEs on attention outputs help resolve the mystery of why models have many seemingly redundant induction heads?
• RQ4 Do SAE features in attention layers correspond to causally meaningful intermediate variables in known circuits, such as the Indirect Object Identification circuit?
• RQ5 Can recursive attribution techniques like RDFA trace the origin of a feature back to specific input tokens and upstream components?

Key findings

• SAEs successfully decompose attention layer outputs into sparse, interpretable features, including long-range context, short-range context, and induction patterns.
• At least 90% of attention heads in GPT-2 Small are polysemantic, performing multiple unrelated functions depending on context.
• RDFA reveals that seemingly redundant induction heads specialize in either long-prefix or short-prefix induction, resolving a key open question.
• The Indirect Object Identification circuit is better understood when analyzed via SAE features, which capture causally meaningful intermediate variables.
• The model routes information about 'Dave' from the 'is' token to 'isn’t' via attention connections, as confirmed by recursive attribution.
• The authors release trained SAEs and a circuit explorer tool to enable community exploration of attention features across prompts.

(a)