Deep models reveal better strategies for superposition

1. Introduction The current, incredible performance of AI models is closely related to their compression capabilities (Language Modeling Is Compression (Delétang et al., 2023); Compression Represents Intelligence Linearly (Huang et al., 2024)). This compression is imposed on them by the architectural choices made by engineers. For example, GPT-2 had a vocabulary of 50,257 tokens, yet its “operational space” was only of size 768. In such a space, only 768 directions can be described fully independently, so the model had to find strategies to efficiently store all 50,257 input tokens in this reduced space. In general, we call the phenomenon in which a model stores more features than it has dimensions superposition. However, superposition comes at a cost. At a given step of computation, models utilizing it have neurons that fire for multiple different inputs (so-called polysemantic neurons), which makes them hard to interpret. One might hope to sidestep this by training models wide enough that every feature gets its own dimension (Engineering Monosemanticity in Toy Models (Jermyn et al., 2022)). Besides the cost, it is unclear that this would yield interpretable models, since whether natural features are cleanly separable at all is itself debated (The ‘strong’ feature hypothesis could be wrong (Smith, 2024)). We take a different route and try to understand the phenomenon itself. Prior work, Toy Models of Superposition (Elhage et al., 2022), studied how superposition forms in two-layer ReLU networks with tied weights. The task was to reconstruct the inputs given to the model (the input features) at its output, after passing them through a narrow bottleneck inside the model. The importance of each feature and the rate at which it occurs (its sparsity) could be varied across experiments, and the quality of the reconstruction was measured with the mean squared error (MSE). The authors suggest that for independent input features of equal importance and sparsity, the best geometries for storing the features may be uniform polytopes. In this work we study an analogous problem, but with a different model class: a linear encoding layer (the encoder) followed by one or more multilayer perceptrons (MLPs) with a bilinear activation. We treat the stack of MLPs as a single object, the decoder, whose task is to reconstruct the network's input from the geometry produced by the encoding layer. With this setup we show that the uniform geometries found in Toy Models of Superposition are not the best possible solution to the feature-reconstruction task: non-symmetric geometries yield lower error. We argue that their result is an artifact of the model class they studied, and that deeper, more expressive models do arrive at these asymmetric solutions. To explain why only sufficiently powerful models find them, we derive two theoretical benchmark decoders: (i) an analytical solution in a simplified case, obtained from the known data-generating process, and (ii) its generalization, approximated from data samples alone. Comparing different models against these benchmarks lets us understand why shallow models default to symmetric solutions while deeper models find the non-symmetric ones. We hope these results bring us closer to understanding the computations performed by LLMs, since our models and transformer architectures both use a linear encoder as the first step of their computation. 2. Setup We will be working with the following toy setup: The input to the models consists of 4 independent features of equal importance. Each feature is active with a small probability mjx-container[jax="CHTML"] { line-height: 0; } mjx-container [space="1"] { margin-left: .111em; } mjx-container [space="2"] { margin-left: .167em; } mjx-container [space="3"] { margin-left: .222em; } mjx-container [space="4"] { margin-left: .278em; } mjx-container [space="5"] { margin-left: .333em; } mjx-container [rspace="1"] { margin-right: .111em; } mjx-container [rspace="2"] { margin-right: .167em; } mjx-container [rspace="3"] { margin-right: .222em; } mjx-container [rspace="4"] { margin-right: .278em; } mjx-container [rspace="5"] { margin-right: .333em; } mjx-container [size="s"] { font-size: 70.7%; } mjx-container [size="ss"] { font-size: 50%; } mjx-container [size="Tn"] { font-size: 60%; } mjx-container [size="sm"] { font-size: 85%; } mjx-container [size="lg"] { font-size: 120%; } mjx-container [size="Lg"] { font-size: 144%; } mjx-container [size="LG"] { font-size: 173%; } mjx-container [size="hg"] { font-size: 207%; } mjx-container [size="HG"] { font-size: 249%; } mjx-container [width="full"] { width: 100%; } mjx-box { display: inline-block; } mjx-block { display: block; } mjx-itable { display: inline-table; } mjx-row { display: table-row; } mjx-row > * { display: table-cell; } mjx-mtext { display: inline-block; text-align: left; } mjx-mstyle { display: inline-block; }