SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Alberto Bietti*	Vivien Cabannes	Diane Bouchacourt	Hervé Jégou	Léon Bottou
Flatiron Institute	FAIR, Meta	FAIR, Meta	FAIR, Meta	FAIR, Meta
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.
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  $[\dots ,\mathsf{a},\mathsf{b},\dots ,\mathsf{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
*Work done while at 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.
- 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.
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} \in [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 \in \mathbb{R}^{d \times N}$ . We will denote the embeddings of tokens  $z_t$  by  $x_t \coloneqq 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) \in \mathbb{R}^d$  are added to each token embedding depending on its position in the sequence, leading to the following input embeddings:
$$
x _ {t} := x _ {t} + p _ {t} = w _ {E} \left(z _ {t}\right) + w _ {P} (t). \tag {1}
$$
Figure 1: Induction head mechanism. Induction heads are a two-layer mechanism that can predict  $b$  from a context  $[\ldots, a, b, \ldots, a]$ . The first layer is a previous token head, which attends to the previous token based on positional embeddings  $(p_t \rightarrow p_{t-1})$  and copies it after a remapping  $(w_E(a) \rightarrow w_1(a) := W_O^1 W_V^1 w_E(a))$ . The second layer is the induction head, which attends based on the output of the previous token head  $(w_E(a) \rightarrow w_1(a))$  and outputs the attended token, remapped to output embeddings  $(w_E(b) \rightarrow w_U(b))$ . Boxes in the diagram represent different embeddings in superposition on each token's residual stream (we omit some irrelevant ones for clarity, e.g., positional embeddings in upper layers), and attention and output associations are shown with the associative memory viewpoint presented in Section 4.
- Attention blocks: given an input sequence  $x_{1:T} \in \mathbb{R}^{d \times T}$  of embeddings, the causal attention block computes, for  $W_K, W_Q, W_V, W_O \in \mathbb{R}^{d \times d}$  (key, query, value, output), and for each  $t$ ,
$$
x _ {t} ^ {\prime} := W _ {O} W _ {V} x _ {1: t} \sigma \left(x _ {1: t} ^ {\top} W _ {K} ^ {\top} W _ {Q} x _ {t}\right) \in \mathbb {R} ^ {d}, \tag {2}
$$
where  $\sigma$  takes the softmax of its elements, leading to an attention of the "values"  $W_{V}x_{t}$  with weights proportional to  $\exp ((W_Kx_s)^\top (W_Qx_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}^{\prime} = \mathcal{A}(x_{1:T};W_{K},W_{Q},W_{V},W_{O})$ , the  $\ell$ -th layer of the transformer applies attention with layer-specific parameters along with a residual connection as follows:
$$
x _ {1: T} := x _ {1: T} + \mathcal {A} (x _ {1: T}; W _ {K} ^ {\ell}, W _ {Q} ^ {\ell}, W _ {V} ^ {\ell}, W _ {O} ^ {\ell})
$$
- 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  $\mathcal{F}(\cdot; W_F)$ , with a residual connection: at layer  $\ell$ , we have
$$
x _ {t} := x _ {t} + \mathcal {F} \left(x _ {t}; W _ {F}\right).
$$
Our simplified setup will linear feed-forward layers:  $\mathcal{F}(x_t;W_F) = W_Fx_t$
- Unembedding: After the last transformer layer, the embeddings are mapped back to the vocabulary space  $\mathbb{R}^N$  through a linear "unembedding" layer  $W_U = [w_U(1),\dots ,w_U(N)]^\top \in \mathbb{R}^{N\times d}$ , where we refer to the  $w_{U}(j)$  as "output embeddings". The output of this layer is then fed into a cross-entropy loss for predicting of  $z_{t + 1}$  from each  $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  $[\dots ,\mathsf{a},\mathsf{b},\dots ,\mathsf{a}]\to \mathsf{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 consists 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]).
2We omit layer indices for simplicity of notation, and use the assignment operator  $\coloneqq$  instead.
Figure 2: Induction head behavior in attention maps observed on a 2-layer transformer trained on two variants of our synthetic dataset. Each row shows the attention pattern for predicting the next token. (left) The first layer head always attends to the previous token. (center) For fixed triggers  $Q = \{a, t\}$ , the second layer head mainly attends to tokens following such triggers. (right) For random triggers, the induction head mechanism is active for any repeated token (here the only trigger is  $L$ ). Red and green boxes highlight tokens following previous occurrences of the query, with red boxes corresponding to "correct" output tokens  $o_k$  following trigger tokens  $q_k$ .
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]$ :  $\pi_b(\cdot | i)$ ,  $\pi_u$ ,  $\pi_o(\cdot | i)$  and  $\pi_q$ , for  $i \in [N]$ .  $\pi_b(\cdot | i)$  are the global bigram conditionals,  $\pi_u$  the global unigram distribution, while  $\pi_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  $\pi_q$ . Each sequence  $z_{1:T}^{n}$  is generated as follows:
- (optional) Sample  $q_{1},\ldots ,q_{K}\sim \pi_{q}$ , i.i.d. without replacement (random triggers)
- Sample  $o_k \sim \pi_o(\cdot | q_k)$ , i.i.d. with replacement.
- Sample  $z_1^n \sim \pi_u$  and  $z_t^n | z_{t - 1}^n \sim p_n(\cdot | z_{t - 1}^n)$  for  $t = 2, \dots, T$ , where
$$
p _ {n} (j | i) = \left\{ \begin{array}{l l} \pi_ {b} (j | i), & \text {i f} i \notin \{q _ {k} \} _ {k} \\ \mathbb {1} \{j = o _ {k} \}, & \text {i f} i = q _ {k}. \end{array} \right.
$$
Experimental setup and initial experiment. Our experiments take  $\pi_u$  and  $\pi_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  $\pi_q = \pi_u$  or fix them to the  $K$  most frequent tokens. We sample uniform outputs  $o_k$  in most cases, but also experiment with  $\pi_o = \pi_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}) \in \{(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}^{\top}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:
$$
W = \sum_ {i, j} \alpha_ {i j} v _ {j} u _ {i} ^ {\top}, \tag {3}
$$
so that the scores  $v_{j}^{\top}Wu_{i}\approx \alpha_{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 nearly-orthonormal 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]
$$
u _ {i} ^ {\top} u _ {i} \approx 1 \quad \text {a n d} \quad u _ {i} ^ {\top} u _ {j} \approx O \left(\frac {1}{\sqrt {d}}\right),
$$
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 \in \mathbb{R}^{d \times 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).
Lemma 1 (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:
$$
L (W) = \mathbb {E} _ {(z, y) \sim p} [ \ell (y, W _ {U} W w _ {E} (z)) ], \tag {4}
$$
with  $\ell$  the cross-entropy loss. The gradients of the population loss  $L$  then take the form
$$
\nabla_ {W} L (W) = \sum_ {k = 1} ^ {N} \mathbb {E} _ {z} [ (\hat {p} _ {W} (y = k | z) - p (y = k | z)) w _ {U} (k) w _ {E} (z) ^ {\top} ], \tag {5}
$$
where  $\hat{p}_W(y = k|x) = \sigma (W_U W w_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  $\alpha_0$  and  $\alpha_{ij}$  that vary with the number of iterations:
$$
\hat {W} = \alpha_ {0} W _ {0} + \sum_ {i, j} \alpha_ {i j} w _ {U} (j) w _ {E} (i) ^ {\top}. \tag {6}
$$
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  $\hat{p}_W(y = k|z) \approx 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  $\alpha_{ij} \neq 0$  will concentrate around zero when  $d$  is large, while the  $(i,j)$  term in the sum will concentrate around  $\alpha_{ij}$ . We also note that the gradient updates described above correspond to a so-called maximal feature learning regime similar to  $\mu \mathrm{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\}} \coloneqq u_i + u_j$  needs to map to  $W u_i + W u_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_{Q}^{1} = W_{Q}^{2} = I_{d}$ , so that  $W_{K}^{1}$  and  $W_{K}^{2}$  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}^{\top}W_{K}^{\ell}x_{k}$ , with  $x_{q}$  (resp.  $x_{k}$ ) the query (resp. key) embeddings, which now directly resembles an associative memory lookup.
- We freeze  $W_V^1$ ,  $W_O^1$ , and  $W_V^2$  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_{O}^{2}$ , 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. We now show how the 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.