Understanding and Steering Llama 3 with Sparse Autoencoders

We're releasing preview.goodfire.ai, a desktop interface to help you understand and steer Llama 3's behavior. To do this, we trained interpreter models (sparse autoencoders) on Llama-3-8B to extract modifiable "features" from Llama 1.

It's commonly assumed that neural networks - particularly the large language models that power most of the advances in modern AI products - are black boxes, with internals we can neither understand nor control. Recent advances in interpretability research have demonstrated that this assumption is incorrect: we can in fact train interpreter models that parse neural network activations into components that are often human-understandable (these components are called "features"). The most successful class of interpreter models so far are known as sparse autoencoders (SAEs) [1][2][3], which we give a short introduction to below. We can also intervene on the "features" discovered by interpreter models in order to steer model behaviour in a precise, quantitative way.

The intuition for why this model creates interpretable representations in the hidden layer is that neural network features 'want' to correspond to individual neurons, but by forcing them into a lower-dimensional representation, we cause them to get squashed together such that they no longer align with individual neurons. Although we think of neural network representations as very high-dimensional objects, in this picture they should be even higher-dimensional! This view is still largely an intuition with relatively limited empirical evidence outside of toy models [4], although the success of sparse autoencoders offers some support.

More formally, we extract LLM activations \(x\) from some point in the model (in our case the residual stream at layer 19) and train a two-layer MLP to autoencode:

where \(\sigma(\cdot)\) is some nonlinearity (we simply use ReLU, though other activation functions like top-k [5], Gated [6] and Jump-ReLU [7] have been tried as well). The hidden layer \(h\) is typically substantially wider than the dimension of the residual stream \(x\), which is a vector of size \(d_{\text{model}}\). The width of \(h\) is determined by the expansion factor \(\alpha\): \(d_{\text{hidden}} = \alpha d_{\text{model}}\).

Sparse autoencoders are typically trained using a reconstruction loss \(L_{\text{recon}}\) and a sparsity-promoting loss \(L_{\text{sparsity}}\). We use a mean squared error reconstruction loss and an L1 sparsity loss:

We tune the value of \(\beta\) but otherwise use the same settings as the Anthropic April update [8]. Our choices here were intentionally conservative: although more recent SAE architectures offer increases in performance along some proxy metrics, we were concerned that TopK and JumpReLU SAEs appear to have some very high-frequency features. Investigating these concerns in sufficient detail would have added considerable time to development of our first research preview, so we chose to prioritise simplicity & velocity. Deeply understanding the tradeoffs of SAE architectural alternatives in practical usage is an important topic for us in the future.

where \(||\cdot||_0\) is the \(L_0\) metric and \(h(x_i)\) is the hidden vector produced on input activation \(x_i\). We track this as an exponential moving average across batches in training. \(L_0\) compares the relative sparsity of different SAEs but doesn't directly say anything about the interpretability of those features, although it's a common belief that SAEs with lower \(L_0\) are more interpretable.

Our second important metric is fidelity, which measures the degree to which the SAE captures functionally-relevant components:

where \(l_{\mathrm{SAE}}(x)\) is the log-loss of the network on input \(x\) with the SAE inserted and \(l_{\mathrm{ablated}}(x)\) is the log-loss of the network on \(x\) with the SAE output set to zero.

Because our use case is primarily chat-focused, the SAE we use in our research preview was trained on activations harvested from Llama-3-8B-Instruct on the LMSYS-Chat-1M dataset [9]. This led to an interesting and diverse range of features relevant to our use case. Training on a non-chat model, or on non-chat datasets such as FineWeb [10], generated a range of interesting features but didn't transfer as well to chat applications. We conjecture that this is because many features that are causally relevant to model behaviour occur on chat-specific special tokens (such as features indicating that the model should respond in a certain style), and missing these by training on non-chat data prevents us from finding these behaviourally-relevant features.

