Title: Birth of a Transformer: A Memory Viewpoint
Venue: NeurIPS 2023
Accepted: yes
Keywords: transformers

language models

deep learning theory

interpretability
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

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.

1 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. A better understanding can help with interpretability of how these models make decisions, and will be crucial to improve these models and mitigate their failure cases, such as hallucinations or reasoning errors.

An important ingredient in the success of recent LLMs is 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 to 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 to be stored in the model.

In order to better understand how transformers develop these capabilities during training, we introduce a synthetic dataset that exhibits both aspects. It consists of sequences generated from a bigram language model, but where some of the bigrams are specific to each sequence. Then, the model needs to rely on in-context learning for good prediction on the sequence-specific bigrams, while the global bigrams can be guessed from global statistics conditioned on the current token. While one-layer transformers fail to reliably predict the in-context bigrams, we find that two-layer transformers succeed by developing an induction head mechanism [16, 40], namely a 'circuit' of two attention heads that allows the transformer to predict b from a context [ · · · , a , b , · · · , a ] , and which appears to be ubiquitous in transformer language models [40, 54].

In order to obtain a fine-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 on attention and feed-forward mechanisms, while avoiding the difficulty of learning representations, which may

Hervé Jégou FAIR, Meta

Léon Bottou FAIR, Meta

require complex nonlinear dynamics [15, 33, 45]. 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 to 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 on this top-down learning process through population gradient dynamics. Despite its simplicity, our setup already provides useful insights on the internal structure of transformer language models and its evolution throughout training, paving the way for a better understanding of LLMs. We hope that our insights may lead to future research and improvements for LLM practitioners, e.g. , for optimization algorithms, data pre-processing and selection, interpretability, fine-tuning, and model editing.

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 weight matrices as associative memories that learn to store specific pairs of embeddings, and use this to 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 in 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 [9, 33, 39, 61]. Several works have attempted to understand internal mechanisms in transformers that are responsible for certain behaviors, an area known as 'mechanistic interpretability' [16, 17, 35, 39, 40, 54]. Memory and neural networks have a long history of connections [5, 18, 19, 21, 26, 27, 30, 34, 50, 55, 56]. The associative memories we consider bear similarity to [29, 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 to probing based on memory recall is related to techniques in [13, 17], though motivated differently. [14, 32, 36] study statistical and approximation properties of transformers, highlighting benefits of sparse attention patterns, but do not consider training dynamics. [25, 31, 49, 51] 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 a 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, autoregressive architectures with a causal attention mask, which are commonly used in large language models trained for next-token prediction [6, 11, 41, 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 z 1: T ∈ [ N ] T of length T , where N is the vocabulary size, the transformer operates as follows:

Token embeddings : each discrete token is mapped to a d -dimensional embedding via an embedding map W E ∈ R d × N . We will denote the embeddings of tokens z t by x t := w E ( z t ) , where w E ( j ) is the j -th column of W E .

Positional embeddings : the positional embeddings p t = w P ( t ) ∈ R d are added to each token embedding depending on its position in the sequence, leading to the following input embeddings:

Layer 0

Sequence

Layer 1

Layer 2

p

t

-

1

w

E

(

a

)

p

t

w

E

(

b

)

p

T

-

1

w

E

(

a

)

a

b

a

b

[

· · ·

]

∗

w

E

(

a

)

w

1

(

a

)

w

E

(

b

)

∗

w

E

(

a

)

∗

∗

∗

∗

w

U

(

b

)

w

E

(

a

)

Attn1

:

∑

s

p

s

-

1

p

/latticetop

s

W

1

O

W

1

V

Residual

Attn2

:

∑

k

w

1

(

k

)

w

E

(

k

)

/latticetop

W

2

O

W

2

V

:

∑

k

w

U

(

k

)

w

E

(

k

)

/latticetop

Prediction

Attention blocks : given an input sequence x 1: T ∈ R d × T of embeddings, the causal attention block computes, for W K , W Q , W V , W O ∈ R d × d (key, query, value, output), and for each t ,

where σ takes the softmax of its elements, leading to an attention of the 'values' W V x t with weights proportional to exp(( W K x s ) ⊤ ( W Q x t )) . Note that the attention operation usually considers multiple 'heads' that each projects the input to a lower dimension. Here we stick to a single head for simplicity, since it will be sufficient for our purposes. Rewriting (2) on each t as x ′ 1: T = A ( x 1: T ; W K , W Q , W V , W O ) , the ℓ -th layer of the transformer applies attention with layer-specific parameters along with a residual connection as follows: 2

Feed-forward blocks : feed-forward blocks operate on individual token embeddings after each attention block, typically by applying a one-hidden-layer MLP to each token, denoted F ( · ; W F ) , with a residual connection: at layer ℓ , we have

Unembedding : After the last transformer layer, the embeddings are mapped back to the vocabulary space R N through a linear 'unembedding' layer W U = [ w U (1) , . . . , w U ( N )] ⊤ ∈ R N × d , where we refer to the w U ( j ) as 'output embeddings'. The output of this layer is then fed into a crossentropy loss for predicting of z t +1 from each x t .

Our simplified setup will linear feed-forward layers: F ( x t ; W F ) = W F x t .

We will sometimes refer to the representations x t for a given token t 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 a particular type of mechanism (or 'circuit') in transformers that allows basic in-context prediction of the form [ · · · , a , b , · · · , a ] → b . These were found to be ubiquitous in transformer language models, playing a key role in enabling various forms of in-context learning. The basic mechanism consist of two attention heads in separate layers (see Figure 1 for an illustration): (i) the first is a previous token head which attends to the previous token using positional information and copies its embedding to the next token; (ii) the 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]).

r

s a b

t

s

L a b

t

s

L

,

a b

t

h b n

t

L & C L & C a b

t h

3 Synthetic Setup

In this section, we introduce our synthetic data setup, which allows us 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 trigger tokens denoted q k are modified in each sequence to always be followed by some output tokens o k . Let K be the number of trigger tokens, and fix the following distributions over the vocabulary [ N ] : π b ( ·| i ) , π u , π o ( ·| i ) and π q , for i ∈ [ N ] . π b ( ·| i ) are the global bigram conditionals, π u the global unigram distribution, while π o is used to sample output tokens at each sequence. The triggers are either fixed to some predefined set of tokens Q , or sampled from π q . Each sequence z n 1: T is generated as follows:

