Paper: Sparse Autoencoders Find Highly Interpretable Features in Language Models (ICLR-2024)

Reviewer comment (Experiments):
The paper claims in the abstract and conclusion that the results show *greater* monosemanticity than other methods but I do not see such comparisons in the paper (it should be in section 5.1). Either I missed something or one of the main claims is not supported.

Author response:
Our claim that we have greater monosemanticity rests on our comparisons in section 3 to the neuron basis, random directions, and directions found by PCA and ICA.

Reviewer comment (Novelty):
Most of the pieces in this paper (the sparse autoencoder and the interpretability score) are pre-existing, and so is the idea and demonstration that sparse features yield greater interpretability. This makes the overall level of originality and significance of the paper pretty low.

Author response:
On the originality of the approach, while we agree that none of the individual elements is novel on its own, the pipeline of using a sparse autoencoder to decompose activations in a large model (section 2), which are then passed to an automatic interpretation protocol (section 3), and then analysed in terms of the circuits that build up later features (section 5) represents a meaningful step in our ability to peer into the inner workings of language models.

Reviewer comment (Writing):
In the first paragraph of sec. 5.1, the authors talk about 'real-word features'. Please clarify what that means. If it is a "human-understable concept", one crucial problem is that such a concept (except for single-token cases studied in the paper) can be represented in many ways as a sequence of words, making any practical extension of the proposed idea (in sec 5.1) not obvious at all (or maybe not even feasible).

Author response:
We appreciate the point about ‘real-world features’ lacking clarity, and we’ve edited the text to amend this.

Reviewer comment (Writing):
What is R? It is an important part of many text captions such as Figure 2 and Figure 3 but never explained in main text. Only explained in Appendix C.

Author response:
> What is R?

R is defined in the second paragraph of section 2: “The autoencoder is a neural network with a single hidden layer of size d_hid = R d_in, where d_in is the dimension of the language model internal activation vectors1, and R is a hyperparameter that controls the ratio of the feature dictionary size to the model dimension.”

Reviewer comment (Reproducibility):
Almost all the results do not mention the details of how they were achieved. In particular, table 1, and figures 2, 4, 5 don't mention the model used at all, and most figures do not mention R and alpha used to train the in question.

Author response:
We appreciate your call for more clarity on how results have been achieved! All tables and figures now indicate which model and hyperparameter values were used.

Reviewer comment (Presentation):
In general the main text needs more discussion of the basic methods and results like reconstruction loss, and some of the content from Appendix B and C should be moved to main text.

Author response:
We think the point that some of the material in Appendices B and C should be in the main text is an interesting one. We’ve brought some of Appendix B forward into a new Figure 2 which hopefully gives the reader better insight into the nature of the models that we trained. While the material in Appendix C is interesting, we feel that the fact that the learned bases are under-complete in many cases represents a significant drawback to them that wouldn’t be appropriate to show in the main text, and which is best shown in the appendix as a very enticing but ultimately incomplete direction for future work to build on. (We also value the range of experiments and don’t have pages to spare!)

Reviewer comment (Presentation):
Figure 2: Is identity ReLU the default basis?

Author response:
> Figure 2: Is identity ReLU the default basis? 

Yes

Reviewer comment (Theory):
Section 3.2: Why is ReLU applied to the default and random basis? It is said it’s applied to make activations nonnegative, but why is this necessary? The actual model also uses negative activations, right?

Author response:
> Section 3.2: Why is ReLU applied to the default and random basis? It is said its applied to make activations nonnegative, but why is this necessary? The actual model also uses negative activations right? 

We restricted all activations to be positive here because we were concerned that negative activations were a source of additional variance that would bias the results in favour of our method. The interpretation score is based on a correlation metric and thus has variance in the denominator. If an explanation could explain the same amount of variance in the positive domain for one of our features and one default basis features, the resulting correlation score would be higher for our feature, thus biasing our results. In our early experiments we found that autointerpretability scores were similar with or without ReLU. It is likely that some features are bidirectional and so would be best explained by allowing a feature to have both positive and negative activations, but this would be penalized equally for our features and default/random/PCA features.

Reviewer comment (Evaluation):
What is the intended way to use these learned features? Adding the autoencoder changes the network behavior, and increases computational cost. Are you supposed to use the expanded network in practice to make explainable decisions, or is this better thought of as an approximation to explain decisions made by the unmodified network?

Author response:
> What is the intended way to use these learned features? Adding the autoencoder changes the network behavior, and increases computational cost. Are you supposed to use the expanded network in practice to make explainable decisions, or is this better thought as an approximation to explain decisions made by the unmodified network. 

The latter is definitely more what we have in mind. A core principle of this research is to find interpretability techniques that don’t come at a performance cost - we want to be able to run the model as normal and use this ‘featurized’ version of a layer to understand what the model is doing as best we can, hopefully eventually being able to detect behaviour that’s deceptive or otherwise not in line with the developer’s intentions. We also think that the utility of the current technique is as a platform for further work in understanding how models compute their outputs - we think we should eventually be able to learn better features, and also to build them up into more interesting units like circuits for particular behaviours.

Another hypothesis, though as yet unproven, is that they may be useful for model editing. For example, model editing often involves learning a vector to add to a residual stream, but finding the right direction across dozens of 10K+ dimensional spaces requires quite a bit of data and these features could provide a useful prior for productive directions to edit.

Reviewer comment (Evaluation):
Overly complex approach, where the auto-encoder will likely add noise to the results

Author response:
We’re sorry that you didn’t find the paper of value though we hope that the other reviews which recommend acceptance suggest that there’s content which is informative for others.

We agree that interpretability is a difficult area to find real benchmarks in, and the benchmarks are not as robust as would be expected in most areas of ML, though through the use of automatic interpretability we have tried hard to give performance measurements which have clear definitions and can be replicated. We note that Bills et al represents a core push by OpenAI to understand their models and so being able to find primitives which are advances according to their own protocols represents a significant step among those working in this area.

We disagree that the approach is overly complex - to us it is both simple and directly motivated by the theoretical considerations in Toy Modes of Superposition. The computational simplicity of the linear autoencoder over iterative methods for sparse coding (which you may be implicitly referring to as the less complex approach?) is likely to be very important in scaling this technique up to cutting edge models.

Reviewer comment (Experiments):
No real baseline comparison

Author response:
Yes, a collaborator recently completed a comparison experiment between a sparse autoencoder and the FISTA method (as used in Yun et al https://arxiv.org/abs/2103.15949). In this comparison experiment, they trained 8 dictionaries on 512-dimensional data, with a feature ratio R=1, on a 1050 Ti. The results: autoencoders completed training in ~15 minutes, compared to ~10 hours for the FISTA dictionaries. In this experiment at least, the autoencoder was 40x faster than FISTA.

Reviewer comment (Evaluation):
The reconstruction loss is not fully minimized, indicating the dictionaries do not capture all information.

Author response:
We agree that the autoencoder’s imperfect reconstruction is a drawback, but getting much higher auto-interpretation scores still marks this out as a significant step towards being able to understand a layer in a large language model.

Reviewer comment (Evaluation):
The method works less well for later model layers and MLPs compared to early residual layers.

Author response:
These are good questions, and while we’ll describe our best current theories, we definitely don’t know for sure. This uncertainty is a big reason why these results are included as an appendix and we’d be excited to see future work that shines more light on this.

Our the three best answers to your questions are:
- There just isn’t a good overcomplete basis in these layers and we need a different framework to make sense of these later layers.
- There is a good overcomplete basis, but the geometry of MLP layers makes learning it hard, which we need to reflect in our training dynamics somehow.
- There is a good overcomplete basis and we just need to train for much much longer.

We’ll go through these in more detail. 

Firstly, as you note, we did find overcomplete bases in the early layers - it’s the later layers which are the problem. This mirrors the fact that improvement in reconstruction loss in the residual stream from doubling the width is much higher in the early layers, so it may be the case that the typical geometry of those layers allows them to better be captured by a smaller number of ‘features’ and/or makes the long tail of features harder to find. This doesn’t mean that there can’t be useful information from learning a high quality overcomplete basis, just that it seems there’s something fundamental about later layers that makes this more difficult. Also, the output of the non-linearity in a GPT-style transformer is 4x wider than the residual stream, but is only a relatively simple transformation of the residual stream, so we might expect to need a smaller ratio R in the MLP. Still, since the reconstruction loss is significant here, we should not be complacent and must continue to look for ways to understand these layers. Perhaps more complicated encoders are needed and we just failed to get them to work, or perhaps we need something more flexible than discrete linear features.

The second possibility is that the MLP’s geometry interferes with our training process. In particular, MLP activations are mostly in the positive orthant of activation space, leading to more dead neurons in the training process. We tried a couple alternative approaches to overcome this but they were unsuccessful. In particular, we tested our technique on the MLP layer with: 1) centered data, and 2) adding the minimum of the GELU function to all activations and then enforcing only positive entries in the weight matrix. It’s still possible that this is a major cause but our attempts to address it weren’t successful.

