Sparse Autoencoders Find Highly Interpretable Features in Language Models — part 2/3
Cleaned-text code points [15478, 31456). Bibliography omitted here; source.md includes it.
where F F is the subset of the features which we intervene on (we describe the selection process for F F later in this section). Let 𝐳 \mathbf{z} denote the output logits of the model when you finish applying it to the patched residual stream 𝐱 1 ′ , … , 𝐱 k ′ \mathbf{x}^{\prime}_{1},...,\mathbf{x}^{\prime}_{k} . Finally, compute the KL divergence D K L ( 𝐳 | | 𝐲 ) D_{KL}(\mathbf{z}||\mathbf{y}) , which measures how close the patched model’s predictions are to the target’s. We compare these interventions to equivalent interventions using principal components found as in Section 3.2 .
To select the feature subset F F , we use the Automated Circuit Discovery (ACDC) algorithm of Conmy et al. (2023) . In particular, we use their Algorithm 4.1 on our features, treating them as a flat computational graph in which every feature contributes an independent change to the D K L D_{KL} output metric, as described above and averaged over a test set of 50 IOI data points. The result is an ordering on the features so that patching the next feature usually results in a smaller D K L D_{KL} loss than each previous feature. Then our feature subsets F F are the first k k features under this ordering. We applied ACDC separately on each decomposition.
### 4.2 Precise Localisation of IOI Dictionary Features
We show in Figure 3 that our sparse feature dictionaries allow the same amount of model editing, as measured by KL divergence from the target, in fewer patches (Left) and with smaller edit magnitude (Right) than the PCA decomposition. We also show that this does not happen if we train a non-sparse dictionary ( α = 0 \alpha=0 ). However, dictionaries with a larger sparsity coefficient α \alpha have lower overall reconstruction accuracy which appears in Figure 3 as a larger minimum KL divergence. In Figure 3 we consider interventions on layer 11 of the residual stream, and we plot interventions on other layers in Appendix F .
## 5 Case Studies
In this section, we investigate individual dictionary features, highlighting several that appear to correspond to a single human-understandable explanation (i.e., that are monosemantic). We perform three analyses of our dictionary features to determine their semantic meanings: (1) Input : We identify which tokens activate the dictionary feature and in which contexts, (2) Output : We determine how ablating the feature changes the output logits of the model, and (3) Intermediate features : We identify the dictionary features in previous layers that cause the analysed feature to activate.
### 5.1 Input: Dictionary Features are Highly Monosemantic
We first analyse our dictionary directions by checking what text causes them to activate. An idealised monosemantic dictionary feature will only activate on text corresponding to a single real-world feature, whereas a polysemantic dictionary feature might activate in unrelated contexts.
To better illustrate the monosemanticity of certain dictionary features, we plot the histogram of activations across token activations. This technique only works for dictionary features that activate for a small set of tokens. We find dictionary features that only activate on apostrophes (Figure 4 ); periods; the token “ the”; and newline characters. The apostrophe feature in Figure 4 stands in contrast to the default basis for the residual stream, where the dimension that most represents an apostrophe is displayed in Figure 11 in Appendix D.1 ; this dimension is polysemantic since it represents different information at different activation ranges.
Although the dictionary feature discussed in the previous section activates only for apostrophes, it does not activate on all apostrophes. This can be seen in Figures 14 and 15 in Appendix D.2 , showing two other apostrophe-activating dictionary features, but for different contexts (such as “[I/We/They]’ll” and “[don/won/wouldn]’t”). Details for how we searched and selected for dictionary features can be found in Appendix D.3 .
### 5.2 Output: Dictionary Features have Intuitive Effects on the Logits
In addition to looking at which tokens activate the dictionary feature, we investigate how dictionary features affect the model’s output predictions for the next token by ablating the feature from the residual stream 5 5 5 Specifically we use less-than-rank-one ablation, where we lower the activation vector in the direction of the feature only up to the point where the feature is no longer active. . If our dictionary feature is interpretable, subtracting its value from the residual stream should have a logical effect on the predictions of the next token. We see in Figure 4 (Right) that the effect of removing the apostrophe feature mainly reduces the logit for the following “s”. This matches what one would expect from a dictionary feature that detects apostrophes and is used by the model to predict the “s” token that would appear immediately after the apostrophe in possessives and contractions like “let’s”.
### 5.3 Intermediate Features: Dictionary Features Allow Automatic Circuit Detection
We can also understand dictionary features in relation to the upstream and downstream dictionary features: given a dictionary feature, which dictionary features in previous layers cause it to activate, and which dictionary features in later layers does it cause to activate?
To automatically detect the relevant dictionary features, we choose a target dictionary feature such as layer 5’s feature for tokens in parentheses which predicts a closing parentheses (Figure 5 ). For this target dictionary feature, we find its maximum activation M M across our dataset, then sample 20 contexts that cause the target feature to activate in the range [ M / 2 , M ] [M/2,M] . For each dictionary feature in the previous layer, we rerun the model while ablating this feature and sort the previous-layer features by how much their ablation decreased the target feature. If desired, we can then recursively apply this technique to the dictionary features in the previous layer with a large impact. The results of this process form a causal tree, such as Figure 5 .
Being the last layer, layer 5’s role is to output directions that directly correspond to tokens in the unembedding matrix. In fact, when we unembed feature 5 2027 5_{2027} , the top-tokens are all closing parentheses variations. Intuitively, previous layers will detect all situations that precede closing parentheses, such as dates, acronyms, and phrases.
## 6 Discussion
### 6.1 Related Work
A limited number of previous works have learned dictionaries of sparsely-activating features in pre-trained models, including Yun et al. (2021) and Sharkey et al. (2023) , the latter of which motivated this work. However, similar methods have been applied in other domains, in particular in understanding neurons in the visual cortex ( Olshausen & Field, 2004 ; Wright & Ma, 2022 ) .
In contrast to our approach, where we try to impose sparsity after training, many previous works have encouraged sparsity in neural networks via changes to the architecture or training process. These approaches include altering the attention mechanism ( Correia et al., 2019 ) , adding ℓ 1 \ell^{1} penalties to neuron activations ( Kasioumis et al., 2021 ; Georgiadis, 2019 ) , pruning neurons ( Frankle & Carbin, 2018 ) , and using the softmax function as the non-linearity in the MLP layers ( Elhage et al., 2022a ) . However, training a state-of-the-art foundation model with these additional constraints is difficult ( Elhage et al., 2022a ) , and improvements to interpretability are not always realized ( Meister et al., 2021 ) .
### 6.2 Limitations and Future Work
While we have presented evidence that our dictionary features are interpretable and causally important, we do not achieve 0 reconstruction loss (Equation 4 ), indicating that our dictionaries fail to capture all the information in a layer’s activations. We have also confirmed this by measuring the perplexity of the model’s predictions when a layer is substituted with its reconstruction. For instance, replacing the residual stream activations in layer 2 of Pythia-70M with our reconstruction of those activations increases the perplexity on the Pile ( Gao et al., 2020 ) from 25 to 40. To reduce this loss of information, we would like to explore other sparse autoencoder architectures and to try minimizing the change in model outputs when replacing the activations with our reconstructed vectors, rather than the reconstruction loss. Future efforts could also try to improve feature dictionary discovery by incorporating information about the weights of the model or dictionary features found in adjacent layers into the training process.
Our current methods for training sparse autoencoders are best suited to the residual stream. There is evidence that they may be applicable to the MLPs (see Appendix C ), but the training pipeline used to train the dictionaries in this paper is not able to robustly learn overcomplete bases in the intermediate layers of the MLP. We’re excited by future work investigating what changes can be made to better understand the computations performed by the attention heads and MLP layers, each of which poses different challenges.
In Section 4 , we show that for the IOI task, behaviour is dependent on a relatively small number of features. Because our dictionary is trained in a task-agnostic way, we expect this result to generalize to similar tasks and behaviours, but more work is needed to confirm this suspicion. If this property generalizes, we would have a set of features which allow for understanding many model behaviours using just a few features per behaviour. We would also like to trace the causal dependencies between features in different layers, with the overarching goal of providing a lens for viewing language models under which causal dependencies are sparse. This would hopefully be a step towards the eventual goal of building an end-to-end understanding of how a model computes its outputs.
### 6.3 Conclusion
Sparse autoencoders are a scalable, unsupervised approach to disentangling language model network features from superposition. Our approach requires only unlabelled model activations and uses orders of magnitude less compute than the training of the original models. We have demonstrated that the dictionary features we learn are more interpretable by autointerpretation, letting us pinpoint the features responsible for a given behaviour more finely, and are more monosemantic than comparable methods. This approach could facilitate the mapping of model circuits, targeted model editing, and a better understanding of model representations.
An ambitious dream in the field of interpretability is enumerative safety ( Elhage et al., 2022b ) : producing a human-understandable explanation of a model’s computations in terms of a complete list of the model’s features and thereby providing a guarantee that the model will not perform dangerous behaviours such as deception. We hope that the techniques we presented in this paper also provide a step towards achieving this ambition.
#### Acknowledgments
We would like to thank the OpenAI Researcher Access Program for their grant of model credits for the autointerpretation and CoreWeave for providing EleutherAI with the computing resources for this project. We also thank Nora Belrose, Arthur Conmy, Jake Mendel, and the OpenAI Automated Interpretability Team (Jeff Wu, William Saunders, Steven Bills, Henk Tillman, and Daniel Mossing) for valuable discussions regarding the design of various experiments. We thank Wes Gurnee, Adam Jermyn, Stella Biderman, Leo Gao, Curtis Huebner, Scott Emmons, and William Saunders for their feedback on earlier versions of this paper. Thanks to Delta Hessler for proofreading. LR is supported by the Long Term Future Fund. RH is supported by an Open Philanthropy grant. HC was greatly helped by the MATS program, funded by AI Safety Support.
## Appendix A Autointerpretation Protocol
The autointerpretability process consists of five steps and yields both an interpretation and an autointerpretability score:
1. On each of the first 50,000 lines of OpenWebText, take a 64-token sentence-fragment, and measure the feature’s activation on each token of this fragment. Feature activations are rescaled to integer values between 0 and 10.
2. Take the 20 fragments with the top activation scores and pass 5 of these to GPT-4, along with the rescaled per-token activations. Instruct GPT-4 to suggest an explanation for when the feature (or neuron) fires, resulting in an interpretation.
3. Use GPT-3.5 6 6 6 While the process described in Bills et al. (2023) uses GPT-4 for the simulation step, we use GPT-3.5. This is because the simulation protocol requires the model’s logprobs for scoring, and OpenAI’s public API for GPT-3.5 (but not GPT-4) supports returning logprobs. to simulate the feature across another 5 highly activating fragments and 5 randomly selected fragments (with non-zero variation) by asking it to provide the per-token activations.
4. Compute the correlation of the simulated activations and the actual activations. This correlation is the autointerpretability score of the feature. The texts chosen for scoring a feature can be random text fragments, fragments chosen for containing a particularly high activation of that feature, or an even mixture of the two. We use a mixture of the two unless otherwise noted, also called ‘top-random’ scoring.
5. If, amongst the 50,000 fragments, there are fewer than 20 which contain non-zero variation in activation, then the feature is skipped entirely.
Although the use of random fragments in Step 4 is ultimately preferable given a large enough sample size, the small sample sizes of a total of 640 tokens used for analysis mean that a random sample will likely not contain any highly activating examples for all but the most common features, making top-random scoring a desirable alternative.
## Appendix B Sparse Autoencoder Training and Hyperparameter Selection
To train the sparse autoencoder described in Section 2 , we use data from the Pile ( Gao et al., 2020 ) , a large, public webtext corpus. We run the model that we want to interpret over this text while caching and saving the activations at a particular layer. These activations then form a dataset, which we use to train the autoencoders. The autoencoders are trained with the Adam optimiser with a learning rate of 1e-3 and are trained on 5-50M activation vectors for 1-3 epochs, with larger dictionaries taking longer to converge. A single training run using this quantity of data completes in under an hour on a single A40 GPU.
When varying the hyperparameter α \alpha which controls the importance of the sparsity loss term, we consistently find a smooth tradeoff between the sparsity and accuracy of our autoencoder, as shown in Figure 6 . The lack of a ‘bump’ or ‘knee’ in these plots provides some evidence that there is not a single correct way to decompose activation spaces into a sparse basis, though to confirm this would require many additional experiments. Figure 7 shows the convergence behaviour of a set of models with varying α \alpha over multiple epochs.
## Appendix C Further Autointerpretation Results
### C.1 Interpretability is Consistent across Dictionary Sizes
We find that larger interpretability scores of our feature dictionaries are not limited to overcomplete dictionaries (where the ratio, R R , of dictionary features to model dimensions is > 1 >1 ), but occurs even in dictionaries that are smaller than the underlying basis, as shown in Figure 8 . These small dictionaries are able to reconstruct the activation vectors less accurately, so with each feature being similarly interpretable, the larger dictionaries will be able to explain more of the overall variance.
### C.2 High Interpretability Scores Are Not an Artefact of Top Scoring