Sparse Autoencoders Find Highly Interpretable Features in Language Models β part 1/3
Cleaned-text code points [0, 15478). Bibliography omitted here; source.md includes it.
# Sparse Autoencoders Find Highly Interpretable Features in Language Models
###### Abstract
One of the roadblocks to a better understanding of neural networksβ internals is polysemanticity , where neurons appear to activate in multiple, semantically distinct contexts. Polysemanticity prevents us from identifying concise, human-understandable explanations for what neural networks are doing internally. One hypothesised cause of polysemanticity is superposition , where neural networks represent more features than they have neurons by assigning features to an overcomplete set of directions in activation space, rather than to individual neurons. Here, we attempt to identify those directions, using sparse autoencoders to reconstruct the internal activations of a language model. These autoencoders learn sets of sparsely activating features that are more interpretable and monosemantic than directions identified by alternative approaches, where interpretability is measured by automated methods. Moreover, we show that with our learned set of features, we can pinpoint the features that are causally responsible for counterfactual behaviour on the indirect object identification task ( Wang et al., 2022 ) to a finer degree than previous decompositions. This work indicates that it is possible to resolve superposition in language models using a scalable, unsupervised method. Our method may serve as a foundation for future mechanistic interpretability work, which we hope will enable greater model transparency and steerability.
## 1 Introduction
Advances in artificial intelligence (AI) have resulted in the development of highly capable AI systems that make decisions for reasons we do not understand. This has caused concern that AI systems that we cannot trust are being widely deployed in the economy and in our lives, introducing a number of novel risks ( Hendrycks et al., 2023 ) , including potential future risks that AIs might deceive humans in order to accomplish undesirable goals ( Ngo et al., 2022 ) . Mechanistic interpretability seeks to mitigate such risks through understanding how neural networks calculate their outputs, allowing us to reverse engineer parts of their internal processes and make targeted changes to them ( Cammarata et al., 2021 ; Wang et al., 2022 ; Elhage et al., 2021 ) .
To reverse engineer a neural network, it is necessary to break it down into smaller units (features) that can be analysed in isolation. Using individual neurons as these units has had some success ( Olah et al., 2020 ; Bills et al., 2023 ) , but a key challenge has been that neurons are often polysemantic , activating for several unrelated types of feature ( Olah et al., 2020 ) . Also, for some types of network activations, such as the residual stream of a transformer, there is little reason to expect features to align with the neuron basis ( Elhage et al., 2023 ) .
Elhage et al. (2022b) investigate why polysemanticity might arise and hypothesise that it may result from models learning more distinct features than there are dimensions in the layer. They call this phenomenon superposition . Since a vector space can only have as many orthogonal vectors as it has dimensions, this means the network would learn an overcomplete basis of non-orthogonal features. Features must be sufficiently sparsely activating for superposition to arise because, without high sparsity, interference between non-orthogonal features prevents any performance gain from superposition. This suggests that we may be able to recover the networkβs features by finding a set of directions in activation space such that each activation vector can be reconstructed from a sparse linear combinations of these directions. This is equivalent to the well-known problem of sparse dictionary learning ( Olshausen & Field, 1997 ) .
Building on Sharkey et al. (2023) , we train sparse autoencoders to learn these sets of directions. Our approach is also similar to Yun et al. (2021) , who apply sparse dictionary learning to all residual stream layers in a language model simultaneously. Our method is summarised in Figure 1 and described in Section 2 .
We then use several techniques to verify that our learned features represent a semantically meaningful decomposition of the activation space. First, we show that our features are on average more interpretable than neurons and other matrix decomposition techniques, as measured by autointerpretability scores (Section 3 ) ( Bills et al., 2023 ) . Next, we show that we are able to pinpoint the features used for a set task more precisely than other methods (Section 4 ). Finally, we run case studies on a small number of features, showing that they are not only monosemantic but also have predictable effects on the model outputs, and can be used for fine-grained circuit detection. (Section 5 ).
## 2 Taking Features out of Superposition with Sparse Dictionary Learning
To take network features out of superposition, we employ techniques from sparse dictionary learning ( Olshausen & Field, 1997 ; Lee et al., 2006 ) . Suppose that each of a given set of vectors { π± i } i = 1 n vec β β d \{\mathbf{x}_{i}\}_{i=1}^{n_{\text{vec}}}\subset\mathbb{R}^{d} is composed of a sparse linear combination of unknown vectors { π j } j = 1 n gt β β d \{\mathbf{g}_{j}\}_{j=1}^{n_{\text{gt}}}\subset\mathbb{R}^{d} , i.e. π± i = β j a i , j β π j \mathbf{x}_{i}=\sum_{j}a_{i,j}\mathbf{g}_{j} where π π’ \mathbf{a_{i}} is a sparse vector. In our case, the data vectors { π± i } i = 1 n vec \{\mathbf{x}_{i}\}_{i=1}^{n_{\text{vec}}} are internal activations of a language model, such as Pythia-70M ( Biderman et al., 2023 ) , and { π j } j = 1 n gt \{\mathbf{g}_{j}\}_{j=1}^{n_{\text{gt}}} are unknown, ground truth network features. We would like learn a dictionary of vectors, called dictionary features, { π k } k = 1 n feat β β d \{\mathbf{f}_{k}\}_{k=1}^{n_{\text{feat}}}\subset\mathbb{R}^{d} where for any network feature π j \mathbf{g}_{j} there exists a dictionary feature π k \mathbf{f}_{k} such that π j β π k \mathbf{g}_{j}\approx\mathbf{f}_{k} .
To learn the dictionary, we train an autoencoder with a sparsity penalty term on its hidden activations. The autoencoder is a neural network with a single hidden layer of size d hid = R β d in d_{\text{hid}}=Rd_{\text{in}} , where d in d_{\text{in}} is the dimension of the language model internal activation vectors 1 1 1 We mainly study residual streams in Pythia-70M and Pythia 410-M, for which the residual streams are of size d in = 512 d_{\text{in}}=512 and d in = 1024 d_{\text{in}}=1024 , respectively ( Biderman et al., 2023 ) , and R R is a hyperparameter that controls the ratio of the feature dictionary size to the model dimension. We use the ReLU activation function in the hidden layer ( Fukushima, 1975 ) . We also use tied weights for our neural network, meaning the weight matrices of the encoder and decoder are transposes of each other. 2 2 2 We use tied weights because (a) they encode our expectation that the directions which detect and define the feature should be the same or highly similar, (b) they halve the memory cost of the model, and (c) they remove ambiguity about whether the learned direction should be interpreted as the encoder or decoder direction. They do not reduce performance when training on residual stream data but we have observed some reductions in performance when using MLP data. Thus, on input vector π± β { π± i } \mathbf{x}\in\{\mathbf{x}_{i}\} , our network produces the output π± ^ \mathbf{\hat{x}} , given by π \displaystyle\mathbf{c} = \displaystyle= ReLU β‘ ( M β π± + π ) \displaystyle\mathrm{ReLU}(M\mathbf{x}+\mathbf{b}) (1) π± ^ \displaystyle\mathbf{\hat{x}} = \displaystyle= M T β π \displaystyle M^{T}\mathbf{c} (2) = \displaystyle= β i = 0 d hid β 1 c i β π i \displaystyle\sum_{i=0}^{d_{\text{hid}}-1}c_{i}\mathbf{f}_{i} (3)
where M β β d hid Γ d in M\in\mathbb{R}^{d_{\text{hid}}\times d_{\text{in}}} and π β β d hid \mathbf{b}\in\mathbb{R}^{d_{\text{hid}}} are our learned parameters, and M M is normalised row-wise 3 3 3 Normalisation of the rows (dictionary features) prevents the model from reducing the sparsity loss term β π β 1 ||\mathbf{c}||_{1} by increasing the size of the feature vectors in M M . . Our parameter matrix M M is our feature dictionary, consisting of d hid d_{\text{hid}} rows of dictionary features π i \mathbf{f}_{i} . The output π± ^ \mathbf{\hat{x}} is meant to be a reconstruction of the original vector π± \mathbf{x} , and the hidden layer π \mathbf{c} consists of the coefficients we use in our reconstruction of π± \mathbf{x} .
Our autoencoder is trained to minimise the loss function
β β‘ ( π± ) = β π± β π± ^ β 2 2 β Reconstruction loss + Ξ± β β π β 1 β Sparsity loss \mathcal{L}(\mathbf{x})=\underbrace{||\mathbf{x}-\mathbf{\hat{x}}||_{2}^{2}}_{\text{Reconstruction loss}}+\underbrace{\alpha||\mathbf{c}||_{1}}_{\text{Sparsity loss}} (4)
where Ξ± \alpha is a hyperparameter controlling the sparsity of the reconstruction. The β 1 \ell^{1} loss term on π \mathbf{c} encourages our reconstruction to be a sparse linear combination of the dictionary features. It can be shown empirically ( Sharkey et al., 2023 ) and theoretically ( Wright & Ma, 2022 ) that reconstruction with an β 1 \ell^{1} penalty can recover the ground-truth features that generated the data. For the further details of our training process, see Appendix B .
## 3 Interpreting Dictionary Features
### 3.1 Interpretability at Scale
Having learned a set of dictionary features, we want to understand whether our learned features display reduced polysemanticity, and are therefore more interpretable. To do this in a scalable manner, we require a metric to measure how interpretable a dictionary feature is. We use the automated approach introduced in Bills et al. (2023) because it scales well to measuring interpretability on the thousands of dictionary features our autoencoders learn. In summary, the autointerpretability procedure takes samples of text where the dictionary feature activates, asks a language model to write a human-readable interpretation of the dictionary feature, and then prompts the language model to use this description to predict the dictionary featureβs activation on other samples of text. The correlation between the modelβs predicted activations and the actual activations is that featureβs interpretability score. See Appendix A and Bills et al. (2023) for further details.
We show descriptions and top-and-random scores for five dictionary features from the layer 1 residual stream in Table 1 . The features shown are the first five under the (arbitrary) ordering in the dictionary.
### 3.2 Sparse dictionary features are more interpretable than baselines
We assess our interpretability scores against a variety of alternative methods for finding dictionaries of features in language models. In particular, we compare interpretability scores on our dictionary features to those produced by a) the default basis, b) random directions, c) Principal Component Analysis (PCA), and d) Independent Component Analysis (ICA). For the random directions and for the default basis in the residual stream, we replace negative activations with zeros so that all feature activations are nonnegative 4 4 4 For PCA we use an online estimation approach and run the decomposition on the same quantity of data we used for training the autoencoders. For ICA, due to the slower convergence times, we run on only 2GB of data, approximately 4 million activations for the residual stream and 1m activations for the MLPs. .
Figure 2 shows that our dictionary features are far more interpretable by this measure than dictionary features found by comparable techniques. We find that the strength of this effect declines as we move through the model, being comparable to ICA in layer 4 and showing minimal improvement in the final layer.
This could indicate that sparse autoencoders work less well in later layers but also may be connected to the difficulties of automatic interpretation, both because by building on earlier layers, later features may be more complex, and because they are often best explained by their effect on the output. Bills et al. (2023) showed that GPT-4 is able to generate explanations that are very close to the average quality of the human-generated explanations given similar data. However, they also showed that current LLMs are limited in the kinds of patterns that they can find, sometimes struggling to find patterns that center around next or previous tokens rather than the current token, and in the current protocol are unable to verify outputs by looking at changes in output or other data.
We do show, in Section 5 , a method to see a featureβs causal effect on the output logits by hand, but we currently do not send this information to the language model for hypothesis generation. The case studies section also demonstrates a closing parenthesis dictionary feature, showing that these final layer features can give insight into the modelβs workings.
See Appendix C for a fuller exploration of different learned dictionaries through the lens of automatic interpretability, looking at both the MLPs and the residual stream.
## 4 Identifying Causally-Important Dictionary Features for Indirect Object Identification
In this section, we quantify whether our learned dictionary features localise a specific model behaviour more tightly than the PCA decomposition of the modelβs activations. We do this via activation patching, a form of causal mediation analysis ( Vig et al., 2020 ) , through which we edit the modelβs internal activations along the directions indicated by our dictionary features and measure the changes to the modelβs outputs. We find that our dictionary features require fewer patches to reach a given level of KL divergence on the task studied than comparable decompositions (Figure 3 ).
Specifically, we study model behaviour on the Indirect Object Identification (IOI) task ( Wang et al., 2022 ) , in which the model completes sentences like βThen, Alice and Bob went to the store. Alice gave a snack to β. This task was chosen because it captures a simple, previously-studied model behaviour. Recall that the training of our feature dictionaries does not emphasize any particular task.
### 4.1 Adapting activation patching to dictionary features
In our experiment, we run the model on a counterfactual target sentence, which is a variant of the base IOI sentence with the indirect object changed (e.g., with βBobβ replaced by βVanessaβ); save the encoded activations of our dictionary features; and use the saved activations to edit the modelβs residual stream when run on the base sentence.
In particular, we perform the following procedure. Fix a layer of the model to intervene on. Run the model on the target sentence, saving the model output logits π² \mathbf{y} and the encoded features π Β― 1 , β¦ , π Β― k \bar{\mathbf{c}}_{1},...,\bar{\mathbf{c}}_{k} of that layer at each of the k k tokens. Then, run the model on the base sentence up through the intervention layer, compute the encoded features π 1 , β¦ , π k \mathbf{c}_{1},...,\mathbf{c}_{k} at each token, and at each position replace the residual stream vector π± i \mathbf{x}_{i} with the patched vector
π± i β² = π± i + β j β F ( π Β― i , j β π i , j ) β π j \mathbf{x}^{\prime}_{i}=\mathbf{x}_{i}+\displaystyle\sum_{j\in F}(\bar{\mathbf{c}}_{i,j}-\mathbf{c}_{i,j})\mathbf{f}_{j}