Interpreting Attention Layer Outputs
with Sparse Autoencoders

Here we provide weights for all trained SAEs (Table 1) as well as the interface for feature dashboards that we used to evaluate feature interpretability discussed in Section 3.3.

For GPT-2 Small, you can find the weights here: https://huggingface.co/ckkissane/attn-saes-gpt2-small-all-layers/tree/main. You can view feature dashboards for 30 randomly sampled feature per each layer here: https://ckkissane.github.io/attn-sae-gpt2-small-viz/. We additionally provide a colab notebook demonstrating how to use the SAEs here: https://colab.research.google.com/drive/1hZVEM6drJNsopLRd7hKajp_2v6mm_p70?usp=sharing

For our GELU-2L SAE trained on Layer 1 (the second layer), you can find weights here: https://huggingface.co/ckkissane/tinystories-1M-SAES/blob/main/concat-z-gelu-21-l1-lr-sweep-3/gelu-2l_L1_Hcat_z_lr1.00e-03_l12.00e%2B00_ds16384_bs4096_dc1.00e-07_rie50000_nr4_v78.pt. You can view dashboards for 50 randomly sampled features here: https://ckkissane.github.io/attn-sae-gelu-2l-viz/. We additionally provide a colab notebook showing how to use the SAEs here: https://colab.research.google.com/drive/10zBOdozYR2Aq2yV9xKs-csBH2olaFnsq?usp=sharing

To view the top 10 features attributed to all 144 attention heads in GPT-2 Small (as in Section 4.1) see here: https://robertzk.github.io/gpt2-small-saes/.

Weights for the Gemma-2B SAE can be found here: https://wandb.ai/jbloom/gemma_2b_hook_z/artifacts/model/sae_group_gemma-2b_blocks.6.attn.hook_z_16384/v1/files.

You can also view similar dashboards for any feature from all of our GPT-2 Small SAEs on neuronpedia [30] here: https://www.neuronpedia.org/gpt2-small/att-kk.

Further, we introduce an interactive tool for exploring several attention SAEs throughout a model at https://robertzk.github.io/circuit-explorer and discuss this more in Appendix N.

Code is available at https://github.com/ckkissane/attention-output-saes.

Important details of SAE training include:

• •

  SAE Widths. Our GELU-2L and Gemma-2B SAEs have width $16384$. All of our GPT-2 Small SAEs have width $24576$, with the exception of layers 5 and 7, which have width $49152$.

• •

  Loss Function. We trained our Gemma-2B SAE with a different loss function than the SAEs from other models. For Gemma-2B we closely follow the approach from Olah et al. [52], while for GELU-2L and GPT-2 Small, we closely follow the approach from Bricken et al. [4].

