SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\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} = \frac {1}{n \sqrt {n}} \sum_ {\nu \kappa} \sum_ {k k ^ {\prime} j} A ^ {\alpha \nu} X _ {k} ^ {\nu} W _ {k k ^ {\prime}} W _ {j k ^ {\prime}} X _ {j} ^ {\kappa} A ^ {\beta \kappa}.
$$
Taking expectation with respect to  $W$ :
$$
\begin{array}{l} \frac {1}{n \sqrt {n}} \left(A _ {\ell} X _ {\ell} W _ {\ell} W _ {\ell} ^ {\top} X _ {\ell} ^ {\top} A _ {\ell} ^ {\top}\right) ^ {\alpha \beta} = \frac {1}{n \sqrt {n}} \sum_ {\nu \kappa} \sum_ {k k ^ {\prime} j} A ^ {\alpha \nu} X _ {k} ^ {\nu} \mathbb {E} [ W _ {k k ^ {\prime}} W _ {j k ^ {\prime}} ] X _ {j} ^ {\kappa} A ^ {\beta \kappa} \\ = \frac {1}{n \sqrt {n}} \sum_ {\nu \kappa} \sum_ {k k ^ {\prime} j} A ^ {\alpha \nu} X _ {k} ^ {\nu} \delta_ {k j} X _ {j} ^ {\kappa} A ^ {\beta \kappa} \\ = \frac {1}{n \sqrt {n}} \sum_ {\nu \kappa} \sum_ {k k ^ {\prime}} A ^ {\alpha \nu} A ^ {\beta \kappa} X _ {k} ^ {\nu} X _ {k} ^ {\kappa} \\ = \frac {1}{\sqrt {n}} \sum_ {\nu \kappa} \sum_ {k} A ^ {\alpha \nu} A ^ {\beta \kappa} X _ {k} ^ {\nu} X _ {k} ^ {\kappa} \\ = \sqrt {n} \sum_ {\nu \kappa} V ^ {\nu \kappa} A ^ {\alpha \nu} A ^ {\beta \kappa}. \\ \end{array}
$$
Taking expectation w.r.t the Softmax weights, we get the desired result.
For second moment, we can take the conditional expectation:
$$
\mathbb {E} _ {\ell} [ \mathcal {T} _ {2} ^ {\alpha \beta} \mathcal {T} _ {2} ^ {\delta \omega} ] = \frac {1}{n ^ {3}} \sum_ {\nu \kappa \nu^ {\prime} \kappa^ {\prime}} \sum_ {k k ^ {\prime} j i i ^ {\prime} j ^ {\prime}} A ^ {\alpha \nu} A ^ {\beta \kappa} A ^ {\delta \nu^ {\prime}} A ^ {\omega \kappa^ {\prime}} X _ {k} ^ {\nu} X _ {j} ^ {\kappa} X _ {i} ^ {\nu^ {\prime}} X _ {j ^ {\prime}} ^ {\kappa^ {\prime}} \mathbb {E} [ W _ {k k ^ {\prime}} W _ {j k ^ {\prime}} W _ {i i ^ {\prime}} W _ {j ^ {\prime} i ^ {\prime}} ],
$$
Using Isserlis Theorem, we have that:
$$
\mathbb {E} \left[ W _ {k k ^ {\prime}} W _ {j k ^ {\prime}} W _ {i i ^ {\prime}} W _ {j ^ {\prime} i ^ {\prime}} \right] = \delta_ {k j} \delta_ {i j ^ {\prime}} + \delta_ {k i} \delta_ {k ^ {\prime} i ^ {\prime}} \delta_ {j j ^ {\prime}} + \delta_ {k j ^ {\prime}} \delta_ {k ^ {\prime} i ^ {\prime}} \delta_ {j i}.
$$
$$
\begin{array}{l} \mathbb {E} _ {\ell} \left[ \mathcal {T} _ {2} ^ {\alpha \beta} \mathcal {T} _ {2} ^ {\delta \omega} \right] = \frac {1}{n ^ {3}} \sum_ {\nu \kappa \nu^ {\prime} \kappa^ {\prime}} A ^ {\alpha \nu} A ^ {\beta \kappa} A ^ {\delta \nu^ {\prime}} A ^ {\omega \kappa^ {\prime}} \sum_ {k k ^ {\prime} j i i ^ {\prime} j ^ {\prime}} X _ {k} ^ {\nu} X _ {j} ^ {\kappa} X _ {i} ^ {\nu^ {\prime}} X _ {j ^ {\prime}} ^ {\kappa^ {\prime}} \Big (\delta_ {k j} \delta_ {i j ^ {\prime}} + \delta_ {k i} \delta_ {k ^ {\prime} i ^ {\prime}} \delta_ {j j ^ {\prime}} + \\ \left. + \delta_ {k j ^ {\prime}} \delta_ {k ^ {\prime} i ^ {\prime}} \delta_ {j i}\right) \\ = \frac {1}{n ^ {3}} \sum_ {\nu \kappa \nu^ {\prime} \kappa^ {\prime}} A ^ {\alpha \nu} A ^ {\beta \kappa} A ^ {\delta \nu^ {\prime}} A ^ {\omega \kappa^ {\prime}} \left(\sum_ {k k ^ {\prime} i i ^ {\prime}} X _ {k} ^ {\nu} X _ {k} ^ {\kappa} X _ {i} ^ {\nu^ {\prime}} X _ {i} ^ {\kappa^ {\prime}} + \sum_ {k k ^ {\prime} j} X _ {k} ^ {\nu} X _ {j} ^ {\kappa} X _ {k} ^ {\nu^ {\prime}} X _ {j} ^ {\kappa^ {\prime}} \right. \\ \left. + \sum_ {k k ^ {\prime} j} X _ {k} ^ {\nu} X _ {j} ^ {\kappa} X _ {j} ^ {\nu^ {\prime}} X _ {k} ^ {\kappa^ {\prime}}\right) \\ = \frac {1}{n ^ {3}} \sum_ {\nu \kappa \nu^ {\prime} \kappa^ {\prime}} A ^ {\alpha \nu} A ^ {\beta \kappa} A ^ {\delta \nu^ {\prime}} A ^ {\omega \kappa^ {\prime}} \left(n ^ {4} V ^ {\nu \kappa} V ^ {\nu^ {\prime} \kappa^ {\prime}} + n ^ {3} V ^ {\nu \nu^ {\prime}} V ^ {\kappa \kappa^ {\prime}} + n ^ {3} V ^ {\nu \kappa^ {\prime}} V ^ {\nu^ {\prime} \kappa}\right) \\ = \sum_ {\nu \kappa \nu^ {\prime} \kappa^ {\prime}} 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}
$$
By taking expectation w.r.t the Softmax parameters, we get the desired result.
Lemma C.8 (Covariance of  $\mathcal{T}_2$ ).
$$
\mathbb {E} _ {\ell} [ \mathcal {T} _ {2} ^ {\alpha \beta} \mathcal {T} _ {2} ^ {\delta \omega} ] - \mathbb {E} _ {\ell} [ \mathcal {T} _ {2} ^ {\alpha \beta} ] \mathbb {E} _ {\ell} [ \mathcal {T} _ {2} ^ {\delta \omega} ] = \frac {n n _ {k}}{\tau^ {2}} \mathcal {A} ^ {\alpha \beta \delta \omega} + V ^ {\alpha \delta} V ^ {\beta \omega} + V ^ {\alpha \omega} V ^ {\beta \delta} + \mathcal {O} \left(\frac {n n _ {k}}{\tau^ {3}} + \frac {n _ {k}}{\tau^ {2}}\right),
$$
$$
\mathcal {A} ^ {\alpha \beta \delta \omega} := \frac {1}{m ^ {2}} \sum_ {\nu \kappa} \left(V ^ {\alpha \kappa} V ^ {\delta \nu} S _ {1} ^ {\beta \kappa , \omega \nu} + V ^ {\alpha \kappa} V ^ {\omega \nu} S _ {1} ^ {\beta \kappa , \delta \nu} + V ^ {\beta \nu} V ^ {\delta \kappa} S _ {1} ^ {\alpha \nu , \omega \kappa} + V ^ {\beta \nu} V ^ {\omega \kappa} S _ {1} ^ {\alpha \nu , \delta \kappa}\right).
$$
Proof. Using Lemma C.7, we have that:
$$
\begin{array}{l} \mathbb {E} _ {\ell} \left[ \mathcal {T} _ {2} ^ {\alpha \beta} \mathcal {T} _ {2} ^ {\delta \omega} \right] - \mathbb {E} _ {\ell} \left[ \mathcal {T} _ {2} ^ {\alpha \beta} \right] \mathbb {E} _ {\ell} \left[ \mathcal {T} _ {2} ^ {\delta \omega} \right] = \sum_ {\nu \kappa \nu^ {\prime} \kappa^ {\prime}} \mathbb {E} _ {\ell} \left[ A ^ {\alpha \nu} A ^ {\beta \kappa} A ^ {\delta \nu^ {\prime}} A ^ {\omega \kappa^ {\prime}} \right] \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) \\ - n \sum_ {\nu \kappa \nu^ {\prime} \kappa^ {\prime}} V ^ {\nu \kappa} V ^ {\nu^ {\prime} \kappa^ {\prime}} \mathbb {E} _ {\ell} \left[ A ^ {\alpha \nu} A ^ {\beta \kappa} \right] \mathbb {E} _ {\ell} \left[ A ^ {\delta \nu^ {\prime}} A ^ {\omega \kappa^ {\prime}} \right] \\ = n \sum_ {\nu \kappa \nu^ {\prime} \kappa^ {\prime}} V ^ {\nu \kappa} V ^ {\nu^ {\prime} \kappa^ {\prime}} \left(\mathbb {E} _ {\ell} \left[ A ^ {\alpha \nu} A ^ {\beta \kappa} A ^ {\delta \nu^ {\prime}} A ^ {\omega \kappa^ {\prime}} \right] - \mathbb {E} _ {\ell} \left[ A ^ {\alpha \nu} A ^ {\beta \kappa} \right] \mathbb {E} _ {\ell} \left[ A ^ {\delta \nu^ {\prime}} A ^ {\omega \kappa^ {\prime}} \right]\right) \\ + \sum_ {\nu \kappa \nu^ {\prime} \kappa^ {\prime}} \mathbb {E} _ {\ell} \left[ A ^ {\alpha \nu} A ^ {\beta \kappa} A ^ {\delta \nu^ {\prime}} A ^ {\omega \kappa^ {\prime}} \right] \left(V ^ {\nu \nu^ {\prime}} V ^ {\kappa \kappa^ {\prime}} + V ^ {\nu \kappa^ {\prime}} V ^ {\nu^ {\prime} \kappa}\right). \\ \end{array}
$$
Now we can use Lemma C.3 and Lemma C.4 to compute the moments of  $A$ . For the second summand, we simply have:
$$
\mathbb {E} _ {\ell} \left[ A ^ {\alpha \nu} A ^ {\beta \kappa} A ^ {\delta \nu^ {\prime}} A ^ {\omega \kappa^ {\prime}} \right] = \delta_ {\alpha \nu} \delta_ {\beta \kappa} \delta_ {\delta \nu^ {\prime}} \delta_ {\omega \kappa^ {\prime}} + \mathcal {O} (n _ {k} \tau^ {- 2}),
$$
hence:
$$
\sum_ {\nu \kappa \nu^ {\prime} \kappa^ {\prime}} \mathbb {E} _ {\ell} \left[ A ^ {\alpha \nu} A ^ {\beta \kappa} A ^ {\delta \nu^ {\prime}} A ^ {\omega \kappa^ {\prime}} \right] \left(V ^ {\nu \nu^ {\prime}} V ^ {\kappa \kappa^ {\prime}} + V ^ {\nu \kappa^ {\prime}} V ^ {\nu^ {\prime} \kappa}\right) = V ^ {\alpha \delta} V ^ {\beta \omega} + V ^ {\alpha \omega} V ^ {\beta \delta} + \mathcal {O} (n _ {k} \tau^ {- 2}).
$$
For the first summand, recall from Lemma C.5 that:
$$
\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[ 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} \left(n _ {k} \tau^ {- 3}\right). \\ \end{array}
$$
$$
\begin{array}{l} n \sum_ {\nu \kappa \nu^ {\prime} \kappa^ {\prime}} V ^ {\nu \kappa} V ^ {\nu^ {\prime} \kappa^ {\prime}} \left(\mathbb {E} _ {\ell} \left[ A ^ {\alpha \nu} A ^ {\beta \kappa} A ^ {\delta \nu^ {\prime}} A ^ {\omega \kappa^ {\prime}} \right] - \mathbb {E} _ {\ell} \left[ A ^ {\alpha \nu} A ^ {\beta \kappa} \right] \mathbb {E} _ {\ell} \left[ A ^ {\delta \nu^ {\prime}} A ^ {\omega \kappa^ {\prime}} \right]\right) \\ = \frac {n n _ {k}}{\tau^ {2} m ^ {2}} \sum_ {\nu \kappa \nu^ {\prime} \kappa^ {\prime}} V ^ {\nu \kappa} V ^ {\nu^ {\prime} \kappa^ {\prime}} \Big (\delta_ {\alpha \nu} \delta_ {\delta \nu^ {\prime}} S _ {1} ^ {\beta \kappa , \omega \kappa^ {\prime}} + \delta_ {\alpha \nu} \delta_ {\omega \kappa^ {\prime}} S _ {1} ^ {\beta \kappa , \delta \nu^ {\prime}} + \delta_ {\beta \kappa} \delta_ {\delta \nu^ {\prime}} S _ {1} ^ {\alpha \nu , \omega \kappa^ {\prime}} + \delta_ {\beta \kappa} \delta_ {\omega \kappa^ {\prime}} S _ {1} ^ {\alpha \nu , \delta \nu^ {\prime}} \Big) \\ + \mathcal {O} (n n _ {k} \tau^ {- 3}) \\ = \frac {n n _ {k}}{\tau^ {2}} \underbrace {\frac {1}{m ^ {2}} \sum_ {\nu \kappa} \left(V ^ {\alpha \kappa} V ^ {\delta \nu} S _ {1} ^ {\beta \kappa , \omega \nu} + V ^ {\alpha \kappa} V ^ {\omega \nu} S _ {1} ^ {\beta \kappa , \delta \nu} + V ^ {\beta \nu} V ^ {\delta \kappa} S _ {1} ^ {\alpha \nu , \omega \kappa} + V ^ {\beta \nu} V ^ {\omega \kappa} S _ {1} ^ {\alpha \nu , \delta \kappa}\right)} _ {\mathcal {A} ^ {\alpha \beta \delta \omega}} \\ + \mathcal {O} (n n _ {k} \tau^ {- 3}). \\ \end{array}
$$
We are now ready to re-state and proof of Theorem 4.2.
$$
d V _ {t} = b (V _ {t}) d t + \Sigma (V _ {t}) ^ {1 / 2} d B _ {t}, \quad V _ {0} = \frac {1}{n} X _ {0} X _ {0} ^ {\top},
$$
the diffusion coefficient is defined by  $\Sigma (V) = \gamma^{2}(2 - \gamma^{2})\Sigma_{lin}(V) + \gamma^{4}\tau_{0}^{-2}[\mathcal{A}^{\alpha \beta ,\delta \omega}]_{\alpha \leq \beta ,\delta \leq \omega},$  and
Proof. Recall that:
Drift. Summing and subtracting  $\mathbb{E}_{\ell}[\mathcal{T}_2]$  (and using Lemma C.7), we have that:
$$
V _ {\ell + 1} ^ {\alpha \beta} = \lambda^ {2} V _ {\ell} ^ {\alpha \beta} + \gamma^ {2} \sum_ {\nu \kappa} V ^ {\nu \kappa} \mathbb {E} _ {\ell} \left[ A ^ {\alpha \nu} A ^ {\beta \kappa} \right] + \frac {\lambda \gamma}{\sqrt {n}} \mathcal {T} _ {1} ^ {\alpha \beta} + \frac {\gamma^ {2}}{\sqrt {n}} \left(\mathcal {T} _ {2} ^ {\alpha \beta} - \mathbb {E} _ {\ell} [ \mathcal {T} _ {2} ]\right),
$$
Re-stating Lemma C.3, we have that:
$$
\mathbb {E} _ {\ell} [ A ^ {\alpha \nu} A ^ {\beta \kappa} ] = \delta_ {\alpha \nu} \delta_ {\beta \kappa} + \frac {n _ {k}}{\tau^ {2} m ^ {2}} S _ {1} ^ {\alpha \nu , \beta \kappa} + \frac {n _ {k}}{2 \tau^ {2} m} (\delta_ {\beta \kappa} S _ {2} ^ {\alpha \nu} + \delta_ {\alpha \nu} S _ {2} ^ {\beta \kappa}) + O (n _ {k} \tau^ {- 3}).
$$
Plugging in the expression, and using  $\lambda^2 +\gamma^2 = 1$  , we have that the drift is:
$$
\begin{array}{l} \lambda^ {2} V _ {\ell} ^ {\alpha \beta} + \gamma^ {2} V _ {\ell} ^ {\alpha \beta} + \gamma^ {2} \frac {n _ {k}}{\tau^ {2} m} \sum_ {\nu \kappa} V ^ {\nu \kappa} \left(\frac {1}{m} S _ {1} ^ {\alpha \nu , \beta \kappa} + \frac {1}{2} (\delta_ {\beta \kappa} S _ {2} ^ {\alpha \nu} + \delta_ {\alpha \nu} S _ {2} ^ {\beta \kappa})\right) + O (n _ {k} \tau^ {- 3}) \\ = V _ {\ell} ^ {\alpha \beta} + \gamma^ {2} \frac {n _ {k}}{\tau^ {2}} \left(\frac {1}{m ^ {2}} \sum_ {\nu \kappa} V ^ {\nu \kappa} S _ {1} ^ {\alpha \nu , \beta \kappa} + \frac {1}{2 m} \sum_ {\nu} (V ^ {\beta \nu} S _ {2} ^ {\alpha \nu} + V ^ {\alpha \nu} S _ {2} ^ {\beta \nu})\right) + O (n _ {k} \tau^ {- 3}). \\ \end{array}
$$
From here, it is evident that in order to have the drift scaling as  $\mathcal{O}(1 / n)$  we need to choose:
$$
\tau^ {2} = \tau_ {0} ^ {2} n n _ {k}, \tag {26}
$$
where  $\tau_0 > 0$  is a constant.
Covariance. Recall that:
$$
V _ {\ell + 1} ^ {\alpha \beta} = \lambda^ {2} V _ {\ell} ^ {\alpha \beta} + \gamma^ {2} \sum_ {\nu \kappa} V ^ {\nu \kappa} \mathbb {E} _ {\ell} \left[ A ^ {\alpha \nu} A ^ {\beta \kappa} \right] + \frac {\lambda \gamma}{\sqrt {n}} \mathcal {T} _ {1} ^ {\alpha \beta} + \frac {\gamma^ {2}}{\sqrt {n}} \left(\mathcal {T} _ {2} ^ {\alpha \beta} - \mathbb {E} _ {\ell} [ \mathcal {T} _ {2} ]\right).
$$
Furthermore, we have set  $\lambda^2 +\gamma^2 = 1$  and  $\tau^2 = \tau_0^2 nn_k$  to have the right scaling for the drift.
What's left to is to compute the conditional covariance for  $V_{\ell +1}$  given  $\mathcal{F}_{\ell}$ . Noting that  $T_1^{\alpha \beta},T_2^{\delta \omega}$  are uncorrelated, i.e.  $\mathbb{E}_{\ell}\left[T_1^{\alpha \beta}T_2^{\delta \omega}\right] = 0$  (similarly to the case of Resnet with shaped ReLU Lemma B.3), we have that:
$$
\begin{array}{l} \operatorname {C o v} _ {\ell} \left(V _ {\ell + 1} ^ {\alpha \beta}, V _ {\ell + 1} ^ {\delta \omega}\right) = \mathbb {E} _ {\ell} \left[ \left(\lambda \gamma \mathcal {T} _ {1} ^ {\alpha \beta} + \gamma^ {2} \mathcal {T} _ {2} ^ {\alpha \beta} - \gamma^ {2} \mathbb {E} _ {\ell} [ \mathcal {T} _ {2} ^ {\alpha \beta} ]\right) \left(\lambda \gamma \mathcal {T} _ {1} ^ {\delta \omega} + \gamma^ {2} \mathcal {T} _ {2} ^ {\delta \omega} - \gamma^ {2} \mathbb {E} _ {\ell} [ \mathcal {T} _ {2} ^ {\delta \omega} ]\right) \right] \\ = \lambda^ {2} \gamma^ {2} \mathbb {E} _ {\ell} \left[ \mathcal {T} _ {1} ^ {\alpha \beta} \mathcal {T} _ {1} ^ {\delta \omega} \right] + \gamma^ {4} \mathbb {E} _ {\ell} \left[ \mathcal {T} _ {2} ^ {\alpha \beta} \mathcal {T} _ {2} ^ {\delta \omega} \right] - \gamma^ {4} \mathbb {E} _ {\ell} \left[ \mathcal {T} _ {2} ^ {\alpha \beta} \right] \mathbb {E} _ {\ell} \left[ \mathcal {T} _ {2} ^ {\delta \omega} \right], \\ \end{array}
$$
where we use  $\mathrm{Cov}_{\ell}$  to denote the conditional covariance given the sigma-algebra generated by  $\mathcal{F}_{\ell} = \sigma (\{X_{k}\}_{k\in [\ell ]})$
Using Lemma C.6 and Lemma C.8, we have that:
$$
\begin{array}{l} \operatorname {C o v} _ {\ell} \left(V _ {\ell + 1} ^ {\alpha \beta}, V _ {\ell + 1} ^ {\delta \omega}\right) = \lambda^ {2} \gamma^ {2} \left(2 \left(V ^ {\alpha \delta} V ^ {\beta \omega} + V ^ {\alpha \omega} V ^ {\beta \delta}\right)\right) + \gamma^ {4} \left(\frac {1}{\tau_ {0} ^ {2}} \mathcal {A} ^ {\alpha \beta \delta \omega} + V ^ {\alpha \delta} V ^ {\beta \omega} + V ^ {\alpha \omega} V ^ {\beta \delta}\right) \\ = \gamma^ {2} (2 - \gamma^ {2}) \left(V ^ {\alpha \delta} V ^ {\beta \omega} + V ^ {\alpha \omega} V ^ {\beta \delta}\right) + \frac {\gamma^ {4}}{\tau_ {0} ^ {2}} \mathcal {A} ^ {\alpha \beta \delta \omega}. \\ \end{array}
$$
Now we can apply Proposition A.2 for locally Lipschitz drift and covariance coefficients, which gives us the desired result in local convergence in the Skorohod topology.
We will also restate and prove Corollary 4.3.
Proof. To combine the results of Theorem 3.2 and Theorem 4.2, it is sufficient to combine the following (simplified) iterated Markov updates into one Markov chain
$$
U _ {\ell} = V _ {\ell} + \frac {b \left(V _ {\ell}\right)}{n} + \frac {\Sigma \left(V _ {\ell}\right) ^ {1 / 2} \xi_ {\ell}}{\sqrt {n}}, \quad V _ {\ell + 1} = U _ {\ell} + \frac {b _ {\text {R e L U}} \left(U _ {\ell}\right)}{n} + \frac {\Sigma_ {\text {l i n}} \left(U _ {\ell}\right) ^ {1 / 2} \xi_ {\ell} ^ {\prime}}{\sqrt {n}}, \tag {27}
$$
where  $\xi_{\ell},\xi_{\ell}^{\prime}$  are independent zero mean and identity covariance random vectors.
Since in the limit, we have that either updates are infinitesimal, i.e.
$$
\left| U _ {\ell} - V _ {\ell} \right| \xrightarrow {n \rightarrow \infty} 0 \quad \text {a l m o s t s u r e l y}, \tag {28}
$$
then we can write  $b_{\mathrm{ReLU}(U_\ell)} = \widehat{b}_{\mathrm{ReLU},n}(V_\ell, \omega_\ell)$ , where eventually we have that
$$
\lim  _ {n \rightarrow \infty} \mathbb {E} _ {\ell} \left[ \widehat {b} _ {n} \left(V _ {\ell}, \omega_ {\ell}\right)\right] = b \left(V _ {\ell}\right), \tag {29}
$$
so it will not contribute towards affecting the limit. Here 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 ]})$  . We can treat  $\Sigma_{\mathrm{lin}}(U_{\ell})$  similarly to get the Markov chain update
$$
V _ {\ell + 1} = V _ {\ell} + \frac {b \left(V _ {\ell}\right) + \widehat {b} _ {\text {R e L U} , n} \left(V _ {\ell} , \omega_ {\ell}\right)}{n} + \frac {\Sigma \left(V _ {\ell}\right) ^ {1 / 2} \xi_ {\ell}}{\sqrt {n}} + \frac {\widehat {\Sigma} _ {\text {l i n} , n} \left(V _ {\ell}\right) ^ {1 / 2} \xi_ {\ell} ^ {\prime}}{\sqrt {n}}, \tag {30}
$$
which converges to the following SDE with two Brownian motions using Proposition A.2.
$$
d V _ {t} = \left[ b \left(V _ {t}\right) + b _ {\text {r e s}} \left(V _ {t}\right) \right] d t + \Sigma \left(V _ {t}\right) ^ {1 / 2} d B _ {t} + \Sigma_ {\text {r e s}} \left(V _ {t}\right) ^ {1 / 2} d B _ {t} ^ {\prime}. \tag {31}
$$
Observe that since the two Brownian motions  $B_{t}, B_{t}^{\prime}$  are independent, it's equivalent to write
$$
\Sigma \left(V _ {t}\right) ^ {1 / 2} d B _ {t} + \Sigma_ {\text {r e s}} \left(V _ {t}\right) ^ {1 / 2} d B _ {t} ^ {\prime} \stackrel {{d}} {{=}} \left[ \Sigma \left(V _ {t}\right) + \Sigma_ {\text {r e s}} \left(V _ {t}\right) ^ {1 / 2} \right] ^ {1 / 2} d B _ {t}. \tag {32}
$$
We recover the desired results from considering a more general form of the iterated Markov updates as in Proposition A.2, which do not hinder the above derivation.
D Preliminary Experiments
We perform preliminary experiments to understand the effect of shaped attention on training deep Transformer architectures. In particular, we consider a pre-training masked language modeling task, where we mask  $15\%$  of the tokens. We use a subset of the English Wikipedia 20220301.en and English bookcorpus datasets [69, 70]. As a baseline, we adopt a Pre-LN Transformer encoder architecture with 18 or 24 blocks. For the residual feedforward layer, we shape the ReLU activation according to Eq. 4, by changing its negative slope to  $s_{-} = 1 - 1 / \sqrt{n}$  instead of 0. We then incorporate our shaped attention by replacing the attention mechanism as dictated in Eq. 9. We also add scalar multipliers  $\gamma_1,\gamma_2\in \mathbb{R}$  both to the identity and centering terms of the shaped attention:
$$
A _ {\ell} = \gamma_ {1} I + \operatorname {S o f t m a x} \left(\tau^ {- 1} Y _ {\ell}\right) - \gamma_ {2} \frac {1}{m} \mathbf {1 1} ^ {\top}, \quad \tau = \tau_ {0} \sqrt {n n _ {k}}, \tag {33}
$$
and propose two ways to set  $\gamma_{1},\gamma_{2}$  during training. In both alternatives we initialize  $\gamma_{1},\gamma_{2} = 1$  , thus leveraging the stability properties of shaped attention at initialization. During training, we either:
1. Recover. Linearly decrease  $\gamma_{1},\gamma_{2}$  to zero with in the first 4000 steps, thus recovering the standard attention layer. Apply the same schedule for the shaped-ReLU slope  $s_{-}$ , recovering the usual ReLU activation. This approach recovers the vanilla Transformer architecture (without LayerNorm).
2. Learn. Learn all the shaping parameters  $\gamma_{1},\gamma_{2}$  and  $s_{-}$
The intuitive rationale behind these choices is that at initialization we want good signal propagation and a non-degenerate covariance (according to our theory, this requires the shaped attention and ReLU). On the other hand, we also allow the model to more dynamically make use of the nonlinearity during training to modify the correlation structure with Recover or Learn. We report that without either adjustment, shaped attention is still trainable but at much slower rates.
We also incorporate the  $\frac{1}{\sqrt{n}}$  factor into the initialization of the queries and keys weights by decreasing the variance of their entries by a factor of  $\frac{1}{n}$ . This allows us to not re-tune the learning rate for the queries and keys. We stress that at at initialization, the two formulations ( $\tau = \sqrt{nn_k}$  and  $\tau = \sqrt{n_k}$  with decreased weight's variance) are equivalent. In both alternatives we train all the skip connection parameters  $\lambda$ ,  $\gamma$ , and initialize  $\gamma$  in the grid (0.05, 0.1, 0.2) and set  $\tau_0 = 1$ . We report training instabilities (loss divergence) for larger values of  $\gamma$ . All models — including the baselines — use Adam [71] with learning rate warmup of 4000 steps, the learning rate is tuned in the grid (0.0001, 0.0005, 0.001, 0.005). We report the train/test loss after  $100K$  optimization steps, averaged over 4 random seeds. All the other experimental details can be found in Appendix D.1.
The Shaped Transformer is Trainable. In Table 1, we compare the train/test loss of the two variant of shaped attention with the baseline Pre-LN model. Notice that our model (in both variants) achieves comparable performance to standard Transformers across all the reported values of  $\gamma$ .
GLUE evaluation Furthermore, we evaluate the trained models on three datasets from the GLUE benchmark [72] and summarize the results in Table 2. These demonstrate that our shaped Transformers hold promise by outperforming our pre-ln baseline.
d=18	d=24
γ	Train Loss	Test Loss	Train Loss	Test Loss
Learn	0.05	2.09±0.02	2.07±0.02	2.03±0.02	2.03±0.01
0.1	2.10±0.02	2.07±0.01	2.02±0.01	2.02±0.03
0.2	2.07±0.05	2.06±0.02	2.09±0.04	2.09±0.03
Recover	0.05	2.05±0.03	2.04±0.02	2.02±0.04	2.00±0.01
0.1	2.05±0.01	2.04±0.01	2.06±0.03	2.00±0.02
0.2	2.01±0.03	2.03±0.01	2.05±0.04	2.00±0.01
Baseline	2.08±0.03	2.07±0.01	2.08±0.03	2.01±0.01
Table 1: Training and test loss of the proposed shaped attention, both in the "Recover" and "Learn" alternatives. Baseline refers to the vanilla Pre-LN Transformer. We report the best run under the learning rates (0.0001, 0.0005, 0.001, 0.005), averaged over 4 random seeds. We include the confidence interval of  $\pm$  one standard deviation.
Model	COLA	MRPC	RTE
d=18	Baseline	0.139±0.008	0.797±0.010	0.504±0.024
Learn, γ = 0.2	0.182±0.041	0.813±0.005	0.528±0.019
Recover, γ = 0.2	0.221±0.042	0.809±0.009	0.520±0.029
d=24	Baseline	0.022±0.055	0.785±0.009	0.48±0.046
Learn, γ = 0.05	0.150±0.026	0.812±0.004	0.513±0.019
Recover, γ = 0.05	0.211±0.039	0.803±0.012	0.529±0.019
BERT (Geiping et al.)	0.103	74.8	0.509
Table 2: Evaluation of the pretrained models on a subset of the GLUE benchmark. We take the checkpoints of the baselines and selected shaped Transformers ( $\gamma = 0.2$  for the shallower model  $d = 18$ ,  $\gamma = 0.05$  for the deeper model  $d = 24$ ). Then, we fine-tune on each of three datasets from GLUE (COLA, MRPC, RTE) for 5 epochs and Adam optimizer  $lr = 5e - 4$  with batch size 16, and report test evaluation metric for the corresponding task.