Superposition, Memorization, and Double Descent

interpretabilitysuperpositionmemorizationoverfittingdouble descenttoy models

Henighan et al. train the toy superposition model on finite datasets and find that with little data it stores individual training examples, rather than features, in superposition, that with enough data it switches to storing generalizing features, and that test loss rises and falls again (double descent) at the transition; invited comments replicate the result and extend it.

Tom Henighan, Shan Carter, Tristan Hume, Nelson Elhage et al. (Anthropic) · 2023-01-05 · original

A model overfits when it fits its training data in ways that do not carry over to new data, and in the extreme it memorizes individual examples. Overfitting is a central problem in machine learning, but in early 2023 there was little mechanistic account of what a network is doing internally when it memorizes. This short paper, published on the Transformer Circuits Thread, takes the toy model from Toy Models of Superposition, which had been studied only with unlimited data, and trains it on datasets of fixed size. Its motivating idea is that memorization suits superposition well: a lookup table that dedicated one neuron to each memorized sequence would be wasteful, but the cases are mutually exclusive, so they could share neurons with little interference. The authors call the investigation "very preliminary".

Terms used in this note

  • Superposition. Representing more features than a layer has dimensions, by giving each feature a direction that overlaps slightly with others (see Toy Models of Superposition).
  • Feature sparsity. The probability that a feature is zero on a given input, written S. Here S = 0.999, so each feature is active on about one input in a thousand.
  • Dimensionality. A measure from the earlier paper of how much of a hidden dimension a feature occupies. This paper defines the same quantity for a training example: its hidden vector's squared length divided by the sum of squared projections of all training examples' hidden vectors onto its direction ("How Do Models Change with Dataset Size?"). If T examples are arranged evenly around a circle in two dimensions, each has dimensionality 2/T.
  • Double descent. The observation, from earlier work by other researchers, that test loss can fall, rise and fall again as model size, data size or training time increases, so that more data can temporarily make a model worse. Data double descent is the version in which dataset size varies.
  • Weight decay. A training penalty that pulls weights toward zero, which limits how large they can grow.

What the paper does

The inputs are synthetic vectors of n = 10,000 features. Each feature is zero with probability 0.999 and otherwise uniform between 0 and 1, and each input vector is rescaled to unit length, which the authors say makes training examples easier to memorize ("Experiment Setup"). The model is the earlier paper's ReLU output model: a weight matrix W compresses the input into m hidden dimensions, and the transpose of W plus a bias, followed by a ReLU, reconstructs it. It is trained to minimize squared reconstruction error, with all features equally important, on training sets of T examples. Training uses 50,000 full-batch steps of the AdamW optimizer, with the learning rate warmed up over 2,500 steps to 0.001 and then decayed to zero along a cosine curve, and a weight decay of 0.01. The authors say the number of steps, full batches and the learning-rate schedule all mattered for the results, and that for some experiments they trained several random initializations and kept the one with the lowest training loss to avoid poor local solutions. Optimization was much less sensitive with more than three hidden dimensions.

For pictures, most experiments use m = 2 hidden dimensions, so that each feature's weight vector and each training example's hidden vector can be drawn as an arrow in the plane.

What it finds

Small datasets: training examples in superposition

With unlimited data, the earlier paper found sparse features arranged as regular polygons, for example a pentagon in two dimensions. With small training sets, the feature vectors look like a disordered scatter, but the hidden vectors of the training examples form a clean regular polygon with one vertex per example ("'Datapoint Features' vs 'Generalizing Features'"). The model has in effect treated each training example as a feature. The authors propose what they call a "naive mechanistic theory" of overfitting: memorization happens when a model uses data point features, each standing for one training example, instead of generalizing features. They say they expect this theory to be too simple. The setup was chosen to produce the clearest version of the effect: it is strongest when training examples are orthogonal and of similar size (footnote 3).

Large datasets: a pentagon of real features

As T grows, the training examples' dimensionalities follow 2/T in the small-data regime, as expected for a T-sided polygon ("How Do Models Change with Dataset Size?"). In the large-data regime, five features take on large dimensionality and are arranged as a pentagon, while the rest of the 10,000 features are largely ignored. The pentagon features' dimensionality is noticeably below the 2/5 a pentagon would give, which the authors attribute to the small contributions of the 9,995 other features adding up. Between the two regimes, the models are messier and harder to interpret.

The size of the weights follows the same pattern. The lengths of the hidden and feature vectors, and the norms of W and the bias, peak at around T = 30, fall rapidly in the middle regime, and level off once the model generalizes. The authors note that a model memorizing T orthonormal examples as a polygon needs larger weights for larger T, which matches the rise below T = 30, and that their simplified solution for the infinite-data case gives a weight-matrix norm of the square root of 5, about 2.2, close to the observed value.

Double descent at the transition

Averaging the test loss over four models trained on different random datasets at each size shows a clear bump at the transition between storing data points and storing features ("Superposition Double Descent"). The authors note that this double descent occurs without label noise, since the targets equal the inputs. The source of error is instead the compression itself: 10,000 features cannot be encoded in 2 dimensions without loss.

To check that the effect is not an artifact of two dimensions or of optimization trouble, they vary both the dataset size T and the number of hidden dimensions m, again averaging four models per setting, and find regions where either more data or a bigger model raises test loss ("The Effect of m on Double Descent"). In an appendix with m = 4, training for more epochs makes the bump larger, and raising weight decay from 0.01 to 1.0 removes it entirely, in line with earlier double-descent studies.

