Sparse Autoencoders Find Highly Interpretable Features in Language Models

In plain words

Neural networks pack many concepts into the same set of numbers, so a single internal unit often responds to several unrelated things. That makes it hard to say what the network is doing by looking at its units one at a time.

This paper looks for better building blocks. It trains a small autoencoder, a network that squeezes a signal down and rebuilds it, on the internal activations of a language model. The autoencoder is pushed to rebuild each activation from only a few active units, and it has more units than the activation has dimensions. Each learned direction in the activations becomes a "feature" that can be inspected.

To check whether the features are easier to understand, the authors use an automated method: a language model writes a short description of what a feature responds to, and then predicts where that feature fires in new text. How well the predictions match the real activations is the score. On this measure the learned features beat single neurons and standard decompositions. A small set of them can also be used to find the parts of a model that drive one specific behaviour.

TL;DR

A sparse autoencoder with a ReLU encoder and an L1 penalty on its hidden code is trained on one layer's activations (residual stream, MLP or attention). Its decoder rows form an overcomplete dictionary of features §2 eq. 1 eq. 2 eq. 4. Author-reported: the features score higher on automated interpretability (top-and-random) than the default basis, random directions, PCA and ICA §3.2 fig. 2. Weaker results: the advantage shrinks in later layers, and the MLP-layer pipeline learns many dead features §3.2 app. E. In the indirect-object-identification (IOI) task, patching a few dictionary features reaches a given change in output with fewer patches than PCA components §4.2 fig. 3. Read this if you want to train a feature dictionary on a model's activations or to reproduce the interpretability comparison. The reconstruction loss in the paper and in the code differ in scale (see Paper vs code).

Problem

Reverse engineering a network needs units that can be analysed in isolation. Individual neurons are often polysemantic, responding to several unrelated feature types, and the residual stream has no reason to align with the neuron basis §1. The paper follows the superposition hypothesis: a layer represents more features than it has dimensions, so the features form an overcomplete, non-orthogonal set that only works if each feature is sparsely active §1. Recovering such features is framed as sparse dictionary learning §1 §2.

Measuring interpretability by hand does not scale to thousands of features, so the paper uses the automated protocol of Bills et al. (an LLM explainer and an LLM simulator) §3.1 app. A.

Method

The method has three parts: training a feature dictionary on cached activations, scoring the features automatically, and testing the features on one causal task fig. 1.

Dictionary training §2. Activations are cached from a chosen layer of Pythia-70M or Pythia-410M over the Pile app. B. The autoencoder has one hidden layer of size d_hid = R·d_in, where R is the overcompleteness ratio and d_in is the residual width (512 for Pythia-70M, 1024 for Pythia-410M) §2. The encoder applies a ReLU to a linear map plus bias, and the decoder is the transpose of the row-normalised dictionary matrix M in the tied case §2 eq. 1 eq. 2. Training minimises squared reconstruction error plus alpha times the L1 norm of the hidden code §2 eq. 4. Each row of M is normalised so the L1 term cannot be reduced by growing the feature vectors §2, footnote on row normalisation.

# Sparse autoencoder loss, tied weights  [§2] [eq. 1] [eq. 2] [eq. 4]
# Input:  activations (B, d_in) from one layer; alpha (L1 weight)
# Output: scalar loss; dictionary rows are the features (d_hid, d_in)
def sae_loss(activations, dictionary, bias, alpha):
    # Normalise each dictionary row to unit length  [§2]
    dictionary = dictionary / norm(dictionary, axis=1)[:, None]
    # Encode: codes = ReLU(M x + b), one code per feature  [eq. 1]
    codes = relu(activations @ dictionary.T + bias)  # (B, d_hid)
    # Decode: x_hat = M^T c, a sparse sum of feature rows  [eq. 2]
    reconstruction = codes @ dictionary              # (B, d_in)
    # Squared error summed over d_in, per example  [eq. 4]
    error = sum((activations - reconstruction) ** 2, axis=-1)
    # L1 penalty on the codes, weighted by alpha  [eq. 4]
    sparsity = alpha * sum(abs(codes), axis=-1)
    # Mean over the batch; the code does this, eq. 4 is per example
    # [code: autoencoders/sae_ensemble.py#L148-L149]
    return mean(error + sparsity)

