SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
T _ {R} := \sup  \left\{t \geq 0 \quad | \quad \| w _ {t} - w _ {0} \| <   R, \text {a n d} \sigma^ {\star} \left(A _ {w _ {t}}\right) > 0 \right\}.
$$
This implies, in particular, that  $w_{t}$  never escapes the ball  $\mathcal{B}(w_0,R)$ . Moreover, using (16) we get that the smallest singular value  $\sigma^{\star}(A_{w_t}^\top)$  is always lower-bounded by a positive constant  $\tilde{\sigma} := a_0^{-2}\big(1 - \zeta C_R\Delta_0a_0^2\big)$ . Therefore,  $T_{R}$  must necessarily be greater or equal to  $T$ , the maximal time over which  $w_{t}$  is defined. Additionally, if  $T$  was finite, then  $w_{t}$  must escape any bounded domain. This contradicts the fact that  $w_{t} \in \mathcal{B}(w_{0},R)$ . Therefore,  $T$  must be infinite and the inequality (6) applies at all times  $t \geq 0$ .
B.5 Auxiliary results
Lemma 1. Let  $J$  be a linear operator from a Hilbert space  $\mathcal{W}$  to a finite-dimensional space  $\mathcal{H}$  and assume that  $JJ^{\top}$  is invertible. Define the projector  $P = J^{\top}(JJ^{\top})^{-1}J$ . Then the following relations hold:
$$
(J ^ {\dagger}) ^ {\top} (I - P) = 0
$$
Lemma 2. Let  $J$  be a linear operator from a Hilbert space  $\mathcal{W}$  to a finite-dimensional space  $\mathcal{H}$  and assume that  $JJ^{\top}$  is invertible. Then for any non-negative number  $\epsilon$ , it holds that:
$$
\left\| J J ^ {\top} \left(J J ^ {\top} + \epsilon I\right) ^ {- 2} \right\| _ {F} \leq \left\| J ^ {\dagger} \right\| _ {F} ^ {2}
$$
Lemma 3. Let  $\mathcal{W}$  be a Hilbert space and  $\mathcal{R}$  be a finite-dimensional euclidean space. Let  $J$  be a linear operator from  $\mathcal{W}$  to  $\mathcal{R}$  and  $H$  an invertible operator from  $\mathcal{R}$  to itself. Further, assume that  $JJ^{\top}$  is invertible. Then the following holds for any non-negative number  $\epsilon$ :
$$
\left\| H ^ {- \frac {1}{2}} \left(J J ^ {\top} + \epsilon H ^ {- 1}\right) ^ {- 1} H ^ {- \frac {1}{2}} \right\| _ {o p} \leq \frac {1}{\epsilon + \sigma^ {\star} (H) \sigma^ {\star} \left(J J ^ {\top}\right)}.
$$
Lemma 4. Let  $J$  and  $K$  be two Hilbert-Schmidt operators between two Hilbert spaces. Then, it holds that:
$$
\left\| J K ^ {\top} \right\| _ {F} \leq \left\| J \right\| _ {o p} \left\| K \right\| _ {F} \leq \left\| J \right\| _ {F} \left\| K \right\| _ {F}
$$
Lemma 5. Let  $C$  and  $r$  be two positive constants. For a given initial condition  $b_0 > 0$ , consider the following differential equation:
$$
\dot {b} _ {t} = C b _ {t} ^ {3} e ^ {- r t}.
$$
Then  $b_{t}$  is defined for any time  $t < T^{+}$  where  $T^{+}$  is given by:
$$
T ^ {+} = \left\{ \begin{array}{l l} r ^ {- 1} \log \left(1 - \frac {r}{2 C b _ {0} ^ {2}}\right), & \quad 2 C r ^ {- 1} b _ {0} ^ {2} > 0 \\ + \infty & \quad O t h e r w i s e. \end{array} \right.
$$
Moreover,  $b_{t}$  is given by:
$$
b _ {t} = b _ {0} \left(1 + 2 C r ^ {- 1} b _ {0} ^ {2} \left(e ^ {- r t} - 1\right)\right) ^ {- \frac {1}{2}}, \qquad t <   T ^ {+}.
$$
Proof. First note that a local solution exists by the Cauchy-Lipschitz theorem. Moreover, it must never vanish since otherwise, it would coincide with the null solution  $b_{t} = 0$  by uniqueness. However, this is impossible, since  $b_{0} > 0$ . We can therefore divide the ODE by  $b_{t}^{3}$  and explicitly integrate it which gives:
$$
b _ {t} ^ {- 2} = b _ {0} ^ {- 2} + 2 r ^ {- 1} C \left(e ^ {- r t} - 1\right).
$$
The solution is defined only for times  $t$  so that the r.h.s. is positive, which is exactly equivalent to having  $t < T^{+}$ .
C.1 Multi-seed experiments
We present result obtained for 5 independent runs. Each run uses a different initialization for the parameters of both student and teacher networks. Additionally, the training and test data are all generated independently for each run. All these results are obtained using SiLU non-linearity. It is clear from Figure 3, that the results display little variability w.r.t. the seed of the experiment, which justifies the single-seed setting chosen in the main paper.
Effect of over-parameterization Figure 4 (Left) shows the effect of over-parameterization of the network on generalization error. It appears that the generalization error remains stable with an increasing parameterization for both GN and GD, as soon as the network has enough overparameterization to fit the data exactly. On the other hand, the test error of RF improves with increasing over-parameterization which is consistent with [22].
Evolution of the smallest singular value of  $A_w$ . Figure 4 (Right) shows that the smallest singular value  $\sigma^*(A_w)$  systematically increases during training, especially in the feature learning regime (bright red colors). Hence, the dynamics are far from blowing up even as the features change substantially. On the other hand,  $\sigma^*(A_w)$  remains nearly constant for large values of  $\tau_0$  (darker colors) which indicates that the features barely change in the kernel regime.
Figure 3: (Left) Test error vs std of the initial weights  $\tau_0$ . All models are trained until the training objective is smaller than  $10^{-6}$  using  $M = 5000$  hidden units and  $N = 500$ . Confidence interval at  $95\%$  estimated from 5 independent runs shown in shaded colors. (Right) Test error vs training sample size. All models are trained until the training objective is smaller than  $10^{-6}$  using  $M = 5000$  hidden units and initial std of the weights  $\tau_0 = 1$ . Confidence interval at  $95\%$  estimated from 5 independent runs shown in shaded colors.
Figure 4: (Left) Test error vs number of hidden units. All models are trained using  $N = 2000$  training samples until the training objective is smaller than  $10^{-6}$ . The initial std of the weights is set to  $\tau_0 = 1$ . (Right) Evolution of  $\sigma^{\star}(A_w)$  during training using Gauss-Newton for different values of the initial std  $\tau_0$  of the weights. All models are trained using  $N = 500$  training samples until the training objective is smaller than  $10^{-6}$ . We used  $M = 5000$  units. Confidence interval at 95% estimated from 5 independent runs shown in shaded colors.
C.2 Effect of the sample-size
Figure 5 shows the evolution of the test error as the training sample size  $N$  varies in  $\{100,200,500,1000,1500\}$  for two choices of activation functions: ReLU (Left) and SiLU (Right). In both cases, we display the results for two different values for the initial std  $\tau_0$  ( $10^{-3}$  and 1). Results are reported for the best performing step size for each of GN and GD. The first observation is that the performance gap between random features RF and learned features (using either GN or GD) keeps increasing with the sample size. This observation suggests that the learned features use training samples more efficiently than random features. The second observation is the almost affine relation between the generalization error  $\mathcal{L}_{test}$  and training sample size in a logarithmic scale with a strong dependence of the slope on the variance at initialization and the optimization method. Such affine relation implies the following upper-bound on  $\mathcal{L}_{test}$  in terms of sample-size  $N$ :
$$
\mathcal {L} _ {t e s t} \leq C \frac {1}{N ^ {\alpha}},
$$
for some positive constants  $C$  and  $\alpha$ . The coefficient  $\alpha$  controls the speed at which the estimator approximates  $f^{\star}$  as one accesses more training data and usually appears in learning theory to describe
the statistical convergence rate of a given estimator [46]. The consistent and strong dependence of  $\alpha$  on the optimization methods (GN vs GD) and regime (kernel vs feature learning) further confirms the implicit bias of both initialization and optimization algorithms.
Figure 5: Final values of various metrics vs training sample size for both ReLU (left) and SiLU (right) networks. Two values for  $\tau_0$  are displayed for each method  $\tau_0 \in \{10^{-3}, 1\}$  while training data ranges from  $N = 100$  to  $N = 1500$ . All models are optimized until the training objective is smaller than  $10^{-6}$  using  $M = 5000$  hidden units. For both GD and GN, results are reported for the best-performing step-size  $\lambda$  selected according to the test loss on a regular logarithmic grid ranging from  $10^{-3}$  to  $10^{3}$ .
C.3 Choice of the non-linearity
Figure 6 shows the effect of step-size on the final performance of a SiLU network with two choices for the temperature parameter:  $\beta = 1$  and  $\beta = 10^6$ . Similarly to Figure 1, it is clear that increasing the step size in Gauss-Newton results in features that do not generalize well, while the opposite is observed in case of gradient descent. On the other hand, depending on the value of  $\beta$ , we observe that GN either outperforms GD ( $\beta = 10^6$ ) or the opposite ( $\beta = 1$ ). This variability suggests that the effectiveness of an optimization method in improving generalization is dependent on the specific characteristics of the problem at hand. Finally, note that for  $\beta = 10^6$ , we recover almost the same results as those obtained for ReLU in Figure 1 (right). This is due to the fact SiLU becomes a tight smooth approximation to ReLU for large values of  $\beta$ . We also notice that when using a larger std for the initial weights (ex.  $\tau_0 = 1$ ), using either ReLU or SiLU with ( $\beta = 1$ ) yields similar results, as shown in Figure 7.
C.4 Experiments on MNIST
We have performed a series of experiments on MNIST to illustrate the behavior of GN on higher dimensional data (notably, with multidimensional targets). To stay as close as possible to the theoretical framework, we modify the original MNIST dataset. Specifically, we construct a student/teacher setup in which the trained network can achieve a zero-loss, which is an assumption made in Section 3.1.
Building datasets based on MNIST by using a teacher. In our theoretical framework, we restricted ourselves to cases where the minimum loss can be achieved on the considered model. To make this condition hold, we build a new dataset by using a teacher trained on the MNIST classification task [16]:
1. we construct a training set for the teacher which consists of  $\mathcal{D}_{tr} \cup \mathcal{D}_{ts}$ , where  $\mathcal{D}_{ts}$  is the original MNIST test dataset and  $\mathcal{D}_{tr}$  is a balanced subset of size 3000 of the original MNIST training set;
2. we then train a teacher neural network with one hidden layer of size 50 and ReLU activation function on training set  $\mathcal{D}_{tr} \cup \mathcal{D}_{ts}$  using a cross-entropy objective;
Figure 6: Final values of various metrics vs the step size for both SiLU network with  $\beta = 1$  (left) and  $\beta = 10^{6}$  (right). (Right figure) The std of the weights at initialization is set to  $\tau_0 = 10^{-3}$ . All models are optimized up to a training error of  $10^{-6}$  or until the maximum number of steps is exceeded,  $(M = 5000, N = 500)$ .
Figure 7: Final values of various metrics vs the step size for both ReLU (left) and SiLU (right) networks. The std of the weights at initialization is set to  $\tau_0 = 1$ . All models are optimized up to a training error of  $10^{-6}$  or until the maximum number of steps is exceeded, ( $M = 5000$ ,  $N = 500$ ).
3. given the trained teacher  $f_{T}$ , we build new datasets  $\bar{\mathcal{D}}_{tr}$  and  $\bar{\mathcal{D}}_{ts}$  for the student network:
$$
\bar {\mathcal {D}} _ {t r} := \left\{\left(x _ {i}, \operatorname {S o f t m a x} \left(f _ {T} \left(x _ {i}\right)\right)\right): \left(x _ {i}, y _ {i}\right) \in \mathcal {D} _ {t r} \right\},
$$
$$
\bar {\mathcal {D}} _ {t s} := \left\{\left(x _ {i}, \operatorname {S o f t m a x} \left(f _ {T} \left(x _ {i}\right)\right)\right): \left(x _ {i}, y _ {i}\right) \in \mathcal {D} _ {t s} \right\}.
$$
Training procedure. Once the pair of training and testing datasets  $(\bar{\mathcal{D}}_{tr},\bar{\mathcal{D}}_{ts})$  has been built, we use them to train and test a student neural network based on (10), with one hidden layer of size 5000 and a ReLU activation function. We use a similar initialization scheme for the student's network and fix the std of the initial weights to  $\tau_0 = 0.001$ . For training, we use a quadratic loss between the student and teacher's outputs and perform a maximum of 100000 iterations of either GN or GD. The training stops whenever the training loss goes below  $10^{-6}$ . We vary the step size  $\lambda$  on a logarithmic range  $[10^{-3},10^{3}]$  for both methods and evaluate the test loss and test loss after linear re-fitting.
Figure 8: Final values of various metrics vs the step size for a ReLU network trained using MNIST data. Std of the initial weights is set to  $\tau_0 = 0.001$ . For GD, results are displayed up to  $\lambda = 10$  the algorithm diverges for higher values.
Results. Figure 8 shows the evolution of the test loss (both with or without re-fitting) as a function of the step size. We observe similar behavior as in the main experiments, with hidden learning occurring for GN when using small step sizes and a degraded generalization error as the step size increases. On the other hand, generalization improves for GD by increasing the step size up to the point where optimization becomes unstable.