Towards Monosemanticity: Decomposing Language Models With Dictionary Learning
Mechanistic interpretability seeks to understand neural networks by breaking them into components that are more easily understood than the whole. By understanding the function of each component, and how they interact, we hope to be able to reason about the behavior of the entire network. The first step in that program is to identify the correct components to analyze.
Unfortunately, the most natural computational unit of the neural network – the neuron itself – turns out not to be a natural unit for human understanding. This is because many neurons are polysemantic: they respond to mixtures of seemingly unrelated inputs. In the vision model Inception v1, a single neuron responds to faces of cats and fronts of cars
-
Feature Visualization [link]
C. Olah, A. Mordvintsev, L. Schubert.
Distill. 2017.
[1]
. In a small language model we discuss in this paper, a single neuron responds to a mixture of academic citations, English dialogue, HTTP requests, and Korean text. Polysemanticity makes it difficult to reason about the behavior of the network in terms of the activity of individual neurons.
One potential cause of polysemanticity is 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.
-
Toy Models of Superposition
N. Elhage, T. Hume, C. Olsson, N. Schiefer, T. Henighan, S. Kravec, Z. Hatfield-Dodds, R. Lasenby, D. Drain, C. Chen, R. Grosse, S. McCandlish, J. Kaplan, D. Amodei, M. Wattenberg, C. Olah.
Transformer Circuits Thread. 2022.
[2, 3, 4, 5]
, a hypothesized phenomenon where a neural network represents more independent “features” of the data than it has neurons by assigning each feature its own linear combination of neurons. If we view each feature as a vector over the neurons, then the set of features form an overcomplete linear basis for the activations of the network neurons. In our previous paper on Toy Models of Superposition
-
Toy Models of Superposition
N. Elhage, T. Hume, C. Olsson, N. Schiefer, T. Henighan, S. Kravec, Z. Hatfield-Dodds, R. Lasenby, D. Drain, C. Chen, R. Grosse, S. McCandlish, J. Kaplan, D. Amodei, M. Wattenberg, C. Olah.
Transformer Circuits Thread. 2022.
[5]
, we showed that superposition can arise naturally during the course of neural network training if the set of features useful to a model are sparse in the training data. As in compressed sensing, sparsity allows a model to disambiguate which combination of features produced any given activation vector.
1
For more discussion of this point, see Distributed Representations: Composition and Superposition.
In Toy Models of Superposition, we described three strategies to finding a sparse and interpretable set of features if they are indeed hidden by superposition: (1) creating models without superposition, perhaps by encouraging activation sparsity; (2) using dictionary learning to find an overcomplete feature basis in a model exhibiting superposition; and (3) hybrid approaches relying on a combination of the two. Since the publication of that work, we’ve explored all three approaches. We eventually developed counterexamples which persuaded us that the sparse architectural approach (approach 1) was insufficient to prevent polysemanticity, and that standard dictionary learning methods (approach 2) had significant issues with overfitting.
In this paper, we use a weak dictionary learning algorithm called a sparse autoencoder to generate learned features from a trained model that offer a more monosemantic unit of analysis than the model’s neurons themselves. Our approach here builds on a significant amount of prior work, especially in using dictionary learning and related methods on neural network activations (e.g.
-
Sparse overcomplete word vector representations
M. Faruqui, Y. Tsvetkov, D. Yogatama, C. Dyer, N. Smith.
arXiv preprint arXiv:1506.02004. 2015.
-
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.
-
Spine: Sparse interpretable neural embeddings
A. Subramanian, D. Pruthi, H. Jhamtani, T. Berg-Kirkpatrick, E. Hovy.
Proceedings of the AAAI Conference on Artificial Intelligence, Vol 32(1). 2018.
-
Word embedding visualization via dictionary learning
J. Zhang, Y. Chen, B. Cheung, B.A. Olshausen.
arXiv preprint arXiv:1910.03833. 2019.
-
Word2Sense: sparse interpretable word embeddings
A. Panigrahi, H.V. Simhadri, C. Bhattacharyya.
Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 5692–5705. 2019.
-
Transformer visualization via dictionary learning: contextualized embedding as a linear superposition of transformer factors
Z. Yun, Y. Chen, B.A. Olshausen, Y. LeCun.
arXiv preprint arXiv:2103.15949. 2021.
[6, 2, 7, 8, 9, 10]
), and a more general allied literature on disentanglement. We also note interim reports
-
[Interim research report] Taking features out of superposition with sparse autoencoders [link]
L. Sharkey, D. Braun, B. Millidge.
2022.
-
[Replication] Conjecture’s Sparse Coding in Toy Models [link]
H. Cunningham, L. Smith. 2023.
-
[Research Update] Sparse Autoencoder features are bimodal [link]
R. Huben. 2023.
-
(tentatively) Found 600+ Monosemantic Features in a Small LM Using Sparse Autoencoders [link]
L. Smith. 2023.
-
Really Strong Features Found in Residual Stream [link]
L. Smith. 2023.
-
AutoInterpretation Finds Sparse Coding Beats Alternatives [link]
H. Cunningham. 2023.
[11, 12, 13, 14, 15, 16]
which independently investigated the sparse autoencoder approach in response to Toy Models, culminating in the recent manuscript of Cunningham et al.
(Excerpt: opening motivation and the polysemanticity/superposition problem statement; full text at source URL.)