SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $\alpha(M)$  is a scaling factor that depends on the number of neurons  $M$ . For simplicity, we assume that the weights  $(v_i, u_i)$  are i.i.d. samples from a Gaussian. We are interested in the dynamics of a gradient flow when  $M$  increases to  $+\infty$  and for specific choices of the scaling factor  $\alpha(M)$ . First, let's introduce some notations. We denote by  $F(w)$  the vector of evaluation of the network on the training data  $F(w) = (f_w(x_1), \ldots, f_w(x_N))$  and by  $J_w$  the Jacobian of  $F(w)$  w.r.t. parameter  $w$ . Moreover, we introduce the NTK matrix  $A_w := J_w J_w^\top$  which is an  $N \times N$  symmetric positive semi-definite matrix. Finally,  $\nabla \mathcal{L}(f_w)$  denotes the vector of size  $N$  representing the gradient of  $\mathcal{L}$  w.r.t.  $F(w)$ . The gradient flow of an objective  $\mathcal{L}(F_w)$  is given by:
$$
\dot {w} _ {t} = - J _ {w _ {t}} \nabla_ {F} \mathcal {L} \left(F _ {w _ {t}}\right),
$$
Additionally, it is often informative to track the time evolution of the vector  $F_{w_t}$ , which is given by the following ODE:
$$
\partial_ {t} F _ {w _ {t}} = - J _ {w _ {t}} J _ {w _ {t}} ^ {\top} \nabla_ {F} \mathcal {L} \left(F _ {w _ {t}}\right) = - A _ {w _ {t}} \nabla_ {F} \mathcal {L} \left(F _ {w _ {t}}\right). \tag {11}
$$
One can see from (11) that the NTK matrix  $A_{w_t}$  arises naturally when expressing the dynamics of the function  $F_{w_t}$ . We can even provide an explicit expression of  $A_{w_t}$  using the NTK kernel, which is a positive semi-definite kernel  $K_w(x,x')$  defined by the inner product:
$$
K _ {w} (x, x ^ {\prime}) = \partial_ {w} f (x) \partial_ {w} f (x ^ {\prime}) ^ {\top} = \alpha (M) ^ {2} \sum_ {i = 1} ^ {M} \left(\gamma (u _ {i} ^ {\top} x) \gamma (u _ {i} ^ {\top} x ^ {\prime}) + (v _ {i}) ^ {2} \gamma^ {\prime} (u _ {i} ^ {\top} x) \gamma^ {\prime} (u _ {i} ^ {\top} x ^ {\prime}) x ^ {\top} x ^ {\prime}\right).
$$
The NTK matrix  $A_w$  is then formed by collecting the values  $(K_w(x_i, x_j))_{1 \leq i, j \leq N}$  on the training data. It is easy to see that when  $\alpha(M) = \frac{1}{\sqrt{M}}$ , then  $K_w(x, x')$  converges to an expectation under the distribution of the weights  $(v_i, u_i)_{1 \leq i \leq M}$  under mild conditions. In this case,  $A_w$  depends on the choice of the distribution of the weights  $(v_i, u_i)_{1 \leq i \leq M}$  and not on the particular samples. On the other hand, when  $\alpha(M) = \frac{1}{M}$ , it is the function  $f_w$  itself that converges to an expectation under the distribution of the weights. This distinction will lead to different qualitative dynamics for the gradient flows as we discuss next.
A.1 The Neural Tangent Kernel limit and the kernel regime  $(\alpha(M) = \frac{1}{\sqrt{M}})$
When choosing  $\alpha(M) = \frac{1}{\sqrt{M}}$  and setting  $M \to +\infty$ , it is shown in [23] that the NTK matrix  $A_{w_t}$  remains constant in time and depends only on the distribution of the weights  $w_0$  at initialization. Hence,  $F_{w_t}$  is effectively the solution of the ODE:
$$
\partial_ {t} F _ {t} = - A _ {w _ {0}} \nabla_ {F} \mathcal {L} (F _ {t}), \qquad F _ {0} = F _ {w _ {0}},
$$
which converges at a linear rate to a global minimizer of  $\mathcal{L}$  whenever  $A_{w_0}$  is invertible and  $\mathcal{L}$  is strongly convex. The NTK limit results in a kernel regime, where the resulting solution is equivalent to the one obtained using a kernel method with the initial NTK kernel  $K_{w_0}(x,x')$  [15, 23].
A.2 The mean-field limit and the feature learning regime  $(\alpha(M) = \frac{1}{M})$
When choosing instead  $\alpha(M) = \frac{1}{M}$ , it is possible to view the limiting function as an expectation over weights drawn from some probability distribution  $\rho$
$$
f _ {\rho} (w) = \int v \gamma (u ^ {\top} x) \mathrm {d} \rho (v, u).
$$
In [48], it is further shown that such a formulation is essentially equivalent to the one considered in (3). As a result, the time evolution in (11) does not simplify as in the NTK limit, since the NTK matrix  $A_{w_t}$  (and NTK kernel  $K_{w_t}$ ) is allowed to evolve in time. This is a limit in which feature learning is possible since the dynamics do not necessarily simplify to a kernel method.
Notations. For an operator  $A$  between two Hilbert spaces, denote by  $\| A\|_{op}$  its operator norm and by  $\| A\| _F$  its Frobenius norm and let  $\sigma^{\star}(A)$  be its smallest singular value whenever it is well-defined. Denote by  $\mathcal{B}(w,R)$  the open ball of radius  $R$  centered at an element  $w$  in a Hilbert space.
B.1 Over-parameterization implies interpolation.
Proof of Proposition 2. Denote by  $F^{\star} = (f^{\star}(x_{1}),\dots,f^{\star}(x_{n}))$ . First, it is easy to see that the matrix  $\partial_v f_w\partial_v f_w^\top$  is exactly equal to the  $N\times N$  matrix  $G(u_0)$  considered in Assumption (A), which is assumed to be invertible. Moreover,  $\partial_v f_w$  is independent of  $v$  since  $f_{(v,u_0)}$  is linear in  $v$ . Hence, we can simply choose:
$$
v ^ {\star} = (\partial_ {v} f _ {(v, u _ {0})}) ^ {\dagger} F ^ {\star} = (\partial_ {v} f _ {(v, u _ {0})}) ^ {\top} (\partial_ {v} f _ {(v, u _ {0})} \partial_ {v} f _ {(v, u _ {0})} ^ {\top}) ^ {\dagger} F ^ {\star} = (\partial_ {v} f _ {(v, u _ {0})}) ^ {\top} G (u _ {0}) ^ {- 1} F ^ {\star}.
$$
Moreover, using again the linearity of  $f$  in  $v$ , we have that
$$
f _ {(v ^ {\star}, u _ {0})} = \partial_ {v} f _ {(v ^ {\star}, u _ {0})} v ^ {\star} = G (u _ {0}) G (u _ {0}) ^ {- 1} F ^ {\star} = F ^ {\star}.
$$
We have therefore shown that  $w^{\star} = (v^{\star}, u_0)$  is a global minimizer of  $w \mapsto \mathcal{L}(f_w)$ . Finally, since  $\nabla \mathcal{L}(f_{w^{\star}}) = \nabla \mathcal{L}(f^{\star}) = 0$ , it follows that  $\phi(w^{\star}) = 0$ .
Proposition 4. Assume that Assumption (A) holds for some vector  $u_0$ . Then, the function  $v \mapsto l(v) = \mathcal{L}(f_{(v,u_0)})$  is convex and satisfies a Polyak-Lojasiewicz inequality:
$$
\left\| \nabla_ {u} l (u) \right\| ^ {2} \geq 2 \mu \sigma_ {0} ^ {2} (l (u) - \mathcal {L} (f ^ {\star})).
$$
with  $\mu$  the strong convexity constant of  $\mathcal{L}$  and  $\sigma_0^2$  the smallest singular value of the matrix  $G(u_0)$  appearing in Assumption (A).
Proof. It is easy to see that  $l(v)$  is convex as a composition of a strongly convex function  $\mathcal{L}$  and a linear function  $v \mapsto f_{(v, u_0)}$ . We will show now that  $l$  satisfies a Polyak-Lojasiewicz (PL) inequality. To this end, we make the following calculations:
$$
\left\| \nabla_ {u} l (u) \right\| ^ {2} = \nabla_ {f} \mathcal {L} \left(f _ {w}\right) ^ {\top} \partial_ {v} f _ {w} \partial_ {v} f _ {w} ^ {\top} \nabla_ {f} \mathcal {L} \left(f _ {w}\right). \tag {12}
$$
It is easy to see that the matrix  $\partial_v f_w\partial_v f_w^\top$  is exactly equal to the  $N\times N$  matrix  $G(u_0)$  considered in Assumption (A), which is assumed to be invertible. Hence, its smallest singular value  $\sigma_0^2$  is positive and the following inequality holds:
$$
\left\| \nabla_ {u} l (u) \right\| ^ {2} \geq \sigma_ {0} ^ {2} \left\| \nabla_ {f} \mathcal {L} \left(f _ {w}\right) \right\| ^ {2}. \tag {13}
$$
On the other hand, since  $\mathcal{L}$  is strongly convex and differentiable, then it satisfies the following PL inequality:
$$
\left\| \nabla_ {f} \mathcal {L} \left(f _ {w}\right) \right\| ^ {2} \geq 2 \mu \left(\mathcal {L} \left(f _ {w}\right) - \mathcal {L} \left(f ^ {\star}\right)\right), \tag {14}
$$
where  $\mu$  is the strong convexity constant of  $\mathcal{L}$ . Hence, we get the desired inequality by combining (13) and (14):
B.2 Global convergence under no blow-up: Proof of Proposition 1
In this section we prove the global convergence result under no blow-up stated in Proposition 1. We start by stating the smoothness assumption on the activation function  $\gamma$  as we will be using it in many places.
(C) The non-linearity  $\gamma$  is twice-continuously differentiable.
As a first step, we need to prove a local existence and uniqueness of the solution to the ODE in (4) which results from Cauchy-Lipschitz theorem applied to the vector field  $\Phi$ . To this end, a key step is to show that  $\Phi$  is locally Lipschitz near initialization when the Jacobian  $J_{w_0}$  is surjective.
Proposition 5 (Regularity of  $\Phi$ ). Assume  $(B)$  and  $(C)$  and that  $w_0$  satisfies Assumption (A). Then, there exists a neighborhood of  $w_0$  so that  $\Phi(w)$  is locally Lipschitz in  $w$ .
Proof of Proposition 5. By application of the Woodbury matrix identity,  $\Phi$  can be equivalently written as:
$$
\Phi (w) = J _ {w} ^ {\top} \left(J _ {w} J _ {w} ^ {\top} + \varepsilon (w) H _ {w} ^ {- 1}\right) ^ {- 1} H _ {w} ^ {- 1} \nabla \mathcal {L} \left(f _ {w}\right). \tag {15}
$$
Moreover, by the smoothness assumption,  $J_{w}$ ,  $H_{w}$ ,  $\epsilon(w)$  and  $\nabla \mathcal{L}(f_w)$  are all differentiable in  $w$ . Additionally, since  $J_{w_0}$  is surjective, it must hold that  $\sigma^2(w) > 0$  in a neighborhood of  $w_0$  so that  $J_{w} J_{w}^{\top}$  remains invertible. Finally, since  $H_{w}$  is always positive, it holds that  $\left(J_{w} J_{w}^{\top} + \varepsilon(w) H_{w}^{-1}\right)^{-1}$  and  $H_{w}^{-1}$  are differentiable in a neighborhood of  $w_0$ . Therefore,  $\Phi$  is Locally Lipschitz in a neighborhood of  $w_0$ .
The next proposition establishes the existence and uniqueness of the dynamics defined in (4).
Proposition 6 (Local existence). Under the same conditions as Proposition 5, there exists a unique solution  $(w_{t})_{t\geq 0}$  to (4) defined up to  $a$ , possibly infinite, maximal time  $T > 0$ .
Proposition 6 is a direct consequence of the Cauchy-Lipschitz theorem which holds since the vector field  $\Phi(w)$  is locally Lipschitz in  $w$  in a neighborhood of  $w_0$  by Proposition 5. The above local existence result does not exclude situations where the dynamics blow up in finite time  $T < +\infty$  and possibly fail to globally minimize the objective  $\mathcal{L}$ . We will show, later, that finite-time blow-ups never occur under additional assumptions on the initial condition.
Proof of Proposition 1. Existence and uniqueness of the continuous-time dynamics up to a maximal time  $T > 0$  follows from Proposition 6. It remains to obtain the global convergence rate when  $T$  is infinite. For simplicity, we will ignore the dependence on  $w_{t}$  in what follows and will write  $b \coloneqq \nabla \mathcal{L}(f_{w_t})$ . By differentiating the objective in time, we get the following expression:
$$
\begin{array}{l} \partial_ {t} \left[ \mathcal {L} \left(f _ {w _ {t}}\right) \right] = - b ^ {\top} J \left(J ^ {\top} H J + \varepsilon I\right) ^ {- 1} J b \\ = - b ^ {\top} H ^ {- 1} b + \varepsilon b ^ {\top} H ^ {- 1} \left(J J ^ {\top} + \varepsilon H ^ {- 1}\right) ^ {- 1} H ^ {- 1} b \\ = - b ^ {\top} H ^ {- 1} b + \varepsilon b ^ {\top} H ^ {- \frac {1}{2}} \left(H ^ {- \frac {1}{2}} \left(J J ^ {\top} + \varepsilon H ^ {- 1}\right) ^ {- 1} H ^ {- \frac {1}{2}}\right) H ^ {- \frac {1}{2}} b \\ \leq - b ^ {\top} H ^ {- 1} b \left(1 - \varepsilon \left\| H ^ {- \frac {1}{2}} \left(J J ^ {\top} + \varepsilon H ^ {- 1}\right) ^ {- 1} H ^ {- \frac {1}{2}} \right\| _ {o p}\right) \\ \leq - b ^ {\top} H ^ {- 1} b \underbrace {\left(1 - \frac {\varepsilon}{\varepsilon + \sigma^ {\star} (H) \sigma^ {\star} (J J ^ {\top})}\right)} _ {\geq \eta : = (1 + \alpha / \mu_ {H}) ^ {- 1}} \\ \leq - \eta b ^ {\top} H ^ {- 1} b \leq - \eta L _ {H} ^ {- 1} \| b \| ^ {2} \leq - 2 \eta \mu L _ {H} ^ {- 1} \mathcal {L} \left(f _ {w _ {t}}\right). \\ \end{array}
$$
We used the Woodbury matrix inversion lemma to get the second line while the third and fourth lines use simple linear algebra. To get the fifth line, we used Lemma 3 to upper-bound the operator norm appearing in the fourth line. The last line follows by definition of  $\varepsilon(w)$ , the fact that  $H \leq L_{H}I$  and the fact that  $\frac{1}{2}\|\nabla\mathcal{L}(f)\|^{2} \geq \mu\mathcal{L}(f)$  by  $\mu$ -strong convexity of  $\mathcal{L}$ . This allows to directly conclude that  $\mathcal{L}(f_{w_t}) \leq \mathcal{L}(f_{w_0})e^{-\mu_{GN}t}$ , with  $\mu_{GN} \coloneqq 2\eta\mu L_H^{-1}$ .
B.3 Control of the singular values of NTK matrix
In this section, we are interested in the evolution of the smallest singular value of the NTK matrix  $A_{w_t} \coloneqq J_{w_t}J_{w_t}^\top$ . For any arbitrary positive radius  $R > 0$ , define the constant  $C_R$  and stopping time  $T_R$  as follows:
$$
C _ {R} := \sup  _ {w \in \mathcal {B} (w _ {0}, R)} \| \partial_ {w} J _ {w} \| _ {o p} <   + \infty
$$
$$
T _ {R} := \sup  \left\{t \geq 0 \quad | | w _ {t} - w _ {0} | | <   R, \text {a n d} \sigma^ {\star} \left(A _ {w _ {t}}\right) > 0 \right\},
$$
The stopping time  $T_R$  characterizes the smallest-time when the dynamics becomes either degenerate  $(\sigma^*(A_{w_t}) = 0)$  or grows too far away from initialization  $(w_t \neq \mathcal{B}(w_0, R))$ .
To control the singular values of  $A_{w_t}$ , we find it useful to consider the dynamics of the Frobenius norm of the pseudo-inverse  $J_w^\dagger$  which we denote by  $a_t \coloneqq \left\| J_t^\dagger \right\|_F$ .
Indeed, simple linear algebra provides the following lower-bound on  $\sigma^{\star}(A_{w_t})$  in terms of  $a_{t}$ :
$$
\sigma^ {\star} \left(A _ {w _ {t}}\right) \geq a _ {t} ^ {- 2}. \tag {16}
$$
Note that, while the space of parameters can even be infinite-dimensional (mean-field limit),  $J_{w}$  has a finite-dimensional range and thus always admits a pseudo-inverse  $J_{w}^{\dagger}$  given by:
$$
J _ {w} ^ {\dagger} = J _ {w} ^ {\top} \left(J _ {w} J _ {w} ^ {\top}\right) ^ {\dagger}.
$$
The next proposition provides a differential inequality for  $a_{t} = \left\| J_{w_{t}}^{\dagger}\right\|_{F}$  in terms of  $\| \dot{w}_t\|$  that holds for all times smaller than  $T_{R}$ .
Proposition 7. Assume  $(\mathbf{B})$  and  $(C)$  and that  $w_0$  satisfies Assumption  $(A)$ , then for any  $R > 0$ , the time  $T_R$  is positive and smaller than the maximal time  $T$ , i.e.  $0 < T_R \leq T$ . Moreover, for any  $t \in [0, T_R)$ ,  $a_t$  satisfies:
$$
\left| \dot {a} _ {t} \right| \leq a _ {t} ^ {2} C _ {R} \| \dot {w} _ {t} \|,
$$
Proof. Positive stopping time  $T_{R}$ . By construction,  $T_{R}$  must be smaller than  $T$  as it requires  $w_{t}$  to be well-defined. Moreover,  $J_{w_0}$  is surjective by assumption and  $J_{w}$  is continuous in  $w$ , therefore  $J_{w_t}$  must also be surjective for small enough positive times  $t > 0$ . Additionally, it must hold that  $\| w_{t} - w_{0}\| \leq R$  for small positive times, by continuity of  $t\mapsto w_t$ . This allows shows that  $T_{R} > 0$ .
Dynamics of  $J_{w_t}^\dagger$ . By smoothness of  $J_w$  it holds that  $J_{w_t}^\dagger$  is differentiable on the open interval  $(0, T_R)$  and satisfies the following ODE obtained by direct differentiation:
$$
\dot {J} _ {t} ^ {\dagger} = - J _ {t} ^ {\dagger} \dot {J} _ {t} J _ {t} ^ {\dagger} + (I - P _ {t}) \dot {J} _ {t} ^ {\top} \left(J _ {t} J _ {t} ^ {\top}\right) ^ {- 1}. \tag {17}
$$
In (17), we introduced notation  $J_{t} \coloneqq J_{w_{t}}$  and its time derivative  $\dot{J}_{t}$  for simplicity and denote by  $P_{t}$ , the projector:  $P_{t} \coloneqq J_{t}^{\top}\left(J_{t}J_{t}^{\top}\right)^{-1}J_{t}$ .
Controlling  $|\dot{a}_t|$ . Taking the time derivative of  $a_{t}$  and recalling the evolution (17) yields:
$$
a _ {t} \dot {a} _ {t} = - \left(J _ {t} ^ {\dagger}\right) ^ {\top} J _ {t} ^ {\dagger} \dot {J} _ {t} J _ {t} ^ {\dagger}.
$$
where we used that  $\left(J_t^\dagger\right)^\top (I - P_t) = 0$  by Lemma 1 in Appendix B.5. Furthermore, taking the absolute value of the above equation, and recalling that the Frobenius norm is sub-multiplicative for the product of operators (Lemma 4 in Appendix B.5), we directly get:
$$
\left. a _ {t} \left| \dot {a} _ {t} \right| \leq \left| \left(J _ {t} ^ {\dagger}\right) ^ {\top} J _ {t} ^ {\dagger} \dot {J} _ {t} J _ {t} ^ {\dagger} \right| \leq a _ {t} ^ {3} \| \dot {J} _ {t} \| _ {F}. \right. \tag {18}
$$
By the chain rule, we have that  $\dot{J}_t = \partial_w J_t \dot{w}_t$  therefore,  $\left\| \dot{J}_t \right\|_F \leq \left\| \partial_w J_t \right\|_{op} \left\| \dot{w}_t \right\|$ . Moreover, since  $t < T_R$ , then it must hold that  $w_t \in B(w_0, R)$  so that  $\left\| \partial_w J_t \right\|_{op} \leq C_R$ . Henceforth,  $\left\| \dot{J}_t \right\|_F \leq C_T \left\| \dot{w}_t \right\|$  and (18) yields the desired inequality after dividing by  $a_t$  which remains positive all times  $t \in [0, T_R)$  since  $J_t$  does not vanish.
The next proposition provides an estimate of the time derivatives  $|\dot{a}_t|$  and  $\| \dot{w}_t\|$  up to time  $T_R$ . For simplicity, we introduce the notation  $\Delta_t = \| \nabla \mathcal{L}(f_{w_t})\|$ .
Proposition 8. Assume  $(\mathbf{B})$  and  $(\mathbf{C})$  and that  $w_{0}$  satisfies Assumption (A). Then, for any  $R > 0$ , the time  $T_{R}$  is positive and smaller than the maximal time  $T$  up to which (4) is defined, i.e.  $0 < T_{R} \leq T$ . Moreover, define  $C \coloneqq L\mu^{-1}\mu_H^{-1}\Delta_0C_R$ . Then  $a_{t}$  and  $\| \dot{w}_t\|$  satisfy the following differential inequality on  $[0,T_R)$ :
$$
\begin{array}{l} | \dot {a} _ {t} | \leq C a _ {t} ^ {3} e ^ {- \frac {\mu_ {G N}}{2} t} \\ \| \dot {w} _ {t} \| \leq C C _ {R} ^ {- 1} a _ {t} e ^ {- \frac {\mu_ {G N}}{2} t} \\ \end{array}
$$
Proof of Proposition 8. By application of Proposition 7 to  $(w_{t})_{0\leq t < T}$ , we directly have that  $T_{R}$  is positive and smaller than  $T$  and the following estimate holds for any  $t < T_R$ :
It remains to upper-bound  $\| \dot{w}_t\|$ .
Controlling  $\| \dot{w}_t\|$ . Ignoring the dependence on  $w_{t}$  in the notations for simplicity and setting  $B\coloneqq H^{\frac{1}{2}}J$ , we have, by definition of the dynamical system (4):
$$
\begin{array}{l} \left\| \dot {w} _ {t} \right\| \leq \left\| \left(J ^ {\top} H J + \varepsilon I\right) ^ {- 1} J ^ {\top} \nabla \mathcal {L} (f) \right\| \\ = \left\| J ^ {\top} \left(J J ^ {\top} + \varepsilon H ^ {- 1}\right) ^ {- 1} H ^ {- 1} \nabla \mathcal {L} (f) \right\| \\ = \left\| B ^ {\top} \left(B B ^ {\top} + \varepsilon\right) ^ {- 1} H ^ {- \frac {1}{2}} \nabla \mathcal {L} (f) \right\| \\ \leq \left\| B B ^ {\top} \left(B B ^ {\top} + \varepsilon I\right) ^ {- 2} \right\| _ {F} ^ {\frac {1}{2}} \left\| H ^ {- \frac {2}{2}} \nabla \mathcal {L} (f) \right\| \\ \leq \left\| B ^ {\dagger} \right\| _ {F} \| \nabla \mathcal {L} (f) \| \\ \leq \mu_ {H} ^ {- 1} \| J ^ {\dagger} \| _ {F} \| \nabla \mathcal {L} (f) \| = \mu_ {H} ^ {- 1} a _ {t} \Delta_ {t}, \\ \end{array}
$$
where the second line follows by the Woodbury matrix identity, and the third line follows by simple linear algebra. For the fourth line, we use the properties of the Frobenius norm. The fifth and last lines are direct consequences of Lemmas 2 and 4 in Appendix B.5.
Concluding. We can combine the upper-bound on  $|\dot{a}_t|$  and  $\| \dot{w}_t\|$  to get:
$$
\left| \dot {a} _ {t} \right| \leq \mu_ {H} ^ {- 1} C _ {R} a _ {t} ^ {3} \Delta_ {t}.
$$
Finally, since  $t \leq T$ , it follows from Proposition 1 that  $\Delta_t \leq \frac{L}{\mu} \Delta_0 e^{-\frac{\mu G N}{2} t}$  which allows to conclude.
The next proposition controls  $a_{t}$  and  $\| w_{t} - w_{0}\|$  in terms of  $T_{R}, C_{R}$ , and other constants.
Proposition 9. Consider the same conditions as in Proposition 8. Let  $\zeta := \frac{4L}{\mu\mu_H\mu_{GN}}$  and consider the (possibly infinite) time  $T^{+}$  defined as:
$$
T ^ {+} = - 2 \mu_ {G N} ^ {- 1} \log \left(1 - (\zeta C _ {R} \Delta_ {0} a _ {0} ^ {2}) ^ {- 1}\right),
$$
with the convention that  $T^{+} = +\infty$  if  $1 \geq \zeta C_{R}\Delta_{0}a_{0}^{2}$ . Then for any time  $t < T^{-} \coloneqq \min (T_R, T^+)$  it holds that:
$$
\left\| w _ {t} - w _ {0} \right\| \leq \zeta \Delta_ {0} a _ {0}, \quad a _ {t} \leq a _ {0} \left(1 + \zeta C _ {R} \Delta_ {0} a _ {0} ^ {2} \left(e ^ {- \frac {\mu_ {G N}}{2} t} - 1\right)\right) ^ {- \frac {1}{2}} \tag {19}
$$
Proof of Proposition 9. First, by Proposition 8, we know that  $a_{t}$  satisfies the following differential inequality:
$$
\left| \dot {a} _ {t} \right| \leq C a _ {t} ^ {3} e ^ {- \frac {\mu_ {G N}}{2} t},
$$
with  $C \coloneqq \frac{L}{\mu_H\mu}\Delta_0C_R$ . We will control  $a_{t}$  by application of Gronwall's inequality. To this end, consider the ODE:
$$
\dot {b} _ {t} = C b _ {t} ^ {3} e ^ {- \frac {\mu_ {G N}}{2} t}, \qquad b _ {0} = a _ {0} > 0.
$$
We know by Lemma 5 in Appendix B.5 that  $b_{t}$  is given by:
$$
\begin{array}{l} b _ {t} = b _ {0} \Big (1 + 4 C \mu_ {G N} ^ {- 1} b _ {0} ^ {2} \Big (e ^ {- \frac {\mu_ {G N}}{2} t} - 1 \Big) \Big) ^ {- \frac {1}{2}}, \\ = a _ {0} \left(1 + \zeta 2 C _ {R} \Delta_ {0} \mu_ {G N} ^ {- 1} a _ {0} ^ {2} \left(e ^ {- \frac {\mu_ {G N}}{2} t} - 1\right)\right) ^ {- \frac {1}{2}} \\ \end{array}
$$
for all times  $t < T^{+}$ . Therefore, by Gronwall's inequality, it must hold that  $a_{t} \leq b_{t}$  for all times  $t < \min(T_{R}, T^{+})$ . Moreover, recalling now that  $\| \dot{w}_{t} \| \leq L \mu^{-1} \mu_{H}^{-1} \Delta_{0} a_{t} e^{-\frac{\mu_{GN}}{2} t}$  by Proposition 8, we directly get that:
$$
\left\| \dot {w} _ {t} \right\| \leq \frac {C}{C _ {R}} a _ {0} e ^ {- \frac {\mu_ {G N}}{2} t} \left(1 + 4 C \mu_ {G N} ^ {- 1} b _ {0} ^ {2} \left(e ^ {- \frac {\mu_ {G N}}{2} t} - 1\right)\right) ^ {- \frac {1}{2}}. \tag {20}
$$
Integrating the above inequality in time yields the following estimate on  $\| w_t - w_0\|$  for any  $t < \min (T_R,T^+)$ :
$$
\begin{array}{l} \left\| w _ {t} - w _ {0} \right\| \leq \int_ {0} ^ {t} \left\| \dot {w} _ {s} \right\| d s \\ \leq \frac {1 - \left(1 + 4 C \mu_ {G N} ^ {- 1} a _ {0} ^ {2} \left(e ^ {- \frac {\mu_ {G N}}{2} t} - 1\right)\right) ^ {\frac {1}{2}}}{C _ {R} a _ {0}} \\ \leq 4 L \left(\mu \mu_ {H} \mu_ {G N}\right) ^ {- 1} \Delta_ {0} a _ {0} = \zeta \Delta_ {0} a _ {0}, \\ \end{array}
$$
where the second line follows by explicitly integrating the r.h.s. of (20) while the last line follows by the concavity of the square-root function.
Proposition 9 shows that the Frobenius norm of the  $J_{t}^{\dagger}$  remains bounded at all times provided that  $\Delta_0 a_0^2$  is small enough, a quantity that depends only on the initial conditions. Moreover, it also shows that making the product  $\Delta_0 a_0$  small enough at initialization ensures that  $w_{t}$  remains in a ball of radius  $R$  around  $w_{0}$ . We exploit these two consequences to prove Proposition 3 in Appendix B.4.
B.4 Absence of blow-up for almost-optimal initial linear layer: proof of Proposition 3
To prove Proposition 3, we rely on the results of Appendix B.3 which allow controlling the evolution of the smallest singular value of the NTK matrix  $A_{w_t}$ . More precisely, Proposition 3 is a direct consequence of Proposition 9 with a particular choice for the initialization  $w_0$ .
Proof of Proposition 3. By assumption on  $w_0$ , it holds that
$$
\Delta_ {0} := \| \nabla \mathcal {L} \left(f _ {w _ {0}}\right) \| <   \epsilon = \frac {\mu \mu_ {H} \mu_ {G N}}{8 L N} \min  \left(R, C _ {R} ^ {- 1}\right) \min  \left(\sigma_ {0}, \sigma_ {0} ^ {2}\right).
$$
Moreover, by definition of  $a_0 \coloneqq \left\| J_0^\dagger \right\|_F$  and of  $\sigma_0^2 = \sigma_{min}(G(u_0))$ , we know that  $a_0^2 < N\sigma_0^{-2}$ . Hence, we get
$$
\Delta_ {0} <   \frac {1}{2 \zeta} \min (R, C _ {R} ^ {- 1}) \min \left(a _ {0} ^ {- 1} N ^ {- \frac {1}{2}}, a _ {0} ^ {- 2}\right),
$$
where we introduced  $\zeta \coloneqq \frac{4L}{\mu\mu_H\mu_{GN}}$ . This above inequality directly yields:
$$
\zeta \Delta_ {0} a _ {0} \leq \frac {R}{2 N ^ {\frac {1}{2}}}, \quad \zeta C _ {R} \Delta_ {0} a _ {0} ^ {2} <   1, \tag {21}
$$
Therefore, we can use the above inequalities in the estimates (19) of Proposition 9 to get:
$$
\left\| w _ {t} - w _ {0} \right\| \leq \frac {R}{2 N ^ {\frac {1}{2}}}, \quad a _ {t} \leq a _ {0} \left(1 - \zeta C _ {R} \Delta_ {0} a _ {0} ^ {2}\right) ^ {- \frac {1}{2}},
$$
for any time  $t\in [0,T_R)$ , where  $T_{R}$  is defined as: