Residual Stream Analysis with Multi-Layer SAEs
Authors: Tim Lawson
Lucy Farnik
Conor Houghton
Laurence Aitchison
Venue: ICLR 2025 Poster (2025)
Keywords: sparse autoencoders
mechanistic interpretability
language models
Primary area: interpretability and explainable AI
TL;DR: We train a single SAE on the residual stream activations from every transformer layer and find latents active at multiple layers for different tokens.

Abstract: Sparse autoencoders (SAEs) are a promising approach to interpreting the internal representations of transformer language models. However, SAEs are usually trained separately on each transformer layer, making it difficult to use them to study how information flows across layers. To solve this problem, we introduce the multi-layer SAE (MLSAE): a single SAE trained on the residual stream activation vectors from every transformer layer. Given that the residual stream is understood to preserve information across layers, we expected MLSAE latents to 'switch on' at a token position and remain active at later layers. Interestingly, we find that individual latents are often active at a single layer for a given token or prompt, but the layer at which an individual latent is active may differ for different tokens or prompts. We quantify these phenomena by defining a distribution over layers and considering its variance. We find that the variance of the distributions of latent activations over layers is about two orders of magnitude greater when aggregating over tokens compared with a single token. For larger underlying models, the degree to which latents are active at multiple layers increases, which is consistent with the fact that the residual stream activation vectors at adjacent layers become more similar. Finally, we relax the assumption that the residual stream basis is the same at every layer by applying pre-trained tuned-lens transformations, but our findings remain qualitatively similar. Our results represent a new approach to understanding how representations change as they flow through transformers. We release our code to train and analyze MLSAEs at https://github.com/tim-lawson/mlsae.

Published as a conference paper at ICLR 2025
RESIDUAL STREAM ANALYSIS WITH MULTI-LAYER SAES
Tim Lawson" Lucy Farnik Conor Houghton Laurence Aitchison
School of Engineering Mathematics and Technology University of Bristol Bristol, UK
ABSTRACT
Sparse autoencoders (SAEs) a promising approach to interpreting the internal representations of transformer language models. However, SAEs are usually trained separately each transformer layer, making it difficult to use them study how information flows across layers. To solve this problem, we introduce the multi-layer SAE (MLSAE): a single SAE trained on the residual stream activation vectors from every transformer layer. Given that the residual stream is understood to preserve information across layers, we expected MLSAE latents 'switch on' a token position and remain active at later layers. Interestingly, we find that individual latents are often active at a single layer for a given token or prompt, but the layer at which an individual latent is active may differ for different tokens or prompts. We quantify these phenomena by defining distribution over layers and considering its variance. We find that the variance of the distributions of latent activations over layers is about two orders of magnitude greater when aggregating over tokens compared with a single token. For larger underlying models, the degree to which latents are active at multiple layers increases, which is consistent with the fact that the residual stream activation vectors adjacent layers become more similar. Finally, we relax the assumption that the residual stream basis is the same at every layer by applying pre-trained tuned-lens transformations, but findings remain qualitatively similar. Our results represent a new approach to understanding how representations change they flow through transformers. We release our code train and analyze MLSAEs at httpt::///// //gi thub. com/ htt///timwwso/mmssaaaa
INTRODUCTION
Sparse autoencoders (SAEs) learn interpretable directions or 'features' in the representation spaces of language models (Elhage et al. 2022; Cunningham et al. 20233 Bricken et al. (003).. Typically, SAES are trained on the activation vectors from a single model layer (Gao et al. 2024 Templeton et al. 2024 Lieberum et al.) 2024). This approach illuminates the representations within a layer. However, Olah (2024) Templeton et al. (2024) believe that models may encode meaningful concepts by simultaneous activations in multiple layers, which SAES trained at a single layer do not address. Furthermore, is not straightforward to automatically identify correspondences between features from SAES trained at different layers, which may complicate circuit analysis (e.g. He et al., 2024).
To solve this problem, we take inspiration from the residual stream perspective, which states that transformers VVaswani et al. 2017) selectively write information to and read information from token positions with self-attention and MLP layers (Elhage et al., 2021 Ferrando et al., 2024h. The results of subsequent circuit analyses, like the explanation of the indirect object identification task presented by WWan et al. (2022), support this viewpoint and cause us to expect the activation vectors at adjacent layers in the residual stream be relatively similar (Lad et al. 2024).
To capture the structure shared between layers in the residual stream, we introduce the multi-layer SAE (A): single SAE trained the residual stream activation vectors from every layer of a transformer language model. Importantly, the autoencoder itself has single hidden layer - it is
"Correspondence tim. Lawson@bristoo uc.uk.

---

Published as a conference paper at ICLR 2025
Mcan Cosine Similarity
0.8
0.4
Pythia-70n 70m
Pythia-160m
Pythia-410m
Pythia-
0.2
Pythia-1.45
Pythia-2.8b
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
Relative Layer
The mean cosine similarities between the residual stream activation vectors at adjacent layers of transformers, over million tokens from the test set. To compare transformers with different numbers of layers, we divide the lower of each pair of adjacent layers by the number of pairs. This 'relative layer' the r-axis the plot. We subtract the dataset mean from the activation vectors at each layer before computing cosine similarities to control for changes in the norm between layers (Heimersheim Turner 2023), which we demonstrate in Figure
Figure 1:
multi-layer only in the sense that it is trained activations from multiple layers of the underlying transformer. In particular, we consider the activation vectors from each layer as separate training examples, which is equivalent training single SAE each layer individually but with the parameters tied across layers. We briefly discuss alternative methods in Section
We show that multi-layer SAEs achieve comparable reconstruction error and downstream loss to single-layer SAES while allowing us to directly identify and analyze features that are active at multiple layers (Section 4.1). When aggregating over large sample of tokens, find that individual latents are likely to be active at multiple layers, and this measure increases with the number of latents. However, for single token, latent activations are more likely to isolated single layer. For larger underlying transformers, we show that the residual stream activation vectors at adjacent layers are more similar and that the degree to which latents are active at multiple layers increases.
Finally, we relax the assumption that the residual stream basis is the same at every layer by applying pre-trained tuned-lens transformations to activation vectors before passing them the encoder. Surprisingly, this does not obviously increase the extent of multi-layer latent activations.
2 RELATED WORK
A sparse code represents many signals, such sensory inputs, by simultancously activating relatively small number of elements, such as neurons (Olshausen Field 1996 Bell & Sejnowski 1997). Sparse dictionary learning (SDL.) approximates each input vector by a linear combination of a relatively small number of learned basis vectors. The learned basis is usually overcomplete: has greater dimension than the inputs. Independent Component Analysis (ICA) achieves this aim by maximizing the statistical independence of the learned basis vectors by iterative optimization or training (Bell & Sejnowski] 1995) (1977 Hyvärinen Oja 2000) Le al., 201T). Sparse autoencoders (SAEs) can be understood as ICA with the addition of noise model optimized by gradient descent (Lee et al. 2006 Ng 201T Makhzani Frey 2014)
The activations of language models have been hypothesized to be a dense, compressed version of a sparse, expanded representation space (Elhage et al. 2021 2022). Under this view, there are interpretable directions in the dense representation spaces corresponding to distinct semantic concepts, whereas their basis vectors (neurons) are 'polysemantic' (Park et al. (003). It has been shown theoretically (Wright & Ma 2022) and empirically (Elhage al.) 2022: Sharkey et al. 2022 Whittington et al. 2023) that SDL recovers ground-truth features in toy models, and that learned dictionary elements more interpretable than the basis vectors of language models (Cunningham
ct al. 2023 Bricken et al. (023) or dense embeddings (O'Neill et al. (024).. Notably, "features' are
necessarily linear (Wattenberg & Viégas 2024 Engels al. 2024 Hernande al. 2024).
2

---

Published as a conference paper at ICLR 2025
Pythia-1660 160m
Layer
0
Pythia-41000
23
Layer
0
Pythia- 1b
15
Layer
Pythia-1.4B
23
Layer
Latent
Figure 2: Heatmaps of the distributions of latent activations over layers when aggregating over 10 million tokens from the test set. Here, we plot the distributions for MISAEs trained on Pythia models with an expansion factor of R 64 and sparsity 32. The latents are sorted in ascending order of the expected value the layer index (Eq. 10
Pythia-
160m 1ythia-160m
11
Layer
0
Pythia-4100 10mm
23
Layer
Pythia lb
Layer
Pythia- 1.4b
23
Layer
0
Latent
Figure 3: Heatmaps of the distributions of latent activations over layers for single example prompt. Here, we plot the distributions for MLSAES trained on Pythia models with an expansion factor of R 64 and sparsity gave" (Wang et al. 2022). We exclude latents with maximum activation below x 10-3 and sort latents in ascending order of the expected value of the layer index (Eq. 10
32. The example prompt is "When John and Mary
the store, John
3

---

Published as a conference paper at ICLR 2025
The standard SAE architecture is single hidden layer with ReLU activation function and L sparsity penalty in the training loss (Bricken et al. (2023), but various activation functions (Makhzani
& Frey 2014 Konda al.) 2013 Rajamanoharan et al. 2024b a) and objectives (Braun al. (2024) have been proposed. The prevailing approach is to train an SAE on the activation vectors from a single transformer layer, except for Kissane al. (2024), who concatenate the outputs of multiple attention heads in single layer, and Yun et al. (2021). who learn a sparse dictionary for the residual stream at multiple layers, albeit by iterative optimization instead of with autoencoder.
Mechanistic interpretabilityy research often attempts to identify circuits: computational subgraphs of neural networks that implement specific behaviors (Olah al. 20200 Wang et al., 2022 Conmy
et al. 2023 Dunefsky et al., 2024) Garcia- Garcia-Carrasco et al. 2024; Marks et al., 2024). Representing
networks terms of SAE latents may help to improve circuit discovery (He et al. 2024 O' O'Neill
& BBi 2024), and these latents can be used to construct steering vectors (Subramani et al. 2022
Templeton al. 2024) Makelov| (002)) but
unclear whether SAEs outperform baselines for causal analysis (Chaudhary & Geiger 2024 Huang et al. 2024). Importantly, SAEs can be scaled
up to the activations of large language models, where we expect the number of distinct semantic
concepts to be extremely large (Tmmlleton et al., 2024 Gao al.
2024 Lieberum et al. 2024).
The 'logit lens' is a method to interpret directions the residual stream by projecting them onto the vocabulary space to elicit token predictions, i.c., multiplying them by the unembedding matrix (nostalgebraist, 2020). However, the residual stream basis is not fixed, so Belrose et al. (2023) introduce the 'tuned lens' approach, where linear transformation is learned for each layer in the residual stream. The objective is to minimize the KL divergence between the probability distribution over tokens generated by the transformed activations and the 'true' distribution of the model. This approach draws on the perspective of iterative inference (Jastrzebski et al. 2018).
The key difference between previous work (Bricken al. 2023] Cunningham 2023 Templeton et al. 2024] Gao et al. 2024) and our work is that we introduce the multi-layer SAE, i.c. we train single SAE at all layers of the residual stream.
METHODS
The key idea with multi-layer SAE is to train a single SAE on the residual stream activation vectors from every layer. In particular, we consider the activations at each layer to be different training examples. Hence, for residual stream activation vectors of model dimension d, the inputs to the multi-layer SAE also have dimension For nu tokens and nz layers, we train the multi-layer SAE nynz activation vectors. We use the terms 'SAE feature' and 'latent' interchangeably.
3.1 SETUP
We train MLSAEs primarily on GPT-style language models from the Pythia suite (Biderman et al. 2023); see Appendix B.3 for others. We are interested in the computation performed by sell-attention and MLP layers on intermediate representations (Valeriani et al. 2023). Hence, we take the residual stream activation vectors after a given transformer block has been applied, excluding the input embeddings before the first block and taking the last-layer activations before the final layer norm.
We use k-sparse autoencoder (Makhzani & Frey 2014 Gao et al. 2024), which directly controls the sparsity of the latent space by introducing Topk activation function that keeps only the largest latents. The A largest latents are almost always positive for & < d, but we follow Gao et al. (2024) in applying a ReLU activation function to guarantee non-negativity. This setup effectively fixes the sparsity (L" norm) of the latents at per activation vector (layer and token) throughout training. For input vectors x e IRd and latent vectors h E R", the encoder and decoder are defined by:
h ReLUU Wen.x bpre))
(1)
8 = Wacch + b
(2)
where W 4x W E xxd and bpre Rd We constrain the pre-encoder bias buue to be the negative of the dest-decoder bias, following Bricken et al. (2023); Gao et al. (2024), and standardize activation vectors to zero mean and unit variance before passing them to the encoder.
4

