Privileged Bases in the Transformer Residual Stream

interpretabilityresidual streamprivileged basisoutlier featuresadamtransformers

Elhage, Lasenby and Olah ask why individual coordinates of a transformer's residual stream carry unusually large values when the architecture gives them no special status; using a kurtosis measure and models trained with random rotations, they rule out layer normalization and floating-point precision and provisionally blame the per-coordinate normalization in the Adam optimizer.

Nelson Elhage, Robert Lasenby, Chris Olah (Anthropic) · 2023-03-16 · original

The residual stream of a transformer is the vector at each token position that every layer reads from and adds its output to (see Assembling the Transformer Block). Anthropic's mathematical account of transformers held that this stream has no privileged basis: every read and write passes through a full linear map, so the stream could be rotated by any invertible transformation, with the compensating transformation folded into every weight matrix, and the network would compute exactly the same function in different coordinates. On that account no single coordinate should mean anything special. Yet Tim Dettmers had reported "emergent outliers" in large transformers: particular residual-stream coordinates taking values up to 20 times larger than any other. This short paper, published on the Transformer Circuits Thread, asks not what those outliers represent but why they line up with the coordinate axes at all. Its abstract states the answer as provisional: "the per-dimension normalizers in the Adam optimizer are to blame".

Terms used in this note

  • Basis and privileged basis. A basis is the set of coordinate directions in which a vector is written down. A layer has a privileged basis when something in the computation treats those particular directions differently from others, most commonly a nonlinear function applied to each coordinate separately. Without one, a rotated copy of the representation is equally valid.
  • Outlier. Following Dettmers' first definition, a single residual-stream value whose absolute value exceeds 6 ("Measuring Outliers").
  • Kurtosis. A measure of how heavy a distribution's tails are. A Gaussian distribution has kurtosis 3; larger values mean more extreme values than a Gaussian would produce.
  • LayerNorm and RMSNorm. Layer normalization subtracts the mean of a vector, divides by its spread and applies a learned scale to each coordinate. The variant used here, close to RMSNorm, only divides by the root-mean-square of the vector and multiplies by a single learned number, which gives the same result in any rotated basis ("LayerNorm").
  • Adam. An optimizer that keeps running averages of each parameter's gradient and squared gradient and scales each parameter's update by them separately (see The Bigram Language Model for its use in a small model). Because the scaling is done parameter by parameter, it treats the coordinates in which the weights are stored differently from other directions.

What the paper does

A model and a measure

The experiments use a 200-million-parameter transformer trained with Anthropic's codebase, with a residual stream of 1,280 dimensions ("The experiments"). Dettmers had suggested that outliers appear inconsistently below larger sizes, but the authors find them frequent enough at this size. Counting, over a test batch of 128 sequences of 1,024 tokens, how many residual-stream dimensions ever exceed the outlier threshold, they find the count growing through the model, with a typical layer having 20 to 60 outlier dimensions.

Because a threshold is arbitrary, they propose a general test ("Activation Kurtosis"). If features were stored in random directions, the coordinates of each feature's direction would look like independent samples from a Gaussian, and so would the coordinates of the activation vector, which is a sum of such directions. Treating one token's activation vector as a sample and computing its kurtosis should then give about 3. Values above 3 indicate that some coordinates are special. The authors average this over many tokens. A footnote adds an asymmetry: kurtosis above 3 implies a privileged basis, but kurtosis of 3 does not prove there is none, since features with Gaussian activations would also give 3 (footnote 1).

Three candidate causes

The paper tests three sources of basis dependence ("Hypotheses for the basis-alignment"):

  • LayerNorm, the one standard operation in the forward pass that depends on the basis, through its mean subtraction and per-coordinate learned weights. The test replaces it with the rotation-invariant RMSNorm variant.
  • Finite-precision arithmetic, a hypothesis attributed to Neel Nanda and Alex Silverstein: floating-point numbers lose precision when values of very different sizes are added, so the model might benefit from keeping features of different magnitudes in separate coordinates. The test rotates the trained model's weights by a random orthonormal matrix and measures the change in loss; it then trains models in which the residual stream is rotated relative to the layers.
  • The optimizer, which remains after the other two are ruled out.

The rotated-training experiment has two versions ("Optimizing in an arbitrary basis"). In the shared rotation model, one fixed random rotation is applied whenever a layer reads from the residual stream (and its inverse whenever a layer writes to it), so all layers compute in one shared basis that differs from the residual stream's. In the independent rotations model, every read and write uses a different random rotation, so each layer computes in its own unrelated basis. Because each rotation sits next to a full linear map, neither change alters what functions the model can express.

