SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
d V _ {t} = b _ {\text {R e L U}} \left(V _ {t}\right) d t + \Sigma_ {\ln} \left(V _ {t}\right) ^ {1 / 2} d B _ {t}, \quad V _ {0} = \frac {1}{n} X _ {0} X _ {0} ^ {\top}, \tag {20}
$$
where the coefficients were given in Theorem 3.2.
We remark that if the output is defined as  $X_{\mathrm{out}} = \frac{1}{\sqrt{n}} W_{\mathrm{out}}X_d \in \mathbb{R}^{n_{\mathrm{out}}}$ , then we can recover the output distribution as
$$
X _ {\text {o u t}} \sim N \left(0, V _ {T} \otimes I _ {n _ {\text {o u t}}}\right), \tag {21}
$$
where we treated  $X_{\mathrm{out}}$  as a vector in  $\mathbb{R}^{mn_{\mathrm{out}}}$ .
B SDE for Residual Network
Recall that we adopt the following model:
$$
X _ {\ell + 1} = \lambda X _ {\ell} + \gamma \frac {1}{\sqrt {n}} \sigma_ {s} \left(\sqrt {\frac {c}{n}} X _ {\ell} W _ {\ell} ^ {\text {p r e}}\right) W _ {\ell} ^ {\text {p o s t}}, \tag {22}
$$
where  $\sigma_s(x) := s_+ \max(x, 0) + s_- \min(x, 0)$  is the shaped ReLU with slopes  $s_\pm := 1 + \frac{c_\pm}{\sqrt{n}}$  for some constants  $c_+, c_- \in \mathbb{R}$ . We let the weights be  $(W_\ell^{\mathrm{pre}})_{ij}(W_\ell^{\mathrm{post}})_{ij} \stackrel{\mathrm{iid}}{\sim} \mathcal{N}(0, 1)$  and  $c^{-1} = \mathbb{E}\sigma_s(g)^2$  for  $g \sim \mathcal{N}(0, 1)$  is the He initialization constant [55].
From here onwards, we will define the filtration  $\mathcal{F}_{\ell} = \sigma (\{X_{k}\}_{k\in [\ell ]})$ , where  $\sigma (\cdot)$  denotes the sigma-algebra generated by the random variable. Furthermore, we will define the conditional expectation  $\mathbb{E}_{\ell}[\cdot ] = \mathbb{E}[\cdot |\mathcal{F}_{\ell}]$ .
We are interested in studying the neural covariance, i.e.  $V_{\ell}^{\alpha \beta} \coloneqq \frac{c}{n}\langle x_{\ell}^{\alpha},x_{\ell}^{\beta}\rangle$  where  $\ell \in [d]$  indexes the layers and  $d$  is the depth. In particular, we are interested in understanding the simultaneous limit  $(n,d)\to \infty$ , where the ratio  $t = d / n$  remains constant.
From now on, we will remove the dependence on  $\ell$ . Defining the residual branch as  $f(X,W) \coloneqq \frac{1}{\sqrt{n}} \sigma_s \left( \sqrt{\frac{c}{n}} X W^{\text{pre}} \right) W^{\text{post}}$ , we can write  $V^{\alpha \beta}$  as:
$$
V ^ {\alpha \beta} = \lambda^ {2} X + \lambda \gamma \langle x ^ {\alpha}, f (X, W) ^ {\beta} \rangle + \lambda \gamma \langle x ^ {\beta}, f (X, W) ^ {\alpha} \rangle + \gamma^ {2} \langle f (X, W) ^ {\alpha}, f (X, W) ^ {\beta} \rangle .
$$
For the cross product terms, we get:
$$
\begin{array}{l} \langle x ^ {\alpha}, f (X, W) ^ {\beta} \rangle = \frac {\sqrt {c}}{n} \langle x ^ {\alpha}, \left(\sigma_ {s} \left(X W ^ {\text {p r e}}\right) W ^ {\text {p o s t}}\right) ^ {\beta} \rangle = \frac {\sqrt {c}}{n} \sum_ {i, j} ^ {n} x _ {i} ^ {\alpha} W _ {j i} ^ {\text {p o s t}} \sigma_ {s} \left(\sum_ {j ^ {\prime}} x _ {j ^ {\prime}} ^ {\beta} W _ {j ^ {\prime} j} ^ {\text {p r e}}\right), \\ \langle x ^ {\beta}, f (X, W) ^ {\alpha} \rangle = \frac {\sqrt {c}}{n} \sum_ {i, j} ^ {n} x _ {i} ^ {\beta} W _ {j i} ^ {\text {p o s t}} \sigma_ {s} \left(\sum_ {j ^ {\prime}} x _ {j ^ {\prime}} ^ {\alpha} W _ {j ^ {\prime} j} ^ {\text {p r e}}\right). \\ \end{array}
$$
For the term in  $\gamma^2$ :
$$
\begin{array}{l} \frac {c}{n ^ {2}} \left\langle \left(\sigma_ {s} (X W ^ {\text {p r e}}) W ^ {\text {p o s t}}\right) ^ {\alpha}, \left(\sigma_ {s} (X W ^ {\text {p r e}}) W ^ {\text {p o s t}}\right) ^ {\beta} \right\rangle = \frac {c}{n ^ {2}} \sum_ {i = 1} ^ {n} \left(\sigma_ {s} (X W ^ {\text {p r e}}) W ^ {\text {p o s t}}\right) _ {\alpha i} \left(\sigma_ {s} (X W ^ {\text {p r e}}) W ^ {\text {p o s t}}\right) _ {\beta i} \\ = \frac {c}{n ^ {2}} \sum_ {i, j, j ^ {\prime} = 1} ^ {n} W _ {j i} ^ {\mathrm {p o s t}} W _ {j ^ {\prime} i} ^ {\mathrm {p o s t}} \sigma_ {s} \left(\sum_ {k} x _ {k} ^ {\alpha} W _ {k j} ^ {\mathrm {p r e}}\right) \sigma_ {s} \left(\sum_ {k ^ {\prime}} x _ {k ^ {\prime}} ^ {\beta} W _ {k ^ {\prime} j ^ {\prime}} ^ {\mathrm {p o s t}}\right). \\ \end{array}
$$
We will define the following terms
$$
\mathcal {T} _ {1} ^ {\alpha \beta} := \frac {c \sqrt {c}}{n \sqrt {n}} \sum_ {i, j = 1} ^ {n} W _ {j i} ^ {\text {p o s t}} \left(x _ {i} ^ {\alpha} \sigma_ {s} \left(\sum_ {j ^ {\prime}} x _ {j ^ {\prime}} ^ {\beta} W _ {j ^ {\prime} j} ^ {\text {p r e}}\right) + x _ {i} ^ {\beta} \sigma_ {s} \left(\sum_ {j ^ {\prime}} x _ {j ^ {\prime}} ^ {\alpha} W _ {j ^ {\prime} j} ^ {\text {p r e}}\right)\right),
$$
$$
\mathcal {T} _ {2} ^ {\alpha \beta} := \frac {c ^ {2}}{n ^ {2} \sqrt {n}} \sum_ {i, j, j ^ {\prime} = 1} ^ {n} W _ {j i} ^ {\mathrm {p o s t}} W _ {j ^ {\prime} i} ^ {\mathrm {p o s t}} \sigma_ {s} \left(\sum_ {k} x _ {k} ^ {\alpha} W _ {k j} ^ {\mathrm {p r e}}\right) \sigma_ {s} \left(\sum_ {k ^ {\prime}} x _ {k ^ {\prime}} ^ {\beta} W _ {k ^ {\prime} j ^ {\prime}} ^ {\mathrm {p r e}}\right).
$$
Hence we get the following update for  $V_{\ell}^{\alpha \beta}\coloneqq \frac{c}{n}\langle x_{\ell}^{\alpha},x_{\ell}^{\beta}\rangle$  ..
$$
V _ {\ell + 1} ^ {\alpha \beta} = \lambda^ {2} V _ {\ell} ^ {\alpha \beta} + \frac {\gamma \lambda}{\sqrt {n}} \mathcal {T} _ {1} ^ {\alpha \beta} + \frac {\gamma^ {2}}{\sqrt {n}} \mathcal {T} _ {2} ^ {\alpha \beta}.
$$
It is easy to see that the cross product terms have (conditional) mean zero. For the term with  $\gamma^2$ :
$$
\mathbb {E} _ {\ell} \left[ \mathcal {T} _ {2} ^ {\alpha \beta} \right] = \sqrt {n} \sqrt {V ^ {\alpha \alpha} V ^ {\beta \beta}} c K _ {1} \left(\rho^ {\alpha \beta}\right), \tag {23}
$$
where  $K_{1}(\rho^{\alpha \beta})\coloneqq \mathbb{E}[\sigma_{s}(g^{\alpha})\sigma_{s}(g^{\beta})]$  where  $g^{\alpha},g^{\beta}\sim \mathcal{N}(0,1)$  and  $\mathbb{E}[g^{\alpha}g^{\beta}] = \frac{\langle x^{\alpha},x^{\beta}\rangle}{\|x^{\alpha}\|\|x^{\beta}\|}$ . 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 ]})$
M. Li et al. [18] showed that if the non-linearity scaling exponent is  $p = 1 / 2$ , then:
$$
c K _ {1} (\rho) = \rho + \frac {\nu (\rho)}{n} + \mathcal {O} (n ^ {- 3 / 2}),
$$
where  $\nu (\rho) = \frac{(c_{+} + c_{-})^{2}}{2\pi}\left(\sqrt{1 - \rho^{2}} +\rho \mathrm{arccos}(\rho)\right)$ . Using this result, and summing and subtracting the mean of  $\mathcal{T}_2$ :
$$
\begin{array}{l} V _ {\ell + 1} ^ {\alpha \beta} = \lambda^ {2} V _ {\ell} ^ {\alpha \beta} + \gamma^ {2} \sqrt {V ^ {\alpha \alpha} V ^ {\beta \beta}} c K _ {1} (\rho^ {\alpha \beta}) + \frac {\gamma \lambda}{\sqrt {n}} \mathcal {T} _ {1} ^ {\alpha \beta} + \frac {\gamma^ {2}}{\sqrt {n}} (\mathcal {T} _ {2} ^ {\alpha \beta} - \mathbb {E} _ {\ell} [ \mathcal {T} _ {2} ^ {\alpha \beta} ]) \\ = \lambda^ {2} V _ {\ell} ^ {\alpha \beta} + \gamma^ {2} V _ {\ell} ^ {\alpha \beta} + \gamma^ {2} \sqrt {V ^ {\alpha \alpha} V ^ {\beta \beta}} \frac {\nu (\rho^ {\alpha \beta})}{n} + \frac {\gamma \lambda}{\sqrt {n}} \mathcal {T} _ {1} ^ {\alpha \beta} + \frac {\gamma^ {2}}{\sqrt {n}} (\mathcal {T} _ {2} ^ {\alpha \beta} - \mathbb {E} _ {\ell} [ \mathcal {T} _ {2} ^ {\alpha \beta} ]) + \mathcal {O} (n ^ {- 3 / 2}). \\ \end{array}
$$
Using  $\lambda^2 +\gamma^2 = 1$
$$
V _ {\ell + 1} ^ {\alpha \beta} = V _ {\ell} ^ {\alpha \beta} + \gamma^ {2} \sqrt {V ^ {\alpha \alpha} V ^ {\beta \beta}} \frac {\nu (\rho^ {\alpha \beta})}{n} + \frac {\gamma \lambda}{\sqrt {n}} \mathcal {T} _ {1} ^ {\alpha \beta} + \frac {\gamma^ {2}}{\sqrt {n}} (\mathcal {T} _ {2} ^ {\alpha \beta} - \mathbb {E} _ {\ell} [ \mathcal {T} _ {2} ^ {\alpha \beta} ]) + \mathcal {O} (n ^ {- 3 / 2}).
$$
Furthermore, we need second order moments of  $\mathcal{T}_1$  and  $\mathcal{T}_2$ . We derive them in the following Lemmas:
Lemma B.1.
$$
\begin{array}{l} \mathbb {E} _ {\ell} \left[ \mathcal {T} _ {1} ^ {\alpha \beta} \mathcal {T} _ {1} ^ {\delta \omega} \right] = V ^ {\alpha \delta} \sqrt {V ^ {\beta \beta} V ^ {\omega \omega}} c K _ {1} \left(\rho^ {\beta \omega}\right) + V ^ {\alpha \omega} \sqrt {V ^ {\beta \beta} V ^ {\delta \delta}} c K _ {1} \left(\rho^ {\beta \delta}\right) \\ + V ^ {\beta \delta} \sqrt {V ^ {\alpha \alpha} V ^ {\omega \omega}} c K _ {1} (\rho^ {\alpha \omega}) + V ^ {\beta \omega} \sqrt {V ^ {\alpha \alpha} V ^ {\delta \delta}} c K _ {1} (\rho^ {\alpha \delta}) \\ = 2 V ^ {\alpha \delta} V ^ {\beta \omega} + 2 V ^ {\alpha \omega} V ^ {\beta \delta} + \mathcal {O} (n ^ {- 1}). \\ \end{array}
$$
Proof. Recall the definition of  $\mathcal{T}_1^{\alpha \beta}$ :
$$
\begin{array}{l} \mathbb {E} _ {\ell} [ \mathcal {T} _ {1} ^ {\alpha \beta} \mathcal {T} _ {1} ^ {\delta \omega} ] = \frac {c ^ {3}}{n ^ {3}} \sum_ {i j i ^ {\prime} j ^ {\prime}} \mathbb {E} \left[ W _ {j i} W _ {j ^ {\prime} i ^ {\prime}} \right] \mathbb {E} _ {\ell} \left[ \left(x _ {i} ^ {\alpha} \left\| x ^ {\beta} \right\| \sigma_ {s} (g _ {j} ^ {\beta}) + x _ {i} ^ {\beta} \left\| x ^ {\alpha} \right\| \sigma_ {s} (g _ {j} ^ {\alpha})\right) \left(x _ {i} ^ {\delta} \left\| x ^ {\omega} \right\| \sigma_ {s} (g _ {j} ^ {\omega}) + x _ {i} ^ {\omega} \left\| x ^ {\delta} \right\| \sigma_ {s} (g _ {j} ^ {\delta})\right) \right] \\ = \frac {c ^ {3}}{n ^ {3}} \sum_ {i j} \mathbb {E} _ {\ell} \left[ x _ {i} ^ {\alpha} x _ {i} ^ {\delta} \left\| x ^ {\beta} \right\| \left\| x ^ {\omega} \right\| \sigma_ {s} (g _ {j} ^ {\beta}) \sigma_ {s} (g _ {j} ^ {\omega}) + x _ {i} ^ {\alpha} x _ {i} ^ {\omega} \left\| x ^ {\beta} \right\| \left\| x ^ {\delta} \right\| \sigma_ {s} (g _ {j} ^ {\beta}) \sigma_ {s} (g _ {j} ^ {\delta}) \right. \\ + x _ {i} ^ {\beta} x _ {i} ^ {\delta} \left\| x ^ {\alpha} \right\| \left\| x ^ {\omega} \right\| \sigma_ {s} (g _ {j} ^ {\alpha}) \sigma_ {s} (g _ {j} ^ {\omega}) + x _ {i} ^ {\beta} x _ {i} ^ {\omega} \left\| x ^ {\alpha} \right\| \left\| x ^ {\delta} \right\| \sigma_ {s} (g _ {j} ^ {\alpha}) \sigma_ {s} (g _ {j} ^ {\delta}) \Bigg ] \\ = V ^ {\alpha \delta} \sqrt {V ^ {\beta \beta} V ^ {\omega \omega}} c K _ {1} \left(\rho^ {\beta \omega}\right) + V ^ {\alpha \omega} \sqrt {V ^ {\beta \beta} V ^ {\delta \delta}} c K _ {1} \left(\rho^ {\beta \delta}\right) \\ + V ^ {\beta \delta} \sqrt {V ^ {\alpha \alpha} V ^ {\omega \omega}} c K _ {1} \left(\rho^ {\alpha \omega}\right) + V ^ {\beta \omega} \sqrt {V ^ {\alpha \alpha} V ^ {\delta \delta}} c K _ {1} \left(\rho^ {\alpha \delta}\right) \\ \end{array}
$$
$$
\mathbb {E} _ {\ell} \left[ \mathcal {T} _ {2} ^ {\alpha \beta} \right] = c \sqrt {n} \sqrt {V ^ {\alpha \alpha} V ^ {\beta \beta}} K _ {1} \left(\rho^ {\alpha \beta}\right) = \sqrt {n} V ^ {\alpha \beta} + \frac {1}{\sqrt {n}} \sqrt {V ^ {\alpha \alpha} V ^ {\beta \beta}} \nu (\rho) + \mathcal {O} \left(n ^ {- 1}\right),
$$
$$
\mathbb {E} _ {\ell} [ \mathcal {T} _ {2} ^ {\alpha \beta} \mathcal {T} _ {2} ^ {\delta \omega} ] - \mathbb {E} _ {\ell} [ \mathcal {T} _ {2} ^ {\alpha \beta} ] \mathbb {E} [ \mathcal {T} _ {2} ^ {\delta \omega} ] = 2 V ^ {\alpha \delta} V ^ {\beta \omega} + 2 V ^ {\alpha \omega} V ^ {\beta \delta} + \mathcal {O} (n ^ {- 1}).
$$
Proof. Recall the definition of  $\mathcal{T}_2^{\alpha \beta}$ :
$$
\begin{array}{l} \mathcal {T} _ {2} ^ {\alpha \beta}: = \frac {c ^ {2}}{n ^ {2} \sqrt {n}} \sum_ {i, j, j ^ {\prime} = 1} ^ {n} W _ {j i} ^ {\text {p o s t}} W _ {j ^ {\prime} i} ^ {\text {p o s t}} \sigma_ {s} \left(\sum_ {k} x _ {k} ^ {\alpha} W _ {k j} ^ {\text {p r e}}\right) \sigma_ {s} \left(\sum_ {k ^ {\prime}} x _ {k ^ {\prime}} ^ {\beta} W _ {k ^ {\prime} j ^ {\prime}} ^ {\text {p r e}}\right) \\ = \frac {c ^ {2}}{n ^ {2} \sqrt {n}} \| x ^ {\alpha} \| \| x ^ {\beta} \| \sum_ {i, j, j ^ {\prime} = 1} ^ {n} W _ {j i} ^ {\text {p o s t}} W _ {j ^ {\prime} i} ^ {\text {p o s t}} \sigma_ {s} (g _ {j} ^ {\alpha}) \sigma_ {s} (g _ {j ^ {\prime}} ^ {\beta}) \\ \end{array}
$$
For the mean, using the independence between  $W^{\mathrm{pre}}$  and  $W^{\mathrm{post}}$ , we have that:
$$
\begin{array}{l} \mathbb {E} _ {\ell} \left[ \mathcal {T} _ {2} ^ {\alpha \beta} \right] = \frac {c ^ {2}}{n ^ {2} \sqrt {n}} \| x ^ {\alpha} \| \| x ^ {\beta} \| \sum_ {i, j, j ^ {\prime} = 1} ^ {n} \mathbb {E} \left[ W _ {j i} ^ {\mathrm {p o s t}} W _ {j ^ {\prime} i} ^ {\mathrm {p o s t}} \right] \mathbb {E} \left[ \sigma_ {s} (g _ {j} ^ {\alpha}) \sigma_ {s} (g _ {j ^ {\prime}} ^ {\beta}) \right] \\ = \frac {c ^ {2}}{n \sqrt {n}} \| x ^ {\alpha} \| \| x ^ {\beta} \| \sum_ {j = 1} ^ {n} \mathbb {E} \left[ \sigma_ {s} \left(g _ {j} ^ {\alpha}\right) \sigma_ {s} \left(g _ {j} ^ {\beta}\right) \right] \\ = \frac {c ^ {2}}{\sqrt {n}} \| x ^ {\alpha} \| \| x ^ {\beta} \| K _ {1} (\rho^ {\alpha \beta}) \\ = c \sqrt {n} \sqrt {V ^ {\alpha \alpha} V ^ {\beta \beta}} K _ {1} (\rho^ {\alpha \beta}), \\ \end{array}
$$
which is the desired result. The final expression is the result of the aforementioned expansion for  $K_{1}$  [18]. For the covariance, we have that:
$$
\begin{array}{l} \mathbb{E}_{\ell}[\mathcal{T}_{2}^{\alpha \beta}\mathcal{T}_{2}^{\delta \omega}] = \frac{c^{4}}{n^{5}}\| x^{\alpha}\| \left\| x^{\beta}\right\| \left\| x^{\delta}\right\| \left\| x^{\omega}\right\| \\ \sum_{i,j,j^{\prime},i^{\prime},k,k^{\prime}}\mathbb{E}\left[W_{j i}^{\mathrm{post}}W_{j^{\prime}i}^{\mathrm{post}}W_{k i^{\prime}}^{\mathrm{post}}W_{k^{\prime}i^{\prime}}^{\mathrm{post}}\right]\mathbb{E}\left[\sigma_{s}(g_{j}^{\alpha})\sigma_{s}(g_{j^{\prime}}^{\beta})\sigma_{s}(g_{k}^{\delta})\sigma_{s}(g_{k^{\prime}}^{\omega})\right] \\ = \frac {c ^ {2}}{n ^ {3}} \sqrt {V ^ {\alpha \alpha} V ^ {\beta \beta} V ^ {\delta \delta} V ^ {\omega \omega}} \sum_ {i, j, j ^ {\prime}, i ^ {\prime}, k, k ^ {\prime}} \left(\delta_ {j j ^ {\prime}} \delta_ {k k ^ {\prime}} + \delta_ {j k} \delta_ {i i ^ {\prime}} \delta_ {j ^ {\prime} k ^ {\prime}} + \delta_ {j k ^ {\prime}} \delta_ {i i ^ {\prime}} \delta_ {j ^ {\prime} k}\right) \mathbb {E} \left[ \sigma_ {s} (g _ {j} ^ {\alpha}) \sigma_ {s} (g _ {j ^ {\prime}} ^ {\beta}) \sigma_ {s} (g _ {k} ^ {\delta}) \sigma_ {s} (g _ {k ^ {\prime}} ^ {\omega}) \right] \\ \end{array}
$$
Let's look at each term of the sum separately. For the first term we have that:
$$
\begin{array}{l} \sum_ {i, j, j ^ {\prime}, i ^ {\prime}, k, k ^ {\prime}} \delta_ {j j ^ {\prime}} \delta_ {k k ^ {\prime}} \mathbb {E} \left[ \sigma_ {s} (g _ {j} ^ {\alpha}) \sigma_ {s} (g _ {j ^ {\prime}} ^ {\beta}) \sigma_ {s} (g _ {k} ^ {\delta}) \sigma_ {s} (g _ {k ^ {\prime}} ^ {\omega}) \right] \\ = n ^ {2} \sum_ {j, k} \mathbb {E} \left[ \sigma_ {s} \left(g _ {j} ^ {\alpha}\right) \sigma_ {s} \left(g _ {j} ^ {\beta}\right) \sigma_ {s} \left(g _ {k} ^ {\delta}\right) \sigma_ {s} \left(g _ {k} ^ {\omega}\right) \right] \\ = n ^ {2} \sum_ {j} \mathbb {E} \left[ \sigma_ {s} \left(g _ {j} ^ {\alpha}\right) \sigma_ {s} \left(g _ {j} ^ {\beta}\right) \sigma_ {s} \left(g _ {j} ^ {\delta}\right) \sigma_ {s} \left(g _ {j} ^ {\omega}\right) \right] + n ^ {2} \sum_ {j \neq k} \mathbb {E} \left[ \sigma_ {s} \left(g _ {j} ^ {\alpha}\right) \sigma_ {s} \left(g _ {j} ^ {\beta}\right) \right] \mathbb {E} \left[ \sigma_ {s} \left(g _ {k} ^ {\delta}\right) \sigma_ {s} \left(g _ {k} ^ {\omega}\right) \right] \\ = n ^ {3} K _ {2} (\rho^ {\alpha \beta \delta \omega}) + n ^ {3} (n - 1) K _ {1} (\rho^ {\alpha \beta}) K _ {1} (\rho^ {\delta \omega}), \\ \end{array}
$$
where we have defined the fourth moment of the shaped activation:  $K_{2}(\rho^{\alpha \beta \delta \omega})\coloneqq \mathbb{E}\left[\sigma_{s}(g^{\alpha})\sigma_{s}(g^{\beta})\sigma_{s}(g^{\delta})\sigma_{s}(g^{\omega})\right]$  for which it holds that [18, Lemma C.2]:
$$
K _ {2} (\rho^ {\alpha \beta \delta \omega}) = \mathbb {E} [ g ^ {\alpha} g ^ {\beta} g ^ {\delta} g ^ {\omega} ] + \mathcal {O} (n ^ {- 1 / 2}) = \rho^ {\alpha \beta} \rho^ {\delta \omega} + \rho^ {\alpha \delta} \rho^ {\beta \omega} + \rho^ {\alpha \omega} \rho^ {\beta \delta} + \mathcal {O} (n ^ {- 1 / 2}).
$$
The other two summands can be solved similarly:
$$
\sum_ {i, j, j ^ {\prime}, i ^ {\prime}, k, k ^ {\prime}} \delta_ {j k} \delta_ {i i ^ {\prime}} \delta_ {j ^ {\prime} k ^ {\prime}} \mathbb {E} \left[ \sigma_ {s} (g _ {j} ^ {\alpha}) \sigma_ {s} (g _ {j ^ {\prime}} ^ {\beta}) \sigma_ {s} (g _ {k} ^ {\delta}) \sigma_ {s} (g _ {k ^ {\prime}} ^ {\omega}) \right] = n ^ {2} K _ {2} (\rho^ {\alpha \beta \delta \omega}) + n ^ {2} (n - 1) K _ {1} (\rho^ {\alpha \delta}) K _ {1} (\rho^ {\beta \omega}),
$$
$$
\sum_ {i, j, j ^ {\prime}, i ^ {\prime}, k, k ^ {\prime}} \delta_ {j k ^ {\prime}} \delta_ {i i ^ {\prime}} \delta_ {j ^ {\prime} k} \mathbb {E} \left[ \sigma_ {s} (g _ {j} ^ {\alpha}) \sigma_ {s} (g _ {j ^ {\prime}} ^ {\beta}) \sigma_ {s} (g _ {k} ^ {\delta}) \sigma_ {s} (g _ {k ^ {\prime}} ^ {\omega}) \right] = n ^ {2} K _ {2} (\rho^ {\alpha \beta \delta \omega}) + n ^ {2} (n - 1) K _ {1} (\rho^ {\alpha \omega}) K _ {1} (\rho^ {\beta \delta}).
$$
Hence, summing the three terms:
$$
\begin{array}{l} \mathbb {E} _ {\ell} \left[ \mathcal {T} _ {2} ^ {\alpha \beta} \mathcal {T} _ {2} ^ {\delta \omega} \right] = c ^ {2} \sqrt {V ^ {\alpha \alpha} V ^ {\beta \beta} V ^ {\delta \delta} V ^ {\omega \omega}} \Big (K _ {2} (\rho^ {\alpha \beta \delta \omega}) + (n - 1) K _ {1} (\rho^ {\alpha \beta}) K _ {1} (\rho^ {\delta \omega}) \\ \left. + \frac {1}{n} \left(K _ {2} \left(\rho^ {\alpha \beta \delta \omega}\right) + (n - 1) K _ {1} \left(\rho^ {\alpha \delta}\right) K _ {1} \left(\rho^ {\beta \omega}\right)\right) + \frac {1}{n} \left(K _ {2} \left(\rho^ {\alpha \beta \delta \omega}\right) + (n - 1) K _ {1} \left(\rho^ {\alpha \omega}\right) K _ {1} \left(\rho^ {\beta \delta}\right)\right)\right) \\ = c ^ {2} \sqrt {V ^ {\alpha \alpha} V ^ {\beta \beta} V ^ {\delta \delta} V ^ {\omega \omega}} \Big (K _ {2} (\rho^ {\alpha \beta \delta \omega}) + n K _ {1} (\rho^ {\alpha \beta}) K _ {1} (\rho^ {\delta \omega}) - K _ {1} (\rho^ {\alpha \beta}) K _ {1} (\rho^ {\delta \omega}) \\ + K _ {1} \left(\rho^ {\alpha \delta}\right) K _ {1} \left(\rho^ {\beta \omega}\right) + K _ {1} \left(\rho^ {\alpha \omega}\right) K _ {1} \left(\rho^ {\beta \delta}\right)\left. \right) + \mathcal {O} (n ^ {- 1}). \\ \end{array}
$$
Now, subtracting  $\mathbb{E}_{\ell}[\mathcal{T}_2^{\alpha \beta}]\mathbb{E}_{\ell}[\mathcal{T}_2^{\delta \omega}]$ , using the aforementioned expansions for  $K_{1}$  and  $K_{2}$ :
$$
\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] = c ^ {2} \sqrt {V ^ {\alpha \alpha} V ^ {\beta \beta} V ^ {\delta \delta} V ^ {\omega \omega}} \left(\rho^ {\alpha \beta} \rho^ {\delta \omega} + \rho^ {\alpha \delta} \rho^ {\beta \omega} + \rho^ {\alpha \omega} \rho^ {\beta \delta}\right) \\ - V ^ {\alpha \beta} V ^ {\delta \omega} + V ^ {\alpha \delta} V ^ {\beta \omega} + V ^ {\alpha \omega} V ^ {\beta \delta} + \mathcal {O} (n ^ {- 1}) \\ = 2 V ^ {\alpha \delta} V ^ {\beta \omega} + 2 V ^ {\alpha \omega} V ^ {\beta \delta} + \mathcal {O} (n ^ {- 1}), \\ \end{array}
$$
where in the final step we have used the fact that  $c = 1 + \mathcal{O}(n^{-1 / 2})$ . This completes the proof.
Finally we also have that  $\mathcal{T}_1$  and  $\mathcal{T}_2$  are uncorrelated, i.e.
$$
\mathbb {E} _ {\ell} \left[ \mathcal {T} _ {1} ^ {\alpha \beta} \mathcal {T} _ {2} ^ {\delta \omega} \right] = 0.
$$
Proof. It is easy to see that  $\mathbb{E}_{\ell}\mathcal{T}_1^{\alpha \beta}\mathcal{T}_2^{\delta \omega}$  involve odd standard Gaussian moments (in particular third order moments), which vanish due to the parity of the centered Gaussian measure.
Proof. We know that:
Using  $\lambda^2 +\gamma^2 = 1$  , and summing and subtracting the mean of  $\mathcal{T}_2$  , we have that:
$$
V _ {\ell + 1} ^ {\alpha \beta} = V _ {\ell} ^ {\alpha \beta} - \gamma^ {2} V _ {\ell} ^ {\alpha \beta} + \frac {\gamma^ {2}}{\sqrt {n}} \mathbb {E} _ {\ell} \left[ \mathcal {T} _ {2} ^ {\alpha \beta} \right] + \frac {\gamma \lambda}{\sqrt {n}} \mathcal {T} _ {1} ^ {\alpha \beta} + \frac {\gamma^ {2}}{\sqrt {n}} \left(\mathcal {T} _ {2} ^ {\alpha \beta} - \mathbb {E} _ {\ell} \left[ \mathcal {T} _ {2} ^ {\alpha \beta} \right]\right).
$$
Using Lemma B.2 for the mean of  $\mathcal{T}_2$ , we have that:
$$
V _ {\ell + 1} ^ {\alpha \beta} = V _ {\ell} ^ {\alpha \beta} + \frac {1}{n} \sqrt {V ^ {\alpha \alpha} V ^ {\beta \beta}} \nu (\rho) + \frac {\gamma \lambda}{\sqrt {n}} \mathcal {T} _ {1} ^ {\alpha \beta} + \frac {\gamma^ {2}}{\sqrt {n}} \left(\mathcal {T} _ {2} ^ {\alpha \beta} - \mathbb {E} _ {\ell} \left[ \mathcal {T} _ {2} ^ {\alpha \beta} \right]\right) + \mathcal {O} (n ^ {- 3 / 2}),
$$
which gives us an expression in the right Markov Chain form with drift term  $\sqrt{V^{\alpha\alpha}V^{\beta\beta}}\nu(\rho)$ . For the diffusion term, we need to compute the covariance between two different neural covariances  $V_{\ell+1}^{\alpha\beta}$ ,  $V_{\ell+1}^{\delta\omega}$ . Using Lemmas B.1 to B.3, we have
$$
\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} \left[ \mathcal {T} _ {2} ^ {\alpha \beta} \right]\right) \left(\lambda \gamma \mathcal {T} _ {1} ^ {\delta \omega} + \gamma^ {2} \mathcal {T} _ {2} ^ {\delta \omega} - \gamma^ {2} \mathbb {E} _ {\ell} \left[ \mathcal {T} _ {2} ^ {\delta \omega} \right]\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] \\ = 2 \lambda^ {2} \gamma^ {2} \left(V ^ {\alpha \delta} V ^ {\beta \omega} + V ^ {\alpha \omega} V ^ {\beta \delta}\right) + 2 \gamma^ {4} \left(V ^ {\alpha \delta} V ^ {\beta \omega} + V ^ {\alpha \omega} V ^ {\beta \delta}\right) + \mathcal {O} (n ^ {- 1}) \\ = 2 \gamma^ {2} (V ^ {\alpha \delta} V ^ {\beta \omega} + V ^ {\alpha \omega} V ^ {\beta \delta}) + \mathcal {O} (n ^ {- 1}), \\ \end{array}
$$
where we use  $\operatorname{Cov}_{\ell}$  to denote the covariance conditioned on the sigma-algebra  $\mathcal{F}_{\ell} = \sigma(\{X_k\}_{k \in [\ell]})$ . Finally, applying Proposition A.2, we get the desired result.
C SDE for Softmax-based Attention
C.1 Dot-Product Attention
Dot-product attention applies the Softmax row-wise to the following matrix:
$$
Y _ {\ell} = \frac {1}{n} X _ {\ell} \underbrace {W _ {\ell} ^ {K} W _ {\ell} ^ {Q , \top}} _ {W _ {\ell} ^ {B}} X _ {\ell} ^ {\top},
$$
where  $W_{\ell}^{K}, W_{\ell}^{Q} \in \mathbb{R}^{n \times n_k}$  are Gaussian matrices with unit variance entries, and  $n_k$  is the queries and keys' dimension. Here we study the first two moments of  $Y_{\ell}$ . In particular, note that  $\mathbb{E}_{\ell}[Y_{\ell}] = 0$  where  $\mathbb{E}_{\ell}[\cdot] = \mathbb{E}[\cdot|\mathcal{F}_{\ell}]$  denotes the conditional expectation given the sigma-algebra generated by  $\mathcal{F}_{\ell} = \sigma(\{X_k\}_{k \in [\ell]})$ .
For the second moment we have the following Lemma.
Lemma C.1. Let
$$
Y _ {\ell} = \frac {1}{n} X _ {\ell} W _ {\ell} ^ {K} W _ {\ell} ^ {Q, \top} X _ {\ell} ^ {\top},
$$
be the dot-product attention parametrized by  $W_{\ell}^{K}, W_{\ell}^{Q} \in \mathbb{R}^{n \times n_k}$ . Then:
$$
\mathbb {E} _ {\ell} \left[ Y ^ {\alpha \beta} Y ^ {\delta \omega} \right] = n _ {k} V ^ {\alpha \delta} V ^ {\beta \omega}. \tag {24}
$$
$$
\mathbb {E} _ {\ell} [ Y ^ {\alpha \beta} Y ^ {\delta \omega} ] = \frac {1}{n ^ {2}} \sum_ {k k ^ {\prime} j j ^ {\prime}} X _ {k} ^ {\alpha} X _ {k ^ {\prime}} ^ {\beta} X _ {j} ^ {\delta} X _ {j ^ {\prime}} ^ {\omega} \mathbb {E} [ W _ {k k ^ {\prime}} ^ {B} W _ {j j ^ {\prime}} ^ {B} ].
$$
For the expectation, we have that: