Toy Models of Superposition

interpretabilitysuperpositionfeaturespolysemanticitytoy models

Elhage et al. train small ReLU networks on synthetic sparse data and show that they store more features than they have dimensions (superposition), that whether they do so changes abruptly with feature sparsity and importance, that the stored features form regular geometric arrangements, and that a network can compute while in superposition; they argue this is the main obstacle to reading features off neurons.

Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer et al. (Anthropic and Harvard) · 2022-09-14 · original

Interpretability research tries to explain a trained neural network in terms of parts a person can understand. The easiest case is when single neurons (the individual units of a layer, each producing one number per input) respond to one recognizable thing, such as a curve in an image. In vision models many neurons behave this way, but in language models such clean neurons are rare: most neurons respond to mixtures of unrelated inputs. This paper, published on the Transformer Circuits Thread on 14 September 2022, asks why neurons are sometimes clean and sometimes not. Its answer is a hypothesis it then demonstrates in small models: networks represent more features than they have neurons, packing them into shared directions at the cost of some interference. The paper calls this superposition. It matters for the rest of the interpretability notes in this library because it explains why later work stopped reading neurons and started extracting features with dictionary learning.

Terms used in this note

  • Feature. A property of the input that the network represents, such as "this token is part of a DNA sequence". The paper admits that defining the term is hard and works with the idea of a property that a large enough network would dedicate a neuron to ("What are Features?").
  • Direction in activation space. A layer's output for one input is a vector of numbers (its activations). The linear representation hypothesis is that each feature corresponds to a direction in that vector space, and the strength of the feature is how far the activation vector points along it ("Features as Directions").
  • Privileged basis. A layer has a privileged basis when something in the architecture, usually a nonlinear activation function applied to each neuron separately, makes the neuron directions special. Word embeddings and the transformer's residual stream (the running sum that each layer reads from and writes to; see Assembling the Transformer Block) have no privileged basis, so there is no reason to expect a single coordinate to mean anything ("Privileged vs Non-privileged Bases").
  • Monosemantic and polysemantic. A neuron that responds to one feature is monosemantic; one that responds to several unrelated features is polysemantic.
  • Superposition. A representation with more features than dimensions, each feature stored as a direction that is almost, but not exactly, orthogonal to the others. Because the directions overlap, an active feature makes others look slightly active; the paper calls this interference.
  • Sparsity. How rarely a feature is active. In the paper's synthetic data, each feature is zero with probability S, so S near 1 means very sparse features.

What the paper does

The core toy model

The main experiment asks whether a network can compress a vector of n features into m < n hidden dimensions and recover it ("Experiment Setup"). Each input feature is zero with probability S and otherwise uniform between 0 and 1, and each feature has an importance that weights its error in the loss. Two models are compared, differing only in the last step:

  • a linear model, which multiplies the input by a weight matrix W to get the hidden vector and multiplies by the transpose of W (plus a bias) to reconstruct it;
  • a ReLU output model, identical except that a ReLU (which sets negative values to zero) is applied to the reconstruction.

The loss is squared reconstruction error weighted by feature importance. The authors visualize the trained model by the product of W's transpose with W, a features-by-features matrix whose off-diagonal entries show which features share directions, and by two per-feature quantities: the length of each feature's weight vector (is the feature represented at all?) and how much other features project onto its direction (does it share space?).

Variations

The same setup is then used to study:

  • phase diagrams, with 2 features in 1 dimension and 3 features in 2 dimensions, varying the sparsity and the relative importance of one extra feature, and comparing trained models with closed-form losses of candidate solutions ("Superposition as a Phase Change");
  • geometry, with 400 features of equal importance in 30 dimensions, measuring a per-feature dimensionality, the fraction of a hidden dimension each feature effectively occupies ("The Geometry of Superposition");
  • correlated and anticorrelated features, where groups of features switch on together or never together;
  • learning dynamics over training;
  • adversarial vulnerability, attacking each trained model with small input perturbations;
  • a ReLU hidden layer model, which adds a ReLU to the hidden layer so that its dimensions become neurons in a privileged basis ("Superposition in a Privileged Basis");
  • computation in superposition, where a model must compute the absolute value of each input feature through a ReLU hidden layer, which requires real nonlinear work rather than storage ("Computation in Superposition").