---

Published as a conference paper at ICLR 2025
3.2 TRAINING
We use the fraction of variance unexplained (FVU) as the reconstruction emror:
|x - 2
FVU(x, 8)
Var(x)
Here, Var is the variance. We chose the FVU because the input vectors from different layers may have different magnitudes; choosing the mean squared error (MSE) would encourage the autoencoder to prioritize minimizing the reconstruction errors of the layers with the greatest magnitudes.
potential issue when training SAEs the occurrence of 'dead' latents, i.e., latent dimensions that are almost always zero. With a k-sparse autoencoder, this means latent dimensions that almost never appear among the largest latent activations. We follow Bricken et al. (023); Cunningham et al. (2023) by considering a latent 'dead' if it is not activated within the last 10 million tokens during training. In the multi-layer setting, a latent may be activated by the input vectors from any layer.
Gao et al. 2024 Appendix A.2) propose auxiliary loss term to minimize the occurrence of dead latents. This AuxK term models the MSE reconstruction error using the haux largest dead latents:
AuxK(x, x) = lle - él12
Here, e & is the reconstruction error of the main model, and è is its reconstruction using the top-kaux dead latents. Let Dead be an 'activation function' that keeps only the dead latents. Then:
h ReLUU (Dead(Wene) bpre)))
(5)
Wacehadadd + bpr
(6)
The full loss the FVU plus the auxiliary loss term, multiplied by small coefficient as
[ =VU x) + a AuxK(xx
Following Gao al. (2024), we choose kaux as a power of close to d/2 and a 1/32.
Our hyperparameters are the expansion factor R = n/d, the ratio of the number of latents to the model dimension, and the sparsity kt the number of largest latents to keep in the TopK activation function. We choose expansion factors as powers of between 1 and 256, yielding autoencoders with between 512 and 131072 latents for Pythia-70m, and as powers of 2 between 16 and 512 (Appendix (B).
The computational expense of training single multi-layer SAE on n. layers of the residual stream approximately equal to training n,, single-layer SAES on the same number of tokens. We trained most MLSAEs single NVIDIA GeForee RTX 3090 GPU for between 12 and 24 hours; trained the largest MLSAEs (c.g., with Pythia-lb or an expansion factor of 256) on a single NVIDIA A100 80GB GPU for up to three days.
The implementation is based on Gao et al. 2023) Belrose 2024 ; see Appendix for details.
3.3
TUNED LENS
In the tuned lens method, an affine transformation is learned from the output space layer to the output space of the final layer, called the translator for layer é (Belrose et al. (223). With our setup, we want to transform the residual stream activation vectors at each layer into more similar bases before passing them to the encoder and invert that transformation after the decoder.
Importantly, the authors note that their implementation uses a residual connection:
x' x (Winss,X le x bjens.)
Here, x is the input vector to the encoder, and x' is the transformed input vector. This parameterization ensures that Le regularization (weight decay) pushes the transformation towards the identity matrix instead of zero. Hence, to invert the transformation, we need:
=
+ Wiens
(x' biens,))
(9)
file: tuned wwws//n//ennseeeeeeeeeeeeeee
5

---

Published as a conference paper at ICLR 2025
Model
FVU
L Norm
Loss f
Pythia-70m
0.097
66.5
0.565
Pythia-160m
0.106
76.1
0.432
Model
FVU
L Norm
Loss
Pythia-410n
0.081
84.6
0.414
Pythia-70m
0.030
61.5
0.274
Pythia- 0.095
109.7
0.404
Pythia-160m
0.088
89.6
-0.080
Pythia- 1.4b
0.106
124.7
0.487
Pythia-410m
0.073
80.5
0.827
Gemma 2B
0.210
255.5
1.483
Llama 3.2 3B
0.299
141.4
0.347
(b) With tuned lens
GPT-2 small
0.093
196.9
0.759
(a) Standard
Table The mean FVU, L norm, and increase the cross-entropy loss (Loss f) for MLSAEs with an expansion factor of R 64 and sparsity le 32, over million tokens from the test set. We provide details the evaluation metrics and transformer architectures Appendix B
In Eq. 9 & is the transformed output vector of the decoder, & is the output vector, Winns E RRdxd and biens, E IRd. With our setup, x' and & replace the input and output vectors that we pass to the encoder and use to compute the loss. Notably, we use the transformed vectors to compute reconstruction errors (Figure 12), and we compute the inverse (I + ei - for each layer once at the start of training. The pre-trained tuned lenses used were provided by the authors of Belrose et al. (2023). which did not include Pythia-I lb at the time of writing
4 RESULTS
4.1
EVALUATION
The key advantage of multi-layer SAE is to be able study how information flows across layers in the residual stream. However, this approach is only useful the MLSAE performs comparably to single-layer SAEs. The FVU reconstruction error term in the loss (Section 3.2) is proxy for the degree to which an SAE explains the behavior of the underlying model. Hence, we also measured the increase in the cross-entropy loss when the residual stream activations a given layer replaced by their reconstruction, following Braun et al. (204)); Gao et al. (2024); Lieberum et al, (2024).
Table summarizes the evaluation results for MLSAEs trained with default hyperparameters. The FVU and delta cross-entropy (CE) loss remain consistent across model sizes for Pythia transformers. In most cases, applying tuned-lens transformations decreases the FVU and delta CE loss (see 4.4 44.4|aaa and Figure 12). We provide results for other hyperparameters. breakdowns by the layer of the input activation vectors, and explicit comparisons single-layer SAEs in Appendix B
Section
4.2 REPRESENTATION DRIFT
Guided by the residual stream perspective (Elhage et al. 2021 Ferrando et al. 2024). expected dense activation vectors to be relatively similar across layers. As an approximate measure of the degree to which information is preserved in the residual stream, computed the cosine similarities between the activation vectors at adjacent layers, similarly to Lad et al. 2024 Appendix A). A similarity of one means that the information represented intervening residual block, whereas a similarity of zero means the activation vectors on either side of the block are orthogonal. We had expected changes in the residual stream become smaller the model size increased, and the mean cosine similarities increased as expected (Figure [..
token position is unchanged by the
Given that the residual stream activation vectors are relatively similar between adjacent layers, we expected to find many MLSAE latents active at multiple layers. We confirmed this prediction over a large sample of 10 million tokens from the test (Figure 2}. Interestingly, we found that for individual prompts, a much greater proportion of latents are active at only a single layer (Figure 3).
https://huuginngface.cmm http:///eeerriiooooooeeee/llr/te

---

Published as a conference paper at ICLR 2025
Pythia-70m 70mm
Pythia 160m
200
Pythia-410m
Pythia- 1b
Mean L2 Norm
Pythia-1.4b
Pythia-2.8b
100
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
Relative Layer
Figure 4: The mean L norm of the residual stream activation vectors at every layer, over 10 million tokens from the test set. To compare transformers with different numbers of layers, we divide the layer index l by the number of layers nt, This 'relative layer' is the x-axis of the plot.
Following Heimersheim & Turner (002)), we verified that the mean L norm of the activation vectors increases across layers, which prompted us to center the vectors at each layer by subtracting the dataset mean before computing the similarities between vectors (Figure
4.3 LATENT DISTRIBUTIONS OVER LAYERS
Given dataset and MLSAE, each combination of token and latent produces distribution of activations over layers. We want to understand the degree to which the variance of that distribution depends the token versus the latent quantify the intuition gleaned from Figures 2 and
Consider the layer index L, token T, and latent index J to be random variables. We take P(J) to be uniform discrete distribution, P(T) J) to be uniform discrete distribution over tokens for which the latent is active (at any layer), and L to be sampled from a conditional distribution proportional to the total latent activation at that layer, aggregating over tokens:
1x(x),))
P = 5) = Ee hz(x)1))
(10)
Here, X is the dense residual stream activation vector at token and layer e, while hy (x+,)) is the activation of the j-th MISAE latent at that token and layer.
We order latents in all heatmaps using the expected value of the layer index for single latent ETL J jj. The variance of the distribution over layers measures the degree to which latent is active at single layer (in which case, it is zero) versus multiple layers (in which case, it is positive). We are interested the following variances of the distribution over layers:
VariL J j, T f), for single latent and token
VariL J = j], for single latent, aggregating over tokens
Var aggregating over both latents and tokens
These quantities related by the law of total variance (see Appendix C.2 For the moment, note that the variance of the distribution over layers naturally depends on the number layers n.. Hence, compare different models, we look ratios between these variances:
Variance for one latent, aggregating over tokens, as a proportion the total variance over latents
E Var(L J]]
=
Var(L) (L)
(11)
Variance for one token and latent as = proportion of the total variance for that latent
E. Var(L J, T))
(12)
E[Var(L J)]
The former measures the degree which latents active multiple layers when aggregating tokens, and the latter compares this to the case for a single token. We explore alternative measures of aggregate multi-layer activity in Appendix
7

---

