SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The above result ensures a global solution can always be found by optimizing the linear weights while keeping the hidden weights fixed. In practice, this can be achieved using any standard optimization algorithm such as gradient descent since the function  $v \mapsto l(v) = \mathcal{L}(f_{(v,u_0)})$  converges to a global optimum as it satisfies a Polyak-Lojasiewicz inequality (see Proposition 4 in Appendix B.1). However, this procedure is not of interest here as it does not allow us to learn the hidden features. Nonetheless, since  $w^{\star} = (v^{\star},u_{0})$  is a fixed point of the Gauss-Newton dynamics, we can expect that an initial condition  $w_{0} = (v_{0},u_{0})$ , where  $v_{0}$  is close to  $v^{\star}$ , results in a well-behaved dynamics. The following proposition, proven in Appendix B.4, makes the above intuition more precise.
Proposition 3. Let  $w_0 = (v_0, u_0) \in \mathcal{W}$  be an initial condition so that Assumption (A) holds, i.e. the gram matrix  $G(u_0)$  is invertible, and denote by  $\sigma_0^2$  the smallest eigenvalue of  $G(u_0)$ . Assume that the activation  $\gamma$  is twice-continuously differentiable and that Assumption (B) holds. Let  $R$  be any arbitrary positive number and define  $C_R = \sup_{w \in \mathcal{B}(w_0, R)} \| \partial_w J_w \|_{op}$  where  $\| . \|_{op}$  denotes the operator norm. If the linear layer is almost optimal  $v_0$  in the sense that:
$$
\left\| \nabla \mathcal {L} \left(f _ {w _ {0}}\right) \right\| <   \epsilon , \quad \epsilon := \left(\mu \mu_ {H} \mu_ {G N} / 8 L N\right) \min  \left(R, C _ {R} ^ {- 1}\right) \min  \left(\sigma_ {0}, \sigma_ {0} ^ {2}\right), \tag {7}
$$
then (4) is defined at all times, i.e.  $T = +\infty$ , the objective  $\mathcal{L}(f_{w_t})$  converges at a linear rate to 0 according to (6) and the parameters  $w_t$  remain within a ball of radius  $R$  centered around  $w_0$ .
Proposition 3 essentially states that the dynamics never blow up, and hence, converge globally at the linear rate in (6) by Proposition 1, provided the hidden features are diverse enough and the initial linear weights optimize the objective well enough. As discussed in Section 3.1, the first condition on the hidden weights  $u_{0}$  typically holds for a Gaussian initialization when the data are non-degenerate. Additionally, the near-optimality condition on the linear weights  $v_{0}$  can always be guaranteed by optimizing the linear weights while keeping the hidden ones fixed.
Remark 1. In the proof of Proposition 3, we show that the occurrence of a finite-time blow-up is tightly related to the Neural Tangent Kernel matrix  $A_{w_t}$  becoming singular at  $t$  increases which causes the vector field  $\Phi$  to diverge. When the NTK matrix  $A_{w_t}$  at time  $t$  is forced to remain close to the initial one  $A_{w_0}$ , as is the case in the NTK limit considered in [11], one can expect that the singular values of  $A_{w_t}$  remain bounded away from 0 since this is already true for  $A_{w_0}$ . Therefore, the main technical challenge is the study, performed in Propositions 7 to 9 of Appendix B.3, of the time evolution of the eigenvalues of  $A_{w_t}$  in the mean-field limit, where  $A_{w_t}$  is allowed to differ significantly from the initial NTK matrix  $A_{w_0}$ . The final step is to deduce that, while the weights  $w_t$  are allowed to be in a ball centered in  $w_0$  of arbitrarily large radius  $R$ , the singular values remain bounded away from 0 provided the initial linear weights  $v_0$  satisfy the condition in (7).
Remark 2. Comparing Proposition 3 to the convergence results for gradient flows in the kernel regime [15, Theorem 2.4], our result also holds for an infinite dimensional space of parameters such as in the mean-field limit of one-hidden layer network. Moreover, while [15, Theorem 2.4] requires the norm  $w_{t}$  to be within a ball of fixed radius  $R_0$  determined by initial smallest singular value  $\sigma (w_0)$  of  $A_{w_0}$ , our result allows  $w_{t}$  to be within arbitrary distance  $R$  regardless of the initial  $\sigma (w_0)$  provided the condition in (7) holds. Proposition 3 is more similar in flavor to the convergence result of the Wasserstein gradient flow for some probability functionals in [13, Theorem 3.3, Corollary 3.4] which requires the initial objective to be smaller than a given threshold.
While the results of this section guarantee global convergence of the training objective, they do not guarantee learning features that generalize well. Next, we empirically show that GN can exhibit regimes where it yields solutions with good generalization properties.
5 Empirical study of generalization for a Gauss-Newton algorithm
We perform a comparative study between gradient descent (GD) and Gauss-Newton (GN) method in the context of over-parameterized networks. Since we observed little variability when changing the
seed, we have chosen to use a single seed for the sake of clarity in the figures when presenting the results of this section, while deferring the analysis for multiple seeds to Appendix C.1. Additional experiments using MNIST dataset [16] are provided in Appendix C.4. The results presented below were obtained by running 720 independent runs, each of which optimizes a network given a specific configuration on a GPU. The total time for all runs was 3600 GPU hours.
We consider a regression task on a synthetic dataset consisting of  $N$  training points  $(X_{n},Y_{n})_{1\leq n\leq N}$ . The objective is to minimize the mean-squared error  $\mathcal{L}(f_w)$ , defined in (1), over the parameters  $w$  of a model  $f_{w}$  that predicts the target values  $Y_{n}$  based on their corresponding input points  $X_{n}$ .
Data generation. We generate  $N$  i.i.d. samples denoted as  $(X_{n})_{1\leq n\leq N}$  from a standard Gaussian distribution of dimension  $d = 10$ . Each corresponding target  $Y_{n}$  is obtained by applying a predefined function  $f^{\star}$ , referred to as the teacher network, to the input  $X_{n}$  (i.e.,  $Y_{n} = f^{\star}(X_{n})$ ). We choose  $f^{\star}$  to be a one-hidden layer network of the form in (10) with  $M^{\star} = 5$  hidden units  $(v_{i}^{\star},u_{i}^{\star})_{1\leq i\leq M^{\star}}$  drawn independently from a standard Gaussian distribution. Furthermore, we consider two nonlinearities  $\gamma$  when defining  $f^{\star}$ : the widely used ReLU activation [37] for the main results and its smooth approximation SiLU for additional ablations in Appendix C.3. This choice of target function  $f^{\star}$  results in a hidden low-dimensional structure in the regression task as  $f^{\star}$  depends only on 5-dimensional linear projection of the input data. In most cases, we take the size of the training data to be  $N = 500$  except when studying the effect of the training size. This choice allows us to achieve a balance between conducting an extensive hyper-parameter search and being able to compute precise GN updates on the complete training data within a reasonable timeframe. Finally, we generate 10000 test samples to measure the generalization error.
Model. We consider a well-specified setting where the model  $w \mapsto f_w$ , referred to as the student network, is also a one-hidden layer with the same activation function  $\gamma$  as the teacher network  $f^{\star}$ . Importantly, the student network possesses a number  $M = 5000$  of hidden units, denoted as  $w := (v_i, u_i)_{1 \leq i \leq M}$ , which is much larger than the teacher's number of units  $M^{\star} = 5$  and allows fitting the training data perfectly (over-parameterized regime). Following [15], we also normalize the output of the student  $f_w$  by  $M$  to remain consistent with the mean-field limit as  $M$  increases.
Initialization. We initialize the student's hidden units according to a centered Gaussian with standard deviation (std)  $\tau_0$  ranging from  $10^{-3}$  to  $10^{3}$ . Varying the initial std  $\tau_0$  allows us to study the optimization dynamics in two regimes: the kernel regime (large values of  $\tau_0$ ) and the feature learning regime (small values of  $\tau_0$ ). Finally, we initialize the weights of the last layer to be 0. This choice allows us to perform a systematic comparison with the minimum norm least squares solution obtained using random features as described next.
Methods and baselines. We consider three optimization methods for the model's parameters: a regularized version of Gauss-Newton (GN), gradient descent (GD), and optimizing the linear layer parameters alone, which can be viewed as a random features (RF) model.
(GN): We use the discrete GN updates in (4) with a constant step-size  $\lambda$  and  $H_{w} = I$ . Each update is obtained using Woodbury's matrix identity by writing  $\Phi(w_{k}) = J_{w_{k}}^{\top}z_{k}$  with  $z_{k}$  the solution of a linear system  $(J_{w_{k}}J_{w_{k}}^{\top} + \epsilon(w_{k})I)z_{k} = \nabla \mathcal{L}(f_{w_{k}})$  of size  $N$ . Here, we use the damping defined in (5) with  $\alpha = 1$  and ensure it never falls below  $\epsilon_{0} = 10^{-7}$  to avoid numerical instabilities. In practice, we found that small values of  $\epsilon_{0}$  had little effect on the results (see Figure 4 (Right) of Appendix C.1).
(GD): The model's parameters  $w$  are learned using gradient descent with a constant step-size  $\lambda$ .
(RF): Instead of optimizing the entire parameter vector  $w = (v, u)$  of the model  $f_w$ , we selectively optimize the parameter  $v$  of the linear layer while keeping the weights  $u$  of the hidden layer constant. This procedure corresponds to computing a minimal-norm least squares solution  $v^{RF}$  for a random features (RF) model, where the features are obtained from the output of the hidden layer at initialization. Specifically, the solution  $v^{RF}$  is obtained as
$$
v ^ {R F} = \Gamma (u) \left(\Gamma (u) ^ {\top} \Gamma (u)\right) ^ {\dagger} \mathbb {Y}, \tag {8}
$$
where  $\Gamma(u) = (\gamma(u^\top X_n))_{1 \leq n \leq N} \in \mathbb{R}^{M \times N}$  is the feature vector computed over the training data, and  $\mathbb{Y}$  is a vector of size  $N$  consisting of the target values  $Y_n$ . The choice of the un-regularized solution as a baseline is supported by recent findings [22], demonstrating its good generalization in the over-parameterized regime.
Stopping criterion. For both (GD) and (GN), we perform as many iterations as needed so that the final training error is at least below  $10^{-5}$ . Additionally, we stop the algorithm whenever the training error drops below  $10^{-7}$  or when a maximum number of iterations of  $K^{GD} = 10^6$  iterations for (GD) and  $K^{GN} = 10^5$  iterations for (GN) is performed. For (RF) we solve the linear system exactly using a standard linear solver.
5.2 Performance metrics
In addition to the training and test losses, we introduce two complementary metrics to assess the quality of the hidden features  $\Gamma(u)$  learned by the student model.
Weighted cosine distance (WCD). Measuring proximity between the student and teacher's hidden weights  $(u_{1},\dots,u_{M})$  and  $(u_{1}^{\star},\dots,u_{M^{\star}}^{\star})$  requires being able to find a correspondence between each student's parameter  $u_{i}$  and a teacher's one  $u_{j}^{\star}$ . When using a positively homogeneous non-linearity such as ReLU, only the alignment between vectors is relevant. However, since the student's model is vastly over-parameterized, there could be vanishing weights  $u_{i} = 0$  which do not influence on the network's output, while the remaining weights are well aligned with the teacher's weights. We introduce the following weighted distance to account for these specificities:
$$
\operatorname {W C D} \left(u, u ^ {\star}\right) = 2 \sum_ {i = 1} ^ {M} p _ {i} \left(1 - \max  _ {1 \leq j \leq M ^ {\star}} \frac {u _ {i} ^ {\top} u _ {j} ^ {\star}}{\| u _ {i} \| \| u _ {j} ^ {\star} \|}\right), \quad p _ {i} = \frac {\left\| u _ {i} \right\| ^ {2}}{\sum_ {k = 1} ^ {M} \| u _ {k} \| ^ {2}}. \tag {9}
$$
Equation (9) finds a teacher's parameter  $u_{j}^{\star}$  that is most aligned with a given student's parameter  $u_{i}$  and downweights its cosine distance if  $u_{i}$  has a small norm. In practice, we found this measure to be a good indicator for generalization.
Test loss after linear re-fitting (Test-LRfit). To evaluate the relevance of the hidden features  $\Gamma(u)$ , we train a new linear model  $\tilde{v}$  on those features and measure its generalization error. In practice, we freeze the weights  $u$  of the hidden layer and fit the last layer  $v$  using the procedure (8) described above for the random features (RF) baseline. In our over-parameterized setting, this new model should always achieve perfect training loss, but its generalization error strongly depends on how the features were learned.
5.3 Numerical results
Implicit bias of initialization. Figure 1(Left) shows that GN and GD generalize more or less well compared to a random features model depending on the variance  $\tau_0$  of the weights at initialization. First, in the case of the RF model, changing  $\tau_0$  barely affects the final training and test error. This is essentially due to the choice of the non-linearity ReLU which is positively homogeneous. Second, for large  $\tau_0$ , the final test error of both GN and GD matches that of RF suggesting that the dynamics were running under the kernel regime/azy regime [15, 23] where the final layer's weights are learned without changing the hidden features much. Finally, for small values of  $\tau_0$ , both GN and GD obtained a better test error than RF which can only be explained by learning suitable hidden features. These results are further confirmed when varying the size of the training set as shown in Appendix C.2. In Appendix C.3, we performed the same experiment using SiLU non-linearity instead of ReLU and observed the same transition between the kernel and feature learning regimes. While prior works such as [11] analyzed Gauss-Newton in the kernel learning regime, our results indicate that Gauss-Newton can also exhibit a feature learning regime for small values of  $\tau_0$ .
Implicit bias of the step size. Figure 1(Right) shows that increasing the step size in Gauss-Newton results in features that do not generalize well. This is unlike gradient descent where larger step sizes yield better-performing features<sup>3</sup>. This behavior is first observed in the top figures showing an increasing weighted cosine distance (WCD) between the student's hidden weights and the teacher's weights as the step size increases in the case of Gauss-Newton. On the contrary, the distance decreases with larger step sizes in the case of gradient descent, indicating that the algorithm learns better features. This effect is also confirmed when computing the test loss after refitting the linear layer to the training data exactly while keeping the learned hidden weights fixed (see linear re-fitting in Section 5.2). Several works, such as [3, 36], highlight the importance of taking large step-sizes
3The step size can be taken as large as possible provided the algorithm does not diverge. In our experiments, we found that going beyond a step size of  $\lambda = 10^{3}$  for GD results in a divergence of the algorithm
Figure 1: Final values of various metrics vs std of the hidden layer's weights at initialization  $\tau_0$  (left) and step size (right) for a ReLU network. (Left figure) The training size is set to  $N = 500$  while  $\tau_0$  ranges from  $10^{-3}$  to  $10^{2}$ . 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}$ . (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$ ).
for gradient descent, to improve generalization, our results show that Gauss-Newton benefits instead of using smaller step-sizes. As shown in Appendix C.3, this behavior persists when using a different non-linearity, such as SiLU with different values for  $\beta$  (1 and  $10^{6}$ ). Interestingly, we observed that GN might underperform GD as the parameter  $\beta$  defining SiLU decreases making SiLU nearly linear and less like ReLU. In all cases, while the final test loss remains large for small step sizes (in both GD and GN), the test loss after linear re-fitting is much lower in the case of Gauss-Newton, indicating that the algorithm was able to learn good features while the last layer remained under-optimized.
Hidden feature learning. Figure 2 (Left) shows the evolution of various metrics, including the WCD, the training, and test errors with the number of iterations for both GN and GD. For each method, results for the best step size and variance at initialization are presented  $(\tau_0,\lambda) = (10^{-3},10^2)$  for GD and  $(\tau_0,\lambda) = (10^{-3},10^{-2})$  for GN). The bottom figure clearly indicates that the early iterations of GN essentially optimize internal features first while barely improving the last layer. This can be seen deduced from the training and test losses which remain almost constant at the beginning of optimization, while the test loss after linear re-fitting steadily decreases with the iterations. The last layer is learned only towards the end of training as indicated by a faster decrease in both train and test objectives. Gradient descent displays a different behavior with test loss before and after linear re-fitting remaining close to each other during optimization.
Better generalization of GN comes at the cost of a slower training. Figure 2 (Right) shows the evolution of the same quantities as in the left figures for  $\tau_0 = 10^{-3}$  as a function of time (in seconds) and for the three values of the step-size  $\lambda$  ( $10^{-2}$  (bright colors) 1 (medium colors) and  $10^{2}$  (dark colors)). The figure shows that, while the training loss converges much faster for larger step sizes in the case of GN, it does not result in the best generalization error as measured by the test loss after linear re-fitting. Hence, better generalization comes at the cost of a slower convergence of the training loss. Again, the behavior is quite different in the case of gradient descent for which larger step sizes result in both faster optimization of the training objective and better generalization.
The results illustrate the importance of the initialization and choice of the step size for the generalization performance of GN when optimizing over-parameterized one-hidden layer networks. While
Figure 2: Evolution of various metrics during training for both GD and GN. Both the variance of the weights at initialization  $\tau_0$  and step-size  $\lambda$  are selected from two sets  $\{10^{-3}, 1, 10^2\}$  and  $\{10^{-2}, 1, 10^2\}$ . In all cases, we report results for the best-performing choice of  $\tau_0$  (achieved for  $\tau_0 = 10^{-3}$ ) based on test error after linear re-fitting (Test-LRfit). Both Top-Left and Bottom-Left figures show the evolution of the WCD defined in (9) and the three metrics (the training loss, the test loss and Test-LRfit) for both GD and GN using the step-size  $\lambda$  achieving the lowest Test-LRfit. Top-Right and Bottom-Right figures show the evolution of WCD, train loss and test loss after linear re-fitting for the various choices of step-size (darker colors correspond to larger step-sizes).
our theoretical analysis shows that GN can reach a global solution of the training objective when starting from a close-to-optimal initialization, this result does not imply anything about the generalization properties of the obtained solution or the quality of the learned features. Our empirical study instead shows that GN may favor feature learning (thus yielding improved generalization) at the cost of a slower optimization of the training objective due to the use of a smaller step size. Providing a theoretical analysis of these phenomena is an interesting direction for future work. Finally, our study shows that the training or test error may not always be a good indicator of the quality of the learned features due to an under-optimized final layer. Instead, a test/Validation error after re-fitting the linear layer may be a better indicator for the quality of the learned features which can be of practical use.
The project leading to this work has received funding from the French National Research Agency (ANR-21-JSTM-0001) in the framework of the "Investissements d'avoir" program (ANR-15-IDEX-02). This work was granted access to the HPC resources of IDRIS under the allocation 2023-AD011013762R1 made by GENCI.
[1] Z. Allen-Zhu, Y. Li, and Z. Song. “A convergence theory for deep learning via overparameterization.” In: International Conference on Machine Learning. PMLR. 2019, pp. 242-252.
[2] S.-i. Amari. "Natural Gradient Works Efficiently in Learning." In: Neural Computation 10.2 (Feb. 1998), pp. 251-276. URL: https://www.mitpressjournals.org/doi/10.1162/089976698300017746 (visited on 08/22/2019).
[3] M. Andriushchenko, A. Varre, L. Pillaud-Vivien, and N. Flammarion. "SGD with large step sizes learns sparse features." In: arXiv preprint arXiv:2210.05337 (2022).
[4] M. Arbel, A. Gretton, W. Li, and G. Montufar. "Kernelized Wasserstein Natural Gradient." In: (Sept. 2019). URL: https://openreview.net/forum?id=HkIz71rYvS.
[45] I. Sutskever, J. Martens, G. Dahl, and G. Hinton. "On the importance of initialization and momentum in deep learning." en. In: International Conference on Machine Learning. Feb. 2013, pp. 1139-1147. URL: http://proceedings.mlr.press/v28/sutskever13.html.
[46] A. B. Tsybakov. Introduction to Nonparametric Estimation. en. Springer Series in Statistics. New York: Springer-Verlag, 2009. URL: https://www.springer.com/gp/book/9780387790510.
[47] P. Vicol, J. Lorraine, D. Duvenaud, and R. Grosse. "Implicit Regularization in Overparameterized Bilevel Optimization." In: ICML 2021 Beyond First Order Methods Workshop. 2021.
[48] S. Wojtowytsch. "On the convergence of gradient descent training for two-layer relu-networks in the mean field regime." In: arXiv preprint arXiv:2005.13530 (2020).
[49] B. Woodworth, S. Gunasekar, J. D. Lee, E. Moroshko, P. Savarese, I. Golan, D. Soudry, and N. Srebro. "Kernel and rich regimes in overparametrized models." In: Conference on Learning Theory. PMLR. 2020, pp. 3635-3673.
[50] N. Yudin and A. Gasnikov. "Flexible modification of gauss-newton method and its stochastic extension." In: arXiv preprint arXiv:2102.00810 (2021).
[51] C. Yun, S. Krishnan, and H. Mobahi. “A unifying view on implicit bias in training linear neural networks.” In: arXiv preprint arXiv:2010.02501 (2020).
[52] G. Zhang, J. Martens, and R. B. Grosse. "Fast convergence of natural gradient descent for over-parameterized neural networks." In: Advances in Neural Information Processing Systems 32 (2019).
A Background: NTK vs mean-field limits, kernel vs feature learning regimes
We provide a brief discussion of the notions of NTK and mean-field limits as well as the related convergence regimes for a gradient flow. Given a non-polynomial point-wise and differentiable activation function  $\gamma$  and some positive integer  $M$ , consider the following one-hidden network:
$$
f _ {w} (x) = \alpha (M) \sum_ {i = 1} ^ {M} v _ {i} \gamma \left(u _ {i} ^ {\top} x\right), \quad w = (v, u) \in \mathcal {W} := \mathbb {R} ^ {M} \times \mathbb {R} ^ {M \times d}, \tag {10}
$$