Enhancing Neural Network Interpretability with Feature-Aligned Sparse Autoencoders

Thank you for your insightful review.

“This paper is about enhancing the interpretability of an internal representation of a neural network. However, there are quantitative values compared with baseline models, but there are no interpretation-related results or explanations in the experiments.”

We acknowledge the absence of a verification that SAEs trained with MFR permit better downstream analysis. We would commit to this experiment for a camera-ready revision:

1. Collect data on the activation of SAE features on tokens for the GPT-2 SAEs
2. Task humans with explaining the features based on the data from (1)
3. Task humans with predicting the activations of SAE features on unseen tokens given the explanations in (2), where more accurate predictions of the activations on unseen tokens indicate more interpretable features

At the reviewers request we could perform this experiment using a language model such as in Bills et al., however as mentioned in Section 2.2, recent work has brought this evaluation into question, showing a lack of causal relevance of language model feature descriptions. Ultimately we feel skeptical that there is a good method of evaluation for the interpretability of SAE features.

“When training SAEs, the hidden size plays an important role, which may cause the feature superposition or feature distribution to occur. In this respect, there should be a comparison among varying hidden sizes in SAEs.”

For the SAEs we trained on the synthetic dataset, we already know the true input feature count, so we can optimally set the SAE hidden size without needing to search over hidden sizes.

“The current experimental results do not support many hypothetical statements well in Section 3.3. The validity of their arguments needs rigorous efforts.”

We politely request which statements the reviewer feels are not well supported. We feel that we have supported the primary claims of the section. We view these as:

1. Showing the correlation between a feature being learned by multiple SAEs and it being an input feature (we show this correlation in Figure 2).
2. Features that are similar across SAEs but not with the input features can often be detected by how active they are (Figure 4).
3. By avoiding the features mentioned in point (2) and encouraging similarity in the decoder features we learn more of the input features (Figure 8).

Per your questions:

1. This is the number of active features per data point we specified when generating the synthetic dataset. Because the data is synthetic, we are able to control this and manually determine the active features per data point.
2. MFR is the combination of our reinitialization technique and auxiliary penalty. We define the auxiliary penalty from lines 312-314, and detail the reinitialization from lines 300-309 of the revised PDF.

Thank you for your time in reviewing our submission. As the rebuttal period is nearly over, we wanted to kindly follow up to see if our rebuttal addressed some of the points in your review, or if there were any further questions we could answer.

Thank you for your insightful review.

1. By “features that are not features of the input”, we do not mean with respect to human interpretability, we mean that the features are fundamentally different from all features of the input. We do not claim that polytopes are the fundamental units of neural networks. Previous literature has chosen to study a particular division of the neural network (e.g., polytopes) and treat the neural network as being composed of that unit, which is what the quote in question is describing.
2. We are concerned that bounding the sparse feature coefficients away from zero could reduce how realistic the sparsity is, and bias the dataset toward over-representing less significant features. We feel that doing this changes the task significantly in a way that makes it less suitable for a sparse autoencoder.
3. It is concerning to us to have a cluster of features similar between sparse autoencoders that was not similar to the input features even if those features are inactive because (1) These features are occupying space in the decoder weights that could be replaced with features more similar to the input features, and (2) It makes the expected gains of encouraging similarity in the decoder weight matrices of multiple sparse autoencoders less, as the correlation between the similarity with the input features and with the features of other sparse autoencoders weaker.
4. We suggest that the asymmetry of the auxiliary penalty is not inherently problematic. At the reviewers request we could experiment with a symmetric penalty (e.g. by modifying the MMCS calculation to be an averaged sum of the MMCS of the first decoder features to the second, and then the second decoder features to the first) in a camera-ready revision.
5. We believe the threshold in Figure 8 is appropriate because it refers to the probability a feature is in the top-K, not to its activation value. If all features were equally likely to be in the top-K, then they would have an activation probability of k/N, where N is the number of features.
6. The initial reconstruction loss is the reconstruction loss after the first batch. The "100 training step cosine warmup for the penalty" refers to a cosine warmup on the coefficient for the auxiliary penalty over the first 100 batches.
7. Because the reconstruction loss is an imperfect measure of how much information the decoder weights preserve about the input features, we expect that the reconstruction losses may not match up perfectly with the performance on downstream applications of the autoencoders.
8. Yes, we vectorize multichannel windows. Each window represents one second of EEG data, and so represents 250 (the sampling rate) samples per channel per window.
9. Please see the updated PDF with a revised Figure 10 that should be much more readable.

Thank you for your insightful review.

We suggest that DCCAE and MFR are not directly comparable. DCCAE leverages multiple input views, but MFR uses a single view as the input to all autoencoders. For example, in Section 2 of Wang, Weiran, et al. (2015) they refer to “paired observations from two views”. In MFR, we essentially create different views of the input by varying the sparsity constraint on the hidden state. It’s not clear to us how we would apply DCCAE to our problem setting, or MFR to the DCCAE setting, as our understanding is that they assume different problem settings (multiple VS a single input view).

We believe that the baseline used in our paper is the most appropriate baseline, as it is a common architecture for autoencoders trained on language model activations.

We agree that the requested qualitative analysis of latents to gauge the positive effect of improved reconstruction loss on interpretability is an important experiment, and commit to having the following experiment complete for a camera-ready revision:

