Scaling and evaluating sparse autoencoders

# Scaling and evaluating sparse autoencoders Leo Gao ∗ Tom Dupré la Tour † Henk Tillman † Gabriel Goh Rajan Troll Alec Radford Ilya Sutskever Jan Leike Jeffrey Wu † OpenAI ## Abstract Sparse autoencoders provide a promising unsupervised approach for extracting in-terpretable features from a language model by reconstructing activations from a sparse bottleneck layer. Since language models learn many concepts, autoencoders need to be very large to recover all relevant features. However, studying the proper-ties of autoencoder scaling is difficult due to the need to balance reconstruction and sparsity objectives and the presence of dead latents.

We propose using k-sparse autoencoders [Makhzani and Frey, 2013] to directly control sparsity, simplifying tuning and improving the reconstruction-sparsity frontier. Additionally, we find modifications that result in few dead latents, even at the largest scales we tried. Us-ing these techniques, we find clean scaling laws with respect to autoencoder size and sparsity. We also introduce several new metrics for evaluating feature qual-ity based on the recovery of hypothesized features, the explainability of activation patterns, and the sparsity of downstream effects. These metrics all generally im-prove with autoencoder size. To demonstrate the scalability of our approach, we train a 16 million latent autoencoder on GPT-4 activations for 40 billion tokens.

We release code and autoencoders for open-source models, as well as a visualizer. 3 ## 1 Introduction Sparse autoencoders (SAEs) have shown great promise for finding features [Cunningham et, 2023, Bricken et, 2023, Templeton et, 2024, Goh, 2016] and circuits [Marks et, 2024] in language models. Unfortunately, they are difficult to train due to their extreme sparsity, so prior work has primarily focused on training relatively small sparse autoencoders on small language models. We develop a state-of-the-art methodology to reliably train extremely wide and sparse autoencoders with very few dead latents on the activations of any language model.

We systematically study the scaling laws with respect to sparsity, autoencoder size, and language model size. To demonstrate that our methodology can scale reliably, we train a 16 million latent autoencoder on GPT-4 [OpenAI, 2023] residual stream activations. Because improving reconstruction and sparsity is not the ultimate objective of sparse autoencoders, we also explore better methods for quantifying autoencoder quality. We study quantities corresponding > ∗ Primary Contributor. Correspondence to com. > † Core Research Contributor. This project was conducted by the Superalignment Interpretability team. Author contributions statement in Appendix I. > 3 Our open source code can be found at

com/openai/sparse_autoencoder and our visualizer is hosted at html > arXiv:2406.04093v1 LG] 6 Jun 2024 10 910 710 5 > FLOPs (as fraction of GPT-4 pretraining) > 0.2 > 0.3 > 0.4 > 0.5 > 0.6 > Normalized MSE > L(C) > L= 0.09 + 0.056 C0.084 > 10 4 > 10 5 > 10 6 > 10 7 > Number of Latents 10 410 5 > Number of Latents > 3 × 10 1 > 4 × 10 1 > 6 × 10 1 > Normalized MSE > L(N, k) > 32 > 64 > 128 > 256 > 512 > k Figure 1: Scaling laws for TopK autoencoders trained on GPT-4 activations.

(Left) Optimal loss for a fixed compute budget. (Right) Joint scaling law of loss at convergence with fixed number of total latents n and fixed sparsity (number of active latents) k. Details in Section 3. to: whether certain hypothesized features were recovered, whether downstream effects are sparse, and whether features can be explained with both high precision and recall. Our contributions: 1. In Section 2, we describe a state-of-the-art recipe for training sparse autoencoders. 2. In Section 3, we demonstrate clean scaling laws and scale to large numbers of latents. 3. In Section 4, we introduce metrics of latent quality and find larger sparse autoencoders are generally better according to these metrics.

We also release code, a full suite of GPT-2 small autoencoders, and a feature visualizer for GPT-2 small autoencoders and the 16 million latent GPT-4 autoencoder. ## 2 Methods 2.1 Setup Inputs: We train autoencoders on the residual streams of both GPT-2 small [Radford et, 2019] and models from a series of increasing sized models sharing GPT-4 architecture and training setup, including GPT-4 itself [OpenAI, 2023] 4. We choose a layer near the end of the network, which should contain many features without being specialized for next-token predictions (see Section 1 for more discussion).

Specifically, we use a layer 56 of the way into the network for GPT-4 series models, and we use layer 8 (34 of the way) for GPT-2 small. We use a context length of 64 tokens for all experiments. We subtract the mean over the dmodel dimension and normalize to all inputs to unit norm, prior to passing to the autoencoder (or computing reconstruction errors). Evaluation: After training, we evaluate autoencoders on sparsity L 0, and reconstruction mean-squared error (MSE). We report a normalized version of all MSE numbers, where we divide by a baseline reconstruction error of always predicting the mean activations.

Hyperparameters: To simplify analysis, we do not consider learning rate warmup or decay unless otherwise noted. We sweep learning rates at small scales and extrapolate the trend of optimal learning rates for large scale. See Appendix A for other optimization details. 2.2 Baseline: ReLU autoencoders For an input vector x ∈ Rd from the residual stream, and n latent dimensions, we use baseline ReLU autoencoders from [Bricken et, 2023]. The encoder and decoder are defined by: z = ReLU (Wenc (x − bpre) + benc)ˆx = Wdec z + bpre (1) > 4 All presented results either have qualitatively similar results on both GPT-2 small and GPT-4 models, or were only ran on one model class.

210 1 10 2 > Sparsity (L0) > 4 × 10 1 > 5 × 10 1 > 6 × 10 1 > Normalized MSE > better > ReLU > ProLU STE > Gated > TopK (ours) (a) At a fixed number of latents (n = 32768), TopK has a better reconstruction-sparsity trade off than ReLU and ProLU, and is comparable to Gated. 10 4 10 5 > Number of Latents > 4 × 10 1 > 5 × 10 1 > 6 × 10 1 > Normalized MSE > ReLU > ProLU STE > Gated > TopK (ours) (b)

