Sparse dictionary learning

Context. Elhage et al. (2021) read attention-only transformers off their weights but stopped at the MLP, whose neurons are mostly polysemantic, and Toy Models of Superposition (Elhage et al., 2022) explained polysemanticity as superposition and proposed two remedies: train models without it, or find an overcomplete basis after training. The first route failed. Even models trained with one-hot activations, which rule superposition out, had polysemantic neurons, because under cross-entropy a binary neuron that fires for two of four mutually exclusive next-token features (loss $\ln2\approx0.69$ ) has lower loss than one that fires for a single feature ( $\tfrac34\ln3\approx0.82$ ). So the paper takes the second route on the smallest model with an MLP, a one-layer transformer with 512 neurons. If MLP activations are sparse nonnegative combinations of many more directions than neurons, a wide ReLU encoder with an $L_1$ penalty, trained on billions of activations to reconstruct them, recovers directions far more interpretable than neurons. The evidence that the model itself uses them is that their effects on the output, which the dictionary never saw, match their activations. This is sparse recovery with an unknown dictionary and one ReLU layer in place of an $\ell_1$ program, the same weak decoder as in the toy models.

Setting. Two one-layer transformers (seeds A and B) with a 512-neuron ReLU MLP, trained on 100 billion tokens of the Pile. The sparse autoencoder acts on MLP activations $x\in\mathbb R^{512}$ : $\bar x=x-b_d,\qquad f=\mathrm{ReLU}(W_e\bar x+b_e),\qquad\hat x=W_df+b_d,\qquad\mathcal L=\frac1{|X|}\sum_{x\in X}\|x-\hat x\|_2^2+\lambda\|f\|_1,$ with unit-norm decoder columns $d_i$ , the feature directions, so that $x\approx b_d+\sum_if_i(x)\,d_i$ . It is trained on 8 billion activation vectors at sizes from 512 to 131,072 features; most analysis is on run A/1 with 4,096 features, excluding dead ones (never active) and an ultralow-density cluster the authors regard as a training artifact. A feature's logit weights are its direction pushed through the MLP output weights, a linearized layer norm and the unembedding, $d_iW_{\rm down}\pi LW_U$ . Every result is empirical.

Main findings. Features are much more interpretable than neurons, by a blinded human rubric, by automated explanation and prediction of activations, and by prediction of logit effects. The reconstruction carries most of what the MLP does, features recur across seeds much more than neurons do, and features split into finer families as the dictionary grows; some chain through the generated tokens into loops that act like finite-state automata.

Step 1

Hand-checked features fire specifically for one context. For four features whose contexts are easy to compute, the authors score text with a proxy, the log-likelihood ratio $\log P(s\mid\text{context})/P(s)$ against a unigram baseline, and compare it with the activations. The Arabic-script feature A/1/3450 is the clearest: Arabic script is 0.13% of training tokens but 81% of the tokens where the feature is active (98% at high activations), and its correlation with the proxy is 0.74. The DNA, Hebrew and base64 features reach 0.8, 0.55 and 0.38, the last low mainly because the proxy is too broad. Gaps in sensitivity are often filled by closely related features. So, at least at high activations, these units do not mix unrelated contexts. Weak activations are less specific, and the four were chosen for having easy proxies.

Step 2

Their effects on the output agree with their activations. The autoencoder sees only activations, never their downstream use, so interpretable effects on the output are a property of the model. The Arabic feature's logit weights single out Arabic-script tokens. Zeroing the feature throughout a context lowers the model's likelihood of Arabic tokens at all activation levels, and pinning it at its maximum while sampling makes the model write Arabic; base64 behaves the same way. On the data side, a dictionary trained on a transformer with shuffled weight entries yields mostly single-token features and nothing interpretable beyond them, so the richer structure is built by training.

Step 3

The features do not correspond to individual neurons. Only one neuron has any Arabic in its top 20 activating examples, and only one such example. The Arabic direction spreads over 27 neurons, and its most correlated neuron responds to many non-English languages, a natural group to share neurons in superposition, since languages rarely co-occur. The Hebrew and base64 features have no good neuron counterpart either. These features would be invisible in the neuron basis.

Step 4

The typical feature is interpretable, mostly at its stronger activations. A blinded annotator scored features and neurons on a rubric for confidence, consistency and specificity, with examples drawn evenly across activation levels so that a unit's score reflects its weak activations as well as its top examples. The median score over feature activation intervals is 12, a confident, specific hypothesis consistent with the logits, and the median neuron 0, no hypothesis at all. Automated scoring agrees: when a language model explains a unit from activating examples and predicts activations on unseen text from that explanation, features beat neurons clearly, and predicting from the explanation which tokens a unit upweights succeeds 74% of the time for features against 58% for neurons, with chance at 50%. The authors give the caveats: interpretability falls off at low activations, where most activations lie, though most of a feature's effect comes from its large ones, and the scored features were drawn uniformly, so an importance-weighted sample could differ.

Step 5

The features account for most of what the MLP does, though not all of it. Replacing the MLP activations by the reconstruction recovers 79% of the MLP's loss reduction for A/1, and 94.5% for the 131,072-feature run with a smaller $L_1$ coefficient. The authors call this an overestimate, since the baseline is the loss with the MLP zero-ablated and a transformer trained from scratch with a small MLP reaches similar percentages, and they expect a long tail of rare features behind the remainder. Feature-density histograms agree: larger dictionaries keep finding rarer features, so even 131,072 does not exhaust this 512-neuron layer.

Step 6

Features recur across seeds and split as the dictionary grows. Matching each A/1 feature to its most correlated B/1 feature gives a median activation correlation of 0.72, against 0.46 for best-matched neurons. Logit weights agree less, because most of them are small interference weights on tokens that do not matter: a PLOS citation feature has activation correlation 0.98 across seeds but negative logit-weight correlation, with only the following "." agreeing. Weighting by the token that actually comes next (attribution similarity) removes those weights and tracks activation similarity, so matched features help predict the same tokens. Within one model, families refine with size. The single base64 feature at 512 becomes three at 4,096, covering letters, single digits (a feature that predicts non-digits, since BPE would have merged two adjacent digits into one token) and base64 that encodes ASCII text. Families form tight clusters of nearby directions, and the odd single-token features of small dictionaries appear to be collapsed families.

Consequences and limits. Dictionary learning gives a better unit of analysis than the neuron, and it changes the picture from Toy Models of Superposition: features cluster in families that share output actions instead of spreading isotropically, and the number of features may be ill-posed, with dictionary size choosing the resolution at which they are seen. The limits are mostly the authors' own. There is no ground truth or trusted metric for a good dictionary, so they steer by inspection, feature density and reconstruction; loss recovered uses a weak baseline; low activations stay polysemantic; many token-in-context features may be a local code in the model (one unit per token-context pair) or an artifact of the $L_1$ penalty; and the downstream checks are unusually easy in a one-layer model, whose path to the logits is linear. Templeton et al. (2024) then showed the method works at scale, extracting millions of features from the residual stream of a production language model.

Terms

Sparse autoencoder: $f=\mathrm{ReLU}(W_e(x-b_d)+b_e)$ , $\hat x=W_df+b_d$ ; MSE plus $\lambda\|f\|_1$
Feature direction $d_i$: unit-norm decoder column; $x\approx b_d+\sum_if_i(x)d_i$
Expansion factor: features per neuron: 1× (512) to 256× (131,072); main run A/1 is 8×
Feature density: fraction of tokens on which a feature is nonzero
Ultralow-density cluster: features active on under one in a million tokens; treated as a training artifact
Logit weights: $d_iW_{\rm down}\pi LW_U$ : linear effect of a feature on the next-token logits
Proxy: $\log P(s\mid\text{context})/P(s)$ , a computable score of the hypothesized context
Loss recovered: share of the MLP's loss reduction kept by the reconstruction; 79% for A/1
Activation similarity: correlation of two features' activations over the same tokens; universality measure
Attribution similarity: correlation of activation times logit weight of the actual next token
Feature splitting: one feature in a small dictionary becomes a family of finer ones in a larger one

Quiz

Why did training models with one-hot activations, which rules out superposition, still give polysemantic neurons?
What does the dictionary trained on a transformer with shuffled weights show?
The A/1 reconstruction recovers 79% of the MLP's loss reduction. Why do the authors call this an overestimate?
As the dictionary grows from 512 to 4,096 features, the single base64 feature becomes three. What does this suggest about counting features?

Sources

Trenton Bricken et al. Towards Monosemanticity: Decomposing Language Models With Dictionary Learning. Transformer Circuits Thread, 2023.

Nelson Elhage et al. A Mathematical Framework for Transformer Circuits. Transformer Circuits Thread, 2021.
Nelson Elhage et al. Toy Models of Superposition. Transformer Circuits Thread, 2022.
Adly Templeton et al. Scaling Monosemanticity. Transformer Circuits Thread, 2024.