1. Collect data on the activation of SAE features on tokens for the GPT-2 SAEs
2. Task humans with explaining the features based on the data from (1)
3. Task humans with predicting the activations of SAE features on unseen tokens given the explanations in (2), where more accurate predictions of the activations on unseen tokens indicate more interpretable features

At the reviewers request we could perform this experiment using a language model such as in Bills et al., however as mentioned in Section 2.2, recent work has brought this evaluation into question, showing a lack of causal relevance of language model feature descriptions. Ultimately we feel skeptical that there is a good method of evaluation for the interpretability of SAE features.

Per your questions:

1. We have ensured there are no overlapping subjects for the EEG data.
2. We intend to lose feature diversity across multiple SAEs, as we encourage the learned features to be similar. However, we do not expect to lose feature diversity in a given SAE.

Thank you for revising your score, we are glad we were able to provide some clarification.

We acknowledge that point 1 is important to address. To ensure this is addressed in a camera-ready revision, we would train many probes on the activations of GPT2-Small, each to detect the activation of a particular feature of the input (this would involve estimating a large quantity of features learned by GPT2-Small). We would then train probes to detect the same features in the outputs of SAEs trained with and without MFR. If SAEs trained with MFR learn more features of the input, more of the probes should successfully train on the MFR SAE outputs than the baseline SAE outputs. This experiment would follow Section 4.2 of Gao et al., and was previously used to justify an architectural improvement in SAE training.

For point 2, we believe that Section 3 of our paper establishes why a bias for similar features is reasonable. This is because of the correlation we find between the similarity of features across SAEs, and the similarity of features with the input features. Empirically, this correlation can be leveraged to train SAEs the achieve a lower reconstruction loss, which we believe is further evidence that this is a good inductive bias.

While there may be reduced diversity in the features learned by all of the SAEs (i.e., the same features should appear across each of the SAEs), we do not believe a given SAE should have less diverse features as a result of being trained with MFR (i.e., the same features should not appear multiple times in a given SAE), as we do not see a component of MFR that incentivizes this.

Thank you again for your engagement throughout the review process.

Thank you for your review, we are glad you found the work of high quality.

We acknowledge the absence of a verification that SAEs trained with MFR permit better downstream analysis. We would commit to this experiment for a camera-ready revision:

1. Collect data on the activation of SAE features on tokens for the GPT-2 SAEs
2. Task humans with explaining the features based on the data from (1)
3. Task humans with predicting the activations of SAE features on unseen tokens given the explanations in (2), where more accurate predictions of the activations on unseen tokens indicate more interpretable features

At the reviewers request we could perform this experiment using a language model such as in Bills et al., however as mentioned in Section 2.2, recent work has brought this evaluation into question, showing a lack of causal relevance of language model feature descriptions. Ultimately we feel skeptical that there is a good method of evaluation for the interpretability of SAE features.

We acknowledge the minor weaknesses you mention, and would hope to uniformly address these in a camera ready version.

For your questions:

1. We speculate that MFR is more useful for the more sparse autoencoders because they are more likely to learn features that are not features of the input. We suggest that the extreme sparsity constraint on these models could incentivize feature composition, or hinder feature learning for the baseline autoencoders. Because there is information transfer between all autoencoders, this could hinder larger autoencoders by encouraging them to learn the features from the extremely sparse autoencoders.
2. Yes.
3. We view cosine similarity not accounting for magnitude as beneficial. We are ultimately interested in the features represented by the decoder weights rather than measuring the general similarity of the weights, so relative orientation seems the most important property to measure.

Dear authors, thank you for your rebuttal. I retain my score because I can't change my score based on an experiment that hasn't been done yet. I do want to encourage the authors to explore this line of research further and I think this work can be impactful and significant with more results that delve more into the sparse features of the EEG and GPT-2 data. A few recommendations I have for the authors are:

• Visualize what each sparse feature corresponds to in the EEG signal, e.g. are these specific brain regions? Generally designing an experiment that tries to look at whether the sparse features for example reduce noise in the EEG signal or are more interpretable than ICA would improve the impact of the paper for EEG research.
• Regarding my questions: adding your answers more clearly in the paper will increase the clarity of the paper, e.g. mention your assumption about cosine similarity when you introduce it as a metric (q3). Additionally, I would try to incorporate an experiment that verifies that "... MFR is more useful for the more sparse autoencoders because they are more likely to learn features that are not features of the input. We suggest that the extreme sparsity constraint on these models could incentivize feature composition, or hinder feature learning for the baseline autoencoders."

I would be happy to give you feedback on the experiment you propose in your rebuttal, but I don't quite understand the experiment the way it is presented in the rebuttal. Could you expand on the experimental design?

Thank you!

Dear authors,

Thank you for sending over the proposed experimental design. I think the experimental design sounds good, the authors could also use this approach to further argue that a smaller k (number of sparse latent dimensions in the autoencoder) is better (in the cited paper, Figure 6 shows that a higher number of sparse latent dimensions in the autoencoder is not necessarily better). However, I would additionally urge the authors to implement an experiment like Section 4.3 from the same paper as well. By showing how quantitative explanations differ across random initializations and are better with MFR the authors can even more clearly motivate their method. The authors could show a quick comparison either in the main text or Appendix across a few random initializations to show that differently initialized models can learn different sparse features, and then show how using MFR makes explanations more stable, i.e. similar across random initializations (I assume this is the case, but the authors could also explicitly show this), and more aligned with human interpretability.