(optional) Sample q 1 , . . . , q K ∼ π q , i.i.d. without replacement ( random triggers )

Sample z n 1 ∼ π u and z n t | z n t -1 ∼ p n ( ·| z n t -1 ) for t = 2 , . . . , T , where

Sample o k ∼ π o ( ·| q k ) , i.i.d. with replacement.

✶ Experimental setup and initial experiment. Our experiments take π u and π b to be unigram and bigram character-level distributions estimated from the tiny Shakespeare dataset, with vocabulary size N = 65 . We generally sample triggers from π q = π u or fix them to the K most frequent tokens. We sample uniform outputs o k in most cases, but also experiment with π o = π b 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 5. 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-way interactions with a single attention layer [44]. We visualize attention maps on test sequences in Figure 2, which shows that the model has learned an induction head mechanism. The sequence in the middle figure has ( q k , o k ) ∈ { ( 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 is active on every repeated token, so 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 we found to 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

a 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 high-dimensional vectors in residual streams, they are often 'collapsed' down to scalar measurements by testing against other representations, using operations of the form v ⊤ j Wu i for some matrix W . For instance, u i and v j could be key and query vectors in an attention head, or input and output embeddings for predicting the next token. If ( u i ) i and ( v j ) j are orthonormal (or nearly-orthonormal) sets of embeddings, a natural way to store desired input-output associations i, j is through the following associative memory:

so that the scores v ⊤ j Wu i ≈ α ij 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 i ) i and ( v j ) j are nearlyorthonormal is to set them to be random high-dimensional vectors, such as Gaussian vectors with variance 1 /d in d dimensions. Indeed, these are known to satisfy [23, 53]

so that (3) is a reasonable way to define an associative memory, without requiring an explicit activation function as employed in end-to-end memory networks [50]. We may also easily create a 'remapping' of an existing embedding u i by multiplying it by a random matrix W 0 ∈ R d × d with Gaussian entries of variance 1 /d , which is commonly used for initializing neural network parameters. The new remapped embedding W 0 u i is near-unit norm, and is near-orthogonal to u i in addition to the other u j . 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 embeddings via gradient descent leads to a weighted associative memory of a form similar to (3).

Lemma1 (Gradients and associative memories) . Let p be a data distribution over input-output tokens, and consider the following loss, where the input and output embeddings W E and W U are fixed:

with ℓ the cross-entropy loss. The gradients of the population loss L then take the form

where ˆ p W ( y = k | x ) = σ ( W U Ww E ( z )) k are the model's predicted probabilities. Running gradient descent (with or without weight decay) from initialization W 0 then leads to estimates of the following form, for some α 0 and α ij that vary with the number of iterations:

̸

Note that (4) is a 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 ˆ p W ( y = k | z ) ≈ p ( y = k | z ) . We remark that if W 0 is a Gaussian random matrix, as if often the case for neural network layers, the first term in (6) plays a minor role: testing W 0 against an input-output pair ( i, j ) with α ij = 0 will concentrate around zero when d is large, while the ( i, j ) term in the sum will concentrate around α ij . 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].

Handling superposition. In Lemma 1, 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 a superposition, and return a new superposition (up to additional noise due to near-orthogonality). In practice, we will see that learned memories often focus on a single embedding and filter out the rest as noise when irrelevant (see also Section 6). We note that linearity can also be limiting, since it makes it difficult to map sets to specific output embeddings: u { i,j } := u i + u j needs to map to Wu i + Wu j , and thus cannot map to a new embedding v { i,j } . Such mappings of sets thus require non-linear associative memories, for instance by leveraging a sparse decoding of which elements are actually present ( e.g. , using compressed sensing), or by using MLPs with non-linear activations [15, 30].

4.2 A simplified two-layer transformer architecture

We consider a 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 ( W E , W U , W P ) to their random initialization throughout training. This brings us to the Gaussian random vector setup presented above.

We fix W 1 Q = W 2 Q = I d , so that W 1 K and W 2 K 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 ⊤ q W ℓ K x k , with x q (resp. x k ) the query (resp. key) embeddings, which now directly resembles an associative memory lookup.

We freeze W 1 V , W 1 O , and W 2 V to random initialization. These play the role of remapping attended tokens into new tokens, since for random W and large d , Wx is nearly orthogonal to x and to any other random embeddings independent of x .

We train W 2 O , since the outputs of the induction head need to be mapped back into appropriate output embeddings in order to predict the output tokens o k correctly.

We use a single linear feedforward layer after the second attention block, with weight matrix W F . 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 [10, 22, 24] since the model is still highly non-linear in its parameters and, as we will see, induces rich non-linear learning dynamics.

Solving the bigram problem with associative memories. Wenowshowhowthe above architecture can solve the synthetic bigram problem from Section 3 with well-chosen weight matrices. While this is only a 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 superpositions of various embeddings including token embeddings, positional embeddings, or 'remapped' versions thereof. These may be viewed as sets, as illustrated in Figure 1, 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:

where Q is the set of triggers when they are fixed, or the support of π q 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 1 O W 1 V , and the output matches a remapping of the input token by W 2 V to the corresponding output token. We remark that one can easily make the attention patterns more peaked on the correct associations by rescaling W 1 K and W 2 K . The global bigram statistics can be encoded in the feed-forward layer as follows:

The question remains of how the model could trade-off predictions from the induction head and from the feed-forward layer, which are added together due to residual connections. With fixed triggers Q , we may simply remove all i ∈ Q from the summation in (8), so that the model exclusively relies on the attention head for all triggers (indeed, the output of W 2 O is in the span of output embeddings, which are nearly orthogonal to the row space of W 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 on context, and W F may try to learn more complex mappings that also use the outputs of W 2 O . 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 factorization W ⊤ K W Q (potentially with low rank matrices W K , W Q ∈ R d h × d with d h < d as in multi-head attention), instead of a single matrix W K . 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 E that they empirically lead to similar observations. Nonetheless, we focus on our simpler architecture in the main paper due to simplicity of exposition and better interpretability thanks to a clear identifiability of the role of each matrix, which is lost in models with more heads and layers.

5 Empirical Study

In this section, we present our empirical analysis of learning dynamics on 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 at https://github.com/albietz/transformer-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 use a fixed learning rate and weight decay. Hyperparameters are given in Appendix E. Unless otherwise noted, we use d = 128 , random triggers with π q = π u 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-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 ∗ = ∑ ( i,j ) ∈M v j u ⊤ i , the corresponding recall metric is computed from the empirical estimate ˆ W as

We use this for each matrix in (7) as target, and additionally test the previous token matrix W 1 K on smaller time windows. For the final feed-forward layer, we measure the average KL divergence between the predicted softmax distribution using only W F and the global bigram distribution π b :

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 feed-forward layer, in order to focus on the learning of attention matrices W 1 K , W 2 K and W 2 O in isolation.

Figure 3 studies the effect of freezing different layers until iteration 300 on the training dynamics. By looking at memory recall probes, we see that training key-query matrices does not lead to any

0

200

400

600

800

1000

0.0

0.2

0.4

0.6

0.8

1.0

in-context accuracy

freeze

W

2

O

0

200

400

600

800

1000

0.0

0.2

0.4

0.6

0.8

1.0

freeze

W

1

K

and

W

2

K

0

200

400

600

800

1000

0.0

0.2

0.4

0.6

0.8

1.0

freeze

W

2

K

0

200

400

600

800

1000

0.0

0.2

0.4

0.6

0.8

1.0

freeze

W

1

K

0

200

400

600

800

1000

iteration

0.0

0.2

0.4

0.6

0.8

1.0

memory recall

0

200

400

600

800

1000

iteration

0.0

0.2

0.4

0.6

0.8

1.0

0

200

400

600

800

1000

iteration

0.0

0.2

0.4

0.6

0.8

1.0

0

200

400

600

800

1000

iteration

0.0

0.2

0.4

0.6

0.8

1.0

Wo2

Wk2

Wk1 (t<64)

Wk1

learning unless W 2 O is learned first, and that W 2 O can learn the correct associations even when trained by itself with key-value matrices at random initialization. Recall that the attention weights are essentially uniform when W K are at random initialization, so that training W 2 O 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 2 O (a similar observation was made in [49] 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 learned before the first, in the sense that W 2 K is easier to learn when W 1 K is frozen compared to the reverse, yet learning them together seems beneficial, possibly due to helpful feedback loops [1]. We also observe that W 1 K 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 q k , 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/right) 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 W 2 O 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 W 2 O and W F later in training, leading to slight degradations of our probe metrics. This may be due to the fact that the input to W F 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 4.2 may no longer be accurate.

Role of the data distribution. We can see in Figure 4(left) 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) a 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 2 O , 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 2 K , thus increasing the data requirements in order to learn the full associative memory. We also show in Figure 4(center) that changing the output token distribution to bigram distributions at training time reduces the in-context accuracy

0

100

200

300

400

500

iteration

0

1

2

3

4

5

loss

in-context vs global loss

K=1, fix q frequent

K=5, fix q rare

K=1, random q

K=5, random q

0

100

200

300

400

500

iteration

0.0

0.2

0.4

0.6

0.8

in-context accuracy

out of distribution

unif (train unif)

bigram (train unif)

unif (train bigram)

bigram (train bigram)

0

100

200

300

400

500

iteration

0.0

0.2

0.4

0.6

0.8

memory recall

Wo2 recall

Wk2 recall

Wk1 recall (t<64)

1.0

1.2

1.4

1.6

1.8

2.0

2.2

KL

(

WF

,

b

)

attention and feed-forward probes

Wf KL

when using out-of-distribution output tokens, while training on uniform outputs performs well on both distributions. This highlights that using a more diverse training distribution can lead to models with better generalization accuracy, with little additional training cost.

Additional experiments. In Appendix E, we provide additional experimental results for varying dimensionality, more complex architectures and training methods, as well as more fine-grained visualizations of the memory associations.

6 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 1, 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 a data distribution on ( x, y ) ∈ R d × [ N ] , and consider the following classification problem, with fixed output embeddings W U :

The gradients take the following form: denoting µ k := E [ x | y = k ] and ˆ µ k := E x [ ˆ p W ( k | x ) p ( y = k ) x ] ,

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, µ k can isolate relevant parts of x that are predictive of a label k , and thus can lead to the right associations.

An illustrative example. To gain more intuition about this result, consider the following example: we would like to predict y from x = w E ( y ) + p t , where p t is a positional embedding at a random position t ∈ [ T ] , which we would like to ignore. Further assume that y 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 W 0 = 0 (so that ˆ p W 0 ( k | x ) = 1 /N ):

with ¯ µ = E [ x ] . We show in Appendix B that when d is large enough to ensure near-orthonormal embeddings, we have

so that for large enough N and T , we obtain a near-perfect classifier that ignores the positional embedding, after just one gradient step (but a 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 [2, 25, 28]. Using a 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 a 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 W 2 O , then W 2 K , followed by W 1 K .

To show this result, we use Lemma 2 in a similar manner to the illustrative example above to show how training W 2 O by itself at initialization, i.e. , when the attention patterns are near-uniform, 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 is guaranteed to be present in the context, while other tokens need not appear more relative to other sequences. See also Figure 9 in Appendix E for numerical experiments verifying this for finite data and dimension. Once W 2 O has learned the correct associations, we show that the gradient with respect to the key-value matrix W 2 K at zero initialization can leverage the correctness of W 2 O to find the right associative memory that focuses the attention on correct triggers. Finally, by linearizing the second-layer attention around W 2 K = 0 , we show how gradients w.r.t. W 1 K 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 a 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 [37, 31], or grokking [33, 39]), factorized key-query and valueoutput 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

Z. Allen-Zhu and Y. Li. Backward feature correction: How deep learning performs deep learning. arXiv preprint arXiv:2001.04413 , 2020.

Z. Allen-Zhu and Y. Li. Towards understanding ensemble, knowledge distillation and selfdistillation in deep learning. In Proceedings of the International Conference on Learning Representations (ICLR) , 2023.

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.

J. L. Ba, J. R. Kiros, and G. E. Hinton. Layer normalization. arXiv preprint arXiv:1607.06450 , 2016.

T. Bricken and C. Pehlevan. Attention approximates sparse distributed memory. Advances in Neural Information Processing Systems (NeurIPS) , 2021.

T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. In Advances in Neural Information Processing Systems (NeurIPS) , 2020.

V. Cabannes, E. Dohmatob, and A. Bietti. Scaling laws for associative memories. arXiv preprint arXiv:2310.02984 , 2023.

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.

F. Charton. What is my math transformer doing?-three results on interpretability and generalization. arXiv preprint arXiv:2211.00170 , 2022.

L. Chizat, E. Oyallon, and F. Bach. On lazy training in differentiable programming. Advances in Neural Information Processing Systems (NeurIPS) , 2019.

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.

A. Damian, J. Lee, and M. Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference on Learning Theory (COLT) , 2022.

G. Dar, M. Geva, A. Gupta, and J. Berant. Analyzing transformers in embedding space. arXiv preprint arXiv:2209.02535 , 2022.

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.

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.

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. Hatfield-Dodds, D. Hernandez, A. 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.

M. Geva, J. Bastings, K. Filippova, and A. Globerson. Dissecting recall of factual associations in auto-regressive language models. arXiv preprint arXiv:2304.14767 , 2023.

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.

A. Graves, G. Wayne, and I. Danihelka. Neural turing machines. arXiv preprint arXiv:1410.5401 , 2014.

S. Gunasekar, B. E. Woodworth, S. Bhojanapalli, B. Neyshabur, and N. Srebro. Implicit regularization in matrix factorization. Advances in Neural Information Processing Systems (NIPS) , 2017.

J. J. Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences , 79(8):2554-2558, 1982.

J. Hron, Y. Bahri, J. Sohl-Dickstein, and R. Novak. Infinite attention: Nngp and ntk for deep attention networks. In Proceedings of the International Conference on Machine Learning (ICML) , 2020.

A. Iscen, T. Furon, V. Gripon, M. Rabbat, and H. Jégou. Memory vectors for similarity search in high-dimensional spaces. IEEE transactions on big data , 4(1):65-77, 2017.

A. Jacot, F. Gabriel, and C. Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in Neural Information Processing Systems (NeurIPS) , 2018.

S. Jelassi, M. Sander, and Y. Li. Vision transformers provably learn spatial structure. In Advances in Neural Information Processing Systems (NeurIPS) , 2022.

Y. Jiang and C. Pehlevan. Associative memory in iterated overparameterized sigmoid autoencoders. In Proceedings of the International Conference on Machine Learning (ICML) , 2020.

A. Joulin and T. Mikolov. Inferring algorithmic patterns with stack-augmented recurrent nets. Advances in Neural Information Processing Systems (NIPS) , 2015.

S. 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.

T. Kohonen. Correlation matrix memories. IEEE Transactions on Computers , 1972.

D. Krotov and J. J. Hopfield. Dense associative memory for pattern recognition. Advances in Neural Information Processing Systems (NIPS) , 2016.

Y. Li, Y. Li, and A. Risteski. How do transformers learn topic structure: Towards a mechanistic understanding. In Proceedings of the International Conference on Machine Learning (ICML) , 2023.

B. Liu, J. T. Ash, S. Goel, A. Krishnamurthy, and C. Zhang. Transformers learn shortcuts to automata. In Proceedings of the International Conference on Learning Representations (ICLR) , 2023.

Z. Liu, O. Kitouni, N. S. Nolte, E. Michaud, M. Tegmark, and M. Williams. Towards understanding grokking: An effective theory of representation learning. In Advances in Neural Information Processing Systems (NeurIPS) , 2022.

R. McEliece, E. Posner, E. Rodemich, and S. Venkatesh. The capacity of the hopfield associative memory. IEEE transactions on Information Theory , 33(4):461-482, 1987.

K. Meng, D. Bau, A. Andonian, and Y. Belinkov. Locating and editing factual associations in gpt. In Advances in Neural Information Processing Systems (NeurIPS) , 2022.

W. Merrill, A. Sabharwal, and N. A. Smith. Saturated transformers are constant-depth threshold circuits. Transactions of the Association for Computational Linguistics , 10:843-856, 2022.

T. Mikolov, K. Chen, G. Corrado, and J. Dean. Efficient estimation of word representations in vector space. In Proceedings of the International Conference on Learning Representations (ICLR) , 2013.

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.

N. Nanda, L. Chan, T. Liberum, J. Smith, and J. Steinhardt. Progress measures for grokking via mechanistic interpretability. In Proceedings of the International Conference on Learning Representations (ICLR) , 2023.

C. Olsson, N. Elhage, N. Nanda, N. Joseph, N. DasSarma, T. Henighan, B. Mann, A. Askell, Y. Bai, A. Chen, T. 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.

A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, et al. Language models are unsupervised multitask learners. Technical report, OpenAI , 2019.

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:2112.11446 , 2021.

Y. Razeghi, R. L. Logan IV, M. Gardner, and S. Singh. Impact of pretraining term frequencies on few-shot reasoning. In Conference on Empirical Methods in Natural Language Processing (EMNLP) , 2022.

C. Sanford, D. Hsu, and M. Telgarsky. Representational strengths and limitations of transformers. In Advances in Neural Information Processing Systems (NeurIPS) , 2023.

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.

I. Schlag, K. Irie, and J. Schmidhuber. Linear transformers are secretly fast weight programmers. In Proceedings of the International Conference on Machine Learning (ICML) , 2021.

J. Schmidhuber. Learning to control fast-weight memories: An alternative to dynamic recurrent networks. Neural Computation , 4(1):131-139, 1992.

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 a large-scale language model. In North American Chapter of the Association for Computational Linguistics (NAACL) , 2022.

C. Snell, R. Zhong, D. Klein, and J. Steinhardt. Approximating how single head attention learns. arXiv preprint arXiv:2103.07601 , 2021.

S. Sukhbaatar, A. Szlam, J. Weston, and R. Fergus. End-to-end memory networks. Advances in Neural Information Processing Systems (NIPS) , 2015.

Y. Tian, Y. Wang, B. Chen, and S. Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. In Advances in Neural Information Processing Systems (NeurIPS) , 2023.

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.