Open questions the authors list

The conclusion asks whether the naive theory applies to real models at all, and, within toy models, what happens when a model learns a mixture of features and data point features, what models in the middle regime are doing, and whether one notion of "feature" can cover both generalizing features and single data points.

Comments, replications and extensions

The article includes comments from researchers who saw the results early. They are separate from the authors' claims and are summarized here under their authors' names.

  • Replication by Adam Jermyn. The three regimes and the shapes of the embeddings reproduced for m = 2, with quantitative differences: the generalizing regime began at T = 20,000 rather than the paper's T = 5,000, the intermediate regime began at about T = 100 rather than about 700, and middle-regime test losses reached up to 5,000 rather than about 1.01. Peak bias norms were about three times larger. The authors reply that some hyperparameter must differ but that they care about the existence of the transition, not where it falls.
  • Replication by Marius Hobbhahn. With sparsity and the number of features each reduced tenfold, double descent and similar arrangements still appeared, and also appeared with other loss functions and with a ReLU between the layers.
  • What sets the generalization scale (Jermyn). Changing weight decay did not move the dataset size at which generalizing features appear, which argues against the idea that the switch happens when weight decay makes memorization too costly. Measuring dataset size by how often each feature is expected to appear, models with very different sparsities switched once each feature had been seen roughly 10 times. Jermyn notes that the curves line up less cleanly with more hidden dimensions.
  • Repeated data points (Jermyn). A data point repeated two or three times behaves like any other, but with more repeats the model learns a mixture: each repeated point takes the place of one of the five generalizing features, and with more than five repeated points all five features are displaced. A rough argument suggests a point is memorized once its number of repeats exceeds T/N, where N is the number of features, up to a constant, independent of sparsity, which the experiments matched at larger T/N. Some models memorized a point and then abandoned it during training.
  • Data dimensionality in MNIST (Chris Olah). In a one-hidden-layer network with 512 hidden units trained on handwritten digits, a variant called maximal data dimensionality picks out a few training examples with values almost 100 times higher than typical, which tend to be unusual images. Olah suggests, while calling it far from conclusive, that these may be memorized special cases, and connects this to detecting when a model is using an unusual mechanism. Hobbhahn reports replicating this.
  • Optimization failures in two dimensions (Olah and Henighan). Some m = 2 models in the memorization regime arranged examples on two circles of different radii and failed to merge them.

Limitations

  • Toy model only. The models compress synthetic sparse vectors; whether the data point picture applies to real networks is the first open question the authors list.
  • Sensitive training. The results needed long full-batch training, an annealed learning rate and, in some cases, selecting the best of several initializations, and the two-dimensional models were prone to optimization failures (setup section and the "Optimization Failures in 2D" comment).
  • Hyperparameter-dependent locations. The replication found the regime boundaries at quite different dataset sizes, and the authors could not identify which difference caused it.
  • Preliminary. The authors describe the work as a very preliminary investigation and the mechanistic theory as naive.

Research questions

  1. Do real language models store memorized sequences as data point features? The paper motivates the work with verbatim memorization but tests only toy models. Applying a data-dimensionality measure to a language model's activations on known memorized and non-memorized text would test it.
  2. Is high maximal data dimensionality a reliable marker of memorized examples? The MNIST comment shows outliers without confirming they are memorized. Comparing the measure with an independent test, such as the change in loss when the example is removed from training, would answer it.
  3. What do models in the middle regime compute? The authors list this as open. Tracking feature and example dimensionalities through training at dataset sizes near the loss bump would show whether models alternate between strategies or mix them.
  4. Does the "each feature seen about 10 times" rule hold with untied weights and more hidden dimensions? Jermyn's comment found it cleanest for small m. Repeating the sweep with untied encoder and decoder weights, which the authors note are more realistic, would test its generality.
  5. How do repeated data interact with generalizing features in larger models? The repeated-points comment shows memorized points displacing features in two dimensions. A similar experiment with controlled repetition in a small language model, measured by sparse-autoencoder features, would show whether the same displacement occurs.

How it connects

  • Toy Models of Superposition (September 2022) is the paper whose model and dimensionality measure this one reuses; its comments section also includes Hobbhahn's report of replicating most of this paper.
  • Towards Monosemanticity (October 2023) observes features that respond to very specific, often-seen text, such as standard license headers, and calls this an example of the mechanistic theory of memorization described here; its related-work section summarizes this paper.
  • Mapping the Mind of a Large Language Model (2024) cites the paper, in its related work on the theory of superposition, as extending toy models of superposition to memorization.
  • Circuits Updates (May 2023) follows up on this article's "Optimization Failures in 2D" comment with a section on the "two circle" phenomenon in memorization.
  • For overfitting and why test loss is measured on data held out from training, with a worked example of the gap between training and test performance, see Evaluating Classifiers.

References

  • Henighan, T., Carter, S., Hume, T., Elhage, N., Lasenby, R., Fort, S., Schiefer, N., and Olah, C. (2023). Superposition, Memorization, and Double Descent. Transformer Circuits Thread, 5 January 2023. https://transformer-circuits.pub/2023/toy-double-descent/index.html. Section names in this note refer to the web version, which includes the comments summarized above.
  • Anthropic. (2023, January 5). Superposition, Memorization, and Double Descent [research page]. https://www.anthropic.com/research/superposition-memorization-and-double-descent