SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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:
$$
W _ {K} ^ {1} = \sum_ {t = 2} ^ {T} p _ {t} p _ {t - 1} ^ {\top}, \quad W _ {K} ^ {2} = \sum_ {k \in Q} w _ {E} (k) \left(W _ {O} ^ {1} W _ {V} ^ {1} w _ {E} (k)\right) ^ {\top}, \quad W _ {O} ^ {2} = \sum_ {k = 1} ^ {N} w _ {U} (k) \left(W _ {V} ^ {2} w _ {E} (k)\right) ^ {\top}, \tag {7}
$$
where  $Q$  is the set of triggers when they are fixed, or the support of  $\pi_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_O^1 W_V^1$ , and the output matches a remapping of the input token by  $W_V^2$  to the corresponding output token. We remark that one can easily make the attention patterns more peaked on the correct associations by rescaling  $W_K^1$  and  $W_K^2$ . The global bigram statistics can be encoded in the feed-forward layer as follows:
$$
W _ {F} = \sum_ {i = 1} ^ {N} \sum_ {j = 1} ^ {N} \log \pi_ {b} (j | i) w _ {U} (j) w _ {E} (i) ^ {\top}. \tag {8}
$$
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 \in Q$  from the summation in (8), so that the model exclusively relies on the attention head for all triggers (indeed, the output of  $W_{O}^{2}$  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_{O}^{2}$ . 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}^{\top}W_{Q}$  (potentially with low rank matrices  $W_{K}, W_{Q} \in \mathbb{R}^{d_{h} \times 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  $\pi_q = \pi_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_{*} = \sum_{(i,j)\in \mathcal{M}}v_{j}u_{i}^{\top}$ , the corresponding recall metric is computed from the empirical estimate  $\hat{W}$  as
$$
R (\hat {W}, W _ {*}) = \frac {1}{| \mathcal {M} |} \sum_ {(i, j) \in \mathcal {M}} \mathbb {1} \left\{\arg \max  _ {j ^ {\prime}} v _ {j ^ {\prime}} ^ {\top} \hat {W} u _ {i} = j \right\}. \tag {9}
$$
We use this for each matrix in (7) as target, and additionally test the previous token matrix  $W_{K}^{1}$  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  $\pi_{b}$ :
$$
d _ {K L} \left(W _ {F}, \pi_ {b}\right) := \frac {1}{N} \sum_ {k = 1} ^ {N} d _ {K L} \left(\sigma \left(W _ {U} W _ {F} w _ {E} (k)\right), \pi_ {b} (\cdot | k)\right). \tag {10}
$$
Emergence of the induction head via top-down learning. We begin our study by only training to minimize the loss on trigger-output token predictions after their first occurrence. This should be predictable with  $100\%$  accuracy using the two-layer induction head mechanism according to Section 4. We also remove the feed-forward layer, in order to focus on the learning of attention matrices  $W_{K}^{1}$ ,  $W_{K}^{2}$  and  $W_{O}^{2}$  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
Figure 3: Learning the induction head alone: in-context accuracy (top) and recall probes (bottom) with some layers frozen until iteration 300. The output matrix  $W_{O}^{2}$  can and must be learned before the key-query matrices, but does not suffice for good accuracy. It is easier to learn  $W_{K}^{2}$  before  $W_{K}^{1}$ , and  $W_{K}^{1}$  stores initial context positions ( $t < 64$ ) much faster than late positions.
learning unless  $W_{O}^{2}$  is learned first, and that  $W_{O}^{2}$  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_{O}^{2}$  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_{O}^{2}$  (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_{K}^{2}$  is easier to learn when  $W_{K}^{1}$  is frozen compared to the reverse, yet learning them together seems beneficial, possibly due to helpful feedback loops [1]. We also observe that  $W_{K}^{1}$  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_{O}^{2}$  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_{O}^{2}$  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_{O}^{2}$ , 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_{K}^{2}$ , 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
Figure 4: Global vs in-context learning and data-distributional effects. (left) Loss on global (dashed) vs in-context (solid) tokens throughout training, for fixed or random trigger tokens  $q_{k}$ . The red curves fix the trigger  $q_{1}$  to the most frequent token, while the fixed triggers in blue curves are less common. (center) In-context accuracy with different training and test distributions  $\pi_{o}$  for output tokens. Uniform leads to better generalization than global bigrams  $\pi_{b}$ . (right) Probe metrics throughout training:  $W_{O}^{2}$  and  $W_{F}$  eventually compete and deviate from our natural estimates.
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) \in \mathbb{R}^d \times [N]$ , and consider the following classification problem, with fixed output embeddings  $W_U$ :
$$
L (W) = \mathbb {E} _ {(x, y) \sim p} [ \ell (y, W _ {U} W x) ].
$$
The gradients take the following form: denoting  $\mu_k \coloneqq \mathbb{E}[x|y = k]$  and  $\hat{\mu}_k \coloneqq \mathbb{E}_x[\frac{\hat{p}_W(k|x)}{p(y = k)} x]$ ,
$$
\nabla_ {W} L (W) = \sum_ {k = 1} ^ {N} p (y = k) w _ {U} (k) \left(\hat {\mu} _ {k} - \mu_ {k}\right) ^ {\top}.
$$
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,  $\mu_{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 \in [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  $\eta$  starting from an initialization  $W_{0} = 0$  (so that  $\hat{p}_{W_0}(k|x) = 1 / N$ ):
$$
W _ {1} = \frac {\eta}{N} \sum_ {k = 1} ^ {N} w _ {U} (k) (\mu_ {k} - \bar {\mu}) ^ {\top},
$$
with  $\bar{\mu} = \mathbb{E}[x]$ . We show in Appendix B that when  $d$  is large enough to ensure near-orthonormal embeddings, we have
$$
w _ {U} (k) ^ {\top} W _ {1} \left(w _ {E} (y) + p _ {t}\right) \approx \frac {\eta}{N} \mathbb {1} \{k = y \} + O \left(\frac {1}{N ^ {2}}\right),
$$
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_{O}^{2}$ , then  $W_{K}^{2}$ , followed by  $W_{K}^{1}$ .
To show this result, we use Lemma 2 in a similar manner to the illustrative example above to show how training  $W_{O}^{2}$  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_{O}^{2}$  has learned the correct associations, we show that the gradient with respect to the key-value matrix  $W_{K}^{2}$  at zero initialization can leverage the correctness of  $W_{O}^{2}$  to find the right associative memory that focuses the attention on correct triggers. Finally, by linearizing the second-layer attention around  $W_{K}^{2} = 0$ , we show how gradients w.r.t.  $W_{K}^{1}$  may learn correct associations for the previous token head.
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 value-output matrices that may induce additional regularization effects [20], and non-linear feedforward layers, which may provide richer associative memories between sets of embeddings. Understanding how transformers leverage such aspects to learn in richer settings is an important next step.
The authors thank Sainbayar Sukhbaatar and Shubham Toshniwal for helpful discussions.
[1] Z. Allen-Zhu and Y. Li. Backward feature correction: How deep learning performs deep learning. arXiv preprint arXiv:2001.04413, 2020.
[2] Z. Allen-Zhu and Y. Li. Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. In Proceedings of the International Conference on Learning Representations (ICLR), 2023.
[3] J. Ba, M. A. Erdogdu, T. Suzuki, Z. Wang, D. Wu, and G. Yang. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems (NeurIPS), 2022.
[58] G. Yang and E. J. Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In Proceedings of the International Conference on Machine Learning (ICML), 2021.
[59] G. Yang, E. J. Hu, I. Babuschkin, S. Sidor, X. Liu, D. Farhi, N. Ryder, J. Pachocki, W. Chen, and J. Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
[60] B. Zhang and R. Sennrich. Root mean square layer normalization. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
[61] Y. Zhang, A. Backurs, S. Bubeck, 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}\in \mathbb{R}^{d}$  with i.i.d. Gaussian  $\mathcal{N}(0,\frac{1}{d})$  entries.
We recall a few facts:
- (Norm) We have  $u_i^\top u_i \approx 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 \neq j$ , we have  $u_i^\top u_j = O(1/\sqrt{d})$ . To see this, denoting  $u_i = d^{-1/2}(\tilde{u}_{ik})_k$ , where  $\tilde{u}_{ik}$  are the unnormalized entries of  $u_i$ , note that we have
$$
\sqrt {d} u _ {i} ^ {\top} u _ {j} = \frac {1}{\sqrt {d}} \sum_ {k = 1} ^ {d} \tilde {u} _ {i k} \tilde {u} _ {j k} \rightarrow \mathcal {N} (0, 1),
$$
by the central limit theorem, since for each  $k$ , the quantities  $\tilde{u}_{ik}\tilde{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.  $\mathcal{N}(0, \frac{1}{d})$  entries, then for any fixed  $x$  we have  $\| Wx \| \approx \| 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 \| \approx 1$ .
- (Remapping: near-orthogonality) Consider a random vector  $x = \frac{1}{\sqrt{d}} \tilde{x}$  and a random matrix  $W = \frac{1}{\sqrt{d}} \tilde{W}$ , where the entries of  $\tilde{x}$  and  $\tilde{W}$  are i.i.d.  $\mathcal{N}(0, 1)$ . Then  $x$  and  $Wx$  are nearly orthogonal. To see this, note that  $\mathbb{E}[x^\top Wx] = \mathbb{E}[x^\top \mathbb{E}[W]x] = 0$ , and the variance is
$$
\begin{array}{l} \mathbb {E} (x ^ {\top} W x) ^ {2} = \mathbb {E} \sum_ {i, j} x _ {i} ^ {2} W _ {i j} ^ {2} x _ {j} ^ {2} = \frac {1}{d ^ {3}} \mathbb {E} \sum_ {i} \tilde {x} _ {i} ^ {4} \tilde {W} _ {i i} ^ {2} + \mathbb {E} \sum_ {i \neq j} \tilde {x} _ {i} ^ {2} \tilde {x} _ {j} ^ {2} \tilde {W} _ {i j} ^ {2} \\ = \frac {1}{d ^ {3}} \left(d M _ {4} M _ {2} + \frac {d (d - 1)}{2} M _ {2} ^ {3}\right) = O \left(\frac {1}{d}\right), \\ \end{array}
$$
where  $M_2$  and  $M_4$  denote the second and fourth moments of the standard Gaussian, respectively. Then, Chebyshev's inequality implies that  $|x^\top W x| = O(1 / \sqrt{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 \in [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  $\eta$  starting from an initialization  $W_{0} = 0$  (so
that  $\hat{p}_{W_0}(k|x) = 1 / N$ :
$$
W _ {1} = \frac {\eta}{N} \sum_ {k = 1} ^ {N} w _ {U} (k) \left(\mu_ {k} - \bar {\mu}\right) ^ {\top}, \tag {11}
$$
with  $\bar{\mu} = \mathbb{E}[x]$
Note that we have  $\mu_{k} = w_{E}(k) + \frac{1}{T}\sum_{t}p_{t}$  and  $\bar{\mu} = \frac{1}{N}\sum_{k}w_{E}(k) + \frac{1}{T}\sum_{t}p_{t}$ , so that (11) becomes
$$
W _ {1} = \frac {\eta}{N} \sum_ {k = 1} ^ {N} w _ {U} (k) \left(w _ {E} (k) - \bar {w} _ {E}\right) ^ {\top}, \tag {12}
$$
with  $\bar{w}_E\coloneqq \frac{1}{N}\sum_{k = 1}^{N}w_E(k)$ . When  $d$  is large enough to ensure near-orthonormal embeddings, we have for any  $y$  and  $t$ ,
$$
W _ {1} (w _ {E} (y) + p _ {t}) \approx \frac {\eta}{N} w _ {U} (y) + O \left(\frac {1}{N ^ {2}}\right).
$$
This implies
as claimed in the main text. The classifier  $\hat{y} = \arg \max_{k} w_{U}(k)^{\top} 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.