Birth of a Transformer: A Memory Viewpoint
Authors: Alberto Bietti
Vivien Cabannes
Diane Bouchacourt
Herve Jegou
Leon Bottou
Venue: NeurIPS 2023 spotlight (2023)
Keywords: transformers
language models
deep learning theory
interpretability
TL;DR: understanding transformers and their learning dynamics through associative memories

Abstract: 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.

Birth of a Transformer: A Memory Viewpoint
Alberto Bietti"
Vivien Cabannes FAIR, Meta
Diane Bouchacourt FAIR, Meta
Hervé Jégou FAIR, Meta
Léon Bottou FAIR, Meta
Flatiron Institute
Abstract
Large language models based on transformers have achieved great empirical suc- cesses. However, they are deployed more widely, there is growing need 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 their context or prompt. We study how transformers balance these two types of knowledge synthetic setup where tokens are generated from either global or context-specific bigram distributions. By careful empirical analysis of the training process 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 ditt-distributional properties.
considering a
Introduction
As large language models (LLMs) are growing in usage and deployment, it is increasingly important to open the black box and understand how they work. better understanding can help with inter- pretability of how these models make decisions, and will be crucial to improve these models and mitigate their failure cases, such as hallucinations reasoning errors.
An important ingredient in the success of recent LLMs their ability to learn and reason from information present in their context (6. These "in-context" learning capabilities are often attributed to the transformer architecture [52], in particular its self-attention blocks, which are able to carefully select parts of the input sequence in order infer plausible next tokens. Additionally, predictions may require "global" knowledge, such as syntactic rules or general facts, which may not appear in the context and thus needs be stored in the model.
In order to better understand how transformers develop these capabilities during training, we introduce synthetic dataset that exhibits both aspects. It consists of sequences generated from bigram language model, but where some of the bigrams are specific to each sequence. Then, the model needs the sequence-specifie bigrams, while the global
rely in-context learning for good prediction bigrams can be guessed from global statistics conditioned on the current token. While one-layer transformers fail to reliably predict the in-context bigrams, find that two-layer transformers succeed by developing an induction head mechanism [16] [0]40]] namely a "circuit" of two attention heads that allows the transformer predict b from context [... .a.,,,, b.... a). and which appears to be ubiquitous in transformer language models 44 54).
In order to obtain a fine- line-grained understanding of how this in-context mechanism emerges during training, we further simplify the two-layer architecture by freezing some of the layers at random initialization, including embeddings and value matrices. This focuses our study attention and feed-forward mechanisms, while avoiding the difficulty learning representations, which may
Work done while at FAIR, Meta.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

This simplification also allows us to introduce a natural model for individual weight matrices as associative memories, which store input-output or key-value pairs of embeddings through their outer products. Random high-dimensional embeddings are particularly well-suited this viewpoint thanks to their near-orthogonality. We provide a detailed empirical study of the training dynamics, by measuring how quickly each weight matrix learns to behave as the desired associative memory, studying how this is affected by data-distributional properties, and investigate the order in which layers are learned: the model first finds the right output associations from the current token and from uniform attention patterns, then the attention heads learn to focus on the correct key-value pairs. We then present theoretical insights this top-down learning process through population gradient dynamics. Despite its simplicity, our setup already provides useful insights the internal structure of transformer language models and its evolution throughout training, paving the way for a better understanding of L.Mss. We hope that our insights may lead to future research and improvements for L.M practitioners, e.g., for optimization algorithms, data pre-processing and selection, interpretabilityy fine-tuning, and model editing.
require complex nonlinear dynamics [ 33 33 45
In summary, we make the following contributions:
We introduce a new synthetic setup to study global vs in-context learning: sequences follow bigram language models, where some bigrams change across sequences and others do not.
We view the transformer's 's weight matrices as associative memories that learn to store specific pairs of embeddings, and use this derive a simplified but more interpretable model for our task.
We empirically study the training dynamics with careful probing: global bigrams are learned first, then the induction head is formed by learning appropriate memories in a top-down fashion.
We give theoretical insights on training dynamics, showing how a few top-down gradient steps on the population loss can recover the desired associative memories by finding signal noisy inputs.
Related work. After the success of transformer language models for in-context learning was found 6). several works have studied how in-context learning may arise in various contexts [8 38 43 48, 57]. Multiple recent papers have introduced synthetic tasks in order to better understand and interpret transformers ( 33
39 61). Several works have attempted understand internal mechanisms in transformers that are responsible for certain behaviors, an area known as "mechanistic interpretability" [66
17
35
39
40
54
Memory and neural networks have long history of
56). The associative memories we consider bear 22 56 though we use continuous input/outputs. The reader may also be interested in Fast Weight programmers 46 47.. The use of random vectors for storing memories is related to 23] Our approach probing based on memory recall related to techniques in [7], though motivated differently. 14
connections [S 18
19
27
30
34
50
55
similarity to
32 36 study statistical and approximation properties of transformers, highlighting benefits of sparse attention patterns, but do not consider training dynamics. 225 49 [1]] provide theoretical analyses of learning dynamics in transformers and other attention models, but consider different data setups and focus on single-layer architectures, while we focus on two-layer models and take different viewpoint based on associative memories.
2 Background
This section provides background on transformer architectures and induction head mechanisms.
Transformer architecture. Transformers 52 operate on sequences of embeddings by alternating self-attention operations and token-wise feed-forward layers. We focus on decoder-only, auto- regressive architectures with a causal attention mask, which are commonly used in large language models trained for next-token prediction 66]
42). We ignore normalization layers in order to simplify the architecture, since its stability benefits are not as crucial in the small models we consider. Given an input sequence of tokens 21.7 E [N] of length T,, where N is the vocabulary size, the transformer operates as follows:
Token embeddings: each discrete token is mapped to d-dimensional embedding via an em- bedding map WE E d We will denote the embeddings of tokens 2, by x; := ug (2)).
where (e) is the j-th column of We.
Positional embeddings: the positional embeddings p. = wp(() E Rd are added to each token embedding depending on its position in the sequence, leading to the following input embeddings:
De == Z₁ + Pt wg(()) + wp(t).
(I)

---

Layer 2
(b) mg((a)
Athm2: Σn (A)
Σn (K)
Layer
(x)
(i(a)
Wx(()
(rg((a
Attnl:
Residual
Prediction
Layer
Pr-1
wry (a)
Pr
wy(()
wy(a)
Sequence
[0]
two-layer mechanism that can predict b from a context [.... a, b, b.....a. a]. The first layer is a previous token head, which attends the previous token based on positional embeddings (pr + Pr-1) and copies after a remapping (w + wi (a) == W 721h W w The second layer is the induction head, which attends based on the output of the previous token head (w) + w/ (a)) and outputs the attended token, remapped to output embeddings (wg(b) - we((6), Boxes in the diagram represent different embeddings in superposition on each token's residual stream (we omit some irrelevant ones for clarity, e.g., positional embeddings in upper layers), and attention and output associations are shown with the associative memory viewpoint presented in Section 4
Figure 1:
Induction head mechanism. Induction heads
given an input sequence I1.7 Rdx of embeddings, the causal attention block computes, for WK, Wo, Wv, Wo E Rdxd (key, query, value, output), and for each I,
Attention blocks:
1/2 WoW WQzz) E R
(2)
where takes the softmax of its elements, leading weights proportional to exp((Wxa) [Wooz)). Note that the attention operation usually con- siders multiple "heads" that each projects the input single head for simplicity, since it will be sufficient for our purposes. Rewriting (2) on each I as "'-r- A(Eirr WK, Wo, Wy, Wo), the l-th layer of the transformer applies attention with layer-specific parameters along with residual connection follows |
attention of the "values" Wyz, with
lower dimension. Here we stick
F1. .1. A(TI:: W K: W Wt, W
blocks operate on individual token embeddings after each attention block, typically by applying one-hidden-layer MLP to each token, denoted F(; F). with residual connection: layer l, have
Feed-forwart blocks:
# x4 + F(r;; Wr).
Our simplified setup will linear foed-forward layers: F(x;; Wr) = Wyzt.
Unembedding: After the last transformer layer, the embeddings are mapped back to the vocabulary space N through a linear "unembedding" layer Wy = (w),()),. ....., we (N)] E N where refer the wo (j) as "output embeddings". The output of this layer is then fed into cross- entropy loss for predicting of 2++1 from each It.
We will sometimes refer to the representations for a given token throughout layers as its residual stream 16], since they consist of sums of embeddings and layer outputs due to residual connections.
Induction head mechanism. Induction heads 16 40 are particular type of mechanism (or
"circuit") in transformers that allows basic in-context prediction of the form [... a, b.... .a + These were found to be ubiquitous in transformer language models, playing a key role in enabling various forms of cn-tontext learning. The basic mechanism consist of two attention heads in separate layers (see Figure T for an illustration): (i) the first is previous token head which attends to the previous token using positional information and copies second is the induction head itself, which attends using the output of the previous token head, and outputs the original token. Our work focuses on this basic copy mechanism, but we note that richer behaviors are possible, particularly when combining multiple such mechanisms (e.g., [54]
embedding to the next token; (ii) the
We omit layer indices for simplicity of notation, and use the assignment operator iistead.
3

---

thbnt CANTEL&CCCCDEE
Figure 2: Induction head behavior in attention maps observed 2-layer transformer trained on two variants of our synthetic dataset. Each row shows the attention pattern for predicting the next token. (left) The first layer head always attends the previous token. (center) For fixed triggers Q = {a, t). the second layer head mainly attends to tokens following such triggers. (right) For random triggers, the induction head mechanism is active for any repeated token (here the only trigger is L). Red and green boxes highlight tokens following previous occurrences of the query, with red boxes corresponding to "correct" output tokens og following trigger tokens qk.
3 Synthetic Setup
In this section, introduce our synthetic data setup, which allows to carefully study how the induction head mechanism develops during training, and how transformers learn to use information from the context vs simple associations from the training data.
Bigram data model. Our model for sequences consists of a generic bigram language model (i.e. Markov chain), but where the transitions for a few frigger tokens denoted qs are modified in each sequence to always be followed by some output tokens og. Let be the number of trigger tokens, and fix the following distributions over the vocabulary [N]: ma(-))), Tu. T-(-li) and T.. for i E [N]. (-(-l) are the global bigram conditionals, Ty the global unigram distribution, while To is used to sample output tokens at each sequence. The triggers either fixed to some predefined set of tokens Q, or sampled from T... Each sequence "i.r is generated as follows:
(optional) Sample QI...., 9K qg. i.d. without replacement (random triggers)
Sample T-(-19)), i.i.d. with replacement.
Sample zf and =/1=0" Pm (-1=/-1)) for L = 2...., T, where
P.. (b)) = (jli),
if $ {Tc}x if i = qk-
1 = oz).
Experimental setup and initial experiment. Our experiments take n. and ny to be unigram and bigram character-level distributions estimated from the tiny Shakespeare dataset, with vocabulary size N = 65. We generally sample triggers from Fe тu or fix them to the K most frequent tokens. We sample uniform outputs oş in most cases, but also experiment with w. n. in Section 5
As a preliminary experiment, we train a two-layer vanilla transformer with single-head attention layers and MLP feed-forward layers, following the training setup described in Section On our synthetic data, with fixed (resp. random) triggers and uniform outputs, the model achieves over 99% accuracy (resp. 95%) on output tokens after the first occurrence, versus around 55% for one layer. This gap may be related to the difficulty of modeling three-war way interactions with a single attention layer [44). We visualize attention maps test sequences in Figure 2] which shows that the model has learned an induction head mechanism. The sequence in the middle figure has (gk. 0₂) E {(a, b), (t, s)}. For fixed triggers, the induction head is only active for the triggers used in training, which suggests the presence of a "memory" in the attention layer. For random triggers, it active on every repeated token, that the model then needs to disambiguate between in-context and global predictions. For instance, the model may choose to use the retrieved token when it is unlikely to be sampled from the global bigram distribution, something which found often be the case in practice.
4 The Associative Memory Viewpoint
In this section, we present our associative memory view on transformers: with nearly orthogonal embeddings, the weight matrices behave as associative memories which store pairs of embeddings as

---

weighted sum of their outer products. We then introduce a simplified transformer model with fixed random embeddings that will yield a precise understanding of learning dynamics using this viewpoint.
4.1 Weight matrices as associative memories
While intermediate representations in the transformer consist of nigh-dimensiinal vectors in residual streams, they are often "collapsed" down to scalar measurements by testing against other representa- tions, using operations of the form 2/ Wu; for some matrix W. For instance, u, and by could be key and query vectors in an attention head, or input and output embeddings for predicting the next token. If (u,); and (2y) are orthonormal (or nearly-orthe orthonormalll sets of embeddings, a natural way to store desired input-output associations i, j through the following associative memory:
W Σ aujoyu
(3)
so that the scores i/ Ww 0ij may be used to assess the relevance of the (i,j) pair, e.g., as part of a softmax operation in attention or next token prediction.
Random embeddings. A simple way to ensure that embeddings (u:); and (2y) are nearly- orthonormal to set them to be random high-dimensional vectors, such as Gaussian vectors with variance d in d dimensions. Indeed, these are known to satisfy 53
1 (
u U4 R 1 and u O
so that (3) is reasonable way to define associative memory, without requiring an explicit activation function as employed end-to-end memory networks [50]. We may also easily create a "remapping" of an existing embedding u; by multiplying it by random matrix Wo E Rdddd with Gaussian entries variance 1/d, which is commonly used for initializing neural network parameters. The new remapped embedding Won, is near-unit norm, and is near-orthogonal to u, in addition to other uj. Note that this fact implies that attention scores at initialization are near-uniform. See Appendix A for more details.
Learning associative memories. We now show that learning associations of input-output embed- dings via gradient descent leads to weighted associative memory of a form similar to (
Lemma 1 (Gradients and associative memories). Let be data distribution over input-output tokens, and consider the following loss, where the input and output embeddings W g and Wu are fixed:
L(W)= 1l
with E the cross-entropy loss. The gradients of the population loss L then take the form
LwW(W) Σ E. (1=) - ply k[2)) (k)wg())))
(5)
where Pw (y==k<))= o (Wo (W)))))()) ((Wl) are the model's predicted probabilities. Running gradient descent (with or without weight decay) from initialization Wo then leads to estimates of the following form, for some ogg and asy that vary with the number of iterations:
W Wo Σ Caywe())wr(() 3.3
(6)
Note that (4) convex problem in W, thus with appropriate step-size and large enough number of steps (with no weight decay) we can expect gradient descent to be close to the global minimum. At the optimum, if the embeddings are nearly orthogonal, then (5) implies fw (y = k|2) p(y k|2). We remark that if Wú is Gaussian random matrix, as if often the case for neural network layers, the first term (6) plays minor role: testing Wo against an input-output pair (i.j) with æi * will
concentrate around zero when d is large, while the (i,j) j) term in the sum will concentrate around osj.
We also note that the gradient updates described above correspond to a so-called maximal feature learning regime similar to μP updates in intermediate layers of deep networks [58 59].
5