At a fixed sparsity level (L0 = 128), scaling laws are steeper for TopK than ReLU. 6 Figure 2: Comparison between TopK and other activation functions. with Wenc ∈ Rn×d, benc ∈ Rn, Wdec ∈ Rd×n, and bpre ∈ Rd. The training loss is defined by L = || x − ˆx|| 22 + λ|| z|| 1, where || x − ˆx|| 22 is the reconstruction MSE, || z|| 1 is an L 1 penalty promoting sparsity in latent activations z, and λ is a hyperparameter that needs to be tuned. 2.3 TopK activation function We use a k-sparse autoencoder [Makhzani and Frey, 2013], which directly controls the number of active latents by using an activation function (TopK) that only keeps the k largest latents, zeroing the rest.

The encoder is thus defined as: z = TopK (Wenc (x − bpre)) (2) and the decoder is unchanged. The training loss is simply L = || x − ˆx|| Using k-sparse autoencoders has a number of benefits: • It removes the need for the L 1 penalty. L 1 is an imperfect approximation of L 0, and it introduces a bias of shrinking all positive activations toward zero (Section 5.1). • It enables setting the L 0 directly, as opposed to tuning an L 1 coefficient λ, enabling simpler model comparison and rapid iteration.

It can also be used in combination with arbitrary activation functions. 5 • It empirically outperforms baseline ReLU autoencoders on the sparsity-reconstruction frontier (Figure 2a), and this gap increases with scale (Figure 2b). • It increases monosemanticity of random activating examples by effectively clamping small activations to zero (Section 4.3). 2.4 Preventing dead latents Dead latents pose another significant difficulty in autoencoder training. In larger autoencoders, an increasingly large proportion of latents stop activating entirely at some point in training. For example, Templeton et al. [2024] train a 34 million latent autoencoder with only 12 million alive latents, and in our ablations we find up to 90% dead latents 7 when no mitigations are applied (Figure 15).

This results in substantially worse MSE and makes training computationally wasteful. We find two important ingredients for preventing dead latents: we initialize the encoder to the transpose of the decoder, and we use an auxiliary loss that models reconstruction error using the top-kaux dead latents (see Section 2 for more details). Using these techniques, even in our largest (16 million latent) autoencoder only 7% of latents are dead. > 5 In our code, we also apply a ReLU to guarantee activations to be non-negative. However, the training curves are indistinguishable, as the k largest activations are almost always positive for reasonable choices of k.

> 6 Non-TopK cannot set a precise L 0, so we interpolated using a piecewise linear function in log-log space. > 7 We follow Templeton et al. [2024], and consider a latent dead if it has not activated in 10 million tokens 320 21 22 23 24 25 > Learning Rate (x 1e-4) > 4 × 10 1 > 5 × 10 1 > Normalized MSE > 7.9B 7.1B 6.3B 3.9B 3.2B 2.2B 2.0B > 4.0B 3.4B 2.9B 2.7B 2.4B 2.2B 1.7B > 4.0B 3.7B 2.9B 3.2B 1.2B 1.3B 1.2B > 2.3B 2.0B 1.8B 1.8B 1.2B 1.1B 1.0B > 1.1B 1.0B 1.0B 1.0B 1.0B 0.9B 0.6B > 1.3B 1.1B 1.1B 0.9B 0.9B 0.9B 1.2B > 211 > 212 > 213 > 214 > 215 >

216 > Number of Latents Figure 3: Varying the learning rate jointly with the number of latents. Number of tokens to con-vergence shown above each point. 10 4 10 5 > Number of Latents > 2 × 10 1 > 3 × 10 1 > 4 × 10 1 > 6 × 10 1 > Normalized MSE > 10 6 > 10 5 > 10 4 > 10 3 > 10 2 > 10 1 > 10 0 > Subject model pretraining compute Figure 4: Larger subject models in the GPT-4 family require more latents to get to the same MSE (k = 32).

## 3 Scaling laws 3.1 Number of latents Due to the broad capabilities of frontier models such as GPT-4, we hypothesize that faithfully representing model state will require large numbers of sparse features. We consider two primary approaches to choose autoencoder size and token budget: 1.1 Training to compute-MSE frontier (L(C)) Firstly, following Lindsey et al. [2024], we train autoencoders to the optimal MSE given the available compute, disregarding convergence. This method was introduced for pre-training language models [Kaplan et, 2020, Hoffmann et, 2022]. We find that MSE follows a power law L(C) of compute, though the smallest models are off trend (Figure 1).

However, latents are the important artifact of training (not reconstruction predictions), whereas for language models we typically care only about token predictions. Comparing MSE across different n is thus not a fair comparison — the latents have a looser information bottleneck with larger n, so lower MSE is more easily achieved. Thus, this approach is arguably unprincipled for autoencoder training. 1.2 Training to convergence (L(N)) We also look at training autoencoders to convergence (within some ϵ). This gives a bound on the best possible reconstruction achievable by our training method if we disregard compute efficiency.

In practice, we would ideally train to some intermediate token budget between L(N) and We find that the largest learning rate that converges scales with 1/√n (Figure 3). We also find that the optimal learning rate for L(N) is about four times smaller than the optimal learning rate for We find that the number of tokens to convergence increases as approximately Θ(n0.6) for GPT-2 small and Θ(n0.65) for GPT-4 (Figure 11). This must break at some point – if token budget continues to increase sublinearly, the number of tokens each latent receives gradient signal on would approach zero.

8 1.3 Irreducible loss Scaling laws sometimes include an irreducible loss term e, such that y = αx β + e [Henighan et, 2020]. We find that including an irreducible loss term substantially improves the quality of our fits for both L(C) and L(N). > 8 One slight complication is that in the infinite width limit, TopK autoencoders with our initialization scheme are actually optimal at init using our init scheme (Section 1), so this allows for an exponent very slightly less than 1; however, this happens very slowly with n so is unlikely to be a major factor at realistic scales.

410 1 10 2 > Sparsity (L0) > 10 1 > Delta cross-entropy > better > ReLU > ProLU ReLU > ProLU STE > Gated > TopK (ours) (a) For a fixed number of latents (n = 2 17 =131072), the downstream-loss/sparsity trade-off is better for TopK autoencoders than for other activa-tion functions. 10 16 × 10 2 2 × 10 1 3 × 10 1 > Normalized MSE > 10 1 > Delta cross-entropy > better > ReLU > ProLU ReLU > ProLU STE > Gated > TopK (ours) (b) For a fixed sparsity level (L0 = 128), a given MSE level leads to a lower downstream-loss for TopK autoencoders than for other activation func-tions.

Figure 5: Comparison between TopK and other activation functions on downstream loss. Comparisons done for GPT-2 small, see Figure 13 for GPT-4. It was initially not clear to us that there should be a nonzero irreducible loss. One possibility is that there are other kinds of structures in the activations. In the extreme case, unstructured noise in the activations is substantially harder to model and would have an exponent close to zero (Appendix G). Existence of some unstructured noise would explain a bend in the power law. 1.4 Jointly fitting sparsity (L(N, K)) We find that MSE follows a joint scaling law along the number of latents n and the sparsity level k (Figure 1b).

Because reconstruction becomes trivial as k approaches dmodel, this scaling law only holds for the small k regime. Our joint scaling law fit on GPT-4 autoencoders is: L(n, k) = exp(α + βk log(k) + βn log(n) + γ log(k) log(n)) + exp(ζ + η log(k)) (3) with α = −0.50, βk = 0.26, βn = −0.017, γ = −0.042, ζ = −1.32, and η = −0.085. We can see that γ is negative, which means that the scaling law L(N) gets steeper as k increases. η is negative too, which means that the irreducible loss decreases with k.

3.2 Subject model size Ls(N) Since language models are likely to keep growing in size, we would also like to understand how sparse autoencoders scale as the subject models get larger. We find that if we hold k constant, larger subject models require larger autoencoders to achieve the same MSE, and the exponent is worse (Figure 4). ## 4 Evaluation We demonstrated in Section 3 that our larger autoencoders scale well in terms of MSE and sparsity (see also a comparison of activation functions in Section 5.2). However, the end goal of autoencoders is not to improve the sparsity-reconstruction frontier (which degenerates in the limit 9), but rather to find features useful for applications, such as mechanistic interpretability.

Therefore, we measure autoencoder quality with the following metrics: 1. Downstream loss: How good is the language model loss if the residual stream latent is replaced with the autoencoder reconstruction of that latent? (Section 4.1) 2. Probe loss: Do autoencoders recover features that we believe they might have? (Section 4.2) > 9 Improving the reconstruction-sparsity frontier is not always strictly better. An infinitely wide maximally sparse (k = 1) autoencoder can perfectly reconstruct by assigning latents densely in Rd, while being completely structure-less and uninteresting. 5(a) Probe loss (b) Logit diff sparsity Figure 6: The probe loss and logit diff metrics as a function of number of total latents n and active latents k, for GPT-2 small autoencoders.

More total latents (higher n) generally improves all metrics (yellow = better). Both metrics are worse at L 0 = 512, a regime in which solutions are dense (see Section 5). 3. Explainability: Are there simple explanations that are both necessary and sufficient for the activation of the autoencoder latent? (Section 4.3) 4. Ablation sparsity: Does ablating individual latents have a sparse effect on downstream logits? (Section 4.5) These metrics provide evidence that autoencoders generally get better when the number of total latents increases. The impact of the number of active latents L 0 is more complicated.

Increasing L 0 makes explanations based on token patterns worse, but makes probe loss and ablation sparsity better. All of these trends also break when L 0 gets close to dmodel, a regime in which latents also become quite dense (see Section 5 for detailed discussion). 4.1 Downstream loss An autoencoder with non-zero reconstruction error may not succeed at modeling the features most relevant for behavior [Braun et, 2024]. To measure whether we model features relevant to language modeling, we follow prior work [Bills et, 2023, Cunningham et, 2023, Bricken et, 2023, Braun et

, 2024] and consider downstream Kullback-Leibler (KL) divergence and cross-entropy loss. 10 In both cases, we test an autoencoder by replacing the residual stream by the reconstructed value during the forward pass, and seeing how it affects downstream predictions. We find that k-sparse autoencoders improve more on downstream loss than on MSE over prior methods (Figure 5a). We also find that MSE has a clean power law relationship with both KL divergence, and difference of cross entropy loss (Figure 5b), when keeping sparsity L 0 fixed and only varying autoencoder size. Note that while this trend is clean for our trained autoencoders, we can observe instances where it breaks such as when modulating k at test time (see Section 5.3).

One additional issue is that raw loss numbers alone are difficult to interpret—we would like to know how good it is in an absolute sense. Prior work [Bricken et, 2023, Rajamanoharan et, 2024] use the loss of ablating activations to zero as a baseline and report the fraction of loss recovered from that baseline. However, because ablating the residual stream to zero causes very high downstream loss, this means that even very poorly explaining the behavior can result in high scores. 11 Instead, we believe a more natural metric is to consider the relative amount of pretraining compute needed to train a language model of comparable downstream loss.

For example, when our 16 million latent autoencoder is substituted into GPT-4, we get a language modeling loss corresponding to 10% of the pretraining compute of GPT-4. > 10 Because a perfect reconstruction would lead to a non-zero cross-entropy loss, we actually consider the difference to the perfect-autoencoder cross-entropy (“delta cross-entropy”). > 11 For completeness, the zero-ablation fidelity metric of our 16M autoencoder is 98.2%. 64.2 Recovering known features with 1d probes If we expect that a specific feature g sentiment, language identification) should be discovered by a high quality autoencoder, then one metric of autoencoder quality is to check whether these features are present.

Based on this intuition, we curated a set of 61 binary classification datasets (details in Table 1). For each task, we train a 1d logistic probe on each latent using the Newton-Raphson method to predict the task, and record the best cross entropy loss (across latents). 12 That is: min > i,w,b E [y log σ (wz i + b) + (1 − y) log (1 − σ (wz i + b))] (4) where zi is the ith pre-activation autoencoder latent, and y is a binary label. Results on GPT-2 small are shown in Figure 6a.

We find that probe score increases and then decreases as k increases. We find that TopK generally achieves better probe scores than ReLU (Figure 23), and both are substantially better than when using directly residual stream channels. See Figure 32 for results on several GPT-4 autoencoders: we observe that this metric improves throughout training, despite there being no supervised training signal; and we find that it beats a baseline using channels of the residual stream. See Figure 33 for scores broken down by component. This metric has the advantage that it is computationally cheap. However, it also has a major limitation, which is that it leans on strong assumptions about what kinds of features are natural.

4.3 Finding simple explanations for features Anecdotally, our autoencoders find many features that have quickly recognizable patterns that suggest explanations when viewing random activations (Section 1). However, this can create an “illusion” of interpretability [Bolukbasi et, 2021], where explanations are overly broad, and thus have good recall but poor precision. For example, Bills et al. [2023] propose an automated interpretability score which disproportionately depends on recall. They find a feature activating at the end of the phrase “don’t stop” or “can’t stop”, but an explanation activating on all instances of “stop” achieves a high interpretability score.

As we scale autoencoders and the features get sparser and more specific, this kind of failure becomes more severe. Unfortunately, precision is extremely expensive to evaluate when the simulations are using GPT-4 as in Bills et al. [2023]. As an initial exploration, we focus on an improved version of Neuron to Graph (N2G) [Foote et, 2023], a substantially less expressive but much cheaper method that outputs explanations in the form of collections of n-grams with wildcards. In the future, we would like to explore ways to make it more tractable to approximate precision for arbitrary English explanations.

To construct a N2G explanation, we start with some sequences that activate the latent. For each one, we find the shortest suffix that still activates the latent. 13 We then check whether any position in the n-gram can be replaced by a padding token, to insert wildcard tokens. We also check whether the explanation should be dependent on absolute position by checking whether inserting a padding token at the beginning matters. We use a random sample of up to 16 nonzero activations to build the graph, and another 16 as true positives for computing recall.

Results for GPT-2 small are found in Figure 25a and 25b. Note that dense token patterns are trivial to explain, thus n = 2048, k = 512 latents are easy to explain on average since many latents activate extremely densely (see Section 5) 14. In general, autoencoders with more total latents and fewer active latents are easiest to model with N2G. We also obtain evidence that TopK models have fewer spurious positive activations than their ReLU counterparts. N2G explanations have significantly better recall (>1.5x) and only slightly worse precision (>0.9x) for TopK models with the same n (resulting in better F1 scores) and similar L 0 (Figure 24).

> 12 This is similar to the approach of Gurnee et al. [2023], but always using k= 1, and regressing on autoencoder latents instead of neurons. > 13 with at least half the original activation strength > 14 This highlights an issue with our precision/recall metrics, which care only about binarized values. We also propose a more expensive metric which uses simulated values Section 4.4 and addresses this issue. 7(a) A feature with precision = 0.97, recall = 0.56 (b) A feature with precision = 0.06, recall = 0.05 > Recall Eval Sequences: \n\nTherefore she’s not going to play pranks but to give a sweet time-a rendezvous choice.

[Warning: Ripe with Wine Puns]\n\nVital Lacerda (September. [Warning: Developed with Evolutionary Puns]\n\nMailbags\n\nTrial & al6606 M mepwn12 Thorny [B]Available voice actors, his picture of President \n\nPresident Prankster\n\nDmitry Rogozin, Precision Eval Sequences: \n\nThese are not all of the % 70% 99%\n\n71 NoPwn1ntended i7-950 3.4 ottonee Whimsicott M 0 Prankster / Infiltrator Prankster / Inf of Chimera players using one-dimensional teams: Prank of Chimera players using one-dimensional teams: Prankster SporeSeed, Shell Smash, Poison Recall Eval Sequences: \nMeal 1 06h30 Oats, double serving of Whey, 6 \n\n04Jul65 to 01Aug65\n\n1/3 un at\n\n2013-11-16 09:47:18 (25) chat: <DiamondCard 0000 - 00260076] 119 pages 5: [00260, calm it2013-11-11 20:00:43 (25) chat: It’s vanity Precision Eval Sequences: 1 Tue Mar 13 22:37:59 EST 2001 i686 Locale: LANG=C Kiwi76 added 01:08 - Apr 28\n\nCheers Clive and hopefully you can

key=248858886|000101&CMC=&PN=&Is key=248858886|000102&CMC=&PN=&Is Figure 7: Qualitative examples of latents with high and low precision/recall N2G explanations. Key: Green = Ground truth feature activation, Underline = N2G predicted feature activation 4.4 Explanation reconstruction When our goal is for a model’s activations to be interpretable, one question we can ask is: how much performance do we sacrifice if we use only the parts of the model that we can interpret? Our downstream loss metric measures how much of the performance we’re capturing (but our features could be uninterpretable), and our explanation based metric measures how monosemantic our features are (but they might not explain most of the model).

This suggests combining our downstream loss and explanation metrics, by using our explanations to simulate autoencoder latents, and then checking downstream loss after decoding. This metric also has the advantage that it values both recall and precision in a way that is principled, and also values recall more for latents that activate more densely. We tried this with N2G explanations. N2G produces a simulated value based on the node in the trie, but we scale this value to minimize variance explained. Specifically, we compute E[sa]/E [s2], where s is the simulated value and a is the true value, and we estimate this quantity over a training set of tokens.

Results for GPT-2 are shown in Figure 8. We find that we can explain more of GPT-2 small than just explaining bigrams, and that larger and sparser autoencoders result in better downstream loss. 4.5 Sparsity of ablation effects If the underlying computations learned by a language model are sparse, one hypothesis is that natural features are not only sparse in terms of activations, but also in terms of downstream effects [Olah et, 2024]. Anecdotally, we observed that ablation effects often are interpretable (see our visualizer). Therefore, we developed a metric to measure the sparsity of downstream effects on the output logits.

At a particular token index, we obtain the latents at the residual stream, and proceed to ablate each autoencoder latent one by one, and compare the resulting logits before and after ablation. This process leads to V logit differences per ablation and affected token, where V is the size of the token vocabulary. Because a constant difference at every logit does not affect the post-softmax probabilities, we subtract at each token the median logit difference value. Finally, we concatenate these vectors together across some set of T future tokens (at the ablated index or later) to obtain a vector of

8Figure 8: Downstream loss on GPT-2 with various residual stream ablations at layer 8. N2G explanations of autoencoder latents improves downstream loss with larger n and smaller k. V · T total numbers. We then measure the sparsity of this vector via (L1 > L2)2, which corresponds to an “effective number of vocab tokens affected”. We normalize by V · T to have a fraction between 0 and 1, with smaller values corresponding to sparser effects. We perform this for various autoencoders trained on the post-MLP residual stream at layer 8 in GPT-2 small, with T = 16.

Results are shown in Figure 6b. Promisingly, models trained with larger k have latents with sparser effects. However, the trend reverses at k = 512, indicating that as k approaches dmodel = 768, the autoencoder learns latents with less interpretable effects. Note that latents are sparse in an absolute sense, having a (L1 > L2)2 of 10-14%, whereas ablating residual stream channels gives 60% (slightly better than the theoretical value of ∼ 2 > π for random vectors). ## 5 Understanding the TopK activation function 5.1 TopK prevents activation shrinkage A major drawback of the L 1 penalty is that it tends to shrink all activations toward zero [Tibshirani, 1996].

Our proposed TopK activation function prevents activation shrinkage, as it entirely removes the need for an L 1 penalty. To empirically measure the magnitude of activation shrinkage, we consider whether different (and potentially larger) activations would result in better reconstruction given a fixed decoder. We first run the encoder to obtain a set of activated latents, save the sparsity mask, and then optimize only the nonzero values to minimize MSE. 15 This refinement method has been proposed multiple times such as in k-SVD [Aharon et, 2006], the relaxed Lasso [Meinshausen, 2007], or ITI [Maleki, 2009].

We solve for the optimal activations with a positivity constraint using projected gradient descent. This refinement procedure tends to increase activations in ReLU models on average, but not in TopK models (Figure 9a), which indicates that TopK is not impacted by activation shrinkage. The magnitude of the refinement is also smaller for TopK models than for ReLU models. In both ReLU and TopK models, the refinement procedure noticeably improves the reconstruction MSE (Figure 9b), and the downstream next-token-prediction cross-entropy (Figure 9c). However, this refinement only closes part of the gap between ReLU and TopK models.

> 15 unlike the “inference-time optimization” procedure in Nanda et al. [2024] 93 2 1 0 1 2 3 > Difference after refinement > 0.0 > 0.2 > 0.4 > 0.6 > 0.8 > 1.0 > 1.2 > 1.4 > Density > ReLU > TopK 10 110 2 > L0 > 10 1 > Normalized MSE > better > ReLU > TopK > (raw) > (refined) 10 110 2 > L0 > 10 2 > 10 1 > Delta cross-entropy > better > ReLU > TopK > (raw) > (refined) Figure 9: Latent activations can be refined to improve reconstruction from a frozen set of latents.

For ReLU autoencoders, the refinement is biased toward positive values, consistent with compensating for the shrinkage caused by the L 1 penalty. For TopK autoencoders, the refinement is not biased, and also smaller in magnitude. The refinement only closes part of the gap between ReLU and TopK. 5.2 Comparison with other activation functions Other recent works on sparse autoencoders have proposed different ways to address the L 1 activation shrinkage, and Pareto improve the L 0-MSE frontier [Wright and Sharkey, 2024, Taggart, 2024, Rajamanoharan et, 2024]. Wright and Sharkey [2024] propose to fine-tune a scaling parameter per latent, to correct for the L 1 activation shrinkage.

In Gated sparse autoencoders [Rajamanoharan et, 2024], the selection of which latents are active is separate from the estimation of the activation magnitudes. This separation allows autoencoders to better estimate the activation magnitude, and avoid the L 1 activation shrinkage. Another approach is to replace the ReLU activation function with a ProLU [Taggart, 2024] (also known as TRec [Konda et, 2014], or JumpReLU [Erichson et, 2019]), which sets all values below a positive threshold to zero Jθ (x) = x · 1(x>θ). Because the parameter θ is non-differentiable, it requires a approximate gradient such as a ReLU equivalent (ProLU-ReLU) or a straight-through estimator (ProLU-STE) [Taggart, 2024].

We compared these different approaches in terms of reconstruction MSE, number of active latents L0, and downstream cross-entropy loss (Figure 2 and 5). We find that they significantly improve the reconstruction-sparsity Pareto frontier, with TopK having the best performance overall. 5.3 Progressive recovery In a progressive code, a partial transmission still allows reconstructing the signal with reasonable fidelity [Skodras et, 2001]. For autoencoders, learning a progressive code means that ordering latents by activation magnitude gives a way to progressively recover the original vector. To study this property, we replace the autoencoder activation function (after training) by a TopK(k′) activation function where k′ is different than during training.

We then evaluate each value of k′ by placing it in the L 0-MSE plane (Figure 10). We find that training with TopK only gives a progressive code up to the value of k used during training. MSE keeps improving for values slightly over k (a result also described in [Makhzani and Frey, 2013]), then gets substantially worse as k′ increases (note that the effect on downstream loss is more muted). This can be interpreted as some sort of overfitting to the value k. 3.1 Multi-TopK To mitigate this issue, we sum multiple TopK losses with different values of k (Multi-TopK).

For example, using L(k) + L(4 k)/8 is enough to obtain a progressive code over all k′ (note however that training with Multi-TopK does slightly worse than TopK at k). Training with the baseline ReLU only gives a progressive code up to a value that corresponds to using all positive latents. 3.2 Fixed sparsity versus fixed threshold At test time, the activation function can also be replaced by a JumpReLU activation, which activates above a fixed threshold θ, Jθ (x) = x · 1(x>θ). In contrast to TopK, JumpReLU leads to a selection of active latents where the number of active latents can vary across tokens.

Results for replacing the activation function at test-time with a JumpReLU are shown in dashed lines in Figure 10. 10 10 0 10 1 10 2 10 3 > Sparsity (L0) > 10 1 > 10 0 > Normalized MSE > better > TopK > Act fn (train) > TopK (test) > JumpReLU (test) > Act fn (train) > TopK (test) > JumpReLU (test) 10 010 110 210 3 > Sparsity (L0) > 10 1 > 10 0 > Normalized MSE > better > Multi-TopK > Act fn (train) > TopK (test) > JumpReLU (test) > Act fn (train) > TopK (test) > JumpReLU (test) 10 010 110 210 3 > Sparsity (L0) > 10 1 > 10 0 > Normalized MSE > better > ReLU > Act fn (train) > TopK (test) > JumpReLU (test) > Act fn (train) > TopK (test) > JumpReLU (test) 10 010 110 210 3 > Sparsity (L0) > 10 1 > 10 0 > Normalized MSE > better > Gated > Act fn (train) > TopK (test) > JumpReLU (test) > Act fn (train) > TopK (test) > JumpReLU (test) Figure 10: Sparsity levels can be changed at test time by replacing the activation function with either TopK(k) or JumpReLU(θ), for a given value k or θ.

TopK tends to overfit to the value of k used during training, but using Multi-TopK improves generalization to larger For autoencoders trained with TopK, the test-time TopK and JumpReLU curves are superimposed only for values corresponding to an L 0 below the training L 0, otherwise the JumpReLU activation is worse than the TopK activation. This discrepancy disappears with Multi-TopK, where both curves are nearly superimposed, which means that the model can be used with either a fixed or a dynamic number of latents per token without loss in reconstruction. The two curves are also superimposed for autoencoders trained with ReLU.

Interestingly, it is sometimes more efficient to train a ReLU model with a low L 1 penalty and to use a TopK or JumpReLU at test time, than to use a higher L 1 penalty that would give a similar sparsity level (a result independently described in Nanda et al. [2024]). ## 6 Limitations and Future Directions We believe many improvements can be made to our autoencoders. • TopK forces every token to use exactly k latents, which is likely suboptimal. Ideally we would constrain E[L0] rather than • The optimization can likely be greatly improved, for example with learning rate scheduling, 16 better optimizers, and better aux losses for preventing dead latents.

• Much more could be done to understand what metrics best track relevance to downstream applications, and to study those applications themselves. Applications include: finding vectors for steering behavior, doing anomaly detection, identifying circuits, and more. • We’re excited about work in the direction of combining MoE [Shazeer et, 2017] and au-toencoders, which would substantially improve the asymptotic cost of autoencoder training, and enable much larger autoencoders. • A large fraction of the random activations of features we find, especially in GPT-4, are not yet adequately monosemantic. We believe that with improved techniques and greater scale 17 this is potentially surmountable.

> 16 Anecdotally, we also found that lowering learning rates helped with decreasing dead latents. > 17 both in number of latents and in training tokens 11 • Our probe based metric is quite noisy, which could be improved by having a greater breadth of tasks and higher quality tasks. • While we use n2g for its computational efficiency, it is only able to capture very simple patterns. We believe there is a lot of room for improvement in terms of more expressive explanation methods that are also cheap enough to simulate to estimate explanation precision.

• A context length of 64 tokens is potentially too few tokens to exhibit the most interesting behaviors of GPT-4. ## 7 Related work Sparse coding on an over-complete dictionary was introduced by Mallat and Zhang [1993]. Ol-shausen and Field [1996] refined the idea by proposing to learn the dictionary from the data, without supervision. This approach has been particularly influential in image processing, as seen for example in [Mairal et, 2014]. Later, Hinton and Salakhutdinov [2006] proposed the autoencoder archi-tecture to perform dimensionality reduction. Combining these concepts, sparse autoencoders were developed [Lee et

, 2007, Le et, 2013, Konda et, 2014] to train autoencoders with sparsity priors, such as the L 1 penalty, to extract sparse features. Makhzani and Frey [2013] refined this con-cept by introducing k-sparse autoencoders, which use a TopK activation function instead of the L 1 penalty. Makelov et al. [2024] evaluates autoencoders using a metric that measures recovery of fea-tures from previously discovered circuits. More recently, sparse autoencoders were applied to language models [Yun et, 2021, Lee Sharkey, 2022, Bricken et, 2023, Cunningham et, 2023], and multiple sparse autoencoders were trained on small open-source language models [Marks, 2023,

Bloom, 2024, Mossing et, 2024]. Marks et al. [2024] showed that the resulting features from sparse autoencoders can find sparse circuits in language models. Wright and Sharkey [2024] pointed out that sparse autoencoders are subject to activation shrinking from L 1 penalties, a property of L 1 penalties first described in Tibshirani [1996]. Taggart [2024] and Rajamanoharan et al. [2024] proposed to use different activation functions to address activation shrinkage in sparse autoencoders. Braun et al. [2024] proposed to train sparse autoencoders on downstream KL instead of reconstruction MSE. Kaplan et al. [2020] studied scaling laws for language models which examine how loss varies with various hyperparameters.

Clark et al. [2022] explore scaling laws related to sparsity using a bilinear fit. Lindsey et al. [2024] studied scaling laws specifically for autoencoders, defining the loss as a specific balance of reconstruction and sparsity (rather than simply reconstruction, while holding sparsity fixed). ## Acknowledgments We are deeply grateful to Jan Leike and Ilya Sutskever for leading the Superalignment team and creating the research environment which made this work possible. We thank David Farhi for supporting our work after their departure. We thank Cathy Yeh for explorations in finding features. We thank Carroll Wainwright for initial insights into clustering latents.

We thank Steven Bills, Dan Mossing, Cathy Yeh, William Saunders, Boaz Barak, Jakub Pachocki, and Jan Leike for many discussions about autoencoders and their applications. We thank Manas Joglekar for some improvements on an earlier version of our activation loading. We thank William Saunders for an early autoencoder visualization implementation. We thank Trenton Bricken, Dan Mossing, and Neil Chowdhury for feedback on an earlier version of this manuscript. We thank Trenton Bricken, Lawrence Chan, Hoagy Cunningham, Adam Goucher, Ryan Greenblatt, Tristan Hume, Jan Hendrik Kirchner, Jake Mendel, Neel Nanda, Noa Nabeshima, Chris Olah, Logan Riggs, Fabien Roger, Buck Shlegeris, John Schulman, Lee Sharkey, Glen Taggart, Adly Templeton, Vikrant Varma for valuable discussions.

12 References Michal Aharon, Michael Elad, and Alfred Bruckstein. K-SVD: An algorithm for designing over-complete dictionaries for sparse representation. IEEE Transactions on signal processing, 54(11): 4311–4322, 2006. David Baehrens, Timon Schroeter, Stefan Harmeling, Motoaki Kawanabe, Katja Hansen, and Klaus-Robert Müller. How to explain individual classification decisions. The Journal of Machine Learning Research, 11:1803–1831, 2010. Steven Bills, Nick Cammarata, Dan Mossing, Henk Tillman, Leo Gao, Gabriel Goh, Ilya Sutskever, Jan Leike, Jeff Wu, and William Saunders. Language models can explain neurons in lan-guage models. OpenAI Blog, 2023. URL net/ Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi.

Piqa: Reasoning about physical commonsense in natural language. In Thirty-Fourth AAAI Conference on Artificial Intelligence, 2020. Joseph Bloom. Open source sparse autoencoders for all residual stream layers of gpt2-small. AI Alignment Forum, 2024. URL com/posts/f9EgfLSurAiqRJySD/ Tolga Bolukbasi, Adam Pearce, Ann Yuan, Andy Coenen, Emily Reif, Fernanda Viégas, and Martin Wattenberg. An interpretability illusion for BERT. arXiv preprint arXiv:2104.07143, 2021. Dan Braun, Jordan Taylor, Nicholas Goldowsky-Dill, and Lee Sharkey. Identifying functionally important features with end-to-end sparse dictionary learning. 2024. Trenton Bricken, Adly Templeton, Joshua Batson, Brian Chen, Adam Jermyn, Tom Conerly, Nick Turner, Cem Anil, Carson Denison, Amanda Askell, Robert Lasenby, Yifan Wu, Shauna Kravec, Nicholas Schiefer, Tim Maxwell, Nicholas Joseph, Zac Hatfield-Dodds, Alex Tamkin, Karina Nguyen, Brayden McLean, Josiah E Burke, Tristan Hume, Shan Carter, Tom Henighan, and Christopher Olah.

Towards monosemanticity: Decomposing language models with dictionary learning. Transformer Circuits Thread, 2023. html. Cjadams, Jeffrey Sorensen, Julia Elliott, Lucas Dixon, Mark McDonald, Nithum, and Will Cukier-ski. Toxic comment classification challenge. Kaggle, 2017. URL com/ Aidan Clark, Diego de Las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann, Bogdan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, et al. Unified scaling laws for routed language models. In International conference on machine learning, pages 4057–4086. PMLR, 2022. Tom Conerly, Adly Templeton, Trenton Bricken, Jonathan Marcus, and Tom Henighan. Update on how we train saes. Transformer Circuits Thread, 2024.

html#training-saes. Hoagy Cunningham, Aidan Ewart, Logan Riggs, Robert Huben, and Lee Sharkey. Sparse autoen-coders find highly interpretable features in language models. arXiv preprint arXiv:2309.08600,2023. Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac Hatfield-Dodds, Robert Lasenby, Dawn Drain, Carol Chen, et al. Toy models of superposition. arXiv preprint arXiv:2209.10652, 2022. N Benjamin Erichson, Zhewei Yao, and Michael W Mahoney. JumpReLU: A retrofit defense strategy for adversarial attacks. arXiv preprint arXiv:1904.03750, 2019. Alex Foote, Neel Nanda, Esben Kran, Ioannis Konstas, Shay Cohen, and Fazl Barez. Neuron to graph: Interpreting language model neurons at scale.

arXiv preprint arXiv:2305.19911, 2023. 13 Gabriel Goh. Decoding the thought vector, 2016. URL io/ ThoughtVectors/. Accessed: 2024-05-24. Antonio Gulli. Ag’s corpus of news articles. it/~gulli/AG_corpus_ html. Accessed: 2024-05-21. Wes Gurnee, Neel Nanda, Matthew Pauly, Katherine Harvey, Dmitrii Troitskii, and Dimitris Bertsimas. Finding neurons in a haystack: Case studies with sparse probing. arXiv preprint arXiv:2305.01610,2023. Dan Hendrycks, Collin Burns, Steven Basart, Andrew Critch, Jerry Li, Dawn Song, and Jacob Steinhardt. Aligning ai with shared human values. arXiv preprint arXiv:2008.02275, 2020. Tom Henighan, Jared Kaplan, Mor Katz, Mark Chen, Christopher Hesse, Jacob Jackson, Heewoo Jun, Tom B Brown, Prafulla Dhariwal, Scott Gray, et al.

Scaling laws for autoregressive generative modeling. arXiv preprint arXiv:2010.14701, 2020. Geoffrey E Hinton and Ruslan R Salakhutdinov. Reducing the dimensionality of data with neural networks. science, 313(5786):504–507, 2006. Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556, 2022. Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism.

Advances in neural information processing systems, 32, 2019. Adam Jermyn and Adly Templeton. Ghost grads: An improvement on resampling. Transformer Circuits Thread, 2024. html#dict-learning-resampling. Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020. Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014. Philipp Koehn. Europarl: A parallel corpus for statistical machine translation. In Proceedings of Machine Translation Summit X: Papers, pages 79–86, Phuket, Thailand, September 13-15 2005.

URL Kishore Konda, Roland Memisevic, and David Krueger. Zero-bias autoencoders and the benefits of co-adapting features. arXiv preprint arXiv:1402.3337, 2014. Quoc V Le, Marc’Aurelio Ranzato, Rajat Monga, Matthieu Devin, Kai Chen, Greg S Corrado, Jeff Dean, and Andrew Y Ng. Building high-level features using large scale unsupervised learning. In 2013 IEEE international conference on acoustics, speech and signal processing, pages 8595–8598. IEEE, 2013. Honglak Lee, Chaitanya Ekanadham, and Andrew Ng. Sparse deep belief net model for visual area v2. Advances in neural information processing systems, 20, 2007. Beren Millidge Lee Sharkey, Dan Braun. Taking features out of su-perposition with sparse autoencoders.

AI Alignment Forum, 2022. URL org/posts/z6QQJbtpkEAX3Aojj/ Stephanie Lin, Jacob Hilton, and Owain Evans. Truthfulqa: Measuring how models mimic human falsehoods. arXiv preprint arXiv:2109.07958, 2021. Jack Lindsey, Tom Conerly, Adly Templeton, Jonathan Marcus, and Tom Henighan. Scaling laws for dictionary learning. Transformer Circuits Thread, 2024. html#scaling-laws. 14 Julien Mairal, Francis Bach, Jean Ponce, et al. Sparse modeling for image and vision processing. Foundations and Trends® in Computer Graphics and Vision, 8(2-3):85–283, 2014. Aleksandar Makelov, George Lange, and Neel Nanda. Towards principled evaluations of sparse autoencoders for interpretability and control, 2024. Alireza Makhzani and Brendan Frey. K-sparse autoencoders.

arXiv preprint arXiv:1312.5663, 2013. Arian Maleki. Coherence analysis of iterative thresholding algorithms. In 2009 47th Annual Allerton Conference on Communication, Control, and Computing (Allerton), pages 236–243. IEEE, 2009. Stéphane G Mallat and Zhifeng Zhang. Matching pursuits with time-frequency dictionaries. IEEE Transactions on signal processing, 41(12):3397–3415, 1993. Sam Marks. Some open-source dictionaries and dictionary learning infrastructure. AI Align-ment Forum, 2023. URL org/posts/AaoWLcmpY3LKvtdyq/ Samuel Marks, Can Rager, Eric J Michaud, Yonatan Belinkov, David Bau, and Aaron Mueller. Sparse feature circuits: Discovering and editing interpretable causal graphs in language models. arXiv preprint arXiv:2403.19647, 2024. Julian McAuley and Jure Leskovec.

Hidden factors and hidden topics: understanding rating dimen-sions with review text. In Proceedings of the 7th ACM conference on Recommender systems, pages 165–172, 2013. Leland McInnes, John Healy, and James Melville. Umap: Uniform manifold approximation and projection for dimension reduction. arXiv preprint arXiv:1802.03426, 2018. Nicolai Meinshausen. Relaxed lasso. Computational Statistics & Data Analysis, 52(1):374–393, 2007. Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. In EMNLP, 2018. Dan Mossing, Steven Bills, Henk Tillman, Tom Dupré la Tour, Nick Cammarata, Leo Gao, Joshua Achiam, Catherine Yeh, Jan Leike, Jeff Wu, and William Saunders.

Transformer debugger. com/openai/transformer-debugger, 2024. Neel Nanda, Arthur Conmy, Lewis Smith, Senthooran Rajamanoharan, Tom Lieberum, János Kramár, and Vikrant Varma. Progress update #1 from the gdm mech interp team: Full update. AI Align-ment Forum, 2024. URL org/posts/C5KAZQib3bzzpeyrg/ Chris Olah, Adly Templeton, Trenton Bricken, and Adam Jermyn. Open problem: Attribution dictionary learning. Transformer Circuits Thread, 2024. html#attr-dl. Bruno A Olshausen and David J Field. Emergence of simple-cell receptive field properties by learning a sparse code for natural images. Nature, 381(6583):607–609, 1996. OpenAI. Gpt-2 output dataset. com/openai/gpt-2-output-dataset/tree/ master, 2019. Accessed: 2024-05-21. OpenAI. GPT-4 technical report. arXiv prepreint arXiv:2303.08774, 2023.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019. Senthooran Rajamanoharan, Arthur Conmy, Lewis Smith, Tom Lieberum, Vikrant Varma, János Kramár, Rohin Shah, and Neel Nanda. Improving dictionary learning with gated sparse autoen-coders. arXiv preprint arXiv:2404.16014, 2024. Anna Rogers, Olga Kovaleva, Matthew