Scaling and evaluating sparse autoencoders

Introduction

Sparse autoencoders (SAEs) have shown great promise for finding features [12, 8, 64, 16] and circuits [40] 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 [50] 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 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.

Methods

Setup

Inputs: We train autoencoders on the residual streams of both GPT-2 small [51] and models from a series of increasing sized models sharing GPT-4 architecture and training setup, including GPT-4 itself [50].⁴ We choose a layer near the end of the network, which should contain many features without being specialized for next-token predictions (see Section F.1 for more discussion). Specifically, we use a layer $\frac{5}{6}$ of the way into the network for GPT-4 series models, and we use layer 8 ($\frac{3}{4}$ of the way) for GPT-2 small. We use a context length of 64 tokens for all experiments. We subtract the mean over the $d_{\mathrm{model}}$ 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.

Baseline: ReLU autoencoders

For an input vector $x \in\mathbb{R}^{d}$ from the residual stream, and $n$ latent dimensions, we use baseline ReLU autoencoders from [8]. The encoder and decoder are defined by:

with $W_{\mathrm{enc}} \in\mathbb{R}^{n \times d}$, $b_{\mathrm{enc}} \in\mathbb{R}^{n}$, $W_{\mathrm{dec}} \in\mathbb{R}^{d \times n}$, and $b_{\mathrm{pre}} \in\mathbb{R}^{d}$. The training loss is defined by $\mathcal{L} = \lVert x-\hat{x}\rVert_{2}^{2}+\lambda\lVert z\rVert_{1}$, where $\lVert x-\hat{x}\rVert_{2}^{2}$ is the reconstruction MSE, $\lVert z\rVert_{1}$ is an $L_{1}$ penalty promoting sparsity in latent activations $z$, and $\lambda$ is a hyperparameter that needs to be tuned.

TopK activation function

We use a $k$-sparse autoencoder [36], 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:

and the decoder is unchanged. The training loss is simply $\mathcal{L} = \lVert x-\hat{x}\rVert_{2}^{2}$.

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 $\lambda$, enabling simpler model comparison and rapid iteration. It can also be used in combination with arbitrary activation functions.⁵

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

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, [64] train a 34 million latent autoencoder with only 12 million alive latents, and in our ablations we find up to 90% dead latents⁶ 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-$k_{\mathrm{aux}}$ dead latents (see Section A.2 for more details). Using these techniques, even in our largest (16 million latent) autoencoder only 7% of latents are dead.

Scaling laws

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:

Training to compute-MSE frontier (L(C))

Firstly, following [33], we train autoencoders to the optimal MSE given the available compute, disregarding convergence. This method was introduced for pre-training language models [25, 22]. 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.

Training to convergence (L(N))

We also look at training autoencoders to convergence (within some $\epsilon$). 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 $L(C)$.

We find that the largest learning rate that converges scales with $1/\sqrt{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 $L(C)$.

We find that the number of tokens to convergence increases as approximately $\Theta(n^{0.6})$ for GPT-2-small and $\Theta(n^{0.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.⁸

Irreducible loss

Scaling laws sometimes include an irreducible loss term $e$, such that $y = \alpha x^{\beta} + e$ [20]. We find that including an irreducible loss term substantially improves the quality of our fits for both $L(C)$ and $L(N)$.

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.

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 1). Because reconstruction becomes trivial as $k$ approaches $d_{\mathrm{model}}$, this scaling law only holds for the small $k$ regime. Our joint scaling law fit on GPT-4 autoencoders is:

with $\alpha= -0.50$, $\beta_k = 0.26$, $\beta_n = -0.017$, $\gamma= -0.042$, $\zeta= -1.32$, and $\eta= -0.085$. We can see that $\gamma$ is negative, which means that the scaling law $L(N)$ gets steeper as $k$ increases. $\eta$ is negative too, which means that the irreducible loss decreases with $k$.

Subject model size $L_s(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).

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⁹), 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)

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 $d_{\text{model}}$, a regime in which latents also become quite dense (see Section E.5 for detailed discussion).

Downstream loss

An autoencoder with non-zero reconstruction error may not succeed at modeling the features most relevant for behavior [7]. To measure whether we model features relevant to language modeling, we follow prior work [3, 12, 8, 7] and consider downstream Kullback-Leibler (KL) divergence and cross-entropy loss.¹⁰ 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 5).

We also find that MSE has a clean power law relationship with both KL divergence, and difference of cross entropy loss (Figure 5), 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 [8, 52] 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.¹¹

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.

Recovering known features with 1d probes

If we expect that a specific feature (e.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).¹² That is:

Task Name	Details
amazon	McAuley and Leskovec [41]
sciq	Welbl et al. [67]
truthfulqa	Lin et al. [32]
mc_taco	Zhou et al. [72]
piqa	Bisk et al. [4]
quail	Rogers et al. [53]
quartz	Tafjord et al. [61]
justice	Hendrycks et al. [19]
virtue	Hendrycks et al. [19]
utilitarianism	Hendrycks et al. [19]
deontology	Hendrycks et al. [19]
commonsense_qa	Talmor et al. [63]
openbookqa	Mihaylov et al. [44]
base64	discrimination of base64 vs pretraining data
wikidata_isalive	Gurnee et al. [18]
wikidata_sex_or_gender	Gurnee et al. [18]
wikidata_occupation_isjournalist	Gurnee et al. [18]
wikidata_occupation_isathlete	Gurnee et al. [18]
wikidata_occupation_isactor	Gurnee et al. [18]
wikidata_occupation_ispolitician	Gurnee et al. [18]
wikidata_occupation_issinger	Gurnee et al. [18]
wikidata_occupation_isresearcher	Gurnee et al. [18]
phrase_high-school	Gurnee et al. [18]
phrase_living-room	Gurnee et al. [18]
phrase_social-security	Gurnee et al. [18]
phrase_credit-card	Gurnee et al. [18]
phrase_blood-pressure	Gurnee et al. [18]
phrase_prime-factors	Gurnee et al. [18]
phrase_social-media	Gurnee et al. [18]
phrase_gene-expression	Gurnee et al. [18]
phrase_control-group	Gurnee et al. [18]
phrase_magnetic-field	Gurnee et al. [18]
phrase_cell-lines	Gurnee et al. [18]
phrase_trial-court	Gurnee et al. [18]
phrase_second-derivative	Gurnee et al. [18]
phrase_north-america	Gurnee et al. [18]
phrase_human-rights	Gurnee et al. [18]
phrase_side-effects	Gurnee et al. [18]
phrase_public-health	Gurnee et al. [18]
phrase_federal-government	Gurnee et al. [18]
phrase_third-party	Gurnee et al. [18]
phrase_clinical-trials	Gurnee et al. [18]
phrase_mental-health	Gurnee et al. [18]
social_iqa	Sap et al. [55]
wic	Wang et al. [66]
cola	Wang et al. [66]
gpt2disc	OpenAI [49]
ag_news_world	Gulli [17]
ag_news_sports	Gulli [17]
ag_news_business	Gulli [17]
ag_news_scitech	Gulli [17]
europarl_es	Koehn [27]
europarl_en	Koehn [27]
europarl_fr	Koehn [27]
europarl_nl	Koehn [27]
europarl_it	Koehn [27]
europarl_el	Koehn [27]
europarl_de	Koehn [27]
europarl_pt	Koehn [27]
europarl_sv	Koehn [27]
jigsaw	Cjadams et al. [9]

where $z_i$ is the ith pre-activation autoencoder latent, and $y$ is a binary label.

Results on GPT-2 small are shown in fig: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.

Image not extracted

Image not extracted

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.

Finding simple explanations for features

Anecdotally, our autoencoders find many features that have quickly recognizable patterns that suggest explanations when viewing random activations (Section E.1). However, this can create an “illusion” of interpretability [6], where explanations are overly broad, and thus have good recall but poor precision. For example, [3] 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 [3]. As an initial exploration, we focus on an improved version of Neuron to Graph (N2G) [15], 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.¹³ 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 Figure 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 E.5)¹⁴. 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.5\mathrm{x}$) and only slightly worse precision ($>0.9\mathrm{x}$) for TopK models with the same $n$ (resulting in better F1 scores) and similar $L_0$ (Figure 24).

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[s^2]$, 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.

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 [47]. 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 $V \cdot T$ total numbers. We then measure the sparsity of this vector via $\left(\frac{L_1}{L_2}\right)^2$, which corresponds to an “effective number of vocab tokens affected”. We normalize by $V \cdot 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 fig:6b. Promisingly, models trained with larger $k$ have latents with sparser effects. However, the trend reverses at $k = 512$, indicating that as $k$ approaches $d_{\mathrm{model}} = 768$, the autoencoder learns latents with less interpretable effects. Note that latents are sparse in an absolute sense, having a $\left(\frac{L_1}{L_2}\right)^2$ of 10-14%, whereas ablating residual stream channels gives 60% (slightly better than the theoretical value of $\sim\frac{2}{\pi}$ for random vectors).