---

Handling superposition. In Lemma we assumed that inputs to the matrix W are embeddings of a single token. Yet, in transformer models, the inputs to weight matrices are often sums, or superpositions of embeddings. For instance, the initial representations of each token are sums of token and positional embeddings, and representations at later layers are sums of the outputs of each previous block, due to residual connections. Outputs of attention layers are also weighted sums of potentially many embeddings, at least initially when attention patterns are spread out. By linearity, associative memories of the form (6) simply operate individually on each embedding of superposition, and return a new superposition (up to additional noise due to saao-oothgoonlllt)). In practice, we will see that learned memories often focus single embedding and filter out the rest as noise when irrelevant (see also Section 6). We note that linearity can also be limiting, since makes it difficult to map sets to specific output embeddings: ((.j) = 14 + u needs to map to W u; W uj, and thus cannot map to a new embedding "(SJ)- Such mappings of sets thus require non-linear associative memories, for instance by leveraging sparse decoding of which elements actually present (e.g., using compressed sensing), or by using MLPs with non-linear activations [5 30
4.2 A simplified two-layer transformer architecture
We consider simpler two-layer transformer which is more interpretable with the memory viewpoint, and will help us analyze learning dynamics both empirically and theoretically.
We freeze input, output and positional embeddings (WE, We, Wp) to their random initialization throughout training. This brings us to the Gaussian random vector setup presented above.
We fix W W3 Ia. so that W k and W play the role of both key and query matrices. This changes the gradient dynamics, but simplifies the model by avoiding the redundancy in (2). The pre-softmax attention scores then take the form x. 3T,, with T. (resp. m²) the query (resp. key) embeddings, which now directly resembles an associative memory lookup.
We freeze W7. W and W/ to random initialization. These play the role of remapping attended tokens into new tokens, since for random W and large d, Wa is nearly orthogonal and to any other random embeddings independent of x.
We train W since the outputs of the induction head need to be mapped back into appropriate output embeddings in order to predict the output tokens o correctly.
We use a single linear feedforward layer after the second attention block, with weight matrix Wr- This is plausibly the layer responsible for learning the global bigram distributions.
We remark that while this model freezes some parameters at initialization, it is richer than a "lazy" or neural tangent kernel approximation IT0 22 24 since the model is still highly non-linear in its parameters and, as we will see, induces rich non non-linear learning dynamics.
Solving the bigram problem with associative memories. We now show how the above architecture can solve the synthetic bigram problem from Section with well-chosen weight matrices. While this is only hypothetical model, we show in Section 5 that it is surprisingly faithful to the learned model.
Recall that due to residual connections, the inputs to the weight matrices typically consist of superpo- sitions of various embeddings including token embeddings, positional embeddings, or "remapped" versions thereof. These may be viewed as sets, as illustrated in Figure and associative memories can easily ignore certain elements of the set, e.g., ignore token embeddings by only focusing on positional embeddings. The induction head mechanism can be obtained by setting:
T W/- Σ PIPLD 1-2
W = Σ wa(H)(WdW)wxZ() KEQ
W3 = Σ wr wn(K)((2wE(K))
where Q is the set of triggers when they are fixed, or the support of ny when they are random. In words, the first attention layer matches a token to the previous tokens using positional embeddings. The second layer matches the trigger token to a remapping of itself by W!W,, and the output matches remapping of the input token by W3 to the corresponding output token. We remark that one can easily make the attention patterns more peaked on the correct associations by rescaling k and W The global bigram statistics can be encoded in the feed-forward layer as follows:
WF
log n. ()
(8)
6

---

The question remains of how the model could trade-off predictions from the induction head and from the reed-forward layer, which are added together due residual connections. With fixed triggers Q. we may simply remove all i E Q from the summation in (8). so that the model exclusively relies on the attention head for all triggers (indeed, the output of WO is in the span of output embeddings, which are nearly orthogonal to the row space of F)). When the triggers can vary across different sequences, choosing between the induction head and the feed-forward layer is more ambiguous as it depends context, and WF may try to learn more complex mappings that also use the outputs of Wo In practice, we observe that the model often prefers the induction head, unless its output agrees with one of the top predictions from the global bigram, in which case it tends to prefer those.
Beyond the simplified architecture. While our simplified architecture already captures the relevant aspects for the bigram model, it lacks some of the components that appear in standard transformers, such as non-linear MLPs, trained embeddings, layer normalization, and joint learning of a factoriza- tion W / WQ (potentially with low rank matrices WK: WQ E RR xd with d, < d in multi attention), instead of a single matrix Wx. In practice, transformers also involve many more layers, as well as multiple heads at each self-attention layer. In Appendix D we discuss how our memory viewpoint naturally extends to such architectural components, and we illustrate in Appendix that they empirically lead to similar observations. Nonetheless, we focus on our simpler architecture the main paper due to simplicity of exposition and better interpretability thanks to clear identifiability of the role of each matrix, which is lost in models with more heads and layers.
Empirical Study
In this section, we present empirical analysis of learning dynamics the bigram data defined in Section 3 for the simplified architecture defined in Section 4.2 See Appendix [E] for additional results.
Our code is available htttts://gitttut http://gttuub..oomaaabtt////t//t awwiaatllbtttaee
com httpp/////.ellee/eeraaaaooorrrr. wttp://tttaasfoorreeeeeeoe birth
Experimental setup. We train our models using mini-batch SGD with momentum, where each batch consists of 512 fresh sequences of length T = 256 sampled from our synthetic model. We a fixed learning rate and weight decay. Hyperparameters are given in Appendix
[ Unless otherwise noted, we use d 128, random triggers with Te Ru and uniform output tokens. The reported accuracies and losses are computed over each fresh batch before it is used for optimization, and are averaged over relevant tokens: "in-context accuracy/loss" numbers only consider predictions of output tokens on triggers starting at the second occurrence (the first is non-deterministic) while "global loss" refers to average loss on non- non-trigger tokens.
Memory recall probes. In addition to loss and accuracy, we consider metrics to check whether individual matrices have learned the desired associative memories: for a desired target memory W. E(.eM uzu. the corresponding recall metric is computed from the empirical estimate W as
R(W. W.) = I.
1arg max J.W.; = = (i.))EM
Σ
(9)
We use this for each matrix in 7 as target, and additionally test the previous token matrix W on smaller time windows. For the final foed-forward layer, we measure the average KL divergence between the predicted softmax distribution using only WF and the global bigram distribution my:
dKL(WF, F, mb)
dxr. (WW mo (-(k)))).
(10)
Emergence of the induction head via top-down learning. We begin our study by only training to minimize the loss on trigger-output token predictions after their first occurrence. This should be predictable with 100% accuracy using the two-layer induction head mechanism according to Section 4 We also remove the foed-forward layer, in order to focus on the learning of attention matrices WA. W and W& in isolation.
Figure 3 studies the effect of freezing different layers until iteration 300 on the training dynamics. By looking at memory recall probes, that training key-query matrices does not lead to any
7

---

freeze Wg
freeze W2 and W
freeze W3
freeze W?
1.2
comou SCCUrSTy
0.0
0.0
3.-
0.-
0.6
2.6
///²//
02-
0.5
0.0
0.0
3.0
202
400
602
1000
2001
607
603
801
1000
200
601
601
001
1000
2000
6:0
602
801
1000
1.0
1.0
1.0
0.0
0.0
2.0
0.5
04
Wed
Wiz
Wed (b-6)))
400 007 a50
400
2000
200
403
002
002
2000
4.0 002
3000
Figure 3: Learning the induction head alone: in-context accuracy (top) and recall probes (bottom) with some layers frozen until iteration 300. The output matrix W can and must be learned before the key-query matrices, but does not suffice for good accuracy. is easier to learn R before W k. and W stores initial context positions (t < 64) much faster than late positions.
learning unless W3 is learned first, and that W3 can learn the correct associations even when trained by itself with key-value matrices at random initialization. Recall that the attention weights essentially uniform when WK are at random initialization, so that training W alone resembles a bag-of-words models that aggregates representations throughout the sequence. While such a model has poor prediction accuracy, it is nevertheless sufficient to recover the correct associations in W (a similar observation was made in [4) in a different setup).
Then, these associations enable learning key-query matrices that focus the attention on relevant tokens, by storing relevant key-query pairs in the form of associative memories, which eventually recovers the desired induction head behavior and leads to near-perfect accuracy. The two rightmost plots suggest that the second layer is leamee before the first, in the sense that W3 is easier to learn when W is frozen compared to the reverse, yet learning them together seems beneficial, possibly due to helpful feedback loops [1. We also observe that W fits previous token associations for early positions much faster than later positions (purple vs gray line). This is likely due to the fact that it should be enough for the previous token head to attend to the first appearance of each trigger qk. which is typically early in the sequence, so that most of the gradient will focus on early positions.
Overall, this provides a fine- grained understanding of the learning dynamics of induction heads. In Section 6 we analyze how a few gradient steps in a top-down fashion may suffice to recover appropriate associative memories in high dimension and with enough data. See Appendix E for additional experiments, including on the role of dimensionality.
Global vs in-context learning. Figure 4 left/ight) shows that when training all layers jointly, the global bigram statistics tend to be learned more quickly than the induction head, as seen from the quick drop in loss and KL in early iterations. The W3 probe also seems to improve quickly initially, but only leads to mild improvements to in-context predictions. The full learning of the in-context mechanism takes longer, likely due to slower dynamics of the key-query matrices. We also observe a tension between W3 and WF later in training, leading to slight degradations of our probe metrics. This may be due to the fact that the input to WF now contains additional signal from the induction head which may be leveraged for better predictions, in particular for disambiguation in the case of random triggers, so that our guess of memories in Section may no longer be accurate.
Role of the data distribution. We can see in Figure lcft) that changes to the data distribution can have a significant effect on the speed of learning the in-context mechanism. We observe that the following may slow down in-context learning: (i smaller number of triggers K, (ii) using only rare fixed triggers, and (iii) using random triggers instead of fixed triggers. By inspecting the individual memory probes (see Figure 5 in Appendix [E). we hypothesize that (i) and (ii) are due to slow learning of W while (iii) is more related to slow learning of key-query matrices. This is reasonable since (i-ii) reduce the number of overall output tokens in the data, while (iii) increases the number of possible trigger tokens that should be stored in W thus increasing the data requirements in order to learn the full associative memory. We also show in Figure F center) that changing the output token distribution to bigram distributions at training time reduces the in-context accuracy
8

---

in-context vs global loss
out distribution
attention and feed-forward probes
Kel, frequent
In-context accuracy
0.n
WK2 recal
W92 recal
WKL recall (x664)
-20
KeG, random
0.6
l085
0.4
0.4
1.6
unift
0.2
unif (train bigram)
bigram (train unif)
0.2
1.2
0.0
bigram ttrain blgrami
1.0
1000
200
300
400
500
100
200
300
400
500
6
100
200
300
400
500
iteration
Iteration
Iteration
Figure 4: Global vs in-context learning and data-distributional effects. (left) Loss global (dashed) vs in-context (solid) tokens throughout training, for fixed or random trigger tokens Ok- The red curves fixes the trigger qi to the most frequent token, while the fixed triggers in bluc curves are less common. (center) In-context accuracy with different training and test distributions To for output tokens. Uniform leads better generalization than global bigrams da- (right) Probe metrics throughout training: W3 and Wr eventually compete and deviate from our natural estimates.
when using outtof-distribution output tokens, while training on uniform outputs performs well on both distributions. This highlights that using more diverse training distribution can lead to models with better generalization accuracy, with little additional training cost.
Additional experiments. In Appendix E provide additional experimental results for varying dimensionality, more complex architectures and training methods, as well as more fine-grained visualizations of the memory associations.
Theoretical Insights on Learning Dynamics
In this section, we provide theoretical insights on how gradients near initialization may allow the emergence of induction heads, and how this behavior is affected by data-distributional properties.
Finding signal in noisy inputs. In Lemma we showed how gradient dynamics on a simple classification task with fixed embeddings of the inputs and outputs lead to associative memories. We now show that when inputs consist of superpositions of multiple embeddings, as is the case in the transformer residual streams, gradients may learn associative memories that filter out irrelevant components of these superpositions, focusing on useful signal instead.
Lemma 2 (Gradient associative memory with noisy inputs). Let p be data distribution (x,3) Rd x [N], and consider the following classification problem, with fixed output embeddings Wo:
L(W)= = B(x.))~pl [[ Wo Wx)).
The gradients lake the following form: denoting pa = E[xly k]] and jux Ezz
Pw (k) ply=k)
wr).
VwL(W) = Σ P(v==k)uu())) (k) - -
The key takeaway from this lemma is that with enough data (here infinite data), the associative memory arising from gradients can learn to filter out noise from inputs, since it only depends on its expectations or conditional expectations. In particular, μı can isolate relevant parts of x that are predictive of label is, and thus can lead the right associations.
An illustrative example. To gain more intuition about this result, consider the following example: we would like to predict y from wg(v) + Pt. where Pt a positional embedding at a random position t E [T]. which we would like to ignore. Further assume that is uniformly distributed with p(y k) 1/N, and consider the matrix obtained after one population gradient step with step-size starting from an initialization Wo = 0 (so that Pw, (k( = 1/N):
N
W
7 N
we (b) - T.

---

with == B[r]. We show in Appendix B that when d is large enough to ensure near-orthonormal embeddings, we have
wy ( (WE(V) + pc) R
7 N
[(k= 1 y} +
1 N2
so that for large enough N and T, we obtain near- pear-perfecct classifier that ignores the positional embedding, after just one gradient step (but highly idealized one). Understanding how this translates to the finite dimension and finite sample regime is an important theoretical question that we leave for future work (see [7] for an initial step in that direction). We note that data models related to the above have been useful to study gradient dynamics of neural networks on continuous data 22 25 28) Using single gradient step to learn representations has also been fruitful in other contexts 3 12
Learning the induction head with gradients. We may extend the arguments above to show how a few gradient steps can learn the induction head mechanism. We show the following in Appendix [B.3
Theorem 3 (Learning induction head via three gradient steps, informal). In simplified setup, the induction head mechanism as constructed in (7) can be learned via sequential gradient steps on the population loss from random initialization, on W2, then W followed by W K-
To show this result, we use Lemma 2 in similar manner to the illustrative example above to show how training W3 by itself at initialization, i.e., when the attention patterns are near-ume ceeuruniform, can recover the desired associative memory. This is possible because when predicting an output token at later occurrences of a trigger, the same output token guaranteed be present the context, while other tokens need not appear more relative to other sequences. See also Figure in Appendix [ for numerical experiments verifying this for finite data and dimension. Once W has learned the correct associations, we show that the gradient with respect to the key-value matrix W? R at zero initialization can leverage the correctness of W to find the right associative memory that focuses the attention on correct triggers. Finally, by linearizing the second-layer attention around W = O, we show how gradients w.r.t. W may learn correct associations for the previous token head.
7 Discussion
In this paper, we studied the question of how transformers develop in-context learning abilities, using a simplified setup that allows fine-grained understanding the model and its training dynamics. While our model already captures rich phenomena at play in the bigram task we consider, more elaborate models are likely needed to understand transformers trained on more complex tasks like language modeling. This includes learning embeddings that are more adapted to the data and more structured (e.g., word embeddings (1). grokking [33 39),, factorized key-query and value- output matrices that may induce additional regularization effects 20), and non-linear feedforward layers, which may provide richer associative memories between sets of embeddings. Understanding how transformers leverage such aspects to learn in richer settings is an important next step.
Acknowledgments and Disclosure of Funding
The authors thank Sainbayar Sukhbaatar and Shubham Toshniwal for helpful discussions.
References
[1] Z. Allen-Zhu and Y. Li. Backward feature correction: How deep learning performs deep learning. arXiv preprint arXiv:2001.0441: 2020.
[2] Z. Allen-Zhu and Y. Li. Towards understanding ensemble, knowledge distillation and self- distillation in deep learning. In Proceedings of the International Conference Learning Representations (ICLR), 2023.
[3] J. Ba, M. A. Erdogdu, T. Suzuki, Z. Wang, D. Wu, and G. Yang. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems (NeurIPS), 2022.
10

---

[4] L. Ba, J. R. Kiros, and G. E. Hinton. Layer normalization. arXiv preprint 160.006607.06450, 2016.
[5] T. Bricken and C. Pehlevan. Attention approximates sparse distributed memory. Advances in Neural Information Processing Systems (NeurIPS), 2021.
[6] Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
[7] V. Cabannes, E. Dohmatob, and A. Bietti. Scaling laws for associative memories. arXiv preprint arXiv:2310.02984, 2023.
[8] S. Chan, A. Santoro, A. Lampinen, J. Wang, A. Singh, P. Richemond, J. McClelland, and F, Hill. Data distributional properties drive emergent in-context learning in transformers. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[9] E. Charton. What is my math transformer doing? crinng?thhre threee results interpretability and general- ization. arXiv preprint 1rXiv:2211.001700 2022.
[10] L. Chizat, E. Oyallon, and F. Bach. On lazy training in differentiable programming. Advances in Neural Information Processing Systems (NeurIPS), 2019.
[11] A. Chowdhery, S. Narang, J. Devlin, M. Bosma, G. Mishra, A. Roberts, P. Barham, H. W. Chung, C. Sutton, S. Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311, 2022.
[12] A. Damian, J. Lee, and M. Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference on Learning Theory (COLT), 2022.
[13] G. Dar, M. Geva, A. Gupta, and J. Berant. Analyzing transformers in embedding space. arXiv preprint arXiv:2209.0255 02555, 2022.
[14] B. L. Edelman, S. Goel, S. Kakade, and C. Zhang. Inductive biases and variable creation in self-attention mechanisms. In Proceedings of the International Conference on Machine Learning (ICML), 2022.
[15] 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, and C. Olah. Toy models of superposition. Transformer Circuits Thread, 2022.
[16] N. Elhage, N. Nanda, C. Olsson, T. Henighan, N. Joseph, B. Mann, A. Askell, Y. Bai, A. Chen, T. Conerly, N. DasSarma, D. Drain, D. Ganguli, Z. HatfieldDoodss, D. Hernandez, Jones, J. Kernion, L. Lovitt, K. Ndousse, D. Amodei, T. Brown, J. Clark, J. Kaplan, S. McCandlish, and C. Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021.
[17] M. Geva, J. Bastings, K. Filippova, and A. Globerson. Dissecting recall of factual associations in auto-regressive language models. arXiv preprint vrXiv:2304.14767 2023.
[18] M. Geva, R. Schuster, J. Berant, and O. Levy. Transformer feed-forward layers are key-value memories. In Conference on Empirical Methods in Natural Language Processing (EMNLP), 2021.
[19] A. Graves, G. Wayne, and I. Danihelka. Neural turing machines. arXiv preprint arXiv: 2014.
[20] S. Gunasckar, B. E. Woodworth, S. Bhojanapalli, B. Neyshabur, and N. Srebro. Implicit regularization in matrix factorization. Advances in Neural Information Processing Systems (NIPS), 2017.
[21] J. Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences, (49)-2555-255 1982.
11

---

[22] J. Hron, Y. Bahri, J. Sohl-Dickstein, and R. Novak. Infinite attention: Nogp and ntk for deep attention networks. In Proceedings of the International Conference on Machine Learning (ICML), 2020.
[23] A. Iscen, T. Furon, V. Gripon, M. Rabbat, and H. Jégou. Memory vectors for similarity search in high-dimensional spaces. IEEE transactions big data, 4(1 (1)5667777, 2017.
[24] A. Jacot, F. Gabriel, and C. Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances Neural Information Processing Systems (NeurIPS), 2018.
[25] S. Jelassi, M. Sander, and Y. Li. Vision transformers provably learn spatial structure. In Advances in Neural Information Processing Systems (NeurlPS), 2022.
[26] Y. Jiang and C. Pehlevan. Associative memory in iterated overparameterized sigmoid au- toencoders. In Proceedings of the International Conference on Machine Learning (ICML), 2020.
[27] A. Joulin and T. Mikolov. Inferring algorithmic patterns with stack-augmented recurrent nets. Advances in Neural Information Processing Systems (NIPS), 2015.
[28] Karp, E. Winston, Y. Li, and A. Singh. Local signal adaptivity: Provable feature learning in neural networks beyond kernels. Advances in Neural Information Processing Systems (NeurIPS), 2021.
[29] Kohonen. Correlation matrix memories. IEEE Transactions on Computers, 1972.
[30] D. Krotov and J. Hopfield. Dense associative memory for pattern recognition. Advances in Neural Information Processing Systems (NIPS), 2016.
[31] Y. Li, Li, and A. Risteski. How do transformers learn topic structure: Towards mechanistic understanding. In Proceedings of the International Conference on Machine Learning (ICML), 2023.
[32] B. Liu, J. Ash, S. Goel, A. Krishnamurthy, and C. Zhang. Transformers learn shortcuts to automata. In Proceedings of the International Conference on Learning Representations (ICLR), 2023.
[33] Z. Liu, O. Kitouni, N. S. Nolte, E. Michaud, M. Tegmark, and M. Williams. Towards un- derstanding grokking: An effective theory of representation learning. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[34] R. McElicce, E. Posner, E. Rodemich, and S. Venkatesh. The capacity of the hopfield associative memory. IEEE transactions on Information Theory, (3(44)441-448) 1987.
[35] K. Meng, D. Bau, A. Andonian, and Y. Belinkov. Locating and editing factual associations in gpt. In Advances Neural Information Processing Systems (NeurIPS), 2022.
[36] W. Merrill, A. Sabharwal, and N. A. Smith. Saturated transformers constant-depth threshold circuits. Transactions of the Association for Computational Linguistics, (0:843-856, 2022.
[37] Mikolov, K. Chen, G. Corrado, and Dean. Efficient estimation of word representations in vector space. In Proceedings of the International Conference on Learning Representations (ICLR), 2013.
[38] S. Min, X. Lyu, A. Holtzman, M. Artetxe, M. Lewis, H. Hajishirzi, and L. Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? In Conference on Empirical Methods in Natural Language Processing (EMNLP), 2022.
[39] N. Nanda, L. Chan, T. Liberum, J. Smith, and J. Steinhardt. Progress measures for grokking via mechanistic interpretabilityy. In Proceedings of the International Conference on Learning Representations (ICLR), 2023.
12

---

[40] C. Olsson, N. Elhage, N. Nanda, N. Joseph, N. DasSarma, T. Henighan. B. Mann, A. Askell, Bai, A. Chen, Conerly, D. Drain, D. Ganguli. Z. Hatfield-Dodds, D. Hernandez, S. Johnston, A. Jones, J. Kernion, L. Lovitt, K. Ndousse, D. Amodei, T. Brown, J. Clark, J. Kaplan, S. McCandlish, and C. Olah, In-context learning and induction heads. Transformer Circuits Thread, 2022.
[41] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, al. Language models are unsupervised multitask learners. Technical report, OpenAl, 2019.
[42] J. W. Rae, S. Borgeaud, T. Cai, K. Millican, J. Hoffmann, F. Song, J. Aslanides, S. Henderson, R. Ring, S. Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:21
2021.
[43] Y. Razeghii L. Logan IV, M. Gardner, and S. Singh. Impact of pretraining term frequencies on few-shot reasoning. In Conference Empirical Methods in Natural Language Processing (EMNLP), 2022.
[44] C. Sanford, D. Hsu, and M. Telgarsky, Representational strengths and limitations of transformers. In Advances Neural Information Processing Systems (NeurIPS), 2023.
[45] A. M. Saxe, J. L. McClelland, and S. Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv: 1312.6120, 2013.
[46] I. Schlag, K. Irie, and J. Schmidhuber. Linear transformers are secretly fast weight programmers. In Proceedings of the International Conference Machine Learning (ICML), 2021.
[47] Schmidhuber. Learning to control fast-weight memories: An alternative to dynamic recurrent networks. Neural Computation, 4(1):131-139 1992.
[48] S. Shin, S.-W. Lee, H. Ahn, S. Kim, H. Kim, B. Kim, K. Cho, G. Lee, W. Park, J.-W. Ha, et al. On the effect of pretraining corpora on in-context learning by large-scale language model. In North American Chapter of the Association for Computational Linguistics (NAACL), 2022.
[49] C. Snell, R. Zhong, D. Klein, and J. Steinhardt. Approximating how single head attention learns. arXiv preprint 1xXiv:20030070/1 2021.
[50] Sukhbaatar, A. Szlam, Weston, and Fergus. End-to-end memory networks. Advances in Neural Information Processing Systems (NIPS), 2015.
[51] Y. Tian, Wang, B. Chen, and S. Du Scan and snap: Understanding training dynamics and token composition in I-layer transformer. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
[52] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems (NIPS), 2017.
[53] R. Vershynin. High-dimenssonal probability: An introduction with applications in data science, volume 47, Cambridge university press, 2018
[54] K. Wang, A. Variengien, A. Conmy, B. Shlegeris, and J. Steinhardt. Interpretability in the wild: a circuit for indirect object identification in gpt-2 small. In Proceedings of the International Conference on Learning Representations (ICLR), 2023.
[55] Weston, S. Chopra, and A. Bordes. Memory networks. In Proceedings of the International Conference on Learning Representations (ICLR), 2015.
[56] D. J. Willshaw, P. Buneman, and H. C. Longuet-Higgins. Non-holographic associative memory. Nature, 222(5197) 1969.
[57] S. M. Xie, A. Raghunathan, P. Liang, and T. Ma. An explanation of in-context learning as implicit bayesian inference. In Proceedings of the International Conference on Learning Representations (ICLR), 2022.
13

---

[58] G. Yang and E. J. Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In Proceedings of the International Conference on Machine Learning (ICML). 2021.
[59] G. Yang, E. J. Hu, I. Babuschkin, S. Sidor, X. Liu, D. Farhi, N. Ryder, J. Pachocki, W. Chen, and J. Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
[60] B. Zhang and R. Sennrich. Root mean square layer normalization. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
[61] Y. Zhang, A. Backurs, S. Bubeck, Eldan, S. Gunasekar, and Wagner. Unveiling transform- ers with lego: synthetic reasoning task. arXiv preprint arXiv:2206. 04301, 2022.
14

---

A Associative Memories with Random Vectors
In this section, we provide basic properties of associative memories based on outer products of random Gaussian embeddings, as described in Section 4
We consider embeddings ug E Id with i.i.d. Gaussian (O, à) entries.
We recall a few facts:
(Norm) We have u,! u, 1. This is standard from the concentration of random vectors in high dimension (see, e.g., [53, Theorem 3.1.1D.
(Near-orthogoalltt)) For i + J. we have un " 0(1/Vd.. To see this, denoting u d=1//2 dūt/)(ūg))R, where ū. are the unnormalized entries of u, note that we have
ddu, U Va
Uu 1),
by the central limit theorem, since for each k the quantities ūي Uyh are zero-mean, unit-variance, i.i.d. random variables.
(Remapping: norm) If W is a Gaussian random matrix with i.i.d. Ao, 2) entries, then for any fixed æ we have ||Wa|| ze || This follows from Johnson-Lindeenstraus (see, e.g., [53] Lemma 5.3.2 and Exercise 5.3.3J). In particular, if is normalized Gaussian embedding above, then || 1.
(Remapping: near-orthen ) Consider a random vector x = va"
and a random matrix W
1 W, where the entries of i and W are i..d. (0, 1). Then x and Wa are nearly orthogonal. To va
see this, note that [x
Wx] TWa]
E[x EW]
and the variance
E(x Wx)²
2
E
2
W
Σ
i
(
1 d3
MM + d(d- ( M 0
2
where M2 and M4 denote the second and fourth moments of the standard Gaussian, respectively. Then, Chebyshev's 's inequality implies that x Wa O(1/Vd) with high probability.
Ensuring appropriate memory lookups then requires such properties to hold for many embeddings and pairs of embeddings. with errors that are small enough to ensure correct associations. This may be achieved with careful union bounds or more powerful concentration results. We do not attempt do this precise manner in this paper, and will generally assume large enough to satisfy the desired associative memory behaviors, noting that a precise analysis is an important direction for future work (see 7 for a first step in this direction).
B Theoretical Insights on Gradient Dynamics
In this section, we provide additional details on the theoretical insights from Section 6 including details the illustrative example (Section B.T. derivations of gradients w.r.t. key-query matri- ces at initialization (Section B.2), as well as a study of how the induction head mechanism may develop in simplified setup, using sequence of single layer-wise wise gradient steps in top-down manner (Section B.3).
B.1 Details on illustrative example
Consider the example discussed in Section b we would like to predict from = wF(y) + Pt-
where is positional embedding at a random position / [T]. which we would like to ignore.
Further assume that y uniformly distributed (p(y k) = 1/N) and consider the matrix obtained
after one population gradient step with step-size 77 starting from initialization Wo (so
15

---

that (k(z)) = 1/N):
W1
7
1e (k)) 10),
(11)
N
with μ = E[
Note that we have u = wk(k) + E, Pt and j = N Ex we (k) + E. Pt. so that 11
becomes
" W1
N Σ
(y(kk(WB(()) (k))
(12)
with we /े Ex=1 (k()). When d is large enough ensure near-orthonormal embeddings, have for any y and f,
7 Wi(we(s) + + = wo(()) + +0 +0 N2 (N²2)
This implies
wy (K) Wi( +p) 7 1(k y}
N2
as claimed in the main text. The classifier arg max wy (k) W (w (v) +p:) then has essentially perfect accuracy, and has learned to ignore the spurious positional embeddings, which are simply exogenous noise.
B.2 Gradients on key-query matrices at initialization
We now derive expressions for population gradients of the attention key-query matrices zero initialization, noting that random initialization behaves similarly to zero initialization. Although the optimization problems involving these matrices are non-convex, these gradients at initialization lead to associative memories, similar to Lemma When output matrices of the previous layer already encode the desired associations, these gradients can lead associative memories that focus the attention on the correct key-value pairs.
We begin with the following lemma, which gives the gradient of the loss w.r.t. W W zero initialization. For simplicity, we drop the d-1/2 factor from the softmax, which only changes gradients by multiplicative factor, and thus does not change its form.
Lemma 4 (Gradient of second attention layer). Consider the following loss for predicting the next token from an attention layer with inputs X [x1,...,,T]] and value-output mairix $2 W
L(W)- Ex
<(X) 22 (XTTT))
(13)
with the cross-entropy loss and (u)) = enc
E. for E IR." is the softmax.
The gradient at W 0 is given by
N
VWL(W)) W
E (x
1 (Dw(k)))- 1(v 1yy= k3)
Σ (k)(k) T22e (xc 2nT))aT
1
Ex liw (klX))r (k) D2 (xt - Bur)aal
T
N
ply = k) Ex (k) Dyt
(x. Tob Txr)) 1y yy= =k]
with 31. = 7 Et=1 Tz.
Now we consider the gradient w.r.t. W = WW at zero initialization, and consier simplification of the second layer attention its linearization around W = 0. We will see that this still provides first-order information that is sufficient for W1l be learned.
16

---

Lemma 5 (Gradient of first attention layer). Consider the following loss for predicting the next token y from a stack of two attention layers, with all parameters fixed except for W W kc the key-query matrix at the first attention layer:
L(W) - x) (((,,(X)))
=(X)= Wo--W/xbxaa0 W24T).
(14)
Here, o(u1:T) - +11+u 7 Eam us) is the linearization of the softmax around O and Z(W)
=
[21 (W) (............................... 2T (T(W) with
z(W) (10.0 Wt)ss:
and 22 = W for { 1,2
The gradient at W = 0 is given by
V L()
W=
N
Σ
Ex
pw (w(x)
1
wol (k) T
T222=
1 (1xx)) W2oT(ps (eaxxssWWWTlle D(:)
T
s=1
N
ply k)) Ex
Σ
we (1:)
( (1 (h)ps))T2z W22r - p0-)pp w
k=1
N
T
1
Σ
Ex pe (klo)
(21xz) Wwwr(P:- - pLcelt
T
1=1
8=1
T
1
Σ k=1
ply 1)) Ex (k((k) D2ū
(1xx))
(T W22r(p. -P.:).) 1y k
T
B.3 Learning the induction head mechanism
In this section, we analyze the training dynamics of the induction head mechanism, in the following simplified setup: we consider single trigger (K 1), and assume that Tas Tg. To and (-(-))) are uniform over [N] for any i.
To further simplify the analysis, we consider a loss that only considers sequences of length T where the last input 2T is the second occurrence of the trigger token, and the label y = 2T+1 is the corresponding output token. We note that this may be easily extended to later occurrencies of the trigger. This is similar to the setup of Figure [ where the loss is only taken on triggers after the second occurrence: in that case, the loss may be written
weighted sum of the here, weighted by the probability of the second (or later) trigger appearing at the given position T.
consider
In practice, when the loss is on all tokens and WF is also learned, we may expect that WF quickly learns the global bigram statistics, as we saw empirically in Section 5 Indeed, the current token embedding, which is included in the input superposition, has strong predictive signal compared to the attention layers, which initially mainly appear as noise. This is then similar to the setup of Lemma [ which provides recovery of bigram statistics when dl is large (though we note that the other information from attention layers in the inputs may eventually be used and bias away from perfect recovery, see Figure 4 rright)). Once such global estimates are obtained, the expected loss will be mainly dominated by trigger tokens, leading to the setup above.
For simplicity, we thus drop the feed-forward layer WP in the remainder of this section, focusing on the learning of W3. W K and WT. in this top-down order. We will consider zero pro-initialization single gradient steps, noting that random initialization should lead to similar associative memory behaviors when the dimension large enough, since leads to a remapping of input embeddings which is near-orthogonal to any output embedding (see Appendix A)
B.3.1 Learning W3
We begin by studying the learning of the second output matrix W In the above data model, we may consider loss in Lemma 2 with onput-outputs (x, v). where y is the output token of the sequence,
17

---

and z depends on the random sequence 21:7
x 7 D W (www((z) + st);
where e, = Pp: + L1 (1(WE(2)) pp) with Φ1 = W W. is a "noise" vector from the residual streams, containing positional embeddings as well as average attention output from the first layer. In practice, the logit predictions are the form (r + ET) due to residual connections, but we ignore the term Woer for simplicity, noting that it is near-zero when d is large.
After a gradient step on W with step-size 7) starting from sero-initialization (so that p(k]a)))
p(y = k) 1/N for all r), Lemma ylields
N we (b) = k] B(x))
7 3 = N
Σ
(15)
Now, consider the random variables q (trigger token), o (output token), to (position of the first occurrence of the output token). In our simplified data model, q and to have the same distribution regardless of the conditioning on y k, while o is equal to k when k, while it is uniform in [M] without this condition. The sequence 21.7 has the same distribution in either p(-) or p(-ly k).
except for the token =t.-
Then, may write:
1
B[xly = k) -][b[
[E(W3Wg(FL.) = 1] 3 wE (zc,))
T
1
E
WBedly=k
W
T
since E is independent of o when t < to. Noting that E, only depends on 2t, and thus on through the first layer attention, we have
Eizly
B[x] (wE((()) wg)
T
T W W33)
+ E
Σ
W 2 wx( k -
W (wE (R:) wg + T
W o (w)(()) wg),
= T
where T = E. Ei=to and üe = 1 EN (k()). Thus, (5) becomes
W3 17 NT we (W)W wE))
1
wy (I) wE)) (16)
NT
so that when is learn enough to ensure ncar-orthonormal embeddings, we have
wy wx)( W3W wy (k:)k:) WW3 (ww()
" 1(k = /) +
NT
N2T
NT [k= +
'
T/T NºT
where the O()) terms are due to the Wg elements. The first line yields behavior that matches desired associative memory in (7) of Section 4.2 when is large. The second line shows additional spurious associations that are stored in W3 due to the output of the first layer attention, but which may be "cleaned up" once the attention layers start focusing on the correct tokens.
Finally, we note that despite the recovery of these useful associations after one gradient step, the predictions with this estimate W/ are still near-random, since the bag-of-words setup with average attention, the output token cannot be distinguished from any other token in the sequence in our model (except perhaps the trigger token, which is guaranteed to appear twice, but does not provide any signal to infer the output token, since the two are independent).
18

---

B.3.2 Learning 3
Now assume that W3 is in 16p. As argued above, the predictions p(k(a) are essentially ran-
dom 1/N in our model for W 0, so that after one gradient step on W with learning rate Lemma yields::
W
Σ
(k)
P224 (xc - 3)0} y k] Te P2rz (0x -))x+))) (17)
TN
R,T
where 2y are the inputs to the second attention layer, given by
24 = Tt,0 + Ft,1
(18)
Ft.0 wg(+1) pt
(19)
Σ $ (we(z)) + ps).
(20)
From now on, we consider a simplified architecture where only te,0 are fed queries and values. while only Tt, are fed as keys. Using the fact that trigger tokens q are sampled uniformly (i.e..
Tg = 1/N). we have
W
n
[B 1 Ex (Ar,e))
(21)
TN
T
N
n
[E[A 1y= y- k 5 - E[A q= jl)
(22)
TN2
where
At. = wy (k)
2200000 (F - f1)=TT..
(23)
with J1 = Σ, 21,1- Now, note that we have wy (k) D2 al = k} with n/TN
by (16). and TT,0 = wp(g) + pr. This yields
W
OF
(wE()) + PTT)
(24)
TN2
with
Akj E
1 1/ k)),1111111111 Bi)lyy k
E
Σ
1(a ( =1)10 }
Di, A Ak.
where the three terms split the sum inside the expectation
4 E
[1
(a Ei)ly k j]
E
[2,, k)( -b1)l0- jl
Ah,-
E
Σ 1 -
E
1(2= =z)la=5
1ET,
1ET,
Mr := E
-===k)((z1 - - - E
-
LET
LET.
where T {to 1, and T. [T] \ (to, 1,7} (recall that to is random variable, corresponding to the first occurrence of the output token, so that these sets are random).
We will now show that Ag carries the desired signal for the appropriate induction head associative memory, while Ai and AE, are negligible, for N large enough.
19

---

Controlling kg For t = to, noting that 2,, = we have
a E [1 yy k} (Zt..1 -1)ly k. = 13]-B[116 [L -
E [rr..l - k. 4
N
N
p Σ aRg D (,
N
with akjja 2 (fTwE(1)
E [0,,.11 - 11 = k = 5]
thanks to near-orthonormality and
to
T
p = Et
P
T
s=1
t=1
a spurious positional mixture.
We then distinguish the following cases:
If j * and i j. since the trigger token j only appears at positions to 1 and T,, we have
Σ
a&ji Er,
TT.
T
t=t.-1
We may expect to to be concentrated around T/2, in which case YT 2 for T' larger than small constant.
2 T
0
T
j = k i, the two occurrences of the trigger happen after the other, must be that to T. Then
T T(T 1) T T + )
2
2
Ob.i ~
for T' larger than small constant.
if * j k all tokens up to position to 2 = T 2 are iid. uniform in [N] \ {i}, that
- 2 (T (T- 2) N + (T- 1) 1) + T ) Ok.ji ~ T(N 1) T
T 2 1
T(N - T(N-1)
N
Iff and k, all tokens except at positions f, 1, and T (we have to < T) are uniform in \ [i}- The triggers do not contribute anything to dx.jj since + j. and the output token may be also randomized by taking the average over [N] (i}. We thus obtain
akjii O
N
N
In summary, obtain
1
10 1
ifi ifi=j.
akja
N
km1
Thus, when N is large, while 7 is moderate, the above sum leads to more signal in the i terms compared to i A j. In particular, this yields
(
ifi ifi
(Q)w)(1))
A
that this component in 24 acts precisely like the desired associative memory in 7 remains to show that the other components are negligible compared this. It then suffices show:
Ag, + AE = (
20

---

Controlling A k.j" Fort E Tg. note that we always have = in the expectations, so that 4 = 0
unless k = j. For lo j, we have 2 ((1), ) that
1 N
2 = 0
N
Controlling 3 Using that ||
21ll < C ((1) for all t, we provide the following crude
bound via the triangle inequality and Hölder inequality:
D = E
Σ
11= - -1110=
Σ
1 = k - =1)l==j
LIET
LEET
1 12 kkk= .1)
2CT N
A SC
E
Σ LIET,
1 - + B
Σ LIET
since 2, independent of y given t T. and thus uniform in [M] (1), and TT, S T. We note, however, that A may be controlled much more finely by leveraging the similarities between the distributions of Zz. E T, with or without conditioning on y-
Overall, we have shown that up some spurious positional embeddings, W behaves the desired associative memory from ( when is large enough, satisfying:
(f)EE()) WkW
{ Ω TN
am
1 T 1 11
}}
T
(25)
We note that one may then amplify the gap between correct and incorrect associations by having large enough step-size, which then makes the softmax more peaked and hence the attention more sparse and focused on correct associations.
B.3.3 Learning W
We now assume that W3 and W have learned the correct associations, and consider one gradient step away from sero-initialization still often near random chance. Indeed, the second layer attention will attend to all tokens starting at the first occurrence of the trigger, since all such tokens contain 4 w (g))g) in their average attention, which activates the second-lay layer attention head. Then the output likely to predict the trigger itself, which will be an incorrect prediction most of the time.
W Note that when W O, the predictions of the model are
We may thus consider p(k|X) 1/N at this stage as well. We also consider simplified architecture where the first layer attention only uses positional embeddings in the key-query matrix, and only token embeddings the value-output matrix. particular, we have x; (2(2).. Lemma then gives the following form for W after one gradient step of step-size ry:
W
7
Ex
We (k)
D2
Σ ( TWRaT(P. PIADD by k
N
T
77
Σ
Ex
We
(k)
(1
TW - pLNDTT
N
T
1
Σ
Ex
(k)
DTJIT
$ (x,)) TW3rr(P. PIo)eT ly
N
T
t=1
#=1
+ Ex
1
M
wo (k) T
1
(12))) T WTr - PL.D
Note that since W3 and W already captured the desired associations at this stage, we have
(k)(k) Dr x a = and (01z.) ( ~ a' 2r).
21

---

for some a, a' > O. Recall that in our model, we have 2T = with probability one (q is the trigger token), and that q only appears twice: once at position E to 'I' and once at position T. We then have, for any 1 1,
WkRe ~ naa'
[Aee Bid Ci,k + Dr,R),
NTI
k=1
with
Ar = [[ H[t 1 tp - p1:))l=kk
(26)
B [1
(p PI+t)l
(27)
C₁ E(rx 11g Sp,, PLol1u= ki
(28)
Di, E(rs [tty 1 (pe, p1c))..
(29)
where Th == T Σ 12 = k}. We have
At, = B ( c = k) (1 fgg = t - 1) - 1ft E 1 ( p1:)l=
= P(ty = f 1 = (p - p1:t) +
1
Σ
P(tg = sly = k) - -P1:))
N
sE|t-2]1(1}
P(l = 1 1)(p-11 p12) N
Σ
P(44 8) p1a))
sE(t-2]u((]
= P(, =t t 1)(p-11 P1a) + +0
N
since the distribution of t is the same regardless of y. We proceed similarly for the other quantities and obtain the following:
BI, = =0 (
1
N
Ct = P(Ly =1-1)
1 N
(p--1 (n-1--1a)+00
T
D1,6 = )
This yields the following associative memory behavior, for >>
Pe WK (T
F(I, -1 - (111)(111880 (1(8-1-1)- 16 (n)) +0 ( )
T24
which matches the desired "previous token head" behavior from 7 when N is large. As in the case of W we may then "saturate" the softmax by choosing a large enough step-size.
C Other Proofs
C.1 Proof of Lemma
Proof. Recall the form of the cross-entropy loss for classification with K classes:
N Σ
e
E(y.E) =
1 = k} log
et
km
Its derivatives take the form
() s(E), 11 k:)
with s(E)
the softmax.
E,
22

---

The gradient of L is then given by
N 0e
Vw L(W) = (xM)
Σ σEk
(y. W.WWW (w (w)x(())
(x,))
(pw (A|2) L (A))
N
B. (,yK 1y k} yw (k) ||]]
(=(((((k())) - [ = ==) (A))r==)))))))))))))))))))))))))
1=1
which yields the desired result.
C.2 Proof of Lemma 2
Proof. Using similar steps as the proof of Lemma
we have
de
V WwL(W) ((a
(y, Ww Wo)
dE
Σ (ow (k|2)) 1 = (k)x
we (k) E. low (k(ar)=]
Eg [1 (y k:wwr (A:) E[a[y] ]
A=1
wo (k) Ex (pw (k|x)=]
Σ
ply 5) 1 (k) 5
kØ=1
N
ply k)wo (b) px)T
with Pu ply k)) Eg [pw (k( and μn E[xly lc).
C.3 Proof of Lemma 4
Proof. To better isolate the role of keys from values, we denote the keys that fed into the matrix W by Z = [ ZT] E ıaxTT, while the query is simply IT. In practice we have Z = X, and both are superpositions of potentially multiple embeddings (if W is part of the second attention layer, these are the token embedding, positional embedding, and the output of the first attention layer).
The gradient of the loss at W 0 writes:
N
at
LWW) w = 2(x,2.))
(y, E) V ( 22 (ar)
(30)
DE
Exxx Σ
(iw (pw (k) (k]X, 2) 1 = =k)) Vw(wu(h 22 xxx
(31)
We have
T
Wxr)) ( W=0 = Ww(ww(k) 12Xo(Z (H)
Σ
wo (k)
D2 Vw(o(ZTWar))
t=1
T
Wer (k) 122, (2. ZI.T)
T
23

---

where 21.7 7 Ex 2e, and we used the fact that
8
(u) o(w)) - 7 -1t= 1[t= = 5)
(32)
ou.
T
The gradient (31 now writes
T
Vw W=
Σ Rxx (x 2) - ( = k}}
Σ
wy wv(kk 42 (2t - -+1r)
t=1
and the result follows.
C.4 Proof of Lemma
Proof. The linearization of the second layer softmax around takes the following form:
a( ZZT W25T): T (1 + W22r - W2TT))
with 2t Exm 1 1 (p W pt)s the output of the first attention layer.
T
we (k)
T 12 W2rT)
I
wy (k) T 12 +
Σ
wy (k) Dzzz ( (p
T 1=1
T
1=1
1
P2Z11T (k)
(12.)
Werro Wp)).
T
1=1
1=1
Then,
LWW(W)
(33)
W=
N
de
= =E(x.)
(y. E)
VW W=0
(34)
de
T
= -E(x))
(v.
Σ 1=1
(k)(k)
P2
(1T,)) Werr(p. pu.wo
(35)
ĐEk
1=1
OE
E(x.))
(y, (k)
(1æx)
Woar PIa)pE
DE
(36)
using
32
at The result follows by using ne.
(v =(k()) 1 = k}.
D Beyond our Simplified Architecture
While the simplified architecture presented Section 4.2 is sufficient to capture the desired induction behavior for our bigram task, transformer architectures used in practice typically involve more components, as well as more heads and layers. In this section, we discuss how our memory viewpoint extends to such architectures.
24

---

Factorizations. In practice, transformers typically involve products of matrices, potentially with low-rank bottleneck. For instance, our key-query matrices Wx should instead be considered as product W We. and the output-value matrices Wo and Wy are typically jointly optimized.
Consider associative memory of the form:
We = Σ yar E Rdxdd
where (r:); and (y:)s are appropriate collections of orthooooooaa embeddings. We now argue that a similar associative memory can be achieved with the factorization d UV, where U E Rd V Rd'xd with d' S d (for instance d' could be the dimension of attention heads), are given by]
2d'
U Up + Σ (Voo))
V = Vo + Σ u)d!
where Un and are random matrices with AN 4) entries. These matrices are similar to those that would arise from single gradient step individually on U and V from initializations Uo and Vo, as in To see why W behaves like Wa, note that we have
Lemma !1
UV UoVo Σ (Voo))))) Vo Σ(vo u))a + + vc(Vo)))) (Un vz)aJ.
2.
It is also easy to check using central limit arguments (similar remapping in Appendix A that i = v 3 Voza E gdd and it = V Uo y are all nearly-orthonormal embeddings. Thus, we have
20 Way = U3 5n Σ Uk y in + Σ Uh D.r In + Σ Uk ya ūyaear
10 1] 1[k=
where the first and last term vanish due to the cross-terms v. By which vanish for any i, f. Thus, W and W. encode the same associations, when d and d' large enough to ensure near-orthogonality
Layer-normalization. Normalization layers 60 are typically used in transformers to improve training stability [52), and applied each token representation, either after [52 or before 6 each block. It may be seen as an operation of the form
LN(x)
||>||
applied to the input output of given block.
In order to obtain basic understanding the role of layer-norm in our associative memory setup, we may consider the setup Lemma [ with normalization applied after the linear operation, leading to the population loss:
L(W)- E(z,))~ (( Wo WN
(37)
The gradients then take the form
LWW))> = ΣE Ex pw = k|r) - ply = k(zz)
(1- (W)) we wo (k)x
(38)
Wa
|
1=1
Whhen d' is the head dimension, the $ scaling can be interpreted as the correct multiplicr to use attention logits, which plays similar role the multiplicr in the jPP scaling [50 for our setup where the variance of the random entries of input embeddings is 1/d instead 1 as 59
*RMSNorm [00] would use the variance instead of the norm in the denominator, leading to additional vd factor in numerator. Here we use the norm, which more natural when embeddings have near-unit norm, contrast the vd norm for the standard parameterization.
25

---

This illustrates that in addition to weighting the updates on class k by the prediction error p(k[z)) p(k(z) the updates are also projected on the orthogonal of the Wz direction. This means that an update on the direction wy(k)x' will occur only to the extent that Wæ is not already aligned with wy (k). Thus, if an association is "stored" once, so that Wr x z wy (y), layer-norm will avoid further updating W in that direction, hopefully avoiding norms that grow too much, and also encouraging frequent and infrequent tokens to be weighted similarly in the final memory (see also ? for more discussion on this).
Note that at random initialization Wa is nearly orthogonal to any wy (k),,), so that layer-norm only starts playing significant role later in training, and does not affect our theoretical analysis based on single gradient steps.
MLP blocks. If we denote by (us), and (w;); collections of near-orthonormal input and output embeddings, an MLP block may encode associations (i,5) E . as follows:
F(x)= = Σ vyo(", -o-bb. b).
(KJ)MM
where a is non-linear activation and bias term. Then, one assumes that uy dj S for i * and (t) = 0 for é < O, then this can help filter out noise that arises from near-orthogonality, and may then lead to additional storage capacity, at the cost of additional computation (see, egg., 30)..
An additional benefit of MLP layers discussed in Section is that they may encode many-to-many associations, which is useful when multiple embeddings are in the residual stream and need to be considered jointly (e.g., subject and relation in factual recall task (35]). This may be achieved, for instance, by considering embeddings uy
V LiEz U, where I are sets of bounded size (e.g..
as obtained using layer-norm over a residual stream). Then, assuming the u, are -rthoyoorthonoomal, we have ut Wz 1 while "z Wr 1 8 if I I', for some 5 that depends on the maximal cardinality of these sets. Although 5 is no-longer vanishingly small, defining
F(x)= = Σ vyolufa- x (I,J)EM
the non-linearity may still succeed at filtering out any I that does not correspond to the query set in æ. We leave the question of how such non-linear associative memories may arise from training dynamics to future work.
Multiple heads and layers. We remark that our view of weights associative memories applies to any parameter other than embedding//unembedding layers, and thus naturally extends to multiple heads (using the low-rank factorizations described above) and multi-layer models.
is important to note, however, that the redundancy introduced by having more heads and layers makes more challenging to identify which syyer/head/weihht will learn certain associations (see, egg. Figure 11 in Appendix F This is in contrast to our simplified architecture of Section 4.2, where may identify the role of each matrix (up some possible redundancy when using the feed-forward layer w).. ). In practice, mechanisms may appear in different heads/layers across different training runs, which makes interpretability more challenging, and typically requires some causal identification techniques, such as mediation analysis [35] 54].
E Experiment Details and Additional Experiments
In this section, we present additional details on the experiments, as well as additional results.
Computing setup. We use Pytorch and each run uses single GPU, along with 60 CPU cores for real-time data generation. We will make our code available upon publication.
Hyperparameters. We now provide the hyperparameters used in each figure. The SGD step-size denoted 7. We fix the momentum parameter to 0.9 and the weight decay parameter to 10-4 U denotes the uniform distribution over [N].
26

---

Wk2
Wkl (tx64)
0.7
0.8
memory rocall
0.8
L. ftx-q frequert K-5, flx rore
0.6
memory
0.6-
memery
0.3
Ke fix. frequant
0.4
random
K-1., fix frequent
0.2
K-5, fix
K-5, fix rare Km1, random random
0.1
random
5 random
0.0
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
Deration
iteration
ieeratoon
Figure 5: Memory recall probes for the setting of Figure 4 left).
(tx64)
train unif
frain bigrum
merrory rocall
0.5
memary racall
0.6
0.4
merrary
0.3
0.4
0.4
0.2
train unif
train unf bigram
100
200
300
400
500
100
200
300
400
500
0
100
200
300
400
500
Reration
iteration
iteration
Figure 6: Memory recall probes for the setting of Figure center).
Figure K 3, Fg = Tu (random triggers) or Q is the K most likely elements of Tu. To = U, d 128, dhiddenn 4 x 128 (hidden dimension of the foed-forward MLPs), 0.2.
Figure 3
K 5, ;, n. (random triggers), To U, d 128, 17 0.2.
Figure 4 left) and Figure n., U, d 128, r For random triggers
n. Tu. For K 1 with fixed frequent trigger, the only trigger is the most probable token according to Ta while for K 5 with fixed rare triggers, the five triggers the 6-th to 10-th most probable
tokens according to TT.
Figure 4 center): K = 3, Ty. = nu. (random triggers), U or To (conditioned on the trigger), d 128, 17 1.
Figure rrihht:: 3, Fg = uw (random triggers), n. = U, d 128, 7 1.
Memory recall probes and data-distributional properties. Figure 5 and Figure 6 show the evolution of the different memory probes for the settings considered in Figure (left,center) Figure 5 highlights that associative memories for the induction head are slower to learn when using few triggers (small K), rare fixed triggers, random triggers (note that the probe for W with fixed triggers only shows recall accuracy on the set of triggers Q which is an easier task). Figure
F
20
Figure 7: Visualization of the weights W (left), W (center), and W3 (right) after training with random triggers, K = 3, म. Ta. To = U. For each of these weight matrices W, if we write the corresponding target memory in 7 as W. = Σ, "zug with appropriate embeddings (us); and (v:); (for instance u, Pt-1 and D = Pt for W/ on the left), and we show all values u W uz.
27

---

Wo2 recall, fixed Ir
Wo2 recall, tuned Ir
1.0
1.0
0.8
0.5
64
0.6
129
128
256
0.2
1024
0.2
1024
0.0
50
100
1.50
200
250
20
40
60
iteration
iteration
Figure 8: Effect of dimension on learning W3 alone, with fixed or tuned learning rate.
one-step classitter
one-step classifier
1.0
batches
0.8
0.8
10 batches
accuracy
0.6
64
128
0.6
0.4
256
0.2
d048
0.0
20 betches
107
dimension
Figure 9: Accuracy of one-step estimate of W3 with varying dimension and number of batches used for computing expectations. Each batch consists of 32 sequences of 256 tokens for a total of 8 192 tokens, with K = 5 random triggers and uniform outputs.
shows that using uniform output tokens can lead to better fitting of W and W compared to using output tokens sampled from bigrams. In addition the increased diversity when using uniform outputs, this may also be due to the fact that bigram outputs are already well predicted using global statisties with just the feed-forward layer, hence the gradient signal on such well -ree-perdd tokens may not propagate through the induction head mechanism. In contrast, the recall accuracy for W K is comparable for both settings, since the previous token head useful at all positions regardless of the output token distribution.
Visualizing memories. Figure shows visualizations of the associative memory behaviors after training. We see that diagonal elements dominate in the plots, which corresponds to correct asso- ciations lead to high 'memory recall'. Nonetheless, we see that some of the diagonal elements are weaker than others, particularly for late positions in k. and for some of the trigger tokens in W while the diagonal for W3 seems to be roughly uniform. We note that characters corresponding to capital letters have token index 13 38, while lowercase letters have index 39 64. The association patterns found in W then seem related to frequencies of appearance of triggers, whereby capital letters appear less frequently in the data, and are also less frequently chosen as triggers, compared to lowercase letters. Similarly, since the first occurrence of triggers is typically early in sequence, it natural that W learns stronger associations at earlier positions. In contrast, diagonal elements for W3 are nearly uniform, which agrees with the fact that output tokens are sampled uniformly in this setup. We refer to the follow-up work 77) for an analysis of how data frequencies affect association strength in such associative memories.
Effect of dimension. Recall that our study of associative memories with random embeddings requires large dimension d in order ensure near-orthogonalitty and thus store input-output pairs more effectively. In Figure 8 we evaluate the recall accuracy for W for varying dimension, when training by itself, and only on the output tokens (as in Figure 3). We see that higher dimension leads to faster learning of the memory, in particular d = 128 seems sufficient for fast learning after just few iterations with a tuned learning rate. If the learning rate isn't 't tuned, we notice that there is a further slowdown for low dimension, is likely due to issues with the fact that our experiments use the standard parameterization of neural networks at initialization, rather than maximal update parameterizations (58) Note that learning W3 alone is convex optimization problem, and we hypothesize that higher dimension makes the problem better conditioned, and hence easier to learn.
In Figure we show "one-step" recall accuracies for classifying output tokens from the average
28