R. Vershynin. High-dimensional probability: An introduction with applications in data science , volume 47. Cambridge university press, 2018.

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.

J. Weston, S. Chopra, and A. Bordes. Memory networks. In Proceedings of the International Conference on Learning Representations (ICLR) , 2015.

D. J. Willshaw, O. P. Buneman, and H. C. Longuet-Higgins. Non-holographic associative memory. Nature , 222(5197):960-962, 1969.

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.

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.

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.

B. Zhang and R. Sennrich. Root mean square layer normalization. In Advances in Neural Information Processing Systems (NeurIPS) , 2019.

Y. Zhang, A. Backurs, S. Bubeck, R. Eldan, S. Gunasekar, and T. Wagner. Unveiling transformers with lego: a synthetic reasoning task. arXiv preprint arXiv:2206.04301 , 2022.

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 u k ∈ R d with i.i.d. Gaussian N (0 , 1 d ) entries.

We recall a few facts:

( Norm ) We have u ⊤ i u i ≈ 1 . This is standard from the concentration of random vectors in high dimension (see, e.g. , [53, Theorem 3.1.1]).

̸

( Near-orthogonality ) For i = j , we have u ⊤ i u j = O (1 / √ d ) . To see this, denoting u i = d -1 / 2 (˜ u ik ) k , where ˜ u ik are the unnormalized entries of u i , note that we have

by the central limit theorem, since for each k , the quantities ˜ u ik ˜ u jk are zero-mean, unit-variance, i.i.d. random variables.

( Remapping: norm ) If W is a Gaussian random matrix with i.i.d. N (0 , 1 d ) entries, then for any fixed x we have ∥ Wx ∥ ≈ ∥ x ∥ . This follows from Johnson-Lindenstrauss (see, e.g. , [53, Lemma 5.3.2 and Exercise 5.3.3]). In particular, if x is a normalized Gaussian embedding as above, then ∥ Wx ∥ ≈ 1 .

( Remapping: near-orthogonality ) Consider a random vector x = 1 √ d ˜ x and a random matrix W = 1 √ d ˜ W , where the entries of ˜ x and ˜ W are i.i.d. N (0 , 1) . Then x and Wx are nearly orthogonal. To see this, note that E [ x ⊤ Wx ] = E [ x ⊤ E [ W ] x ] = 0 , and the variance is

̸

where M 2 and M 4 denote the second and fourth moments of the standard Gaussian, respectively. Then, Chebyshev's inequality implies that | x ⊤ Wx | = O (1 / √ d ) 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 to do this in a precise manner in this paper, and will generally assume d 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 on the illustrative example (Section B.1), derivations of gradients w.r.t. key-query matrices at initialization (Section B.2), as well as a study of how the induction head mechanism may develop in a simplified setup, using a sequence of single layer-wise gradient steps in a top-down manner (Section B.3).

B.1 Details on illustrative example

Consider the example discussed in Section 6: we would like to predict y from x = w E ( y ) + p t , where p t is a positional embedding at a random position t ∈ [ T ] , which we would like to ignore. Further assume that y is uniformly distributed ( p ( y = k ) = 1 /N ) and consider the matrix obtained after one population gradient step with step-size η starting from an initialization W 0 = 0 (so

that ˆ p ( k x ) = 1 /N

with ¯ µ = E [ x ] .

Note that we have µ k = w E ( k ) + 1 T ∑ t p t and ¯ µ = 1 N ∑ k w E ( k ) + 1 T ∑ t p t , so that (11) becomes with ¯ w E := 1 N ∑ N k =1 w E ( k ) . When d is large enough to ensure near-orthonormal embeddings, we have for any y and t ,

This implies

as claimed in the main text. The classifier ˆ y = arg max k w U ( k ) ⊤ W 1 ( w E ( y )+ p t ) 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 at 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 2. When output matrices of the previous layer already encode the desired associations, these gradients can lead to 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 2 K at zero initialization. For simplicity, we drop the d -1 / 2 factor from the softmax, which only changes gradients by a 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 y from an attention layer with inputs X = [ x 1 , . . . , x T ] , and value-output matrix Φ 2 := W 2 O W 2 V :

with ℓ the cross-entropy loss and σ ( u ) t = e ut ∑ s e us for u ∈ R T is the softmax.

The gradient at W = 0 is given by

Now we consider the gradient w.r.t. W = W 1 K at zero initialization, and consier a simplification of the second layer attention to its linearization around W 2 K = 0 . We will see that this still provides first-order information that is sufficient for W 1 K to be learned.

with ¯ x 1: T = 1 T ∑ T t =1 x t .

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 1 K , the key-query matrix at the first attention layer:

Here, ¯ σ ( u 1: T ) t = 1 T (1+ u t -1 T ∑ T s =1 u s ) is the linearization of the softmax around 0 , and Z ( W ) = [ z 1 ( W ) , . . . , z T ( W )] with and Φ ℓ = W ℓ O W ℓ V for ℓ = 1 , 2 .

The gradient at W = 0 is given by

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 a single trigger ( K = 1 ), and assume that π u , π q , π o and π b ( ·| i ) 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 z T is the second occurrence of the trigger token, and the label y = z T +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 1, where the loss is only taken on triggers after the second occurrence: in that case, the loss may be written as a weighted sum of the one we consider here, weighted by the probability of the second (or later) trigger appearing at the given position T .

In practice, when the loss is on all tokens and W F is also learned, we may expect that W F 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 1, which provides recovery of bigram statistics when d 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(right)). 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 W F in the remainder of this section, focusing on the learning of W 2 O , W 2 K and W 1 K , in this top-down order. We will consider zero-initialization a single gradient steps, noting that random initialization should lead to similar associative memory behaviors when the dimension is large enough, since it leads to a remapping of input embeddings which is near-orthogonal to any output embedding (see Appendix A).

B.3.1 Learning W 2 O

We begin by studying the learning of the second output matrix W 2 O . In the above data model, we may consider a loss as in Lemma 2 with input-outputs ( x, y ) , where y is the output token of the sequence,

and x depends on the random sequence z 1: T as