What it finds

Superposition appears when features are sparse

With n = 20 features and m = 5 dimensions, the linear model always learns only the five most important features, the analogue of keeping the top principal components ("Basic Results"; compare the low-rank approximation in Latent Semantic Indexing). The ReLU output model does the same when features are dense. As sparsity increases, it starts representing additional features in directions that are not orthogonal, beginning with the least important features: first as antipodal pairs (two features stored as exact opposites of one direction), then in other arrangements. A model with 20 dimensions and 80 features behaves the same way at larger scale.

The section "Mathematical Understanding" explains why. For a linear model, the loss has a benefit term for representing a feature and an interference term for overlap, and it is never worth representing more features than dimensions. With a ReLU on the output and very sparse inputs, negative interference is removed for free by the ReLU, and a negative bias can turn small positive interference into negative interference, so packing in extra features becomes worthwhile. The authors note a resemblance between the resulting objective and the Thomson problem of spacing charged points on a sphere.

The change is abrupt

In the phase-diagram experiments, a feature ends up in one of three states: not learned, learned in superposition, or given its own dimension. Which state wins depends on sparsity and relative importance, and the switch between states is sharp in both the trained models and the closed-form "toy model of the toy model"; the authors confirm analytically that it is a first-order phase change ("Superposition as a Phase Change").

The features form regular shapes

With 400 features in 30 dimensions, the number of dimensions per learned feature, plotted against sparsity, sticks at 1 and at 1/2, the latter corresponding to antipodal pairs ("Uniform Superposition"). Splitting this by feature shows further clusters of dimensionality at 3/4, 2/3, 1/2, 2/5 and 3/8, which the paper matches to a tetrahedron, a triangle, an antipodal pair, a pentagon and a square antiprism ("Why these geometric structures?"). Unequal importance or sparsity deforms these shapes smoothly until they snap to a different shape ("Non-Uniform Superposition"). Correlated features tend to be placed orthogonally to each other, and when there is not enough room the model represents their shared principal component instead, so that superposition and a PCA-like solution trade off; anticorrelated features prefer to share space with negative interference ("Correlated and Anticorrelated Features"). The authors say this geometry seems "too elegant to be true" and may be partly an artifact of the toy model.

Training proceeds in jumps, and superposition costs robustness

During training, features jump between dimensionalities, and each jump coincides with a drop in the loss ("Phenomenon 1: Discrete 'Energy Level' Jumps"). In the adversarial experiments, vulnerability to small worst-case perturbations rose sharply as superposition formed, by more than three times, and tracked the number of features per dimension ("Relationship to Adversarial Robustness"). Adversarial training reduced superposition, but removing it fully required attacks of 80% of the input's norm. The authors decline to claim that superposition explains adversarial examples in real models.

Neurons can be monosemantic or polysemantic, sometimes in the same model

In the ReLU hidden layer model (10 features, 5 neurons), dense features give neurons that each represent one feature; as sparsity rises, polysemantic neurons appear, and both kinds can coexist in one model ("Superposition in a Privileged Basis"). The authors note a weakness of this model: it gains nothing from the hidden ReLU and avoids using it when it can.

Computation can happen in superposition

In the absolute-value task, a model with 100 features and 40 neurons dedicates neurons to single features when inputs are dense and uses superposition to handle more features than neurons when inputs are sparse ("Superposition vs Sparsity"). Neurons for the most important features tend to stay monosemantic, and many neurons combine a large weight for one primary feature with small weights for secondary ones, so their strongest activations look clean and weaker ones look mixed. The authors point out that this matches observations of language-model neurons that appear interpretable from their top examples but respond to other things at lower activation. They also describe a specific mechanism, asymmetric superposition, in which one neuron stores two features at different scales and a second neuron inhibits the resulting interference ("The Asymmetric Superposition Motif").

Why it matters for safety, in the authors' framing

The section "The Strategic Picture of Superposition" argues that being able to list all of a model's features would let researchers say that certain kinds of circuit do not exist, and that superposition blocks this. It names three routes out: train models without superposition (for example with an L1 penalty on activations, at a cost in performance), find an overcomplete basis after training with sparse coding or dictionary learning (learning a larger set of directions such that each activation is a sparse combination of them), or combine the two. The second route is the one later notes follow.

Limitations and open questions

  • Toy models only. The models are small ReLU networks on synthetic data. The authors expect some findings (superposition itself, the mix of monosemantic and polysemantic neurons, perhaps the link to adversarial examples) to generalize and regard the geometry and learning-dynamics results as much more uncertain ("To What Extent Does Superposition Exist in Real Models?"). In the introduction they say it is "very unclear what to generalize to real networks".
  • Evidence for real models is indirect. The case that real networks use superposition rests on consistency with observations, such as polysemantic neurons and InceptionV1 having more polysemantic neurons in later layers, not on a direct test.
  • Definitions are unsettled. The paper offers three working definitions of a feature and says it is probably important not to commit to one too early.
  • Open questions listed by the authors ("Open Questions") include a statistical test for superposition, how to control whether it forms, closed-form solutions, how to estimate the importance and sparsity of real features, whether superposition shrinks or grows with scale, and how much computation can be done in superposition.
  • Later work. The overcomplete-basis route was applied to a real one-layer transformer a year later in Towards Monosemanticity, which reports that real features cluster more densely than the isotropic picture in this paper suggested.

Research questions

  1. Can the importance and sparsity curves of a real model be estimated? The toy model predicts behavior from these two curves, but they are unknown for real networks. A method that infers them from models of different sizes or regularization strengths, checked against features later found by dictionary learning, would test whether the toy theory transfers.
  2. Does the phase change appear in real models? The paper shows sharp transitions in toy settings only. Tracking whether particular features in a real model move abruptly between unrepresented, superposed and dedicated states across model sizes would answer it.
  3. How much does superposition contribute to adversarial vulnerability in practice? The toy result is a more than threefold increase. Comparing vulnerability with a measure of superposition across real models of matched accuracy would show whether the relationship holds outside the toy setting.
  4. Is the observed geometry specific to these models? The polytope structure may be an artifact of tied weights and synthetic uniform features. Repeating the geometry analysis with untied weights, other activation functions and non-uniform feature distributions would show how much survives.
  5. What kinds of computation can be done in superposition? Absolute value is a simple, sparse computation. Testing functions that combine several features at once would show whether the capacity gain persists when features must interact.
  6. Can superposition be removed at acceptable cost? The paper suggests L1 penalties and mixture-of-experts architectures. Measuring the loss penalty of superposition-free models at increasing scale would settle whether the first route out is practical.

How it connects

References

  • Elhage, N., Hume, T., Olsson, C., Schiefer, N., Henighan, T., Kravec, S., Hatfield-Dodds, Z., Lasenby, R., Drain, D., Chen, C., Grosse, R., McCandlish, S., Kaplan, J., Amodei, D., Wattenberg, M., and Olah, C. (2022). Toy Models of Superposition. Transformer Circuits Thread. https://transformer-circuits.pub/2022/toy_model/index.html (also arXiv:2209.10652). Section names in this note refer to the web version.
  • Anthropic. (2022, September 14). Toy Models of Superposition [research page]. https://www.anthropic.com/research/toy-models-of-superposition
  • Code for the toy models: https://github.com/anthropics/toy-models-of-superposition