Residual Stream Analysis with Multi-Layer SAEs

Belrose et al., 2023: https://arxiv.org/abs/2303.08112

Braun et al., 2024: https://arxiv.org/abs/2405.12241

Gao et al., 2024: https://arxiv.org/abs/2406.04093

Lieberum et al., 2024: https://arxiv.org/abs/2408.05147

Marks et al., 2024: https://arxiv.org/abs/2403.19647

Sharkey et al., 2022: https://www.alignmentforum.org/posts/z6QQJbtpkEAX3Aojj/interim-research-report-taking-features-out-of-superposition

I was surprised that in Figure 3 so many latents are active at only a single layer. Is it possible the MLSAE has effectively learned the same latent multiple times for different layers?

Thank you for calling attention to this. In the original manuscript, we included an analysis of the ‘mean max cosine similarity’ between decoder weight vectors (Sharkey et al., 2022) in Appendix B.3 and Figure 13 as a preliminary measure of the pairwise similarities between latents, and the results did not suggest a large number of ‘duplicate’ latents. However, we agree that this point merits further investigation.

In the revised paper, we have added comparisons between the distribution of pairwise cosine similarities between decoder weight vectors and (a) an equal number of independent and identically normally distributed vectors of equal length as a negative control and (b) a smaller number of i.i.d. normal vectors copied for each model layer plus a small amount of noise. Figure 19 shows that the distributions for decoder weight vectors are slightly heavier-tailed and right-shifted, i.e., a pair of MLSAE latents are slightly more likely to have high cosine similarity than a pair of i.i.d. normal vectors. However, we do not find evidence for many pairs of latents that are highly similar but active at different layers.

Thanks for your positive comments!

The paper does not evaluate the quality of the resulting MLSAEs using standard metrics, such as variance explained and cross-entropy loss when the SAE is inserted into the model. It is very easy to train a poor-quality SAE, and it is not clear if the MLSAEs in the paper are decent SAEs by these metrics.

Thank you for pointing this out. While we included such an evaluation in Appendix B of the original manuscript, we agree that we did not signpost this clearly in the main text (previously L079–080). As such, in the revised paper, we added a top-level summary of our evaluation results to Section 4.1 of the main text as Table 1 (p. 6). Specifically, we report the FVU and MSE reconstruction errors, the $L^1$ norm, and the increase in the cross-entropy loss and KL divergence when we replace the residual stream activations at a given layer by their reconstruction (‘downstream loss’ metrics, following Gao et al., 2024).

There is also no performance comparison to standard single-layer SAEs to see if MLSAEs result in a performance decrease when evaluated on a single layer.

Thank you for pointing this out. In Appendix B of the original manuscript, we compared the results of our evaluation to single-layer SAEs trained by Marks et al. (2024), Gao et al. (2024), Braun et al. (2024), and Lieberum et al. (2024). Specifically, we found that the FVU (normalized) reconstruction error and increase in cross-entropy loss were comparable. However, we agree with Reviewers g44V and d1f6 that further comparison to single-layer SAEs was needed, and we have now trained single-layer SAEs on each layer of Pythia-70m and 160m and evaluated them on data from every layer. This analysis appears in Appendix B.4 and Figures 15–17 in the revised paper. It confirms that, for a given layer of test data, a multi-layer SAE achieves very similar performance to a single-layer SAE trained at that layer and much better performance than a single-layer SAE trained at a different layer. Interestingly, we also find that applying tuned-lens transformations to the input vectors from each layer during training and evaluation does not improve the ability of single-layer SAEs to generalize to test data from other layers.

The paper only evaluates on Pythia models and topk SAEs - it would be beneficial to study other models and SAE architectures if possible.

We chose the Pythia suite of transformers because it is highly reproducible, as noted by Reviewer g44V, and offers a controlled setting to study how our results scale across model sizes. We chose TopK SAEs due to the simplicity of directly controlling the sparsity ($L^0$ norm) of the latent space (rather than indirectly, by tuning the $L^1$ penalty). However, we agree that confirming our findings for other transformers and SAE architectures would strengthen the paper, so we have now trained a multi-layer TopK SAE with our default hyperparameters on GPT-2 small (similar in size to Pythia-160m). In the revised paper, the results for GPT-2 small are described in Appendix B.3, along with aggregate and single-prompt heatmaps of the distributions of latent activations over layers (Figures 13 and 14). The qualitative results are similar to the Pythia models (e.g., Figures 2 and 3).

