SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Euler-discretization
$$
V _ {\ell + 1} = V _ {\ell} + \frac {b (V _ {\ell})}{n} + \frac {\Sigma (V _ {\ell}) ^ {1 / 2} \xi_ {\ell}}{\sqrt {n}} + O \left(n ^ {- 3 / 2}\right), \tag {7}
$$
where  $b, \Sigma$  are deterministic coefficients, and  $\xi_{\ell}$  are random vectors with zero mean and identity covariance. From here onwards, we use  $O(n^{-p})$  to denote a random variable  $Z$  such that  $n^p Z$  has all moments bounded by universal constants (i.e. independent of  $n$ ). Since the update can be interpreted as discretization with step size  $n^{-1}$ , naturally the Markov chain converges to an SDE. We again note that a stable SDE implies a stable covariance structure for finite size networks.
To achieve the same goal for modified attention mechanisms, we consider a similar approach as M. Li et al. [18] for smooth activation functions, and Taylor expand the Softmax function in terms of a large temperature parameter  $\tau$ . To this end, let  $Y_{\ell}$  to be the matrix of dot-products between queries, and keys, i.e.  $Y_{\ell} := X_{\ell} \frac{1}{\sqrt{n}} W_{\ell}^{Q} \frac{1}{\sqrt{n}} W_{\ell}^{K,\top} X_{\ell}^{\top}$ .
More specifically, given a row  $y^\alpha \in \mathbb{R}^{1 \times m}$  of the logits  $Y_\ell \in \mathbb{R}^{m \times m}$ , we can Taylor expand the row-wise Softmax in terms of  $\tau^{-1}$ :
$$
\operatorname {S o f t m a x} \left(\tau^ {- 1} y ^ {\alpha}\right) = \frac {1}{m} \mathbf {1} ^ {\top} + \frac {1}{\tau m} \left(y ^ {\alpha} - \bar {y} ^ {\alpha}\right) + \frac {1}{2 \tau^ {2} m} \left[ \left(y ^ {\alpha} - \bar {y} ^ {\alpha}\right) ^ {2} - \left(\bar {y} ^ {\alpha} ^ {2} - \overline {{(y ^ {\alpha}) ^ {2}}}\right) \right] + O \left(\tau^ {- 3}\right), \tag {8}
$$
where  $\overline{y^{\alpha}}\coloneqq \frac{1}{m}\sum_{\beta}y^{\alpha \beta}\mathbf{1}^{\top}$  and  $(y^{\alpha})^{2}$  is the vector with squared entries of  $y^\alpha$ , and  $\mathbf{1}\in \mathbb{R}^{m\times 1}$  is the (column) vector of ones. We note in practice  $\tau$  is often set to  $\sqrt{n_k}$ , which is often quite large and allows for asymptotic analysis [9].
We observe that the zero-th order term  $m^{-1}\mathbf{1}^{\top}$  is independent of  $\tau$ . Considering the form of the attention block as  $A_{\ell}X_{\ell}\frac{1}{\sqrt{n}} W_{\ell}^{V}$ , this yields an update that is no longer a small perturbation of  $V_{\ell}$ , regardless of how  $\tau$  is chosen. Therefore, to form a Markov chain like Eq. 7, we actually require  $A_{\ell}$  to be approximately the identity matrix.
4.2 Shaped Attention
To shape the Softmax-attention mechanism as a perturbation of the identity matrix, we propose the following modifications which we call the shaped attention
$$
A _ {\ell} = I + \operatorname {S o f t m a x} \left(\tau^ {- 1} Y _ {\ell}\right) - \frac {1}{m} \mathbf {1 1} ^ {\top}, \quad \tau = \tau_ {0} \sqrt {n n _ {k}}. \tag {9}
$$
The shaped attention presents three modifications to the Softmax attention in Eq. 2. Firstly, the zero-order term  $m^{-1} \mathbf{11}^{\top}$  of the Taylor expansion (Eq. 8) is removed as it causes a non-infiniteal drift in the Markov Chain that ultimately leads to instability in the covariance (see Section 4.1). Secondly, we also observe that when  $\tau$  is very large, the centered Softmax is a perturbation around zero. To recover an approximate Euler-update like in Eq. 7, we simply add back the identity matrix. By biasing the attention matrix towards the identity, we encourage each token to self-attend. This type of modification was also previously considered by B. He et al. [46]. Finally, the Softmax's temperature is chosen to scale as  $\tau = \tau_0 \sqrt{nn_k}$ , for some constant  $\tau_0 > 0$ , which guarantees a non-degenerate limit as  $(d, n) \to \infty$  (Theorem 4.2). Note that the extra  $\sqrt{n}$  term is a departure from the standard parameterization.
In Figure 4, we show how removing any of the proposed changes individually alters the neural covariance structure, which becomes degenerate for large depths, while the proposed modifications remain stable. We stress that here for simplicity we focus on attention without masking. Shaped attention can be extended to include masking (e.g. casual masking) by centering each i-th row of the Softmax matrix by a different factor  $1 / m_{i}$ , where  $m_{i}$  is the number of un-masked tokens in the i-th row.
In principle, it could be possible to have a close-to-identity Softmax matrix when the logits are large. However, this regime also corresponds to a very saturated Softmax, thus making training unstable [57]. As a result, we will avoid this direction in this work.
Figure 4: Mean correlation (left) and covariance (right) (in absolute value) under various interventions on the proposed shaped attention. In particular, we remove either one or two of the three modifications from the shaped attention in Eq. 9. For instance  $\tau^2 = nn_k$ , center" indicates that we use the proposed temperature, and we center by  $m^{-1}$ , but we do not add the identity matrix, while in "only id" we add the identity matrix but use  $\tau = \sqrt{n_k}$  and do not center. We note in this "only id" case, the covariance remains unstable due to incorrect scaling. Due to exploding covariance, we choose to not include the cases "id,  $\tau^2 = nn_k$ " and "only id" in the correlation plot (but only in the covariance plot). Simulated with  $n = 300$ ,  $d = 150$ ,  $\rho_0^{\alpha \beta} = 0.2$ ,  $\gamma = 1 / \sqrt{2}$  and  $2^{13}$  samples.
4.3 Main Result – Neural Covariance SDEs for Shaped Attention Models and Shaped Transformers
Before we state our main results, we will first define a weakened notion of convergence, which is required whenever the drift and covariance coefficients are not Lipschitz. This was also required for the case of shaped MLPs with smooth activations [18].
Definition 4.1 (Local Convergence). We say the covariance  $V^{(n)}$  converges locally to  $V$  if the stopped process  $\{V_{t\wedge T_r}^{(n)}\}_{t\geq 0}$  converges to  $\{V_{t\wedge T_r}\}_{t\geq 0}$  in the sense of Definition 3.1 for all stopping times of the form  $T_{r} = \inf \{t > 0:\| V_{t}\| \geq r\}$  with  $r > 0$ .
Let the covariance with respect to the average token be defined as  $V^{\alpha \bar{x}} \coloneqq m^{-1} \sum_{\nu = 1}^{m} V^{\alpha \nu}$ , and the average trace be  $\bar{V} \coloneqq m^{-1} \sum_{\nu = 1}^{m} V^{\nu \nu}$ . We will need to compute a couple of important moments from the Taylor expansion terms of the Softmax (Lemma C.2)
$$
\begin{array}{l} S _ {1} ^ {\alpha \delta , \beta \omega} := n _ {k} ^ {- 1} \mathbb {E} (Y ^ {\alpha \delta} - \overline {{y ^ {\alpha}}}) (Y ^ {\beta \omega} - \overline {{y ^ {\beta}}}) = 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} := n _ {k} ^ {- 1} \mathbb {E} \left[ (Y ^ {\alpha \delta} - \bar {y} ^ {\alpha}) ^ {2} - \left(\overline {{(Y ^ {\alpha}) ^ {2}}} - \bar {y} ^ {\alpha^ {2}}\right) \right] = V ^ {\alpha \alpha} \left(V ^ {\delta \delta} - 2 V ^ {\delta \bar {x}} + 2 V ^ {\bar {x} \bar {x}} - \bar {V}\right). \tag {10} \\ \end{array}
$$
We are now ready to state our main result.
Theorem 4.2. Let  $X_{\ell}$  be the hidden layers of a residual attention network defined in Eq. 1 with shaped attention in Eq. 9, parameters  $\lambda^2 + \gamma^2 = 1$  and  $\tau = \tau_0 \sqrt{nn_k}$ , where  $\lambda, \gamma, \tau_0$  all do not depend on  $d, n$ . Then the feature covariance  $V_{\ell}$  converges locally to the solution of the following SDE (in the sense of Definition 4.1)
$$
d V _ {t} = b (V _ {t}) d t + \Sigma (V _ {t}) ^ {1 / 2} d B _ {t}, V _ {0} = \frac {1}{n} X _ {0} X _ {0} ^ {\top},
$$
where the drift has the following form
$$
b (V) = \frac {\gamma^ {2}}{\tau_ {0} ^ {2}} \left[ \frac {1}{m ^ {2}} \sum_ {\nu , \kappa = 1} ^ {m} V ^ {\nu \kappa} S _ {1} ^ {\alpha \nu , \beta \kappa} + \frac {1}{2 m} \sum_ {\nu = 1} ^ {m} (V ^ {\beta \nu} S _ {2} ^ {\alpha \nu} + V ^ {\alpha \nu} S _ {2} ^ {\beta \nu}) \right] _ {\alpha \leq \beta},
$$
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
$$
\mathcal {A} ^ {\alpha \beta , \delta \omega} := \frac {1}{m ^ {2}} \sum_ {\nu , \kappa = 1} ^ {m} \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).
$$
The drift depends on the shaped attention mechanism through  $S_{1}^{\alpha \delta, \beta \omega}$  and  $S_{2}^{\alpha \delta}$ , the moments of the first and second order terms of the Softmax's Taylor expansion. On the other hand, the diffusion term depends on the attention solely through  $S_{1}$ , present in the additional term  $A^{\alpha \beta, \delta \omega}$ . The presence of  $A^{\alpha \beta, \delta \omega}$  is an intriguing difference compared to shaped ReLU networks, where the diffusion is not affected by the activation function. Both components of the SDE depend on averages over the tokens, reflecting the mixing property of the self-attention mechanism, in which every pair of tokens is compared through dot products to form the attention weights. Finally, notice how the residual branch parameter  $\gamma^{2}$  has a dampening effect on the scale of both the drift and the diffusion in a similar way as in fully-connected residual network.
We are now ready to introduce the full shaped Transformer architecture, where we combine the attention and residual layers:
$$
Z _ {\ell} = \lambda X _ {\ell} + \gamma A _ {\ell} X _ {\ell} \frac {1}{\sqrt {n}} W _ {\ell} ^ {V}, \quad X _ {\ell + 1} = \lambda Z _ {\ell} + \gamma \sigma_ {s} \left(Z _ {\ell} \frac {1}{\sqrt {n}} W _ {\ell} ^ {\text {p r e}}\right) \sqrt {\frac {c}{n}} W _ {\ell} ^ {\text {p o s t}}, \tag {11}
$$
where  $A_{\ell}$  is the shaped attention defined by Eq. 9. We note that covariance SDE handle stacking of different layer types very conveniently by simply adding the drift and covariance of the diffusion coefficients, which we summarize in the Corollary below.
Corollary 4.3 (Shaped Transformer Covariance SDE). Let  $X_{\ell}$  be the hidden layers of a shaped transformer defined in Eq. 11 with parameters  $\lambda^2 + \gamma^2 = 1$  and  $\tau = \tau_0 \sqrt{nn_k}$ , where  $\lambda, \gamma, \tau_0$  all do not depend on  $d, n$ . Then the feature covariance  $V_{\ell}$  converges locally to the solution of the following SDE (in the sense of Definition 4.1)
$$
d V _ {t} = \left[ b \left(V _ {t}\right) + b _ {r e s} \left(V _ {t}\right) \right] d t + \left[ \Sigma \left(V _ {t}\right) + \Sigma_ {r e s} \left(V _ {t}\right) \right] ^ {1 / 2} d B _ {t}, \tag {12}
$$
where the coefficients are defined in Theorem 3.2 and Theorem 4.2.
4.4 On Finite Time Stability of the SDE and Shaped Attention Networks
Although we did not observe numerical instability in majority of our simulations of the shaped attention networks and the corresponding SDE, we did observe that the drift component  $b(V_{t})$  in Theorem 4.2 is cubic in the entries of  $V_{t}$ . Whenever the drift is not Lipschitz as in this case, we do not have general guarantees for the existence of a solution for all time (see the Feller test for explosions [58, Theorem 5.5.29]). In fact, MLPs with smooth activations also yield non-Lipschitz drift coefficients as seen in M. Li et al. [18].
However, locally Lipschitz coefficients are sufficient to guarantee the existence of local solutions, in the sense of up to a stopping time [59, Proposition 6.9]. Not only does this fact help us establish a precise notion of convergence (Definition 4.1), we can also study the practical implications of this for finite sized attention networks. More specifically, we can inspect the effect of architectural changes to a stopping time.
To demonstrate the potential numerical instabilities, we had to choose an adversarial set of parameters: in particular, an unrealistically large norm (approx.  $10\sqrt{n}$ ) for the initial tokens  $X_0$ , which enlarges the eigenvalues of  $V_0$  to the order of 100. Given these initial conditions and a large residual connection weight  $\gamma$ , we were able to consistently generate numerically unstable behaviour in shaped attention networks (see Figure 5 (left)).
That being said, it is very straight forward to stabilize the network by tweaking parameters such as  $\gamma, \tau_0$  and the depth-to-width ratio of the network. We demonstrate the effect of tuning  $\gamma$  on both sample trajectories of the maximum eigenvalue of  $V_{\ell}$  and the stopping time in Figure 5. As we may intuitively expect, tuning  $\gamma$  smaller will delay the time scale of numerical instabilities, hence allowing for larger depth networks to remain stable.
Architecture Design and Hyperparameter Tuning. Previous work have demonstrated the practical impact scaling limits can have on designing activation functions [16, 17] and tuning hyperparameters [19]. We follow this line of motivations and proposed a novel attention mechanism, which successfully stabilizes the covariance structure in arbitrarily deep Transformers (e.g. Figure 1). The natural next
Figure 5: Left: Trajectories of the maximum eigenvalue of the covariance matrix in a shaped attention network, with adversarially large initial condition. Right: Stopping time of the shaped attention neural network, capped at 1. Stopping time is defined as  $t^* = d^* / n$  with  $d^*$  the maximum depth beyond which one of the eigenvalues of the covariance matrix exceeds  $10^4$  or drops below  $10^{-4}$ . Simulated with  $n = d = 200$ ,  $\tau_0 = 1$ , and 100 samples used for median and 10th percentile.
step is to investigate the scaling of gradients in the infinite-depth-and-width limit. As Yang et al. [19] illustrated, the existence of an infinite-width limit for the gradient implies the optimal hyperparameters for the training algorithm will also converge. This type of results allows for tuning of hyperparameters on networks with a much smaller width, yet extends easily to arbitrarily large networks that approximates the same limit, saving massive amounts of computing cost in the process. Given the existence of an infinite-depth-and-width limit for the forward pass, we believe it's possible to extract optimal hyperparameters from networks with not only a much smaller width, but smaller depth as well.
Preliminary Experiments. Although this work is primarily theoretical, it is important to consider whether or not the proposed architecture is useful in practice. Given limited computing resources, we chose to only briefly test the feasibility of training the shaped Transformer. Nevertheless, our preliminary experiments show promising results when it comes to training stability. In particular, the shaped Transformer (without LayerNorm) does indeed train at approximately the same speed as well tuned Transformer architectures. Full details of the experiment and results can be found in Appendix D. A more comprehensive set of experiments with different tasks, datasets, and larger networks will be required to confidently determine the practical feasibility of the shaped Transformer, which we defer to future work.
Training Dynamics and Generalization. As mentioned in the introduction, the limitations of infinite-width NTK theories motivate our study of the proportional infinite-depth-and-width limit. In particular, to address many of the open problems in deep learning theory, we need a faithful and tractable description of training dynamics. Given the results at initialization, the proportional limit holds the potential for such a theory of training as well. Another promising indicator is that deep networks learn features in the proportional regime [38], which has been identified as a key advantage of neural networks over kernel methods [25, 60-66]. A precise theory of training will help us understand other types of instabilities during training and improve existing optimization methods. Furthermore, determining the network which training converges to is a necessary step towards a theory of generalization, as demonstrated by the infinite-width approach [67]. In light of our results, we believe that our theory sets the stage for future work on training and generalization in deep learning.
CL and ML would like to thank Keiran Paster for insightful discussions. LN would like to thank Sotiris Anagnostidis for support in pre-processing the dataset used for the training experiments of this manuscript. ML is supported by the Ontario Graduate Scholarship and Vector Institute. DMR is supported in part by Canada CIFAR AI Chair funding through the Vector Institute, an NSERC Discovery Grant, Ontario Early Researcher Award, a stipend provided by the Charles Simonyi Endowment, and a New Frontiers in Research Exploration Grant.
[67] P. L. Bartlett, A. Montanari, and A. Rakhlin. “Deep learning: a statistical viewpoint”. In: Acta numerica 30 (2021), pp. 87–201.
[68] O. Kallenberg. Foundations of Modern Probability. Probability theory and stochastic modelling. Springer, 2021. ISBN: 9783030618728.
[69] HuggingFace. Wikipedia data set. https://huggingface.co/datasets/wikipedia.
[70] HuggingFace. Bookcorpus data set. https://huggingface.co/datasets/bookcorpus.
[71] D. P. Kingma and J. Ba. “Adam: A method for stochastic optimization”. In: arXiv preprint arXiv:1412.6980 (2014).
[72] A. Wang, A. Singh, J. Michael, F. Hill, O. Levy, and S. R. Bowman. "GLUE: A multi-task benchmark and analysis platform for natural language understanding". In: arXiv preprint arXiv:1804.07461 (2018).
[73] HuggingFace. Hugging face Bert implementation. https://github.com/huggingface/transformers/blob/main/src/transformers/models/bert/modeling_bert.py.
A Preliminaries: Covariance SDE Framework
In this section, we will review existing results on Markov chain convergence to an SDE, as well as existing results for the covariance SDE. See also the Appendix of M. Li et al. [18] for more details.
Firstly, we will define the Skorohod  $J_{1}$ -topology, or just the Skorohod topology for short [68, Appendix 5]. The topology is used to describe convergence of continuous time processes with discontinuities, in particular, Markov chains with a continuous time interpolation fits in this category. Let  $S$  be a complete separable metric space, and  $D_{\mathbb{R}_{+}, S}$  be the space of càdlàg functions (right continuous with left limits) from  $\mathbb{R}_{+} \to S$ . We use  $x_{n} \xrightarrow{ul} x$  to denote locally uniform convergence (uniform on compact subsets of  $\mathbb{R}_{+}$ ), and consider the class of bijections  $\lambda$  on  $\mathbb{R}_{+}$  so that  $\lambda$  is strictly increasing with  $\lambda_{0} = 0$  (can be interpreted as a time change).
Definition A.1. We define Skorohod convergence  $x_{n}\stackrel {s}{\to}x$  on  $D_{\mathbb{R}_{+},S}$  if there exists a sequence of bijections  $\lambda_{n}$  satisfying above conditions and
$$
\lambda_ {n} \xrightarrow {u l} I d, \quad x _ {n} \circ \lambda_ {n} \xrightarrow {u l} x. \tag {13}
$$
In particular, we call the topology that induces this convergence the Skorohod topology [68, Theorem A5.3].
On a heuristic level, if we have sequences of Markov chains  $Y^n$  that satisfy the following type of Euler updates
$$
Y _ {\ell + 1} ^ {n} = Y _ {\ell} ^ {n} + \frac {b \left(Y _ {\ell} ^ {n}\right)}{n} + \frac {\sigma \left(Y _ {\ell} ^ {n}\right) \xi_ {\ell}}{\sqrt {n}}, \tag {14}
$$
where  $\xi_{\ell}$  are iid random variables with zero mean and identity covariance, then we can interpolate this process in continuous time with  $X_{t}^{n} = Y_{[tn]}^{n}$  and show that as  $n\to \infty$  , we have that  $X^n$  converges to the solution of the following SDE (weakly with respect to the Skorohod topology)
$$
d X _ {t} = b \left(X _ {t}\right) d t + \sigma \left(X _ {t}\right) d B _ {t}, \quad X _ {0} = \lim  _ {n \rightarrow \infty} Y _ {0} ^ {n}. \tag {15}
$$
Our next theorem essentially weakens this result in several ways. Firstly, we don't need to take the step size  $n^{-1}$ , but instead replace it with  $n^{-2p}$  for all  $p > 0$ . Next, we can allow the update to contain higher order terms that vanish, in particular, any terms of order  $O(n^{-2p - \delta})$  for all  $\delta > 0$ . Here, we remind the readers that  $O(n^{-p})$  denotes a random variable  $Z$  such that  $n^p Z$  has all moment bounded by a universal constant (independent of  $n$ ). Thirdly, we can allow  $b(y)$  to be random, or more precisely replace it with  $\widehat{b}(y, \omega_n)$  such that  $\mathbb{E}\widehat{b}(y, \omega_n) = b(y)$ . Fourthly, we can also allow  $b, \sigma$  to be a sequence  $b_n, \sigma_n$  such that they converge to  $b, \sigma$  on uniformly on compact sets. Finally, we can also weaken the topology of convergence to locally, that is all processes are stopped by a stopping time as in Definition 4.1.
Now we will state the main technical result in this section.
Proposition A.2 (Convergence of Markov Chains to SDE, Proposition A.6, [18]). Let  $Y^n$  be a discrete time Markov chain on  $\mathbb{R}^N$  defined by the following update for  $p, \delta > 0$
$$
Y _ {\ell + 1} ^ {n} = Y _ {\ell} ^ {n} + \frac {\widehat {b} _ {n} \left(Y _ {\ell} ^ {n} , \omega_ {\ell} ^ {n}\right)}{n ^ {2 p}} + \frac {\sigma_ {n} \left(Y _ {\ell} ^ {n}\right)}{n ^ {p}} \xi_ {\ell} ^ {n} + O \left(n ^ {- 2 p - \delta}\right), \tag {16}
$$
where  $\xi_{\ell}^{n}\in \mathbb{R}^{N}$  are iid random variables with zero mean, identity covariance, and moments uniformly bounded in  $n$ . Furthermore,  $\omega_{\ell}^{n}$  are also iid random variables such that  $\mathbb{E}[\widehat{b}_n(Y_\ell^n,\omega_\ell^n)|Y_\ell^n = y] = b_n(y)$  and  $\widehat{b}_n(y,\omega_\ell^n)$  has uniformly bounded moments in  $n$ . Finally,  $\sigma_{n}$  is a deterministic function, and the remainder terms in  $O(n^{-2p - \delta})$  have uniformly bounded moments in  $n$ .
Suppose  $b_n, \sigma_n$  are uniformly Lipschitz functions in  $n$  and converges to  $b, \sigma$  uniformly on compact sets, then in the limit as  $n \to \infty$ , the process  $X_t^n = Y_{[tn^{2p}]}^n$  converges in distribution to the solution of the following SDE in the Skorohod topology of  $D_{\mathbb{R}_+,\mathbb{R}^N}$
$$
d X _ {t} = b \left(X _ {t}\right) d t + \sigma \left(X _ {t}\right) d B _ {t}, \quad X _ {0} = \lim  _ {n \rightarrow \infty} Y _ {0} ^ {n}. \tag {17}
$$
Suppose otherwise  $b_n, \sigma_n$  are only locally Lipschitz (but still uniform in  $n$ ), then  $X^n$  converges locally to  $X$  in the same topology (see Definition 4.1). More precisely, for any fixed  $r > 0$ , we consider the stopping times
$$
\tau^ {n} := \inf  \left\{t \geq 0: \left| X _ {t} ^ {n} \right| \geq r \right\}, \quad \tau := \inf  \left\{t \geq 0: \left| X _ {t} \right| \geq r \right\}, \tag {18}
$$
then the stopped process  $X_{t\wedge \tau^n}^n$  converges in distribution to the stopped solution  $X_{t\wedge \tau}$  of the above SDE in the same topology.
We will briefly recall the main result of M. Li et al. [18] next. As mentioned earlier in the text, the setting is for MLPs defined as follows
$$
X _ {\ell + 1} = \sigma_ {s} \left(X _ {\ell}\right) \sqrt {\frac {c}{n}} W _ {\ell}, \tag {19}
$$
where  $\sigma_s(x) = s_+ \max(x, 0) + s_- \min(x, 0)$  with  $s_\pm = 1 + \frac{c_\pm}{\sqrt{n}}$ ,  $c^{-1} = \mathbb{E}\sigma_s(g)^2$  for  $g \sim \mathcal{N}(0, 1)$ ,  $n$  is the width of the network, and  $W_{\ell, ij} \sim \mathcal{N}(0, 1)$  are iid random weights.
Then it was shown that in the limit as  $n, d \to \infty$  with  $\frac{d}{n} = T > 0$ , the covariance matrix  $V_{\ell} = \frac{1}{n} X_{\ell}X_{\ell}^{\top}$  (or equivalent the post activation  $\frac{c}{n}\sigma_s(X_\ell)\sigma_s(X_\ell)^{\top}$  since the activation becomes infinitesimal) converges the solution of the following SDE in the Skorohod topology [18, Theorem 3.2]