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

Summary

This work describes two new type of sparse autoencoders (SAEs) that are trained with two new loss functions, L_{e2e} and L_{e2e + ds}. L_{e2e} loss penalizes the KL divergence of the original model vs. when the SAE is inserted, as opposed to existing SAE training techniques which instead optimize MSE with the original activations. L_{e2e + downstream} additionally seeks to minimize the MSE of downstream layers after the SAE is inserted vs. the original model. Both types of new SAE have a better L0 vs. CE tradeoff than naive SAEs, and e2e + downstream additionally has similar downstream reconstruction loss. Interpretability is no worse for the new SAEs. Finally, e2e+ds has similar geometric properties to naive SAEs, including discovering similar features.

Strengths

- The new SAE techniques that the authors introduce are a Pareto improvement over existing techniques on L0 vs. CE loss. Furthermore, the techniques are exciting because they is in a different direction from recent improvements in training SAEs: it is possible that the new techniques can be combined with recent techniques like gated SAEs and the top-k activation function. - The investigation of the geometry of the different SAE features is extremely interesting. Most notably, e2e + ds finds similar features to local and is reproducible over different seeds, and e2e has potentially less feature splitting. - The new techniques result in features of similar interpretability as local, as determined by standard automated interpretably tests.

Weaknesses

- The main evaluation metrics that show improvement (CE loss for both new SAEs and later layer reconstruction error for e2e + downstream) are the same ones that are being optimized. This is noted by the authors when they mention Goodhart's law, but it would greatly improve the paper if there were additional metrics for downstream SAE quality that were tested (e.g. some of the methods described in section 4.4). - Feature shrinkage (shown in Table 5) seems much worse with the new SAE training methods. Moreover, even with e2e + ds SAEs, later layer reconstructions in Figure 12 are much worse than local. This seems potentially a problem for e.g. finding circuits or understanding interactions between multiple SAEs via feature similarity. This is another reason I think additional downstream metrics would be useful.

Questions

- Building off the above weaknesses, could the authors add a measure of downstream SAE functional importance that is not CE loss? - 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. - 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. - 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. - 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? Small points: - 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. - 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. I would be happy to raise my score if some of these questions/concerns were addressed.

Rating

6

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

I believe the authors should more explicitly discuss the Goodhart's law problem I've discussed above and ideally the increased feature suppression as well. Otherwise, I think the limitations and impacts are well addressed.

Summary

This paper proposes to train Sparse Autoencoders (SAEs) with an additional loss function term: The KL divergence between the original output distribution and the output distribution obtained when using model with the inserted SAEs. This additional loss pushes the SAEs to focus on functionally important features The authors find that this setup requires fewer total features and fewer simultaneously active features per data point.

Strengths

- the authors point out a weakness of previous SEAs (feature splitting) and propose an intuitive solution: learning functionally important features by minimizing the KL divergence between original output and output with inserted SAEs - writing is mostly clear and easy to follow - the authors discuss limitations of their approach like robustness to random seeds, reconstruction loss, training time

Weaknesses

- I think the major weakness is that the contribution is not very groundbreaking. - There are no new findings wrt interpretability, and the e2e SAEs require more compute to be trained - while the authors discuss several findings (robustness to random seeds, reconstruction loss, training time) the reader is left waiting a bit for a conclusion or take home message - the comparison to local SAEs is sometimes a bit unclear: - 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" - the features found by local SAEs are not the functionally important ones , but SAElocal and SAEe2e+ds features are somewhat similar so local SAEs can be used as initializations (line 233)? - notation could be improved - N_dict_elements and d_hidden seem to belong more into pseudo code so it looks a bit odd to mix it with more traditional mathy notation - $\lambda$ is hidden in the $\phi$ parameter but then you show values for lambda in Table 1 and looking back at i.e. Eq. (1) it's not immediately clear what $\lambda$ is and why $\phi$ is not mentioned. I would suggest either putting $\lambda$ directly in the equations or citing values for $\phi$

Questions

line 50: but we no longer optimize activation reconstruction. -> but you do use activation reconstruction for $SAE_{e2e+ds}$, don't you? Maybe I just missed i but where in the network do you plug in the SAEs? the residual stream/decoder layer output? 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? 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. 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 - line 5: datatset -> dataset - line 44: a feature is functional important -> a feature is functionally important - line 113: ")" missing in equation

Rating

5

Confidence

2

Soundness

3

Presentation

3

Contribution

2

Limitations

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).

Summary

The authors present their observations on modified sparse autoencoder (SAE) learning. In the proposed method SAE is trained to reconstruct original model weights (features). SAE is optimized with the KL-divergence loss between the model output and the output at the same location in the original model. An additional variant is training with the reconstruction loss computed after each layer between the original model and SAE output. The sparsity loss term is added to increase the number of weights that will converge to zero. The authors claim that their method explains the model features in a better way with fewer dictionary elements.

Strengths

- Proposed a new training method for SAE with end-to-end loss that allows to use of fewer features from the original model - The authors provided extensive analysis of their results - The related work is sufficient - The proposed method decreases the number of active dictionary elements twice while preserving the same CE loss decrease as a baseline SAE.

Weaknesses

- 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) - 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 - The interpretability is not well defined, in addition, the claim that fewer features means more interpretability should be proved by empirical evidence, white the authors showed that interpretability was not changed or changed significantly (appendix A7). If I get it correctly, the main goal of the proposed method is interpretability but a few experiments were done to evaluate it. Also, qualitative results are not provided (anonymous url) - Std values are missing (for different random initialization seeds) in Figure 1 - The writing should be improved, Sec. 3.4 should be more self-contained - It is not clear which method e2e or e2e-ds  is the best performer overall?

