SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} A _ {t, k} = \mathbb {E} [ \mathbb {1} \{z _ {t} = k \} (\mathbb {1} \{t _ {q} = t - 1 \} + \mathbb {1} \{t _ {q} \in [ t - 2 ] \cup \{t \}) \}) (p _ {t _ {q}} - \bar {p} _ {1: t}) | y = k ] \\ = \mathbb {P} (t _ {q} = t - 1 | y = k) (p _ {t - 1} - \bar {p} _ {1: t}) + \frac {1}{N} \sum_ {s \in [ t - 2 ] \cup \{t \}} \mathbb {P} (t _ {q} = s | y = k) (p _ {s} - \bar {p} _ {1: t}) \\ = \mathbb {P} (t _ {q} = t - 1) \left(p _ {t - 1} - \bar {p} _ {1: t}\right) + \frac {1}{N} \sum_ {s \in [ t - 2 ] \cup \{t \}} \mathbb {P} (t _ {q} = s) \left(p _ {s} - \bar {p} _ {1: t}\right) \\ = \mathbb {P} (t _ {q} = t - 1) \left(p _ {t - 1} - \bar {p} _ {1: t}\right) + O \left(\frac {1}{N}\right), \\ \end{array}
$$
since the distribution of  $t_q$  is the same regardless of  $y$ . We proceed similarly for the other quantities and obtain the following:
$$
B _ {t, k} = O \left(\frac {1}{N}\right)
$$
$$
C _ {t, k} = \frac {\mathbb {P} (t _ {q} = t - 1)}{T} (p _ {t - 1} - \bar {p} _ {1: t}) + O \left(\frac {1}{N}\right)
$$
$$
D _ {t, k} = O \left(\frac {1}{N}\right).
$$
This yields the following associative memory behavior, for  $t > 1$ :
$$
p _ {s} ^ {\top} W _ {K} ^ {1} p _ {t} \approx \frac {\eta \alpha \alpha^ {\prime} (T - 1)}{T ^ {2} t} \left\{\mathbb {P} (t _ {q} = t - 1) \left(\mathbb {1} \{s = t - 1 \} - \frac {1}{t} \mathbb {1} \{s \in [ t ] \}\right) + O \left(\frac {1}{N}\right) \right\},
$$
which matches the desired "previous token head" behavior from (7) when  $N$  is large. As in the case of  $W_{K}^{2}$ , 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:
$$
\ell (y, \xi) = - \sum_ {k = 1} ^ {N} \mathbb {1} \{y = k \} \log \frac {e ^ {\xi_ {k}}}{\sum_ {j} e ^ {\xi_ {j}}}.
$$
Its derivatives take the form
$$
\frac {\partial \ell}{\partial \xi_ {k}} (y, \xi) = s (\xi) _ {k} - \mathbb {1} \{y = k \},
$$
with  $s(\xi)_k = \frac{e^{\xi_k}}{\sum_j e^{\xi_j}}$  the softmax.
The gradient of  $L$  is then given by
$$
\begin{array}{l} \nabla_ {W} L (W) = \mathbb {E} _ {(z, y)} \left[ \sum_ {k = 1} ^ {N} \frac {\partial \ell}{\partial \xi_ {k}} (y, W _ {U} W w _ {E} (z)) \nabla_ {W} (w _ {U} (k) ^ {\top} W w _ {E} (z)) \right] \\ = \mathbb {E} _ {(z, y)} \left[ \sum_ {k = 1} ^ {N} (\hat {p} _ {W} (k | z) - \mathbb {1} \{y = k \}) w _ {U} (k) w _ {E} (z) ^ {\top} \right] \\ = \sum_ {k = 1} ^ {N} \mathbb {E} _ {z} \left[ \mathbb {E} _ {y} \left[ (\hat {p} _ {W} (k | z) - \mathbb {1} \{y = k \}) w _ {U} (k) w _ {E} (z) ^ {\top} \mid z \right] \right] \\ = \sum_ {k = 1} ^ {N} \mathbb {E} _ {z} [ (\hat {p} _ {W} (k | z) - \mathbb {E} _ {y} [ \mathbb {1} \{y = k \} | z ]) w _ {U} (k) w _ {E} (z) ^ {\top} ], \\ \end{array}
$$
which yields the desired result.
C.2 Proof of Lemma 2
Proof. Using similar steps as the proof of Lemma 1, we have
$$
\begin{array}{l} \nabla_ {W} L (W) = \mathbb {E} _ {(x, y)} \left[ \sum_ {k = 1} ^ {N} \frac {\partial \ell}{\partial \xi_ {k}} (y, W _ {U} W x) \nabla_ {W} (w _ {U} (k) ^ {\top} W x) \right] \\ = \mathbb {E} _ {(x, y)} \left[ \sum_ {k = 1} ^ {N} (\hat {p} _ {W} (k | x) - \mathbb {1} \{y = k \}) w _ {U} (k) x ^ {\top} \right] \\ = \sum_ {k = 1} ^ {N} w _ {U} (k) \mathbb {E} _ {x} [ \hat {p} _ {W} (k | x) x ] ^ {\top} - \sum_ {k = 1} ^ {N} \mathbb {E} _ {y} [ \mathbb {1} \{y = k \} w _ {U} (k) \mathbb {E} [ x | y ] ^ {\top} ] \\ = \sum_ {k = 1} ^ {N} w _ {U} (k) \mathbb {E} _ {x} [ \hat {p} _ {W} (k | x) x ] ^ {\top} - \sum_ {k, j = 1} ^ {N} p (y = j) \mathbb {1} \{j = k \} w _ {U} (k) \mathbb {E} [ x | y = j ] ^ {\top} \\ = \sum_ {k = 1} ^ {N} p (y = k) w _ {U} (k) (\hat {\mu} _ {k} - \mu_ {k}) ^ {\top}, \\ \end{array}
$$
with  $\hat{\mu}_k = p(y = k)^{-1}\mathbb{E}_x[\hat{p}_W(k|x)x]$  and  $\mu_k = \mathbb{E}[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},\ldots ,z_{T}]\in \mathbb{R}^{d\times 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:
$$
\begin{array}{l} \left. \left. \nabla_ {W} L (W) \right| _ {W = 0} = \mathbb {E} _ {(X, Z, y)} \left[ \sum_ {k = 1} ^ {N} \frac {\partial \ell}{\partial \xi_ {k}} (y, \xi) \cdot \nabla_ {W} \left(w _ {U} (k) ^ {\top} \Phi_ {2} X \sigma \left(Z ^ {\top} W x _ {T}\right)\right) \right| _ {W = 0} \right] (30) \\ = \mathbb {E} _ {(X, Z, y)} \left[ \sum_ {k = 1} ^ {N} \left(\hat {p} _ {W} (k | X, Z) - \mathbb {1} \{y = k \}\right) \cdot \nabla_ {W} \left(w _ {U} (k) ^ {\top} \Phi_ {2} X \sigma \left(Z ^ {\top} W x _ {T}\right)\right) \mid_ {W = 0} \right]. (31) \\ \end{array}
$$
$$
\begin{array}{l} \left. \nabla_ {W} \left(w _ {U} (k) ^ {\top} \Phi_ {2} X \sigma \left(Z ^ {\top} W x _ {T}\right)\right) \right| _ {W = 0} = \sum_ {t = 1} ^ {T} w _ {U} (k) ^ {\top} \Phi_ {2} x _ {t} \cdot \nabla_ {W} \left(\sigma \left(Z ^ {\top} W x _ {T}\right) _ {t}\right) \\ = \frac {1}{T} \sum_ {t = 1} ^ {T} w _ {U} (k) ^ {\top} \Phi_ {2} x _ {t} \cdot (z _ {t} - \bar {z} _ {1: T}) x _ {T} ^ {\top}, \\ \end{array}
$$
where  $\bar{z}_{1:T} = \frac{1}{T}\sum_{t}z_{t}$ , and we used the fact that
$$
\left. \frac {\partial}{\partial u _ {s}} \sigma (u) _ {t} \right| _ {u = 0} = \frac {1}{T} \mathbb {1} \{t = s \} - \frac {1}{T ^ {2}}. \tag {32}
$$
The gradient (31) now writes
$$
\left. \nabla_ {W} L (W) \right| _ {W = 0} = \sum_ {k = 1} ^ {N} \mathbb {E} _ {(X, Z)} [ (\hat {p} _ {W} (k | X, Z) - \mathbb {1} \{y = k \}) \frac {1}{T} \sum_ {t = 1} ^ {T} w _ {U} (k) ^ {\top} \Phi_ {2} x _ {t} \cdot (z _ {t} - \bar {z} _ {1: T}) x _ {T} ^ {\top} ],
$$
C.4 Proof of Lemma 5
Proof. The linearization of the second layer softmax around zero takes the following form:
$$
\bar {\sigma} (Z ^ {\top} W _ {2} x _ {T}) _ {t} = \frac {1}{T} (1 + z _ {t} ^ {\top} W _ {2} x _ {T} - \bar {z} _ {1: T} ^ {\top} W _ {2} x _ {T}),
$$
with  $z_{t} = \sum_{s = 1}^{t}\Phi_{1}x_{s}\sigma (p_{1:t}^{\top}Wp_{t})_{s}$  the output of the first attention layer.
$$
\begin{array}{l} \xi_ {k} = \sum_ {t = 1} ^ {T} w _ {U} (k) ^ {\top} \Phi_ {2} x _ {t} \bar {\sigma} (Z ^ {\top} W _ {2} x _ {T}) \\ = \frac {1}{T} \sum_ {t = 1} ^ {T} w _ {U} (k) ^ {\top} \Phi_ {2} x _ {t} + \frac {1}{T} \sum_ {t = 1} ^ {T} w _ {U} (k) ^ {\top} \Phi_ {2} x _ {t} \sum_ {s = 1} ^ {t} (\Phi_ {1} x _ {s}) ^ {\top} W _ {2} x _ {T} \sigma (p _ {1: t} ^ {\top} W p _ {t}) _ {s} \\ - w _ {u} (k) ^ {\top} \Phi_ {2} \bar {x} _ {1: T} \cdot \frac {1}{T} \sum_ {t = 1} ^ {T} \sum_ {s = 1} ^ {t} \left(\Phi_ {1} x _ {s}\right) ^ {\top} W _ {2} x _ {T} \sigma \left(p _ {1: t} ^ {\top} W p _ {t}\right) _ {s}. \\ \end{array}
$$
$$
\begin{array}{l} \left. \nabla_ {W} L (W) \right| _ {W = 0} (33) \\ = \mathbb {E} _ {(X, y)} \left[ \sum_ {k = 1} ^ {N} \frac {\partial \ell}{\partial \xi_ {k}} (y, \xi) \cdot \left. \nabla_ {W} \xi_ {k} \right| _ {W = 0} \right] (34) \\ = \mathbb {E} _ {(X, y)} \left[ \sum_ {k = 1} ^ {N} \frac {\partial \ell}{\partial \xi_ {k}} (y, \xi) \frac {1}{T} \sum_ {t = 1} ^ {T} w _ {U} (k) ^ {\top} \Phi_ {2} x _ {t} \frac {1}{t} \sum_ {s = 1} ^ {t} \left(\Phi_ {1} x _ {s}\right) ^ {\top} W _ {2} x _ {T} \left(p _ {s} - \bar {p} _ {1: t}\right) p _ {t} ^ {\top} \right] (35) \\ - \mathbb {E} _ {(X, y)} \left[ \sum_ {k = 1} ^ {N} \frac {\partial \ell}{\partial \xi_ {k}} (y, \xi) w _ {U} (k) ^ {\top} \Phi_ {2} \bar {x} _ {1: T} \cdot \frac {1}{T} \sum_ {t = 1} ^ {T} \frac {1}{t} \sum_ {s = 1} ^ {t} \left(\Phi_ {1} x _ {s}\right) ^ {\top} W _ {2} x _ {T} \left(p _ {s} - \bar {p} _ {1: t}\right) p _ {t} ^ {\top} \right], (36) \\ \end{array}
$$
using (32). The result follows by using  $\frac{\partial\ell}{\partial\xi_k} (y,\xi) = \hat{p} (k|\xi) - \mathbb{1}\{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}^{\top}W_{Q}$ , and the output-value matrices  $W_{O}$  and  $W_{V}$  are typically jointly optimized.
Consider an associative memory of the form:
$$
W _ {*} = \sum_ {i} y _ {i} x _ {i} ^ {\top} \in \mathbb {R} ^ {d \times d},
$$
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 = \frac{d}{2d'} UV$ , where  $U \in \mathbb{R}^{d \times d'}$ ,  $V \in \mathbb{R}^{d' \times d}$  with  $d' \leq d$  (for instance  $d'$  could be the dimension of attention heads), are given by:
$$
U = U _ {0} + \sum_ {i} y _ {i} (V _ {0} x _ {i}) ^ {\top}
$$
$$
V = V _ {0} + \sum_ {i} (U _ {0} ^ {\top} y _ {i}) x _ {i} ^ {\top},
$$
where  $U_{0}$  and  $V_{0}$  are random matrices with  $\mathcal{N}(0, \frac{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
$$
U V = U _ {0} V _ {0} + \sum_ {i} y _ {i} (V _ {0} x _ {i}) ^ {\top} V _ {0} + \sum_ {i} (U _ {0} ^ {\top} y _ {i}) x _ {i} ^ {\top} + \sum_ {i, j} y _ {i} (V _ {0} x _ {i}) ^ {\top} (U _ {0} ^ {\top} y _ {j}) x _ {j} ^ {\top}.
$$
It is also easy to check using central limit arguments (similar to remapping in Appendix A) that  $\tilde{x}_i := \sqrt{\frac{d}{d'}} V_0 x_i \in \mathbb{R}^{d'}$  and  $\tilde{y}_i := \sqrt{\frac{d}{d'}} U_0^\top y_i$  are all nearly-orthonormal embeddings. Thus, we have
$$
\begin{array}{l} 2 y _ {k} ^ {\top} W x _ {l} = \tilde {y} _ {k} ^ {\top} \tilde {x} _ {l} + \sum_ {i} y _ {k} ^ {\top} y _ {i} \tilde {x} _ {i} ^ {\top} \tilde {x} _ {l} + \sum_ {i} \tilde {y} _ {k} ^ {\top} \tilde {y} _ {i} x _ {i} ^ {\top} x _ {l} + \sum_ {i, j} y _ {k} ^ {\top} y _ {i} \tilde {x} _ {i} ^ {\top} \tilde {y} _ {j} x _ {j} ^ {\top} x _ {l} \\ \approx 0 + \mathbb {1} \{k = l \} + \mathbb {1} \{k = l \} + 0, \\ \end{array}
$$
where the first and last term vanish due to the cross-terms  $\tilde{y}_i^\top \tilde{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
$$
L N (x) = \frac {x}{\| x \|},
$$
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:
$$
L (W) = \mathbb {E} _ {(x, y) \sim p} [ \ell (y, W _ {U} L N (W x)) ]. \tag {37}
$$
The gradients then take the form
$$
\nabla_ {W} L (W) = \sum_ {k = 1} ^ {N} \mathbb {E} _ {x} \left[ \frac {\hat {p} _ {W} (y = k | x) - p (y = k | x)}{\| W x \|} \left(I - \frac {(W x) (W x) ^ {\top}}{\| W x \| ^ {2}}\right) w _ {U} (k) x ^ {\top} \right]. \tag {38}
$$
3When  $d^{\prime}$  is the head dimension, the  $\frac{d}{d^{\prime}}$  scaling can be interpreted as the correct multiplier to use in attention logits, which plays a similar role to the  $\frac{1}{d^{\prime}}$  multiplier in the  $\mu \mathrm{P}$  scaling [59], for our setup where the variance of the random entries of input embeddings is  $1 / d$  instead of 1 as in [59].
$^4$ RMSNorm [60] would use the variance instead of the norm in the denominator, leading to an additional  $\sqrt{d}$  factor in the numerator. Here we use the norm, which is more natural when embeddings have near-unit norm, in contrast to the  $\approx \sqrt{d}$  norm for the standard parameterization.
This illustrates that in addition to weighting the updates on class  $k$  by the prediction error  $\hat{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^\top$  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 \approx 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)\in \mathcal{M}$  as follows:
$$
F (x) = \sum_ {(i, j) \in \mathcal {M}} v _ {j} \sigma (u _ {i} ^ {\top} x - b),
$$
where  $\sigma$  is a non-linear activation and  $b$  a bias term. Then, if one assumes that  $u_{i}^{\top}u_{j}\leq b$  for  $i\neq j$  and  $\sigma (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_{\mathcal{I}} = \frac{1}{\sqrt{\left|\mathcal{I}\right|}}\sum_{i\in \mathcal{I}}u_i$ , where  $\mathcal{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_{\mathcal{I}}^{\top}u_{\mathcal{I}}\approx 1$  while  $u_{\mathcal{I}}^{\top}u_{\mathcal{I}'}\lesssim 1 - \delta$  if  $\mathcal{I} = \mathcal{I}'$ , for some  $\delta$  that depends on the maximal cardinality of these sets. Although  $1 - \delta$  is no-longer vanishingly small, defining
$$
F (x) = \sum_ {(\mathcal {I}, \mathcal {J}) \in \mathcal {M}} v _ {\mathcal {J}} \sigma \left(u _ {\mathcal {I}} ^ {\top} x - b\right),
$$
the non-linearity may still succeed at filtering out any  $\mathcal{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  $\eta$ . We fix the momentum parameter to 0.9 and the weight decay parameter to  $10^{-4}$ .  $U$  denotes the uniform distribution over  $[N]$ .
Figure 5: Memory recall probes for the setting of Figure 4(left).
Figure 6: Memory recall probes for the setting of Figure 4(center).
- Figure 2:  $K = 3$ ,  $\pi_q = \pi_u$  (random triggers) or  $Q$  is the  $K$  most likely elements of  $\pi_u$ ,  $\pi_o = U$ ,  $d = 128$ ,  $d_{hidden} = 4 \times 128$  (hidden dimension of the feed-forward MLPs),  $\eta = 0.2$ .
- Figure 3:  $K = 5$ ,  $\pi_q = \pi_u$  (random triggers),  $\pi_o = U$ ,  $d = 128$ ,  $\eta = 0.2$ .
- Figure 4(left) and Figure 5:  $\pi_{o} = U$ ,  $d = 128$ ,  $\eta = 1$ . For random triggers we use  $\pi_{q} = \pi_{u}$ . For  $K = 1$  with fixed frequent trigger, the only trigger is the most probable token according to  $\pi_{u}$ , while for  $K = 5$  with fixed rare triggers, the five triggers are the 6-th to 10-th most probable tokens according to  $\pi_{u}$ .
- Figure 4(center):  $K = 3$ ,  $\pi_q = \pi_u$  (random triggers),  $\pi_o = U$  or  $\pi_o = \pi_b$  (conditioned on the trigger),  $d = 128$ ,  $\eta = 1$ .
- Figure 4(right):  $K = 3$ ,  $\pi_q = \pi_u$  (random triggers),  $\pi_o = U$ ,  $d = 128$ ,  $\eta = 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_{K}^{2}$  with fixed triggers only shows recall accuracy on the set of triggers  $Q$ , which is an easier task). Figure 6
Figure 7: Visualization of the weights  $W_{K}^{1}$  (left),  $W_{K}^{2}$  (center), and  $W_{O}^{2}$  (right) after training with random triggers,  $K = 3$ ,  $\pi_q = \pi_u$ ,  $\pi_o = U$ . For each of these weight matrices  $W$ , if we write the corresponding target memory in (7) as  $W_{*} = \sum_{i} v_{i} u_{i}^{\top}$  with appropriate embeddings  $(u_i)_i$  and  $(v_i)_i$  (for instance  $u_t = p_{t-1}$  and  $v_t = p_t$  for  $W_{K}^{1}$  on the left), and we show all values  $v_j^{\top} W u_i$ .
Figure 8: Effect of dimension on learning  $W_{O}^{2}$  alone, with fixed or tuned learning rate.
Figure 9: Accuracy of one-step estimate of  $W_{Q}^{2}$  with varying dimension and number of batches used for computing expectations. Each batch consists of 32 sequences of 256 tokens for a total of 8192 tokens, with  $K = 5$  random triggers and uniform outputs.
shows that using uniform output tokens can lead to better fitting of  $W_{O}^{2}$  and  $W_{K}^{2}$  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_{K}^{1}$  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_K^1$ , and for some of the trigger tokens in  $W_K^2$  while the diagonal for  $W_O^2$  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_K^2$  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_K^1$  learns stronger associations at earlier positions. In contrast, diagonal elements for  $W_O^2$  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_O^2$  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_O^2$  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
Figure 10: Training of a more realistic architecture with (i) ReLU MLP instead of linear layer for the second feed-forward layer, (ii) all parameters trained, including embeddings, (iii) pre-layer normalization. The loss, in-context accuracy and memory recall probes are similar to the simplified architecture (see, e.g., Figure 4).
Figure 11: Attention maps for a two-layer model with 4 attention heads. In the first layer (top), the previous token mechanism is mostly achieved by one of the four heads, while the induction behavior at the second layer (bottom) is distributed across the different heads.
attention input to  $W_{O}^{2}$ , given by
$$
R _ {1} = \frac {1}{N} \sum_ {k = 1} ^ {N} \mathbb {1} \left\{k = \arg \max  _ {k ^ {\prime}} \left(W _ {V} ^ {2} w _ {E} \left(k ^ {\prime}\right)\right) ^ {\top} \left(\mu_ {k} - \mu\right) \right\},
$$
where  $\mu_{k} = \mathbb{E}[x|y = k]$  and  $\mu = \mathbb{E}[x]$ , for  $x = \frac{1}{t}\sum_{s=1}^{t}W_{V}^{2}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_{O}^{2}$  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.