---

in-context vs global
in-context accuracy
attention feed-forward probes
K-5 random
K-5, tandom
1.0
Wez recal
W42 recal
2.
accuracy
Wil recal (::54)
0.6
1.6
0.4
0.4
0.2
0.2
0.0
1.0
100
200
300
400
503
100
200
300
400
5.00
200
400
600
Reeratoo
iteration
iteration
Figure 10: Training of a more realistic architecture with (i) ReLU MLP instead of linear layer for the second feed-forward layer, (ii) all parameters trained, including embeddings, (iii) pre-layer normalization. The loss, in-context accuracy and memory recall probes are similar to the simplified architecture (see, e.g., Figure
wwwThee_merrec...ee K: see
<BTThee _eer F_me
Figure 11: Attention maps for a two-layer model with 4 attention heads. In the first layer (top), the previous token mechanism is mostly achieved by one of the four heads, while the induction behavior at the second layer (bottom) is distributed across the different heads.
attention input to W3. given by
arg we (uL - w)]
R
where μlu = E[aly k]] and μ = E(x). for = +. WBwn(Zz) and y = 2e411 when ž is a trigger token after its first occurrence. Expectations are computed over batches of data of varying sizes and in different dimensions. We call this "one-step" since it is related to the classifier obtained after performing single gradient step W from initialization (see Lemma 2 and Appendix BB..1].. The plots illustrate that this simple one-step model is already able to extract relevant signal from the noisy average attention, after handful of batches of data, corresponding to tens of thousands of tokens, and that this gets easier as the dimension increases.
More complex architectures. Figure shows training behavior for a more complex model than the simplified one considered in Section namely where we train parameters, replace linear Wr feedforward layer by two-layer MI.P, and were (pre-)layer- -layer-normalizztion is added. Despite these changes, we see similar behavior for the memory recall probes (which now involve embeddings that may change over time), suggesting that the model is still identifying the same memory associations, despite the additional redundancies in parameters and modified training dynamics.
Figure I1 shows the attention maps obtained when training a multi-hear version of our two-layer model, with four attention heads per layer. We see that the redundancy of multiple heads creates difficulties in identifiability: only one of the first layer heads leams the previous token behavior, while the induction behavior is shared across different heads at the second layer. This illustrates challenges of interpretability in the presence of redundant models, which then require additional work identify which of the layers and heads are performing given behavior, egg., through interventions and causal mediation analysis 35 54].
29


Reviews:
[{"review_id": "vANhn57Hqd", "reviewer": "Reviewer_hu6p", "summary": "Authors perform an in-depth study of the toy case of learning associate recall task using causal Transformers with the aim to understand the emergence of in-context learning abilities during training. Informally, they propose a modified bigram distribution where after sampling a sequence, for a special set $Q$ of pre-determined \"trigger\" tokens, at every occurance of a token $q$ from $Q$, they replace its following token with the token that followed $q$ at its first occurrence in the sequence. I.e. $[...q\\ r ....q\\ s] \\mapsto [...q\\ r ....q\\ r]$. Hence, for $q$, the model is required to determine $r$ first from the context and, at every following occurance of $q$ output $r$. The authors prove that a simplified 2-layer Transformer (the only non-linearity is softmax) can learn this task via gradient descent and verify this empirically.", "questions": ".", "limitations": "There is no discussion of the limitations.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The problem is well-motivated as large language models demonstrate impressive in-context learning abilities and hence it is beneficial to understand how this ability develops.\n\n2. The theoretical statements are non-trivial, interesting and not straightforward to prove.\n", "weaknesses": "1. The associative recall task has been studied in the context of in-context learning of attention-based models before and its not fully clear what novel contributions the authors have made.\n\n2. Given the simplicity of the toy task/model the practical implications and the generality of the findings are unclear (other tasks / architectures).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thoughtful review.\n\n*“The associative recall task has been studied… novel contributions”.*\n\nIndeed, several works have looked at similar tasks. Nevertheless, to our knowledge we are the first to have a precise picture of (pre-)training dynamics and as a result, a precise understanding of the form of the weights, in a way that extends to multiple layers. The structure of our bigram data model we introduce is also simple enough that it is amenable to theoretical analysis. Overall, we believe that our study provides many new insights about the internal structure of transformers during pre-training, and can pave the way for new improvements, e.g. for optimization algorithms, data selection, fine-tuning, model editing, and interpretability.\n\n*“…practical implications and the generality of the findings are unclear...”*\n\nIndeed, this work is a first step. That said, our associative memory viewpoint for weights naturally extends to more complex architectures since gradients will take similar forms, and we chose a simple task precisely to make the viewpoint more transparent and understandable. We also simplified the architecture in order to make the role of each component identifiable, while for a more complex model, there may be many different solutions (e.g. the previous token head could happen in many different layers and attention heads).\n\nOther tasks are definitely interesting but might involve very different mechanisms which we believe will require separate studies, yet, copying mechanisms like ours are likely crucial in many different tasks (as the induction head papers illustrate), and enable more complex reasoning operations, e.g. with a semantic hierarchy that may be learned in multiple layers.\n\nWe expand on this point much more in the general response to all reviewers, and there we discuss additional experiments/theory that shows our viewpoint extends to more complex settings.\n\nPlease let us know if this response helps with your assessment of the paper, we’d be happy to answer any additional questions."}, {"review_id": "4HwlwVhykA", "reviewer": "Reviewer_uvp3", "summary": "This paper studies the dynamics of how induction heads emerge in LLMs during training. The authors describe a simple synthetic task to test their hypotheses on, outline a plausible implementation of an induction head based on associative memory, and describe both empirical and theoretical observations.", "questions": "* Did you check that the actual weights learned by your toy model match the solution you made in any meaningful way? Maybe exactly, modulo a simple transformation or something? Did you find any interesting surprises?\n* Did you have a chance to test any hypotheses about superposition experimentally, even preliminary results? If so, what did you find? Any unexpected things? Negative results ok! I'd imagine it's quite relevant, and is a better representation of what really happens in transformers.", "limitations": "n/a", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "* The paper is well-written and generally easy to follow. Motivations, conclusions, and analysis are crystal clear.\n* The authors study an interesting question: what mechanism do LLMs implement induction heads with, and how might it emerge during training?\n* The methodology is clean and nice! I like the inclusion of \"memory recall probes,\" as they are a direct measure of what you're claiming. The results are also appropriately framed in the context of the merits & limitations of the experimental setup.", "weaknesses": "* Can you better explain the significance of the results? My interpretation is that you framed the $W$ matrices as associative memories and found an order in which they seem to be learned. I'd love to hear what you think the \"so what\" of this is! What interesting things can we do with this understanding? What more can we learn? Can we make the induction heads better? Implement them manually? Do you expect the learnings here to help us understand larger, more complex models?\n* It's sort of implied that the framing of induction heads as associative memory is new. What are the other competing mental frameworks? How does this one compare?\n* It'd be interesting to see experiments that vary the dimensionality of the vectors/matrices. In my experience, superposition is rampant, and LLMs often try to stuff a lot into a not-so-large $d$. How do you expect the findings to change when you can't assume near-orthonormality? Other than a degradation in performance?\n\nMinor things:\n* Figure 1 is nice! It'd suggest labeling the words in the caption with colors or letters, so it's immediately clear what refers to what.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the encouraging review and helpful suggestions. We are glad you found our work interesting, well written, and that you liked our methodology!\n\n*“Can you better explain the significance of the results? …”*\n\nThanks for asking this, it is definitely something we should have discussed more in the submission -- we will comment on it more in the next version. At its core, we hope our work can provide a new language to reason about the internals of transformers, and how they are affected by learning dynamics. Here are examples of areas where it could be useful:\n* architecture/algorithms: improve optimization algorithms and architectures/initialization by studying how they impact the form of the learned memories\n* choice of pre-training data: induction heads can be learned more quickly when [trigger, output] pairs are more frequent and more diverse -- perhaps this could extend to more general “reasoning” problems: does synthetic and diverse logic data help? is this why training on code seems useful?\n* interpretability: presumably, knowing the structure of weight matrices can allow a more fine-grained understanding of what each transformer block is doing\n* fine-tuning vs model editing: if gradient updates are just adding/reweighting outer products to our memories, can we do fine-tuning in a more controlled and targeted way by manually changing appropriate weights?\n\n*“Do you expect the learnings here to help us understand larger, more complex models?”*\n\nYes! The nice thing about this model is that it just assumes that the inputs and outputs of a weight matrix are embeddings (or superpositions thereof), which is basically true anywhere in a transformer other than for embedding/unembedding layers. The main difficulty is redundancy: many different parts of a large model may implement similar mechanisms, so that they may become harder to pin-point and identify during training. This is what motivated us to simplify the architecture into an identifiable model, and similar modifications may work on more complex tasks, but we expect that more general problems will require more work on the interpretability/identification side of things.\n\nSee our general response for more discussion on this, and additional results in this direction.\n\n*“It's sort of implied that the framing of induction heads as associative memory is new…”*\n\nThanks for the question. A perhaps more natural model, which we were initially inclined to believe when we started this work, would be that weight matrices found by training were a bit more “idealized”: for instance, copying could be done with an identity value matrix at the second layer. While recovering matrices like the identity could happen in some settings, we found that in general it is difficult to obtain. One thing we tried in order to see why this intuition fails is to use two different sets of output tokens at training vs test time. Even with tied input/output embeddings (so that the identity approach might actually work on the test tokens), the accuracy on test output tokens is much worse than on train tokens, particularly when the dimension is large [in low dimension, the outer products could span the entire space and actually give something close to the identity, but the overall accuracies are then lower since the embeddings are “less” orthogonal]. We’d be happy to include this discussion in the appendix.\n\n*“…experiments that vary the dimensionality of the vectors/matrices…”*\n\nFigure 7 in the current appendix shows some basic results in this direction, and we’ll include more experiments that study the effect of d and sample size for the one-gradient-step scenario. In practice, we see that higher d definitely helps store things “more quickly” (in terms of iterations and amount of data), but even d=64 is sufficient for our setup.\n\nRegarding superposition, you’re right that there are typically lots of things in the residual stream, but we would argue that d is quite large in LLMs (at least in the thousands), and you can have lots of near-orthogonal embeddings (in fact, exponential in d). The dimension of attention heads is usually much smaller, but we expect that the Q/K matrices filter out a lot of irrelevant embeddings from the residual streams (using a similar process to what we describe in sec. 6), so that the low-d vectors contain just one or a handful of directions.\n\n*“Figure 1 is nice! It'd suggest labeling the words”*\n\nThanks for the suggestion! We will change this in the final version if the paper is accepted.\n\n*“Did you check that the actual weights…”*\n\nBased on your suggestion, we looked a more granular maps illustrating $(v_j^\\top M u_i)_{ij}$ for a desired memory M, instead of just the recall metrics. These do indeed show that the desired associations have much larger values than the remaining items, though we also notice that different input tokens may lead to different magnitudes, especially for $W_K^2$, reflecting different frequencies of triggers in the data.\n\n*“Did you have a chance to test any hypotheses about superposition..”*\n\nThanks for the suggestion. In our controlled setup where embeddings are just nearly-orthonormal random vectors (including \"remapped\" random vectors), it is quite easy to check which embeddings are present in a given representation/superposition, by just taking the inner product with each such embedding. Of course this detection becomes more difficult/noisy when there are many elements (e.g. for the initial average attention in a long sequence), but we found that it works reliably for a handful of elements in large enough dimension. We'd be happy to include this in the appendix.\n\nAgain, many thanks for your comments and suggestions. We hope these clarifications and improvements may help increase your score. Please let us know if you have any more questions or concerns."}, {"review_id": "zW0wBS1LRA", "reviewer": "Reviewer_hHJQ", "summary": "Given the blackbox that large language models are, this paper tries to use a small simplified view of a 2 layer transformer model, and a synthetic task to understand how global and in-context language statistics are learned by the transformer model.\n\nBy freezing specific layers, The authors show how the memory recall and in-context accuracy varies.", "questions": "- In the multi-head attention, was any experiment perfomed to understand whether for a fixed model dimension, is it better to have larger number of heads or smaller number and how that affects memorization?\n\n- Was there any experiment performed where we just have attention (no feed forward layer), and how that affects global memorization?", "limitations": "Discussions could mention that this analysis may not extend to larger versions of the model, or provide empirical proof that it does.", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper presents a nice setup with the simplified transformer model, and the synthetic task for analysis of transformer architecture. The current analysis are great, but just the start and the foundation laid could be useful for future explorations.\n\nThe proposal around how the global memory and the in-context induction work is supported (partially) through the experiments by freezing various layers.\n\nWith the experiments, they are able to show that the global bigram statistics are learned faster and the induction heads for in-context bigrams are developed in subsequent training steps. Additionally, the conclusion that having diverse training distribution leads to better generalization is great. So this is a good started to understanding the transformer LLM blackboxes, and the results may apply on not just text, but other modalities as well.", "weaknesses": "- The results are still very preliminary. The transformer model is simplified, and so is the dataset. It is unclear whether a larger transformer model would use a similar paradigm for learning. Essentially, it might be possible that the results do not extend to larger models and additional experiment needs to be done to verify that.\n\n- Transformer models have a lot of other components, and it would be nice to study how they impact the memory of these models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the helpful review and for the positive assessment of our work.\n\n*“The results are still very preliminary…” \"Transformer models have a lot of other components...\"*\n\nThis is the main topic of our general response post, which we hope provides a valid justification of our approach, as well as as new results in this direction.\n\nIn summary, we chose a simple setup in order to provide a more accessible and intuitive description of the weight matrices and their learning dynamics. That said, the “memory” viewpoint for the weights, as well as the insights on training dynamics, extend naturally to more complex architectures, since they are basically a consequence of the fact that inputs and outputs of a weight matrix are embeddings. We will include new experiments and theory justifying this, as presented in the general response. The main difficulty with more complex tasks and architectures is that different mechanisms can be implemented in different layers/heads, making them harder to identify and track. Tackling this in well-chosen or general problems is an important future direction.\n\n\n*\"In the multi-head attention...\"*\n\nWe did not experiment extensively with multi-head attention, since a single head at each layer is sufficient for the induction head mechanism needed for our task. Nevertheless, we ran basic experiments using multiple heads on our task, and observed that the associations made by the second-layer attention tends to distribute across different heads (i.e. different tokens are handled by different heads), while the previous-token mechanism at the first layer seemed to use only one of the heads.\n\n*\"Was there any experiment performed where we just have attention...\"*\n\nThanks for the suggestion. We tried an attention-only model, and noticed that the loss on global tokens decreases more slowly than with the feed-forward layer, and ends up at a higher loss number (but note that training embedding layers could circumvent this). For fixed triggers, the task becomes a bit easier and looking at attention maps for the second layer shows that the attention usually focuses on the current token, except for the fixed triggers which need the induction mechanism.\n\nThanks again for your encouraging review. Please let us know if you have any other questions of concerns."}, {"review_id": "Grc4CiX8zR", "reviewer": "Reviewer_PfRr", "summary": "This paper provides a detailed analysis of how in-context learning behavior emerges in a simplified version of the transformer architecture on a toy task. This work can provide important insight into how in-context learning emerges in LLMs. The toy task is this: given a sequence of tokens of the form $\\ldots, a, b, \\ldots, a$, predict $b$. The token $a$ is a \"trigger\" token. If token $b$ (the \"output\" token) comes after the first occurrence of $a$, then the model must predict $b$ after the second occurrence of $a$. The model is expected to use in-context learning to predict the correct $b$. The rest of the sequence follows a bigram language model distribution. The authors zoom in on two aspects of the model: the \"induction head\" mechanism that implements in-context learning, and the \"associative memory\" mechanism in the key/query weight matrices that allows the model to learn the global bigram distribution. They give mathematical justification for the emergence of both mechanisms in their simplified transformer model trained on infinite data. They provide experimental results that provide evidence that these mechanisms are learned in a particular order, and that learnability depends on attributes of the training data distribution, such as the number of trigger tokens per sequence, or whether the trigger token types are randomized.", "questions": "Most important questions:\n\n1. 77, 90, 93, 97: Does the decision to ignore layer normalization make the findings of this paper less applicable to real transformers and LLMs? Same question for using only a single head. What about dropout? Same question about using linear transformations instead of feed-forward networks.\n1. Fig 1: I had a very hard time understanding this figure, and it's essential for understanding the paper. I don't know what Attn1 and Attn2 are supposed to signify. In Layer 1 column 2, $w_1(a)$ and $w_E(b)$ aren't the only values in superposition there, right? It would also have $p_t$. It's not possible to completely isolate $w_E(b)$ from $p_t$, because they are in superposition in Layer 0, right? I think it would really help to give a symbolic representation of the query, key, and value functions at each layer. If I understand correctly, the inputs at Layer 1 are essentially $(x_t, t)$, where $x_t$ is the input token at timestep $t$. Then $\\mathrm{query}(x_t, t) = t-1$ (assuming you can compute this with a linear layer), $\\mathrm{key}(x_t, t) = t$, and $\\mathrm{value}(x_t, t) = x_t$. This adds $x_{t-1}$ to the inputs of Layer 2. Then in Layer 2, $\\mathrm{query}(x_t, t, x_{t-1}) = x_t$, $\\mathrm{key}(x_t, t, x_{t-1}) = x_{t-1}$, and $\\mathrm{value}(x_t, t, x_{t-1}) = x_t$. Then if $x_T = x_{t-1}$, this lets you predict $x_t$. Is this right? How do you ensure that this only applies when $x_T = a$, and not other token types in the vocabulary?\n1. Lemma 1: Could you include a proof sketch in the main text? Could you include a longer discussion of the limitations of your simplifying assumptions in Lemma 1 and how they apply to real transformers? The main problem with this section is that I cannot verify that Lemma 1 is true, and that it applies to real transformers in the way you claim that it does. Won't it totally change the dynamics if the other weights are not fixed? I think it's probably a useful result either way, but you should talk about this more.\n1. 197: Could you justify this decision a bit more? How do we know this doesn't drastically change the situation from real transformers?\n1. 308: Can you include a proof sketch for Lemma 2?\n\nOther questions:\n\n1. Can you clarify which key-value associations you expect the linear layers to learn? Is it the global bigram distribution?\n1. 51: If all the examples are generated by a single bigram language model, isn't all of that \"global\" knowledge? That is, how will you know if the model isn't just memorizing the training data? Are you counting on the fact that the transformer isn't big enough to memorize so many examples? Do you test on out-of-distribution examples that verify that this is the case?\n1. 98: What about tied embeddings?\n1. 199: I think this assumption is pretty reasonable. But I think it's harder to understand how the induction head would work without a separate query matrix. Could you spend a little time explaining how this would affect the construction?\n1. In general it would be very helpful to explain in detail how to design an induction head under these constraints.\n1. 208: Again, it would be useful to discuss how W_F would implement this.\n1. 219: How?\n1. 220: Where do $p_t$ and $p_{t-1}$ come from? Are they sinusoidal positional encodings?\n1. Eq 9: Can you explain in words what this is doing?\n1. Fig 5: Why is it easier to learn 5 random q than 1 random q? It would be helpful to be reminded that K is the number of triggers per sequence, not the number of token types used as triggers overall.\n1. 310: What is $E[x|y = k]$? What is $\\hat{p}_W$?\n\nClarity issues:\n\n1. 43: At this point, I was wondering if you would supplement your empirical analysis of the training dynamics with a mathematical explanation. It would be useful to point out earlier that you will do this.\n1. Eq 1: Having $x_t$ on both sides of the equation is confusing. I'm not a big fan of the $:=$ notation.\n1. How would the induction head mechanism work in the presence of multiple instances of a in the past context?\n1. 120: These sentences seem out of order and are hard to read. It's not clear if $b$ is a variable or just a tag to distinguish $\\pi_b$. I don't understand what $\\pi_o$ is. I was confused by the difference between $\\pi_b$ and $\\pi_o$.\n1. 123: What is the variable $n$?\n1. 124: How is $K$ chosen? Is it randomly sampled or set to a constant value for the experiments?\n1. Using i and j as variables for token types instead of input positions is confusing, especially since k is used as an index\n1. 126: What is the first token in the sequence? How is that sampled?\n1. 126: It would be helpful to provide an intuitive explanation in words of the process that this equation implements.\n1. Where is $\\pi_u$ used? (I see it is used at 131, but this seems out of order.)\n1. 129: What is the tiny Shakespeare dataset?\n1. 133: Referring to training details in Section 5 out of order makes this part hard to read. Could you move the training details earlier?\n1. 134: Is this on the training data, or on a held-out test set? Do you use a validation set?\n1. 134: It would be easier to read this results if they were in a table.\n1. Fig 2: The labels for the axes on the first plot are missing. It would be helpful to label the y axes. For the left plot, which version of the dataset was this trained on? Is there a reason you can't show the first layer for both models? I don't understand the significance of the red and green boxes. What are the \"previous occurrences of the query\"? It would help to highlight the triggers and outputs in the axis labels.\n1. 134: At this point I was confused about the significance of using fixed vs. random triggers. But after re-reading 124-126, I understand that for random triggers, the set of triggers is randomly sampled for every sequence, so *every* token type has the potential to be a trigger. The only way they are distinguished as triggers in the training data is that the same bigram appears twice in the same sequence. But this could also happen by chance for non-trigger tokens, right? It would be helpful to add a discussion of this to the main text.\n1. 153: What does this notation mean?\n1. 159: What does O mean?\n1. Eq 5: What does z range over?\n1. I didn't realize until Eq 6 that $w_E(z)$ means the $z$th column of $E$. Could you define this notation beforehand?\n1. 213: I see that this section answers some of my questions above. It would be useful perhaps to introduce or mention this earlier in some way.\n1. 238: Just to clarify, you're using a cross-entropy language modeling loss function, right? I see this is partly answered at 254.\n1. 247: In the equation for $W_{*}$, what are $v_j$ and $u_i$?\n1. Eq 9: What is $M$? How are the $(i, j)$ picked?\n1. 307: This seems to answer one of my earlier questions.", "limitations": "As mentioned above, I think the authors should discuss the limitations of their simplifying assumptions about the transformer architecture more.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This paper provides a very useful case study of how in-context learning behavior emerges in transformers. In-context learning is a hot topic, and it is important to understand architecturally how this behavior emerges, as well as the conditions that are conducive to it. This paper is a good step towards better understanding this phenomenon. I think their bigram task is an appropriate choice of case study. The authors analyze two mechanisms and provide mathematical and empirical justification for both.\n\nOriginality: Good. I think the choice of task is very useful for the purposes of studying in-context learning.\n\nQuality: The mathematical analysis and experimental design appear to be sound.\n\nClarity: The paper seems to be well-contextualized with respect to previous work.\n\nSignificance: High. After reading this paper, I feel I have a better grasp on how in-context learning works in the transformer architecture, and this has important implications for any NLP applications that use transformers.", "weaknesses": "**Edit:** I have read the rebuttal, and it addressed my biggest concerns.\n\nI have two major criticisms of this paper, which is why I did not immediately assign it a higher score:\n\n1. The transformer architecture used in this paper is drastically simplified from real transformers used in LLMs. There is no layer normalization or dropout, and the transformer only has 2 layers. They use linear layers instead of feedforward layers. The key and value matrices are merged into one matrix. Many of the parameters are frozen during experiments and during gradient analysis. I think the paper requires a much more detailed, readable discussion justifying why the authors expect their analysis on this simplified transformer model to carry over to real transformers. I think there is some discussion scattered throughout the paper, but if so, I think it needs to be organized more clearly. I think the results are useful regardless of these simplifying assumptions, but the authors need to talk about this more.\n2. Clarity. This was a very difficult paper to read, and this negatively impacted my ability to interpret the results. As I explain in more detail in the Questions section, a recurring issue I had while reading this paper was that information is frequently presented in *reverse* order; I often needed to read several lines, paragraphs, or sections ahead in order to clarify something I was confused about. I think the whole paper would benefit from a round of editing that alleviates these issues. See the Questions section for more details.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the very detailed and encouraging review.\n\n*\"The transformer architecture used in this paper is drastically simplified ... the paper requires a much more detailed, readable discussion ...\"*\n\nThanks for raising this point. We hope that our general response to all reviewers provides useful elements to address this concern. In particular, we view simplicity as beneficial to illustrate our model. We also show that various aspects of our model can easily be extended to more complex architectures, although multi-head attention and multiple layers may lead to identifiability issues which make a clean analysis more challenging.\n\nWe will do our best to provide a clear discussion of these points in the paper, and will add an appendix with additional experiments and technicalities for extending to more standard architectures.\n\n*\"clarity\"*\n\nWe will do our best to improve clarity and address the points you raise in the questions.\n\n*layer-norm, multi-head, MLPs*\n\nPlease see our general response, which addresses these. We will include results on these in the appendix.\n\n*dropout*\n\nWe did not consider dropout as it is often not used in modern LLMs (see, e.g., the PaLM paper), but leave its study as an interesting question for future work.\n\n\n*Figure 1*\n\nWe apologize for the lack of clarity with this Figure, and will improve this in the updated version. Attn1 and Att2 refer to the desired associative memories for the product $W_Q^\\top W_K$ (or just $W_K$ in the simplified model with $W_Q = I$) at each layer.\nYou should think of key-query associations as considering two tokens at different positions (instead of separately at each position as you write) and checking if there is a match. E.g. for the first layer, we want $(x_t + p_t)^\\top W_Q^\\top W_K (x_s + p_s) \\approx 1$ for $s = t-1$ and zero otherwise, which is the case for the given choice of associative memory (see also eq.7).\nWe will do our best to better represent the full superpositions at each layer, instead of just two elements.\n\n\n*Lemma 1*\n\nWe will add a brief proof sketch and add a pointer to the right appendix with the proof.\nThe form of the gradient is valid regardless of the current values of the input/output embeddings, so it could potentially apply at any time during training. In particular, if somehow the embeddings stop moving because they converged to a good place, then the weight matrix will end up learning an associative memory with these new embeddings. In practice (see above) we saw that associations are still correct when training embeddings. This could mean the embeddings converge quickly to something stable, and the weight matrices use those final directions to store associations.\n\n*197 (why freeze layers)*\n\nSee our general response. Basically our motivation was to get the simplest possible model for which we know what the solution should look like (\"identifiable\"), in order to have a clean and concise description of training dynamics, both empirically and theoretically. Also, if a good \"lazy\" solution exists where some of the weights don't need to move, it will likely be an easier solution for GD to find! For attention matrices, this still works in our theory as well, as mentioned in the general response, and we will include these results in the appendix.\n\n*\"..key-value associations you expect the linear layers..\"*\n\nIndeed, we expect $W_F$ to learn the global bigrams (see eq.8 and the comments below it), at least initially. This would be the case approximately if the two attention layers were not present, at least for non-trigger tokens (following Lemma 1). Our hypothesis for why this also holds initially with attention layers is that all elements in the input superposition to $W_F$ that are not the current token will be mostly noise, since all previous tokens are independent of the next token given the current token (assuming we're on a non-trigger token), so they will get filtered out in the update with enough data (as in Lemma 2).\n\n*\"global knowledge\"*\n\nWhile it is possible that a transformer with very large MLP layers could memorize all sequences in order to guess the correct output token after a trigger, this would likely be intractable since the number of such possible sequences grows exponentially with sequence length. Yet our model gets near-perfect prediction accuracy on tokens even near the end of the sequence, suggesting it uses a better approach. The empirical study also shows that the model is making use of the induction mechanism for predictions. We also tried replacing output tokens at test time, and saw correct predictions (though this generally only works for outputs seen during training, in fact if a token isn't used as an output, we see that the corresponding outer product simply isn't stored in $W_O^2$, see our response to uvp3 for more on this).\n\n*other questions*\n* tied embeddings: these also work in our formalism.\n* L199: it is only the product $W_Q^\\top W_K$ that affects predictions, which is why we preserve the same expressivity by fixing one to the identity.\n* 208: this is shown in eq. (8)\n* L219 \"how?\": thanks to near-orthogonality, any embedding in the superposition that does not appear in the outer products gets filtered out\n* L220: $p_t = w_P(t)$ by definition, and is fixed at random initialization\n* eq 9, fig 5: we will clarify these. $K$ is indeed the total number of tokens used as triggers (see L124). Larger $K$ makes it easier to learn the induction head because there is a lot more data involving trigger-output pairs. The values of K used in experiments, if not shown, are given in Appendix D.\n* L310: we will clarify ($E[x|y=k]$ is the conditional expectation of $x$ given $y = k$, and $\\hat p_W$ are softmax model predictions for parameter $W$ respectively)\n\nThank you for the many other comments and suggestions regarding clarity, we will happily incorporate them and will do our best to improve clarity in the updated version of the paper, if it is accepted."}, {"review_id": "3jANmfTEzO", "reviewer": "Reviewer_k5yo", "summary": "The paper studies in detail two-layer transformers and extend the setting of pure associative recall based on data in-context with mixing these tasks with tasks coming from a global birgram model. They then study these two-layer Transformers by freezing layers and probing. ", "questions": "In general I think the papers focus on the tension between memory and in-context learning is great. \nNevertheless, I think that a couple of very interesting experiments, ablations and analyses about design decisions are missing. \n\n1) Can you provide analyses in both extremes i.e. when there are only \"triggers\" i.e. when all sequences can be learned by quick associative learning and the other way around. This should lead to quite different circuits that can be contrasted against each other. \n\n2) Why do you include a single (linear) feedforward layer in the architecture? The problem should be solvable without (as studied in the induction head paper). If you want to include this layer, then I think the right thing to do would be to include it also after the second layer. \n\n3) Why are you so heavily relying on the freezing of some layers? I actually ran your experiments (with only triggers) and I believe that my matrices do actually show very different behavior. I think that the freezing of the layers, forces you to focus on the superposition too much. If you wouldn't freeze your vocabulary and the input matrix, the model could (and I see this in my experiments) actually produce embeddings which leave token memory free to use for layer computations i.e. essentially concatenate x_i = w_E(z_i) + w_P(i) \\approx [z_i, p_i] since the positional encoding is basically alternating 0 and 1 after some dimension.  By that, the circuits and mechanism that you find could be very different and therefore not stable wrt. to design decisions. \n\n4) Building on 3): I am quite convinced that all design decisions will change the circuits and mechanism of how the two layers solve your problem. Therefore I think it is very interesting and actually needed to investigate these and make suggestion what these design decisions control. E.g. Add or concat frozen/learnable random/sinusoidal PEs, control the dimension of the vocab and the token dimension, include feedforward layer, etc. \n\n5) The theoretical insights on Learning Dynamics seem a bit vague and less supported by empirical study (didnt look in the appendix). As I understand, in-context learning kicks in later in training, so why study on the random init? At least comment on this shortcoming, if I understood this correctly. \n \nMinor things: \n\n5) Can you please better your probing experiments, how exactly do you do that? Which vectors are probed? \n6) What do you mean with discrete tokens line 81?\n7)  Equation 7 - its confusing for me why exactly these weights recover the induction head mechanism. Please elobatore, I can somehow see it but please make this more rigorous.\n8) Please explain a bit better your data generation and the statistics. How large is the possibility of conflicts i.e. a -> based on global knowledge and -> based on in-context. \n9) Figure 4 left. Why cant the in-context learning loss go to zero? Is it because of the conflicts? Please make the Figures a bit less crowded, hard to see all the lines. Would be great to have scaling plots wrt to K imo. \n\n\n\n  \n\n\n\n", "limitations": "I dont think the limitations of design decisions are well discussed or investigated. See Questions. ", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper studies an interesting problem i.e. in-context learning within Transformers, that is believed to be an important characteristic of language models in general. Based on recent work, they aim to mechanistically interpret / reverse engineer the two layer Transformer which in my opinion is an interesting direction to better understand and get intuition on how Transformers work. \n\nI like the mixed bigram data setup that the authors propose and study and I share the opinion of the authors that the tension between storing knowledge and quick adaptation / learning from in-context is very interesting and not explored. Progress in this direction is important. ", "weaknesses": "It feels that the paper is very rushed. I think therefore that the presentation can be vastly improved as well as results and analyses refined. \nFurthermore, I think that the proposed method of freezing the vast majority of layers in the network leads to very biased results and therefore results which might not hold when training on all weights. \n\nA couple of design decisions seem also quite arbitrary and it would be I think the job of the paper to justify and investigate them. See Questions. Therefore, I think that the paper is very promising but can and should be improved.  \n\nMinor things / Language:\nKey, Query, Value Matrix must not be square and the positional encodings are usually not randomly initialized. Please at least comment on these design decisions. Also you do not use layer norm, I would also expect at least a comment on this. Please use LaTex in your Figures, quite hard to read the legends etc (maybe different background color as well?). The sum in Figure 2 over the positional encodings should be over t? Please describe more precisely Figure 3&4 in the caption. These are not concise i.e. it was not possible for me to understand what you are plotting what the Figures show.  Also please increase the font size. \n\nAt least one interesting citation is missing: https://arxiv.org/abs/2212.07677\nThe paper studies in details 1) single and two layer Transformer and provide evidence about copying in the first layer and 2) in-context learning in the second layer by gradient descent. This is, I believe, equivalent to the Hebbian-rule / associative recall with orthogonal inputs.  ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review and for the interest to try out our code! We hope our response can provide more perspective about the motivations behind our work, and may help you reconsider your score.\n\n*freezing layers*\n\nSee our general response. In particular, we chose to simplify the architecture as much as possible to have a simpler setup for a clear study. We also found empirically that training all parameters, even when adding layer-norm and MLPs, still leads to correct memory recall associations.\n\n\n*Minor Things*\n\n* \"Key, Query, Value...\": with a single head, V needs to be squared, but we agree that Q/K need not be squared. We will include the case of separate and potentially non-square Q/K matrices, which still can lead to the desired memories for the product $W_Q^\\top W_K$.\n* positional encodings: to our knowledge, when using learned positional encodings, these are randomly initialized.\n* layer-norm: see general response (it still works!)\n* Figures: thanks for your suggestions! We will improve this in the next version\n* sum in Figure 2: we use the variable $s$ instead of $t$ (as used in Eq. 7) in the sum to avoid confusion with the $t$ of the tokens below.\n\n\n*Missing reference*\n\nThanks for pointing out this very relevant paper! We will cite it and compare to it in the next version.\n\nThe mechanism they describe is indeed similar to ours, and to induction heads in general. However, their construction relies on idealized choices of weights such as the identity matrix (see A.1 in their paper), which could easily be quite different from the weights learned by training, particularly in a setup with discrete data like ours. In contrast, our work (i) provides a precise description of what the weights look like with gradient updates, (ii) verifies this form empirically through memory probes, and (iii) shows theoretically that gradient dynamics on our bigram task indeed recovers these associative memories.\n\n\n*\"..analyses in both extremes...\"*\n\nOur model already covers all cases ($K=0$ means no triggers, $K=N$ means all tokens are triggers). For $K=0$, the model just doesn't learn the induction mechanism since the gradients to the weights of the attention blocks are essentially dominated by noise. For $K \\geq 1$, the same induction head mechanism is learned as long as triggers appear in the data (which is why our experiments consider triggers that are frequent tokens). The only difference is that $W_K^2$ in eq.7 only learns associations for trigger tokens that appear in the data. This is described at the end of Section 4.2, but we will clarify it further.\n\n*\"..single (linear) feedforward layer...\"*\n\nYou are right that global bigrams may be learned without feedforward layers (as in the induction head work), however that would require learning embeddings. With fixed random embeddings, it is much harder to encode global bigrams without feed-forward layers (see response to hHJQ).\nWe prefer encoding this in a linear feedforward layer because its training fits the associative memory viewpoint, as for all internal weight matrices, and might thus capture how other feed-forward layers may store global knowledge in larger networks. Input/output embedding layers likely exhibit quite different learning dynamics that go beyond the message of our work. As a side note, even when training all parameters including embeddings, we observed that the KL probe for $W_F$ still decreases quite quickly, which suggests that gradient dynamics may prefer storying global bigrams in the feed-forward layer as opposed to embedding layers.\n\n*\"Why are you so heavily relying on the freezing... I actually ran your experiments ...\"*\n\nThank you for taking the time to run our code!\n\nWe tried training all layers including embeddings, and found that the recall probes actually reach perfect recall accuracy (see general response), even when using only triggers (i.e. using the option `--data_args.k 65`). Thus, we believe the model still learns the same mechanism. \n\nIf however you are using fixed sinusoidal positional encodings, then the first layer might end up learning something different than eq. 7 for $W_K^1$, since you no longer have the near-orthogonality behavior, and in fact you may get a previous-token head behavior without the need to store outer products for all t thanks to the sinusoidal structure (but the other layers still do get behave like eq. 7). This relies on specific embedding structures that depart from our memory viewpoint based on near-orthogonality, which is why we did not include sinusoidal PEs in our paper, but we agree it is an interesting direction!\n\n\n*\"The theoretical insights...\"\n\nThe detailed analysis of how the induction head mechanism emerges is deferred to Appendix B.3 as it is quite technically involved. We instead focused on the key ingredient behind the proofs, which is simpler to state, and is about how gradient updates on a lot of data can filter out irrelevant parts elements of input superpositions.\n\nRegarding \"later in training\": we show that if you learn sequentially $W_O^2$, then $W_K^2$, then $W_K^1$, each with a single gradient step from its initialization and on enough data, then you *do* learn the in-context learning/induction head mechanism (i.e., there is no shortcoming). The key point is \"enough data\" and also what kind of data: learning the global bigrams is typically very fast because there is a lot of such bigrams and the signal is quite \"clean\" and accessible from the current token's residual stream. In contrast, for the induction head, say, to learn $W_O^2$, this requires finding signal among *all tokens* seen so far, which is much harder and requires more data (in addition to the fact that trigger-output pairs are often less frequent overall compared to global bigrams).\nWe'd be happy to include these intuitions in the paper, if it is accepted.\n\nThank you for your other comments and suggestions, we will happily incorporate them in the updated version."}]

Author rebuttal:
We would like to thank all reviewers for their insightful feedback and valuable comments.
We are happy that most reviewers found our work relevant and significant. Indeed, we believe that our insights on the internals of transformers can pave the way for improved methods in several aspects of LLMs, including better optimization algorithms, data selection, fine-tuning, model editing, and interpretability.

We provide a response below to a shared concern regarding our simplified setup, and respond to individual reviewers in separate replies.

**Why a simplified architecture**

The simplified architecture compared to common transformers/LLMs was a concern for multiple reviews.

Our goal was to simplify the model as much as possible to ease the understanding of what is happening, while ensuring that the model is still rich enough to capture the relevant phenomena to solve the task. It would definitely be much more cumbersome to illustrate the memory viewpoint, and theoretically study training dynamics on a model where many more components are trained. We hope the simplicity of our architecture can help provide better intuition for what we believe to be a key internal mechanism in all transformer models.


**More components**

In line with the reviewers' suggestions, we ran additional experiments to check if we obtain similar mechanisms with more trained components.
In particular, we trained a similar two-layer model with the following modifications:
* train all parameters (including input/output embeddings and all four attention matrices at both layers)
* use a ReLU MLP feed-forward layer at the second layer (instead of linear)
* add pre-normalization layers

We found that despite this added complexity, the "memory probes" described in Section 5 still display the same associations empirically. Concretely, we replace $W_K^\ell$ by $W_Q^{\ell\top} W_K^\ell$ in eq.(7), and~$W_F x$ by $MLP(x)$ for the feed-forward probe (note that these can still be defined even if the parameters are changing). The three recall probes still converge to 1, while the KL probe on the feed-forward layer decreases.

This shows that this more realistic model still identifies the same mechanisms empirically.
On the theory side, we can also easily show that gradient steps on each attention matrix from random initialization still recovers similar associations despite the redundancies (in the pairs Q/K and V/O). Additionally, we can show that layer-norm and MLPs preserve the associative memory form of the weights:
* Adding layer-norm essentially adds a projection operation to the rank-one terms in the gradients. The projection essentially drops outer product terms for which the association was already present. In particular, it plays no significant role at random initialization, thus does not change our single-gradient-step analyses, but is likely important to study optimization stability.
* For the MLPs weights, the outer product terms in gradients involve non-linear mappings of superpositions, which may more easily capture interactions between different elements of the superpositions.


We'd be happy to include these additional experiments and technical details in the appendix, if the paper is accepted.

**More heads and layers**

While the experiments above show some robustness to common complexities, the single-head, two-layer architecture remains very simplified compared to large models.
Unfortunately, understanding more complex architectures with multiple heads and more layers is more challenging since there is a lot of redundancy and it is unclear which layer or head may be implementing different mechanisms.

For instance, we tried using 4 heads instead of one in the two-layer model, and the second layer induction mechanism appears to distribute across all the heads, each head taking care of different tokens (though interestingly the first layer previous-token mechanism seems to only use a single head). This is in contrast to the single-head induction head mechanism we describe, which becomes identifiable (i.e. we know what each matrix ends up doing) thanks to the reduced redundancies.

Extending our work to general models and tasks will thus require more work--in the spirit of interpretability--to first identify which parts of the model are implementing different mechanisms, before understanding dynamics. Nonetheless, our results suggest that within each block, gradient dynamics naturally lead to weight matrices that may be interpreted as associative memories.

We hope this post clarifies our motivations for the simple model, and we will do our best to make this discussion clearer in the paper.

Meta-review: This work conducts a careful mechanistic analysis of how in-context learning capabilities emerge during training of a simplified version of the transformer architecture, on a toy task. The authors focus on two complementary mechanisms of (a) in-context learning through induction-heads, and (b) an "associative memory" in the attention weight matrices that captures the global data distribution. They give both mathematical and experimental justifications for the emergence of these mechanisms, and show concrete connections between learnability and attributes of the training data distribution. 

The main criticism raised by the reviewers was about the simplicity of the analyzed architecture compared to standard language models (LMs) and of the evaluation setup compared to more challenging natural tasks. These concerns were addressed during the rebuttal, as the authors extended their experiments to consider more components from the standard transformer architecture, and explained the non-trivial effort required to go beyond this, which seems reasonable to leave for future work.

Overall, this is a rigorous study that provides valuable insights into fundamental questions of how transformer-based LMs obtain ICL capabilities during training and balance between ICL and generation from parametric knowledge.
Decision: Accept (spotlight)