SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We begin with the following lemma, which gives the gradient of the loss w.r.t.  $W = W_{K}^{2}$  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},\ldots ,x_{T}]$ , and value-output matrix  $\Phi_2 \coloneqq W_O^2 W_V^2$ :
$$
L (W) = \mathbb {E} _ {(X, y)} [ \ell (y, \xi (X)) ], \quad \xi (X) = W _ {U} \Phi_ {2} X \sigma \left(X ^ {\top} W x _ {T}\right), \tag {13}
$$
with  $\ell$  the cross-entropy loss and  $\sigma(u)_t = \frac{e^{u_t}}{\sum_s e^{u_s}}$  for  $u \in \mathbb{R}^T$  is the softmax.
The gradient at  $W = 0$  is given by
$$
\begin{array}{l} \left. \nabla_ {W} L (W) \right| _ {W = 0} = \sum_ {k = 1} ^ {N} \mathbb {E} _ {(X, y)} \left[ (\hat {p} _ {W} (k | X) - \mathbb {1} \{y = k \}) \frac {1}{T} \sum_ {t = 1} ^ {T} w _ {U} (k) ^ {\top} \Phi_ {2} x _ {t} \cdot (x _ {t} - \bar {x} _ {1: T}) x _ {T} ^ {\top} \right] \\ = \frac {1}{T} \sum_ {k = 1} ^ {N} \sum_ {t = 1} ^ {T} \mathbb {E} _ {X} \left[ \hat {p} _ {W} (k | X) w _ {U} (k) ^ {\top} \Phi_ {2} x _ {t} \cdot \left(x _ {t} - \bar {x} _ {1: T}\right) x _ {T} ^ {\top} \right] \\ - \frac {1}{T} \sum_ {k = 1} ^ {N} \sum_ {t = 1} ^ {T} p (y = k) \mathbb {E} _ {X} \left[ w _ {U} (k) ^ {\top} \Phi_ {2} x _ {t} \cdot \left(x _ {t} - \bar {x} _ {1: T}\right) x _ {T} ^ {\top} \mid y = k \right] \\ \end{array}
$$
with  $\bar{x}_{1:T} = \frac{1}{T}\sum_{t=1}^{T}x_t$ .
Now we consider the gradient w.r.t.  $W = W_{K}^{1}$  at zero initialization, and consider a simplification of the second layer attention to its linearization around  $W_{K}^{2} = 0$ . We will see that this still provides first-order information that is sufficient for  $W_{K}^{1}$  to be learned.
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_{K}^{1}$ , the key-query matrix at the first attention layer:
$$
L (W) = \mathbb {E} _ {(X, y)} [ \ell (y, \xi (X)) ], \quad \xi (X) = W _ {U} \Phi_ {2} X \bar {\sigma} (Z (W) ^ {\top} W _ {2} x _ {T}). \tag {14}
$$
Here,  $\bar{\sigma}(u_{1:T})_t = \frac{1}{T} (1 + u_t - \frac{1}{T}\sum_{s=1}^T u_s)$  is the linearization of the softmax around 0, and  $Z(W) = [z_1(W),\dots,z_T(W)]$  with
$$
z _ {t} (W) = \sum_ {s = 1} ^ {t} \Phi_ {1} x _ {s} \sigma (p _ {1: t} ^ {\top} W p _ {t}) _ {s},
$$
and  $\Phi_{\ell} = W_{O}^{\ell}W_{V}^{\ell}$  for  $\ell = 1,2$
$$
\begin{array}{l} \left. \nabla_ {W} L (W) \right| _ {W = 0} \\ = \sum_ {k = 1} ^ {N} \mathbb {E} _ {X} \left[ \hat {p} _ {W} (k | X) \frac {1}{T} \sum_ {t = 1} ^ {T} w _ {U} (k) ^ {\top} \Phi_ {2} x _ {t} \cdot \frac {1}{t} \sum_ {s = 1} ^ {t} (\Phi_ {1} x _ {s}) ^ {\top} W _ {2} x _ {T} (p _ {s} - \bar {p} _ {1: t}) p _ {t} ^ {\top} \right] \\ - \sum_ {k = 1} ^ {N} p (y = k) \mathbb {E} _ {X} \left[ \frac {1}{T} \sum_ {t = 1} ^ {T} w _ {U} (k) ^ {\top} \Phi_ {2} x _ {t} \cdot \frac {1}{t} \sum_ {s = 1} ^ {t} (\Phi_ {1} x _ {s}) ^ {\top} W _ {2} x _ {T} (p _ {s} - \bar {p} _ {1: t}) p _ {t} ^ {\top} | y = k \right] \\ - \sum_ {k = 1} ^ {N} \mathbb {E} _ {X} \left[ \hat {p} _ {W} (k | X) 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] \\ + \sum_ {k = 1} ^ {N} p (y = k) \mathbb {E} _ {X} \left[ 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} (\Phi_ {1} x _ {s}) ^ {\top} W _ {2} x _ {T} (p _ {s} - \bar {p} _ {1: t}) p _ {t} ^ {\top} | y = k \right] \\ \end{array}
$$
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  $\pi_u, \pi_q, \pi_o$  and  $\pi_b(\cdot | 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 occurrences 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_{O}^{2}$ ,  $W_{K}^{2}$  and  $W_{K}^{1}$ , 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_{O}^{2}$
We begin by studying the learning of the second output matrix  $W_{O}^{2}$ . 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
$$
x = \frac {1}{T} \sum_ {t = 1} ^ {T} W _ {V} ^ {2} (w _ {E} (z _ {t}) + \varepsilon_ {t}),
$$
where  $\varepsilon_{t} = p_{t} + \frac{1}{t}\sum_{s = 1}^{t}\Phi_{1}(w_{E}(z_{s}) + p_{s})$  with  $\Phi_1 = W_O^1 W_V^1$ , 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_{O}^{2}x + \varepsilon_{T})$  due to residual connections, but we ignore the term  $W_{U}\varepsilon_{T}$  for simplicity, noting that it is near-zero when  $d$  is large.
After a gradient step on  $W_{O}^{2}$  with step-size  $\eta$ , starting from zero-initialization (so that  $\hat{p}(k|x) = p(y = k) = 1/N$  for all  $x$ ), Lemma 2 yields
$$
W _ {O} ^ {2} = \frac {\eta}{N} \sum_ {k = 1} ^ {N} w _ {U} (k) \left(\mathbb {E} [ x | y = k ] - \mathbb {E} [ x ]\right) ^ {\top}. \tag {15}
$$
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(\cdot)$  or  $p(\cdot | y = k)$ , except for the token  $z_{t_{o}}$ .
Then, we may write:
$$
\begin{array}{l} \mathbb {E} [ x | y = k ] - \mathbb {E} [ x ] = \frac {1}{T} \left(\mathbb {E} [ W _ {V} ^ {2} w _ {E} (z _ {t _ {o}}) | y = k ] - \mathbb {E} [ W _ {V} ^ {2} w _ {E} (z _ {t _ {o}}) ]\right) \\ + \frac {1}{T} \left(\mathbb {E} \left[ \sum_ {t = t _ {o}} ^ {T} W _ {V} ^ {2} \varepsilon_ {t} | y = k \right] - \mathbb {E} \left[ \sum_ {t = t _ {o}} ^ {T} W _ {V} ^ {2} \varepsilon_ {t} \right]\right), \\ \end{array}
$$
since  $\varepsilon_{t}$  is independent of  $o$  when  $t < t_{o}$ . Noting that  $\varepsilon_{t}$  only depends on  $z_{t_o}$  and thus on  $y$  through the first layer attention, we have
$$
\begin{array}{l} \mathbb {E} [ x | y = k ] - \mathbb {E} [ x ] = \frac {1}{T} W _ {V} ^ {2} (w _ {E} (k) - \bar {w} _ {E}) \\ + \frac {1}{T} \left(\mathbb {E} \left[ \sum_ {t = t _ {o}} ^ {T} \frac {1}{t} W _ {V} ^ {2} \Phi_ {1} w _ {E} (z _ {t _ {o}}) | y = k \right] - \mathbb {E} \left[ \sum_ {t = t _ {o}} ^ {T} \frac {1}{t} W _ {V} ^ {2} \Phi_ {1} w _ {E} (z _ {t _ {o}}) \right]\right) \\ = \frac {1}{T} W _ {V} ^ {2} (w _ {E} (k) - \bar {w} _ {E}) + \frac {\tau}{T} W _ {V} ^ {2} \Phi_ {1} (w _ {E} (k) - \bar {w} _ {E}), \\ \end{array}
$$
where  $\tau := \mathbb{E}\left[\sum_{t = t_0}^T\frac{1}{t}\right]$ , and  $\bar{w}_E = \frac{1}{N}\sum_{k = 1}^{N}w_E(k)$ . Thus, (15) becomes
$$
W _ {O} ^ {2} = \frac {\eta}{N T} \sum_ {k = 1} ^ {N} w _ {U} (k) \left(W _ {V} ^ {2} \left(w _ {E} (k) - \bar {w} _ {E}\right)\right) ^ {\top} + \frac {\eta \tau}{N T} \sum_ {k = 1} ^ {N} w _ {U} (k) \left(W _ {V} ^ {2} \Phi_ {1} \left(w _ {E} (k) - \bar {w} _ {E}\right)\right) ^ {\top}, \tag {16}
$$
so that when  $d$  is learn enough to ensure near-orthonormal embeddings, we have
$$
w _ {U} (k) ^ {\top} W _ {O} ^ {2} W _ {V} ^ {2} w _ {E} (j) \approx \frac {\eta}{N T} \mathbb {1} \{k = j \} + O \left(\frac {\eta}{N ^ {2} T}\right)
$$
$$
w _ {U} (k) ^ {\top} W _ {O} ^ {2} W _ {V} ^ {2} \Phi_ {1} w _ {E} (j) \approx \frac {\eta \tau}{N T} \mathbb {1} \{k = j \} + O \left(\frac {\eta \tau}{N ^ {2} T}\right),
$$
where the  $O(\cdot)$  terms are due to the  $\bar{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_O^2$  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_{O}^{2}$  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_{K}^{2}$
Now assume that  $W_{O}^{2}$  is as in (16). As argued above, the predictions  $\hat{p}(k|x)$  are essentially random  $1/N$  in our model for  $W_{K}^{2} = 0$ , so that after one gradient step on  $W_{K}^{2}$  with learning rate  $\eta$ , Lemma 4 yields:
$$
W _ {K} ^ {2} = \frac {\eta}{T N} \sum_ {k, t} \left(\mathbb {E} \left[ w _ {U} (k) ^ {\top} \Phi_ {2} x _ {t} \cdot \left(x _ {t} - \bar {x}\right) x _ {T} ^ {\top} \mid y = k \right] - \mathbb {E} \left[ w _ {U} (k) ^ {\top} \Phi_ {2} x _ {t} \cdot \left(x _ {t} - \bar {x}\right) x _ {T} ^ {\top} \right]\right), \tag {17}
$$
where  $x_{t}$  are the inputs to the second attention layer, given by
$$
x _ {t} = x _ {t, 0} + x _ {t, 1} \tag {18}
$$
$$
x _ {t, 0} = w _ {E} \left(z _ {t}\right) + p _ {t} \tag {19}
$$
$$
x _ {t, 1} = \frac {1}{t} \sum_ {s = 1} ^ {t} \Phi_ {1} \left(w _ {E} \left(z _ {s}\right) + p _ {s}\right). \tag {20}
$$
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.,  $\pi_q = 1 / N$ ), we have
$$
\begin{array}{l} W _ {K} ^ {2} = \frac {\eta}{T N} \sum_ {k = 1} ^ {N} \sum_ {t = 1} ^ {T} \left(\mathbb {E} \left[ A _ {t, k} \mid y = k \right] - \mathbb {E} _ {X} \left[ A _ {t, k} \right]\right) (21) \\ = \frac {\eta}{T N ^ {2}} \sum_ {k = 1} ^ {N} \sum_ {t = 1} ^ {T} \sum_ {j = 1} ^ {N} \left(\mathbb {E} \left[ A _ {t, k} \mid y = k, q = j \right] - \mathbb {E} \left[ A _ {t, k} \mid q = j \right]\right) (22) \\ \end{array}
$$
$$
A _ {t, k} = w _ {U} (k) ^ {\top} \Phi_ {2} x _ {t, 0} \cdot \left(x _ {t, 1} - \bar {x} _ {1}\right) x _ {T, 0} ^ {\top}, \tag {23}
$$
with  $\bar{x}_1 = \frac{1}{T}\sum_t x_{t,1}$ . Now, note that we have  $w_U(k)^\top \Phi_2x_{t,0} \approx \alpha \mathbb{1}\{z_t = k\}$  with  $\alpha = \eta /TN$  by (16), and  $x_{T,0} = w_E(q) + p_T$ . This yields
$$
W _ {K} ^ {2} \approx \frac {\alpha \eta}{T N ^ {2}} \sum_ {j = 1} ^ {N} \sum_ {k = 1} ^ {N} \Delta_ {k, j} \left(w _ {E} (j) + p _ {T}\right) ^ {\top}, \tag {24}
$$
$$
\begin{array}{l} \Delta_ {k, j} := \mathbb {E} \left[ \sum_ {t = 1} ^ {T} \mathbb {1} \{z _ {t} = k \} \left(x _ {t, 1} - \bar {x} _ {1}\right) | y = k, q = j \right] - \mathbb {E} \left[ \sum_ {t = 1} ^ {T} \mathbb {1} \{z _ {t} = k \} \left(x _ {t, 1} - \bar {x} _ {1}\right) | q = j \right] \\ = \Delta_ {k, j} ^ {o} + \Delta_ {k, j} ^ {q} + \Delta_ {k, j} ^ {r}, \\ \end{array}
$$
where the three terms split the sum inside the expectation
$$
\Delta_ {k, j} ^ {o} := \mathbb {E} \left[ \mathbb {1} \left\{z _ {t _ {o}} = k \right\} \left(x _ {t _ {o}, 1} - \bar {x} _ {1}\right) | y = k, q = j \right] - \mathbb {E} \left[ \mathbb {1} \left\{z _ {t _ {o}} = k \right\} \left(x _ {t _ {o}, 1} - \bar {x} _ {1}\right) | q = j \right]
$$
$$
\Delta_ {k, j} ^ {q} := \mathbb {E} \left[ \sum_ {t \in \mathcal {T} _ {q}} \mathbb {1} \{z _ {t} = k \} (x _ {t, 1} - \bar {x} _ {1}) | y = k, q = j \right] - \mathbb {E} \left[ \sum_ {t \in \mathcal {T} _ {q}} \mathbb {1} \{z _ {t} = k \} (x _ {t, 1} - \bar {x} _ {1}) | q = j \right]
$$
$$
\Delta_ {k, j} ^ {r} := \mathbb {E} \left[ \sum_ {t \in \mathcal {T} _ {r}} \mathbb {1} \{z _ {t} = k \} (x _ {t, 1} - \bar {x} _ {1}) | y = k, q = j \right] - \mathbb {E} \left[ \sum_ {t \in \mathcal {T} _ {r}} \mathbb {1} \{z _ {t} = k \} (x _ {t, 1} - \bar {x} _ {1}) | q = j \right],
$$
where  $\mathcal{T}_q = \{t_o - 1,T\}$  and  $\mathcal{T}_r = [T]\setminus \{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  $\Delta_{k,j}^{o}$  carries the desired signal for the appropriate induction head associative memory, while  $\Delta_{k,j}^{q}$  and  $\Delta_{k,j}^{r}$  are negligible, for  $N$  large enough.
Controlling  $\Delta_{k,j}^{o}$ . For  $t = t_{o}$ , noting that  $z_{t_o} = y$ , we have
$$
\begin{array}{l} \Delta_ {k, j} ^ {o} = \mathbb {E} \left[ \mathbb {1} \left\{y = k \right\} \left(x _ {t _ {o}, 1} - \bar {x} _ {1}\right) | y = k, q = j \right] - \mathbb {E} \left[ \mathbb {1} \left\{y = k \right\} \left(x _ {t _ {o}, 1} - \bar {x} _ {1}\right) | q = j \right] \\ = \left(1 - \frac {1}{N}\right) \mathbb {E} \left[ x _ {t _ {o}, 1} - \bar {x} _ {1} | y = k, q = j \right] \\ = \frac {N - 1}{N} \left(\bar {p} + \sum_ {i = 1} ^ {N} a _ {k, j, i} \Phi_ {1} w _ {E} (i)\right), \\ \end{array}
$$
with  $a_{k,j,i} \approx (\Phi_1 w_E(i))^{\top} \mathbb{E}[x_{t_o,1} - \bar{x}_1 | y = k, q = j]$  thanks to near-orthonormality, and
$$
\bar {p} = \mathbb {E} _ {t _ {o}} \left[ \frac {1}{t _ {o}} \sum_ {s = 1} ^ {t _ {o}} p _ {s} - \frac {1}{T} \sum_ {t = 1} ^ {T} \frac {1}{t} \sum_ {s = 1} ^ {t} p _ {s} \right]
$$
is a spurious positional mixture.
We then distinguish the following cases:
- If  $j \neq k$  and  $i = j$ , since the trigger token  $j$  only appears at positions  $t_{o} - 1$  and  $T$ , we have
$$
a _ {k, j, i} \approx \mathbb {E} _ {t _ {o}} \left[ \frac {1}{t _ {o}} - \frac {1}{T} \sum_ {t = t _ {o} - 1} ^ {T} \frac {1}{t} - \frac {1}{T ^ {2}} \right] =: \gamma_ {T}.
$$
We may expect  $t_{o}$  to be concentrated around  $T / 2$ , in which case  $\gamma_{T} \gtrsim \frac{2}{T} - \frac{1}{T} - \frac{1}{T^{2}} \geq \frac{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
$$
a _ {k, j, i} \approx \frac {2}{T} - \frac {1}{T (T - 1)} - \frac {2}{T ^ {2}} = \frac {2}{T} + O \left(\frac {1}{T ^ {2}}\right),
$$
for  $T$  larger than a small constant.
- If  $i \neq j = k$ , all tokens up to position  $t_{o} - 2 = T - 2$  are i.i.d. uniform in  $[N] \setminus \{j\}$ , so that
$$
a _ {k, j, i} \approx \frac {T - 2}{T (N - 1)} - \frac {1}{T} \left((T - 2) \cdot \frac {1}{N - 1} + \frac {T - 2}{(T - 1) (N - 1)} + \frac {T - 2}{T (N - 1)}\right) = O \left(\frac {1}{N}\right)
$$
- If  $i \neq j$  and  $j \neq k$ , all tokens except at positions  $t_{o} - 1$ ,  $t_{o}$  and  $T$  (we have  $t_{o} < T$ ) are uniform in  $[N] \setminus \{j\}$ . The triggers do not contribute anything to  $a_{k,j,i}$  since  $i \neq j$ , and the output token may be also randomized by taking the average over  $k \in [N] \setminus \{j\}$ . We thus obtain
$$
\frac {1}{N - 1} \sum_ {k \neq j} a _ {k, j, i} \approx O \left(\frac {1}{N}\right).
$$
In summary, we obtain
$$
\frac {1}{N} \sum_ {k = 1} ^ {N} a _ {k, j, i} \approx \left\{ \begin{array}{l l} O \left(\frac {1}{N}\right), & \text {i f} i \neq j \\ \Omega \left(\frac {1}{T}\right), & \text {i f} i = j. \end{array} \right.
$$
Thus, when  $N$  is large, while  $T$  is moderate, the above sum leads to more signal in the  $i = j$  terms compared to  $i \neq j$ . In particular, this yields
$$
\big (\Phi_ {1} w _ {E} (i) \big) ^ {\top} \left(\frac {1}{N} \sum_ {k = 1} ^ {N} \Delta_ {k, j} ^ {o}\right) \approx \left\{ \begin{array}{l l} O \left(\frac {1}{N}\right), & \text {i f} i \neq j \\ \Omega \left(\frac {1}{T}\right), & \text {i f} i = j, \end{array} \right.
$$
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:
$$
\frac {1}{N} \sum_ {k = 1} ^ {N} \left(\Delta_ {k, j} ^ {q} + \Delta_ {k, j} ^ {r}\right) \approx o \left(\frac {1}{T}\right).
$$
Controlling  $\Delta_{k,j}^{q}$ . For  $t\in \mathcal{T}_q$ , note that we always have  $z_{t} = j$  in the expectations, so that  $\Delta_{k,j}^{q} = 0$  unless  $k = j$ . For  $k = j$ , we have  $\Delta_{k,j}^{q} = O(1)$ , so that
$$
\frac {1}{N} \sum_ {k = 1} ^ {N} \Delta_ {k, j} ^ {q} = O \left(\frac {1}{N}\right).
$$
Controlling  $\Delta_{k,j}^{r}$ . Using that  $\| x_{t,1} - \bar{x}_1\| \leq C = O(1)$  for all  $t$ , we provide the following crude bound via the triangle inequality and Hölder inequality:
$$
\Delta_ {k, j} ^ {r} = \mathbb {E} \left[ \sum_ {t \in \mathcal {T} _ {r}} \mathbb {1} \{z _ {t} = k \} (x _ {t, 1} - \bar {x} _ {1}) | y = k, q = j \right] - \mathbb {E} \left[ \sum_ {t \in \mathcal {T} _ {r}} \mathbb {1} \{z _ {t} = k \} (x _ {t, 1} - \bar {x} _ {1}) | q = j \right]
$$
$$
\| \Delta_ {k, j} ^ {r} \| \leq C \left(\mathbb {E} \left[ \sum_ {t \in \mathcal {T} _ {r}} \mathbb {1} \{z _ {t} = k \} | y = k, q = j \right] + \mathbb {E} \left[ \sum_ {t \in \mathcal {T} _ {r}} \mathbb {1} \{z _ {t} = k \} | q = j \right]\right) \leq \frac {2 C T}{N},
$$
since  $z_{t}$  is independent of  $y$  given  $t \in \mathcal{T}_r$  and thus is uniform in  $[N] \setminus \{j\}$ , and  $|\mathcal{T}_r| \leq T$ . We note, however, that  $\Delta_{k,j}^{r}$  may be controlled much more finely by leveraging the similarities between the distributions of  $z_{t}, t \in \mathcal{T}_{r}$  with or without conditioning on  $y$ .
Overall, we have shown that up to some spurious positional embeddings,  $W_{K}^{2}$  behaves as the desired associative memory from (7) when  $N$  is large enough, satisfying:
$$
\left(\Phi_ {1} w _ {E} (i)\right) ^ {\top} W _ {K} ^ {2} w _ {E} (j) \approx \frac {\alpha \eta}{T N} \left\{\Omega \left(\frac {1}{T}\right) \mathbb {1} \{i = j \} + O \left(\frac {T}{N}\right) \right\} \tag {25}
$$
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_{K}^{1}$
We now assume that  $W_{O}^{2}$  and  $W_{K}^{2}$  have learned the correct associations, and consider one gradient step away from zero-initialization on  $W_{K}^{1}$ . Note that when  $W_{K}^{1} = 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  $\Phi_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  $\hat{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_{K}^{1}$  after one gradient step of step-size  $\eta$ :
$$
\begin{array}{l} W _ {K} ^ {1} = \frac {\eta}{N} \sum_ {k = 1} ^ {N} \mathbb {E} _ {X} \left[ \frac {1}{T} \sum_ {t = 1} ^ {T} w _ {U} (k) ^ {\top} \Phi_ {2} x _ {t} \cdot \frac {1}{t} \sum_ {s = 1} ^ {t} (\Phi_ {1} x _ {s}) ^ {\top} W _ {K} ^ {2} x _ {T} (p _ {s} - \bar {p} _ {1: t}) p _ {t} ^ {\top} | y = k \right] \\ - \frac {\eta}{N} \sum_ {k = 1} ^ {N} \mathbb {E} _ {X} \left[ \frac {1}{T} \sum_ {t = 1} ^ {T} w _ {U} (k) ^ {\top} \Phi_ {2} x _ {t} \cdot \frac {1}{t} \sum_ {s = 1} ^ {t} \left(\Phi_ {1} x _ {s}\right) ^ {\top} W _ {K} ^ {2} x _ {T} \left(p _ {s} - \bar {p} _ {1: t}\right) p _ {t} ^ {\top} \right] \\ - \frac {\eta}{N} \sum_ {k = 1} ^ {N} \mathbb {E} _ {X} \left[ 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 _ {K} ^ {2} x _ {T} \left(p _ {s} - \bar {p} _ {1: t}\right) p _ {t} ^ {\top} | y = k \right] \\ + \frac {\eta}{N} \sum_ {k = 1} ^ {N} \mathbb {E} _ {X} \left[ 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} (\Phi_ {1} x _ {s}) ^ {\top} W _ {K} ^ {2} x _ {T} (p _ {s} - \bar {p} _ {1: t}) p _ {t} ^ {\top} \right]. \\ \end{array}
$$
Note that since  $W_{O}^{2}$  and  $W_{K}^{2}$  already captured the desired associations at this stage, we have
$$
w _ {U} (k) ^ {\top} \Phi_ {2} x _ {t} \approx \alpha   \mathbb {1} \{z _ {t} = k \} \quad \text {a n d} \quad (\Phi_ {1} x _ {s}) ^ {\top} W _ {K} ^ {2} x _ {T} \approx \alpha^ {\prime}   \mathbb {1} \{z _ {s} = z _ {T} \},
$$
for some  $\alpha, \alpha' > 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$ ,
$$
W _ {K} ^ {1} p _ {t} \approx \frac {\eta \alpha \alpha^ {\prime}}{N T t} \sum_ {k = 1} ^ {N} (A _ {t, k} - B _ {t, k} - C _ {t, k} + D _ {t, k}),
$$
$$
A _ {t, k} = \mathbb {E} \left[ \mathbb {1} \left\{z _ {t} = k \right\} \mathbb {1} \left\{t _ {q} \leq t \right\} \left(p _ {t _ {q}} - \bar {p} _ {1: t}\right) | y = k \right] \tag {26}
$$
$$
B _ {t, k} = \mathbb {E} \left[ \mathbb {1} \left\{z _ {t} = k \right\} \mathbb {1} \left\{t _ {q} \leq t \right\} \left(p _ {t _ {q}} - \bar {p} _ {1: t}\right) \right] \tag {27}
$$
$$
C _ {t, k} = \mathbb {E} \left[ r _ {k} \mathbb {1} \left\{t _ {q} \leq t \right\} \left(p _ {t _ {q}} - \bar {p} _ {1: t}\right) | y = k \right] \tag {28}
$$
$$
D _ {t, k} = \mathbb {E} \left[ r _ {k} \mathbb {1} \left\{t _ {q} \leq t \right\} \left(p _ {t _ {q}} - \bar {p} _ {1: t}\right) \right], \tag {29}
$$
where  $r_k \coloneqq \frac{1}{T} \sum_{t=1}^{T} \mathbb{1}\{z_t = k\}$ . We have