Identifying Functionally Important Features with End-to-End Sparse Dictionary Learning

Thank you for your evidently detailed read of our paper and your thoughtful comments.

Comments on Weaknesses

Regarding evaluation metrics, we've taken your criticism on board and run a set of evaluations on a set of downstream tasks. These are documented in the Author Rebuttal at the top (results in the attached pdf).

As for your observation about feature shrinkage appearing worse for e2e SAEs. We agree. It turns out you can shrink activations in the residual stream without changing the output much and e2e SAEs take advantage of this. Note that this is partially because of layernorm (as shown in Figure 12). Fortunately, one can address the feature shrinkage issue with some of the SAE variants that were released following our submission (and that you mention in the summary of your review), e.g. topk SAEs, Gated SAEs, JumpReLU. We agree with your summary that these improvements should be orthogonal to the improvements we find with e2e SAEs.

Comments on Questions

Why doesn't the original SAE learn the "left block" of features that SAE e2e + ds does? It would be great to see if e.g. these have a small or negative impact on the variance explained.

If you’re referring to what appears to be the left mode of the somewhat bimodal distribution in Figure 8, we could not find any meaningful similarities between these features unfortunately. Once the anonymity is lifted, we will be able to publish an interactive dashboard with these features.

Did the authors try training an SAE with KL divergence + the local reconstruction error (i.e. MSE of layer X reconstruction for an SAE trained on layer X)? It is a little concerning to me that the layer X variance explained is so strongly negative because it means the SAE's reconstruction is very far off distribution.

Please see the comment about variations of e2e+ds in the Author Rebuttal at the top.

Can you quantify how much farther the <e2e ds to local> difference is vs. the inter seed distance? It is hard to know how to interpret Appendix G without this context, as inter-seed SAEs have different features learned too.

Good point. As for the means of those distributions, the <e2e ds to local> cross-type similarity (Figure 3c bottom) has a mean of 0.73, and the <e2e ds> cross-seed similarity (Figure 3b bottom) has a mean of 0.78. Visually, we do see the mode to the left of the cross-type plots that doesn’t exist in cross-seed plots, indicating that there is a set of features which are learned by e2e+ds but not local (as mentioned in response to your earlier question, we weren’t able to interpret any common properties of this set of features). We’ve plotted a UMAP comparison for different seeds on layer 6. Unlike our cross-type UMAP plots in the paper (Figure 19), we don’t see clear regions where one seed is represented but the other isn’t (though, of course, reading deeply into UMAP plots is a source of potential confusion). We’ve attached the layer 6 cross-seed UMAP to the pdf in Figure 3 of the Author Rebuttal.

The ~0.15 CE loss e2e+downstream SAE looks closer to the other chosen SAEs than the 0.125 CE loss one, does it change things to run with the ~0.15 one?

Keen eye. The reason we didn’t select this run for our “similar CE loss” set was simply because it was run after our other analyses (including the time-consuming interpretations in Appendix G and dashboard creation) merely to fill out the pareto frontier for a nicer plot. After your comment, we reran the analysis (not the interpretations) with this run and did not find any significant differences. We posted a sample comparing the cross-type similarity with the different runs in Figure 2 of the uploaded pdf in the Author Rebuttal at the top.

Line 41 states "a significant amount of [causal effects of some features is] mediated by the reconstruction residual errors", but I believe the cited work shows that this is true for circuits as a whole, not features individually.

We agree that it could technically be possible that the circuit does have this property but that individual features in the circuit do not. Though we think it’s reasonable to extrapolate that the same is true for individual features. Nonetheless, we’ve updated the text to be more accurate, thanks for pointing this out.

Lines 133-136 are confusing. Why is it okay to "clear out" features from layer X, but not layers X + 1...? In general, a more clear discussion of the tradeoff between learning functionally important features and maintaining model behavior would help I think.

Our initial concern was that adding a reconstruction loss term at the current layer would have a much larger impact on reconstructing non-functional features than if the loss term was added at downstream layers. Though we have come to agree that this difference is merely quantitative rather than qualitative. As mentioned in the "Other variations of e2e+ds" section in the Author Rebuttal, we unfortunately have not run the extensive sweeps on the e2e+ds variations needed to give an insight to this question empirically.

Summary

Thanks again for your valuable comments. Hopefully we've addressed your concerns/questions, and are keen to hear if there are cases where we have not.

Thanks for your comments and your typo spotting and notation improvement suggestions!

Addressing main misunderstanding

After reading your review, it appears we failed to convey the core benefit of our method. You state

the authors state that "locally trained SAEs are capturing information about dataset structure that is not maximally useful for explaining the algorithm implemented by the network." but do not show that their e2e SAEs are significantly more interpretable or helpful at explaining the algorithm implemented by the network, instead the e2e SAEs are merely "at least as interpretable as SAElocal"

Indeed, we did not find that individual features found by e2e SAEs were more interpretable than local SAEs. Though we didn't expect them to be, nor was it the purpose of this work. The important contribution of this work is that in order to explain the activations at a particular point in the network, our method requires fewer than half of the features that other methods require. We illustrate this with the pareto curves in Figure 1 (left). This dramatically reduces the description length of any feature or circuit that uses SAEs. Shortness of description is an essential aspect of interpretability.

In addition, we show that we need far fewer features in total to explain the network activations over the whole dataset (Figure 1 (right)).

Response to Limitation

Authors briefly discuss the lack of a ground truth to evaluate SAEs. I think discussion of evaluation metrics could be a bit more divided (maybe into metrics to assess interpretability and metrics to assess how close to the original network a network with inserted SAEs is).

We do have a section-level separation between both types of metrics (Sections 3.1-3.3 for the latter type and Section 3.4 for the former), and we link to the sections alongside a short explainer of them at the bottom of the introduction (lines 75-86). We’re open to alternative division if preferred.

In addition, we have run more experiments in response to other reviews that evaluate the SAE setups on a set of downstream tasks.

Other responses to Weaknesses and Questions

line 50: but we no longer optimize activation reconstruction. -> but you do use activation reconstruction for SAEe2e+ds, don't you?

Yes. We train for downstream activation reconstruction SAEe2e+ds, but at all subsequent layers, rather than at the current layer. In this paragraph we are introducing SAEe2e, which does not involve any training for activation reconstruction.

Maybe I just missed i but where in the network do you plug in the SAEs? the residual stream/decoder layer output?

All of our SAEs are inserted in the residual stream (lines 160). We'll add this to the main figure and table for further clarity.

You say that taking different computational pathways through subsequent layers of the network might be a problem, but do you have any evidence for this happening? is this just section 3.3.2?

Primarily the evidence for this is in Section 3.2, especially in Figure 2, as well as in Section 3.3.2. Figure 2 directly shows the e2e SAE’s output activations taking a very different pathway through the network, as the activation reconstruction MSE is much higher and significantly increases downstream through the layers of the network, rather than converging. Section 3.3.2 provides additional evidence that the pathway through the network is under-constrained by the KL divergence loss of e2e SAEs alone. These problems are resolved by e2e+ds SAE training.

While reading, I was wondering if taking different computational pathways through subsequent layers of the network is still a problem if only one SAE is trained at a time. There should be little incentive for computing additional information since the original network post SAE insertion would not have capacity to handle these anyway, so it would only matter for the output directly.

Even with one SAE trained at a time, it turns out that there is capacity for the original network to handle the SAE adding "new" features to the residual stream. Evidence for this can be seen in Figure 3b (middle), which shows that with two different seeds, a pure e2e SAE can learn a very different set of features (and still perform as well in the L0-CE loss pareto). While we're unsure how much of a problem this is, our e2e+ds formulation appears to resolve this issue at very little cost.

small things: line 68: not sure if it's necessary to mention goodhart's law since the target is not cross-entropy loss difference but interpretability so I would be more worried about this if you start optimizing interpretability measures

This is reasonable, and is our argument too. Though it's worth mentioning that reviewer 2 and reviewer 5 both expressed concerns about Goodharting.

[notation improvement suggestions]
line 5: datatset -> dataset.
line 44: a feature is functional important -> a feature is functionally important.
line 113: ")" missing in equation

Thanks for spotting these errors and notation ideas! We'll incorporate these.

Summary

Thanks again for your comments and questions. We hope to have adequately clarified any uncertainties or misunderstandings w.r.t the core contributions of this work and other listed limitations or weaknesses. If so, we would ask you to reassess your review score, taking this response into account.

Thanks for you comments and suggestions

Comments on Weaknesses

The interpretability is not well defined...

The interpretability provided by an SAE contains two main components: The average number of SAE features required to interpret any particular model output (L0 - lower is better), and the individual interpretability of each of those features. Improving either one of these without hurting the other or the amount of model performance explained constitutes an improvement to interpretability by shortening the description length required to explain the model’s behavior. Shortness of description is an essential component of interpretability.

As we show in Figure 1, e2e and e2e+ds SAEs are able to significantly reduce L0, requiring fewer active features compared to local SAEs explaining the same amount of model performance (as measured by the CE loss increase).

While the interpretability of each feature is in general somewhat subjective and hard to measure, we put significant effort into gaining unbiased and rigorous measures by implementing automated interpretability scoring (Bills 2023).

Overall, we believe the evidence is clear that our method produces SAEs which are significantly more useful for interpretability than SAEs produced by the baseline local training process.

The numeric results in Figure 1 are very incremental: e.g. reduction of 0.05 in CE loss increase L0=100 is ~1.6% of the evaluation CE loss.

We believe that this is an inappropriate baseline. At the L0=100 position quoted, our method more than halves the CE loss increase, relative to the local SAE baseline. We would not call a halving of the error incremental. Recall that we are not trying to reduce the CE loss of the original model - merely to explain as much of the original model’s performance as possible.

Second, although an increase of 0.05 in CE loss may seem relatively small, it corresponds to a much larger effective drop in model size or training FLOP via scaling laws between these quantities and loss. The drop in CE loss from the model learning more complex behaviors (which emerge at larger model sizes or training FLOP) generally reduces as the behaviors get more complex (relative to the drop in loss from the model learning simple bigrams or trigrams), and yet we are still especially interested in interpreting how the model performs these more complex behaviors.

Finally, the quoted rises in CE loss are for one single SAE. in practice, an explanation of the model’s behavior will require SAEs in many layers throughout the model, so seemingly small errors can be multiplied in this way, leading to unfaithful descriptions (this is illustrated in Marks2024 (https://arxiv.org/abs/2403.19647)).

The modification is very simple and raises the question of whether additional variants could work too, e.g. local SAE with MSE end-to-end loss (in addition to the layer reconstruction MSE)

Please see our response to reviewer TS92, where we agree that this would be an interesting variant to try, along with a long list of others.

Std values are missing (for different random initialization seeds) in Figure 1

We provide some evidence in Figure 18 that the results of Figure 1 are robust to different random initialization seeds. Unfortunately, training enough SAEs to get robust estimates of the standard deviations for each datapoint would be prohibitively expensive.

It is not clear which method e2e or e2e-ds is the best performer overall?

We recommend e2e-ds as the best performer overall. It achieves a similar performance explained to e2e SAEs while maintaining activations that follow similar pathways through later layers compared to the original model (highlighted in section 3.2 and the conclusion).

Responses to Questions

Can you please explain the diagram in Figure 1 why some blocks are of different shapes?

The different shapes are to distinguish qualitatively different operations, such as the unembed (which projects up from the residual stream to the logits), and the SAE (which we train). These are labelled in the diagram.

It looks like that the Sec. 3.4 result is trivial, if we don’t encourage the SAE to reconstruct after wach activation, that will be the result.

We believe you may have confused Section 3.4 with Section 3.3.2, or some other section. If so, Section 3.3.2 was not intended to be counterintuitive. If not, see our response to questions on evaluating interpretability above.

Are W_e and D the matrices defined for each layer? If they are - an appropriate indexing should be applied.

We do only ever train one SAE at a time, but we will add some layer indices for clarity, thanks for bringing this to our attention.

What is the number of parameters you train in SAE in total for all layers?

As mentioned on line 166, the number of dictionary elements in each SAE is fixed at 60 times the size of the GPT-2 small residual stream (60 x 768 = 46080 dictionary elements), so W_e and D both have sizes 46080 x 768, b_e is a 46080 dimensional vector, and b_d is a 768 dimensional vector, meaning that each SAE has just over 70 million parameters.

Summary

We thank you again for your comments and questions. If we have adequately addressed your concerns, we would kindly ask you to reassess your review score, taking this rebuttal into account.

Thank you for engaging with our paper and giving well-considered feedback.

Regarding the main limitation you’ve mentioned, we agree that it would be helpful for the paper’s narrative to have other metrics of functional importance than KLDiv, including evaluations on an individual feature level. However, there is an issue with this in practice, which relates to your comment

To some extent the paper shows such evaluations. For example, in Appendix G.3 it is shown that vanilla SAEs represent well a feature that is not causally important, and the e2e SAEs in turn represent it poorly. However, what is really required to establish superiority is the opposite: to exhibit a causally important feature not represented by vanilla SAEs, but represented by e2e ones.

The extent to which features exist in one dictionary and do not exist in another is, unfortunately, not a binary concept in practice. We often find that there are directionally similar, but not identical features shared across dictionaries. Or there are clusters of features that are overrepresented in one dictionary and underrepresented (but not absent entirely) in another. Therefore demonstrating that “a causally important feature not represented by vanilla SAEs, but represented by e2e ones” may not reflect the metric that matters overall: the global functional importance of all the features.

Nevertheless, we agree that it would be desirable to be able to identify compelling narratives about individual features across dictionaries and their improved functional importance. Unfortunately the differences are matters of degree – quantitative rather than qualitative, which makes qualitative stories more difficult to tell. We were able to identify one such qualitative story. However, despite efforts to identify qualitative stories in the direction suggested by the reviewer (“what is really required to establish superiority is the opposite: to exhibit a causally important feature not represented by vanilla SAEs, but represented by e2e ones.”), we found that the data did not suggest these qualitative narratives despite the quantitative differences being present on the global level.

As for other external evaluations that don’t directly measure the CE loss over the same distribution used to train the SAE, we have run some additional experiments on a selection of downstream tasks. These are described in the Author Rebuttal at the top of this page

When not using same-layer reconstruction... it is unclear whether the problem has been fully overcome or rather restricted to a more narrow part of the model.

We agree that our existing e2e+ds implementation does not eliminate the problem but rather restricts it to a narrower part of the model. Our initial concern with adding a reconstruction loss term at the current layer was that it would reconstruct too many non-functional features present in the residual stream at the current layer. Though we have come to agree that the difference between reconstructing the current layer and a subsequent layer is merely quantitative rather than qualitative. As mentioned in the "Other variations of e2e+ds" section in the Author Rebuttal, we unfortunately have not run the extensive sweeps on the e2e+ds variations needed to give an insight to this question empirically, so it is possible that additionally reconstructing the current layer's activations would perform just as well on the L0-CE Loss Pareto while further restricting the computational pathways available to the model.

Regarding the question

did you try the "next layer only" SAE variant described in the weaknesses?

please see our comment on this in our Author Rebuttal at the top of the page.

We thank the reviewer for their comments and questions. We believe there may be a core misunderstand of our method which we clarify in the response to Weakness 2.

Response to Weaknesses

Weakness 1

W1: The authors may further explain the intuition between the evaluation metrics. The current version of the evaluations is not explicitly clear to me why the proposed method is advantageous in comparison to the baselines.

In the paper we present three main components to evaluating the performance of an SAE. Below we give some more intuition behind them.

1. If we are to sum up dictionary elements of our SAE and use those in place of the real model activations, how many dictionary elements do we need to use to get similar performance to the real model? Figure 1 (Left) shows that we require fewer dictionary elements on average (L0) to explain the same amount of CE loss degradation (which is the difference between the CE loss of the original model and the model when splicing in SAE features).
2. In Figure 1 (right), we measure the total number of dictionary elements needed in the SAE to explain the network’s behavior over the whole dataset. While less important than the metric above, having fewer dictionary elements in total means that there are fewer variables needed to describe the behavior of the network on your distribution. Note that “alive dictionary elements” is the effective measure of the number of variables needed, as some dictionary elements become forever inactive during training (it’s worth noting that the recent work of Gao2024 (https://cdn.openai.com/papers/sparse-autoencoders.pdf) has reduced the killing of dictionary elements with better initialization and an augmented loss term, though we don’t have a reason to expect that employing these techniques will change the structure of the pareto curves presented in our figure.
3. We want to know whether we can understand each SAE dictionary element individually. We implemented automated interpretability scoring, which provides a scalable and unbiased quantitative measure of interpretability. For an introduction to this technique, see Bills et al. [2023] https://openaipublic.blob.core.windows.net/neuron-explainer/paper/index.html. In short, to score the interpretability of a single SAE feature, GPT-4-turbo is prompted to write a short english-language explanation of the feature, then GPT-3.5-turbo is prompted to try to predict the specific activations of the SAE feature on each token, as we explain in section 3.4. From the automated interpretability scoring, we find strong evidence that the features in our e2e+ds SAEs are not less interpretable than the features in the local SAEs at a similar amount of model performance explained. We even find that they are significantly more interpretable in some layers, such as in layers 2 (p = 0.0053) and 6 (p = 0.0005).

It may be useful to note that several other works in the area have used the metric outlined in point 1 to evaluate SAEs e.g. Figure 5 in Bills2024 (https://cdn.openai.com/papers/sparse-autoencoders.pdf), Rajamanoharan2024 (https://arxiv.org/abs/2404.16014)(note that this uses a “loss recovered” term which is calculated from the average CE loss of the model when splicing in the SAE), Kissane2024 (https://arxiv.org/abs/2406.17759). Additionally, many works have used automated interpretability to evaluate the interpretability of SAE dictionary elements (e.g. Cunningham2023 (https://arxiv.org/abs/2309.08600), Templeton2024 (https://transformer-circuits.pub/2024/scaling-monosemanticity/),

In addition, in the Author Rebuttal at the top, we've run additional evaluations. The results measure how much task performance degradation there is when SAE dictionary elements are spliced into the model on some of the subtasks used in Marks2024 (https://arxiv.org/abs/2403.19647).

Weakness 2

However, in this problem, it is unclear how the distribution loss between the two deterministic feature outputs are computed.

The KL divergence is computed for the output of the LLM. This output is a probability distribution over next tokens, motivating our use of KL divergence.

The motivation of the replacement between KL divergence and the L2 reconstruction is also not very clear to me, and these two losses should be ultimately equivalent and optimal at the same parameters global optima, although the loss changes the optimization landscapes.

For an SAE with perfect reconstruction, both L2 and KL loss would be zero. However, the sparsity_loss and reconstruction_loss terms trade off against one another (given a fixed dictionary size). This means that the global optimum does not result in an SAE with perfect reconstruction.

We use KL divergence rather than L2 reconstruction at the output as the outputs are probability distributions (and the original network was trained with the standard language modeling loss which is the KL divergence between its outputs and the one-hot labels). For SAE_local and SAE_e2e+ds, we have a loss term at intermediate layer(s) in the network. Here, we use the L2 reconstruction as these activations are not probability distributions.

Summary

We thank the reviewer again for their comments and questions, and kindly ask that they reassess the review score if their concerns were adequately addressed by this response.