SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \mathbb {E} [ W _ {k k ^ {\prime}} ^ {B} W _ {j j ^ {\prime}} ^ {B ^ {\prime}} ] = \sum_ {i i ^ {\prime}} ^ {n _ {k}} \mathbb {E} \left[ W _ {k i} ^ {K} W _ {j i ^ {\prime}} ^ {K} W _ {k ^ {\prime} i} ^ {Q} W _ {j ^ {\prime} i ^ {\prime}} ^ {Q} \right] \\ = \sum_ {i i ^ {\prime}} \mathbb {E} \left[ W _ {k i} ^ {K} W _ {j i ^ {\prime}} ^ {K} \right] \mathbb {E} \left[ W _ {k ^ {\prime} i} ^ {Q} W _ {j ^ {\prime} i ^ {\prime}} ^ {Q} \right] \\ = \sum_ {i i ^ {\prime}} \delta_ {k j} \delta_ {k ^ {\prime} j ^ {\prime}} \delta_ {i i ^ {\prime}} \\ = n _ {k} \delta_ {k j} \delta_ {k ^ {\prime} j ^ {\prime}}, \\ \end{array}
$$
where we recall  $W^{B} = W^{K}W^{Q,\top}$ .
$$
\mathbb {E} _ {\ell} \left[ Y ^ {\alpha \beta} Y ^ {\delta \omega} \right] = \frac {n _ {k}}{n ^ {2}} \sum_ {k k ^ {\prime}} X _ {k} ^ {\alpha} X _ {k ^ {\prime}} ^ {\beta} X _ {k} ^ {\delta} X _ {k ^ {\prime}} ^ {\omega} = n _ {k} V ^ {\alpha \delta} V ^ {\beta \omega}.
$$
C.2 Shaping the Softmax
Recall that we have the following model:
$$
X _ {\ell + 1} = \lambda X _ {\ell} + \gamma A _ {\ell} X _ {\ell} \frac {1}{\sqrt {n}} W _ {\ell} ^ {V}
$$
For the above model,  $V_{\ell}$  has the following form:
$$
V _ {\ell + 1} = \lambda^ {2} V _ {\ell} + \frac {\lambda \gamma}{n \sqrt {n}} \left(X _ {\ell} W _ {\ell} ^ {\top} X _ {\ell} ^ {\top} A _ {\ell} ^ {\top} + A _ {\ell} X _ {\ell} W _ {\ell} X _ {\ell} ^ {\top}\right) + \frac {\gamma^ {2}}{n ^ {2}} A _ {\ell} X _ {\ell} W _ {\ell} W _ {\ell} ^ {\top} X _ {\ell} ^ {\top} A _ {\ell} ^ {\top}.
$$
In order to have infinitesimal updates in  $V_{\ell}$ , we would intuitively like the attention matrix to be of the form  $A^{\alpha \beta} = \delta_{\alpha \beta} + \mathcal{O}(n^{-1})$ . In the case of the usual Softmax based attention, we have:
$$
A _ {\ell} := \operatorname {S o f t m a x} \left(\tau^ {- 1} Y _ {\ell}\right),
$$
where  $\tau^{-1}$  is a temperature parameter that regulates the entropy of the resulting categorical distribution;  $\tau$  is often chosen to be large (scale as a power of  $n$  or  $d$ ), as low temperature results in unstable training.
To transform the Softmax into the desired form of  $A^{\alpha \beta} = \delta_{\alpha \beta} + \mathcal{O}(n^{-1})$ , we first center the attention matrix around the identity matrix:
$$
A _ {\ell} = I + \operatorname {S o f t m a x} \left(\tau^ {- 1} Y _ {\ell}\right),
$$
and then examine the Taylor-expansion of the Softmax w.r.t  $\tau^{-1}$ :
$$
A _ {\ell} = I + \frac {1}{m} 1 1 ^ {\top} + \mathcal {O} (\tau^ {- 1}). \tag {25}
$$
The first few terms of the expansion provides a good approximation of  $A_{\ell}$  when  $\tau$  is large. Observe that for fixed  $m$ , the zero order term  $\frac{1}{m} 11^{\top}$  is not of  $\mathcal{O}(n^{-1})$ . In particular, suppose that we use the attention model in Eq. 25. We can show that the expectation of the term in  $\gamma^{2}$  (responsible for the drift) has the following form:
$$
\gamma^ {2} \sum_ {\delta , \omega = 1} ^ {m} \mathbb {E} _ {\ell} [ A ^ {\alpha \delta} A ^ {\beta \omega} ] V ^ {\delta \omega} = \gamma^ {2} V ^ {\alpha \beta} + \frac {\gamma^ {2}}{m} \sum_ {\omega} V ^ {\alpha \omega} + \frac {\gamma^ {2}}{m} \sum_ {\delta} V ^ {\delta \beta} + \frac {\gamma^ {2}}{m ^ {2}} \sum_ {\delta \omega} V ^ {\delta \omega} + \mathcal {O} (\tau^ {- 1}),
$$
where the terms scaled by  $\frac{1}{m}$  leads to large incremental updates in expected value of  $V_{\ell}$  w.r.t the previous layer, and precludes the possibility of convergence to an SDE. Hence, we choose to re-center the Softmax by removing the term  $\frac{1}{m} 11^{\top}$ , as follows:
$$
A _ {\ell} = I + \operatorname {S o f t m a x} \left(\tau^ {- 1} Y _ {\ell}\right) - \frac {1}{m} 1 1 ^ {\top},
$$
which admits the following Taylor-expansion:
$$
[ \mathrm {S o f t m a x} (\tau^ {- 1} Y _ {\ell}) - m ^ {- 1} 1 1 ^ {\top} ] = \frac {1}{\tau m} [ Y ^ {\alpha \beta} - \overline {{Y ^ {\alpha}}} ] _ {\alpha , \beta} + \frac {1}{2 \tau^ {2} m} \left[ (Y ^ {\alpha \beta} - \overline {{Y}} ^ {\alpha}) ^ {2} - (\overline {{(Y ^ {\alpha}) ^ {2}}} - \overline {{Y ^ {\alpha}}} ^ {2}) \right] _ {\alpha , \beta} + O (\tau^ {- 3}),
$$
where  $\overline{Y^{\alpha}}\coloneqq \frac{1}{m}\sum_{\nu}Y^{\alpha \nu}$ , and  $\overline{(Y^{\alpha})^{2}}\coloneqq \frac{1}{m}\sum_{\nu}(Y^{\alpha \nu})^{2}$ .
Hence, up to third order:
$$
A _ {\ell} = I + \frac {1}{\tau m} [ Y ^ {\alpha \beta} - \overline {{Y ^ {\alpha}}} ] _ {\alpha , \beta} + \frac {1}{2 \tau^ {2} m} \left[ (Y ^ {\alpha \beta} - \overline {{Y}} ^ {\alpha}) ^ {2} - (\overline {{(Y ^ {\alpha}) ^ {2}}} - \overline {{Y ^ {\alpha}}} ^ {2}) \right] _ {\alpha \beta} + O (\tau^ {- 3}).
$$
For sufficiently large  $\tau$ , the formulation above allows infinitesimal updates in  $V_{\ell}$ , and as we will show rigorously in the rest of this section, permits convergence to an SDE.
C.3 Lemmas on Moments of Shaped Attention
$$
F _ {1} ^ {\alpha \beta} = Y ^ {\alpha \beta} - \overline {{Y ^ {\alpha}}},
$$
$$
F _ {2} ^ {\alpha \beta} = (Y ^ {\alpha \beta} - \overline {{Y}} ^ {\alpha}) ^ {2} - (\overline {{(Y ^ {\alpha}) ^ {2}}} - \overline {{Y ^ {\alpha}}} ^ {2}).
$$
Hence,  $A_{\ell}$  can be written as:
$$
A _ {\ell} ^ {\alpha \beta} = \delta_ {\alpha \beta} + \frac {1}{\tau m} F _ {1} ^ {\alpha \beta} + \frac {1}{2 \tau^ {2} m} F _ {2} ^ {\alpha \beta} + O (\tau^ {- 3}).
$$
We now compute the moments of  $A_{\ell}$ . We define the following quantities:
$$
S _ {1} ^ {\alpha \delta , \beta \omega} := \frac {1}{n _ {k}} \mathbb {E} _ {\ell} (Y ^ {\alpha \delta} - \overline {{Y ^ {\alpha}}}) (Y ^ {\beta \omega} - \overline {{Y ^ {\beta}}}) = \frac {1}{n _ {k}} \mathbb {E} _ {\ell} F _ {1} ^ {\alpha \delta} F _ {1} ^ {\beta \omega},
$$
$$
S _ {2} ^ {\alpha \delta} := \frac {1}{n _ {k}} \mathbb {E} _ {\ell} \left[ (Y ^ {\alpha \delta} - \overline {{Y}} ^ {\alpha}) ^ {2} - (\overline {{(Y ^ {\alpha}) ^ {2}}} - \overline {{Y}} ^ {\alpha^ {2}}) \right] = \frac {1}{n _ {k}} \mathbb {E} _ {\ell} F _ {2} ^ {\alpha \delta},
$$
where we recall  $\mathbb{E}_{\ell}[\cdot ] = \mathbb{E}[\cdot |\mathcal{F}_{\ell}]$  is the conditional expectation given the sigma-algebra generated by  $\mathcal{F}_{\ell} = \sigma (\{X_{k}\}_{k\in [\ell ]})$
Lemma C.2 (Moments of Taylor Expansion).
$$
S _ {1} ^ {\alpha \delta , \beta \omega} = V ^ {\alpha \beta} \left(V ^ {\delta \omega} - V ^ {\delta \bar {x}} - V ^ {\omega \bar {x}} + V ^ {\bar {x} \bar {x}}\right),
$$
$$
S _ {2} ^ {\alpha \delta} = V ^ {\alpha \alpha} \left(V ^ {\delta \delta} - 2 V ^ {\delta \bar {x}} + 2 V ^ {\bar {x} \bar {x}} - \bar {V}\right),
$$
where  $\bar{V} = \frac{1}{m}\sum_{\nu}V^{\nu \nu}$  and  $\bar{x} = \frac{1}{m}\sum_{\nu}x^{\nu}$  is the average token.
Proof. Using Lemma C.1 and linearity of expectation:
$$
\begin{array}{l} S _ {1} ^ {\alpha \delta , \beta \omega} = \frac {1}{n _ {k}} \left(\mathbb {E} _ {\ell} [ Y ^ {\alpha \delta} Y ^ {\beta \omega} ] - \mathbb {E} _ {\ell} [ Y ^ {\alpha \delta} \overline {{Y ^ {\beta}}} ] - \mathbb {E} _ {\ell} [ Y ^ {\beta \omega} \overline {{Y ^ {\alpha}}} ] + \mathbb {E} _ {\ell} [ \overline {{Y ^ {\alpha}}} \overline {{Y ^ {\beta}}} ]\right) \\ = \left(V ^ {\alpha \beta} V ^ {\delta \omega} - V ^ {\alpha \beta} V ^ {\delta \bar {x}} - V ^ {\alpha \beta} V ^ {\omega \bar {x}} + V ^ {\alpha \beta} V ^ {\bar {x} \bar {x}}\right) \\ = V ^ {\alpha \beta} \left(V ^ {\delta \omega} - V ^ {\delta \bar {x}} - V ^ {\omega \bar {x}} + V ^ {\bar {x} \bar {x}}\right), \\ \end{array}
$$
$$
\begin{array}{l} S _ {2} ^ {\alpha \delta} = \frac {1}{n _ {k}} \left(\mathbb {E} _ {\ell} \left[ (Y ^ {\alpha \delta} - \overline {{Y}} ^ {\alpha}) ^ {2} \right] - \mathbb {E} _ {\ell} \left[ (\overline {{(Y ^ {\alpha}) ^ {2}}} - \overline {{Y ^ {\alpha}}} ^ {2}) \right]\right) \\ = \frac {1}{n _ {k}} \left(\mathbb {E} _ {\ell} \left[ (Y ^ {\alpha \delta}) ^ {2} \right] - 2 \mathbb {E} _ {\ell} \left[ Y ^ {\alpha \delta} \overline {{Y ^ {\alpha}}} \right] + 2 \mathbb {E} _ {\ell} \left[ \overline {{Y ^ {\alpha}}} ^ {2} \right] - \mathbb {E} _ {\ell} \left[ (\overline {{Y ^ {\alpha}}} ^ {2}) \right]\right) \\ = \left(V ^ {\alpha \alpha} V ^ {\delta \delta} - 2 V ^ {\alpha \alpha} V ^ {\delta \bar {x}} + 2 V ^ {\alpha \alpha} V ^ {\bar {x} \bar {x}} - V ^ {\alpha \alpha} \bar {V}\right) \\ = V ^ {\alpha \alpha} \left(V ^ {\delta \delta} - 2 V ^ {\delta \bar {x}} + 2 V ^ {\bar {x} \bar {x}} - \bar {V}\right), \\ \end{array}
$$
where  $\bar{V} = \frac{1}{m}\sum_{\beta}V^{\beta \beta}$
$$
\mathbb {E} _ {\ell} [ A ^ {\alpha \delta} A ^ {\beta \omega} ] = \delta_ {\alpha \delta} \delta_ {\beta \omega} + \frac {n _ {k}}{\tau^ {2} m ^ {2}} S _ {1} ^ {\alpha \delta , \beta \omega} + \frac {n _ {k}}{2 \tau^ {2} m} (\delta_ {\beta \omega} S _ {2} ^ {\alpha \delta} + \delta_ {\alpha \delta} S _ {2} ^ {\beta \omega}) + O (n _ {k} \tau^ {- 3}).
$$
$$
\begin{array}{l} \mathbb {E} _ {\ell} [ A ^ {\alpha \delta} A ^ {\beta \omega} ] = \delta_ {\alpha \delta} \delta_ {\beta \omega} + \frac {\delta_ {\beta \omega}}{\tau m} \mathbb {E} _ {\ell} [ Y ^ {\alpha \delta} - \overline {{Y ^ {\alpha}}} ] + \frac {\delta_ {\alpha \delta}}{\tau m} \mathbb {E} _ {\ell} [ Y ^ {\beta \omega} - \overline {{Y ^ {\beta}}} ] + \frac {1}{\tau^ {2} m ^ {2}} \mathbb {E} _ {\ell} (Y ^ {\alpha \delta} - \overline {{Y ^ {\alpha}}}) (Y ^ {\beta \omega} - \overline {{Y ^ {\beta}}}) \\ + \frac {\delta_ {\beta \omega}}{2 \tau^ {2} m} \mathbb {E} _ {\ell} \left[ (Y ^ {\alpha \delta} - \overline {{Y}} ^ {\alpha}) ^ {2} - (\overline {{(Y ^ {\alpha}) ^ {2}}} - \overline {{Y}} ^ {\alpha^ {2}}) \right] \\ + \frac {\delta_ {\alpha \delta}}{2 \tau^ {2} m} \mathbb {E} _ {\ell} \left[ (Y ^ {\beta \omega} - \overline {{Y}} ^ {\beta}) ^ {2} - (\overline {{(Y ^ {\beta}) ^ {2}}} - \overline {{Y ^ {\beta}}} ^ {2}) \right] + O (n _ {k} \tau^ {- 3}) \\ = \delta_ {\alpha \delta} \delta_ {\beta \omega} + \frac {n _ {k}}{\tau^ {2} m ^ {2}} S _ {1} ^ {\alpha \delta , \beta \omega} + \frac {n _ {k}}{2 \tau^ {2} m} (\delta_ {\beta \omega} S _ {2} ^ {\alpha \delta} + \delta_ {\alpha \delta} S _ {2} ^ {\beta \omega}) + O (n _ {k} \tau^ {- 3}). \\ \end{array}
$$
Lemma C.4.
$$
\begin{array}{l} \mathbb {E} _ {\ell} \left[ A ^ {\alpha \alpha^ {\prime}} A ^ {\beta \beta^ {\prime}} A ^ {\delta \delta^ {\prime}} A ^ {\omega \omega^ {\prime}} \right] = \delta_ {\alpha \alpha^ {\prime}} \delta_ {\beta \beta^ {\prime}} \delta_ {\delta \delta^ {\prime}} \delta_ {\omega \omega^ {\prime}} \\ + \frac {n _ {k}}{\tau^ {2} m ^ {2}} \bigg (\delta_ {\alpha \alpha^ {\prime}} \delta_ {\beta \beta^ {\prime}} S _ {1} ^ {\delta \delta^ {\prime}, \omega \omega^ {\prime}} + \delta_ {\alpha \alpha^ {\prime}} \delta_ {\delta \delta^ {\prime}} S _ {1} ^ {\beta \beta^ {\prime}, \omega \omega^ {\prime}} + \delta_ {\alpha \alpha^ {\prime}} \delta_ {\omega \omega^ {\prime}} S _ {1} ^ {\beta \beta^ {\prime}, \delta \delta^ {\prime}} \\ \left. + \delta_ {\beta \beta^ {\prime}} \delta_ {\delta \delta^ {\prime}} S _ {1} ^ {\alpha \alpha^ {\prime}, \omega \omega^ {\prime}} + \delta_ {\beta \beta^ {\prime}} \delta_ {\omega \omega^ {\prime}} S _ {1} ^ {\alpha \alpha^ {\prime}, \delta \delta^ {\prime}} + \delta_ {\omega \omega^ {\prime}} \delta_ {\delta \delta^ {\prime}} S _ {1} ^ {\alpha \alpha^ {\prime}, \beta \beta^ {\prime}}\right) \\ + \frac {n _ {k}}{2 \tau^ {2} m} \Bigl (\delta_ {\alpha \alpha^ {\prime}} \delta_ {\beta \beta^ {\prime}} \delta_ {\delta \delta^ {\prime}} S _ {2} ^ {\omega \omega^ {\prime}} + \delta_ {\alpha \alpha^ {\prime}} \delta_ {\beta \beta^ {\prime}} \delta_ {\omega \omega^ {\prime}} S _ {2} ^ {\delta \delta^ {\prime}} \\ \left. + \delta_ {\alpha \alpha^ {\prime}} \delta_ {\omega \omega^ {\prime}} \delta_ {\delta \delta^ {\prime}} S _ {2} ^ {\beta \beta^ {\prime}} + \delta_ {\omega \omega^ {\prime}} \delta_ {\beta \beta^ {\prime}} \delta_ {\delta \delta^ {\prime}} S _ {2} ^ {\alpha \alpha^ {\prime}}\right) + O (n _ {k} \tau^ {- 3}). \\ \end{array}
$$
$$
\begin{array}{l} \mathbb {E} _ {\ell} \left[ A ^ {\alpha \alpha^ {\prime}} A ^ {\beta \beta^ {\prime}} A ^ {\delta \delta^ {\prime}} A ^ {\omega \omega^ {\prime}} \right] = \mathbb {E} _ {\ell} \left[ \left(\delta_ {\alpha \alpha^ {\prime}} + \frac {1}{\tau m} F _ {1} ^ {\alpha \alpha^ {\prime}} + \frac {1}{2 \tau^ {2} m} F _ {2} ^ {\alpha \alpha^ {\prime}}\right) \left(\delta_ {\beta \beta^ {\prime}} + \frac {1}{\tau m} F _ {1} ^ {\beta \beta^ {\prime}} + \frac {1}{2 \tau^ {2} m} F _ {2} ^ {\beta \beta^ {\prime}}\right) \right. \\ \left. \left(\delta_ {\delta \delta^ {\prime}} + \frac {1}{\tau m} F _ {1} ^ {\delta \delta^ {\prime}} + \frac {1}{2 \tau^ {2} m} F _ {2} ^ {\delta \delta^ {\prime}}\right) \left(\delta_ {\omega \omega^ {\prime}} + \frac {1}{\tau m} F _ {1} ^ {\omega \omega^ {\prime}} + \frac {1}{2 \tau^ {2} m} F _ {2} ^ {\omega \omega^ {\prime}}\right) \right] + O (\tau^ {- 3}) \\ = \mathbb {E} _ {\ell} \left[ \delta_ {\alpha \alpha^ {\prime}} \delta_ {\beta \beta^ {\prime}} \delta_ {\delta \delta^ {\prime}} \delta_ {\omega \omega^ {\prime}} + \frac {1}{\tau^ {2} m ^ {2}} \delta_ {\alpha \alpha^ {\prime}} \delta_ {\beta \beta^ {\prime}} F _ {1} ^ {\delta \delta^ {\prime}} F _ {1} ^ {\omega \omega^ {\prime}} + \frac {1}{\tau^ {2} m ^ {2}} \delta_ {\alpha \alpha^ {\prime}} \delta_ {\delta \delta^ {\prime}} F _ {1} ^ {\beta \beta^ {\prime}} F _ {1} ^ {\omega \omega^ {\prime}} \right] \\ + \frac {1}{\tau^ {2} m ^ {2}} \delta_ {\alpha \alpha^ {\prime}} \delta_ {\omega \omega^ {\prime}} F _ {1} ^ {\delta \delta^ {\prime}} F _ {1} ^ {\beta \beta^ {\prime}} + \frac {1}{\tau^ {2} m ^ {2}} \delta_ {\beta \beta^ {\prime}} \delta_ {\delta \delta^ {\prime}} F _ {1} ^ {\alpha \alpha^ {\prime}} F _ {1} ^ {\omega \omega^ {\prime}} \\ + \frac {1}{\tau^ {2} m ^ {2}} \delta_ {\beta \beta^ {\prime}} \delta_ {\omega \omega^ {\prime}} F _ {1} ^ {\alpha \alpha^ {\prime}} F _ {1} ^ {\delta \delta^ {\prime}} + \frac {1}{\tau^ {2} m ^ {2}} \delta_ {\delta \delta^ {\prime}} \delta_ {\omega \omega^ {\prime}} F _ {1} ^ {\alpha \alpha^ {\prime}} F _ {1} ^ {\beta \beta^ {\prime}} \\ + \frac {1}{2 \tau^ {2} m} \delta_ {\alpha \alpha^ {\prime}} \delta_ {\beta \beta^ {\prime}} \delta_ {\delta \delta^ {\prime}} F _ {2} ^ {\omega \omega^ {\prime}} + \frac {1}{2 \tau^ {2} m} \delta_ {\alpha \alpha^ {\prime}} \delta_ {\beta \beta^ {\prime}} \delta_ {\omega \omega^ {\prime}} F _ {2} ^ {\delta \delta^ {\prime}} \\ \left. \right.\left. + \frac {1}{2 \tau^ {2} m} \delta_ {\alpha \alpha^ {\prime}} \delta_ {\omega \omega^ {\prime}} \delta_ {\delta \delta^ {\prime}} F _ {2} ^ {\beta \beta^ {\prime}} + \frac {1}{2 \tau^ {2} m} \delta_ {\omega \omega^ {\prime}} \delta_ {\beta \beta^ {\prime}} \delta_ {\delta \delta^ {\prime}} F _ {2} ^ {\alpha \alpha^ {\prime}} \right] + O (\tau^ {- 3}) . \\ \end{array}
$$
Using the linearity of expectation:
$$
\begin{array}{l} \mathbb {E} _ {\ell} \left[ A ^ {\alpha \alpha^ {\prime}} A ^ {\beta \beta^ {\prime}} A ^ {\delta \delta^ {\prime}} A ^ {\omega \omega^ {\prime}} \right] = \delta_ {\alpha \alpha^ {\prime}} \delta_ {\beta \beta^ {\prime}} \delta_ {\delta \delta^ {\prime}} \delta_ {\omega \omega^ {\prime}} \\ + \frac {n _ {k}}{\tau^ {2} m ^ {2}} \Bigl (\delta_ {\alpha \alpha^ {\prime}} \delta_ {\beta \beta^ {\prime}} S _ {1} ^ {\delta \delta^ {\prime}, \omega \omega^ {\prime}} + \delta_ {\alpha \alpha^ {\prime}} \delta_ {\delta \delta^ {\prime}} S _ {1} ^ {\beta \beta^ {\prime}, \omega \omega^ {\prime}} + \delta_ {\alpha \alpha^ {\prime}} \delta_ {\omega \omega^ {\prime}} S _ {1} ^ {\beta \beta^ {\prime}, \delta \delta^ {\prime}} \\ + \delta_ {\beta \beta^ {\prime}} \delta_ {\delta \delta^ {\prime}} S _ {1} ^ {\alpha \alpha^ {\prime}, \omega \omega^ {\prime}} + \delta_ {\beta \beta^ {\prime}} \delta_ {\omega \omega^ {\prime}} S _ {1} ^ {\alpha \alpha^ {\prime}, \delta \delta^ {\prime}} + \delta_ {\omega \omega^ {\prime}} \delta_ {\delta \delta^ {\prime}} S _ {1} ^ {\alpha \alpha^ {\prime}, \beta \beta^ {\prime}}) \\ + \frac {n _ {k}}{2 \tau^ {2} m} \left(\delta_ {\alpha \alpha^ {\prime}} \delta_ {\beta \beta^ {\prime}} \delta_ {\delta \delta^ {\prime}} S _ {2} ^ {\omega \omega^ {\prime}} + \delta_ {\alpha \alpha^ {\prime}} \delta_ {\beta \beta^ {\prime}} \delta_ {\omega \omega^ {\prime}} S _ {2} ^ {\delta \delta^ {\prime}} \right. \\ \left. + \delta_ {\alpha \alpha^ {\prime}} \delta_ {\omega \omega^ {\prime}} \delta_ {\delta \delta^ {\prime}} S _ {2} ^ {\beta \beta^ {\prime}} + \delta_ {\omega \omega^ {\prime}} \delta_ {\beta \beta^ {\prime}} \delta_ {\delta \delta^ {\prime}} S _ {2} ^ {\alpha \alpha^ {\prime}}\right) + O (n _ {k} \tau^ {- 3}). \\ \end{array}
$$
Note that the terms above can be computed using Lemma C.2.
Lemma C.5.
$$
\mathbb {E} _ {\ell} \left[ A ^ {\alpha \alpha^ {\prime}} A ^ {\beta \beta^ {\prime}} A ^ {\delta \delta^ {\prime}} A ^ {\omega \omega^ {\prime}} \right] - \mathbb {E} _ {\ell} \left[ A ^ {\alpha \alpha^ {\prime}} A ^ {\beta \beta^ {\prime}} \right] \mathbb {E} _ {\ell} \left[ A ^ {\delta \delta^ {\prime}} A ^ {\omega \omega^ {\prime}} \right]
$$
$$
= \frac {n _ {k}}{\tau^ {2} m ^ {2}} \Big (\delta_ {\alpha \alpha^ {\prime}} \delta_ {\delta \delta^ {\prime}} S _ {1} ^ {\beta \beta^ {\prime}, \omega \omega^ {\prime}} + \delta_ {\alpha \alpha^ {\prime}} \delta_ {\omega \omega^ {\prime}} S _ {1} ^ {\beta \beta^ {\prime}, \delta \delta^ {\prime}} + \delta_ {\beta \beta^ {\prime}} \delta_ {\delta \delta^ {\prime}} S _ {1} ^ {\alpha \alpha^ {\prime}, \omega \omega^ {\prime}} + \delta_ {\beta \beta^ {\prime}} \delta_ {\omega \omega^ {\prime}} S _ {1} ^ {\alpha \alpha^ {\prime}, \delta \delta^ {\prime}} \Big) + \mathcal {O} (n _ {k} \tau^ {- 3}).
$$
Proof. The results are an immediate consequence of Lemma C.4 and Lemma C.3, where only the terms that do not cancel out are kept.  $\square$
C.4 Neural Covariance SDE for Stable Attention
$$
\mathcal {T} _ {1} ^ {\alpha \beta} := \frac {1}{n} \left(X _ {\ell} W _ {\ell} ^ {\top} X _ {\ell} ^ {\top} A _ {\ell} ^ {\top} + A _ {\ell} X _ {\ell} W _ {\ell} X _ {\ell} ^ {\top}\right) ^ {\alpha \beta},
$$
$$
\mathcal {T} _ {2} ^ {\alpha \beta} := \frac {1}{n \sqrt {n}} \left(A _ {\ell} X _ {\ell} W _ {\ell} W _ {\ell} ^ {\top} X _ {\ell} ^ {\top} A _ {\ell} ^ {\top}\right) ^ {\alpha \beta}.
$$
Hence, the expression for  $V_{\ell}$  simplifies to:
$$
V _ {\ell + 1} ^ {\alpha \beta} = \lambda^ {2} V _ {\ell} ^ {\alpha \beta} + \frac {\lambda \gamma}{\sqrt {n}} \mathcal {T} _ {1} ^ {\alpha \beta} + \frac {\gamma^ {2}}{\sqrt {n}} \mathcal {T} _ {2} ^ {\alpha \beta}.
$$
We need to compute the moments for these two quantities:
Lemma C.6 (Moments of  $\mathcal{T}_1$ ).
$$
\mathbb {E} _ {\ell} [ \mathcal {T} _ {1} ^ {\alpha \beta} ] = 0,
$$
$$
\mathbb {E} _ {\ell} \left[ \mathcal {T} _ {1} ^ {\alpha \beta} \mathcal {T} _ {1} ^ {\delta \omega} \right] = 2 (V ^ {\alpha \delta} V ^ {\beta \omega} + V ^ {\alpha \omega} V ^ {\beta \delta}) + \mathcal {O} (n _ {k} \tau^ {- 3}).
$$
$$
\begin{array}{l} \mathcal {T} _ {1} ^ {\alpha \beta} \mathcal {T} _ {1} ^ {\delta \omega} = \frac {1}{n ^ {2}} \left[ \left(X _ {\ell} W _ {\ell} ^ {\top} X _ {\ell} ^ {\top} A _ {\ell} ^ {\top} + A _ {\ell} X _ {\ell} W _ {\ell} X _ {\ell} ^ {\top}\right) ^ {\alpha \beta} \left(X _ {\ell} W _ {\ell} ^ {\top} X _ {\ell} ^ {\top} A _ {\ell} ^ {\top} + A _ {\ell} X _ {\ell} W _ {\ell} X _ {\ell} ^ {\top}\right) ^ {\delta \omega} \right] \\ = \frac {1}{n ^ {2}} \Bigg [ \left(X _ {\ell} W _ {\ell} ^ {\top} X _ {\ell} ^ {\top} A _ {\ell} ^ {\top}\right) ^ {\alpha \beta} \left(X _ {\ell} W _ {\ell} ^ {\top} X _ {\ell} ^ {\top} A _ {\ell} ^ {\top}\right) ^ {\delta \omega} + \left(X _ {\ell} W _ {\ell} ^ {\top} X _ {\ell} ^ {\top} A _ {\ell} ^ {\top}\right) ^ {\alpha \beta} \left(A _ {\ell} X _ {\ell} W _ {\ell} X _ {\ell} ^ {\top}\right) ^ {\delta \omega} \\ \left. + \left(A _ {\ell} X _ {\ell} W _ {\ell} X _ {\ell} ^ {\top}\right) ^ {\alpha \beta} \left(X _ {\ell} W _ {\ell} ^ {\top} X _ {\ell} ^ {\top} A _ {\ell} ^ {\top}\right) ^ {\delta \omega} + \left(A _ {\ell} X _ {\ell} W _ {\ell} X _ {\ell} ^ {\top}\right) ^ {\alpha \beta} \left(A _ {\ell} X _ {\ell} W _ {\ell} X _ {\ell} ^ {\top}\right) ^ {\delta \omega} \right]. \\ \end{array}
$$
Let's look at the first summand:
$$
\frac {1}{n ^ {2}} \left(X _ {\ell} W _ {\ell} ^ {\top} X _ {\ell} ^ {\top} A _ {\ell} ^ {\top}\right) ^ {\alpha \beta} \left(X _ {\ell} W _ {\ell} ^ {\top} X _ {\ell} ^ {\top} A _ {\ell} ^ {\top}\right) ^ {\delta \omega} = \frac {1}{n ^ {2}} \sum_ {\nu \kappa} \sum_ {k k ^ {\prime} j j ^ {\prime}} X _ {k} ^ {\alpha} W _ {k ^ {\prime} k} X _ {k ^ {\prime}} ^ {\nu} A ^ {\beta \nu} X _ {j} ^ {\delta} W _ {j ^ {\prime} j} X _ {j ^ {\prime}} ^ {\kappa} A ^ {\omega \kappa}.
$$
Hence, in expectation with respect to  $W$ :
$$
\begin{array}{l} \frac {1}{n ^ {2}} \sum_ {\nu \kappa} \sum_ {k k ^ {\prime} j j ^ {\prime}} X _ {k} ^ {\alpha} \mathbb {E} \left[ W _ {k ^ {\prime} k} W _ {j ^ {\prime} j} \right] X _ {k ^ {\prime}} ^ {\nu} A ^ {\beta \nu} X _ {j} ^ {\delta} X _ {j ^ {\prime}} ^ {\kappa} A ^ {\omega \kappa} = \frac {1}{n ^ {2}} \sum_ {\nu \kappa} \sum_ {k k ^ {\prime} j j ^ {\prime}} X _ {k} ^ {\alpha} \delta_ {k j} \delta_ {k ^ {\prime} j ^ {\prime}} X _ {k ^ {\prime}} ^ {\nu} A ^ {\beta \nu} X _ {j} ^ {\delta} X _ {j ^ {\prime}} ^ {\kappa} A ^ {\omega \kappa} \\ = \frac {1}{n ^ {2}} \sum_ {\nu \kappa} \sum_ {k k ^ {\prime}} X _ {k} ^ {\alpha} X _ {k ^ {\prime}} ^ {\nu} A ^ {\beta \nu} X _ {k} ^ {\delta} X _ {k ^ {\prime}} ^ {\kappa} A ^ {\omega \kappa} \\ = \sum_ {\nu \kappa} V ^ {\alpha \delta} V ^ {\nu \kappa} A ^ {\beta \nu} A ^ {\omega \kappa}. \\ \end{array}
$$
An identical argument can be made for the remaining three summands. Hence, taking expectation with respect to the Softmax weights:
$$
\mathbb {E} _ {\ell} \left[ \mathcal {T} _ {1} ^ {\alpha \beta} \mathcal {T} _ {1} ^ {\delta \omega} \right] = \sum_ {\nu \kappa} \left(V ^ {\alpha \delta} V ^ {\nu \kappa} \mathbb {E} _ {\ell} [ A ^ {\beta \nu} A ^ {\omega \kappa} ] + V ^ {\alpha \omega} V ^ {\nu \kappa} \mathbb {E} _ {\ell} [ A ^ {\beta \nu} A ^ {\delta \kappa} ] + V ^ {\beta \delta} V ^ {\nu \kappa} \mathbb {E} _ {\ell} [ A ^ {\alpha \nu} A ^ {\omega \kappa} ] + V ^ {\beta \omega} V ^ {\nu \kappa} \mathbb {E} _ {\ell} [ A ^ {\alpha \nu} A ^ {\delta \kappa} ]\right).
$$
Now, using Lemma C.3:
$$
\begin{array}{l} \mathbb {E} _ {\ell} \left[ \mathcal {T} _ {1} ^ {\alpha \beta} \mathcal {T} _ {1} ^ {\delta \omega} \right] = \sum_ {\nu \kappa} V ^ {\nu \kappa} \left(V ^ {\alpha \delta} \mathbb {E} _ {\ell} \left[ A ^ {\beta \nu} A ^ {\omega \kappa} \right] + V ^ {\alpha \omega} \mathbb {E} _ {\ell} \left[ A ^ {\beta \nu} A ^ {\delta \kappa} \right] + V ^ {\beta \delta} \mathbb {E} _ {\ell} \left[ A ^ {\alpha \nu} A ^ {\omega \kappa} \right] + V ^ {\beta \omega} \mathbb {E} _ {\ell} \left[ A ^ {\alpha \nu} A ^ {\delta \kappa} \right]\right) \\ = \sum_ {\nu \kappa} V ^ {\nu \kappa} \left(V ^ {\alpha \delta} \delta_ {\beta \nu} \delta_ {\omega \kappa} + V ^ {\alpha \omega} \delta_ {\beta \nu} \delta_ {\delta \kappa} + V ^ {\beta \delta} \delta_ {\alpha \nu} \delta_ {\omega \kappa} + V ^ {\beta \omega} \delta_ {\alpha \nu} \delta_ {\delta \kappa}\right) + \mathcal {O} (n _ {k} \tau^ {- 2}) \\ = V ^ {\beta \omega} V ^ {\alpha \delta} + V ^ {\alpha \omega} V ^ {\beta \delta} + V ^ {\beta \delta} V ^ {\alpha \omega} + V ^ {\beta \omega} V ^ {\alpha \delta} + \mathcal {O} \left(n _ {k} \tau^ {- 2}\right) \\ = 2 \left(V ^ {\alpha \delta} V ^ {\beta \omega} + V ^ {\alpha \omega} V ^ {\beta \delta}\right) + \mathcal {O} \left(n _ {k} \tau^ {- 2}\right). \\ \end{array}
$$
Lemma C.7 (Moments of  $\mathcal{T}_2$ ).
$$
\begin{array}{l} \mathbb {E} _ {\ell} [ \mathcal {T} _ {2} ^ {\alpha \beta} ] = \sqrt {n} \sum_ {\nu \kappa} V ^ {\nu \kappa} \mathbb {E} _ {\ell} \left[ A ^ {\alpha \nu} A ^ {\beta \kappa} \right], \\ \mathbb {E} _ {\ell} [ \mathcal {T} _ {2} ^ {\alpha \beta} \mathcal {T} _ {2} ^ {\delta \omega} ] = \sum_ {\nu \kappa \nu^ {\prime} \kappa^ {\prime}} \mathbb {E} _ {\ell} [ A ^ {\alpha \nu} A ^ {\beta \kappa} A ^ {\delta \nu^ {\prime}} A ^ {\omega \kappa^ {\prime}} ] \left(n V ^ {\nu \kappa} V ^ {\nu^ {\prime} \kappa^ {\prime}} + V ^ {\nu \nu^ {\prime}} V ^ {\kappa \kappa^ {\prime}} + V ^ {\nu \kappa^ {\prime}} V ^ {\nu^ {\prime} \kappa}\right). \\ \end{array}
$$