Understanding the TopK activation function

TopK prevents activation shrinkage

A major drawback of the L$_1$ penalty is that it tends to shrink all activations toward zero [?]. 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.¹⁵ This refinement method has been proposed multiple times such as in $k$-SVD [1], the relaxed Lasso [43], or ITI [37]. 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 (fig: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 (fig:9b), and the downstream next-token-prediction cross-entropy (fig:9c). However, this refinement only closes part of the gap between ReLU and TopK models.

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 [68, 62, 52]. [68] propose to fine-tune a scaling parameter per latent, to correct for the $L_1$ activation shrinkage. In Gated sparse autoencoders [52], 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 [62] (also known as TRec [28], or JumpReLU [14]), which sets all values below a positive threshold to zero $J_\theta(x) = x \cdot1_{(x>\theta)}$. Because the parameter $\theta$ is non-differentiable, it requires a approximate gradient such as a ReLU equivalent (ProLU-ReLU) or a straight-through estimator (ProLU-STE) [62].

We compared these different approaches in terms of reconstruction MSE, number of active latents $L_0$, and downstream cross-entropy loss (Figure 2 and Figure 5). We find that they significantly improve the reconstruction-sparsity Pareto frontier, with TopK having the best performance overall.

Progressive recovery

In a progressive code, a partial transmission still allows reconstructing the signal with reasonable fidelity [59]. 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 [36]), 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$.

Multi-TopK

To mitigate this issue, we sum multiple TopK losses with different values of $k$ (Multi-TopK). For example, using $\mathcal{L}(k) + \mathcal{L}(4k)/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.

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 $\theta$, $J_\theta(x) = x \cdot1_{(x>\theta)}$. 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.

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 [46]).

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 $\mathbb{E}[L_0]$ rather than $L_0$.

• The optimization can likely be greatly improved, for example with learning rate scheduling,¹⁶ 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 [56] and autoencoders, 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¹⁷ this is potentially surmountable.

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

Related work

Sparse coding on an over-complete dictionary was introduced by [38] . [48] 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 [34]. Later, [21] proposed the autoencoder architecture to perform dimensionality reduction. Combining these concepts, sparse autoencoders were developed [30, 29, 28] to train autoencoders with sparsity priors, such as the L$_1$ penalty, to extract sparse features. [36] refined this concept by introducing $k$-sparse autoencoders, which use a TopK activation function instead of the L$_1$ penalty. [35] evaluates autoencoders using a metric that measures recovery of features from previously discovered circuits.

More recently, sparse autoencoders were applied to language models [70, 31, 8, 12], and multiple sparse autoencoders were trained on small open-source language models [39, 5, 45]. [40] showed that the resulting features from sparse autoencoders can find sparse circuits in language models. [68] pointed out that sparse autoencoders are subject to activation shrinking from L$_1$ penalties, a property of L$_1$ penalties first described in [65]. [62] and [52] proposed to use different activation functions to address activation shrinkage in sparse autoencoders. [7] proposed to train sparse autoencoders on downstream KL instead of reconstruction MSE.

[25] studied scaling laws for language models which examine how loss varies with various hyperparameters. [10] explore scaling laws related to sparsity using a bilinear fit. [33] 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 Jog le kar 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.

Optimization

Initialization

We initialize our autoencoders as follows:

• We initialize the bias $b_{pre}$ to be the geometric median of a sample set of data points, following [8].

• We initialize the encoder directions parallel to the respective decoder directions, so that the corresponding latent read/write directions are the same¹⁸ Directions are chosen uniformly randomly.

• We scale decoder latent directions to be unit norm at initialization (and also after each training step), following [8].

• For baseline models we use torch default initialization for encoder magnitudes. For TopK models, we initialized the magnitude of the encoder such that the magnitude of reconstructed

vectors match that of the inputs. However, in our ablations we find this has no effect or a weak negative effect (Figure 16).¹⁹

Auxiliary loss

We define an auxiliary loss (AuxK) similar to “ghost grads” [24] that models the reconstruction error using the top-$k_{\mathrm{aux}}$ dead latents (typically $k_{\mathrm{aux}} = 512$). Latents are flagged as dead during training if they have not activated for some predetermined number of tokens (typically 10 million). Then, given the reconstruction error of the main model $e = x - \hat{x}$, we define the auxiliary loss $L_{\mathrm{aux}} = \lVert e - \hat{e} \rVert_2^2$, where $\hat{e} = W_{\mathrm{dec}}z$ is the reconstruction using the top-$k_{\mathrm{aux}}$ dead latents. The full loss is then defined as $L + \alpha L_{\mathrm{aux}}$, where $\alpha$ is a small coefficient (typically $1/32$). Because the encoder forward pass can be shared (and dominates decoder cost and encoder backwards cost, see Appendix D), adding this auxiliary loss only increases the computational cost by about 10%.

We found that the AuxK loss very occasionally NaNs at large scale, and zero it when it is NaN to prevent training run collapse.

Optimizer

We use the Adam optimizer [26] with $\beta_1 = 0.9$ and $\beta_2 = 0.999$, and a constant learning rate. We tried several learning rate decay schedules but did not find consistent improvements in token budget to convergence. We also did not find major benefits from tuning $\beta_1$ and $\beta_2$.

We project away gradient information parallel to the decoder vectors, to account for interaction between Adam and decoder normalization, as described in [8].

Adam epsilon

By convention, we average the gradient across the batch dimension. As a result, the root mean square (RMS) of the gradient can often be very small, causing Adam to no longer be loss scale invariant. We find that by setting epsilon sufficiently small, these issues are prevented, and that $\varepsilon$ is otherwise not very sensitive and does not result in significant benefit to tune further. We use $\varepsilon= 6.25 \times10^{-10}$ in many experiments in this paper, though we reduced it further for some of the largest runs to be safe.

Gradient clipping

When scaling the GPT-4 autoencoders, we found that gradient clipping was necessary to prevent instability and divergence at higher learning rates. We found that gradient clipping substantially affected $L(C)$ but not $L(N)$. We did not use gradient clipping for the GPT-2 small runs.

Batch size

Larger batch sizes are critical for allowing much greater parallelism. Prior work tends to use batch sizes like 2048 or 4096 tokens [8, 11, 52]. To gain the benefits of parallelism, we use a batch size of 131,072 tokens for most of our experiments.

While batch size affects $L(C)$ substantially, we find that the $L(N)$ loss does not depend strongly on batch size when optimization hyperparameters are set appropriately (Figure 12).

Weight averaging

We find that keeping an exponential moving average (EMA) [54] of the weights slightly reduces sensitivity to learning rate by allowing slightly higher learning rates to be tolerated. Due to its low cost, we use EMA in all experiments. We use an EMA coefficient of 0.999, and did not find a substantial benefit to tuning it.

We use a bias-correction similar to that used in [26]. Despite this, the early steps of EMA are still generally worse than the original model. Thus for the $L(C)$ experiments, we take the min of the EMA model’s and non-averaged model’s validation losses.

Other details

• For the main MSE loss, we compute an MSE normalization constant once at the beginning of training, and do not do any loss normalization per batch.

• For the AuxK MSE loss, we compute the normalization per token, because the scale of the error changes throughout training.

• In theory, the $b_{\mathrm{pre}}$ lr should be scaled linearly with the norm of the data to make the autoencoder completely invariant to input scale. In practice, we find it to tolerate an extremely wide range of values with little impact on quality.

• Anecdotally, we noticed that when decaying the learning rate of an autoencoder previously training at the $L(C)$ loss, the number of dead latents would decrease.

Other training details

Unless otherwise noted, autoencoders were trained on the residual activation directly after the layernorm (with layernorm weights folded into the attention weights), since this corresponds to how residual stream activations are used. This also causes importance of input vectors to be uniform, rather than weighted by norm²⁰.

TopK training details

We select $k_{\mathrm{aux}}$ as a power of two close to $\frac{d_{\mathrm{model}}}{2}$ (e.g. 512 for GPT-2 small). We typically select $\alpha= 1/32$. We find that the training is generally not extremely sensitive to the choice of these hyperparameters.