where ε t = p t + 1 t ∑ t s =1 Φ 1 ( w E ( z s ) + p s ) with Φ 1 = W 1 O W 1 V , is a 'noise' vector from the residual streams, containing positional embeddings as well as an average attention output from the first layer. In practice, the logit predictions are of the form W U ( W 2 O x + ε T ) due to residual connections, but we ignore the term W U ε T for simplicity, noting that it is near-zero when d is large.

After a gradient step on W 2 O with step-size η , starting from zero-initialization (so that ˆ p ( k | x ) = p ( y = k ) = 1 /N for all x ), Lemma 2 yields

Now, consider the random variables q (trigger token), o (output token), t o (position of the first occurrence of the output token). In our simplified data model, q and t o have the same distribution regardless of the conditioning on y = k , while o is equal to k when y = k , while it is uniform in [ N ] without this condition. The sequence z 1: T has the same distribution in either p ( · ) or p ( ·| y = k ) , except for the token z t o .

Then, we may write:

since ε t is independent of o when t < t o . Noting that ε t only depends on z t o and thus on y through the first layer attention, we have

where τ := E [ ∑ T t = t o 1 t ] , and ¯ w E = 1 N ∑ N k =1 w E ( k ) . Thus, (15) becomes so that when d is learn enough to ensure near-orthonormal embeddings, we have

✶ where the O ( · ) terms are due to the ¯ w E elements. The first line yields a behavior that matches desired associative memory in (7) of Section 4.2 when N is large. The second line shows additional spurious associations that are stored in W 2 O 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 2 O are still near-random, since in 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).

B.3.2 Learning W 2 K

Now assume that W 2 O is as in (16). As argued above, the predictions ˆ p ( k | x ) are essentially random 1 /N in our model for W 2 K = 0 , so that after one gradient step on W 2 K with learning rate η , Lemma 4 yields:

where x t are the inputs to the second attention layer, given by

From now on, we consider a simplified architecture where only x t, 0 are fed as queries and values, while only x t, 1 are fed as keys. Using the fact that trigger tokens q are sampled uniformly ( i.e. , π q = 1 /N ), we have

where with ¯ x 1 = 1 T ∑ t x t, 1 . Now, note that we have w U ( k ) ⊤ Φ 2 x t, 0 ≈ α ✶ { z t = k } with α = η/TN by (16), and x T, 0 = w E ( q ) + p T . This yields

with

where the three terms split the sum inside the expectation

where T q = { t o -1 , T } and T r = [ T ] \ { t o , t o -1 , T } (recall that t o is a random variable, corresponding to the first occurrence of the output token, so that these sets are random).

We will now show that ∆ o k,j carries the desired signal for the appropriate induction head associative memory, while ∆ q k,j and ∆ r k,j are negligible, for N large enough.

Controlling ∆ o k,j . For t = t o , noting that z t o = y , we have

with a k,j,i ≈ (Φ 1 w E ( i )) ⊤ E [ x t o , 1 -¯ x 1 | y = k, q = j ] thanks to near-orthonormality, and is a spurious positional mixture.

We then distinguish the following cases:

If j = k and i = j , since the trigger token j only appears at positions t o -1 and T , we have

̸

We may expect t o to be concentrated around T/ 2 , in which case γ T ≳ 2 T -1 T -1 T 2 ≥ C T > 0 for T larger than a small constant.

If j = k = i , the two occurrences of the trigger happen one after the other, so it must be that t o = T . Then

for T larger than a small constant.

If i = j = k , all tokens up to position t o -2 = T -2 are i.i.d. uniform in [ N ] \ { j } , so that

̸

̸

In summary, we obtain

̸

̸

̸

If i = j and j = k , all tokens except at positions t o -1 , t o and T (we have t o < T ) are uniform in [ N ] \ { j } . The triggers do not contribute anything to a k,j,i since i = j , and the output token may be also randomized by taking the average over k ∈ [ N ] \ { j } . We thus obtain

̸

̸

Thus, when N is large, while T is moderate, the above sum leads to more signal in the i = j terms compared to i = j . In particular, this yields

̸

so that this component in (24) acts precisely like the desired associative memory in (7).

It remains to show that the other components are negligible compared to this. It then suffices to show:

Controlling ∆ q k,j . For t ∈ T q , note that we always have z t = j in the expectations, so that ∆ q k,j = 0 unless k = j . For k = j , we have ∆ q k,j = O (1) , so that

Controlling ∆ r k,j . Using that ∥ x t, 1 -¯ x 1 ∥ ≤ C = O (1) for all t , we provide the following crude bound via the triangle inequality and Hölder inequality:

Overall, we have shown that up to some spurious positional embeddings, W 2 K behaves as the desired associative memory from (7) when N is large enough, satisfying:

since z t is independent of y given t ∈ T r and thus is uniform in [ N ] \ { j } , and |T r | ≤ T . We note, however, that ∆ r k,j may be controlled much more finely by leveraging the similarities between the distributions of z t , t ∈ T r with or without conditioning on y .

We note that one may then amplify the gap between correct and incorrect associations by having a 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 1 K

We now assume that W 2 O and W 2 K have learned the correct associations, and consider one gradient step away from zero-initialization on W 1 K . Note that when W 1 K = 0 , the predictions of the model are 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 Φ 1 w E ( q ) in their average attention, which activates the second-layer attention head. Then the output is likely to predict the trigger itself, which will be an incorrect prediction most of the time.

We may thus consider ˆ p ( k | X ) = 1 /N at this stage as well. We also consider a simplified architecture where the first layer attention only uses positional embeddings in the key-query matrix, and only token embeddings in the value-output matrix. In particular, we have x t = w E ( z t ) . Lemma 5 then gives the following form for W 1 K after one gradient step of step-size η :

Note that since W 2 O and W 2 K already captured the desired associations at this stage, we have

for some α, α ′ > 0 . Recall that in our model, we have z T = q with probability one ( q is the trigger token), and that q only appears twice: once at position t q := t o -1 < T and once at position T . We then have, for any t > 1 ,

✶ where r k := 1 T ∑ T t =1 { z t = k } . We have

since the distribution of t q is the same regardless of y . We proceed similarly for the other quantities and obtain the following:

This yields the following associative memory behavior, for t > 1 :

