Source-linked AI summary
Birth of a Transformer: A Memory Viewpoint
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, Leon Bottou
TL;DR
The paper asks how transformers balance globally stored knowledge with information learned from context, a question relevant to understanding and improving deployed language models. It studies this balance with synthetic bigram sequences and a simplified two-layer transformer, finding that global bigrams are learned before a slower induction-head mechanism for in-context bigrams.
Problem
The paper addresses the need to understand how transformers store global knowledge and adapt to information in context, supporting interpretability and mitigation of model failures.
Method
The paper combines a synthetic dataset separating global and sequence-specific bigrams with a simplified two-layer transformer, associative-memory analysis, empirical training probes, and population-gradient theory.
Results
Global bigrams are learned first, followed by top-down formation of an induction head that recovers in-context bigram behavior through learned associative memories.
Takeaways & Limitations
Weight matrices can be analyzed as associative memories, providing an interpretable account of how attention mechanisms learn to use context in this simplified setting.
Takeaways & Limitations
The simplified setup does not fully represent complex language-model training, which may require learned structured embeddings, factorized matrices, and nonlinear feed-forward layers.
Abstract
from arXiv · showhide
Large language models based on transformers have achieved great empirical successes. However, as they are deployed more widely, there is a growing need to better understand their internal mechanisms in order to make them more reliable. These models appear to store vast amounts of knowledge from their training data, and to adapt quickly to new information provided in their context or prompt. We study how transformers balance these two types of knowledge by considering a synthetic setup where tokens are generated from either global or context-specific bigram distributions. By a careful empirical analysis of the training process on a simplified two-layer transformer, we illustrate the fast learning of global bigrams and the slower development of an "induction head" mechanism for the in-context bigrams. We highlight the role of weight matrices as associative memories, provide theoretical insights on how gradients enable their learning during training, and study the role of data-distributional properties.
1 Introduction
The paper studies how transformers balance globally stored knowledge with information learned from context, using a synthetic setup to make their training dynamics interpretable. It finds that global bigrams are learned before an induction-head mechanism for sequence-specific bigrams.
- The paper targets mechanistic understanding of transformers to improve interpretability and help mitigate failures such as hallucinations and reasoning errors.
- The synthetic task separates global bigrams from sequence-specific bigrams, requiring global statistics for the former and in-context learning for the latter.
- Two-layer transformers develop an induction-head circuit that predicts b from contexts containing a repeated a followed earlier by b, whereas one-layer transformers fail reliably on these in-context bigrams.
- The analysis freezes selected embeddings and value matrices to focus on how attention and feed-forward mechanisms emerge during training.
- The paper models weight matrices as associative memories, empirically studies training order, and uses population-gradient analysis to explain recovery of desired memories.
2 Background
The background introduces transformer components and the induction-head circuit that supports basic in-context copying. The circuit uses two attention heads: one copies a previous token, and the other retrieves the associated following token.
- Transformer architecture: Transformers alternate causal self-attention and token-wise feed-forward layers in decoder-only autoregressive architectures.
- Transformer architecture: Token and positional embeddings form the inputs, while residual streams accumulate embeddings and layer outputs across the network.
- Transformer architecture: Attention scores queries against keys and combines value transformations, with residual connections carrying the resulting representations forward.
- Induction head mechanism: An induction head predicts b from a context of the form [..., a, b, ..., a] through a two-head mechanism across separate layers.
- Induction head mechanism: The first head attends to the previous token and copies its embedding, enabling the second head to attend by that remapped representation and output the original token.
- Induction head mechanism: Attention maps show previous-token attention in the first layer and trigger-following attention in the second layer for fixed or repeated-token triggers.
3 Synthetic Setup
The synthetic setup generates sequences from global bigram statistics while modifying a few trigger transitions on a per-sequence basis. This lets the experiments compare global associations with in-context prediction and observe the emergence of induction heads.
- Bigram data model: The dataset uses a bigram language model whose transitions for selected trigger tokens are replaced by sequence-specific output tokens.
- Bigram data model: Global bigram conditionals and unigram statistics provide information shared across sequences, while trigger transitions vary by sequence.
- Experimental setup: Triggers may be sampled independently without replacement from a distribution or fixed as the most frequent vocabulary tokens.
- Experimental setup: Experiments use character-level unigram and bigram distributions estimated from tiny Shakespeare with vocabulary size N = 65.
- Experimental setup: Two-layer transformers achieve over 99% accuracy for fixed triggers and 95% for random triggers, compared with around 55% for one layer.
4 The Associative Memory Viewpoint
The paper models transformer weight matrices as associative memories over nearly orthogonal embeddings and uses this perspective to construct an interpretable two-layer architecture for global and in-context bigrams.
- Weight matrices as associative memories: Nearly orthogonal embeddings let weight matrices store input-output associations as weighted sums of embedding outer products.Random high-dimensional embeddings provide a practical way to approximate the required near-orthogonality.
- Weight matrices as associative memories: Gradient descent learns associative memories whose coefficients reflect the trained input-output pairs when embeddings are fixed.The resulting matrix has an initialization term plus weighted outer products, with coefficients varying over training.
- A simplified two-layer transformer architecture: The simplified model freezes embeddings and selected value matrices, focusing analysis on attention and feed-forward mechanisms while retaining nonlinear parameter dynamics.The model trains the output matrix and feed-forward layer while using tied key-query matrices and frozen remapping matrices.
- A simplified two-layer transformer architecture: The induction head matches tokens through positional and remapped-token associations, while the feed-forward layer stores global bigram statistics.Residual connections combine the induction-head and global-bigram outputs, requiring the model to trade off their predictions.
- A simplified two-layer transformer architecture: The model often prefers induction-head outputs unless they agree with top global-bigram predictions, which can then receive preference.This describes the observed interaction between the two prediction mechanisms under residual addition.
- Scope of the architecture: The simplified architecture omits nonlinear MLPs, trained embeddings, layer normalization, joint factorization learning, and the depth and multiplicity of standard transformers.The authors extend the memory viewpoint to richer components but focus on the simpler model for interpretability.
5 Empirical Study
The empirical study probes how the simplified transformer learns global and in-context bigrams. Global statistics emerge early, whereas the induction mechanism develops more slowly through staged associative-memory learning shaped by the data distribution.
- Emergence of the induction head: The output matrix is learned before the key-query matrices and is necessary but insufficient for good in-context accuracy.Output associations can form even when attention remains nearly uniform, but key-query learning is required to focus attention on relevant tokens.
- Emergence of the induction head: Associations learned by the output matrix enable later key-query learning, which recovers induction-head behavior and near-perfect accuracy.The study characterizes this as top-down learning of the second layer before the first.
- Emergence of the induction head: Previous-token associations are learned faster for early sequence positions than for later positions.The authors attribute this to gradients concentrating on the first occurrence of each trigger, which is typically early in the sequence.
- Global vs in-context learning: Global bigram statistics are learned more quickly than the induction head during joint training.The global loss and KL decrease rapidly early in training, while full in-context learning takes longer because key-query matrices develop more slowly.
- Role of the data distribution: Fewer triggers, rare fixed triggers, and random rather than fixed triggers can slow in-context learning through different memory-learning bottlenecks.The first two conditions are associated with slower output-matrix learning, while random triggers increase the key-query memory requirements.
- Role of the data distribution: Training with uniform output tokens generalizes better across output distributions than training with global bigram distributions.Training on bigram-distributed outputs reduces accuracy on out-of-distribution outputs, whereas uniform training performs well on both distributions.
6 Theoretical Insights on Learning Dynamics
The paper explains how gradients can form associative memories that filter noise and, through sequential updates, learn an induction-head mechanism. In an idealized setting, one gradient step can ignore positional noise, while later steps progressively establish the attention associations.
- Gradient associative memories: Gradients can learn associative memories that filter irrelevant components from superposed inputs by using conditional expectations tied to predictive labels.The resulting memory depends on μ_k and can isolate input components relevant to label k.
- Gradient associative memories: A single population gradient step yields a near-perfect classifier that ignores positional embeddings when the embedding dimension and vocabulary and sequence sizes are sufficiently large.This result is explicitly described as highly idealized and leaves finite-dimensional, finite-sample analysis for future work.
- Learning the induction head: The induction head can be learned through a few sequential gradient steps in a simplified setup, beginning from random initialization.The analysis uses a top-down order in which later-layer associations are learned before earlier attention associations.
- Learning the induction head: After the output associations are learned, gradients can form key-query memories that focus attention on the correct trigger and then learn associations for the previous-token head.The argument relies on linearizing the second-layer attention around zero initialization.
7 Discussion
The paper’s simplified bigram setup enables fine-grained analysis of transformer training dynamics, but richer language-modeling settings require more elaborate models.
- Scope and future work: The simplified setup captures rich phenomena in the bigram task while leaving more complex transformer behavior for future study.The discussion identifies structured embeddings, factorized attention matrices, and nonlinear feed-forward layers as important extensions.
A Associative Memories with Random Vectors
The appendix develops the random-vector assumptions and illustrative gradient calculations underlying the associative-memory analysis. It also explains how the simplified training setup isolates the emergence of induction-head components.
- Random-vector properties: Random Gaussian embeddings support associative-memory analysis through concentration, near-orthogonality, and approximately norm-preserving random remappings.These properties are used to treat outer-product weight matrices as storing embedding associations.
- Random-vector properties: The analysis assumes the embedding dimension is large enough for the required memory lookups to work across many embeddings and pairs.A precise concentration and union-bound analysis is left for future work.
- Induction-head analysis: The theoretical analysis studies gradients at initialization and a top-down sequence of layer-wise updates for induction-head formation.The setup uses simplified trigger distributions and may omit the feed-forward layer because global bigram statistics are expected to be learned quickly.
- Illustrative gradient example: A one-step population gradient update can produce an essentially perfect classifier that ignores random positional embeddings.The illustrative calculation starts from zero initialization and assumes uniformly distributed labels and sufficiently well-separated embeddings.
B.3.1 Learning W 2
The analysis shows that the second output matrix can acquire useful token associations after one gradient step, although average attention initially leaves predictions near-random and introduces spurious associations.
- Learning W 2_O: One gradient step on the second output matrix recovers the desired associative memory in the simplified bag-of-words setup.The derivation starts from zero initialization, where predictions are uniform, and uses the sequence’s output-token statistics.
- Learning W 2_O: The learned matrix also stores spurious associations arising from residual-stream and first-layer attention components.These additional associations may be cleaned up as attention layers begin focusing on the correct tokens.
- Learning W 2_O: Despite recovering useful associations, the one-step estimate still produces near-random predictions because average attention cannot identify the output token’s position.The repeated trigger provides no information about the output token in this simplified model.
B.3.2 Learning W 2
The second-layer matrix W 2 is learned as an associative memory for induction, with gradients emphasizing the desired trigger-to-output associations while suppressing spurious terms. Once W 2 is learned, the model can form the previous-token mechanism in W 1.
- Learning W 2: The gradient’s main signal comes from positions where the trigger and relevant association align, while other positional components are negligible for large N.The desired component behaves precisely like the target associative memory, whereas the remaining terms are treated as spurious or negligible.
- Learning W 2: Up to spurious positional embeddings, W 2 behaves as the desired associative memory when N is sufficiently large.A sufficiently large step size can further increase the gap between correct and incorrect associations by making attention more focused.
- Learning W 2: W 2 is analyzed from zero initialization as learning an associative memory for trigger-to-output associations.The analysis assumes the simplified architecture and examines one gradient step under near-uniform attention.
- Learning W 1: After W 2 captures the correct associations, a gradient step on W 1 produces the desired previous-token associative memory.The construction uses positional information in the first layer’s key-query matrix and token embeddings in its value-output matrix.
- Learning W 2: Before W 2 is learned, second-layer attention activates at all occurrences of the trigger, often causing near-random predictions.The head tends to predict the trigger itself, which is incorrect most of the time.
C.1 Proof of Lemma 1
The proof derives gradient expressions for a softmax classifier and applies them to a matrix receiving keys and a query. These expressions show how a gradient step can encode input-output associations in the weight matrix.
- Gradient derivation: For cross-entropy classification, the derivative with respect to each logit is the predicted probability minus the target indicator.This identity is used to obtain the gradient of the loss with respect to the matrix parameters.
- Matrix-gradient setup: The matrix-gradient analysis separates keys Z from the query x_T, while noting that practical attention inputs may superpose token, positional, and remapped embeddings.This abstraction isolates how the matrix maps keys to outputs in the associative-memory analysis.
- Gradient derivation: Linearizing the second-layer softmax around zero yields the expression used to evaluate the initial gradient and establish the lemma.The proof substitutes the softmax linearization into the loss gradient and concludes with the stated result.
D Beyond our Simplified Architecture
The associative-memory viewpoint extends beyond the simplified transformer to factored matrices, normalization, MLPs, and multiple heads or layers. These extensions preserve useful associations but introduce capacity, optimization, and identifiability trade-offs.
- Factorizations: Products of lower-dimensional matrices can implement the same associations as a full matrix when embeddings are sufficiently high-dimensional and near-orthogonal.The factorization uses attention-head-sized intermediate dimensions and shows that the resulting product behaves like the target memory.
- Layer-normalization: Layer normalization projects updates away from the current Wx direction, reducing repeated updates once an association is already stored.At random initialization, this effect is weak because Wx is nearly orthogonal to output embeddings.
- MLP blocks: Nonlinear MLP blocks can filter near-orthogonality noise and may increase storage capacity, at the cost of additional computation.They can also encode many-to-many associations when multiple embeddings coexist in the residual stream.
- MLP blocks: The training dynamics of nonlinear associative memories are left for future work.The paper discusses how the nonlinearity may filter mismatched embedding sets but does not analyze how those memories arise during training.
- Multiple heads and layers: The memory interpretation applies naturally to multiple heads and layers, but redundancy makes it harder to identify which component learned an association.In the four-head example, one first-layer head carries the previous-token behavior while second-layer induction is distributed across heads.
E Experiment Details and Additional Experiments
Additional experiments probe learned memories, data-distribution effects, dimensionality, and architectural realism. They show that association strength and learning speed depend on trigger statistics, embedding dimension, and model redundancy, while the qualitative memory behavior persists in more complex models.
- Memory recall probes and data-distributional properties: Induction-head memories learn more slowly with few triggers, rare fixed triggers, or random triggers.The authors attribute the first two effects primarily to slower learning of W 2_O and the random-trigger effect to slower key-query learning.
- Memory recall probes and data-distributional properties: Uniform output tokens fit W 2_K better than bigram-sampled outputs, while W 1_K recall is comparable across the two settings.Bigram outputs may already be predicted by global statistics, reducing gradient signal through the induction mechanism.
- Effect of dimension: Higher dimension leads to faster learning of W 2_O, with d = 128 sufficient for fast learning after a few iterations when the learning rate is tuned.Without tuning, low dimension causes additional slowdown, likely because of conditioning issues.
- One-step estimate: A one-step estimate from average attention extracts relevant output-token signal after a handful of batches, and this becomes easier as dimension increases.The estimate corresponds to the classifier obtained after one gradient step from zero initialization.
- More complex architectures: Training a model with a ReLU MLP, learned embeddings, and pre-layer normalization produces memory probes similar to the simplified architecture.This suggests that the same memory associations remain identifiable despite the added components.
- More complex architectures: With four heads per layer, one first-layer head learns the previous-token behavior while induction is shared across second-layer heads.The distribution of mechanisms across redundant components creates an identifiability challenge.