Questions

- Can you please explain the diagram in Figure 1 why some blocks are of different shapes? - 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. - What is the number of parameters you train in SAE in total for all layers? - Again, if the main benefit of the method is interpretability, I do not understand how this work addresses this property and what is its main contribution in improving interpretability. - Are $W_e$ and $D$ the matrices defined for each layer? If they are - an appropriate indexing should be applied.

Rating

5

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

not discussed

Summary

This paper proposes a way to train sparse autoencoders (SAEs) that encourages them to learn features that are causally relevant to the model's output. This is done by replacing the usual SAE reconstruction objective - an L2 penalty between original activations and their reconstructions - by the KL divergence between the model's output distribution and the distribution when activations are replaced by their reconstructions. Additional loss terms encourage the reconstructions to lead to downstream layer activations similar to the original downstream layer activations. The sparsity penalty on feature activations is kept as in the "vanilla" SAE. The paper finds that this approach is a clear Pareto improvement over vanilla SAEs w.r.t. the tradeoff between sparsity of the feature activations (L0) and loss recovered when using reconstructions. An automatic interpretability pipeline is used to compare the features learned to vanilla SAE features. A statistically significant advantage for e2e+ds SAEs is found. The paper also attempts to avoid a potential failure mode of the approach, whereby the SAE may learn to exploit different pathways through the model in order to drive down the KL divergence to the original output distribution, because this is easier from an optimization point of view (as the reconstructions are no longer optimized to match the original activations at the layer where the SAE is applied). This is why one of the SAE variants considered includes an L2 penalty for downstream layer activations too. Results here show that using only the KL penalty in the loss produces very different downstream activations compared to adding these additional terms, suggesting that the

Strengths

- it is an important open problem whether SAEs trained solely using activations from some LLM layer will learn all causally important features for the LLM at this layer if trained without supervision from downstream activations. The paper makes some progress on this problem, suggesting that "vanilla" SAEs may struggle to learn such causally relevant features. - the methodology is careful and often considers alternative hypotheses or potential pitfalls in the analyses. - the paper is clearly written

Weaknesses

- given that the KL divergence to the output distribution is incorporated in the loss function, it is not that surprising that the methods in this paper have better values of the loss recovered metric (to their credit, the authors acknowledge this). To truly conclude the superiority of the suggested method for surfacing *individual* causally important features, it would be extremely helpful to have some external (to the KL metric) evaluation. - 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 key problem when not using the same-layer activations as a reconstruction target for the SAE is (as the authors describe) that the optimization may prefer to learn features that achieve a good KL divergence value eventually, but do so through "pathological" pathways in the model. To fix this, the authors encourage closeness with activations starting from the next layer up. However, such "pathological" pathways may exist in a single layer (this was shown e.g. here https://arxiv.org/abs/2311.17030). So it is unclear whether the problem has been fully overcome or rather restricted to a more narrow part of the model. Furthermore, even if reconstructions of downstream layers are similar to the true activations of these layers, this does not establish that the individual features themselves are not "pathological". Again, some additional, *per-feature* evaluation is needed to make the (strong) claims of the paper more believable. - out of the three kinds of SAEs considered - vanilla, e2e, e2e+ds - there seems to be a missing one implied by the others: an SAE that encourages faithful reconstructions only of the *next* layer activations, and does not involve the KL divergence loss. Would such an SAE lead to similar improvements? If so, this would be stronger evidence, as the KL divergence w.r.t. the final output distribution won't be a part of the loss.

Questions

- did you try the "next layer only" SAE variant described in the weaknesses? - how can we get a more fine-grained picture of how individual features change between the vanilla SAE and the e2e ones?

Rating

6

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

- I think the main limitation that I would have loved to see addressed more in the paper is that the evaluations are somewhat indirect w.r.t. the main claim of the paper. Even if in some average sense the e2e SAEs' reconstructions lead to better KL divergence, we still don't know how this plays out on the level of individual features.

Summary

The paper proposes new framework to improve the interpretability of large language models. Inspired by the work proposed in Anthropic's "", the proposed work replaces the original reconstruction loss in the sparse auto-encoder (SAE) with KL-div loss, and additionally add further constraint to minimize the error between reconstructed downstream activations and original activations in the LLM. Empirical results demonstrates the adavantage of the proposed new SAE_{e2e} architecture series.

Strengths

S1: The paper considers new formulation to further improve the interpretability of the large language model.

Weaknesses

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. W2: It is now clear how the KL divergence is implemented to replace the reconstruction loss. KL is only practically useful when the distributions are available. However, in this problem, it is unclear how the distribution loss between the two deterministic feature outputs are computed. 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.

Questions

I appreciate it if the authors could please see the above weakness for my questions.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes the authors have adequately discussed the limitations of the work.

Valuable, if negative, results

Thank you for the detailed and thoughtful rebuttal, as well as the overall response. I appreciate the additional evaluations of your proposed method. My interpretation of the additional experiments is that they provide at best mixed evidence for the superiority of the proposed SAE variants. That being said, I appreciate the honesty of the authors, and I believe these results will be valuable (and perhaps surprising) to the interpretability community, so I am in favor of disseminating these results widely. This is why I maintain my recommendation to accept this work (though I will not be changing my score at present).

Thanks for your rebuttal

Thanks for the rebuttal and the further clarifications on my questions. Given the explanations on the motivations behind using KL to replace L2, I decide to increase my score. I also encourage the authors to include the discussions on these differentiating factors between L2 and KL into the final version of their paper.