• •

  Training Data. We use activations from hundreds of millions to billions of activations from LM forward passes as input data to the SAE. Following Nanda [41], we use a shuffled buffer of these activations, so that optimization steps don’t use data from highly correlated activations. For GELU-2L we use a mixture of 80% from the C4 Corpus [56] and 20% code (https://huggingface.co/datasets/NeelNanda/c4-code-tokenized-2b). For GPT-2 Small we use OpenWebText (https://huggingface.co/datasets/Skylion007/openwebtext). For Gemma-2B we use https://huggingface.co/datasets/HuggingFaceFW/fineweb. The input activations have sequence length of 128 tokens for all training runs.

• •

  Resampling. For our GELU-2L and GPT-2 Small SAEs we used resampling, a technique which at a high-level reinitializes features that activate extremely rarely on SAE inputs periodically throughout training. We mostly follow the approach described in the ‘Neuron Resampling’ appendix of Bricken et al. [4], except we reapply learning rate warm-up after each resampling event, reducing learning rate to 0.1x the ordinary value, and, increasing it with a cosine schedule back to the ordinary value over the next 1000 training steps. Note we don’t do this for Gemma-2B.

• •

  Optimizer hyperparameters. For the GELU-2L and GPT-2 Small SAEs we use the Adam optimizer with $\beta_{2}=0.99$ and $\beta_{1}=0.9$ and a learning rate of roughly $0.001$. For Gemma-2B SAEs we also use the Adam optimizer with $\beta_{2}=0.999$ and $\beta_{1}=0.9$ and a learning rate of $0.00005$.

In Section 3 we claimed that our Attention Output SAEs are sparse, faithful, and interpretable and we provide evaluations of each SAE in Table 1 to support this claim. In this section we further discuss nuances of the fidelity evaluation, and how our SAEs compare to trained SAEs from other work.

We note that we evaluated fidelity with the cross entropy loss relative to zero ablation (4), which has a few potential pitfalls. First, some would argue that zero ablation may be too harsh a baseline, and that alternative baselines using mean ablation or resample ablation may be more principled. We choose to use zero ablation to stay consistent with prior work from Bricken et al. [4], which made our preliminary results easier to evaluate.

Second, the zero ablation baseline makes it’s hard to compare the quality of SAEs between other sites. Intuitively, zero ablating the residual stream should degrade performance much more than ablating a single attention layer or MLP, so we expect that SAEs trained on the residual stream will have much higher % CE recovered metrics, even if splicing in the residual stream SAE causes a much bigger jump in cross entropy loss. See Rajamanoharan et al. [57] for thorough evaluations of trained SAEs across multiple sites. For this reason, we recommend practitioners additionally record the raw cross entropy loss numbers with and without the SAE spliced in.

We also note that there is a trade off between sparsity and fidelity, and due to limited compute, we are likely far from pareto optimal. Recent work [3, 61] has had success interpreting SAEs with higher numbers of features firing, although it’s not clear what L0 we should target. For example, we might expect more features in the residual stream compared to an attention head, and we might expect larger models to compute more features than smaller models.

With this in mind, it’s hard to compare our SAEs to across work that uses on different models and activation sites. When we trained our SAEs, we closely followed Bricken et al. [4] as a reference. The MLP SAE from their work had a % CE recovered of 79%. They claimed that they generally targeted an L0 norm that is less than 10 or 20. Our SAEs have similar metrics, where we generally targeted and L0 of 20 with 80% CE loss recovered,

To evaluate interpretability for Attention Output SAE features, we manually rate the interpretability of a set of randomly sampled SAE features. For each SAE, the two raters (paper authors) collectively inspected 30 randomly sampled live features.

To assess a feature, the rater determined if there was a clear explanation for the feature’s behavior. The rater viewed the top 20 maximum activating dataset examples for that feature, approximate direct logit effects (i.e. $W_{U}W_{O}\textbf{d}_{i}$), and randomly sampled activating examples from lower activation ranges (as in Bricken et al. [4]). For each max activating dataset example, we also show the corresponding source tokens with the top direct feature attribution by source position (Section 2), and additionally show the weight-based head attribution for all heads in that layer (Section 2). The raters used an interface based on an open source SAE visualizer library [33] modified to support attention layer outputs (see Appendix A). Note that we filter out dead features (features that don’t activate at least once in 100,000 inputs, sometimes also referred to as "ultra low frequency cluster") from our interpretability analysis. These features were excluded from the denominator in reporting percentage interpretable in Table 1.

The raters had a relatively high bar for labeling a feature as interpretable (e.g. noticing a clear pattern with all 20 max activating dataset examples, as well as throughout the randomly sampled activations). However, we note that this methodology heavily relies on subjective human judgement, and thus there is always room for error. We expect both false positives (e.g. the raters are overconfident in their interpretations, despite the feature actually being polysemantic) and false negatives (e.g. the raters might miss more abstract features that are hard to spot with our feature dashboards).

In this section we display in Figure 6 two random examples of false negatives identified during the sensitivity analysis from Section 3.3. To recap, these are examples where our proxy identified a case of board induction (i.e. "<token> board … <token>), but the board induction feature did not fire. We generally notice that while they technically satisfy the board induction pattern, "board" should clearly not be predicted as the next token. This is often because there are even stronger cases of induction for another token (fig. 6).

In this section we automatically detect and quantify a large “<token> is next by induction” feature family from our GELU-2L SAE trained on layer 1. This represents  5% of the non-dead features in the SAE. This is notable, as if there are many “one feature per vocab token” families like this, we may need extremely wide SAEs for larger models.

Based on the findings of the “‘board’ is next by induction” feature (see Section 3.3), we surmised that there might exist more features with this property for different suffixes. Guided by this motivation, we were able to find 586 additional features that exhibited induction-like properties from our GELU-2L SAE. We intend this as a crude proof of concept for automated SAE feature family detection, and to show that there are many induction-like features. We think our method could be made significantly more rigorous with more time, and that it likely has both many false positives and false negatives.

While investigating the “board” feature, we confirmed that attention head 1.6 was an induction head. For each feature dashboard, we also generated a decoder weights distribution that gave an approximation of how much each head is attributed to a given feature. We then chose the following heuristic to identify additional features that exhibited induction-like properties:

Induction Selection Heuristic. For each feature, we compute the weight-based head attribution score (2) to head 1.6. We consider features that have a head attribution score of at least 0.6 as induction feature candidates.

Intuitively, given the normalized norms sum to 1, we expect features satisfying this property to primarily be responsible for producing induction behavior for specific sets of suffix tokens. In our case, we found 586 features that pass the above induction heuristic and are probably related to induction. We note that this is a conservative heuristic, as head 1.4 gets a partial score on the random tokens induction metric, and other heads may also play an induction-like role on some tokens, yet fail the random tokens test [54].

We verified that these are indeed behaviorally related to induction using the following behavioral heuristic:

Induction Behavior Heuristic. For each feature, consider the token corresponding to the max positive boosted logit through the direct readout from $W_{U}W_{O}\textbf{d}_{i}$. For a random sample of 200 examples that contain that token, identify which proportion satisfy:

1. 1.

   For any given instance of the token corresponding to the max positive boosted logit for that feature, the feature does not fire on the first prefix of that token (i.e., the first instance of an “AB” pattern).

2. 2.

   For any subsequent instances of the token corresponding to the max positive boosted logit for that feature occurring in the example, the feature activates on the preceding token (i.e. subsequent instances of an “AB” pattern).

We call the proportion of times the feature activates when it is expected to activate (on instances of A following the first instance of an AB pattern) the induction pass rate for the feature. The heuristic passes if the induction pass rate is > 60%.

With the “board” feature, we saw that the token with the top positive logit boost passed this induction behavior heuristic: for almost every example and each bigram that ends with “board”, the first such bigram did not activate the feature but all subsequent repeated instances did.

We ran this heuristic on the 586 features identified by the Induction Selection Heuristic against 500 features that have attribution < 10% to head 1.6 as a control group (i.e., features we would not expect to display induction-like properties as they are not attributed to the induction head). We found the Induction Behavior Heuristic to perform well at separating the features, as 450/586 features satisfied the > 60% induction pass rate. Conversely, only 3/500 features in the control group satisfied the > 60% induction pass rate (Figure 10).

We now consider an “In questions starting with ‘Which’” feature. We categorized this as one of many “local context” features: a feature that is active in some context, but often only for a short time, and which has some clear ending marker (e.g. a question mark, closing parentheses, etc).

Unlike the induction feature (Section 3.3), we also find that it’s computed by multiple attention heads. The fact that our Attention SAEs extracted a feature relying on multiple heads, and made progress towards understanding it, suggests that we may be able to use Attention Output SAEs as a tool to tackle the hypothesized phenomenon of attention head superposition [51].

We first show that our interpretation is faithful over the entire distribution. We define a crude proxy that checks for the first 10 tokens after "Which" tokens, stop** early at punctuation. Similar to the induction feature, we find that this feature activates with high specificity to this context in the upper activation ranges, although there is polysemanticity for lower activations (Figure 11(a)).

We now show that the feature is computed by multiple heads in layer 1. Over 10 million tokens, we compute the direct feature attribution by head (3) for this feature. We find that head 3 heads have non-trivial (>10%) fraction of variance explained (Figure 12).

Despite this, we still get traction on understanding this feature, motivating attention SAEs as a valuable tool to deal with attention head superposition. We first understand the causal downstream effects of this feature. We find that it primarily "closes the question", by directly boosting the logits of question mark tokens (Figure 9(b)).

We also show that the heads in aggregate are moving information from prior "Which" tokens to compute this feature. We apply DFA by source position (aggregated across all heads) (see Section 2) for all feature activations over 10 million tokens and record aggregate scores for each source token. We find that “Which” source tokens explain >50% the variance, and over 95% of the variance if we filter for feature activations greater than $2$, suggesting that the heads are moving this "Which" to compute the feature.

We now consider an “in a text related to pets” feature. This is one example from a family of ‘high-level context features’ extracted by our SAE. High-level context features often activate for almost the entire context, and don’t have a clear ending marker (like a question mark). To us they appear qualitatively different from the local context features, like “in a question starting with ‘Which’”, which just activate for e.g. all tokens in a sentence.

We first show our interpretation of this feature is faithful. We define a proxy that checks for all tokens that occur after any token from a handcrafted set of pet related tokens (’dog’, ’ pet’, ‘ canine’, etc), and compare the activations of our feature to the proxy. Though the proxy is crude, we find that this feature activates with high specificity in this context in Figure 11(b).

We show that we can understand the downstream effects of this feature. The feature directly boosts logits of pet related tokens (’dog’, ’ pet’, ‘ canine’, etc) in Figure 9(c).

We were able to use techniques like direct feature attribution to learn that high-level context features are natural to implement with a single attention head: the head can just look back for past “pet related tokens” (‘dog’, ‘ pet’, ‘ canine’, ‘ veterinary’, etc) , and move these to compute the feature.

We find that the top attention head is using the pet source tokens to compute the feature. We track the direct feature contributions from source tokens in a handcrafted set of pet related tokens (’dog’, ’pet’, etc) and compute the fraction of variance explained from these source tokens. We confirm that “pet” source tokens explain the majority of the variance, especially when filtering by higher activations, with over 90% fraction of variance explained for activations greater than 2.

In this section we present new feature families that we found in GPT-2 Small, but did not find in the GELU-2L SAE⁶⁶6Note we didn’t exhaustively check every GELU-2L feature. However we never came across these in all of our analysis, whereas we quickly discovered these when looking at random features from GPT-2 Small. This suggests that SAEs are a useful tool that can provide hints about fundamentally different capabilities as we apply them to bigger models.

While the technique from Section 4.1 is not sufficient to prove that a head is monosemantic, we believe that having multiple unrelated features attributed to a head is evidence that the head is doing multiple tasks (i.e. exhibit polysemanticity [13]). We also note that there is a possibility we missed some monosemantic heads due to missing patterns at certain levels of abstraction (e.g. some patterns might not be evident from a small sample of SAE features, and in other instances an SAE might have mistakenly learned some red herring features).

During our investigations of each head, we found 14 monosemantic candidates (i.e. all of the top 10 attributed features for these heads were closely related). This suggests that over 90% of the attention heads in GPT-2 small are performing at least two different tasks. In Section J.1, we list notable heads that are plausibly monosemantic or have suggested roles based on this technique.

In this section we evaluate all of our GPT-2 Small attention SAEs on the IOI task. For each layer, we replace attention output activations with their SAE reconstructed activations and observe the effect on the average logit difference [66] between the correct and incorrect name tokens (as in Makelov et al. [31]). We also measure the KL divergence between the logits of the original model and the logits of the model with the SAE spliced in. We compare the effect of splicing in the SAEs to mean ablating these attention layer outputs from the ABC distribution (as described in Wang et al. [66], this is the IOI distribution but with three different names, rather than one IO and two subjects) to also get a rough sense of how necessary these activations are for the circuit.

We find that splicing in our SAEs at each of the early-middle layers [1, 6] maintains an average logit difference roughly equal to the clean baseline, suggesting that these SAEs are sufficient for circuit analysis. On the other hand, we see layers {0, 7, 8} cause a notable drop in logit difference. The later layers actually cause an increase in logit difference, but we think that these are likely breaking things based on the relatively high average KL divergence, illustrating the importance of using multiple metrics that capture different things (Figure 14). We suspect that these late layer SAEs might be missing features corresponding to the Negative Name Mover (Copy Suppression [34]) heads in the IOI circuit, although we don’t investigate this further.

Wang et al. [66] identify many classes of attention heads spread across multiple layers. To investigate if our SAEs are systematically failing to capture features corresponding to certain heads, we splice in our SAEs for each of these cross-sections (similar to Makelov et al. [31]).

For each role classified by Wang et al. [66], we identify the set of attention layers containing all of these heads. We then replace the attention output activations for all of these layers with their reconstructed activations. Note that we recompute the reconstructed activations sequentially rather than patching all of them in at once. We do this for the following groups of heads:

• •

  Duplicate Token Heads {0, 3}

• •

  Previous Token Heads {2, 4}

• •

  Induction Heads {5, 6}

• •

  S-inhibition Heads {7, 8}

• •

  (Negative) Name Mover Heads {9, 10, 11}

We again see promising signs that the early-middle layer SAEs (corresponding to the Induction and Previous Token Heads) seem sufficient for analysis at the feature level (Figure 15). Unfortunately, it’s also clear that our SAEs are likely not sufficient to analyze the outputs of Layer 0 and the later layers (S-inhibition Heads and (Negative) Name Mover Heads). Thus we are unable to study a full end-to-end feature circuit for IOI.

Why is there such a big difference between cross-sections? It is not clear from our analysis, but one hypothesis is that the middle layers contain more general features such as “I am a duplicate token”, whereas the late layers contain niche name-specific features such as “The name X is next”. Not only do we expect a much greater number of per-name features, but we also expect these features to be relatively rare, and thus harder for the SAEs to learn during training. We are hopeful that this will be improved by ongoing work on the science and scaling of SAEs [46, 57, 52].

Here we provide two additional lines of evidence to show that in GPT-2 Small, 5.1 specializes in "long prefix induction", while 5.5 does "short prefix induction". Note we that we do not use SAEs in this section, but the original hypothesis was motivated by our SAEs (see Section 4.2).

We first check each head’s average direct logit attribution (DLA) [54] to the correct next token as a function of prefix length. We again see that head 5.1’s DLA sharply increases as we enter the long prefix regime, while head 5.5’s DLA remains relatively constant (Figure 19(a)).

We then confirmed that these results hold on a random sample of the training distribution. We first filter for examples where the heads are attending non-trivially to some token⁷⁷7We show a threshold of 0.3. The results generally hold for a range of thresholds. (i.e. not just attending to BOS), and check how often these are examples of n-prefix induction. We find that head 5.1 will mostly attend to tokens in long prefix induction, while head 5.5 is mostly doing normal 1-prefix induction (Figure 19(b)).

As a continuation of Section 4.1, we describe the results of manually inspecting the most salient features for all 144 attention heads to examine the role of every attention head in GPT-2 Small. As in Section 2, we apply equation 2 to identify the top ten features by decoder weight attribution to determine which features are most attributed to a given head. We then identify conceptual grou**s that are exhibited in these features.

In this section we describe the Recursive Direct Feature Attribution technique from Section 2 in more detail. We use Attention Output SAEs from Section 3.2 and residual stream SAEs from Bloom [3] to repeatedly attribute SAE feature activation to upstream SAE feature outputs, all the way back to the input tokens for an arbitrary prompt. The key idea is that if we freeze attention patterns and LayerNorm scales, we can decompose the SAE input activations, $\textbf{z}_{\text{cat}}$, into a linear function of upstream activations. Then we recursively decompose those upstream activations into linear contributions.

In Table 6, we provide a full description of the recursive direct feature attribution (RDFA) algorithm, accompanied by equations for the key linear decomposition.

We now provide a few examples of using the Circuit Explorer tool available at https://robertzk.github.io/circuit-explorer.

Example 1: Decomposing information about name. Consider the prompt: "Amanda Heyman, professional photographer. She". In Figure 20, starting with Attention Output SAE feature L3.F15566, we observe that performing a DFA decomposition along source position and then along residual features highlights:

• •

  a residual feature (3.19755) that maximally activates on names ending with "anda": https://www.neuronpedia.org/gpt2-small/3-res-jb/19755

• •

  a residual feature (3.14186) that maximally activates on "Amanda" and boosts last names: https://www.neuronpedia.org/gpt2-small/3-res-jb/14186

Example 2: Routing "Dave" through "is" to "isn’t". Consider the prompt: "So Dave is a really good friend isn’t" as highlighted in Conmy et al. [9]. Focusing on layer 10, the top Attention Output SAE feature is L10.F14709. In Figure 21, we observe that performing a recursive DFA decomposition along source position and then to upstream attention components shows that the model is routing information about "Dave" via the "is" token to the final "[isn]’t" position.