which matches the desired 'previous token head' behavior from (7) when N is large. As in the case of W 2 K , we may then 'saturate' the softmax by choosing a large enough step-size.

C Other Proofs

C.1 Proof of Lemma 1

Proof. Recall the form of the cross-entropy loss for classification with K classes:

Its derivatives take the form

with s ( ξ ) k = e ξ k ∑ j e ξ j the softmax.

with

The gradient of L is then given by which yields the desired result.

C.2 Proof of Lemma 2

Proof. Using similar steps as the proof of Lemma 1, we have with ˆ µ k = p ( y = k ) -x [ˆ p W ( k | x ) x ] and µ k = [ x | y = k ] .

C.3 Proof of Lemma 4

Proof. To better isolate the role of keys from values, we denote the keys that are fed into the matrix W by Z = [ z 1 , . . . , z T ] ∈ R d × T , while the query is simply x T . 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:

We have

where ¯ z 1: T = 1 T ∑ t z t , and we used the fact that

The gradient (31) now writes

C.4 Proof of Lemma 5

Proof. The linearization of the second layer softmax around zero takes the following form:

with z t = ∑ t s =1 Φ 1 x s σ ( p ⊤ 1: t Wp t ) s the output of the first attention layer.

Then,

using (32). The result follows by using ∂ℓ ∂ξ k ( y, ξ ) = ˆ p ( k | ξ ) -✶ { y = k } .

D Beyond our Simplified Architecture

While the simplified architecture presented in 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.

Factorizations. In practice, transformers typically involve products of matrices, potentially with a low-rank bottleneck. For instance, our key-query matrices W K should instead be considered as a product W ⊤ K W Q , and the output-value matrices W O and W V are typically jointly optimized.

Consider an associative memory of the form:

where ( x i ) i and ( y i ) i are appropriate collections of near-orthonormal embeddings.

We now argue that a similar associative memory can be achieved with the factorization W = d 2 d ′ UV , where U ∈ R d × d ′ , V ∈ R d ′ × d with d ′ ≤ d (for instance d ′ could be the dimension of attention heads), are given by: 3

where U 0 and V 0 are random matrices with N (0 , 1 d ) entries. These matrices are similar to those that would arise from a single gradient step individually on U and V from initializations U 0 and V 0 , as in Lemma 1. To see why W behaves like W ∗ , note that we have

It is also easy to check using central limit arguments (similar to remapping in Appendix A) that ˜ x i := √ d d ′ V 0 x i ∈ R d ′ and ˜ y i := √ d d ′ U ⊤ 0 y i are all nearly-orthonormal embeddings. Thus, we have

✶ ✶ where the first and last term vanish due to the cross-terms ˜ y ⊤ i ˜ x i ′ which vanish for any i, i ′ . Thus, W and W ∗ encode the same associations, when d and d ′ are large enough to ensure near-orthogonality.

Layer-normalization. Normalization layers [4, 60] are typically used in transformers to improve training stability [52], and are applied on each token representation, either after [52] or before [6] each block. It may be seen as an operation of the form 4

applied to the input or output of a given block.

In order to obtain a basic understanding of the role of layer-norm in our associative memory setup, we may consider the setup Lemma 1, with a normalization applied after the linear operation, leading to the population loss:

The gradients then take the form

This illustrates that in addition to weighting the updates on class k by the prediction error ˆ p ( k | x ) -p ( k | x ) , the updates are also projected on the orthogonal of the Wx direction. This means that an update on the direction w U ( k ) x ⊤ will occur only to the extent that Wx is not already aligned with w U ( k ) . Thus, if an association is 'stored' once, so that Wx ≈ w U ( 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 [7] for more discussion on this).

Note that at random initialization Wx is nearly orthogonal to any w U ( k ) , so that layer-norm only starts playing a significant role later in training, and does not affect our theoretical analysis based on single gradient steps.

MLP blocks. If we denote by ( u i ) i and ( v i ) i collections of near-orthonormal input and output embeddings, an MLP block may encode associations ( i, j ) ∈ M as follows:

̸

where σ is a non-linear activation and b a bias term. Then, if one assumes that u ⊤ i u j ≤ b for i = j and σ ( t ) = 0 for t < 0 , 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, e.g. , [7, 30]).

An additional benefit of MLP layers discussed in Section 4 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. , a subject and a relation in a factual recall task [35]). This may be achieved, for instance, by considering embeddings u I = 1 √ |I| ∑ i ∈I u i , where I are sets of bounded size ( e.g. ,

as obtained using layer-norm over a residual stream). Then, assuming the u i are nearly-orthonormal, we have u ⊤ I u I ≈ 1 while u ⊤ I u I ′ ≲ 1 -δ if I = I ′ , for some δ that depends on the maximal cardinality of these sets. Although 1 -δ is no-longer vanishingly small, defining the non-linearity may still succeed at filtering out any I that does not correspond to the query set in x . 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 as 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.

It is important to note, however, that the redundancy introduced by having more heads and layers makes it more challenging to identify which layer/head/weight will learn certain associations (see, e.g. , Figure 11 in Appendix E). This is in contrast to our simplified architecture of Section 4.2, where we may identify the role of each matrix (up to some possible redundancy when using the feed-forward layer W F ). 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 a 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 is denoted η . We fix the momentum parameter to 0 . 9 and the weight decay parameter to 10 -4 . U denotes the uniform distribution over [ N ] .

0

100

200

300

400

500

iteration

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

memory recall

Wo2

K=1, fix q frequent

K=5, fix q rare

K=1, random q

K=5, random q

0

100

200

300

400

500

iteration

0.0

0.2

0.4

0.6

0.8

1.0

memory recall

Wk2

K=1, fix q frequent

K=5, fix q rare

K=1, random q

K=5, random q

0

100

200

300

400

500

iteration

0.0

0.2

0.4

0.6

0.8

memory recall

Wk1 (t<64)

K=1, fix q frequent

K=5, fix q rare

K=1, random q

K=5, random q

0

100

200

300

400

500

iteration

0.0

0.1

0.2

0.3

0.4

0.5

0.6

memory recall

Wo2

0

100

200

300

400

500

iteration

0.0

0.2

0.4

0.6

0.8

memory recall

Wk2

0

100

200

300

400