Published as a conference paper at ICLR 2025
Pythia-70m
Standard Tuned lens
Pythia-1600 160m
Pythia-410m
Pythia-Ib
Pythia-1.4B 1,4b
Gemma 2B
Llama
GPT-2 small
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
Variance for one latent, aggregating over tokens, a proportion the total variance over latents
Pythia-70m
Standard
Pythia-160m
Tuned lens
Pythia-410m
Pythia-1b lb
Pythia- 1.4ia-1,.4b
Gemma 2B
Llama 3.2 3B
GPT-2 small
0.01
0.02
0.03
0.04
0.05
0.06
0.07
0.08
0.09
0.1
0.11
0.12
Variance for one token and latent as a proportion of the total variance for that latent
Figure 5: The fraction of the total variance explained by individual latents and the fraction of the variance for an individual latent explained by individual tokens (Eqs. [ and [2) for MLSAEs with expansion factor of R 64 and sparsity h 32, 10 million tokens from the test set. The absence of bars for tuned-lens MISAES indicates the absence of results, not that the values are zero.
The degree to which latents are active multiple layers when aggregating over tokens relatively large, between 54 and 88%, and broadly increases with the model size for fixed hyperparameters (Figure 5). This measure quantifies the observation that, in the aggregate heatmaps (Figure 2 the distributions of latent activations over layers become more 'spread out' the model size increases. Conversely, we find that the fraction of the variance for an individual latent explained by individual tokens is relatively small, the order of to 10%. This quantifies the observation that, in the prompt heatmaps (Figure 3,, the distributions over layers are much less 'spread out" than in the aggregate heatmaps.
4.4
TUNED LENS
Thus far, have assumed that the residual stream basis is the same every layer. We relaxed this assumption by applying pre-trained tuned-lens transformations to the residual stream activations each layer before the encoder (Section B.3). We had expected that these transformations would increase the degree to which latents were active at multiple layers because they translate the activations every layer into a basis more similar to the basis of the output layer. The aggregate and single- prompt heatmaps (Figures 6 and 7) indicate modest increase in the degree which latents active at multiple layers compared with the standard approach.
The variance ratios in Figure clarify that the tuned-lens approach decreases the degree to which latents are active at multiple layers when aggregating over tokens. This ratio remains approximately constant the expansion factor increases, between 37% and 41% (Figure 27). Conversely, the variances for single token relative to a single latent are larger, i.c., the single-prompt heatmaps are more 'spread out' compared with the standard approach.
8

---

Published as a conference paper at ICLR 2025
Pythia-TUe 70m with tuned lens
Layer
Pythia- 160m with tuned lens
11
Layer
Pythia-410 410m with tuned lens
23
Layer
0
Latent
Figure 6: Heatmaps of the distributions of latent activations over layers when aggregating over 10 million tokens from the test set. Here, we plot the distributions for tuned-lens MLSAEs trained on Pythia models with an expansion factor R = 64 and sparsity Figure
For standard MLSSASs, see We note that pre-trained tuned lens not available for Pythia- (Section 3.3
Pythia- Pythia-70m 70m with tuned lens
5
Layer
Pythia- 160m with tuned lens
11
Layer
Pythia-4 410maa-410m with tuned lens
23
Layer
Latent
Figure 7: Heatmaps of the distributions of latent activations over layers for single example prompt. Here, plot the distributions for tuned-lens MISAEs trained Pythia models with expansion factor of R = 64 and sparsity 1 = 32. The example prompt is "When John and Mary went to the store, John gave" (Wang et 2022). For standard MLSAES, see Figure 3 We that a pre-trained tuned lens was not available for Pythia-Ib (Section 3.3

---

Published as a conference paper at ICLR 2025
5 DISCUSSION
We considered the activation vectors from different layers as different training examples, so we passed RrT vectors of length d to the autoencoder, where nr is the number of tokens, is the number of layers, and d is the dimension of the residual stream. This approach might be called a "data-stacked" MLSAE. An alternative would be 'feature-stakedd MLSAE, i.c., concatenate the activation vectors from different layers into a single vector of dimension ny.d This alternative might be better suited to capturing the notion of "cross-layer superposition," which take mean a small number of simultaneously active sparse features at multiple layers encoding a single meaningful concept OOlah 2024 Templeton et al. 2024).
We began by pursuing the feature-stacked approach but discarded it. The essential issue is that single set of sparse features descrihes the residual stream activations at every layer, which makes it difficult to understand how information flows through a transformer. For example, it would not be possible to plot the activations of sparse features across layers. Moreover, to compute this set of features, one must first compute the activations at every layer, which makes it more difficult to evaluate performance by traditional measures like single-layer reconstruction errors. Finally, the information encoded at one token position may differ substantially between layers due to ceff-attention. In the early layers, the representation likely to primarily encode the input token and position embedding, whereas in the later layers, the representation may encode more complex properties of the surrounding context. It is not immediately apparent that jointly encoding this information by a single SAE is sensible. Instead, one might wish to separately capture the different information present at token position across layers, which allowed with our data-stacked approach.
CONCLUSION
We introduced the multi-layer SAE (MLSAE), where we train a single SAE on the activations at every layer of the residual stream. This allowed to study both how information is represented within a single transformer layer and how information flows through the residual stream.
We confirmed that residual stream activations are relatively similar across layers by looking at cosine similarities before considering the distributions of latent activations over layers. When aggregating over a large sample of ten million tokens, we observed that most latents were active at multiple layers, but for single prompt, most latent activations were isolated to single layer. To quantify these observations, we computed the fraction of the total variance explained by individual latents and the fraction of the variance for an individual latent explained by individual tokens. This analysis confirmed that the degree to which latents are active at multiple layers when aggregating over tokens was large, increasing with the model size and expansion factor, and that the fraction of the variance explained by individual tokens was small.
Understanding how representations change as they flow through transformers is critical identifying meaningful circuits, which is a core task of mechanistic interpretability. Despite the utility of the residual stream perspective, our results demonstrate that representation drift, and perhaps the increasing magnitude of changes to the residual stream across layers, is a significant obstacle to identifying meaningful computational variables with SAEs. Nevertheless, we argue that an approach such as the MLSAE, which considers the representations at multiple layers in parallel, is necessary for future methods that seek to interpret the internal computations of transformer language models.
7 ACKNOWLEDGEMENTS
Tim Lawson and Lucy Farnik were supported by the UKRI Centre for Doctoral Training in Interactive Artificial Intelligence (EP/S022937/1). This work carried out with HPC systems provided by the Advanced Computing Research Centre at the University of Bristol. We also thank Dr. Stewart, whose philanthropy supported the compute resources used.
10

---

Published as a conference paper at ICLR 2025
8 REPRODUCIBILITY STATEMENT
We release our code to train and analyze MLSAEs at https://github. http::///gtttuu...ommmmmmm//s/// and the models described in the paper athttpss:/hugg. ingi ace www..cooppaaaaeee wwtpapors////ccecccooc 2469. 04185 Section and Appendix A describe the training setup. Section 4.1 and Appendix B describe the evaluation metries, and Section 4.3 and Appendix
Appendix((eesrrib deescrib our analyses.
REFERENCES
Anthony J. Bell and Terrence Sejnowski. An Information Information-Maxxiizzthiee Maximization Approach to Blind Separation and Blind Deconvolution. Neural Computation, 7(6):1129-11553 November 1995. ISSN 0899-7667. doi: 10. 622/ncco. (at.nec...999..... 1995.7.6. 1129. WRL.https::/ieeexpllee leee. org/ab
stract/6ocument/7799222 Conference Name: Neural Computation.
Anthony J. Bell and Terrence Sejnowski. The "independent components" of natural scenes edge filters. Vision Research, 37(23) 10.101 10:1016633042--69999900000011100
December 1997. ISSN 0042-6989. doi:
e/ e/p11/S0042699377001111
Nora Belrose. ElcutherAl/sac, May 2024. URL httl:tppp//giihib...eeeeeeeeeeee
Nora Belrose, Zach Furman, Logan Smith, Danny Halawi, Igor Ostrovsky, Lev McKinney, Stella Biderman, and Jacob Steinhardt. Eliciting Latent Predictions from Transformers with the Tuned Lens, November 2023. /
08112 arXiv:2303.08 12 [cs].
Stella Biderman. Hailey Schoelkopf. Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien. Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, Usvsn Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar Van Der Wal. Pythia: Suite for Analyzing Large Language Models Across Training and Scaling. In Proceedings of the 40th International Conference on Machine Learning, pp. 2397-2430, PMLR, July 2023. URL/https://proceel/
nns. mls.m ggs.mmr.prres//vvee wwtp:/prewese/eeeeeetee http://222//i/ee 2ttpss///22//bidrmma22..... html ISSN: 2640-3498.
Dan Braun, Jordan Taylor, Nicholas Goldowsky-Din and Lee Sharkey. Identifying Functionally Important Features with End-to-End Sparse Dictionary Learning, May 2024. /tthtt/////arx
org/abs/24055 12241 arXiv:2405.1 12241 [cs].
Trenton Bricken, Adly Templeton, Joshua Batson, Brian Chen, Adam Jermyn, Tom Conerly, Nick Turner, Cem Anil, Carson Denison, and Amanda Askell. Towards Monosemanticity: Decomposing Language Models With Dictionary Learning. 2023. Utu.httsss//taaafffmr--rrrrrrr b/2203//oonseemaiiic-eaaauu
Maheep Chaudhary and Atticus Geiger. Evaluating Open- Open-Source Sparse Autoencoders on Disentan- gling Factual Knowledge in GPT-2 Small, September 2024. URL httt:tttt///iv..ot/ hbsg:abs/2444
04478 arXiv:2409.0444 044477 [cs].
Arthur Conmy, Augustine Mavor-Parker, Aengus Lynch, Stefan Heimersheim, and Adrià Garriga- Alonso. Towards Automated Circuit Discovery for Mechanistic Interpretability. Advances in
Neural Information Processing Systems. 06:161338.. 16352, December 2023. /RLhhttps://pro
2be-Abstractt-oonfeeeee..
Hoagy Cunningham, Aidan Ewart, Logan Riggs, Robert Huben, and Lee Sharkey. Sparse Au- toencoders Find Highly Interpretable Features in Language Models, October 2023. URL
http:://arxiv.org/abbssss 0866 08600 arXiv:2309.08600 [cs].
Jacob Dunefsky, Philippe Chlenski, and Neel Nanda. Transcoders Find Interpretable LLM Feature Circuits, June 2024. URL http/htt:///arii..om httrgg/bs//4000 11944 arXiv:2406.1 1944 [cs].
Nelson Elhage, Neel Nanda, Catherine Olsson, T. Henighan, N. Joseph, B. Mann, A. Askell, Y. Bai, A. Chen, and T. Conerly. Mathematical Framework for Transformer Circuits, 2021. URL
tttts://trannsfommrr-icrcoo putplls.pppppee//eee http:////////aaa/oo///////////aa html

---

Published as a conference paper at ICLR 2025
Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac Hatfield-Dodds, Robert Lasenby, Dawn Drain, Carol Chen, Roger Grosse, Sam McCandlish, Jared Kaplan, Dario Amodei, Martin Wattenberg, and Christopher Olah. Toy Models of Super- position, September 2022. URL [es].
10652 arXiv:2209.1665
Joshua Engels, Isaac Liao, Eric J. Michaud, Wes Gurnee, and Max Tegmark. Not All Language Model Features Are Linear, May 2024. URL http:///aaxii..oorra//////////// 14860 arXiv:2405.14800
[cs].
Javier Ferrando, Gabriele Sarti, Arianna Bisazza, and Marta R. Costa-jussà. A Primer the Inner Workings of Transformer-bas based Language Models, May 2024. URL http:///axxiv org/orgg///..oo....o.... 24 05.00208 arXiv:2405.0022 000008 [cs].
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace Hc, Anish Thite. Noa Nabeshima, Shawn Presser, and Connor Leahy, The Pile: An 800GB Dataset of Diverse Text for Language Modeling, December 2020. URL http:
//axxiv.org/abs//1111 00027, arXiv:2101.00027 [es].
Leo Gao, Tom Dupré la Tour, and Jeffrey Wu. httpaa//sppasc_auuttoooooo December 2023. URL https://gittub. com/ htt.///eeaaaaa/aaa http://paareeeeeeeeeeeeeeeeeeeee
Leo Gao, Tom Dupré Tour, Henk Tillman, Gabriel Goh, Rajan Troll, Alec Radford, Ilya Sutskever, Jan Leike, and Jeffrey Wu. Scaling and evaluating sparse autoencoders, June 2024. URLhttt
//arxiv org. abs 2406. 04093, arXiv:2406.04093 [cs].
Jorge
Alejandro Maté, and Juan Carlos Trujillo. How does GPT-2 Predict Acronyms? Extracting and Understanding a Circuit via Mechanistic Interpretability. In Proceed- ings of The 27th International Conference Artificial Intelligence and Statistics, pp. 3322-3330. PMLR, April 2024. URL https:: http::/mooeeeen........rssss/sss /
tml
ISSN: 2640-3498.
Aaron Grattafiori Abhimanyu Dubey, Abhinav Jauhri, et al. The Llama 3 Herd of Models, November 2024. httpt::://aaxiv.oor/bbssss/ 21788
Zhengfu He, Xuyang Ge, Qiong Tang, Tianxiang Sun, Qinyuan Cheng, and Xipeng Qiu. Dictionary Learning Improves Patch-Frie Free Circuit Discovery in Mechanistic Interpretability: A Case Study on Othello-GPT, February 2024.
12201 arXiv:2402.12200
[es].
Stefan Heimersheim and Alex Tumner, Residual stream norms grow exponentially over the forward pass, 2023. URL http:///wwww....oooooototttttto/
Evan Hernandez, Arnb Sen Sharma, Tal Haklay, Kevin Meng, Martin Wattenberg, Jacob Andreas, Yonatan Belinkov, and David Bau. Linearity of Relation Decoding Transformer Language Models, February 2024. URL
abrrg//bb/2288...m 09124 arXiv:2308.09124 [cs].
Jing Huang, Zhengxuan Wu, Christopher Potts, Mor Geva, and Attieus Geiger. RAVEL: Evaluating Interpretabilityy Methods on Disentangling Language Model Representations, August 2024. URL http://arxiv. http//bb//22005 17700 arXiv:2402.17700 [cs].
A. Hyvärinen and E. Oja. Independent component analysis: algorithms and applications. Neural Networks, 13(4):411-430 June 2000. ISSN 0893-6080. doi: 10. 10100066/009334400000000006
www. UR. wttpiieeerrrttcccceeeeeeeeeeeeee com/
Stanislaw Jastrzębski, Devansh Arpit, Nicolas Ballas, Vikas Verma, Tong Che, and Yoshua Bengio. Residual Connections Encourage Iterative Inference, March 2018. httt:://axxi..orrt
s/1710. 04773 arXiv:1710.04773 [cs].
Diederik P. Kingma and Jimmy Ba. Adam: Method for Stochastic Optimization, January 2017. h http::/ttt/iiii httt:::/aaxxi..rr///b///11111 6980 arXiv:1412.6968 [cs].
12

---

Published as a conference paper at ICLR 2025
Connor Kissane, Robert Krzyzanowski, Joseph Isaac Bloom, Arthur Conmy, and Neel Nanda.
Interpreting Attention Layer Outputs with Sparse Autoencoders. June 2024. URL https
Kishore Konda, Roland Memisevic, and David Krueger. Zero-bias autoencoders and the benefits of arXiv:1402.333.
co-adapting features, April 2015. URL httpp://arx.w..rgrabb/////a [es, stat].
Vedang Lad, Wes Gurnee, and Max Tegmark. The Remarkable Robustness of LLMs: Stages of
Inference?, June 2024. URL
htrrgabb/24066 19384 arXiv:2406.. 19384 [cs].
Quoc Le, Alexandre Karpenko, Jiquan Ngiam, and Andrew Ng. ICA with Reconstruction Cost for Efficient Overcomplete Feature Learning. In Advances Neural Information Processing Systems, volume 24, Curran Associates, Inc., 2011. URL http:/ttp/p//ccccenniiii/ newr httpsscccpppp////// 01 0117hhhh/23550077444477772222333
Honglak Lee, Alexis Battle, Rajat Raina, and Andrew Ng. Efficient sparse coding algorithms. In Advances Neural Information Processing Systems, volume 19. MIT Press, 2006. URL
15877c9911
2cobbdde2bbdd--bsttrrtt....
Tom Lieberum, Senthooran Rajamanoharan, Arthur Conmy, Lewis Smith, Nicolas Sonnerat, Vikrant Varma, János Kramár, Anca Dragan, Rohin Shah, and Neel Nanda. Gemma Scope: Open Sparse Autoencoders Everywhere All Once on Gemma 2, August 2024. JRL.htttp:/arrxi.orrg s/2408. 05147 arXiv:2408.05147 [cs].
Alcksandar Makelov. Sparse Autoencodeers Match Supervised Features I0I Task. June 2024. http://::...ppoeeeeeeeeee//////a
Features for Model Steering on the
Alireza Makhzani and Brendan Frey. k-Sparse Autoencoders, March March 2014. URL http:://txxii....//////////// g/abs/1312.5663 arXiv:1312.5663 [cs].
Samuel Marks, Can Rager, Eric J. Michaud, Yonatan Belinkov. David Bau, and Aaron Mueller. Sparse Feature Circuits: Discovering and Editing Interpretable Causal Graphs Language Models, March 2024. hRtphhh/////rriiorrrr abs 2403. 19647, art arX1v:2403.19647 [cs].
Andrew Ng. Sparse autoencoder,
3-21 2222=srrnnn//feerreenddaaasssss
nostalgebraist. Interpretnn GPT: the logit lens, August 2020. URL https://ww.leeswrrong pott//AABBooaalirrr//eeeeeeeeeep
Chris Olah. The Next Five Hurdles, July 2024. URL htt.:ttss/ttaassooeeerrrrrrrttrs
/july-uudate/indexx httpr////inndee.ttttthuddddd... htw.##urdlesss
Chris Olah, Nick Cammarata, Ludwig Schubert, Gabriel Goh, Michael Petrov, and Shan Carter.
Zoom In: An Introduction to Circuits. Distill, 5(3), March 2020. ISSN 2476-0757. doi: 10.23915 diistil.00022.001.
pub.:82//ccrrrutts/ooo URL/https: ott.lpppss/iiillllooooooaooarrro
Bruno A. Olshausen and David J. Field. Emergence of simple-cell receptive field properties by
learning a sparse code for natural images. Nature, 381 (653))666776 June 1996. ISSN 1476- 4687. doi: 00111008/8667/00 URL httt http://wwwaatre.ccooooo httt:://www....tttrrrcm///////ss
Publisher:
Nature Publishing Group.
Charles O'Neill and Thang Bui. Sparse Autoencoders Enable Scalable and Reliable Circuit Iden- tification in Language Models, May 2024. URL http://arxlv.org//bb//222 12522 arXiv:2405. 12522 [cs].
Charles O'Neill, Christine Ye, Kartheik Iyer, and John F. Wu. Disentangling Dense Embed- dings with Sparse Autoencoders, August 2024. URL http: http:://arxx..orr/ab////a00 00657 arXiv:2408. 00650 [es].
13

---

Published as a conference paper at ICLR 2025
Kiho Park, Yo Joong Choe, and Victor Veitch. The Linear Representation Hypothesis and the Geometry of Large Language Models, November 2023. URL http:://axiii..og//aasssss/
3658 arXiv:2311 1ww:23111.03588 [cs, stat].
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language
Models are Unsupervised Multitask Leamerr, 2019. URthhttps:/cch http:://cch..pppaaa..cooe/eeeeee com/better-1 miguugge-modelss//ngggggee http://www....aoooooooo.oououorr
pdf
Senthooran Rajamanoharan, Arthur Conmy, Lewis Smith, Tom Lieberum, Vikrant Varma, Janos Kramar, Rohin Shah, and Neel Nanda. Improving Sparse Decomposition of Language Model Activations with Gated Sparse Autoencoders. June 2024a. URRUhttss///opeeeeiiee... net/for
um?id=Pppj5KvZU8Q
Senthooran Rajamanoharan, Tom Lieberum, Nicolas Sonnerat, Arthur Conmy, Vikrant Varma, János Kramár, and Neel Nanda. Jumping Ahead: Improving Reconstrruction Fidelity with JumpRel.U Sparse Autoencoders, July 2024b. URL. /tt..tttp//arrvv..r///////////// 14435, #xX442 [cs].
Morgane Riviere, Shreya Pathak, Pier Giuseppe Sessa, et al. Gemma 2: Improving Open Language Models a Practical Size, October 2024. htttp://aaxwi.org//aaaaaaa
00010 18
Lee Sharkey, Dan Braun, and Beren Millidge. Taking features out of superposition with sparse autoencoders, December 2022. URL http::
//tt:p:::wwww/ //www.
Nishant Subramani, Nivedita Suresh, and Matthew Peters. Extracting Latent Steering Vectors from Pretrained I Language Models. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.), Findings of the Association for Computational Linguistics: ACL 2022, pp. 566-581, Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 1/t.65///12222..nnnnnmg... org/ 2022. findings acl. 48
URL
Adly Templeton, Tom Conerly, Jonathan Marcus, Jack Lindsey, Trenton Bricken, Brian Chen, Adam Pearce, Craig Citro, Emmanuel Ameisen, Andy Jones, Hoagy Cunningham, Nicholas L. Turner, Callum McDougall, Monte MacDiarmid, Alex Tamkin, Esin Durmus, Tristan Hume, Francesco Mosconi, C. Daniel Freeman, Theodore R. Sumers, Edward Rees, Joshua Batson, Adam Jermyn, Shan Carter, Chris Olah, and Tom Henighan. Scaling Monosemanticity: Extracting Interpretable ppbp
Features from Claude 3 Sonnet, May 2024. URL ieselinggmmooosmmmniicttttiiiiin html
Lucrezia Valeriani, Diego Doimo, Francesca Cuturello, Alessandro Laio, Alessio Ansuini, and Alberto Cazzaniga. The geometry of hidden representations of large transformer models. Advances
in Neural Information Processing Systems, 36:51234 51252, December 2023. URL https neurips http://ww..c.popooe////////////e
Addoo55ddeeAAsttaat--Crrrreeeeee html
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones. Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is All you Need. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL httpl//htttt////ett htttpps:/ppppss....ttt cc aewoffodeee31111153aa3333333 c4a845aa Abstract.
Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Inter- pretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small, November 2022. URL httt:://arxxv.org/abbb////// hTgg/bbs/2211.11 2211 00593 arXiv:2211.00590 [cs].
Martin Wattenberg and Fernanda Viégas. Relational Composition in Neural Networks: A Survey and Call to Action. June 2024. URL http::/ppppeeeeeeeeeeeee//o/r/nr
James C. R. Whittington. Will Dorrell, Surya Ganguli, and Timothy E. Behrens. Disentanglement with Biological Constraints: A Theory of Functional Cell Types, March 2023. URL ht 7/arxiv T/arxiv.org/abs//211..oo 01768. arXiv:2210.01768 [cs, q-bio).
14

---

Published as a conference paper at ICLR 2025
John Wright and Yi Ma. High- Diigh-Diimnsiooaal Data Analysis with Low- Low-Dimensional Models: Principles, Computation, and Applications. Cambridge University Press, edition, January 2022. ISBN
978-1-108--7993-22988888 178---18--977 doi: 10.1017/978 0
URL //w::///wwwww camb
Zeyu Yun, Yubei Chen, Bruno Olshausen, and Yann LeCun. Transformer visualization via dictionary learning: contextualized embedding as linear superposition of transformer factors. In Eneko Agirre, Marianna Apidianaki. and Ivan Vulić (eds.), Proceedings of Deep Learning Inside Out (DeelJO): The 2nd Workshop on Knowledge Extraction and Integration for Deep Learning Architectures, pp. 1-10, Online, June 2021. Association for Computational Linguistics. doi: 10.18653/v deelio- 1.1. URL https: / http:/aalaanttc /rrp://llaanoooloooor/////2/// deelio-1.1 2o1..oooo.................
15

---

Published as a conference paper at ICLR 2025
A TRAINING
We train each autoencoder on 1 billion tokens from the Pile (Gao et al. 2020). excluding the copyrighted Books3 dataset for a single epoch. Specifically, we concatenate a batch of 1024 text samples with the med-ef-sentence token, tokenize the concatenated text, and divide the output into sequences of 2048 tokens, discarding the final incomplete sequence. We use an effective batch size of 131072 tokens (64 sequences) for all experiments.
We do not compute activation vectors and cache them to disk before training, which minimizes storage overhead at the expense of repeated computation. We construct a batch of activation vectors to input to the autoencoder by performing the forward pass of the underlying transformer for sequence of tokens, collecting the residual stream activation vectors at every layer, and stacking them together. Following Lieberum et al. (2024), we exclude activation vectors corresponding to special tokens (end-of-sentence, beginning- of-sentence, and padding). Hence, each batch has equal number of activation vectors from each layer, which is the number of non-special tokens.
Following the optimization guidelines in Bricken et al. (2023); Gao et al. (2024). initialize the pre-encoder bias bpee to the geometric median of the first training batch; we initialize the decoder weight matrix Wiec the transpose of the encoder Wenei norm at initialization and after each training step: and we remove the component of the gradient of the decoder weight matrix parallel its weight vectors after each training step.
scale the decoder weight vectors unit
We use the Adam optimizer (Kingma & Ba 2017 with the default B3 parameters, a constant learning rate of 1 x 10-1, and e 6.25 10-10 Unlike Gao al. (2024), we do not use gradient clipping or weight averaging. and we use FP16 mixed precision to reduce memory use.
B EVALUATION
B.1 RECONSTRUCTION ERROR AND SPARSITY
While we use the FVU instead of MSE as the reconstruction error in the training loss, we record both metrics for the inputs from each transformer layer and the mean over all layers (Figure 88. The I norm of the latents is fixed at k per activation vector (layer and token), but we record the L norm (Figure 9). We report the values of these metrics over one million tokens from the test set.
For Pythia-70m, the FVU at each layer is comparable to Marks ef al. 2024 p. 21), who trained separate SAES with 32768 and I norms between 54 and 108, as well as Cunningham al. (2023) p. 13). For Pythia-160m the FVU is similar to Gao et al. 2024) who report the normalized MSE on layer 8 of GPT-2 small.
B.2 DOWNSTREAM LOSS
In addition to the increase in cross-entropy (CH) loss, record the Kullhack-Leibler (KI.) divergence between probability distributions when the residual stream activations at a given layer are replaced by
their reconstruction (Section 4.1) We report the values of these metrics
million tokens from
the test set (Figure IO). The increase in cross-entropy loss is comparable to MMakss et al. (2024) p. 21)
for Pythia-70m. Gao et al. (2024) p. 5) and Braun et al. (0044) for GPT-2 small, and Lieberum et al. 7-8) for layer 20 of Gemma 2 2B and 9B.
(2024
B.3 OTHER TRANSFORMER ARCHITECTURES
We predominantly study transformers from the Pythia suite because it affords a controlled setting to study how our results scale across model sizes (Biderman al.) (023). While we do not expect our results to depend strongly on the underlying transformer architecture, we validated this assumption by training MLSAEs on Gemma 2 2B (Riviere et al. 2024), Llama 3.2 3B (Grattafiori et al. 2024). and GPT-2 small (Radford et al. 2019) with our default perrpaaaaeterss i.c., an expansion factor of R 64 and sparsity k 32. We note that Gemma 2B and Llama 3.2 3B are similar in size to 2ythia-2.8b, and GPT-2 small is similar in size to 160m.
https::/hugging: http:///www..t.aaaaaaa////////aa
16

---

Published as a conference paper at ICLR 2025
Sparsity = 32
Expansion Factor I = 64
10-
FVU
10
Layer
Layer
10
Layer
Layer
Layer
Layer
Mean
MSE
10-3
L.I.
10"
10
10a
10
10
10
Expansion Factor
Sparsity
(a) Pythia-70m
Sparsity 32
Expansion Factor R
10 10-0.5
Layer 0
PVU
10-1
Layer
Layer
Layer 3
Layer 4
10
Layer
Layer
100
Layer
Layer
Layer
MSE
10
Layer 10
Layer 11
Mean
10-2
10
10
10a 10
10a
10a
Expansion Factor R
Sparsity &
(b) Pythia-160m
Figure 8: With fixed sparsity k = 32, the FVU and MSE generally decrease as the expansion factor R increases. For inputs from the last layer, they increase for the largest expansion factors, which attribute to fluctuations in the percentage of dead latents (Figure TT). With fixed expansion factor R 64. the FVU and MSE decrease as the sparsity increases. While all inputs are standardized before passing them to the encoder, the decoder outputs are rescaled afterward. Hence, the MSE increases across layers because is not divided by the variance of the inputs.
17

---

Published as a conference paper at ICLR 2025
Sparsity to
32
Expansion Factor
100
10a
90
Layer 0
80
Layer
100
Layer 2
Layer 3
Layer 4
70
Layer
Mean
10a
60
10
10
10
10a
10a
10a
Expansion Factor R
Sparsity &
(a) Pythia-70m
Sparsity & 32
Expansion Factor R 64
120
10
Layer 0
Layer
Layer
Layer
Layer 4
morm
Layer 5
Layer
Layer 7
80
Layer 8
Layer 9
10a
Layer 10
Layer 11
Mean
60
10
10a
10a
10a
10a
10a
Expansion Factor R
Sparsity
Pythia-
Pythia-160m 160m
Figure 9: With fixed sparsity & = 32, the I norm per token (the sum of absolute activations) generally decreases as the expansion factor increases. With fixed expansion factor R 64, the L norm increases as the sparsity increases. Recall that the L norm per token (the count of non-zero activations) is fixed at k.
18

---

Published as a conference paper at ICLR 2025
Sparsity & 32
Expansion Factor R = 64
10
Delta CB loss
10-2
Layer
Layer
Layer 2
Layer 3
10
Layer
KL divergence
Layer
Mean
10a
10a
10
10
10
10
10
10a
Expansion Factor R
Sparsity &
(a) Pythia 70m
Sparsity A = 32
Expansion Factor R 64
Delta CE loss
10
Layer 0
Layer
Layer 2
10-1
Layer
Layer 4
Layer 5
Layer
10
Layer
divergence
Layer 8
Layer 9
Layer 10
10
Layer 11
Mean
10a
10"
10
10a
10'
10a
10a
Expansion Factor
Sparsity
(b) Pythia-160m
Figure 10: With fixed sparsity A: = 32, the delta CE loss and KL divergence generally decrease the expansion factor increases, except for inputs from the last layer. With fixed expansion factor 64, both metries decrease the sparsity & increases, similarly to the FVU and MSE (Figure 8
19

---

Published as a conference paper at ICLR 2025
0.15
Layer FVU
0.1
0.05
R
Auxiliary Loss
R = 4
0.06
R
0.04
0.02
R 32
R 64
0
R 128
256
10%
Dead Latents
5%
2 x 10
4 x 10a
6 10
8 x 10
1 10
Number of Tokens ner
Figure 11: An illustration of the FVU for inputs from the last layer, compared to the auxiliary loss and percentage of dead latents, for MLSAEs trained on Pythia-70m with fixed sparsity k = 32. An increase in dead latents correlates with a decrease in the auxiliary loss and an increase in the FVU at the last layer. We attribute this to the increased scale of the inputs because the auxiliary loss depends on the MSE (Figure 8 The auxiliary loss is multiplied by its coefficient a 1/32 the training loss.
We include quantitative results for Gemma 2 2B, Llama 3.2 3B, and GPT-2 small in Table 1 Figure 5 and throughout Appendix B We include heatmaps of the distributions of latent activations over layers in Figures [3
15 16 [7, and [8] which qualitatively similar to the Pythia models. Interestingly, we observed that our validation metrics improved at a slower rate for Gemma 2 2B, so we continued train this MLSAE up to total of approximately 2.5 billion tokens. Nevertheless, the increase in the cross-entropy loss remains larger than other models of similar sizes (Table
14
B.4 SINGLE-LAYER SAES
While we compare the performance of our multi-layer SAEs single-layer SAEs from the literature in Appendix B.I and B.2 we also trained multiple single-layer SAES on Pythia-70m, 160m, and 410m with our default hyperparameters, leaving the remainder of the experimental setup unchanged.
Predictably, we find that a layere-layer SAE trained on data from a given layer performs best on test data from the same layer (Figures 20 2T]. A multi-layer SAE trained on data from every layer performs comparably to the corresponding single-layer SAF, and more consistently across test data from different layers. Interestingly, applying the corresponding tuned-lens transformation to the input activations from each layer during training and evaluation degrades the performance of single-layer SAES on test data from different layers of Pythia-70ne unlike multi-layer SAEs (Figure (2).
Importantly, the results for the last layer are excluded from these figures. This is because we take the residual stream activation vectors after a given layer has been applied (Section (.I). such that the last-layer activations represent only the next-token predictions of the model and not intermediate computational variables. Hence, we expect these activations to have a significantly different structure to the preceding layers, which could distort comparisons between layers (Lad et al. 2024
20

---

Published as a conference paper at ICLR 2025
Sparsity &
32
Expansion Factor
10
10-1
FVU
10-2
10
10
MSE
10
10
10a
L' norm
70
10a
Delta CE loss
2
0
10
10a
KL divergence
10a
10a
10a
10a
10a
10A
10a
10a
10a
10a
Expansion Factor R
Sparsity k
Pythia-70m
Pythia- 70m with tuned lens
Figure 12: For Pythia-70m, applying tuned-lens transformations decreases the mean FVU and delta cross-entropy loss but not the KL divergence. Importantly, compute reconstruction errors before applying the inverse transformation and downstream loss metrics afterward (Section 3.3). Unlike Figure 10 linear scale for the delta cross-entropy loss because, surprisingly, it is negative for tuned-lens MLSAEs with large expansion factor or sparsity k; see also Figure 22
21

---

Published as a conference paper at ICLR 2025
GPT-2 small
Layer
Latent
Figure 13: Heatmaps of the distributions of latent activations over layers when aggregating over 10 million tokens from the test set. Here, we plot the distributions for MLSAEs trained on GPT-2 small with expansion factor of R = 64. We provide further details in Figure 2
GPT-2 small
11
Layer
0
Latent
Figure 14: Heatmaps of the distributions of latent activations over layers for single example prompt.
Here, we plot the distributions for MLSAFs trained on GPT-2 small with an expansion factor of
R 64. The example prompt is "When John and Mary went to the store, John gave" (Wang et al.
2022
We provide further details Figure 3
Gemma 2B
25
Layer
0
Latent
Figure 15: Heatmaps of the distributions of latent activations over layers when aggregating over 10 million tokens from the set. Here, we plot the distributions for MLSAEs trained on Gemma 2 2B with an expansion factor of R = 64. We provide further details in Figure 2
Gemma 2B
Layer
0
Latent
Figure 16: Heatmaps of the distributions of latent activations over layers for single example prompt.
Here, we plot the distributions for MLSAEs trained on Gemma 2B with an expansion factor of R 64. The example prompt is "When John and Mary went to the store, John gave" (Wang et al.)
2022 We provide further details in Figure
22

---

Published as a conference paper at ICLR 2025
Llama 3.2 3B
27
Layer
0
Latent
Figure 17: Heatmaps of the distributions of latent activations over layers when aggregating over 10 million tokens from the 3B with an expansion factor of R = 64. We provide further details in Figure 2 Notably, greater proportion of latents are only active at the first layer compared with other transformer architectures.
Here, we plot the distributions for MLSAEs trained on Llama 3.2
Llama 3.2 3B
27
Layer
0
Latent
Figure 18: Heatmaps of the distributions of latent activations over layers a single example prompt.
Here, we plot the distributions for MLSAEs trained on Llama 3.2 3B with an expansion factor of
R 64. The example prompt is "When John and Mary went to the store, John gave" (Wang et al.)
2022 We provide further details in Figure
FVU
Delta Cross Cooss-Entropy Loss
0.8
0.067
0.380
0.634
0.660
0.676
0.182
1.302
4.466
4.282
3.666
4
Single- Single-Layer Layer SAES
0.296
0.101
0.473
0.576
0.645
0.6
0.872
0.254
2.305
3.159
3.010
3
0.653
0.427
0.102
0.439
0.609
1.511
0.882
0.542
1.873
2.563
0.4
0.709
0.582
0.338
0.091
0.474
2.560
1.715
1.369
0.499
1.662
0.2
0.865
0.778
0.570
0.404
0.136
3.713
2.826
2.986
1.647
0.309
MLSAE
0.069
0.112
0.121
0.104
0.137
0.161
0.230
0.644
0.502
0.432
0
2
3
4
0
2
3
4
Test Layer
Test Layer
Figure 19: The FVU reconstruction error and delta cross-entropy loss for single-layer SAEs trained each layer of Pythia-70m, compared with single multi-layer SAE trained every layer. The single-layer SAEs trained on data from given layer perform best on test data from the same layer (the diagonal elements of the matrix plot); multi-layer SAE trained data from every layer performs comparably to the corresponding single-layer SAEs (the row beneath the matrix plot).
23

---

Published as a conference paper at ICLR 2025
FVU
Delta Cross-Entrope Loss
1.2
Single- Single-Layel Layer SAEs
1
0.8
0.6
4
6
0.4
2
0.2
10
MLSAE
3
4
5
10
5
10
Test Layer
Test Layer
Figure 20: The FVU reconstruction eтrоr and delta cross-entrope loss for single-layer SAEs trained on each layer of Pythia-166 160m, compared with single multi-layer SAE trained on every layer. The single-layer SAEs trained on data from given layer perform best on test data from the same layer (the diagonal elements of the matrix plot); multi-layer SAE trained on data from every layer performs comparably to the corresponding single-layer SAEs (the row beneath the matrix plot).
FVU
Delta Pross-Entropy Loss
0
2
1.5
4
Single- Layer SABs
6
8
10
12
14
16
0.5
18
2
20
MLSAE
22
0
10 12 Test Layer
18 20 22
12 Test Layer
16 18 20
22
Figure 21: The FVU reconstruction error and delta cross-entropy loss for single-layer SAES trained on each layer of Pythia-410m. compared with a single multi-layer SAE trained on every layer. The single-layer SAEs trained on data from given layer perform best test data from the same layer (the diagonal elements of the matrix plot); a multi-layer SAE trained on data from every layer performs comparably the corresponding single-layer SAES (the row beneath the matrix plot). Interestingly, there is sharp change in the FVU between single-layer SAEs trained on layers and 5.
24

---

Published as a conference paper at ICLR 2025
FVU
Delta Cross- Coos-EEntopyy Loss
0.8
0.068
0.402
0.650
0.706
0.704
0.181
1.307
6.568
7.844
5.712
Single- Layer SAES
0.301
0.105
0.492
0.631
0.681
0.6
0.250
5.769
5.935
4.763
6
Single-Layer
0.647
0.437
0.110
0.492
0.639
5.282
1.071
0.602
2.194
3.230
0.4
0.709
0.592
0.359
0.094
0.483
7.667
5.666
2.705
0.513
1.707
0.2
2
0.867
0.797
0.593
0.409
0.138
8.728
7.838
5.736
1.934
0.396
0
MLSAE
0.031
0.014
0.017
0.022
0.036
1.677
0.052
1.180
0.177
0.473
0
1
2
3
4
0
1
2
3
4
Test Layer
Test Layer
Figure 22: The FVU reconstruction error and delta cross-entropy loss for single-layer SAEs trained on each layer of Pythia- 70m compared with single multi-layer SAE trained on every layer, applying tuned-lens transformations during training and evaluation (Section 3.3).
B.5 MEAN MAX COSINE SIMILARITY
Sharkey al. (2022) define the Mean Max Cosine Similarity (MMCS) between learned dictionary X and a ground-truth dictionary X'. There is no ground-truth dictionary for language models, so larger learned dictionary or the k nearest neighbors each dictionary element are commonly used.
1
MMCS(X,X))
max sim(x,x') x'EX'
(13)
X
The MMCS serves as a proxy measure for 'feature splitting' (Bricken et al. 2023 Braun et al. 2024); as the number of features increases, expect the decoder weight vectors be more similar to their nearest neighbors. We compute the MMCS with & = 1 after training, finding decreases slightly as the model size increases with fixed hyperparameters (Table 2
B.6 PAIRWISE COSINE SIMILARITIES
A potential issue when training multi-layer SAEs is learning multiple versions of 'the same' latent, i.e., multiple latents with similar interpretable functions but which are active at different layers. In this case, we would expect to find pairs of latents with relatively large cosine similarities between their decoder weight vectors but different observed distributions of activations over layers (Section 4.3). We investigated this possibility by comparing the distributions of pairwise cosine similarities between decoder weight vectors for trained MLSAEs reference distributions.
As negative control, we generated equal number (the number of latents n) of normal independently and identically distributed (i.i.d.) vectors x AO I) of the same length (the model dimension d). In this case, the pairwise cosine similarities follow a normal distribution cos sim(x, x') _N(0, 1/a). As positive control, we generated smaller number of normal i.i.d. vectors (the number of latents divided by the number of layers nt). copied the vectors n, times, and added noise ~/ 1) to each copy. In this case, we expect an additional frequency peak for large positive cosine similarities.
Figure shows that the distributions of pairwise cosine similarities for decoder weight vectors slightly
and right-shifted compared with the negative control, i.e., a pair of MLSAE latents are slightly more likely to have high cosine similarity than pair of i.i.d. normal vectors. However, the number of pairs with large positive cosine similarities is small compared to the positive control, which has second peak around 0.5 (visible only with the logarithmic y-axis scale).
25

---

Published as a conference paper at ICLR 2025
Model
Mean
Std. Dev.
Pythia-70m
0.275
0.0843
Pythia-160m
0.250
0.0928
Model
Mean
Std. Dev.
Pythia-410m
0.221
0.0868
Pythia-70m
0.261
0.0763
Pythia-1b
0.201
0.0989
Pythia- 160m
0.206
0.0734
Pythia-1.4b
0.180
0.0861
Pythia-410n 40m
0.216
0.0864
Gemma 2 2B
0.249
0.1052
Llama 3.2 3B
0.215
0. 0.1187
(b) Tuned lens
GPT-2 small
0.258
0.0703
(a) Standard
Table 2: The mean and standard deviation of the maximum cosine similarity between decoder weight vectors for MLSAEs with an expansion factor It 64 and sparsity k 32. The MMCS decreases as the model size increases for Pythia transformers.
Sparsity
Expansion Factor R 64
Mean Max Cosine Similarity
0.2
0.15
10"
10
10a
10
10a Sparsity &
10a
Expansion Factor R
Pythia-70m
Pythia-70m with tuned lens
Pythia-160n
Pythia- with tuned lens
Pythia-410m
Pythia-410m with tuned lens
Pythia-lb
Pythia-
Gemma 2 2B
Llama 3.2 3B
GPT-2 small
Figure 23: The Mean Max Cosine Similarity between decoder weight vectors for standard and tuned-lens MLSAEs. The MMCS increases as the expansion factor R increases and decreases as the sparsity k increases. Applying tuned-lens transformations slightly decreases the MMCS relative to standard MLSAES.
26

---

Published as a conference paper at ICLR 2025
-10
Pythia-70m 70m
Pythia-70m
10
Frequency
10a
2
10
-10
Pythia-160m
Pythia-160m
10
Frequency
10a
0.5
10
-108
Pythia-410m
Pythia-410m
10
Frequency
10a
ky
109
-109
Pythia- 1b
Pythia- 1b
10
1.5
Frequency
10a
0.5
0
10
MMA
-109
Pythia-1.44 1.4b
Pythia- 1.4b
10
1.5
Frequency
10a
0.5
10a
0.1
0.2
0.5
0
0.5
Cosine Similarity
Cosine Similarity
Decoder weight vectors
Paired normal i.i.d. vectors plus noise (positive control)
Normal i..dd. vectors (negative control)
Figure 24: Histograms of the frequencies of pairwise cosine similarities between decoder weight vectors, compared to an equal number of normal i.i.d. vectors of the same length, and n, copies of smaller number normal i.i.d. vectors with added noise N/0, 1). Here, we report the frequencies for MLSAEs trained on Pythia models with an expansion factor of R. 64 and sparsity i 32. The left-hand y-axis scale is linear, the hand is logarithmic.
27

---

Published as a conference paper at ICLR 2025
C MEASURES OF LATENTS ACTIVE AT MULTIPLE LAYERS
C.I HEATMAP NORMALIZATION
we plot the distributions of latent activations over layers, which we take be proportional the total activations when aggregating over large sample of tokens (Eq. I0). We normalized the latent activations in this way visually compare the aggregate and single-prompt heatmaps, as well as individual latents within heatmap, due to the wide range of activation counts and totals across latents.
In the aggregate and single-prompt heatmaps such as Figures 2 and 3
Normalizing the activations discards the relative frequencies and magnitudes of activations for different latents, so we reproduce Figures and and with the un-normalized totals of latent activations We use power-law normalization for the colormaps, i.e., y = x²', to account for the wide range of values; all other heatmaps have linear colormaps. As with all other single-prompt heatmaps, we exclude latents from Figure 26 that never activate. In both cases, the un-normalized results are qualitatively similar to the normalized results.
in Figures and
C.2 VARIANCE OF THE LAYER INDEX
Recall that we consider the layer L, token T, and latent index as random variables (Section 4.3) For single latent, we have, by the law of total variance:
Var [4] E[Var [L []] + Var ELL T]]
(14)
We are interested in the first two terms:
Var [L] is the variance of the distribution over layers, aggregating over tokens;
V[a [& [T] is the mean variance of the distributions over layers for each token; and
Var [B[4 [] is the variance of the mean layers for each token.
Aggregating over latents, we have:
E[Var [L] | J]] = E[Var [L TJ E|var EIL TT
(15)
C.3 NUMBER OF LAYERS ABOVE A THRESHOLLD
The count of layers at which a latent is non-zero ('active') does not necessarily positively correlate with the variance of the layer index considered in Section 4.3 For example, the variance of and (two distinct values) is greater than the variance of 2, 3, and 4 (three distinct values). Strietly speaking, the layer index is ordinal data, we implicitly treat as interval data by taking the arithmetic mean and variance. We chose this approach because we expected latents to be active over contiguous range of layers, which is validated qualitatively by the heatmaps such as Figures 2 and
For comparison, we computed the number of layers at which each latent has count of non-zero activations above threshold ('active layers') divided by the total number of model layers vız, We selected threshold count of 10k tokens (0.1% of sample of 10M tokens). When aggregating over latents, the relative mean active layers decreases as the model size increases for Pythia models (Figure 28a and the number of latents increases relative to the model dimension (Figure 28b). Importantly, this measure depends strongly on the choice of threshold, unlike our variance ratios.
C.4 ENTROPY
A further measure of the degree to which latent is active at multiple layers is the statistical distance between the observed discrete distribution of activations over layers (Eq, [T]] and a reference distribution. At one extreme is Dirac distribution with probability mass 1 for single layer index and 0 elsewhere, in which case the latent active at a single layer. The other extreme is the discrete uniform distribution (0, nt), ), in which case the latent is equally active at every layer. Hence, the entropy of the observed distribution ranges between O and In n,.. Notably, the entropy of the observed distribution is agnostic with respect to the numeric values of the layer indices and their ordering.
28

---

Published as a conference paper at ICLR 2025
Pythia- 160m 160m
Layer
Pythia-410m
23
Layer
Pythia lb
15
Layer
Pythia- 1.4b
23
Layer
Latent
Figure 25: Heatmaps of the total latent activations over layers when aggregating over 10 million tokens from the test set. Here, we plot the totals for MLSAEs trained on Pythia models with an expansion factor of R = 64 and sparsity i = 32. We provide further details in Figure 2 The colormaps use power-law normalization with 7 1/4.
Pythia- 160m
Layer
0
Pythia-4100 -10m
23
Layer
Pythia 1b
15
Layer
Pythia-1.4e 1.44
23
Layer
Latent
Figure 26: Heatmaps of the total latent activations over layers for single example prompt. Here, we
plot the totals for MLSAEs with an expansion factor of 64 and sparsity 32. The example
prompt is "When John and Mary went to the store, John gave" (Wang et al. 2022) We provide
further details in Figure The colormaps use pawer-law normalization with 7 1/2.
29

---

Published as a conference paper at ICLR 2025
Sparsity 32
Expansion Factor R 64
10
10 10-0.2
Var ( J)
Var (L)
10 10-0.4
10 0.6
10 0.8
(LJ,, T)
Var Var((
E Var(L]]))
10-1
10
10"
10a
10a
10a
10a Sparsity
10
Expansion Factor
Pythia-70m
Pythia 70m with tuned lens
Pythia- 160m
Pythia- Fithia-160n 160m with tuned lens
Pythia-410m
Pythia-410m with tuned lens Pythia- ..bb
Pythia- 1b
Gemma 2B
Ilama 3.2 3B
GPT-2 small
Figure 27: The fraction of the total variance explained by individual latents and the fraction of the
variance for an individual latent explained by individual tokens (Eqs. 11 and 12 Here, we plot variance ratios for standard and tuned-lens MLSAEs over 10 million tokens from the test set.
30

---

Published as a conference paper at ICLR 2025
Pythia-70m
Pythia-160m
Pythia-410m
Pythia-Ib
Pythia- 1.4b
Gemma 2B
Llama 3.2 3B
Standard
GPT-2 small
Tuned lens
0
0.02
0.04
0.06
0.08
0.1
0.12
0.14
0.16
0.18
0.2
Mean Active Layers Total Layers
(a) Varying the model with expansion factor of 64 and sparsity
32
Sparsity 32
Expansion Factor
64
10"
Mean Active Layers
10-1
Total Layers
10
10
10"
10a
10a
10a
10a
10
Expansion Factor R
Sparsity
Pythia-70m
Pythia-70m with tuned lens
Pythia-160m
Pythia-160m with tuned lens
Pythia-410m
Pythia-410m with tuned lens
Pythia-Ib
Pythia- Aithia-1.4b 1.4b
Gemma 2B
Llama 3.2 3B
GPT-2 small
(b) Varying the expansion factor with sparsity à 32 and & with 64
Figure 28: The mean number of layers at which latents have a count of non-zero activations above threshold divided by the total number of model layers, over million tokens from the test set. The threshold is 10 thousand tokens (0.1%) As in Figure 5 the absence of bars for tuned-lens MISAFs indicates the absence results, not that the values are zero. The mean active layers decreases as the expansion factor R increases, and increases the sparsity & increases up point. There is clear trend with respect to the model size. Notably, applying tuned-lens transformations to the input activations from each layer decreases the mean active layers (Section 3.3
31

---

Published as a conference paper at ICLR 2025
Pythia-70m
Standard Tuned lens
Pythia-160m
Pythia-410m
Pythia- lb
Pythia-1.4p
Gemma 2 2B
Ilama 3.2 3B
GPT-2 small
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
Mcan Entropy In n..
(a) Varying the model with an expansion factor of 64 and sparsity
Sparsity
Expansion Factor
0.8
Mean Entropy
0.7
0.6
0.5
10
10
102
10
10a Sparsity
10a
Expansion Factor
Pythia-70m
Pythia-70m with tuned lens
Pythia- 160m
Pythia- 160m with tuned lens Pythia-410m with tuned lens Pythia-1.4b Llama 3.2 3B
Pythia-410m
Pythia- 1b
Gemma 2B
GPT-2 small
(b) Varying the expansion factor with sparsity 32 and k with 64
The mean entropy of the observed discrete distributions of latent activations over layers (Eq. 10 divided by the maximum entropy In n, over 10 million tokens from the test set. As in Figure 5 the absence of bars for tuned-lens MLSAEs indicates the absence of results, not that the values are zero. The entropy tends to decrease as the expansion factor increases, to increase as the sparsity ie increases, and to increase the model size increases for Pythia transformers.
Figure 29:
We computed the entropy of the observed distributions of activations over layers, took the mean over latents, and divided it by In ny to compare models with different numbers of layers. The normalized mean entropy increases slightly the model size increases for Pythia models (Figure (9g), similarly to the variance of the layer index (Section 4.3). However, decreases as the number latents increases relative the model dimension, similarly the mean active layers (Figure 29h).
32

---

Published as a conference paper at ICLR 2025
D ADDITIONAL HEATMAPS
For completeness, we include equivalent aggregate and single-prompt heatmaps to Figures 2 and for different models and combinations of hyperparameters:
Varying R for Pythia-70m and & 32 (Figures 30 and
3
Varying & for Pythia-70m and R = 64 (Figures
32 and 33
Varying R for Pythia-160m and 32 (Figures and 35
Varying & for Pythia-160n and 64 (Figures 36 and 37)
Varying R for Pythia-70m with tuned lens and 32 (Figures 38
39
Varying & for Pythia-70m with tuned lens and R 64 (Figures
40 and 41
33

---

Published as a conference paper at ICLR 2025
R
OINSS C1
Layer
R
Layer
R
OINSA C1
Layer
R
Layer
R 16
OIiSSS C1
Layer
R 32
OINGS G1
Layer
R 64
Layer
R = 128
OINGICI
Layer
R 256
OI2SSCO
Layer
Latent
Figure 30: Heatmaps of the distributions of latent activations over layers when aggregating over 10 million tokens from the test set. Here, we plot the distributions for MLSAEs trained on Pythia-70m with sparsity & = 32. We provide further details in Figure 2
34

---

Published as a conference paper at ICLR 2025
Layer
OIIEE
Layer
CINSAS
Layer
8
OIISTS
Layer
16
OHNSAS
Layer
32
OIISES
Layer
OINSPS
Layer
R
OUNOAS PINGPS
Layer
256
Layer
Latent
Figure 31: Heatmaps of the distributions of latent activations over layers for single example prompt.
Here, we plot the distributions for MLSAES trained on Pythia 70m with sparsity k 32. The example prompt is "When John and Mary went the store, John gave" (Wang et al. 2022). We provide further details in Figure
35

---

Published as a conference paper at ICLR 2025
16
5
1
Layer
3
2
k k= 32
5
Layer
3
2
64
5
Layer
3
= 128
5
4
Layer
3
k = 256
4
Layer
3
2
k = 512
4
Layer
3
2
Latent
Figure 32: Hcatmaps of the distributions of latent activations over layers when aggregating over 10 million tokens from the test set. Here, we plot the distributions for MLSAES trained on Pythia-70m with an expansion factor of R = 64. We provide further details in Figure 2
36

---

Published as a conference paper at ICLR 2025
16
5
Layer
3
2
32
5
Layer
3
64
Layer
3
k = 128
Layer
3
2
= 256
5
4
Layer
3
2
k 512
5
4
Layer
3
2
Latent
Figure 33: Heatmaps of the distributions of latent activations over layers for single example prompt.
Here, we plot the distributions for MLSAEs trained on Pythia-70m with an expansion factor of
R 64. The example prompt is "When John and Mary went to the store, John gave" (Wang et al.
2022
We provide further details in Figure
37

---

Published as a conference paper at ICLR 2025
11
Layer
Layer
R
11
Layer
Layer
16
11
Layer
Layer
64
11
Layer
R 128
11
Layer
R 256
11
Layer
Latent
Figure 34: Heatmaps of the distributions of latent activations over layers when aggregating over 10 million tokens from the test set. Here, we plot the distributions for MLSAEs trained Pythia- 160m with sparsity & = 32. We provide further details in Figure 2
38

---

Published as a conference paper at ICLR 2025
Layer
Layer
Layer
Layer
16
Layer
32
Layer
=
Layer
R = 128
Layer
R 256
Layer
0
Latent
Figure 35: Heatmaps of the distributions of latent activations over layers for single example prompt.
Here, we plot the distributions for MLSAEs trained on Pythia- 160m with sparsity 32. The example prompt is "When John and Mary went to the store, John gave" (Wang et al. 2022). We provide further details in Figure
39

---

Published as a conference paper at ICLR 2025
16
11
Layer
= 32
Layer
64
Layer
128
Layer
256
Layer
k = 512
Layer
Latent
Figure 36: Hcatmaps of the distributions of latent activations over layers when aggregating over 10 million tokens from the test set. Here, we plot the distributions for MLSAEs trained on 1ythia-160m with an expansion factor of At = 64. We provide further details in Figure] 2
40

---

Published as a conference paper at ICLR 2025
16
Layer
32
Layer
= 64
Layer
k = 128
Layer
k = 256
Layer
k = 512
Layer
Latent
Figure 37: Heatmaps of the distributions of latent activations over layers for single example prompt. Here, we plot the distributions for MLSAEs trained on Pythia- 160m with an expansion factor of
R 64. The example prompt is "When John and Mary went to the store, John gave" (Wang et al.
2022
We provide further details in Figure
41

---

Published as a conference paper at ICLR 2025
R
OINEII
Layer
R 2
OINOS or
Layer
OIZSIS
Layer
R = 8
Layer
R 16
OIiSS C1
Layer
R = 32
Layer
R = 64
101000 01
Layer
R = 128
Layer
R 256
Layer
Latent
Figure 38: Heatmaps of the distributions of latent activations over layers when aggregating over 10 million tokens from the test set. Here, we plot the distributions for tuned-lens MLSAEs trained on Pythia-70n 70m with sparsity 32. We provide further details in Figure
42

---

Published as a conference paper at ICLR 2025
OINSTS
Layer
OIIEES
Layer
=
CINEAS
Layer
OIISSS
Layer
16
OINSAS
Layer
32
OIISES
Layer
OINGAS
Layer
R
SINGAS
Layer
256
OINGSS
Layer
Latent
Figure 39: Heatmaps of the distributions of latent activations over layers for single example prompt.
Here, we plot the distributions for tuned-lens MLSAEs trained on Pythia 70m with sparsity 32.
The example prompt is "When John and Mary went the store, John gave" (Wang al. 2022 We
provide further details in Figure
43

---

Published as a conference paper at ICLR 2025
16
5
Layer
3
2
= =33
5
Layer
3
2
0
64
5
Layer
3
2
0
= 128
5
4
Layer
3
k- = 256
5
4
Layer
3
2
k = 512
5
4
Layer
3
2
Latent
Figure 40: Hcatmaps of the distributions of latent activations over layers when aggregating over 10 million tokens from the test set. Here, we plot the distributions for tuned-lens MLSAEs trained on 7ythia-70m with an expansion factor of R = 64. We provide further details in Figure 2
44

---

Published as a conference paper at ICLR 2025
16
5
Layer
3
2
32
Layer
3
64
Layer
3
= 128
Layer
3
2
= 256
5
4
Layer
3
2
k 512
5
4
Layer
3
2
Latent
Figure 41: Heatmaps of the distributions of latent activations over layers for single example prompt. Here, we plot the distributions for tuned-lens MLSAEs trained on Pythia-70m with an expansion factor of R 64. The example prompt "When John and Mary went to the store, John gave" (Wang
2022 et al.
We provide further details in Figure 3
45


Reviews:
[{"review_id": "mKHzeFOC1A", "reviewer": "Reviewer_g44V", "summary": "The paper introduces and analyzes a multi-level sparse autoencoder (MLSAE), which is a sparse autoencoder (SAE) trained on latent representations of the residual stream at all layers. The work experiments with Pythia models on language modeling tasks. The paper further reinforces the technique by adding tuned-lens transformations, allowing analysis of the residual stream even if the information is encoded on different dimensions at different layers. The work then studies the behavior and similarity of representation at different layers in Transformers of varying sizes.\n[After the discussion period: I have increased the score from 6 to 8, and increased the confidence from 3 to 4.]", "questions": "Please refer to the weaknesses section, specifically:\n* regarding (1 and 2) - can you provide direct comparison to other methods or explain why such comparisons are absent?\n* regarding (4)  what is the additional overhead when training/using MLSAE?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The work experiments with Pythia models (Transformers) of different sizes, making the work highly reproducible. The experiments are sound and most of the doubts I had while reading were answered already in the text. While the paper doesn't offer revolutionary insights on its own, it is an excellent improvement to sparse autoencoders and may be very useful for mechanistic interpretability. The model also tests the relative performance of MLSAE at different expansion factors and sparsity levels and runs good ablations with \"tuned lens\". While I wondered about some design decisions, in particular \"feature-stacked\" MLSAE, when reading the paper, the discussion provided in Section 5 answered my potential questions before I even wrote them down. All in all, the writing is excellent, and the experiments look solid and reproducible. In general, I am leaning towards acceptance to the conference, and I am willing to adjust the score during the discussion.", "weaknesses": "1. 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.\n2. 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.\n3. 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. The method also is quite simple.\n4. 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MHlU3UFrmT", "reviewer": "Reviewer_WQNj", "summary": "The paper proposes a method based on sparse autoencoders (SAEs) for recovering features and analyzing the distributions of activations of the same features (SAE latents) across layers in Pythia language transformers of different sizes. The SAE is trained in a manner that treats per-token feature vectors (potentially transformed with tuned lens) from all layers as independent training examples. Hence, the method is named Multi-layer SAE.\n\nContrary to the authors' hypothesis that features often remain active at later layers once they are activated, the experiments show activations of the latents are often more localized and active at a single layer.", "questions": "Questions:\n1. How much of computational resources did the experiments require?\n2. L301: \"we found the mean cosine similarities increased accordingly\". This is not true for smaller models. Do you have any intuitions about that?\n3. The analysis is normalizing histograms of activations over layers to get distributions. This loses information about 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?\n4. The paper uses metrics based on 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?\n5. 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?\n\nIf the questions make sense, I would like the final paper to address them.\n\nSuggestions:\n1. Eq (8): Note that the residual connection matters only because of L2 regularization (and because if it effectively changes the initialization).", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well written.\n- The experimental evaluation is well done.\n- The research area of the paper is mechanistic interpretability, which is a potentially very useful for understanding complex models and important for AI safety.", "weaknesses": "1. Some things might be unclear (on the first reading):\n\t- this layer may differ for different tokens or prompts -> which layer it is might differ ... ?\n\t- Under multiple figures: \"expected value of the layer\" -> \"average layer index\" or \"expected value of the layer index\"\n\t- Eq. (3): The meaning of $\\operatorname{Var}$ could be clarified.\n\t- In some places, the notation is a bit inconsistent or non-standard.\n\t\t- In section 3.1, $\\bf x\\in\\mathbb R^d$, but in Eq. (3), $\\bf x$ represents a sequence of vectors $\\bf x_\\ell$.\n\t\t- 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.\n2. Some training details seem missing: How are layers for mini-batches sampled?\n3. 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.\n4. The paper does not state whether the source code for reproducing the experiments will be published.\n\nThere are a few errors:\n- L245: \"ratio between the model dimension and the number of latents\" should be the other way around.\n- L353: Should $\\mathbb E_T[L\\mid J=j]$ be $\\mathbb E[L\\mid J=j]$ (or, equvalently, $\\mathbb E_L[L\\mid J=j]$)?\n\nI apologize for any errors on my part.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CWEDGwgTaa", "reviewer": "Reviewer_F8UN", "summary": "This study introduces a novel approach called MLSAE for analyzing the evolution of representations across various layers in Transformers. The author presents comprehensive experimental evidence to elucidate the statistical characteristics of representations at different layers. Nonetheless, how these statistical analyses can help the understanding of the inherent mechanisms within Transformers not that clear.", "questions": "Please refer to the Weakness Section.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This work proposes the MLSAE method to analyze the interpretability of information transmitted through residual flow in Transformers. \n2. The author only uses one SAE in multi-layer Transformers. \n3. The author analyzes the fraction and variance of latents of models of different sizes through detailed experiments.", "weaknesses": "1. As indicated in both the Abstract and Conclusion sections, the author investigates the transformation of representations across various layers within Transformers. However, the analysis is limited to a statistical perspective, focusing solely on the fractions and variances of representations at different layers. The linkage between this statistical overview and a deeper understanding of the underlying mechanisms of Transformers is not explicitly articulated. It remains unclear to me how these statistical descriptions facilitate insight into the working mechanism of Transformers.\n2. The motivation for the author's selection of Stacked Autoencoders (SAE) as the dimensionality reduction technique is unclear for me. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UGKrRzETxW", "reviewer": "Reviewer_d1f6", "summary": "This paper introduces a novel technique for training a SAE which can be applied to any layer of a model, called a MLSAE, by using training samples drawn from all layers of the model. This is beneficial because it is desirable to be able to see how the active latents change from layer to layer. To address differences in norm between layers, the paper normalizes activations and optionally applies tuned lens to help ensure that activations share the same basis. The paper studies to what extent latents are active across multiple layers as a function of model size, SAE sparsity, and SAE width using Pythia models of various sizes.", "questions": "- The paper mentions not having a pre-trained TunedLens in some cases - where are these pre-trained TunedLenses from?\n- 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\n- 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?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "The problem of studying SAE latents across multiple layers is an important and understudied problem, and the contribution of training a single SAE on activations from multiple layers is novel. Using TunedLens as a pre-processing step is also novel. The paper is clear and understandable.", "weaknesses": "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. There is also no performance comparison to a standard single-layer SAEs to see if MLSAEs result in a performance decrease when evaluated on a single layer. The paper only evaluates on Pythia models and topk SAEs - it would be beneficial to study other models and SAE architectures if possible.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)