What it finds

The residual stream is far from rotation-invariant in practice

In the baseline model, mean kurtosis is well above 3 in every layer ("Activation Kurtosis"). Applying a fixed random rotation to the same activations brings the measure to almost exactly 3, as the argument predicts, which confirms that the measure detects basis alignment rather than something about the values themselves.

LayerNorm is not the cause

The RMSNorm model trained as well as the LayerNorm baseline, and its activations were at least as heavy-tailed, if anything more so ("LayerNorm"). The authors conclude that the small basis dependence of standard LayerNorm does not cause the outliers.

Floating-point precision is not the cause

With RMSNorm, the model should be fully rotation-invariant, so rotating every weight matrix should leave it unchanged up to rounding error. Random rotations did tend to raise the loss on average, but by tiny amounts, against a baseline test loss of about 3.02 nats and much larger variation between random initializations ("Verifying the model is basis-independent"). The authors conclude that the forward pass is rotation-invariant even with floating-point error.

Both rotated-training models trained essentially as well as the baseline; their loss curves differ slightly early in training and converge, within noise, after about a billion tokens ("Optimizing in an arbitrary basis"; "Loss curves"). The authors conclude that transformers do not need a privileged basis to train and work well, and count this as moderate evidence against the precision hypothesis. In both rotated models, the residual stream's kurtosis was almost exactly 3. The stronger evidence came from looking inside the layers: in the shared rotation model, the activations entering each MLP layer, in the shared computation basis, were about as heavy-tailed as the baseline's. The basis alignment therefore follows the basis in which computation happens, not the basis in which information is passed between layers, which the authors call strong evidence against the precision hypothesis. In the independent rotations model, tails were small, mostly in early layers, which suggests that the effect arises when all the layers share one basis and reinforce it together.

The optimizer remains

Adam normalizes each parameter's update by that parameter's own gradient statistics, so it favors the basis in which the weights are stored. Having ruled out the other candidates, the authors call it the strongest basis-dependent operation they know of in the model and conclude that the experiments "push in the direction of suspecting the optimizer" ("Conclusion").

Limitations and open questions

  • Not a direct test of Adam. The authors write that "we cannot claim to have conclusively put the blame on Adam" and that another, unidentified mechanism could be involved ("Conclusion"). They describe two follow-up experiments they chose not to run: training a transformer with plain SGD or another optimizer that has no basis dependence, possibly with a very low learning rate and many steps, and training with Adam while storing its moment estimates in a different basis from the weights.
  • One small model. The experiments use a 200-million-parameter model; Dettmers' outliers were reported mainly in larger models.
  • What the outliers represent is left open. The paper addresses why outliers are basis-aligned, not what they encode or what purpose they serve.
  • The kurtosis test is one-sided. As footnote 1 notes, a value of 3 does not by itself show that a representation has no privileged basis.

Research questions

  1. Do the outliers disappear when the optimizer has no basis dependence? The authors' own proposed experiment: train matched models with Adam and with SGD (or Adam with moments stored in a rotated basis) and compare kurtosis by layer.
  2. Does the effect grow with model size? The study uses one 200-million-parameter model. Measuring kurtosis across a range of sizes trained with the same recipe would show whether basis alignment strengthens at the scales where Dettmers reported outliers.
  3. What do the basis-aligned outlier dimensions encode? The paper sets this question aside. Correlating outlier dimensions with token types or positions, or ablating them and measuring which predictions change, would begin to answer it.
  4. Does residual-stream basis alignment affect interpretability methods that assume none? Later sparse-autoencoder work treats the residual stream as roughly basis-free. Comparing features learned on raw and randomly rotated residual activations would show whether alignment from the optimizer changes what such methods find.
  5. Why do layers that share a basis reinforce alignment? The independent rotations result suggests a collective effect. Training models in which only some layers share a basis would show how the effect builds up.

How it connects

References

  • Elhage, N., Lasenby, R., and Olah, C. (2023). Privileged Bases in the Transformer Residual Stream. Transformer Circuits Thread, 16 March 2023. https://transformer-circuits.pub/2023/privileged-basis/index.html. Section names in this note refer to the web version.
  • Anthropic. (2023, March 16). Privileged Bases in the Transformer Residual Stream [research page]. https://www.anthropic.com/research/privileged-bases-in-the-transformer-residual-stream
  • Dettmers, T., Lewis, M., Belkada, Y., and Zettlemoyer, L. (2022). LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale. arXiv:2208.07339. (The source of the "emergent outliers" the paper investigates, as cited and described by the paper.)