Our final hypothesis is that we simply need to train for orders of magnitude longer, potentially in conjunction with reinitializing dead neurons. In this answer we take inspiration from a recent Anthropic paper (https://transformer-circuits.pub/2023/monosemantic-features/) which used a technique similar to ours on the MLP layer of a one-layer transformer. They used hundreds of times more examples in training, along with reinitialization, and found that they were able to learn a very highly overcomplete basis. Since their model had 512 MLP dimensions rather than our 2048, this suggests we should try increasing our training time by perhaps 1000x. However, the fact that their language model has only a single layer leaves some uncertainty about whether these techniques assist here: is that single layer most analogous to our layer 0 (where our existing technique work) or our layer 5 (where problems arise)? Finally, in our experiments we chose not to extend our training runs because the marginal improvement seemed to be minimal. In particular, we found that at the upper end of our training time, the number of commonly active features was still increasing, but at a very slow rate such that it would not meaningfully change with another doubling of training time, so we didn’t push further.

Reviewer comment (Experiments):
There is limited comparison to other interpretability and disentanglement techniques. For example, beta-VAE, infoGAN, FactorVAE had similar disentanglement goals. Michaud 2023, "The Quantization Model of Neural Scaling," suggested a spectral clustering approach for identifying "monogenic" signals in large models. The paper would be strengthened if it provided a more complete comparison to previous disentanglement approaches.

Author response:
We think that betaVAEs are interesting points of comparison though their motivation and structure are sufficiently different that they can’t be directly compared to our approach.

Reviewer comment (Novelty):
The authors could mention previous works in the field of sparse autoencoding and dictionary learning for word embeddings. Subramanian et al. (2018) similarly found linear factors for word embeddings using a sparse autoencoder. Zhang et al. (2019) solved a similar problem using methods from dictionary learning, discovering elementary structures beyond existing word vectors.

Author response:
The work from Subramanian et al and others working on word embeddings is indeed an important precursor to our work that we weren’t aware of, and we really appreciate these comments and have updated the related works section accordingly.

Reviewer comment (Experiments):
How well do the features transfer across different models and architectures?

Author response:
> How well do the features transfer across different models and architectures? 

We did not test this, but there is no a-priori reason to expect these features to transfer, since they are just linear directions and a model’s residual stream is (in theory) linear, so therefore two models will have “arbitrarily transformed” residual streams relative to each other.

Reviewer comment (Experiments):
Could you incorporate model weights or adjacent layer features to improve dictionary learning?

Author response:
> Could you incorporate model weights or adjacent layer features to improve dictionary learning? 

While we have not attempted this, a recent work (https://www.lesswrong.com/posts/YJpMgi7HJuHwXTkjk/taking-features-out-of-superposition-with-sparse) has explored using model information to initialize the dictionaries to improve training speeds.

Reviewer comment (Experiments):
The authors claim the method is scalable, but have not demonstrated it on very large models. Experiments on models with billions of parameters could better support the scalability claims.

Author response:
We think all of your questions are interesting points and look a lot like our plans for future work, especially incorporating model weights and trying harder to learn overcomplete bases in the residual stream (we suspect much longer training times and reinitialization are likely both to be key). As for scalability, we note that our subject models, Pythia70M and less often Pythia410M are larger than other examples, with the largest being Yun et al using a BERT model with ~100M parameters. The fact that we use autoencoders rather than iterative methods such as FISTA makes scaling significantly easier as only a single pass is needed to compute the coefficients. Nonetheless, the compute required is quadratic in the width of the layer (and possibly even larger if larger ratios and longer training runs are required), so scaling to multi-billion parameters while maintaining quality models requires engineering beyond the scope of this paper.

Reviewer comment (Evaluation):
Auto-interpretation score hasn't been shown to correlate with actual interpretability use cases yet.

Author response:
We agree that it would be very valuable to see more work tying automatic interpretation to more concrete results and we acknowledge some of the limitations in the auto-interpretation section. We still think that its underlying method (measuring interpretability by the correlation between the actual activation of a neuron or feature and the activation as estimated from a potential description of the meaning of that feature) is well-founded, and we hope that by giving people better primitives than neurons to work with, we’ll be able to see more foundational work in this area soon.

Reviewer comment (Experiments):
Activation patching experiment was only conducted on one relatively synthetic dataset.

Author response:
While the activation patching experiment was only conducted on a single, synthetic dataset, the IOI task that we chose is the most studied example of an LLM circuit that we’re aware of. It was therefore the first thing that we tried, so we think this demonstrates the practicality of the method, as you mention.

Reviewer comment (Experiments):
Any comments on how the proposed method scales with complexity of the underlying task / number of possible features?

Author response:
One of the big benefits of the approach is that it’s task agnostic - ideally a dictionary should contain all the relevant features - we certainly can’t claim that we reached this point we found all of them, if that’s even a well defined notion, but we don’t see any fundamental blocker to capturing features relevant to any task the model is capable of in a single training run.  
The method scales somewhere above linearly with the number of features - if you double the width of the autoencoder, you’ve twice as many parameters and there’s some additional slowness to converge though we didn’t try to quantify exactly how much extra time is needed.

Reviewer comment (Experiments):
Is the approach bottlenecked by auto-interpretation method?

Author response:
At the scale of this paper (Pythia70M, ~100M datapoints of training) it’s easier to train lots of autoencoders than to evaluate lots of them via automatic interpretation - a single feature evaluated using the procedure from Bills et al cost around £0.20 so getting a good sample size of around 200 features interpreted cost about £40 in August 2023 OpenAI credits, significantly more than the training cost per encoder (and meant that we couldn’t run automatic interpretation on most of the dictionaries we trained). Training the autoencoder is more than quadratic in the width of the layer being examined though, but interpretation is unaffected by layer width, so at larger scales the cost of the autoencoder will dominate.