The untied variant, used for MLP layers, learns a separate encoder matrix and decoder matrix instead of the transpose app. C eq. 5 eq. 6. The paper says this retains more features in MLP layers.

Alpha sweep app. B. Varying alpha traces a smooth sparsity-versus-unexplained-variance curve with no knee, which the authors read as weak evidence that there is no single correct decomposition app. B.

Automated interpretability score §3.1 app. A. The protocol has five steps:

1. Take the first 50,000 lines of OpenWebText. Cut each into a 64-token fragment and record the feature's activation per token, rescaled to integers from 0 to 10 app. A, step 1.
2. Take the 20 fragments with the highest activations. Give 5 of them, with per-token activations, to the explainer LLM, which writes an explanation app. A, step 2.
3. Ask the simulator LLM to predict per-token activations on 5 highly activating fragments and 5 random fragments with non-zero variation app. A, step 3.
4. The score is the correlation between predicted and real activations app. A, step 4. The paper uses this top-and-random mixture unless noted app. A.
5. Skip any feature with fewer than 20 fragments showing non-zero variation app. A, step 5.

# Automated interpretability score for one feature  [§3.1] [app. A]
# Input:  per-token activations on 50,000 fragments of 64 tokens
# Output: correlation in [-1, 1], or None if the feature is skipped
def interpretability_score(activations, explainer, simulator):
    # Skip features with too few fragments that vary  [app. A, step 5]
    if count_varying(activations) < 20:
        return None
    # Pick the 20 fragments with the highest activation  [app. A, step 2]
    top = top_fragments(activations, 20)
    # The explainer reads 5 of them and writes a description  [app. A, step 2]
    explanation = explainer(top[:5])
    # Simulate on 5 top and 5 random fragments  [app. A, step 3]
    predictions = simulator(explanation, top[5:10] + random_fragments(5))
    # Correlate predicted and real activations  [app. A, step 4]
    return correlation(predictions, real_activations(top[5:10] + random))

Causal test on IOI §4.1. The model runs on a counterfactual sentence, and the dictionary codes of the target sentence replace the base sentence's codes on the chosen features. The residual vector at each position becomes x' = x + sum over chosen features j of (cbar_j − c_j) f_j §4.1, unnumbered display. The KL divergence between the patched and target output logits is the metric. Features are ordered with Algorithm 4.1 of ACDC §4.1, treating each feature as an independent contributor averaged over 50 IOI prompts. The first k features form the patch set §4.1.

Case studies §5. Single features are checked for their input tokens (monosemanticity), their effect on output logits when the feature is ablated, and their upstream features found by ablation §5.1 §5.2 §5.3. Ablation uses less-than-rank-one editing: the activation is lowered along the feature direction only until the feature turns off §5.2.

Hyperparameters

Results

Author-reported. Setup: Pythia-70M and Pythia-410M residual streams; the Pile for training; GPT-4 explainer and GPT-3.5 simulator for scoring; baselines are the default basis (with negatives set to zero), random directions, PCA and ICA §3.2 fig. 2.

Paper vs code

The official repository is https://github.com/HoagyC/sparse_coding, owned by HoagyC (the first author, Hoagy Cunningham). The link is in a footnote the arXiv build does not show (it sits inside an \ificlrfinal block) §1. Checked at commit 69c5ae06813ee77bafa679dfacee33b21395d8e4. The repository has no LICENSE file, so code may be cited by path and line but not quoted verbatim.

Differences:

• Reconstruction scale. Eq. 4 sums the squared error over the d_in dimensions eq. 4. The code takes the mean over batch and dimensions: untied sae_ensemble.py, tied sae_ensemble.py. For the same alpha the code's sparsity term is therefore relatively stronger than Eq. 4 by a factor of d_in (512 or 1024 here). The paper does not say which convention produced alpha = 8.6e-4 or 3.2e-4.
• Bias decay. The code adds an optional L2 penalty on the encoder bias sae_ensemble.py sae_ensemble.py. Eq. 4 has no such term. The sweep script varies bias decay over 0, 0.05 and 0.1 big_sweep_experiments.py. The paper does not mention bias decay, and which value the reported results used is not stated.
• Simulator model. The paper's appendix says the simulator is GPT-3.5 app. A. The code sets the simulator to text-davinci-003 interpret.py. The explainer, GPT-4, matches interpret.py.
• Input centring. The code has an optional dataset-centring flag, off by default activation_dataset.py activation_dataset.py, and the tied model carries a centring transform that defaults to identity sae_ensemble.py. The paper does not describe centring. Whether the reported runs centred their data is not stated.
• IOI ordering. The ACDC ordering (Algorithm 4.1 of Conmy et al.) used in §4.1 was not found in the repository. A search for "acdc" and for KL-divergence patching found only plotting scripts that load precomputed kl_div_scores_layer_*.pt files plot_kl_div.py.
• Matches. ReLU applied after the bias sae_ensemble.py and sae_ensemble.py; row normalisation of the dictionary sae_ensemble.py and sae_ensemble.py; Adam with learning rate 1e-3 basic_l1_sweep.py. No discrepancy found for these.

Pitfalls

• Loss scale. Choose one convention for the reconstruction term and state it. Eq. 4 with the paper's alpha values cannot be reproduced exactly from the paper alone (see Paper vs code).
• Batch size and optimiser details (warmup, decay, number of steps per epoch) are not stated app. B. The code's default batch size is 256 basic_l1_sweep.py.
• Dead features. In MLP layers many features never activate, even at alpha = 0 app. E. In a dictionary twice the neuron count, fewer than half may be active enough to score app. C. The authors plan to reinitialise dead features app. E, which is not in this work.
• Overcompleteness is not guaranteed. Mid and late MLP layers do not yield a truly overcomplete basis with this pipeline §6.2 app. C.
• Lossy reconstruction. Reconstruction does not reach zero error, and substituting it into a model raises perplexity (25 → 40 in one layer) §6.2. The paper suggests that the objective should target output change, not reconstruction error, but does not implement this §6.2.
• Correlation mismatch. The text says feature skew and kurtosis correlate with top-and-random scores at 0.19 and 0.24. Table 2 lists 0.20 and 0.15 app. C tab. 2. Use the table, but check the source before relying on either.
• Score depends on the simulator. The score is a correlation from an LLM simulator that outputs predicted activations, and the simulator needs token log-probabilities app. A. Reproduction needs the same models, or the scores will not be comparable. The paper and the code name different simulators (see Paper vs code).
• Sampling. Features are heavy-tailed, so random-only scoring with limited samples underestimates correlation app. C. Top-and-random is the headline score app. A.
• Feature ordering is arbitrary. The five features in Table 1 are the first five in an arbitrary dictionary order, not a selection §3.1.
• Cost. One training run takes under an hour on one A40 app. B. Scoring calls an explainer and a simulator LLM for each feature app. A, and the figures use 150 features per dictionary fig. 2.
• Not in the compiled paper. The concept-erasure section and its appendix are excluded from the arXiv build with \iffalse §4.1 app. D. The repository has erasure plots, but this card makes no claims from them.
• Only one IOI layer in the main text. Layer-wise IOI results are in app. F and fig. 3 shows layer 11 only.

Lineage

• Builds on: the superposition hypothesis of Elhage et al. (2022) §1; sparse dictionary learning §2; the sparse-autoencoder work of Sharkey et al. (2023), which motivated this paper §1 §6.1.
• Uses: automated interpretability (Bills et al., 2023) §3.1 app. A; activation patching and ACDC (Conmy et al., 2023) §4.1.
• Compared against: default basis, random directions, PCA and ICA §3.2; the transformer-factor work of Yun et al. (2021), which applies sparse dictionary learning to all residual layers §6.1.
• Superseded by: not known from content already in this wiki.