We find empirically that using AuxK eliminates almost all dead latents by the end of training.

Unfortunately, because of compute constraints, we were unable to train our 16M latent autoencoder to $L(N)$, which made it not possible to include the 16M as part of a consistent $L(N)$ series.

Baseline hyperparameters

Baseline ReLU autoencoders were trained on GPT-2 small, layer 8. We sweep learning rate in [5e-5, 1e-4, 2e-4, 4e-4], L1 coefficient in [1.7e-3, 3.1e-3, 5e-3, 1e-2, 1.7e-2] and train for 8 epochs of 6.4 billion tokens at a batch size of 131072. We try different resampling periods in [12.5k, 25k] steps, and choose to resample 4 times throughout training. We consider a feature dead if it does not activate for 10 million tokens.

For Gated SAE [52], we sweep L$_1$ coefficient in [1e-3, 2.5e-3, 5e-3, 1e-2, 2e-2], learning rate in [2.5e-5, 5e-5, 1e-4], train for 6 epochs of 6.4 billion tokens at a batch size of 131072. We resample 4 times throughout training.

For ProLU autoencoders [62], we sweep L$_1$ coefficient in [5e-4, 1e-3, 2.5e-3, 5e-3, 1e-2, 2e-2], learning rate in [2.5e-5, 5e-5, 1e-4], train for 6 epochs of 6.4 billion tokens at a batch size of 131072. We resample 4 times throughout training. For the ProLU gradient, we try both ProLU-STE and ProLU-ReLU. Note that, consistent with the original work, ProLU-STE autoencoders all have $L_0 < 25$, even for small L$_1$ coefficients.

We used similar settings and sweeps for autoencoders trained on GPT-4. Differences include: replacing the resampling of dead latents with a L$_1$ coefficient warm-up over 5% of training [11]; removing the decoder unit-norm constraint and adding the decoder norm in the L$_1$ penalty [11].

Our baselines generally have few dead latents, similar or less than our TopK models (see Figure 14).

Training ablations

Dead latent prevention

We find that the reduction in dead latents is mostly due to a combination of the AuxK loss and the tied initialization scheme.

Initialization

We find that tied initialization substantially improves MSE, and that our encoder initialization scheme has no effect when tied initialization is being used, and hurts slightly on its own.

$b_{\mathrm{enc}}$

We find that $b_{\mathrm{enc}}$ does not affect the MSE at convergence. With $b_{\mathrm{enc}}$ removed, the autencoder is equivalent to a JumpReLU where the threshold is dynamically chosen per example such that exactly $k$ latents are active. However, convergence is slightly slower without $b_{\mathrm{enc}}$. We believe this may be confounded by encoder learning rate but did not investigate this further.

Decoder normalization

After each step we renormalize columns of the decoder to be unit-norm, following [8]. This normalization (or a modified L1 term, as in [11]) is necessary for L1 autoencoders, because otherwise the L1 loss can be gamed by making the latents arbitrarily small. For TopK autoencoders, the normalization is optional. However, we find that it still improves MSE, so we still use it in all of our experiments.

Systems

Scaling autoencoders to the largest scales in this paper would not be feasible without our systems improvements. Model parallelism is necessary once parameters cannot fit on one GPU. A naive implementation can be an order of magnitude slower than our optimized implementation at the very largest scales.

Parallelism

We use standard data parallel and tensor sharding [57], with an additional allgather for the TopK forward pass to determine which $k$ latents should are in the global top $k$. To minimize the cost of this allgather, we truncate to a capacity factor of 2 per shard—further improvements are possible but would require modifications to NCCL. For the largest (16 million) latent autoencoder, we use 512-way sharding. Large batch sizes (Section A.4) are very important for reducing the parallelization overhead.

The very small number of layers creates a challenge for parallelism - it makes pipeline parallelism [23] and FSDP [71] inapplicable. Additionally, opportunities for communications overlap are limited because of the small number of layers, though we do overlap host to device transfers and encoder data parallel comms for a small improvement.

Kernels

We can take advantage of the extreme sparsity of latents to perform most operations using substantially less compute and memory than naively doing dense matrix multiplication. This was important when scaling to large numbers of latents, both via directly increasing throughput and reducing memory usage.

We use two main kernels:

• DenseSparseMatmul: a multiplication between a dense and sparse matrix

