SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The proof is given in Appendix D. The first term on the right-hand side of Eq. (3.11) corresponds to the burn-in time for the network to reach a region around where the Polyak-Lojasiewicz condition holds. We divide our analysis of this burn-in phase into two phases, Phase 1 and Phase 2, and we obtain tight control on the factor by which the signal component,  $q_{\star}^{\top}u$ , grows during Phase 1, while Phase 2 takes place for a comparatively short period of time. (This tight control is critical for our sample complexity bounds where we must show that the coupling error does not blow up too much — see more discussion below Lemma 5.4.) Phases 1 and 2 are mostly governed by the quadratic components in the activation and target functions, and the dynamics behave similarly to a power method update. After the burn-in phase, the dynamics operate for a short period of time (Phase 3) in a regime where the Polyak-Lojasiewicz condition holds. We explicitly prove the dynamics stay away from saddle points during this phase, as further discussed in Section 4.2.
We note that Theorem 3.3 provides a concrete polynomial runtime bound for projected gradient flow which is not achievable by prior mean-field analyses [21, 56, 2] using Wasserstein gradient flow techniques. For instance, while the population dynamics of Abbe et al. [2] only trains the hidden layer for  $O(1)$  time, our projected gradient flow updates the hidden layer for  $O(\log d)$  time. The main challenge in the proof is to deal with the saddle points that are not strict-saddle [32] in the loss landscape which cannot be escaped simply by adding noise in the parameter space [45, 49, 24].<sup>3</sup> Our analysis develops a fine-grained analysis of the dynamics that shows the iterates stay away from saddle points, which allows us to obtain the running time bound in Theorem 3.3 which can be translated to a polynomial-width guarantee in Theorem 3.4. In contrast, Wei et al. [76] escape the saddles by randomly replacing an exponentially small fraction of neurons, which makes the network require exponential width.
Next, we state the main theorem on projected gradient flow on empirical, finite-width dynamics.
Theorem 3.4 (Empirical, finite-width dynamics). Suppose Assumption 3.2 holds. Suppose  $\epsilon = \frac{1}{\log\log d}$  is the target error and  $T_{*,\epsilon}$  is the running time defined in Theorem 3.3. Suppose  $n \geq d^{\mu}(\log d)^{\Omega(1)}$  for any constant  $\mu > 3$ , and the network width  $m$  satisfies  $m \gtrsim d^{2.5 + \mu/2}(\log d)^{\Omega(1)}$  and  $m \leq d^{C}$  for some sufficiently large universal constant  $C$ . Let  $\hat{\rho}_t$  be the projected gradient flow on the empirical loss, initialized with  $m$  uniformly sampled weights, defined in Eq. (3.10). Then, with probability at least  $1 - \frac{1}{d^{\Omega(\log d)}}$  over the randomness of the data and the initialization of the finite-width network, we have that  $L(\hat{\rho}_{T_{*,\epsilon}}) \lesssim (\hat{\sigma}_{2,d}^2 + \hat{\sigma}_{4,d}^2)\epsilon^2$ .
Plugging in  $\mu = 3.1$ , Theorem 3.4 suggests that, when the network width is at least  $d^{4.05}$  (up to logarithmic factors), the empirical dynamics of gradient descent could achieve  $\hat{\sigma}_{\max}^2(1 / \log \log d)^2$  population loss with  $d^{3.1}$  samples (up to logarithmic factors).
In our analysis, we will establish a coupling between neurons in the empirical dynamics and neurons in the population dynamics. The main challenge is to bound the coupling error during Phase 1 where the population dynamics are similar to the power method. During this phase, we show that the coupling error (i.e., the distance between coupled neurons) remains small by showing that it can be upper bounded in terms of the growth of the signal  $q_{\star}^{T}u$  in the population dynamics. Such a delicate relationship between the growth of the error and that of the signal is the main challenge to proving a sample complexity bound better than NTK. Even an additional constant factor in the growth rate of the coupling error would lead to an additional poly(d) factor in the sample complexity. Prior work (e.g. Mei et al. [57], Abbe et al. [2]) also establishes a coupling between the population and empirical dynamics. However, these works obtain bounds on the coupling error which are exponential in the running time — specifically, their bounds have a factor of  $e^{KT}$  where  $K$  is a universal constant and  $T$  is the time. In many settings, including ours,  $\Omega (\log d)$  time is necessary to achieve a non-trivial error, (as we further discuss in Section 4.2) and thus this would lead to a  $d^{O(1)}$  bound on the sample complexity, where  $O(1)$  is a unspecified and likely loose constant, which cannot outperform NTK. Our work addresses this challenge by comparing the growth of the coupling error with the growth of the signal, which enables us to control the constant factor in the growth rate of the coupling error.
Projected Gradient Descent with  $1 / \mathrm{poly}(d)$  Step Size. As a corollary of Theorem 3.4, we show that projected gradient descent with a small learning rate can also achieve low population loss in  $\mathrm{poly}(d)$  iterations. We first define the dynamics of projected gradient descent as follows. As before, we let  $\chi_1,\ldots ,\chi_m\stackrel {\mathrm{i.i.d}}{\sim}\mathbb{S}^{d - 1}$  be the initial weight vectors, and we let  $\tilde{\rho} = \mathrm{unif}(\{\chi_1,\dots ,\chi_m\})$  denote the uniform distribution over these vectors. Thus, we can write a single neuron at initialization as  $\tilde{u}_0(\chi) = \chi$  for  $\chi \in \{\chi_1,\dots ,\chi_m\}$ . The dynamics of projected gradient descent, on a finite-width neural network and using the empirical loss  $\widehat{L}$ , can be then described as
$$
\tilde {u} _ {t + 1} (\chi) = \frac {\tilde {u} _ {t} (\chi) - \eta \cdot \operatorname {g r a d} _ {\tilde {u} _ {t} (\chi)} \widehat {L} \left(\tilde {\rho} _ {t}\right)}{\| \tilde {u} _ {t} (\chi) - \eta \cdot \operatorname {g r a d} _ {\tilde {u} _ {t} (\chi)} \widehat {L} \left(\tilde {\rho} _ {t}\right) \| _ {2}}, \tag {3.12}
$$
$$
\text {a n d} \tilde {\rho} _ {t} = \text {d i s t r i b u t i o n} \tilde {u} _ {t} (\chi) \left(\text {w h e r e} \chi \sim \tilde {\rho} _ {0}\right). \tag {3.13}
$$
where  $\eta > 0$  is the learning rate. We show that projected gradient descent with a  $1 / \mathrm{poly}(d)$  learning rate can achieve a low population loss in  $\mathrm{poly}(d)$  iterations:
3There are a long list of prior works on the loss landscape of neural networks [55, 64, 72, 46, 16, 15, 11, 66, 71, 33, 80, 18].
See Chapter 3, page 20 of Boumal [17], accessed on August 21, 2023.
Theorem 3.5 (Projected Gradient Descent). Suppose we are in the setting of Theorem 3.4. Let  $\tilde{\rho}_t$  be the discrete-time projected gradient descent on the empirical loss, initialized with  $m$  weight vectors sampled uniformly from  $\mathbb{S}^{d - 1}$ , defined in Eq. (3.13). Finally, assume that  $\eta \leq \frac{1}{(\hat{\sigma}_{2,d}^2 + \hat{\sigma}_{4,d}^2)d^B}$  for a sufficiently large universal constant  $B$ . Then,  $L(\tilde{\rho}_{T_{*,\epsilon} / \eta}) \lesssim (\hat{\sigma}_{2,d}^2 + \hat{\sigma}_{4,d}^2)\epsilon^2$ .
Theorem 3.5 follows from Theorem 3.4 together with a standard inductive argument to bound the discretization error. The full proof is in Appendix F.
The following theorem states that in the setting of Theorem 3.4, kernel methods with any inner product kernel require  $\Omega(d^4 (\ln d)^{-6})$  samples to achieve a non-trivial population loss. (Note that the zero function has loss  $\mathbb{E}_{x\sim \mathbb{S}^{d - 1}}[y(x)^2 ]\geq (\hat{h}_{4,d})^2.$ )
Theorem 3.6 (Sample complexity lower bound for kernel methods). Let  $K$  be an inner product kernel. Suppose  $d$  is larger than a universal constant and  $n \lesssim d^4 (\ln d)^{-6}$ . Then, with probability at least  $1/2$  over the randomness of  $n$  i.i.d data points  $\{x_i\}_{i=1}^n$  drawn from  $\mathbb{S}^{d-1}$ , any estimator of  $y$  of the form  $f(x) = \sum_{i=1}^n \beta_i K(x_i, x)$  of  $\{x_i\}_{i=1}^n$  must have a large error:  $\mathbb{E}_{x \sim \mathbb{S}^{d-1}}(y(x) - f(x))^2 \geq \frac{3}{4} (\hat{h}_{4,d})^2$ .
Theorem 3.4 and Theorem 3.6 together prove a clear sample complexity separation between gradient flow and NTK. When  $\hat{h}_{4,d} \asymp \hat{\sigma}_{\max}$  (i.e.,  $\gamma_{2}, \gamma_{4} \asymp 1$ ), gradient flow with finite-width neural networks can achieve  $(\hat{h}_{4,d})^{2} (\log \log d)^{-2}$  population error with  $d^{3.1}$  samples, while kernel methods with any inner product kernel (including NTK) must have an error at least  $(\hat{h}_{4,d})^{2} / 2$  with  $d^{3.9}$  samples.
On a high level, Abbe et al. [2, 3] prove similar lower bounds in a different setting where the target function is drawn randomly and the data points can be arbitrary (also see Kamath et al. [47], Hsu et al. [42], Hsu [41]). In comparison, our lower bound works for a fixed target function by exploiting the randomness of the data points. In fact, we can strengthen Theorem 3.6 by proving a  $\Omega(\hat{h}_{k,d}^2)$  loss lower bound for any universal constant  $k \geq 0$  (Theorem H.2). We defer the proof of Theorem 3.6 to Appendix H.
4 Analysis of Population Dynamics
4.1 Symmetry of the Population Dynamics
A key observation is that due to the symmetry in the data, the population dynamics  $\rho_{t}$  has a symmetric distribution in the subspace orthogonal to the vector  $q_{\star}$ . As a result, the dynamics  $\rho_{t}$  can be precisely characterized by the dynamics in the direction of  $q_{\star}$ .
Recall that  $w = q_{\star}^{\top}u$  denotes the projection of a weight vector  $u$  in the direction of  $q_{\star}$ . Let  $z = (I - q_{\star}q_{\star}^{\top})u$  be the remaining component. For notational convenience, without loss of generality, we can assume that  $q_{\star} = e_1$  and write  $u = (w,z)$ , where  $w \in [-1,1]$  and  $z \in \sqrt{1 - w^2} \cdot \mathbb{S}^{d-2}$ . We will use this convention throughout the rest of the paper. We will use  $w_t(\chi)$  to refer to the first coordinate of the particle  $u_t(\chi)$ , and  $z_t(\chi)$  to refer to the last  $(d-1)$  coordinates.
Definition 4.1 (Rotational invariance and symmetry). We say a neural network  $\rho$  is rotationally invariant if for  $u = (w,z)\sim \rho$ , the distribution of  $z\mid w$  is uniform over  $\sqrt{1 - w^2}\mathbb{S}^{d - 2}$  almost surely. We also say  $\rho$  is symmetric w.r.t a variable  $w$  if the density of  $w$  is an even function.
We note that any polynomial-width neural network (e.g.,  $\hat{\rho}$ ) is very far from rotationally invariant, and therefore the definition is specifically used for population dynamics. If  $\rho$  is rotationally invariant and symmetric, then  $L(\rho)$  has a simpler form that only depends on the marginal distribution of  $w$ .
Lemma 4.2. Let  $\rho$  be a rotationally invariant neural network. Then, for any target function  $h$  and any activation  $\sigma$ ,
$$
L (\rho) = \frac {1}{2} \sum_ {k = 0} ^ {\infty} \left(\hat {\sigma} _ {k, d} \mathbb {E} _ {u \sim \rho} [ P _ {k, d} (w) ] - \hat {h} _ {k, d}\right) ^ {2}. \tag {4.1}
$$
In addition, suppose  $h$  and  $\sigma$  satisfy Assumption 3.2 and  $\rho$  is symmetric. Then
$$
L (\rho) = \frac {\hat {\sigma} _ {2 , d} ^ {2}}{2} \left(\mathbb {E} _ {u \sim \rho} [ P _ {2, d} (w) ] - \gamma_ {2}\right) ^ {2} + \frac {\hat {\sigma} _ {4 , d} ^ {2}}{2} \left(\mathbb {E} _ {u \sim \rho} [ P _ {4, d} (w) ] - \gamma_ {4}\right) ^ {2}. \tag {4.2}
$$
The proof of Lemma 4.2 is deferred to Appendix D.1. Eq. (4.1) says that if  $\rho$  is rotationally invariant, then  $L(\rho)$  only depends on the marginal distribution of  $w$ . Eq. (4.2) says that if the distribution of  $w$  is additionally symmetric and  $\sigma$  and  $h$  are quartic, then the terms corresponding to odd  $k$  and the higher order terms for  $k > 4$  in Eq. (4.1) vanish. Note that  $P_{2,d}(s) \approx s^2$  and  $P_{4,d}(s) \approx s^4$  by Eq. (C.4). Thus,  $L(\rho)$  essentially corresponds to matching the second and fourth moments of  $w$  to some desired values  $\gamma_2$  and  $\gamma_4$ . Inspired by this lemma, we define the following key quantities: for any time  $t \geq 0$ , we define  $D_{2,t} = \mathbb{E}_{u \sim \rho_t}[P_{2,d}(w)] - \gamma_2$  and  $D_{4,t} = \mathbb{E}_{u \sim \rho_t}[P_{4,d}(w)] - \gamma_4$ , where  $\rho_t$  is defined according to the population, infinite-width dynamics (Eq. (3.6), Eq. (3.7) and Eq. (3.8)).
We next show that the rotational invariance and symmetry properties of  $\rho_{t}$  are indeed maintained:
Lemma 4.3. Suppose we are in the setting of Theorem 3.3. At any time  $t \in [0, \infty)$ ,  $\rho_t$  is symmetric and rotationally invariant.
Rotational invariance follows from the rotational invariance of the data. To show symmetry, we use Eq. (4.1), and the facts that  $P_{k,d}$  is an odd polynomial for odd  $k$  and  $\rho_t$  is symmetric at initialization. Using Lemma 4.2 and Lemma 4.3, we obtain a simple formula for the dynamics of  $w_t$ :
Lemma 4.4 (1-dimensional dynamics). Suppose we are in the setting of Theorem 3.3. Then, for any  $\chi \in \mathbb{S}^{d - 1}$ , writing  $w_{t} \coloneqq w_{t}(\chi)$ , we have
$$
\frac {d w _ {t}}{d t} = \underbrace {- \left(1 - w _ {t} ^ {2}\right) \cdot \left(P _ {t} \left(w _ {t}\right) + Q _ {t} \left(w _ {t}\right)\right)} _ {\triangleq v \left(w _ {t}\right)}, \tag {4.3}
$$
where for any  $w \in [-1, 1]$ , we have  $P_{t}(w) = 2\hat{\sigma}_{2,d}^{2}D_{2,t}w + 4\hat{\sigma}_{4,d}^{2}D_{4,t}w^{3}$ , and  $Q_{t}(w) = \lambda_{d}^{(1)}w + \lambda_{d}^{(3)}w^{3}$ , where  $|\lambda_{d}^{(1)}|, |\lambda_{d}^{(3)}| \lesssim \frac{\hat{\sigma}_{2,d}^{2}|D_{2,t}| + \hat{\sigma}_{4,d}^{2}|D_{4,t}|}{d}$ . More specifically,  $\lambda_{d}^{(1)} = 2\hat{\sigma}_{2,d}^{2}D_{2,t} \cdot \frac{1}{d - 1} - 2\hat{\sigma}_{4,d}^{2}D_{4,t} \cdot \frac{6d + 12}{d^{2} - 1}$  and  $\lambda_{d}^{(3)} = 4\hat{\sigma}_{4,d}^{2}D_{4,t} \cdot \frac{6d + 9}{d^{2} - 1}$ .
Eq. (4.3) is a properly defined dynamics for  $w$  because the update rule for  $w_{t}$  only depends on  $w_{t}$  and the quantities  $D_{2,t}$  and  $D_{4,t}$  — additionally,  $D_{2,t}$  and  $D_{4,t}$  only depend on the distribution of  $w$ . We henceforth refer to the  $w_{t}(\chi)$  as particles — this is well-defined by Lemma 4.4.
4.2 Analysis of One-Dimensional Population Dynamics
We also use  $\iota \in [-1, 1]$  to refer to the first coordinate of  $\chi$ , the initialization of a particle under the population dynamics. We note that  $\iota = \langle \chi, e_1 \rangle$  and therefore, the distribution of  $\iota$  is  $\mu_d$ . For any time  $t \geq 0$ , we use  $w_t(\iota)$  to refer to  $w_t(\chi)$  for any  $\chi \in \mathbb{S}^{d-1}$ . This notation is also well-defined by Lemma 4.4. We divide our proof of Theorem 3.3 into three phases, which are defined as follows. For ease of presentation, we will only consider particles  $w_t(\iota)$  for  $\iota > 0$  in the following discussion. Our argument also applies for  $\iota < 0$  by the symmetry of  $\rho_t$  (Lemma 4.3).
Definition 4.5 (Phase 1). Let  $w_{\mathrm{max}} = \frac{1}{\log d}$  and  $\iota_{\mathrm{U}} = \frac{\log d}{\sqrt{d}}$ . Let  $T_{1} > 0$  be the minimum time such that  $w_{T_1}(\iota_{\mathrm{U}}) = w_{\mathrm{max}} := \frac{1}{\log d}$ . We refer to the time interval  $[0, T_1]$  as Phase 1.
Note that essentially all of the particles are less than  $\iota_{\mathrm{U}} = \frac{\log d}{\sqrt{d}}$  at initialization by tail bounds for  $\mu_{d}$ . During Phase 1, the term corresponding to  $D_{2,t}$  in the velocity dominates, and all particles grow by a  $\frac{\sqrt{d}}{(\log d)^2}$  factor. During Phase 1, the loss does not decrease much, but a large portion of the particles have grown by a large factor. During Phase 2, the particles and their velocity will become large.
Definition 4.6 (Phase 2). Let  $T_{2} > 0$  be the minimum time such that either  $D_{2,T_{2}} = 0$  or  $D_{4,T_{2}} = 0$ . We refer to the time interval  $[T_1, T_2]$  as Phase 2. Note that  $T_{2} > T_{1}$  by Lemma D.8.
Definition 4.7 (Phase 3). Phase 3 is defined as the time interval  $[T_2, \infty)$ .
We divide this phase into two cases, which are both more challenging to analyze than Phases 1 and 2: (i)  $D_{2,T_2} = 0$  and  $D_{4,T_2} < 0$ , and (ii)  $D_{2,T_2} < 0$  and  $D_{4,T_2} = 0$ . Here we discuss Case 1 — the analyses of Cases 1 and 2 are in Appendix D.5 and Appendix D.6 respectively. Suppose Case 1 holds, i.e.  $D_{2,T_2} = 0$  and  $D_{4,T_2} < 0$ . Then, for all  $t \geq T_2$ , we will have  $D_{2,t} \geq 0$  and  $D_{4,t} \leq 0$  (Lemma D.19 and Lemma D.20), meaning that for  $t \geq T_2$  a poly  $(\gamma_2)$  fraction of particles will be far from 0 and 1 (Lemma D.17). However, this does not guarantee a large average velocity — unlike
in Phases 1 and 2, the velocity  $v(w) = P_t(w) + Q_t(w)$  defined in Lemma 4.4 may have a positive root  $r \in (0,1)$ , and the root  $r$  can give rise to bad stationary points because for certain values of  $r$ , if  $\rho$  degenerates to the singleton distribution on the root  $r$ , the velocity at  $r$  would be exactly 0. Indeed, we can construct such an  $r$ , intuitively because  $D_{2,t}$  and  $D_{4,t}$  have different signs in  $P_t(w)$  and  $Q_t(w) \approx O(1 / d)$  is a lower order term (Lemma D.35).
Thus, we must leverage some information about the trajectory to show that the average velocity is large when  $D_{2,t}, D_{4,t}$  have different signs. To this end, we prove that  $\mathbb{E}_{u\sim \rho_t}[(w - r)^2] \gtrsim \mathbb{E}_{w,w' \sim \rho_t}(w - w')^2$  is large by designing a novel potential function  $\Phi(w) \coloneqq \log \left(\frac{w}{\sqrt{1 - w^2}}\right)$ . We show that  $|\Phi(w) - \Phi(w')|$  is always increasing for any two particles  $w, w'$  (Lemma D.13). Because  $\Phi$  is Lipschitz on an interval bounded away from 0 and 1, a lower bound on  $|\Phi(w) - \Phi(w')|$  also leads to a lower bound on  $(w - w')^2$ , and thus the particles away from 0 and 1 will have a large variance. Recall that when  $D_{2,t} > 0$  and  $D_{4,t} < 0$ , a large portion of particles are away from 0 and 1, and hence, the average velocity is large. In Appendix B, we include simulations to illustrate the effects of Phases 1, 2 and 3.
5 Analysis of Empirical, Finite-Width Dynamics
Coupling Between Empirical and Population Dynamics. To analyze the difference between the empirical and population dynamics, we define an intermediate process  $\bar{\rho}_t$ , where the initial particles are from  $\hat{\rho}_0$ , but the dynamics of the particles then follow the population trajectories:
$$
\bar {\rho} _ {0} = \hat {\rho} _ {0} = \operatorname {u n i f} \left(\left\{\chi_ {1}, \dots , \chi_ {m} \right\}\right),
$$
$$
\bar {\rho} _ {t} = \text {d i s t r i b u t i o n} u _ {t} (\chi) \left(\text {w h e r e} \chi \sim \bar {\rho} _ {0}\right). \tag {5.1}
$$
For  $\chi \sim \mathrm{unif}(\{\chi_1,\dots ,\chi_m\})$ , let  $\bar{u}_t(\chi) = u_t(\chi)$ . Let  $\Gamma_{t}$  be the joint distribution of  $(\bar{u}_t(\chi),\hat{u}_t(\chi))$ . Then,  $\Gamma_{t}$  forms a natural coupling between  $\bar{\rho}_{t}$  and  $\hat{\rho}_{t}$ . We will use  $\bar{u}_t$  and  $\hat{u}_t$  as shorthands for the random variables  $\bar{u}_t(\chi),\hat{u}_t(\chi)$  respectively in the rest of this section. We define the average distance  $\overline{\Delta}_t^2 \coloneqq \mathbb{E}_{(\hat{u}_t,\bar{u}_t)\sim \Gamma_t}[[\| \hat{u}_t - \bar{u}_t\| _2^2 ]]$ . Intuitively,  $f_{\rho_t}(x)$  and  $f_{\hat{\rho}_t}(x)$  are close when  $\hat{u}_t$  and  $\bar{u}_t$  are close, which is formalized by the following lemma:
Lemma 5.1. In the setting of Theorem 3.4, let  $T \leq \frac{d^{O(1)}}{\hat{\sigma}_{2,d}^2 + \hat{\sigma}_{4,d}^2}$ . Then, with probability at least  $1 - \exp(-d^2)$  over the initialization, we have for all  $t \in [0,T]$  that  $\mathbb{E}_{x \sim \mathbb{S}^{d-1}}[(f_{\hat{\rho}_t}(x) - f_{\bar{\rho}_t}(x))^2] \lesssim \frac{(\hat{\sigma}_{2,d}^2 + \hat{\sigma}_{4,d}^2)d^2(\log d)^{O(1)}}{m}$  and  $\mathbb{E}_{x \sim \mathbb{S}^{d-1}}[(f_{\hat{\rho}_t}(x) - f_{\bar{\rho}_t}(x))^2] \lesssim (\hat{\sigma}_{2,d}^2 + \hat{\sigma}_{4,d}^2)\overline{\Delta}_t^2$ .
The proof of Lemma 5.1 is deferred to Appendix E. As a simple corollary of Lemma 5.1, we can upper bound  $\mathbb{E}_{x\sim \mathbb{S}^{d - 1}}[(f_{\rho_t}(x) - f_{\hat{\rho}_t}(x))^2 ]$  by the triangle inequality. Thus, so long as  $\overline{\Delta}_{T_{*,\epsilon}}$  is small, we can show that the empirical and population dynamics achieve similar test error at time  $T_{*,\epsilon}$ .
Upper Bound for  $\overline{\Delta}_{T_{*,\epsilon}}$ . The following lemma gives our upper bound on  $\overline{\Delta}_{T_{*,\epsilon}}$ :
Lemma 5.2 (Final Bound on  $\overline{\Delta}_{T_{*,\epsilon}}$ ). In the setting of Theorem 3.4, we have  $\overline{\Delta}_{T_{*,\epsilon}} \leq d^{-\frac{\mu - 1}{4}}$ .
We prove Lemma 5.2 by induction over the time  $t \leq T_{*,\epsilon}$ . Let us define the maximal distance  $\Delta_{\max,t} := \max_{(\hat{u}_t, \bar{u}_t) \in \mathrm{supp}(\Gamma_t)} \| \hat{u}_t - \bar{u}_t \|_2$  where the max is over the finite support of the coupling  $\Gamma_t$ . We will control  $\overline{\Delta}_t$  and  $\Delta_{\max,t}$  simultaneously using induction. The inductive hypothesis is:
Assumption 5.3 (Inductive Hypothesis). For some universal constant  $1 > \phi > 1/2$  and  $\psi \in (0, \phi - \frac{1}{2})$ , we say that the inductive hypothesis holds at time  $T$  if, for all  $0 \leq t \leq T$ , we have  $\overline{\Delta}_t \leq d^{-\phi}$  and  $\Delta_{\max,t} \leq d^{-\psi}$ .
Showing that the inductive hypothesis holds for all  $t$  requires studying the growth of the error  $\| \hat{u}_t - \bar{u}_t\|$ . We analyze it using the decomposition  $\frac{d}{dt}\| \hat{u}_t - \bar{u}_t\|_2^2 = A_t + B_t + C_t$ , where  $A_t \coloneqq -2\langle \mathrm{grad}_{\hat{u}}L(\rho_t) - \mathrm{grad}_{\bar{u}}L(\rho_t), \hat{u}_t - \bar{u}_t\rangle$ ,  $B_t \coloneqq -2\langle \mathrm{grad}_{\hat{u}}L(\hat{\rho}_t) - \mathrm{grad}_{\bar{u}}L(\rho_t), \hat{u}_t - \bar{u}_t\rangle$ , and  $C_t \coloneqq -2\langle \mathrm{grad}_{\hat{u}}\widehat{L}(\hat{\rho}_t) - \mathrm{grad}_{\hat{u}}L(\hat{\rho}_t), \hat{u}_t - \bar{u}_t\rangle$ . Intuitively,  $A_t$  captures the growth of the coupling error due to the population dynamics,  $B_t$  the growth due to the discrepancy between the gradients from the
<sup>5</sup>Note that  $\hat{u}_t - \bar{u}_t$  and  $A_t, B_t$  and  $C_t$  implicitly depend on the initialization  $\chi$ , and could be written as  $A_t(\chi), B_t(\chi), C_t(\chi)$ , and  $\delta_t(\chi)$ , but we omit this dependence for ease of presentation.
finite-width and infinite-width networks, and  $C_t$  the growth due to the discrepancy between finite samples and infinite samples. The following lemma gives upper bounds on these terms:
Lemma 5.4 (Bounds on  $A_t, B_t, C_t$ ). In the setting of Theorem 3.4, suppose the inductive hypothesis (Assumption 5.3) holds up to time  $t$ , for some  $t \leq T_{*,\epsilon}$ . Let  $\delta_t \coloneqq \hat{u}_t - \bar{u}_t$ . Let  $T_1$  be the runtime of Phase 1 (where Phase 1 is defined in Definition 4.5). Then, we have
$$
A _ {t} \leq \left\{ \begin{array}{l l} 4 \hat {\sigma} _ {2, d} ^ {2} \left| D _ {2, t} \right| \| \delta_ {t} \| ^ {2} + O \left(\frac {\hat {\sigma} _ {2 , d} ^ {2}}{\log d} \left| D _ {2, t} \right| \| \delta_ {t} \| ^ {2}\right) & i f 0 \leq t \leq T _ {1} \\ O (1) \cdot \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) \gamma_ {2} \| \delta_ {t} \| ^ {2} & i f T _ {1} \leq t \leq T _ {*}, \epsilon \end{array} . \right. \tag {5.2}
$$