The paper mentions not having pre-trained TunedLens in some cases - where are these pre-trained TunedLenses from?

The reviewer rightly points out that we should have included the source of the pre-trained lenses we use in the paper: the HuggingFace repository https://huggingface.co/spaces/AlignmentResearch/tuned-lens. The authors of Belrose et al. (2023) and the implementation we describe in Section 3.2 made these models publicly available. We have added a note to clarify this point in Section 3.3 (L297–298) of the revised paper.

Why does Figure 6 have only Pythia 70m and 160m, and only TunedLens for 70m? The rest of the paper has many more sizes evaluated with and without TunedLens

Thank you for pointing this out. Most figures in the main text describe MLSAEs trained on different models with fixed hyperparameters (an expansion factor of $R=64$ and sparsity $k=32$). In contrast, many figures in the appendix vary the expansion factor $R$ between 1 and 256 and the sparsity $k$ between 16 and 512 for the smallest models (Pythia-70m and 160m). In the original manuscript, we intended Figure 6 to demonstrate the trends in the variance ratios when varying these hyperparameters. Unfortunately, it was prohibitively computationally expensive to perform the same hyperparameter sweeps for all models. In the revised paper, we have moved Figure 6 to the appendix; it is now Figure 22.

We were wondering whether you had time to consider our response?

In short, we agree that missing quality metrics such as variance explained and cross-entropy loss would be a worrying oversight. However, these quality metrics were available in Appendix B in the original paper, but we accept that that wasn't flagged clearly enough. We have therefore expanded Appendix B, and included a summary of these results in the main text of the revised manuscript.

Given the clarifications in our response, along with your otherwise excellent comments ("an important and understudied problem, and the contribution of training a single SAE on activations from multiple layers is novel") and ratings:

Soundness: 3: good

Presentation: 3: good

Contribution: 4: excellent

Would consider raising your score to an accept level?

Bell and Sejnowski, 1995: https://pubmed.ncbi.nlm.nih.gov/7584893/

Cunningham et al., 2023: https://arxiv.org/abs/2309.08600

Elhage et al., 2022: https://arxiv.org/abs/2209.10652

Gao et al., 2024: https://arxiv.org/abs/2406.04093

Park et al., 2023: https://arxiv.org/abs/2311.03658

Thanks for your comments!

It remains unclear to me how these statistical descriptions facilitate insight into the working mechanism of Transformers.

Please see our global response above for our understanding of our mechanistic contributions.

The motivation for the author's selection of Stacked Autoencoders (SAE) as the dimensionality reduction technique is unclear for me.

Importantly, SAE stands for "Sparse Autoencoders," not "Stacked Autoencoders," and SAEs are not a dimensionality reduction technique. Instead, SAEs learn a higher-dimensional or "overcomplete" basis. In the context of language-model interpretability, the use of SAEs starts from the observation that the dimension of the latent vectors in LLMs is relatively small (e.g., 1024), given that LLMs can represent, in some sense, all possible token "meanings." This observation motivated the hypothesis that we can understand LLM activations as lower-dimensional projections of a sparsely activated, higher-dimensional representation space whose basis vectors are easier for humans to interpret. This hypothesis has been explored in depth by, e.g., Elhage et al., 2022; Park et al., 2023. Critically, the resulting interpretable, sparse basis is overcomplete: it has more dimensions than the original LLM activation vectors.

The advantages of employing an SAE, which necessitates training, over more traditional statistical methods such as Principal Component Analysis (PCA) or t-distributed Stochastic Neighbor Embedding (t-SNE) for analyzing the latent space, are not clearly delineated. This choice warrants further clarification.

As described above, the goal is to find an overcomplete basis in which the original LLM activation vectors become sparser and more interpretable. In contrast, PCA and t-SNE give a basis smaller than (or the same size as) the input vectors and do not promote sparsity. Indeed, t-SNE is typically used to visualize high-dimensional data by embedding it into two or three dimensions and prioritizes visual separation over the interpretability of the embedded dimensions. While researchers have used techniques like t-SNE and UMAP to analyze the representations learned by SAEs (e.g., Gao et al., 2024; Appendix E.7), we cannot use them to find higher-dimensional, sparse, interpretable representations.