• MatmulAtSparseIndices: a multiplication of two dense matrices evaluated at a set of sparse indices

Then, we have the following optimizations:

1. The decoder forward pass uses DenseSparseMatmul

2. The decoder gradient uses DenseSparseMatmul

3. The latent gradient uses MatmulAtSparseIndices

4. The encoder gradient uses DenseSparseMatmul

5. The pre-bias gradient uses a trick of summing pre-activation gradient across the batch dimension before multiplying with the encoder weights.

Theoretically, this gives a compute efficiency improvement of up to 6x in the limit of sparsity, since the encoder forward pass is the only remaining dense operation. In practice, we indeed find the encoder forward pass is much of the compute, and the pre-activations are much of the memory.

To ensure that reads are coalesced, the decoder weight matrix must also be stored transposed from the typical layout. We also use many other kernels for fusing various operations for reducing memory and memory bandwidth usage.

Qualitative results

Subjective latent quality

Throughout the project, we stumbled upon many subjectively interesting latents. The majority of latents in our GPT-2 small autoencoders seemed interpretable, even on random positive activations. Furthermore, the ablations typically had predictable effects based on the activation conditions. For example, some features that are potentially part of interesting circuits in GPT-2:

• An unexpected token breaking a repetition pattern (A B C D ... A B X!). This upvotes future pattern breaks (A B Y!), but also upvotes continuation of the pattern right after the break token (A B X → D).

• Text within quotes, especially activating when within two nested sets of quotes. Upvotes tokens that close the quotes like " or ’, as well as tokens which close multiple sets of quotes at once, such as "’ and ’".

• Copying/induction of capitalized phrases (A B ... A → B)

In GPT-4, we tended to find more complex features, including ones that activate on the same concept in multiple languages, or on complex technical concepts such as algebraic rings.

You can explore for yourself at our viewer.

Finding features

Typical features are not of particular interest, but having an autoencoder also lets one easily find relevant features. Specifically, one can use gradient-based attribution to quickly compute how relevant latents are to behaviors of interest [2, 58].

Following the methodologies of [64] and [45], we found features by using a hand-written prompt, with a set of "positive" and "negative" token predictions, and back-propagating from logit differences to latent values. We then consider latents sorted by activation times gradient value, and inspect the latents.

We tried this methodology with a $n = 2^{17} = 131072$, $k = 32$ GPT-2 autoencoder and were able to quickly find a number of safety relevant latents, such as ones corresponding to profanity or child sexual content. Clamping these latents appeared to have causal effect on the samples. For example, clamping the profanity latent to negative values results in significantly less profanity (some profanity can still be observed in situations where it copies from the prompt).

Latent activation distributions

We anecdotally found that latent activation distributions often have multiple modes, especially in early layers.

Latent density and importance curves

We find that log of latent density is approximately Gaussian (Figure 19). If we define the importance of a feature to be the expected squared activation (an approximation of its marginal impact on MSE), then log importance looks a bit more like a Laplace distribution. Modal density and modal feature importance both decrease with number of total latents (for a fixed $k$), as expected.

Solutions with dense latents

One option for a sparse autoencoder is to simply learn the $k$ directions that explain the most variance. As $k$ approaches $d_{\mathrm{model}}$, this “principal component”-like solution may become competitive with solutions where each latent is used sparsely. In order to check for such solutions, we can simply measure the average density of the $d_{\mathrm{model}}$ densest latents. Using this metric, we find that for some hyperparameter settings, GPT-2 small autoencoders find solutions with many dense latents (Figure 20), beginning around $k = 256$ but especially for $k = 512$.

This coincides with when the scaling laws from Section 3 begin to break - many of our trends between $k$ and reconstruction loss bend significantly. Also, MSE becomes significantly less sensitive to $n$, at $k = 512$.

Recurring dense features in GPT-2 small

We manually examined the densest latents across various GPT-2 small layer 8 autoencoders, trained in different ways (e.g. differing numbers of total latents).

The two densest latents are always the same feature: the latent simply activates more and more later in the context ($\sim40\%$ active), and one that activates more and more earlier in the context excluding the first position ($\sim35\%$ active). Both features look like they want to activate more than they do, with TopK probably preventing it from activating with lower values.