To score feature explanations generated by automated interpretability, we use the generated explanation to distinguish between a sample that activates the feature and a 'distractor' sample, on which the feature is not activated. The proportion of correctly identified samples gives a score for the feature, and aggregating these provides a score for the automated interpretability method. The contrastive approach (which Anthropic have also tried [11]) has the additional effect of weakly penalising polysemanticity, which makes it an interesting measure for comparing both SAEs and automated interpretability methods. This approach differs from the 'simulation' approach in the original automated interpretability paper [12], and is similar to recent work from Eleuther [13].

Interventions

A key part of our research preview is the ability to surface causally-relevant aspects of the model's computation and intervene on them. These computational elements significantly impact the model's output, such that modifying them would lead to meaningful changes in result. To understand this process, we first need to explain how we introduce the SAE into model computations and perform interventions.

Remember that an SAE autoencodes a model's activations \(x\) at some layer (layer 19, in our case). The SAE prediction \(\hat{x}\) is imperfect, so there's an error term \(\epsilon\). We want to intervene on SAE features, so let's make the dependence on \(h\) explicit:

So now if we change \(h \rightarrow \tilde{h}\) (for instance by changing the value of a feature) we change the output of the SAE to \(\tilde{x} = f_{\mathrm{dec}}(\tilde{h}) + \epsilon\). The error term \(\epsilon\) is unchanged, so everything the SAE hasn't captured is unaltered by our intervention on \(h\). We now insert the modified activations \(\tilde{x}\) back into the model and continue inference through all the remaining layers to the model output.

Attribution

We surface good intervention candidates by doing gradient-based attribution to SAE features, which is easy to do with backprop (this approach shows the effect of features at all token positions to a single output token). The loss we found most effective is the logit of the predicted token minus the mean of the logits:

where \(\tau_t\) is the t-th token, \(\tau_{\lt t}\) are the tokens up to position t, \(\pi(\tau)\) is the logit of token \(\tau\), and \(\bar{\pi}\) is the mean of the logits. The intuition for using logits rather than the log-loss is that the gradient of \(\log p(\tau_t)\) combines both the gradient for features that promoted the chosen token, and for features that suppressed predictions of other tokens, which was confusing to interpret. We also obtained interesting results using contrastive explanations (i.e. gradient for the difference between a pair of tokens) but the UX flow was more complex. The reason we use the logit mean \(\bar{\pi}\) is to avoid surfacing features that increase the logit of many or all tokens.

Because the SAE is applied at every token position, computing \(\nabla_h\mathcal{L}_{\mathrm{attrib}}\) leads to a gradient matrix of shape [seq_len, n_features]. Summing across token position reliably highlights effective and interpretable features for interventions, whereas other approaches like taking the maximum or showing token positions separately are much less reliable. We then show the top k features by summed attribution and allow you to intervene on them. We scale all interventions to be between the maximum value seen in our autointerp dataset \(h_{\mathrm{max}}\) and \(-h_{\mathrm{max}}\) (although SAE encoders can't output a negative value, this allows 'blocking' of a feature).

Intervention phenomenology

Playing with attribution and intervention rapidly and at scale has surfaced some interesting observations. Some features are easy to reliably intervene on, whereas other very similar-looking ones have little or no effect. We conjecture that this could be due to cross-layer distributed representations (AKA cross-layer superposition): if a feature is split across multiple layers spanning the point at which the SAE was trained, then the SAE will only see a portion of the feature vector, with the remainder of the vector being unchanged as they get computed in layers after the SAE has been used. This means that the majority of the feature isn't getting intervened upon.

As with other SAEs, we often find repeated 'echo' features, which seem to activate on similar inputs. This is relatively unsurprising under an L1 sparsity penalty, as two features of activation strength 1/2 will have the same L1 penalty as a single feature of strength 1 (though using \(p\)-norms with \(p < 1\) would alleviate this). When we intervene on features with echoes we find that 'positive' interventions (i.e. turning on or strengthening a feature) normally work well with only a single feature intervention, whereas 'negative' interventions (trying to prevent a behaviour) frequently require more than one feature to be set to a negative value. This phenomenon could well be explained by self-repair, as a single positively-influenced feature may get amplified, whereas a single ablated feature could get restored. Deeply understanding the dynamics of feature interventions will be an important step towards more reliably steerable models.