Importantly, there is a classical algorithm with strong practical and theoretical connections to SAEs: independent component analysis (ICA). We can understand SAEs as ICA with a noise model optimized by gradient descent. We note that Cunningham et al. (2023) compared PCA, ICA, and SAEs in the context of language-model interpretability, finding that SAEs produced more interpretable directions by their quantitative measures. Finally, it is important to note that ICA also requires iterative optimization or training (e.g., Bell and Sejnowski 1995, Sec. 2.4). In the revised paper, we have expanded on the connection between ICA and SAEs on L092–096.

In summary, the choice of SAEs was motivated by their popularity and apparent effectiveness for learning interpretable, higher-dimensional representations for LLM activation vectors, but one could adapt the key idea with our method to other algorithms.

Thank you for carefully considering our response!

In response to your most recent comment, we trained a multi-layer SAE on Llama-3.2-3B, which gave qualitatively very similar results to those shown in the paper.

This is expected: qualitative findings in mechanistic interpretability (particularly those derived using SAEs) have historically generalized remarkably well across model sizes. For example, the initial SAE papers (Cunningham et al., 2023; Bricken et al., 2023) used much smaller models than we do, and yet their qualitative findings were subsequently replicated in the largest commercial models (Gao et al., 2024; Templeton et al., 2024).

We plan to additionally run a few 7b or 8b models for the camera ready.

Unfortunately, OpenReview does not currently allow revisions. As such, we will, of course, put these results in the camera-ready. We hope you are excited to see them! Given that, and given your kind statement that "I believe your work has a great contribution to the community", we were wondering whether you would consider further raising your score to accept?

Thank you for your positive comments, acknowledging that our method "is an excellent improvement to sparse autoencoders and may be very useful for mechanistic interpretability" and "the writing is excellent, and the experiments look solid and reproducible."

The biggest weakness that I see is the lack of direct experimental comparison to previous techniques. I'd expect a paper to compare MLSAE and existing SAE techniques, similar to the comparison in Figure 5 of MLSAE w/ and w/o tuned lens. For example, I think it would be possible to show a version of Figure 1 (and almost all other figures, too) using previous SAE techniques, even without training on most of those layers. This would give the reader a qualitative understanding of whether MLSAE provided additional value.

Similarly, a quantitative measure (e.g., how much different methods agree on a similarity of different representations) could be made. In my eyes, the paper lacks a clear baseline to show improvement over.

Thank you for calling attention to this. We agree with reviewers g44V and d1f6 that understanding whether MLSAEs provide additional value over standard, single-layer SAEs requires we make explicit comparisons between the two. We argue that the natural baseline to compare against is a standard SAE, trained on the activation vectors from a single layer of the residual stream, but applied to the activation vectors from every layer. Hence, we have now trained single-layer SAEs on each layer of Pythia-70m and 160m and evaluated their performance in terms of the FVU reconstruction error and the increase in the cross-entropy loss when we replace the input activations at a given layer by their SAE reconstruction. In the revised paper, we have included evaluation results for single-layer SAEs in Appendix B.4 and Figures 15–17. These figures show that an MLSAE performs similarly to a single-layer SAE when evaluated on the same layer as its training data but that an MLSAE performs much better than a single-layer SAE evaluated on a different layer. Importantly, the poor performance of single-layer SAEs when applied to activations from other layers means that, while it is possible to reproduce Figures 2 and 3 (for example) by applying a single-layer SAE to the activation vectors from every layer, the single-layer SAE may not faithfully represent the original LLM activations, so it is unclear whether the distributions of latent activations over layers can be considered meaningful.

The above is especially important as the paper on its own doesn't offer that many additional insights into Transformer - apart from the surprising result of isolated latent activations.

Please see our global response above for our understanding of our mechanistic contributions.

The method also is quite simple.

We believe the simplicity of the method is to its advantage. The success of existing SAE techniques is partly due to their simplicity, which has enabled researchers to apply them to varied problems and neural network architectures. We hope MLSAEs will have similar advantages (albeit limited to architectures with residual connections, like transformers).

I would also like to see an analysis of the method's weaknesses - for example, it seems like it should have a higher computational cost of training, but it is not clear what is the exact impact of that (I assume the impact is quite small).