The third densest latent is always a first-token-position feature ($\sim30\%$ active), which has a modal activation value in a narrow range between 14.6-14.8. Most of its activation values are significantly smaller values, at tokens after the first position; the large value is always at the first token. These smaller values appear uninterpretable; we conjecture these are simply interference with the first position direction. (Sometimes there are two of these latents, the second with smaller activation values.)

Finally, there is a recurring “repetition” feature that is $\sim20\%$ dense. Its top activations are mostly highly repetitive sequences, such as series of dates, chapter indices, numbers, punctuations, repeated exact phrases, or other repetitive things such as Chess PGN notation. However, like the first-token-position latents, random activations of this latent are typically appear unrelated and uninterpretable.

Often in the top ten densest latents, we find opposing latents, which have decoder cosine similarity close to $-1$. In particular, the first-token-position feature and the repetition latent both seems to always have an opposite latent. The less dense of the two opposite latents always seems to appear uninterpretable. We conjecture that these are symptoms of optimization failure - the opposite latents cancel out spurious activations in the denser latent.

Clustering latents

[13] discuss how underlying features may lie in distinct sub-spaces. If such sub-spaces exists, we hypothesize that the set of latent encoding vectors $W \in\mathbb{R}^{n \times d}$ can be written as a block-diagonal matrix $W' = PW R$, where $P \in\mathbb{R}^{n \times n}$ is a permutation matrix, and $R \in\mathbb{R}^{d \times d}$ is orthogonal. We can then use the singular vector decomposition (SVD) to write $W = U\Sigma V^{\mathsf{T}}$ and $W' = U'\Sigma'V'^{\mathsf{T}}$, noting that $U'$ is also block diagonal. Finally, we write $W = P^{\mathsf{T}}U'\Sigma'V'^{\mathsf{T}}R^{\mathsf{T}} = U\Sigma V^{\mathsf{T}}$, and because the SVD is unique up to a column permutation $P'$, we get $U = P^{\mathsf{T}}U'P'$. In other words, if $W$ is block-diagonal in some unknown basis, $U$ is also block diagonal up to a permutation of rows and columns.

To find a good permutation of rows, we sorted the rows of $U$ based on how similarly they project on all elements of the singular vector basis. Specifically, we normalized each row to unit norm $\widetilde{U}_i = U_i/\lVert U_i\rVert$ and considered the pairwise euclidean distances $d_{i,j} = \lVert\widetilde{U}_i^{2} - \widetilde{U}_j^{2}\rVert$. These pairwise distances were then reduced to a single dimension with a UMAP algorithm [42]. The obtained 1-dimensional embedding was then used to order the projections $\widetilde{U}_i^{2}$ (Figure 22a), which reveals two fuzzily separated sub-spaces. These two sub-spaces use respectively about 25% and 75% of the dimensions of the entire vector space.

Interestingly, ordering the columns by singular values is fairly consistent with these two sub-spaces. One reason for this result might be that latents projecting to the first sub-space have different encoder norms than latents projecting to the second sub-space Figure 22. This difference in norm can significantly guide the SVD to separate these two sub-spaces.

To further interpret these two sub-spaces, we manually looked at latents from each cluster. We found that latents from the smaller cluster tend to activate on relatively non-diverse vocabulary tokens. To quantify this insight, we first estimated $A_{i,v}$, the average squared activation of latent $i$ on vocabulary token $v$. Then, we normalized the vectors $A_i$ to sum to one, $\widetilde{A}_{i,v} = A_{i,v} / \sum_w A_{i,w}$ and computed the effective number of token $m_i = \exp(\sum_v \widetilde{A}_{i,v} \log(\widetilde{A}_{i,v}))$. The effective number of token is a continuous metric with values in $[1, n_{\mathrm{vocab}}]$, and it is equal to $k$ when a latent activates equally on $k$ vocabulary tokens. With this metric, we confirmed quantitatively Figure 22 that latents from the smaller cluster all activate on relatively low numbers of vocabulary tokens (less than 100), whereas latents from the larger cluster sometimes activate on a larger numbers of vocabulary tokens (up to 1000).

Miscellaneous small results

Impact of different locations

In a sweep across locations in GPT-2 small, we found that the optimal learning rate varies with layer and location type (MLP delta, attention delta, MLP post, attention post), but was within a factor of two.

Number of tokens needed for convergence is noticeably higher for earlier layers. MSE is lowest at early layers of the residual stream, increasing until the final layer, at which point it drops. Residual stream deltas have MSE peaking around layer 6, with attention delta falling sharply at the final layer.

