Toy Models of Superposition
It would be very convenient if the individual neurons of artificial neural networks corresponded to cleanly interpretable features of the input. For example, in an “ideal” ImageNet classifier, each neuron would fire only in the presence of a specific visual feature, such as the color red, a left-facing curve, or a dog snout. Empirically, in models we have studied, some of the neurons do cleanly map to features. But it isn’t always the case that features correspond so cleanly to neurons, especially in large language models where it actually seems rare for neurons to correspond to clean features. This brings up many questions. Why is it that neurons sometimes align with features and sometimes don’t? Why do some models and tasks have many of these clean neurons, while they’re vanishingly rare in others?
In this paper, we use toy models — small ReLU networks trained on synthetic data with sparse input features — to investigate how and when models represent more features than they have dimensions. We call this phenomenon superposition
-
Linear algebraic structure of word senses, with applications to polysemy
S. Arora, Y. Li, Y. Liang, T. Ma, A. Risteski.
Transactions of the Association for Computational Linguistics, Vol 6, pp. 483–495. MIT Press. 2018.
-
Decoding The Thought Vector [link]
G. Goh. 2016.
-
Zoom In: An Introduction to Circuits
C. Olah, N. Cammarata, L. Schubert, G. Goh, M. Petrov, S. Carter.
Distill. 2020.
[1, 2, 3]
. When features are sparse, superposition allows compression beyond what a linear model would do, at the cost of “interference” that requires nonlinear filtering.
Consider a toy model where we train an embedding of five features of varying importance
1
Where “importance” is a scalar multiplier on mean squared error loss. in two dimensions, add a ReLU afterwards for filtering, and vary the sparsity of the features. With dense features, the model learns to represent an orthogonal basis of the most important two features (similar to what Principal Component Analysis might give us), and the other three features are not represented. But if we make the features sparse, this changes:
This figure and a few others can be reproduced using the toy model framework Colab notebook in our Github repo
Not only can models store additional features in superposition by tolerating some interference, but we’ll show that, at least in certain limited cases, models can perform computation while in superposition. (In particular, we’ll show that models can put simple circuits computing the absolute value function in superposition.) This leads us to hypothesize that theneural networks we observe in practice are in some sense noisily simulating larger, highly sparse networks. In other words, it’s possible that models we train can be thought of as doing “the same thing as” an imagined much-larger model, representing the exact same features but with no interference.
Feature superposition isn’t a novel idea. A number of previous interpretability papers have considered it
-
Linear algebraic structure of word senses, with applications to polysemy
S. Arora, Y. Li, Y. Liang, T. Ma, A. Risteski.
Transactions of the Association for Computational Linguistics, Vol 6, pp. 483–495. MIT Press. 2018.
-
Decoding The Thought Vector [link]
G. Goh. 2016.
-
Zoom In: An Introduction to Circuits
C. Olah, N. Cammarata, L. Schubert, G. Goh, M. Petrov, S. Carter.
Distill. 2020.
-
Softmax Linear Units
N. Elhage, T. Hume, C. Olsson, N. Nanda, T. Henighan, S. Johnston, S. ElShowk, N. Joseph, N. DasSarma, B. Mann, D. Hernandez, A. Askell, K. Ndousse, A. Jones, D. Drain, A. Chen, Y. Bai, D. Ganguli, L. Lovitt, Z. Hatfield-Dodds, J. Kernion, T. Conerly, S. Kravec, S. Fort, S. Kadavath, J. Jacobson, E. Tran-Johnson, J. Kaplan, J. Clark, T. Brown, S. McCandlish, D. Amodei, C. Olah.
Transformer Circuits Thread. 2022.
[1, 2, 3, 4]
, and it’s very closely related to the long-studied topic of compressed sensing in mathematics
-
Compressed sensing
D.L. Donoho.
IEEE Transactions on information theory, Vol 52(4), pp. 1289–1306. IEEE. 2006.
[5]
, as well as the ideas of distributed, dense, and population codes in neuroscience
-
Local vs. Distributed Coding
S.J. Thorpe.
Intellectica, Vol 8, pp. 3–40. 1989.
[6]
and deep learning
-
Representation learning: A review and new perspectives
Y. Bengio, A. Courville, P. Vincent.
IEEE transactions on pattern analysis and machine intelligence, Vol 35(8), pp. 1798–1828. IEEE. 2013.
[7]
. What, then, is the contribution of this paper?
For interpretability researchers, our main contribution is providing a direct demonstration that superposition occurs in artificial neural networks given a relatively natural setup, suggesting this may also occur in practice. That is, we show a case where interpreting neural networks as having sparse structure in superposition isn’t just a useful post-hoc interpretation, but actually the “ground truth” of a model. We offer a theory of when and why this occurs, revealing a phase diagram for superposition. This explains why neurons are sometimes “monosemantic” responding to a single feature, and sometimes “polysemantic”
-
Feature Visualization [link]
C. Olah, A. Mordvintsev, L. Schubert.
Distill. 2017.
[8]
responding to many unrelated features. We also discover that, at least in our toy model, superposition exhibits complex geometric structure.
But our results may also be of broader interest. We find preliminary evidence that superposition may be linked to adversarial examples and grokking, and might also suggest a theory for the performance of mixture of experts models. More broadly, the toy model we investigate has unexpectedly rich structure, exhibiting phase changes, a geometric structure based on uniform polytopes, “energy level”-like jumps during training, and a phenomenon which is qualitatively similar to the fractional quantum Hall effect in physics, among other striking phenomena. We originally investigated the subject to gain understanding of cleanly-interpretable neurons in larger models, but we’ve found these toy models to be surprisingly interesting in their own right.
Key Results From Our Toy Models
(Excerpt: opening motivation and main result through the introduction of the toy model; full text at source URL.)