To the best of our knowledge, the strengths and weaknesses of our method are directly inherited from the underlying SAE. As for computational cost, an MLSAE is ultimately an SAE trained on different data (i.e., the activation vectors from every layer of the residual stream instead of a single layer). Hence, we expect the cost per activation vector encountered during training to be the same. Naturally, if you are training on 1 million tokens, then you have 1 million activation vectors for a single-layer SAE versus $L$ million for a multi-layer SAE (and $L$ layers). So, it is cheaper to train a single standard SAE on a given number of tokens than a single MLSAE. However, it is more common to train $L$ standard SAEs, one on each layer, at which point the computational cost becomes very similar to training a single MLSAE. In the revised paper, we briefly describe the computational cost of training MLSAEs at the end of Section 3.2.

Elhage et al., 2021: https://transformer-circuits.pub/2021/framework/index.html

Ferrando et al., 2024: https://arxiv.org/abs/2405.00208

The analysis is normalizing histograms of activations over layers to get distributions. This loses information about the relative frequencies of activations of different latents or about the number of layers in which each latent is active. Do you have any comments on this?

Thanks for the suggestion! In the revised paper, we added Figures 20 and 21 to Appendix D, which are the equivalent aggregate and single-prompt heatmaps to Figures 2 and 3, except with the un-normalized total latent activations instead of normalized distributions over layers. Surprisingly, the single-prompt heatmap contrasts with the normalized version, suggesting a ‘bimodal’ distribution of activations over layers for some latents (i.e., that some latents are active at non-adjacent layers). Unlike the seemingly unimodal normalized heatmaps, this suggests that latents with relatively large total activations are more likely to be bimodal. We are currently investigating trends in the bimodality of larger latent activations, e.g., how it correlates with the layer index, activation magnitude, etc., and will include these in the camera-ready.

The paper uses metrics based on the variance of the layer index. When looking at a single token, they depend on the number of activated layers as well as their distance. Could measuring the number of activated layers be more useful or easy to interpret for some purposes?

Thanks for the suggestion! We investigated the number of layers at which each latent is active, which we take to mean that it has a count of non-zero activations above a threshold value (e.g., 10K) over a large sample of input tokens (e.g., 10M). In the revised paper, we have added this analysis to Appendix E.2. This analysis suggests that the mean number of active layers tends to decrease as the number of latents increases relative to the model dimension. Given the preliminary finding of a bimodal structure in the larger latent activations discussed above, it is important to investigate these properties further, and we intend to pursue this for the camera-ready.

What is the intuition behind the hypothesis that features often remain active at later layers once they are activated? Is the hypothesis based on analogous observations when the SAE is trained only on the last layer?

Our intuition derives from the structure of the network: the residual connection means that the activations from previous layers are maintained, and a layer "writes" new information to the residual stream activations at subsequent layers (Elhage et al., 2021; Ferrando et al., 2024). We also suspected that "deleting" information would be difficult because the network would have to "write" in precisely the opposite direction in the residual stream. Looking at Figure 4, however, we note that a much easier way to effectively "delete" information is to overwhelm it with activation vectors with much greater magnitudes, such that the information written at previous layers becomes negligible.

Eq (8): Note that the residual connection matters only because of L2 regularization (and because it effectively changes the initialization).

Agreed; we have added a note below Eq. 8 (L265–267) to explain this.

Thanks for your positive comments and detailed suggestions!

this layer may differ for different tokens or prompts -> which layer it is might differ ...?

Thanks for pointing this out. We have changed this sentence on L019–020 in the revised paper to read: "Interestingly, we find that individual latents are often active at a single layer for a given token or prompt, but the layer at which the latent is active may differ for different tokens or prompts."

Under multiple figures: "expected value of the layer" -> "average layer index" or "expected value of the layer index"

Thanks! In the revised paper, we have replaced the phrase “expected value of the layer” with “expected value of the layer index” on L133, L161, and L370.

Eq. (3): The meaning of $\operatorname{Var}$ could be clarified.

Thanks! We have clarified this below Eq. 3 with the sentence, “Here, $\operatorname{Var}$ is the variance.”

In section 3.1, $\bf x \in \mathbb{R}^d$, but in Eq. (3), $\bf x$ represents a sequence of vectors $\bf x_\ell$.

Thanks! Our use of subscripts to indicate the layer index of activation vectors needed to be clarified. We have removed the layer indices from the subscripts of vectors (e.g., $\bf x_\ell$) in Eqs. 3, 8, and 9 because they are unnecessary to convey the meaning.