When ablating to reconstructions, downstream loss and KL get strictly worse with layer. This is despite normalized MSE dropping at late layers. However, there appears to be an exception at final layers (layer 11/12 and especially 12/12 of GPT-2 small), which can have better normalized MSE than earlier layers, but more severe effects on downstream prediction (Figure 27).

We can also see that choice of layer affects different metrics differently (Figure 28). While earlier layers (unsurprisingly) have better N2G explanations, later layers do better on probe loss and sparsity.

In early results with autoencoders trained on the last layer of GPT-2 small, we found the results to be qualitatively worse than the layer 8 results, so we use layer 8 for all experiments going forwards.

Impact of token position

We find that tokens at later positions are harder to reconstruct (Figure 29). We hypothesize that this is because the residual stream at later positions have more features. First positions are particularly egregiously easy to reconstruct, in terms of normalized MSE, but they also generally have residual stream norm more than an order of magnitude larger than other positions (Figure 30 shows that in GPT-2 small, the exception is only at early layers and the final layer of GPT-2 small). This phenomenon was explained in [[60], [69]], which demonstrate that these activations serve as crucial attention resting states.

First token positions have significantly worse downstream loss and KL after ablating to autoencoder reconstruction at layers with these large norms (Figure 27), despite having better normalized MSE. This is consistent with the hypothesis that the large norm directions at the first position are important for loss on other tokens but not the current token. This position then potentially gets subtracted back out at the final layer as the model focuses on current-token prediction.

Irreducible loss term

In language models, the irreducible loss exists because text has some intrinsic unpredictableness—even with a perfect language model, the loss of predicting the next token cannot be zero. Since an arbitrarily large autoencoder can in fact perfectly reconstruct the input, we initially expected there to be no irreducible loss term. However, we found the quality of the fit to be substantially less good without an irreducible loss.

While we don’t fully understand the reason behind the irreducible loss term, our hypothesis is that the activations are made of a spectrum of components with different amount of structure. We expect less structured data to also have a worse scaling exponent. At the most extreme, some amount of the activations could be completely unstructured gaussian noise. In synthetic experiments with unstructured noise (see Figure 31), we find an $L(N)$ exponent of -0.04 on 768-dimensional gaussian data, which is much shallower than the approximately -0.26 we see on GPT-2-small activations of a similar dimensionality.

Further probe based evaluations results

Image not extracted

Contributions

Leo Gao implemented the autoencoder training codebase and basic infrastructure for GPT-4 experiments. Leo worked on the systems, including kernels, parallelism, numerics, data processing, etc. Leo conducted most scaling and architecture experiments: TopK and AuxK, tied initialization, number of latents, subject model size, batch size, token budget, optimal lr, $L(C)$, $L(N)$, random data, etc., and iterated on many other architecture and algorithmic choices. Leo designed and iterated on the probe based metric. Leo did investigations into downstream loss, feature recall/precision, and some early explorations into circuit sparsity.

Tom Dupré la Tour studied activation shrinkage, progressive recovery, and Multi-TopK. Tom implemented and trained the Gated and ProLU baselines, and trained and analyzed different layers and locations in GPT-2 small. Tom helped refine the scaling laws for $L(N, K)$ and $L_s(N)$. Tom discovered the latent sub-spaces (refining an idea and code originally from Carroll Wainwright).

Henk Tillman worked on N2G explanations and LM-based explainer scoring. Henk worked on infrastructure for scraping activations. Henk worked on finding qualitatively interesting features, including safety-related features.

Jeff Wu studied ablation effects and sparsity, and ablation reconstruction. Jeff managed infrastructure for metrics, and wrote the visualizer data collation and website code. Jeff analyzed overall cost of having a fully sparse bottleneck, and analyzed recurring dense features. Cathy and Jeff worked on finding safety-relevant features using attribution. Jeff managed core researchers on the project.

Gabriel Goh suggested using TopK, and contributed intuitions about TopK, AuxK, and optimization.

Rajan Troll contributed intuitions about and advised on optimization, scaling, and systems.

Alec Radford suggested using the irreducible loss term and contributed intuitions about and advised on the probe based metric, optimization, and scaling.

Jan Leike and Ilya Sutskever managed and led the Superalignment team.