500

iteration

0.0

0.2

0.4

0.6

0.8

memory recall

Wk1 (t<64)

Figure 2: K = 3 , π q = π u (random triggers) or Q is the K most likely elements of π u , π o = U , d = 128 , d hidden = 4 × 128 (hidden dimension of the feed-forward MLPs), η = 0 . 2 .

Figure 3: K = 5 , π q = π u (random triggers), π o = U , d = 128 , η = 0 . 2 .

Figure 4(left) and Figure 5: π o = U , d = 128 , η = 1 . For random triggers we use π q = π u . For K = 1 with fixed frequent trigger, the only trigger is the most probable token according to π u , while for K = 5 with fixed rare triggers, the five triggers are the 6-th to 10-th most probable tokens according to π u .

Figure 4(center): K = 3 , π q = π u (random triggers), π o = U or π o = π b (conditioned on the trigger), d = 128 , η = 1 .

Figure 4(right): K = 3 , π q = π u (random triggers), π o = U , d = 128 , η = 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 4(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, or random triggers (note that the probe for W 2 K with fixed triggers only shows recall accuracy on the set of triggers Q , which is an easier task). Figure 6

10

10

20

20

20

40

30

30

60

40

40

50

50

80

60 -

60

0

20

40

60

80

0

10

20

30

40

50

60

0

10

20

30

40

50

60

0

50

100

150

200

250

iteration

0.0

0.2

0.4

0.6

0.8

1.0

Wo2 recall, fixed lr

32

64

128

256

512

1024

0

20

40

60

80

iteration

0.2

0.4

0.6

0.8

1.0

Wo2 recall, tuned lr

32

64

128

256

512

1024

0

10

20

30

40

batches

0.0

0.2

0.4

0.6

0.8

1.0

accuracy

one-step classifier

d = 32

d = 64

d = 128

d = 256

d = 512

d = 1024

d = 2048

10 2

10 3

dimension d

0.2

0.4

0.6

0.8

accuracy

one-step classifier

5 batches

10 batches

20 batches

shows that using uniform output tokens can lead to better fitting of W 2 O and W 2 K compared to using output tokens sampled from bigrams. In addition to the increased diversity when using uniform outputs, this may also be due to the fact that bigram outputs are already well predicted using global statistics with just the feed-forward layer, hence the gradient signal on such well-predicted tokens may not propagate through the induction head mechanism. In contrast, the recall accuracy for W 1 K is comparable for both settings, since the previous token head is useful at all positions regardless of the output token distribution.

Visualizing memories. Figure 7 shows visualizations of the associative memory behaviors after training. We see that diagonal elements dominate in the plots, which corresponds to correct associations lead to high 'memory recall'. Nonetheless, we see that some of the diagonal elements are weaker than others, particularly for late positions in W 1 K , and for some of the trigger tokens in W 2 K , while the diagonal for W 2 O seems to be roughly uniform. We note that characters corresponding to capital letters have token index 13 to 38, while lowercase letters have index 39 to 64. The association patterns found in W 2 K 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 a sequence, it is natural that W 1 K learns stronger associations at earlier positions. In contrast, diagonal elements for W 2 O are nearly uniform, which agrees with the fact that output tokens are sampled uniformly in this setup. We refer to the follow-up work [7] 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 to ensure near-orthogonality, and thus store input-output pairs more effectively. In Figure 8, we evaluate the recall accuracy for W 2 O for varying dimension, when training it 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 a few iterations with a tuned learning rate. If the learning rate isn'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 W 2 O alone is a convex optimization problem, and we hypothesize that higher dimension makes the problem better conditioned, and hence easier to learn. In Figure 9, we show 'one-step' recall accuracies for classifying output tokens from the average

0

100

200

300

400

500

iteration

1

2

3

4

loss

in-context vs global loss

K=5, random q

0

100

200

300

400

500

iteration

0.0

0.2

0.4

0.6

0.8

accuracy

in-context accuracy

K=5, random q

0

200

400

600

iteration

0.0

0.2

0.4

0.6

0.8

1.0

recall accuracy

Wo2 recall

Wk2 recall

Wk1 recall (t<64)

1.0

1.2

1.4

1.6

1.8

2.0

2.2

KL divergence

attention vs feed-forward probes

Wf KL

K : [l] T h e e _ m e

r

F _ m e

K : [l] T h e e _ m e

r

F _ m e

K : [l] T h e e _ m e

r

F _ m e

K : [l] T h e e _ m e

r

F _ m e

attention input to W 2 O , given by

where µ k = E [ x | y = k ] and µ = E [ x ] , for x = 1 t ∑ t s =1 W 2 V w E ( z s ) and y = z t +1 , when z t 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 a single gradient step on W 2 O from zero initialization (see Lemma 2 and Appendix B.3.1). The plots illustrate that this simple one-step model is already able to extract relevant signal from the noisy average attention, after a 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 10 shows training behavior for a more complex model than the simplified one considered in Section 5, namely where we train all parameters, replace the linear W F feedforward layer by a two-layer MLP, and were (pre-)layer-normalization 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 11 shows the attention maps obtained when training a multi-head 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 learns the previous token behavior, while the induction behavior is shared across different heads at the second layer. This illustrates the challenges of interpretability in the presence of redundant models, which then require additional work to identify which of the layers and heads are performing a given behavior, e.g. , through interventions and causal mediation analysis [35, 54].

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nIt 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.\n\nQUESTIONS:\nIn 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\nLIMITATIONS:\nI dont think the limitations of design decisions are well discussed or investigated. See Questions.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\n**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.\n\nQUESTIONS:\nMost 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.\n\nLIMITATIONS:\nAs mentioned above, I think the authors should discuss the limitations of their simplifying assumptions about the transformer architecture more.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nGiven 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nDiscussions could mention that this analysis may not extend to larger versions of the model, or provide empirical proof that it does.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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.\n\nLIMITATIONS:\nn/a", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-19"}

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

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.

2. The theoretical statements are non-trivial, interesting and not straightforward to prove.

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.

2. Given the simplicity of the toy task/model the practical implications and the generality of the findings are unclear (other tasks / architectures).

QUESTIONS:
.

LIMITATIONS:
There is no discussion of the limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-26