Eq. (9): $\hat{\bf x}\ell$ is the same as $\bf x\ell$. It is a bit confusing that two symbols are used for the same thing.

Thanks! We have added a sentence to explain our notation after Eq. 9: $\hat{\bf x}$ is the vector output by the autoencoder, i.e., the reconstruction, whereas $\bf x$ is the vector input to the autoencoder. Both the output and input vectors have the same shape but are not equal. As mentioned above, we have removed the layer indices from Eqs. 3, 8, and 9.

Some training details seem missing: How are layers for mini-batches sampled?

Thanks for pointing this out. We have added an explanation to Appendix A in the revised paper (L764–766): we construct a batch of input activation vectors by performing the forward pass of the transformer over a sequence of tokens, collecting the residual stream activation vectors from every layer, and stacking them together. Hence, each batch of activations has an equal number of vectors from each layer of the transformer (the number of tokens in the sequence).

It is not very clear to me how the paper is advancing the research in mechanistic interpretability. Perhaps commenting more on what motivated the investigations and what the results suggest for future work directions would be valuable.

Please see our global response above for our understanding of our mechanistic contributions.

The paper does not state whether the source code for reproducing the experiments will be published.

We have included anonymized source code as supplementary material. We will provide a more convenient link in the camera-ready.

L245: “ratio between the model dimension and the number of latents” should be the other way around.

Fixed!

L353: Should $E_T[L|J=j]$ be $E[L|J=j]$ (or, equivalently, $E_L[L|J=j]$?

Fixed!

How much of computational resources did the experiments require?

Thanks for pointing this out. The revised paper describes the computational cost of training an MLSAE at the end of Section 3.2 (L241–245). In brief, we ran most experiments on a single NVIDIA GeForce RTX 3090 GPU for between 12 and 24 hours; we ran the largest experiments (e.g., with Pythia-1b or an expansion factor of $R=256$) on a single NVIDIA A100 80GB GPU for up to three days.

L301: “we found the mean cosine similarities increased accordingly.” This is not true for smaller models. Do you have any intuitions about that?

Thanks for pointing out our confusing phrasing. We intended by this statement to claim that the mean cosine similarity between residual stream activation vectors increases as the size of the underlying transformer increases, which is true in Figure 1. However, we agree that the phrase "the mean cosine similarities increased accordingly" is unclear: one could take it to mean that the mean cosine similarity increased across layers. In Figure 1, this is indeed false for smaller models. We have clarified this statement, replacing it with "we found the mean cosine similarities increased as the model size increased" (L346).

Thanks for reading the revision so carefully!

We agree about the variance: we've updated the text to read, "Here, Var is the variance, treating x as a random vector, where the randomness is induced by randomizing the token, producing different activation vectors."

We will include all color bars for the camera-ready (there are about 10 figures to update, so I don't believe we're going to be able to finish before the imminent deadline for revisions).

As regards the sorting: good catch! It does seem there was an error in the sorting. We have revised Figure 21, and it now looks very similar to the normalised plots (e.g. Figure 3), as we would expect.

We were wondering whether you would consider raising the score to reflect the "additional experiments, investigations, and insightful explanations!" Needless to say, score changes at this point are likely to be pivotal in deciding the paper's acceptance.

Regarding Figure 21, I am happy this got sorted out.

I agree with your priority assessments about the color bars. I am particularly interested in a color bar for the cases where $\gamma\neq 1$. You might also consider mentioning the name of the color map.

Taking the overall discussion and paper updates into account, I have increased my rating. Thank you for your effort!

Some more comments about notation confusion

After looking at the code, I see that I misinterpreted the FVU equation (Eq. (3)). Is the following interpretation correct? The variance is computed over inputs (from the mini-batch?), token representations (from all tokens and layers), and features (with indices $1..d$). That is, given a set $D$ of representation vectors from multiple inputs, tokens and layers, it is the variance of the sample $\{\mathbf x[i] : i\in \{1..d\}, \mathbf x\in D\} \subset \mathbb R$, where $\mathbf x[i]$ is the $i$-the feature of the vector $\mathbf x$.

What confused me was that (1) variance is usually not defined for random vectors (covariance is), (2) the function in the paper is not computed based on a single pair of vectors, but based on a set of vectors.

For the camera-ready version of the paper, I suggest carefully checking the meaning of the symbols and making the notation more clear.

Thank you! It seems that i had misread the dim argument – I apologize. This makes more sense.