Rethinking Gauss-Newton for learning over-parameterized models
Authors: Michael Arbel
Romain Menegaux
Pierre Wolinski
Venue: NeurIPS 2023 poster (2023)
Keywords: implicit bias
gauss newton
TL;DR: global convergence and generalization of gauss-newton methods for over-parameterized one-hidden layer networks

Abstract: This work studies the global convergence and implicit bias of Gauss Newton's (GN) when optimizing over-parameterized one-hidden layer networks in the mean-field regime. 
We first establish a global convergence result for GN in the continuous-time limit exhibiting a faster convergence rate compared to GD due to improved conditioning. 
We then perform an empirical study on a synthetic regression task to investigate the implicit bias of GN's method.
While GN is consistently faster than GD in finding a global optimum, the learned model generalizes well on test data when starting from random initial weights with a small variance and using a small step size to slow down convergence. 
Specifically, our study shows that such a setting results in a hidden learning phenomenon, where the dynamics are able to recover features with good generalization properties despite the model having sub-optimal training and test performances due to an under-optimized linear layer. This study exhibits a trade-off between the convergence speed of GN and the generalization ability of the learned solution.

Rethinking Gauss-Newton for learning over-parameterized models
Michael Arbel, Romain Menegaux; and Pierre Wolinski*t
Univ. Grenoble Alpes, Inria, CNRS, Grenoble INP, LJK,38000 Grenoble, France firstname. lastnane@inria.fr
Abstract
This work studies the global convergence and implicit bias of Gauss Newton's s (GN) when optimizing over-parameterized one-hidden layer networks in the mean-field regime. We first establish a global convergence result for GN the continuous-time time limit exhibiting a faster convergence rate compared to GD due improved conditioning. We then perform an empirical study on a synthetic regres- sion task to investigate the implicit bias of GN's method. While GN is consistently faster than GD in finding a global optimum, the learned model generalizes well on data when starting from random initial weights with a small variance and using small step size to slow down convergence. Specifically, our study shows that such setting results in hidden learning phenomenon, where the dynamics are able to recover features with good generalization properties despite the model having sub-optimal training and test performances due an under-optimized lin- ear layer. This study exhibits a trade-ofi off between the convergence speed of GN and the generalization ability of the learned solution.
Introduction
Gauss-Newton (GN) and related methods, like the Natural Gradient (NG) can offer faster conver- gence rates compared to gradient descent due to the improved dependence the conditioning of the problem [9, 21, 30, 31]. For these reasons, these methods have attracted attention in machine learning as alternative gradient descent when optimizing ill-conditioned problems arising from the use of vastly over- pece-paraettrizeed networks and large training sets [29, 45]. Unfortunately, GN's high computational cost per iteration, which involves solving expensive linear system, restricts its applicability in large-scale deep learning optimization. Addressing this challenge has been a primary focus, with extensive efforts dedicated crafting computationally efficient approximations for GN/NG methods where the aim is to strike a delicate balance between computational cost and optimization speed, thereby enhancing scalability [4, 11, 21, 31, 42].
While substantial effort has been put into making GN scalable, little known about the global convergence of these methods and their generalization properties when optimizing neural networks, notably compared to gradient descent. Understanding these properties is particularly relevant know- werrparametterize networks for which
ing that state-of-tt of-the-art machine learning methods rely over infinitely many solutions interpolating the data exist albeit exhibit varying generalization proper- ties that depend on implicit bias of the optimization procedure 39]. To date, the majority of studies on generalization in pver-parameterized networks have primarily focused on analyzing the dynamics of gradient descent and gradient flows their continuous- -imee limit making impor- tant progress in understanding the implicit bias of these optimization procedures. Specifically, [15]
"Equal contribution.
Now affiliated with the LMO, Université Paris-Saclay, Orsay, France.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

showed that gradient flow can run under two regimes which yield solutions with qualitatively dif- ferent generalization properties: the kernel regime, where the internal features of network remain essentially constant during optimization, often yields poor generalization, and the feature learning regime, which allows the internal features of network to adapt to the data, seems behind the im- pressive generalization properties observed in practice. In addition, the same work precisely defines conditions for these regimes: The kernel regime systematically occurs, for instance, when taking the Neural Tangent limit (NTK), i.e. scaling the output of network of width M by a factor 1/vM while the feature learning regime can arise when instead taking the mean--feld limit, i.e. scaling the output by a factor of 1/M. However, these insights rest upon attaining a global solution, which is far from obvious given the non-convexity of the training objective. While global linear convergence of gradient flows is known in the kernel regime [1, 23], obtaining such convergence guarantees in the feature learning regime remains an active research area [8. 10, 12, 14]. Importantly, all these works have focused solely on gradient flows and gradient descent, given their prevalent usage when optimizing neural networks.
In contrast, GN/NG methods have received limited attention in the study of both generalization and global convergence for over wvrr-parameterized networks. The added complexity of these methods makes their analysis intricate, particularly in the stochastic setting where inaccurate estimates of the pre-conditionnn matrix can alter the dynamics in highly non-trivial manner [32, Section 12.2]. Recent works studied the convergence and generalization of GN in the kernel regime [11, 20, 26, 52], while others focused on linear networks [27]. However, the behavior of these methods for non- linear over-paue parameterized models remains understudied in the feature learning regime which is a regime of interest it is more likely to result in good generalization.
Contributions. The present work aims to deepen our understanding of the properties of Gauss- Newton (GN) methods for over-parameterized networks, building upon recent advancements in neural network optimization [15]. Unlike prior works that primarily addressed scalability issues of GN/NG or studied the impact of stochastic noise optimization dynamics, we uniquely focus on studying the convergence and implicit bias of an exact GN method. Specifically, our research investigates GN in the context of regression problem using over wee-parameteriz one-hidden layer networks in a deterministic (full-batch) setting. By examining an exact, deterministic GN method, rather than approximate stochastic version, we eliminate potential regularization effects stem- ming from approximation and stochasticity that might influence convergence and generalization. This strategy, in turn, enables to provide both theoretical and empirical insights, shedding light on the properties of GN methods when optimizing neural networks:
1. We provide a global convergence result in Proposition 3 with a linear rate for the continuous-timg limit of GN's dynamics in the over-paranne parameterized regime that holds under the mean-fi mean-field limit. The achieved rate shows better conditioning than the gradient flow with the same objective, emphasizing the method's advantage from a pure optimization perspective. While similar rates can be achieved under strong non-degeneracy conditions as in [50], these conditions are never met over- wer--aramettriieed networks. Instead, our result relies parameters of a similar nature as in [13, Thm 3.3]: (i) that the hidden weights are diverse enough so that an interpolating solution can be obtained simply by fitting the linear weights, and (ii) that the linear weights result from a simple fine-tuning procedure to reach near optimality. The result shows that GN can reach the data interpolation regime, often required for studying the implicit bias of optimization (Section 4).
simple condition the initial
2. The empirical study in Section 5 complements the convergence result in Proposition 3 by further investigating the implicit bias of GN on student/teacher regression task using synthetic data. We show that, while GN attains an interpolating solution, as predicted by Proposition 3, such a solution generalizes more or less well depending on the regime under which the GN dynamics is running. Specifically, we show that the GN dynamics can exhibit both kernel and feature learning regimes depending on the choice of step size and variance of the weights at initializa- tion. Our experiments also indicate that a hidden learning phenomenon occurs in GN dynamics where hidden features that generalize well are obtained even when both raw training and test performances are sub-optimal due to an under-optimized linear layer. Quite surprisingly, these features are found when using small step sizes and, sometimes, generalize better than those learned by gradient descent. This is contrary to the common prescription of using larger step sizes in gradient descent for better generalization [3]. Our results suggest a tradeoff between optimization speed and generalization when using GN which is of practical interest.

---

2 Related work
Convergence of Gauss-Newion's methods. A large body of work on inverse problems studied the local convergence of Gauss-Newton and its regularized versions such as the Levenberg-Maree Marquardd method [25, 38]. Many such works focus on continuous-timg analysis of the methods it allows for a simplified study [24, 41]. In this work, we also consider continuous-time dynamics and leave the study of riscrete-time algorithms to future work. More recently, increasing attention was given to global convergence analysis. [50] studied the global convergence of GN with an inexact oracle and in stochastic setting establishing linear convergence. The analysis requires that the smallest singular value of the Jacobian is positive near initialization, a situation that can never occur in the over-per perrraaamtteeeeed regime. [34] studied the convergence of the Levenberg-Marquuee
dynamics, a regularized version of Gauss-Newton's 's method, in a non-convex setting under a cubic growth assumption and showed global albeit sub-linear convergence rates. is unclear, however, if such cubic growth assumption holds in our setting of interest. Closest to our work is [11] which aims to accelerate optimization by solving kernel regression in the NTK scaling limit. There, the authors show global convergence of a regularized version of GN in the kernel regime. We aim to investigate the behavior of GD the feature learning regimes.
The implicit bias of gradient flows. In the over--arameee parameterized setting, the training objective often possesses several global minima that perfectly interpolate the training data. The generalization error thus heavily depends on the solution selected by the optimization procedure. [39] highlighted the importance of understanding the properties of the selected solutions for given initialization and several works studied this problem in the case of linear neural networks [6. 33, 40, 49, 51]. For non- linear networks, there is still no general characterization of the implicit bias the regression setting, although [10] recently made important progress in the case of one-hidden layer network with ReLUU activation and orthogonal inputs showing that gradient flows select a minimum norm interpolating solution when the initial weights are close to zero. Recently, the implicit bias of gradient flows was shown to play a crucial role in other optimization problems such as non-convex bi-level optimization [5, 47], therefore illustrating the ubiquity of such phenomenon in deep learning. The present work empirically studies the implicit bias of Gauss-Newton which is vastly understudied.
3 Setup and preliminaries
3.1 Regression using over-parameterized one-hidden layer networks
We consider regression problem where the goal is to approximate an unknown real-valued function f* defined over a subset X' of Rd from a pair of i.i.d. inputs/outputs (Øn, VnlIgnsN where each on is a sample from some unknown distribution P. and yn. S (Ix,). We assume that I* belongs to (I)), the set of square-integrable functions w.r.t. P and note that f* always belongs L2(P),, where P denotes the empirical distribution defined by the samples (21...., .NN......N................ .NN). For simplicity, we assume the data are non-degenerate, meaning that Tn ADw whenever n * n'. In this case (2(P) is isomorphic to R- (i.e., (PP) = R-N) and we can identify any function f E L2(P)) with the evaluation vector (f(x1). (f(x))...... --.... (xN)). We interested in approximating /* using one-hidden layer network fw with parameter vector w belonging to some, possibly infinite, Hilbert space W. The network's parameters are learned by minimizing an objective of the form C(fw) where C is a non-nega,ive, L-smooth, and μ-strongly convex function defined over (zP)) and achieving a minimum value at ʃ* A typical example that we consider in Section 5 is the mean-squared error over the training set:
N Σ (fm( ynn ==1
min C(fw). WEW
C(fw)) = 2N
(1)
For the convergence results, we will consider both finite-width one-hidden layer networks and their mean-field infinitewwidth limit. In the experiments, we will restrict to finite-width networks although, we will be using relatively large number of units M compared to the size of the training data, therefore approximating the meean-field limit.
3

---

Finite-width one-hidden layer networks. Given non-polynomial point-wise activation function y and some positive integer M, these networks take the form:
M Σ
1 fa( M
un(u, x),
w (v,u) : I Rxd,
where v E RM the linear weight vector while E RM the hidden weight matrix. Popular choices for y include ReLU [37] and its smooth approximation SiLU(z) xx1 (-)= which can be made arbitrary close to ReLU by increasing / [19]. The above formulation can also account for bias term provided the input vector z is augmented with non-zero constant component.
Mean-field inftnitewwidth limit. We consider some base probability yr measure with full support on A' and finite second moment and denote by L2( and L2 the set of square integrable
functions w.r.t. and taking values in I and Rd. Given a non-polynomial point-wise non- non-linearity, we define infinitely-win -ide one-hidden layer networks fw as:
(x) = v(c)r(u()) T2) du(c),
w = ( ;; L2(n) x R²).
(3)
Functions of the form (3) correspond to the mean-field limit of the network defined in (10) when the number of units M grows to infinity and appear in the Lagrangian formulation of the Wasserstein gradient flow of infinitely wide networks [48]. Existing global convergence results of GN methods were obtained in the context of the NTK limit [11] which does not allow feature learning. The mean- field limit we consider here can result in optimization dynamies that exhibit feature learning which is why we consider it in our global convergence result analysis. For completeness, we provide a brief discussion on mean-field NTK limits and kernel feature learning regimes in Appendix A.
We consider an over- peer-parameteriee setting where the network has enough parameters to be able to fit the training data exactly, thus achieving 0 training error. To this end, we define the gram matrix G which is an x matrix taking the following forms depending on whether are using smite-width network (G*) infinite-wied width network (G'):
ever-parameterizane
G (u)
1 Σ ru, zn)Yuu knr), (u) f /(u(e)] Tn))(u(c)TTr a,)) du(c). M
We say that network is over-parameterrized if G(uo) is invertible for some hidden-layer weight uo:
(A) (Over-parameterizatii G(uo) is invertible for an initial parameter wn=(1o, E.
When uo is sampled randomly from a Gaussian, common choice to initialize neural network, Assumption (A) holds for infinitely wide networks as long as the training data are non-degenerate (see [18, Theorem 3.1] for ReLU, and [17, Lemma F.1] for analytic non-polynomial activations). There, the non-degeneracy assumption simply means that the inputs are not parallel to each other (L.e. 2, If xy whenever * j for any < i,j N), property that holds almost surely if the data distribution IP has a density w.r.t. Lebesgue measure. In the case of finite-width networks with M > N, the result still holds with a high probability when u is sampled from Gaussian [28].
3.2 Generalized Gauss-Newton dynamics
To solve (I), we consider optimization dynamics based on a generalized Gauss-Newton method. To this end, we denote by Ju the Jacobian of (f=( fuo (x(()) w.r.t. parameter w which can be viewed as a linear operator from W to R.,. Moreover, VL(fm) denotes the vector of size N representing the gradient of L w.r.t. the outputs Gauss-Newton vector field & : W + W defined as:
Su We can then introduce the
(ww) (J Ha + --
where H. a symmetric positive operator on (P) and E(w) is an optional non non-negative (possibly vanishing) damping parameter. Starting from some initial condition wo E W and for given positive step-size y, the Gauss-Newton updates and their continuous-tig limit are given by:
Wh++1 = W/ - >
# = - -(w.).
(4)

---

The continuous-tiin time limit is obtained by taking the step-size y to O and rescaling time appropri- ately. When H is given by the Hessian of C. the dynamics (4) recovers the generalized Gauss- Newton dynamics in continuous time when no damping is used [44] and recovers the continuous- time evenherg-Maquardt dynamics when a positive damping term is added [35]. When the matrix H. is the identity H. = I, the resulting dynamics is tightly related to the natural gradient [2, 30]. More generally, we only require the following assumption on Hw:
(B) He is continuously differentiable in w with cigenvalues in [μH, LH] for positive Lg, PH
Optional damping. Amongst possible choices for the damping, we focus on E(w) of the form:
=(w) = (w)(w)
(5)
where a is a non-negative number 0, and 0² (w)) is the smallest eigenvalue of the Neural Tangent Kernel (NTK) A.. Ju which is invertible whenever J. is surjective. This choice allows to study the impact of scalar damping on the convergence rate of the GN dynamics. While computing of (w) is challenging in practice, emphasize that the damping is only optional and that the case when no damping is used, i.c. when a = 0, is covered by our analysis.
4 Convergence analysis
4.1 Global convergence under no-blow up
We start by showing that the continuous-time dynamics (4) converge to a global optinuu provided that it remains defined at all times.
Proposition Let wh = (Vo; un) E W be an initial condition so that un satisfies Assumption (A). Under Assumption (B) and assuming that the activation is twice-continuously differentiable, there exists a unique solution to the continuuu-
dynamics in (4) defined up to some maximal, possibly infinite, positive time T. TT is finite, we say that the dynamics blows-up in finite time. If. instead is infinite, then fw. converges to a global minimizer of Moreover, denoting by μ the strong convexity constant of C and defining PGN == 24/(( + a/Pu)) ), it holds that:
((f S. [((mo)e-uoNt,
(6)
Proposition 1 is proved Appendix B.2. It relies on the Cauchy-Lipshin theorem to show that the dynamics are well-defined up to positive time T and uses strong convexity of L to obtain the rate in (6). The result requires a smooth non-linearity, which excludes ReLU. However, holds for any smooth approximation to ReLU, such as SiLU. As we discuss in Section 5, this does not make large difference practice as the approximation becomes tighter. The rate (6) is only useful when the dynamics is defined at all times which is far from obvious as the vector field & can diverge in finite time causing the dynamics to explode. Thus, the main challenge is to find conditions ensuring the vector field & remains well-behaver which is the object of Section 4.2.
constant jr of C, the smoothness constant of H. and damping strength with the fastest rate achieved when no damping is used, i.e. a O. For instance, when choosing H the identity and a 0, the linear rate becomes μGN = 2p. This is in contrast with gradient flow of w + L(fm) for which linear convergence result, in the kernel regime, follows from [15, Theorem 2.4]:
Comparaison with gradient flow. The rate in (6) only depends on the
L(fw.) ) S C(fwo) )e
with MGF := por (w where (w)) is the smallest singular value of the initial NTK matrix Awo- In practice, PGF & PGN since the linear rate per for gradient flow is proportional to a awo)) which can get arbitrarily small as the training sample size N increases thus drastically slowing down convergence [43].
4.2 Absence of blow-up for almost-optimal initial linear layer
In this section, we provide a condition on the initialization win = (vo, un) E W Newton dynamics never blows up. We start with simple result, with proof in Appendix B. showing the existence of a vector so that w* := (v", up) W interpolates the training data whenever the over-parameterrizaton assumption Assumption (A) holds for given ug.
ensuring the Gauss-
5

---

Proposition 2. Let wo = (vo, uo) E W be an initial condition so that uy satisfies Assumption (A). Define T(u)= = (7(4² xx)) ISNSN and F == (J)( (xx)) and let o² be given by:
v¹ = =u))))
Then the vector w² = (v*, (w", un) E W is a global minimizer of w + ((s()) and a fixed point of the Gauss-Newton dynamics, i.e. ((w=)) = 0.
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 opti- mization algorithm such as gradient descent since the function + 1(v) = ((,(),)))) converges to a global optimum as it satisfies a Lojssa--.ojssieiicz/ inequality (see Proposition 4 Appendix I.... However, this procedure is not of interest here as it does not allow us to learn the hidden features. Nonetheless, since w' = (u) is a fixed point of the Gauss-Newton dynamies, we can expect that an initial condition wo (vo: uo). where Uo is close to v², results in well-behaved dynamics. The following proposition, proven in Appendix B.4, makes the above intuition more precise.
Proposition 3. Let wp (Vo, up) E W be an initial condition so that Assumption (A) holds, i.e. the gram matrix G(ue) is invertible, and denote by of the smallest eigenvalue of G(uo). Assume that the activation 7 is wiiceconttnuously differentiable and that Assumption (B) holds. Let R be any arbitrary positive number and define CR = sup (B operator norm. If the linear layer is almost optimal Vg in the sense that:
R) 1w0l where II-Ilop denotes the
[V <,
::= (upnPaN /8LN) min( R, Cn²) min (an,
(7)
then (4) is defined all times, i.e. T= + co, the objective [(fw.) converges a linear rate according to (6) and the parameters remain within ball of radius R centered around Wo-
Proposition 3 essentially states that the dynamics never blow up, and hence, converge globally 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 un typically holds for a Gaussian initialization when the data are non-degenerate. Additionally, the near-optimality condition on the linear weights vo can always be guaranteed by optimizing the linear weights while keeping the hidden ones fixed.
Remark 1. In the proof of Proposition 3, show that the occurrence of finite-time blow-up is tightly related to the Neural Tangent Kernel matrix Aw, becoming singular at increases which causes the vector field d to diverge. When the NTK matrix Aw, at time t is forced to remain close to the initial one Aw,. as is the case in the NTK limit considered in [11], one can expect that the singular values of Aur, remain bounded away from O since this is already true for Awg. 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 Au, in the mean-field limit, where Aw, is allowed to differ significantly from the initial NTK matrix Awg. The final step is to deduce that, while the weights wy are allowed be in ball centered in wn of arbitrarily large radius R, the singular values remain bounded away from 0 provided the initial linear weights up 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 We to be within a ball of fixed radius Ry determined by initial smallest singular value ow)) of Awe, result allows w, to be within arbitrary distance R regardless of the initial a (wo) 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- peeepaaramteri networks. Since we observed little variability when changing the

---

seed, we have chosen to use 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.I. 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.
5.1 Experimental setup
We consider regression task on synthetic dataset consisting of N training points (X,: Snsss The objective is to minimize the mean-squared error L(fw), defined in (1), over the parameters w of model fw that predicts the target values Yn based on their corresponding input points X...
Data generation. We generate N i.i.d. samples denoted as (X,)isng: from a standard Gaussian distribution of dimension d = 10. Each corresponding target Y, is obtained by applying predefined function f", referred to as the teacher network, to the input X,, (i.e., Y,, = (X)).. We choose f to be one-hidden layer network of the form in (10) with M 5 hidden units (u, uz IS-SsM drawn independently from a standard Gaussian distribution. Furthermore, we consider two non- linearities & when defining f": 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 results in a hidden low-dimensional structure in the regression task as f 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 achieve balance between conducting an extensive
search and being able to compute precise GN updates on the complete training data within reasonable timeframe. Finally, we generate 10000 test samples to measure the generalization error.
Model. We consider well-specified setting where the model fw, referred as the student network, is also a one-hidden layer with the same activation function 7 as the teacher network f² Importantly, the student network possesses number M 5000 of hidden units, denoted as w: w: (e;, IiSS. which is much larger than the teacher's number of units M = 5 and allows fitting the training data perfectly pavattarraeteriized regime). Following [15], we also normalize the output of the student S.e by M to remain consistent with the mean-field limit as M increases.
Initialization. We initialize the student's hidden units according
centered Gaussian with stan- dard deviation (std) To ranging from 10-3 3 to 10 Varying the imitial std To allows us to study the optimization dynamics in regimes: the kernel regime (large values of Tn) and the feature learn- ing regime (small values of TD). Finally, we initialize the weights of the last layer to be O. This choice allows us perform 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: 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 and H. = I. Each update is obtained using Woodbury's 's matrix identity by writing b ((wx)=I,, 2 with 24 the solution of linear system (Jw, JT, +e(ww)))=x
of size N. Here, we use the damping defined in (5) with a=1 and ensure never falls below 6o 10 7 avoid numerical instabilities. In practice, we found that small values of 6n had little effect on the results (see Figure 4 (Right) of Appendix C.I). (GD): The model's parameters w are learned using gradient descent with a constant step-size X. (RF): Instead of optimizing the entire parameter vector w (v. u) of the model fw. we selectively optimize the parameter of the linear layer while keeping the weights u of the hidden layer con- stant. This procedure corresponds to computing a minimal-norm least squares solution o of for random features (RF) model, where the features are obtained from the output of the hidden layer at initialization. Specifically, the solution R. is obtained as
y RF ((u)((())TT()))
(
where I(w) = (7(wT Xn))IsHSN E M is the feature vector computed over the training data, and W is vector of size N consisting of the target values Y,. The choice of the un-regularized solution as a baseline is supported by recent findings [22], demonstrating its good generalization in the over-paramettrrzed regime.
7

---

Stopping criterion. For both (GD) and (GN), we perform many iterations as needed so that the final training error is at least below 10 Additionally, we stop the algorithm whenever the training enror drops below 10-7 or when maximum number of iterations of K KGD 10 iterations for (GD) and KGN 10 iterations for (GN) 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 metries to assess the quality of the hidden features I'(u) learned by the student model.
Weighted cosine distance (WCD). Measuring proximity between the student and teacher's hidden weights (#1,..,, uAr) and (ui, "iz-) ) requires being able to find a correspondance between each student's parameter uy and teacher's one uj. When using a positively homogeneous non-linearity such as ReLU, only the alignment between vectors is relevant. However, since the student's model vastly over-paramettrizze, there could be vanishing weights u, 0 which do not influence on the network's 's output, while the remaining weights are well aligned with the teacher's weights. We introduce the following weighted distance account for these specificities:
M
WCD(u, u²) = 2
pi 1- max 1SJSM
||
P.
M Kk=1
|
Equation (9) finds a teacher's parameter u3 that most aligned with a given student's parameter u; and downweights its cosine distance if u; has 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 I(w), we train a new linear model 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. 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 ((Left) shows that GN and GD generalize more or less well compared a random features model depending on the variance To of the weights at initialization. First, in the case of the RF model, changing To barely affects the final training and test error. This essentially due to the choice of the non-linyarity ReLU which is positively homogeneous. Second, for large To, the final test eтror of both GN and GD matches that of RF suggesting that the dynam- ies were running under the kernel regime/azzy regime [15, 23] where the final layer's weights are learned without changing the hidden features much. Finally, for small values of To, both GN and GD obtained a better test error than RF which can only explained by learning suitable hidden features. These results are further confirmed when varying the size of the training set shown in Appendix C.2. In Appendix C.3, we performed the same experiment using SiLU non-lineen 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
can also exhibit feature learning regime for small values of TD-
Implicit bias of the step size. Figure (Right) shows that increasing the step size in
results in features that do not generalize well. This is unlike gradient descent where larger step sizes yield better-performe perforrien features'. This behavior first observed in the top figures showing an in- creasing weighted cosine distance (WCD) between the student's 's hidden weights and the teacher's weights as the step size increases in the case of Gauss-Newton. On the contrary, the distance de- creases 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- re-fitting in Section 5.2). Several works, such as [3, 36], highlight the importance of taking large step-sizes
'The step size can be taken as large as possible prosiddee provided the algorithm does not diverge. In our experiments, found that going beyond step size of À = 10 for GD results in a divergence of the algorithm
8

---

WCD/Loss vs std at initialization
WCD/Loss vs step size (std=0.001)
10a
WCD
WCD
10-1
10:
10
10"4
10
100
10
10a
10-4
10-1
10-2
10-1
160
101
102
107
10
10"2
Loss
10-4
Loss
10
10a4
10-4
loss (Linear re-fitting)
GD
Test loss (Linear re-fitting)
Test loss
GN
10a4
Test loss
10.4.
Train loss
RF
10a
10
10-3
10a2
10
Std of the initial weights
Step size A
Figure 1: Final values of various metrics vs std of the hidden layer's weights at initialization To (left) and step size (right) for a ReLU network. (Left figure) The training size is set N = 500 while To ranges from 10-3 to 102. For both GD and GN, results are reported for the best-performing step-size selected according to the test loss regular logarithmic grid ranging from 10-3 10 (Right figure) The std of the weights at initialization is set to To 10-3 All models are optimized up to a training error of 10 or until the maximum number of steps exceeded, (M 5000 N 500).
for gradient descent, to improve generalization, our results show that Gruss-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 3 (1 and 10%). Interestingly, GN might underperform GD as the parameter 3 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 Gauss-Newton indicating that the algorithm was able to learn good features while the last layer remained under-optimized.
observed that
Hidden feature learning. Figure (Left) shows the evolution of various metries, 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 initialization are presented ((To. À) (10= for GD and (Tu, A) = (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 learned only towards the end of training 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- rittitgn remaining close to each other during optimization.
Better generalization of GN comes the cost of slower training. Figure (Right) shows the evolution of the same quantities as in the left figures for To 10-3 as function of time (in seconds) and for the three values of the step-size (10-2 (bright colors) (medium colors) and 102 (dark colors)). The figure shows that, while the training loss converges much faster for larger step sizes 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 the case of gradient descent for which larger step sizes result in both faster optimization of the training objective and better generalization.
6 Conclusion
The results illustrate the importance of the initialization and choice of the step size for the general- ization performance of GN when optimizing over-parameterized one-hidden layer networks. While

---

Loss/WCD Iteartions
Loss/WCD vs Time (s) for different À
10a
10a
GN
Step size A
WCD
WCD
GD
10
1e0
10
10
10
10
10'
30
107
10a
10
201
10'
1+2
10a
10"4
Loss
10
Loss
10-4
GD
GN
lee
Step size
10-1
Train loss
10:4
10-1
Test loss
Train loss
Test loss (Linear re- -etfiiting)
10-
Test loss (Linear re- re-tiinnss)
1e-2
10
10'
10a
10
10-1
30
207
10
Iterations
Time (s)
Figure 2: Evolution of various metrics during training for both GD and GN. Both the variance of the weights at initialization to and step-size À are selected from two sets {10-3 1, 102) and 1, (,0-,,1,10% 102]. In all cases, we report results for the best-performing choice of 7n (achieved for TO = 10 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 metries (the training loss, the test loss and Test-LRfit) for both GD and GN using the step-size À achieving the lowest Test-LRfit. Top-Right and Bottom-Righe figures show the evolution of WCD, train loss and 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 close-to-optimal initialization, this result does not imply anything about the general- ization 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 theoretical analysis of these phenomena is an interesting direction for future work. Finally, study shows that the training or test error may not always be a good indicator of the quality of the learned features due the linear layer may be a better indicator for the quality of the learned features which can be of practical use.
under-optimized final layer. Instead, test/validation error after re-fitting
Acknowledgments
The project leading to this work has received funding from the French National Research Agency (1R-21--SSMM-000) in the framework of the "Investissements d' d'avenir" program (ANR-15- IDEX-02). This work was granted access to the HPC resources of IDRIS under the allocation ND2-AAD01001362H made by GENCI.
References
[1] Z. Allen-Zhu, Li, and Z. Song. "A convergence theory for deep learning via over- parameterization." 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: h/wp::///www..// http:://wwww.wniesjssseoooooo 089976698300017746 (visited on 08/22/2019).
org/doi//0/ 1162/
[3] M. Andriushchenko. A. Varre, L. Pillaud-Vivien, and N. Flammarion. "SGD with large step sizes learns sparse features." In: arXiv preprint arXiv:2210.055 (2022).
[4] M. Arbel, Gretton, W. Li, and G. Montufar. "Kernelized Wasserstein Natural Gradient."
In: (Sept. 2019). URL: https:///// hwt.er/wwww..ee/toommm/....c.//
10

---

[5] M. Arbel and J. Mairal. "Non-Conver Bilevel Games with Critical Point Selection Maps." In: Advances in Neural Information Processing Systems (NeurIPS) 2022 (2022).
[6] S. Arora, S. Du, W. Hu, Z. Li, and R. Wang. "Fine-grained analysis of optimization and gen- eralization for overparameterized two-layer neural networks." In: International Conference on Machine Learning. PMLR. 2019, pp. 322-332.
[7] M. Belkin, A. Rakhlin, and A. B. Tsybakov. "Does data interpolation contradict statistical optimality?" In: The 22nd International Conference on Artificial Intelligence and Statistics. PMLR. 2019, pp. 1611-1619.
[8] G. Blanc, N. Gupta, G. Valiant, and P. Valiant. "Implicit regularization for deep neural net- works driven by an ornstein-uhlenber. like process." In: Conference on learning theory. PMLR. 2020, pp. 483-513.
[9] A. Botev, H. Ritter, and D. Barber. "Practical gauss-newton optimisation for deep learning." In: International Conference on Machine Learning. PMLR. 2017, pp. 557-565.
[10] E. Boursier, I. Pillaud-Vivien, and N. Flammarion. "Gradient flow dynamics of shallow ReLU networks for square loss and orthogonal inputs." " In: arXiv preprint arXiv:2206.000 00939 (2022).
[11] T. Cai, R. Gao, J. Hou, S. Chen, D. Wang, D. He, Z. Zhang, and L. Wang. "Gram-gauss- newton method: Learning overparameterized neural networks for regression problems." In: arXiv preprint arXiv: 1675 (2019).
[12] Z. Chen, E. Vanden Vanden-Eiinden, and Bruna. "On feature learning in neural networks with global convergence guarantees." In: arXiv preprint arXiv:2204. J0782 (2022).
[13] L. Chizat. "Sparse Optimization on Measures with Over pver-parrmettriieed Gradient Descent." URR: http: arxiv.
In: ar0iv 10300 [math, stat] (July 2019). arXiv: 1907. org/ abs/1907 10300.
[14] L. Chizat and F. Bach. "On the Global Convergence of Gradient Descent for Over- parameterized Models using Optimal Transport." May 2018. URL: https / hal archives-ouvertee
fr frrffr//eeeeee wtrrhal-000 0111-9199
[15] L. Chizat, E. Oyallon, and F. Bach. "On lazy training in differentiable programming." Advances in Neural Information Processing Systems 32 (2019).
[16] L. Deng. "The mnist database of handwritten digit images for machine learning research." In: IEEE Signal Processing Magazine 29.6 (2012), pp. 141-142.
[17] S. Du, J. Lee, H. Li, L. Wang, and X. Zhai. "Gradient descent finds global minima of deep neural networks." In: International conference on machine learning. PMLR. 2019. pp. 1675- 1685.
[18] S. S. Du, X. Zhai, B. Poczos, and A. Singh. "Gradient descent provably optimizes over- parameterized neural networks." In: arXiv preprint wrXiv:1800.02000 (2018).
[19] S. Elfwing, E. Uchibe, and K. Doya. "Sigmoid-Weiee Weighted Linear Units for Neural Net- work Function Approximation in Reinforcement Learning. ArXiv e-prints (2017)." In: arXiv preprint arXiv: 03118 (2017).
[20] J. R. Garcia, F. Freddi, S. Fotiadis, M. Li, S. Vakili, A. Bernacchia, and G. Hennequin. "Fisher-Leggende (FishLeg) optimization of deep neural networks." In: The Eleventh Interna- tional Conference Learning Representations. 2023.
[21] R. Grosse and J. Martens. "A cronecker-factorfe Approximate Fisher Matrix for Convolution Layers." In: Proceedings of the 33rd International Conference on International Conference on Machine Learning Volume 48. ICML'16. event-place: New York, NY, USA. JMLR.org, acm. org/citation cfm?id=3045390 3045452. [22] T. Hastie, A. Montanari, S. Rosset, and R. J. Tibshirani. "Surprises in high-dimensional ridge- less least squares interpolation." In: The Annals of Statistics 50.2 (2022). pp. 949-986.
2016, pp. 573-582. URL: http://t
[23] A. Jacot, Gabriel, and C. Hongler. "Neural Tangent Kernel: Convergence and Gener- alization in Neural Networks." In: arXiv: 10l.::880007577 fcs, math, stat] (June 2018). arXiv: 1806.07572. URL: hattp://arxia www. abs/ 1806. 07572.
[24] B. Kaltenbacher. A. Neubauer, and A. Ramm. "Convergence rates of the continuous regu- larized Gauss Newton CeussnNNewton method." In: Journal of Inverse and III- ill-Posed Problems 10.3 (2002), pp. 261-280.

---

[25] B. Kaltenbacher, A. Neubauer, and O. Scherzer. "Iterative regularization methods for nonlin- ear ill-posed problems." In: Iterative Regularization Methods for Nonlinear Ill-Posed Prob- lems. de Gruyter, 2008.
[26] R. Karakida and K. Osawa. "Understanding approximate fisher information for fast conver- gence of natural gradient descent in wide neural networks." In: Advances in neural informa- tion processing systems 33 (2020), pp. 10891-10901.
[27] A. Kerekes, A. Mészáros, and Huszár. "Depth Without the Magic: Inductive Bias of Natural Gradient Descent." In: arXiv preprint arXiv:2111.115:5 (2021).
[28) C. Liu, L.. Zhu, and M. Belkin. "Loss landscapes and optimization in wver-parameterized non- linear systems and neural networks." In: Applied and Computational Harmonic Analysis 59 (2022), pp. 85-116.
[29] J. Martens. "Deep Learning via Hessian-free Optimization." In: Proceedings of the 27th In- ternational Conference on International Conference on Machine Learning. ICML' 10. event- place: Haifa, Israel. USA: Omnipress, 2010, pp. 735-742. URL: http: : //dl acm org/ citation. cfm7id-3104322. (tdd--1143222.1144411
[30] J. Martens. "New insights and perspectives on the natural gradient method." In: The Journal
of Machine Learning Research 21.1 (2020), pp. 5776-5851.
[31] J. Martens and R. Grosse. "Optimizing Neural Networks with Fronecker-factored Approx- imate Curvature." In: arXiv: /503.0567 [es, stat] (Mar. 2015). arXiv: 1503.05671. URL.: attp://arxixs org/ aaap//50..055.. 05671.
[32] J. Martens and I. Sutskever. "Training Deep and Recurrent Networks with Hessian Opti- mization." en. In: Neural Networks: Tricks the Trade: Second Edition. Ed. by G. Montavon, G.B. Orr, and K.-R. Müller. Lecture Notes in Computer Science. Berlin, Heidelberg: Springer Berlin Heidelberg, 2012, pp. 479-535.
URL: https: /doi. org/ 10. 1007 / 978 3-
35289-8 8______7
[33] H. Min, S. Tarmoun, R. Vidal, and E. Mallada. "On the explicit role of initialization the convergence and implicit bias of overparametrized linear networks." In: International Con- ference on Machine Learning. PMLR. 2021, pp. 7760-7768.
[34] K. Mishchenko. "Regularized Newton Method with Global O(1/k * +2) Convergence." In: arXiv preprint arXiv:21 12.02089 (2021).
[35] J. J. Moré. "The Levenberg- Marquardt algorithm: implementation and theory." In: Numerical Analysis: Proceedings of the Biennial Conference Held at Dundee, June 28-July 1, 1977. Springer. 2006, pp. 105-116.
[36] R. Mulayoff, T. Michaeli, and D. Soudry. "The implicit bias of minima stability: A view from function space." In: Advances in Neural Information Processing Systems 34 (2021),
pp. 17749-17761
[37] V. Nair and G. E. Hinton. "Rectified linear units improve restricted boltzmann machines." In: Proceedings of the 27th international conference machine learning (ICML-10). 2010, pp. 807-814.
[38] Y. Nesterov. "Modified Gauss-Newton scheme with worst case guarantees for global perfor- mance." In: Optimisation methods and software 22.3 (2007), pp. 469-483.
[39] B. Neyshabur, R. Tomioka, and Srebro. "In search of the real inductive bias: On the role of implicit regularization in deep learning." In: arXiv preprint arXiv: Nel::411.6666 441.66144444 (2014). and N. Flammarion. "Implicit bias of sgd for diagonal linear networks: a provable benefit of stochasticity." In: Advances in Neural Information Processing Systems 34 (2021), pp. 29218-29233
[40] S. Pesme, L. Pillaud-
[41] A. G. Ramm. "Dynamical systems method for solving operator equations." In: Communica- tions in Nonlinear Science and Numerical Simulation 9,4 (2004), pp. 383-402.
[42] Y. Ren and D. Goldfarb. "Efficient subsampled gauss-newton and natural gradient methods for training neural networks." In: arXiv preprint arXiv: 02353 (2019).
[43] D. A. R. Robin, K. Scaman, and M. Lelarge. "Convergence beyond the over-parameterized regime using Rayleigh quotients." In: Advances in Neural Information Processing Systems. Ed. by A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho. 2022. URL:
[44] N. N. Schraudolph. "Fast curvature matrix-vector products for second-order gradient de- scent." In: Neural computation 14.7 (2002), pp. 1723-1738.
12

---

[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. http:htttp//poeeeiiiiiee
13.hmll. html. [46] A. B. Tsybakov. Introduction Nonparametric Estimation. Springer Series in Statis- tics. New York: Springer- Verlag, 2009. URL: https: / www. springer com / gp/ gp/book/
2013, pp. 1139-1147. URI.: http: / /tpp//oeedings..... mmlr press//28/sutttkerrr13 httss/sssseeerr
9780387790510.
[47] P. Vicol, J. Lorraine, D. Duvenaud, and Grosse. "Implicit Regularization in Overparame- terized 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. 2000:0000..3330 (2020).
[49] B. Woodworth, S. Gunasekar, J. D. Lee, E. Moroshko, Savarese, I. Golan, D. Soudry, and
N. Srebro. "Kernel and rich regimes in overparametrizee 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.008 000.000000/000 (2021).
[51] C. Yun, 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, Martens, and R. B. Grosse. "Fast convergence of natural gradient descent for werrparameterizee neural networks." " In: Advances in Neural Information Processing Systems 32 (2019).
13

---

A Background: NTK vs mean-field limits, kernel vs feature learning regimes
We provide brief discussion of the notions of NTK and mean-field limits well the related convergence regimes for a gradient flow. Given a non-polynomial point-wise and differentiable activation function y and some positive integer M, consider the following one-hidden network:
fu(x)= =
aMM) a
Σ uy
w (v, W = x 1x
(10)
where (M) is a scaling factor that depends on the number of neurons M. For simplicity, we assume that the weights (1, u,) are i.i.d. samples from a Gaussian. We are interested in the dynamics of a gradient flow when M increases let's ss introduce some notations. We denote by F(w) the vector of evaluation of the network on the training data F(w) = (fw (z1)...... fu.(:) and by Ju the Jacobian of F(w) w.r.t. parameter w. Moreover, we introduce the NTK matrix Aw == Jw which is N N symmetric positive semi-definite matrix. Finally, (L(fw) denotes the vector of size N representing the gradient of L w.r.t. Fww). The gradient flow of an objective ((Fw) is given by:
+00 and for specific choices of the scaling factor (MM). First,
ي/ -u, ((F.(F)).
Additionally, it is often informative to track the time evolution of the vector Fw,, which is given by the following ODE:
OrFw,, =-w [
= Aw. VFC(F.).
(11)
One can see from (11) that the NTK matrix Aw, arises naturally when expressing the dynamics of the function Fu. We can even provide an explicit expression of Aw, using the NTK kernel, which is a positive semi-definite kernel K (x, x²) defined by the inner product:
M
(w = Bw fz)) a(M)² Σ (1(4) x)r( x²) (v:))// 2) x za)xTa))
The NTK matrix Aw is then formed by collecting the values (K B,))ss on the training data. It is easy to see that when a(M) _ then K (z, converges to an expectation under the distribution of the weights (u:, u))iggs under mild conditions. In this case, Aw depends on the choice of the distribution of the weights (v,, us)ssi M and not on the particular samples. On the other hand, when a(M) M. is the function fwr itself that converges an expectation under the distribution of the weights. This distinction will lead to different qualitative dynamics for the gradient flows as we discuss next.
VM
A.1 The Neural Tangent Kernel limit and the kernel regime (a(M)
VM
When choosing a(M) = VM and setting M + +oc, it is shown in [23] that the NTK matrix A., remains constant in time and depends only the distribution of the weights wo at initialization. Hence, Fu,, is effectively the solution of the ODE:
O,F, = Awo VFF(F).
F
which converges at linear rate to a global minimizer of L whenever Aw is invertible and L is strongly convex. The NTK limit results in kernel regime, where the resulting solution is equivalent to the one obtained using kernel method with the initial NTK kernel Kwo (I, x') [15, 23].
A.2 The mean-field limit and the feature learning regime (a(M)
37
When choosing instead a(M) = M. it is possible to view the limiting function as an expectation over weights drawn from some probability distribution p
fp(w) f vo(u" x) u)(
In [48], it further shown that such 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 As, (and NTK kernel Kw,) is allowed to evolve in time. This is a limit in which feature learning is possible since the dynamics do not necessarily simplify to kernel method.
14

---

B Proofs
Notations. For an operator A between two Hilbert spaces, denote by ||Allop its operator norm and by ||A|] p its Frobenius norm and (A) be its smallest singular value whenever it well-defined. Denote by B(w, R) the open ball of radius R centered at element wo in a Hilbert space.
B.1 Over-parameterization implies interpolation.
Proof of Proposition 2. Denote by F
First, is easy to see that the matrix
O.f is exactly equal the N x N matrix G(uu) considered in Assumption (A), which is assumed to be invertible. Moreover, fw is independent of v since f(o,o)) is linear in e. Hence, we simply choose:
v² (Hof()))))) (DoI()))))) (Oo (O0 Oe F (Oof(),a)) "G(uo)-FF"
Moreover, using again the lincarity of f in v, we have that
f(u).wo) = Öof(v" "N))"" = G(uo)G(uo) We have therefore shown that w un) a global minimizer of w + L(Sw). Finally, since VL(fw)) = V((ft)= =0 it follows that odw") = O.
F*.
Proposition 4. Assume that Assumption (A) holds for some vector Then, the function e + ((v) L(f(e,o)) is convex and satisfies a Polyak-Lojasiewicc inequality:
[ 2 2po L(f*)).
with μ the strong convexity constant of and of the smallest singular value of the matrix G(
appearing in Assumption (A).
Proof. It is easy to see that I(v) is convex as composition of strongly convex function and linear function v + [(u,No)- We will show now that / satisfies Polyak- Lojassewj (PL) inequality. To this end, we make the following calculations:
|[ V.L(fu)
C(fw).
(12)
is easy to see that the matrix 0, I is exactly equal to the N x N matrix G(un) considered in
Assumption (A), which is assumed be invertible. Hence, its smallest singular value of is positive and the following inequality holds:
(.)(2i) 2 <f
(13)
On the other hand, since is strongly convex and differentiable, then satisfies the following PL inequality:
, 2p(((sm) L(S*)),
(14)
where y is the strong convexity constant L. Hence, we get the desired inequality by combining (13) and (14):
[V > ( L(
B.2 Global convergence under no blow-up: Proof Proposition
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 activation function 7 as we will be using it in many places.
The non-linearity y is twice-continuously differentiable.
As first step, we need to prove local existence and uniqueness of the solution to the ODE in (4) which results from Cauchy-Lipschitz theorem applied the vector field To this end, a key step is to show that & is locally Lipschitz near initialization when the Jacobian Jwn is surjective.
15

---

Proposition 5 (Regularity of Φ.). Assume (B) and (C) and that wo satisfies Assumption (A). Then, there exists neighborhood of wo so that (ww) locally Lipschitz in w.
Proof of Proposition 5. By application of the Woodbury matrix identity, & can be equivalently writ- ten as:
((w)= = J. J. JJ + ((W) H (L
(15)
Moreover, by the smoothness assumption, Ju. Hw. <(w) and V V[((w) are all differentiable in w. Additionally, since Juo is surjective, it must hold that o > 0 in a neighborhood of we so that Jw remains invertible. Finally, since Hu is always positive, it holds that (Jw. + H(") and H are differentiable in neighborhood of wo- Therefore, 4 is Locally Lipschitz in neigh- borhood of wp.
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 (wx, )220 to (4) defined up to a, possibly infinite, maximal time T > 0.
Proposition 6 is a direct consequence of the field ((w) is locally Lipschitz in in neighborhood of wh by Proposition 5. The above local existence result does not exclude situations where the dynamics blow up in finite time T < +o0 and possibly fail globally minimize the objective L. We will show, later, that finite-time blow-ups never occur under additional assumptions on the initial condition.
theorem which holds since the vector
Proof of Proposition I. Existence and uniqueness of the continuous-tims dynamics up to maximal time T > 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; in what follows and will write b v.( By differentiating the objective in time, we get the following expression:
a[[((f..] = -b JJJJJT HJ ++1+ - Jbb
= -b- H (JJT + + H
= - -bT + =b H 1 H (JJ ( H H
-
< -T 4
1 - 8 H 3 + H
< DTH
t (HD
2n=(Lto/ux)-
S S Lnnnn-LLllll < 2np L(fw:).
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 e(w), the fact that H s L I and the fact that |/l ||V[(II 2 μL(() by p-strong convexity of L. This allows to directly conclude that w...)) S C(fwo)e - with pGN := 2
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 Aw. := Jan, JT For any arbitrary positive radius R > 0, define the constant CR and stopping time
TR as follows:
Cn = sup
11 llap < +00
Tk = sup 220 |||, woll < R, and 0² "AAm)) > 07,
16

---

The stopping time Ta characterizes the smallest time when the dynamics becomes either degenerate (o* (A..) or grows too far away from initialization (w, B( R)).
To control the singular values of Awo we find it useful to consider the dynamics of the Frobenius norm the pseudo-inverse J which we denote by aj := J
Indeed, simple linear algebra provides the following lower-bound on o² o*Aww,) in terms az:
o² a'(Aw,) 2 at"
(16)
Note that, while the space of parameters can even be infinie-diimeniionll (mean-field limit), Jw has a finiteddimensional range and thus always admits a pseudo-inverse Jh given by:
J J (JWJJ
The next proposition provides a differential inequality for a = 10 |r in terms of || | that holds for all times smaller than Tn-
Proposition 7. Assume (B) and (C) and that we satisfies Assumptiin (A), then for any R R> O, the time TR is positive and smaller than the maximal time T, Lee. 0 TR S T. Moreover, for any t &E [0, TR). at satisfies:
l S CRCRleilll
Proof. Positive stopping time Tn- By construction, TR must be smaller than T it requires we to be well-defined. Moreover, Jw, is surjective by assumption and J.. continuous in w, therefore Ju, must also be surjective for small enough positive times t 0. Additionally, it must hold that ||w; woll - R for small positive times, by continuity of t + wt, This allows shows that T'k > 0.
Dynamics of J By smoothness of Ju it holds that Jh, differentiable on the open interval (0, TR) and satisfies the following ODE obtained by direct differentiation:
J = -f (I P.).T (J.J.)
(17)
In (17), we introduced notation Jj := Ju, and its time derivative J, for simplicity and denote by Pt. the projector: Pt J." Controlling | Taking the time derivative of and recalling the evolution (17) yields:
orå; = -J(J)) J J.J.
where we used that (J! T (1- = by Lemma 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 in Appendix B.5), we directly get:
aclãel s (JI) 1 snoo! cill[illl
(18)
By the chain rule, we have that Je = Owoer therefore, | S 1e@ Moreover, since é < TR. then it must holds that we E B(wo, It) so that 11owJellop S CR. Henceforth, | F S Crlli and (18) yields the desired inequality after dividing by a, which remains positive all times t [0, Tn) since Jr does not vanish.
The next proposition provides an estimate of the time derivatives [ and | up to time Tn. For simplicity, introduce the notation A. = [[VV[((,,,
Proposition 8. Assume (B) and (C) and that we satisfies Assumption (A). Then, for any R > 0, the time TR is positive and smaller than the maximal time T up which (4) is defined, i.e. < TR <
T. Moreover, define C == Lp PH An Then and |[ satisfy the following differential inequality on [0, TR):
| S Cafe "GN"
|| < CC R a,e
HaN
17

---

Proof of Proposition 8. By application of Proposition 7 to (wp)ogi<T, we directly have that Ta is positive and smaller than and the following estimate holds for any TRi
ldel S anlliell
It remains to upper-bound ||
Controlling || Ignoring the dependence on we in the notations for simplicity and setting B : : H.J, we have, by definition of the dynamical system (4):
(J [JTJJ + (
| |l S
JJ JJJJ + H +
(BBB +6) - - --
BB (BB) + + - |
1 S SH | [[ PH -
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 Appendix B.5.
Concluding. We combine the upper-bound on |åg| and | to get:
l S HH CRaA A..
Finally, since < T, it follows from Proposition that A, S 1 Aoe "q which allows to con- clude.
The next proposition controls a, and l|o, woll in terms of TR. CR. and other constants.
Proposition Consider the same conditions as Proposition 8. Let Ç the (possibly infinite) time T defined as:
AL parPaN and consider
T -2uaN log (1 (SCnApag))
with the convention that T = to0 if1 2 (CRAuag Then for any time t < T = mim T*) holds that:
MOO 1)) 1))=
lw: woll < ÇAgao,
at S ao (1+ (CRAaaa
(19)
Proof of Proposition 9. First, by Proposition 8, we know that satisfies the following differential inequality:
làel S Cafe- -FN
with = L CR. We will control a, by application of Grönwall's inequality. To this end, ииn consider the ODE:
by a0 0. We know by Lemma 5 in Appendix B.5 that by is given by:
b. Ob.e "@N
egN: -1) - b. bo 4Cpa (e-0@((()))))
-v (1 (1+ of -gx+ (6-0-000
for all times t < T Therefore, by Grönwall's inequality, it must hold that as b, for all times I < minnTn, T+). Moreover, recalling now that ||xee|| < Lp- " Anage "GNi we directly get that:
- by Proposition 8,
C || 5 -age age "GN: 19p2(1+ 4C (0-00())))11)
(20)
CR
18

---

Integrating the above inequality in time yields the following estimate on ||0, woll for any t < min(Tn. T)
/ ||ū0||1 ds
/
|| woll 1
1 (1 4C " -1)
CRaO
3 AA(ppHppN) = ÇAoeo,
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 Anag is small enough, a quantity that depends only on the initial conditions. Moreover, it also shows that making the product Agag small enough at initialization ensures that w, remains in a ball of radius R around wo. We exploit these two consequences to prove Proposition 3 in Appendix B.4.
of the J. remains bounded at all times provided that
B.4 Absence of blow-up for almost-optimal initial linear layer: proof of Proposition
To prove Proposition 3, we rely the results of Appendix B.3 which allow controlling the evolution of the smallest singular value of the NTK matrix Aw,. More precisely, Proposition 3 is a direct consequence of Proposition 9 with particular choice for the initialization Wo-
Proof of Proposition 3. By assumption on Wp. it holds that
D VL( )ll <
min R, C min (do, oō).
8LN
Moreover, by definition of =
JJ and of oi min omin(G(uu)) we know that ag < No
Hence, we get
min R. C min ( - a0 ²)
Ao
2
where we introduced Ç :
41 рииKON
This above inequality directly yields:
R
(Aoan
(CRApag <1,
(21)
2N
Therefore, we can use the above inequalities in the estimates (19) of Proposition 9 to get:
R
ww woll
a S aa (CnD)aa))
2
for any time E [0, Tn), where Tn is defined as:
TR = sup {t 2 0 [|w, woll < R. and o² (Ap,, > 0}.
This implies. in particular, that we never escapes the ball B(wo, R). Moreover, using (16) get that the smallest singular value d'(AT, is always lower-bounded by positive constant = un (1 (C))) Therefore, TR must necessarily be greater or equal to T, the maximal time over which w, is defined. Additionally, if was finite, then wy must escape any bounded do- main. This contradicts the fact that W, E B(wg, R). Therefore, must be infinite and the incquality (6) applics at all times > 0.
B.5 Auxiliary results
Lemma Let J. be a linear operator from a Hilbert space to a finite-dimensiinn spuce H. and assume that JJJ is invertible. Define the projector P J JT (JJJ)-1) Then the following relations hold:
(J+) P
19

---

Lemma 2. Let J. be linear operator from a Hilbert space W to minite-immensooal space H. and assume that JJJ is invertible. Then for any non son-negative number e, it holds that:
JJ (JJ + + 5
Lemma 3. Let W be Hilbert space and R be a finite-dimentn dimensional euclidean space. Let be linear operator from W R and H an invertible operator from R to itself. Further, assume that JJ is invertible. Then the following holds for any non non-negative number E.
+JJJ+CHHT - S e
Lemma 4. Let and K be two Hilbert-Schmid operators between two Hilbert spaces. Then, holds that:
|| S |1l s ||
Lemma's 5. Let C and r be two positive constants, For given initial condition bo > 0, consider the following differential equation:
by = -t
Then by is defined for any time & < T where T is given by:
T 1--
(1-2(1))
2Cr Otherwise.
to0
Moreover, by is given by:
b, 20r -1))
t t<T
Proof. First note that a local solution exists by the Cauchy-Lipsctee thcorem. Morcover, it must never vanish since otherwise, it would coincide with the null solution by 0 by uniqueness. How- ever, this impossible, since by > 0. We can therefore divide the ODE by bg and explicilly integrate it which gives:
b 2=²= 1 2r 2r
The solution is defined only for times t so that the r.h.s. is positive, which is exactly equivalent to having T
C
Additional experimental results
C.1 Multi-seed experiments
We present result obtained for 5 independent runs. Each uses 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 obtained using SiLU non-linearity. It is clear from Figure 3, that the results display little variability w.r.t. the seed the experiment, which justifies the single-seed setting chosen in the main paper.
Effeet of over- parameterization Figure 4 (Left) shows the effect of over-paramettrizztio of network on generalization error. appears that generalization error remains stable with an increasing parameterization for both GN and GD, soon the network has enough over- parameterization to the data exactly. On the other hand, the test error of RF improves with increasing prreaaaaamterzaiionnn which is consistent with [22].
Evolution of the smallest singular value of A.,. Figure 4 (Right) shows that the smallest singu- lar value a² (A.:) systematically increases during training, especially the feature learning regime (bright red colors). Hence, the dynamics are far from blowing up even the features change sub- stantially. On the other hand, a² (A,,) remains nearly constant for large values of To (darker colors) which indicates that the features barely change in the kernel regime.
20

---

Test loss at convergence
Test loss at convergence
Loss
Loss
10
10
Random Features
10"4
Random Features
1-HL: Gradient Descent
1-HL: Gradient Descent
1-HL: Gauss-Newton
1-HL: Gauss-Newtor
10
10-1
10a
10t
10a
10'
10a Training sample size N
Initial std of the weights
Figure 3: (Left) Test error vs std of the initial weights TO- All models are trained until the training objective is smaller than 10-6 using M = 5000 hidden units and NN = 500. Confidence interval 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 10-6 using M = 5000 hidden units and initial std of the weights To 1. Confidence interval at 95% estimated from independent runs shown in shaded colors.
Test loss at convergence
10
10-1
10.0
Loss
10-0
Random Features
5,0
1-HL: Gradient Descent
1.0 To
1-HL: Gauss-Newton
10-1
20.5
10-1
10-7
0.01
2x
x 10
4x
10a
102
102
10a
10a
10
10F
10'
Number of neurons/Rfs: M
Training iterations
Figure 4: (Left) Test error number of hidden units, All models are trained using N = 2000 training samples until the training objective is smaller than 10 The initial std of the weights is set to = (Right) Evolution of o (AAw) during training using Gauss-Newton for different values of the initial std To of the weights. All models are trained using the training objective is smaller than 10-6 We used M 5000 units. Confidence interval 95% estimated from independent runs shown in shaded colors.
500 training samples until
C.2 Effect of the sample-size
Figure 5 shows the evolution of the test error 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, display the results for two different values for the initial std (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 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 Ctest and training sample size in a logarithmic scale with strong dependence of the slope on the variance initialization and the optimization method. Such affine relation implies the following upper-bound on Liest in terms of sample-size N:
Liest S No.
for some positive constants and a. The coefficient a controls the speed which the estimator ap- proximates fr as accesses more training data and usually appears in learning theory to describe
21

---

the statistical convergence rate of given estimator [46]. The consistent and strong dependence of on the optimization methods (GN GD) and regime (kemel feature learning) further confirms the implicit bias of both initialization and optimization algorithms.
Loss vs training sample size ('ReLU")
Loss vs training sample size ('SiLU')
10
10a
10
10
4
10
GD (T. =1
GD GD 1)
Lo55
(T, 100
(To 10
1
10
(to 1)
(To
10
GN (To = 10-2)
GN (To = 10-²)
RF (To 1)
RF (To 1)
(T==10-9)
(To=10-9)
10
10"
10
10a 10
10a
10a
10'
Training sample size N
Training sample size N
Figure 5: Final values of various metries training sample size for both ReLU (left) and SiLU (right) networks. Two values for To are displayed for each method To E {10-3, 3, 1 while training data ranges from N 100 to N 1500. All models are optimized until the training objective smaller than 10-6 using M = 5000 hidden units. For both GD and GN, results are reported for the best-performing step-size À selected according to the test loss on regular logarithmic grid ranging from 10 to 10
C.3 Choice of the non-linearity
Figure 6 shows the effect of step-size the final performance of a SiLU network with two choices for the temperature parameter: / = 1 and 3 = 100 Similarly to Figure 1, it 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 3, we observe that GN either outperforms GD (3 10%) the opposite (3 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 8 10 we recover almost the same results as those obtained for ReLU in Figure I (right). This is due to the fact SiLU becomes a tight smooth approximation ReLU for large values of We also notice that when using larger std for the initial weights (ex. 1), using either ReLU or SiLU with (B = 1) yields similar results, as shown in Figure 7.
C.4 Experiments on MNIST
We have performed series of experiments MNIST to illustrate the behavior of GN higher dimensional data (notably, with multidimensional targets). To stay as close as possible to the theoret- ical framework, modify the original MNIST dataset. Specifically, construct student/teacher setup in which the trained network can achieve a zero-loss, which is an assumption made in Sec- tion 3.1.
Building datasets based on MNIST by using teacher. In our theoretical framework, we stricted ourselves to cases where the minimum loss can be achieved on the considered model. To make this condition hold, we build new dataset by using teacher trained on the MNIST classifi- cation task [16]:
1. we construct training set for the teacher which consists Der where Deo is the orig- inal MNIST test dataset and Der 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 Drr U Dro using cross-entropy objective;
22

---

WCD/Loss Vs step size (beta=1)
WCD/Loss vs step size (beta= 1e6)
109
WCD
WCD
10-1
10
10
10
10
10:1
107
101
107
107
10
10-1
10
107
10a
10
10-
Loss
L055
10
Test loss (Linear re-fitting)
GD
Test loss loss- re
GD
Test loss
GN
Toss
GN
Train loss
RF
10-4
Train loss
RF
10
10
10a²
10
10
Step size À
Step size À
Figure 6: Final values of various metrics vs the step size for both SiLU network with 8 = (left) and / 10 (right). (Right figure) The std of the weights initialization is set to To 10 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).
WCD/Loss vs step size (std=1.0)
WCD/Loss step size (std=1.0)
105
1.05x100
9x
10*
8x10-1
9.5
10
10a
107
Logs
10
Loss
10"4
10-4
10-1
10
Test loss (Linear fittinng)
GD
10:4_
Test loss (Linear re-fiting)
GD
Test loss
GN
Test loss
GN
10
Train loss
10-a
Train loss
RF
10-4
10
10'
Step size A
Step size À
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 = 1. All models are optimized up to training error of 10-6 or until the maximum number of steps exceeded, (M 5000 N 500).
3. given the trained teacher fr, we build new datasets D.r and Dis for the student network:
D : {I Softmmax/((((x)))) (2x, vx) Dir). D,s { (v., Soflmmax((x(())))) (Ii, vx) Dro).
Training procedure. Once the pair of training and testing datasets (Der. Des) has been built, we use them to train and test student neural network based on (10), with one hidden layer of size 5000 and ReLU activation function. We use similar initialization scheme for the student's network and fix the std of the initial weights to TQ 0.001. For training, we use a quadratic loss between the student and teacher's outputs and perform maximum of 100000 iterations of either GN or GD. The training stops whenever the training loss goes below 10-6 6 We vary the step size on a logarithmic range [10 [10-3, 100) for both methods and evaluate the test loss and test loss after linear re- re-fiiting.
23

---

3x
Loss
10a
GN
Test loss (Linear re-fitting)
GD
10
loss
10 3
10
10-1
10
10
102
20
Step size
Figure 8: Final values of various metries vs the step size for ReLU network trained using MNIST data. Std of the initial weights is set to 0.001. For GD, results are displayed up to À = 10 the algorithm diverges for higher values.
Results. Figure shows the evolution of the test loss (both with or without re-fitting) as a function of the step size. We observe similar behavior 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 the point where optimization becomes unstable.
24


Reviews:
[{"review_id": "mYDQ1V7YCj", "reviewer": "Reviewer_6UDU", "summary": "The authors study the Gauss-Newton optimization algorithm in the overparameterized setting. They derive conditions for convergence of the Gauss-Newton algorithm with parameter-dependent damping in terms of the convexity of the loss, the smoothness of the loss Hessian, and the damping constant.\n\nThey then carry out an empirical study with a student-teacher, one hidden layer setup (overparameterized student), and show that Gauss Netwon in this setting requires small step size and small initialization to be successful. They also show that GN requires a larger number of steps but seems to induce more feature learning.", "questions": "What is the connection of the mean field analysis to the regimes from these two papers?\n\nhttps://proceedings.neurips.cc/paper_files/paper/2022/hash/d027a5c93d484a4312cc486d399c62c1-Abstract-Conference.html\nhttps://proceedings.mlr.press/v139/yang21c", "limitations": "Yes", "rating": 3, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 1, "strengths": "The idea of studying GN in the feature learning limit and trying to understand its implicit biases is an interesting one. The convergence analysis is new to my knowledge, and the experiments focus nicely on key issues of learning dynamics from second-order methods.", "weaknesses": "The algorithm as presented requires a parameter-dependent damping constant, which requires computation of the smallest non-zero NTK eigenvalues. These values are very difficult to compute in practice, and the algorithm is generally impractical. The theory consists of a convergence result for the aforementioned algorithm, and it’s not clear how relevant it is in practice.\n\nPart of the issue with second order methods like GN are their impracticality in large model, large data settings. The experiments are not on any realistic datasets, and don’t suggest any promising ways forward for development of GN methods.\n\nThere may be connections to other mean field regimes which are not discussed in the text.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We are glad you find our study of the generalization and implicit bias of GN to be interesting. We believe the following clarifications should address all your concerns. If you still have concerns, we are more than happy if you could kindly share them as this would help us improve this work.\n\n\n1. **The damping is optional:** Theoretical results in Proposition 2 cover the case without damping ($\\alpha = 0$).  This is stated in L180 where we say \"\\alpha is non-negative\". However, for more clarity, we will discuss the particular case when $\\alpha =0$.  The proof does not rely on alpha>0. In our experiments, the smallest singular value comes for free after performing an SVD on a matrix of the size of a batch of data.  \n\n 2. **Practicality and computational cost:** GN requires solving a system of  size N (N is the size of a mini-batch). More importantly, this does not require computing a matrix of size M (number of parameters). This is thanks to the Woodbury matrix identity as we discuss in L282-284. In our setting, we use a full batch, but using smaller batches is also possible and works well in practice see ref [11].\nIn this work, we focus on the ideal algorithm GN without introducing additional biases or stochasticity in the estimation which can interfere with the interpretation of the results. This approach allows us to understand the ideal performance that a scalable approximation of GN can aim towards.\n\n3. **Connection with mean field analysis:**  \nThank you for the reference, we will include it. As pointed out in their work, the mean-field regime they discuss is the same as the one studied in Chizat 2018, all of which we discussed in the introduction and L140-149. Note that, all these works consider the gradient flow dynamics while, here, we consider the dynamics of Gauss-Newton in the mean-field limit. Hence, the analysis is not directly applicable to GN, although the tools considered might be useful for future work. \n\n"}, {"review_id": "DhMjItKsVT", "reviewer": "Reviewer_c9sb", "summary": "This paper investigates theoretically and empirically the implicit biases of the Gauss-Newton (GN) optimization algorithm on over-parametrized one hidden layer-models (e.g. the capacity of the model is much higher than that of the ground-truth function to approximate). The main theoretical contribution of the paper pertains to global optimization:  a global convergence rate is derived when dynamics do not blow up in finite time, and to provide sufficient conditions on the initial state of the model to prevent this blow up. The experimental section explores generalization: it compares GN and Gradient Descent (GD) for various initial hidden weights variances and learning rates. The key take-away of this last section (and empirical prescription thereof) is that smaller learning rates are better for GN (while the converse holds for GD) and that GN under optimizes the linear layer (e.g. opposite of the lazy regime) while learning good hidden features.\n\nMore precisely:\n\n- Section 3.1 introduces the hypothesis at use on the ground-truth data, the learning objective (Eq. 1) and the model at use, in the finite width (Eq. 2) and mean-field infinite width (Eq. 3) cases. Overparametrization (A) is defined in terms of the invertibility of the feature kernel.\n\n- Section 3.2 defines the generalized GN vector field (Eq. 5) and defines associated discrete and continuous weight dynamics (Eq. 6). Hypothesis on the parameters coming into play in the GN dynamics are introduced (B and Eq. 5) in order to derive subsequent global convergence results.\n\n- Section 4.1 presents Proposition 1 which states that under the aforementioned assumptions, the GN dynamics either blow up or converge to a global minimizer of the loss with an explicit rate which only depends on specific constants characterizing the functions at play, which is in stark contrast with SGD convergence rate which is controlled by the small singular value of the Neural Tangent Kernel (NTK).\n\n- Section 4.2 presents Proposition 2 which states that it is sufficient, for the dynamics not to blow up and converge that the rate derived previously, that the linear layer is near optimal (e.g. $\\\\|\\nabla_v \\mathcal{L}(f(u_0, v_0))\\\\|$ sufficiently small at initialization).\n\n- Section 5.1 presents the teacher-student experimental setup, where the ground-truth data $X$ is Gaussian (with dimension $d=10$) with associated labels $Y$ generated by a shallow one hidden-layer model (5 hidden units) and the student model is a wide/overparametrized one hidden-layer model (5000), both with a ReLU activation. The three optimization algorithms at use are presented: GN, GD and random features (RF) whereby the hidden weights are random and the linear weights are taken as the minimal-norm least square solution (Eq. 9). \n\n- Section 5.2 presents the performance metrics, i.e. the weighted cosine distance (WCD) between teacher and student features (which allows for teacher and student features of different dimensions, promotes feature alignment and penalizes feature norm) and the test loss after linear re-fitting (Test-LRFit). Test-LRFit is particularly useful to disentangle the analysis of the whole model performance *versus* the quality of the learned hidden features, and subsequent identify when GD and GN operate in the kernel/lazy or feature-learning regime.\n\n- Section 5.3 presents and analyzes training experiments, with varying standard deviations of the visible-to-hidden weights (which we denote here $\\sigma_u$ for simplicity) and learning steps. Key results are as follows:\n   + Upon increasing values of $\\sigma_u$, GD and GN both exhibit a transition between the *feature learning regime* (performance is better than the RF baseline) and the *kernel regime* (performance matches that of the RF baseline).\n   + In the feature-learning regime (small $\\sigma_u$), GN operates best (in terms of the resulting test loss) with a small learning rate ($\\approx 10^{-3}$) with performance degrading onwards while GD operates best as the learning rate increases ($\\approx 10^{2}$). GN best performance is better than GD best performance.\n   + GN under-optimizes the output linear layer when employing small learning rates while learning good hidden features (as measured by the gap in test loss before and after Linear re-fitting).\n   + GN starts learning the internal features (as measured by test loss after linear re-fitting until iteration $\\approx 10^4$) and output features onwards (as measured by the test loss). Conversely, last layer and hidden layer are learned simultaneously with GD (i.e. test loss and test loss after linear re-fitting coincide throughout training).\n   + GN converges faster with larger learning rates, but achieves best generalization when learning rate is small. Therefore good generalization with GN comes at the cost of slower convergence.\n   + Finally, across all experiments, the WCD metric is a good proxy to the test loss.\n\nFinally, the Appendix provides very detailed mathematical proofs and extra experimental results (e.g. use of the SiLU activation function instead of ReLU, experiments on MNIST). \n", "questions": "Important questions/remarks:\n\n- My most important remark (hilighted above) is that I suspect your MNIST experiments are not included in the main because they directly contradict the analysis performed on the regression experiments: GD and GN exhibit the exact same trends, GN works best with large learning rates and optimizes the last layer even better than GD. If you have time by the end of the rebuttal phase, I would be very happy if you could \"fix\" the experiments, or make sense of the observed discrepancies between the regression problem and MNIST. \n\n- Could you please define precisely what you mean exactly by \"blow-up\"? You directly say in the main that if w is defined over a finite horizon ($T < \\infty$), then it is said to \"blow-up\". Then you say that blow up happens when the minimal singular value of the NTK is zero. My impression is that there is not enough intuition in the main as to why this is the case. Looking at the appendix, my own intuition is that when $\\sigma^\\star(A_w)$ approaches zero, $J_w^\\dagger$ diverges (Eq. 15 in the Appendix), and since the GN vector field is \"roughly\" $\\Phi(w) \\approx J_w^\\dagger \\cdot \\nabla \\mathcal{L}$, the GN vector field itself diverges. Is this intuition correct? Could you please include such intuition in the main to better define what \"blow up\" is precisely?\n\n- You say in the main that \"GN provably faster than GD\". Could you please write down explicitly the global convergence rate of GD to be able to compare it directly with that of GN derived in Proposition 1 (even if it is a well-known result)? However, the comparison does not seem this straightforward when reading L.196-197. I'm just wondering how much writing GD global convergence rate would enlighten this part. \n\n- I would have been very happy to see an empirical validation of Proposition 2 on a hand-crafted toy model, even simpler than the teacher-student setting, e.g. $\\mathcal{L}=\\frac{1}{2}(\\theta - \\theta_\\star)^\\top \\cdot Q \\cdot (\\theta - \\theta_\\star)$? On such problems the properties of the spectrum of $\\mathcal{L}$ could be well controlled. Taking $dim(\\theta) = 2$ and visualizing trajectories in case of convergence, divergence, and showing the trajectory of GD as well would have been nice to see!\n\n- L.223, you write: \"The near-optimality condition on v0 can always be guaranteed by optimizing the second variable v while keeping u0 fixed\". So why didn't you put it in practice in your own experiments? Especially because you observed that the last layer was not training so well. I'm not saying though that all GN experiments should have been performed with v-pretraining, but extra experiments could have been included. To make myself clearer, here is the gist of some experiment I would have been happy to see:\n  - identify a failure mode of GN in the student-teacher setting -- i.e. the weight dynamics diverge.\n  - show that by pre-training $v$, the problem vanishes (i.e. the weight dynamics converge to an optimum).\n  - demonstrate it on both the teacher-student problem, and on more complex problem like MNIST.\n\n- L. 297: why is the stopping criterion for GN lesser than that of GD ($K^{GD}=10^6$, $K^{GN}=1.5 \\times 10^5$)? Is this because you observed blow-up? This choice is glaring in Fig. 2, bottom left panel and looks very strange: the GN learning trajectory stops at $10^5$ iterations and as the loss was gently decreasing, I would have been happy to see it go even lower! The bottom line is that the stopping criterion for GN looks quite arbitrary and cherry-picked. Why not using an $\\epsilon-$relative-change-of-loss stopping criterion, namely: stop the simulation where the parameters have converged up to precision $\\epsilon$?\n\n- L.336: \"this is unlike GD where larger step size yield better performing features\". Looking closer at Fig. 1: indeed, the test loss decreases until $\\lambda=10^{-3}$, but what happens beyond? Here again, the cut seems arbitrary. And as we expect, when looking at the blue curves (GD) Fig. 6 in the Appendix, the test loss for GD does increase beyond $\\lambda=10^{-3}$. It is not a major issue, however I'm pointing out that the experimental analysis in the main is skewed again by arbitrary choices. \n\nMinor questions/remarks:\n- For non-experts like me: could you please an intuition why overparametrization (A) is defined this way? Perhaps this is something obvious for researchers of this community. \n-  You say in the main that your theoretical results hold both in the kernel and feature learning regime. However, I do not understand why - here I am certainly missing some subtleties. Perhaps it would have been helpful to add background knowledge in the appendix about the exact definition of the feature and kernel limit. Is this about normalizing the model output by $1/\\sqrt{M}$ (kernel) or $1/M$ (mean-field, as done in this paper)?\n- In the same vein, I'm confused about whether the kernel/feature-learning *regime* is the same as (or related to) the kernel/mean-field *limit*. For the paper to be self-contained, it would be useful (at least for me) to add more background knowledge in the Appendix. \n- L.207: could you explain the intuition?", "limitations": "The main limitation of the paper is the quality of the presented experiments, the fact that they are not consistent across two different learning problems, and that the results presented in the paper do not lead to a concrete prescription. For the sake of having the proposed work prescribe concrete takeaway prescriptions to use the GN algorithm in practice, I would recommend accept by the end of the rebuttal phase, once:\n- the results you obtain on MNIST (on a single hidden layer architecture, consistently with the student-teacher setting) are consistent with the results presented in the main.\n- the choice of the stopping criterion for GN is justified.\n\n**POST-REBUTTAL UPDATE**: the authors have thoroughly addressed my concerns, so I increased my score to accept.", "rating": 7, "confidence": 3, "soundness": 2, "presentation": 4, "contribution": 2, "strengths": "- This is a beautifully written paper that is readable for non-super experts of this literature (as I am).\n- The maths in Appendix are extremely rigorous and neat (I carefully checked the proofs).\n- The conditions under which the derived theoretical guarantees hold are apparently less stringent than equivalent studies on GD (e.g. any activation function, initial weight not constrained to lie within a radius determined by the smallest singular value of the NTK).\n- The experimental set up for the teacher-student study is sound and the RF baseline is extremely informative about what is going on.", "weaknesses": "I should state upfront that I am not an expert of this literature, so I may not be aware of the expected theoretical and experimental standards that lead to acceptance. This being said, I'm overall disappointed by the **experiments**:\n\n- **Most importantly**: the MNIST experiments (last page of the appendix) seem to contradict the main findings that GN has to operate in small learning rate regime to work best and that GN would fail to optimize the last (linear) layer -- the gap between test loss and test loss after linear re-fitting is even greater for GD than GN! In spite of this glaring observation in stark contrast with the results appearing in the main, the authors simply do not comment upon this phenomenon at all (L.744-747). \n\n- Apart from the choice of the model which sticks to the mean-field limit introduced in Section 3 (e.g. large number of hidden neurons, normalizing the model output by M), the theoretical and experimental parts are very disconnected. While I do understand that the intention was to cover optimization and generalization in a complimentary fashion (the former theoretically, the latter empirically), theory does not drive much the design of the experiments, and experiments are not crafted to validate the theory. I make some suggestions in the questions below.\n\n- For the experiments: while I do understand that the choice of a one-hidden layer teacher-student setting seems to be a standard in this literature, it lies in between something that could be even more theoretically tractable with controllable analytical quantities (e.g. a toy quadratic optimization problem with known curvature) and something more realistic (e.g. multiple hidden layers model). So I'm twice frustrated that the theory hasn't been checked in an even simpler setting where everything is analytically well controlled, nor empirically tested against a model that is more complicated, to see if the proposed empirical prescriptions hold beyond this simple setting.\n\n- I'm skeptic about the stopping criterion used for GN and worried that **it might skew all the experimental results**.\n\n- In spite of all the rigor and clarity of the proposed study (which I do appreciate a lot), there is unfortunately no concrete practical prescription to use the GN algorithm. Even more so because the results on MNIST contradict those obtained in the teacher-student setting. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your extremely precise and outstanding review! We are very pleased to read that you found the paper well-written and that you took the time to carefully check the proofs.  We hope our answer clarifies all the remaining points.  \n\nI. **Main points:**\n\n   1. **MNIST experiments:** We agree Figure 8 exhibits a different behavior than the experiment in the main. This is attributed to the training objective that is under-optimized for small step-size and thus ‘mechanically’ results in lower performance compared to larger step sizes. This would illustrate the trade-off between good generalisation vs computational cost due to smaller step sizes. Due to time constraints, we are unfortunately unable to provide further experiments on MNIST, but we agree that it would be valuable to strengthen those results. \n\n   2. **Connection between theory and experiments:**\n      - While Proposition 2 provides a global convergence result for the training loss, it says nothing about feature learning and is “happy” with the features staying near initialization as long as the loss reaches a global solution. The experiments illustrate a limitation of proposition 2 by showing that: It is not enough to globally optimize the training loss to get good generalization (as suggested by prop 2). Instead, the amount of feature learning during optimization is what essentially drives generalization performance.\n      - Pre-training the last layer to near optimality: Please refer to the section reserved to this point. \n\t\n3. **Motivation for the 1 hidden-layer model:** This is one of the simplest intractable models for which the dynamics of GN can diverge and can have non-global solutions. We will clarify this in the text and in particular discuss simpler and more complex models as follows:\n   - Linear models with a strongly convex objective (like a quadratic problem) always converge to a global solution. The dynamics is indeed much simpler to analyze since the NTK matrix A_w remains constant by construction. In that case, there is no need for the initial condition of prop 2 and a similar convergence rate can be obtained. Moreover, on a quadratic problem GN converges to the exact solution in a single iteration. As this behavior is well documented, we decided not to include this case. \n   - More complex networks: We do not make claims on deeper/more complex networks as this would different analysis that goes beyond the scope of this paper. We believe this work is a first step towards studying GN for more complex models. \n\n4. **Stopping criterion for GN:** Apologies for the confusion. We had a maximal time budget of 24h per job (we run 720 jobs in total) to control the overall resources allocated to this project. We are currently running GN with a longer budget. We expect it to decrease smoothly and will update the result as soon as we get it.\n   - In all cases, the time limit ensured the training loss (relative to its initial value) was below 10^{-5} and, in most cases, it was below 10-7  (see figure 2 right).\n   - Since the cost per iteration of GN is large, the time budget was exhausted at an error of 10^{-5} when using small step sizes for GN.\n\n5. **Practical prescription:** \n   - A commonly used prescription for GN is to use the largest possible step size as it achieves faster convergence of the training objective. We show that this results in poor generalization (similar to RF).\n   - Instead, we show that smaller step sizes for GN allow better feature learning and are thus desirable but they come at a higher computational cost.\n   - Ultimately, there must be a tradeoff in the step size to achieve the best generalization at the lowest computational cost. An interesting research direction would be to quantify this optimal choice theoretically, but this is beyond the scope of this work. \n\nII. **Major questions:**\n\n1. **Definition of “Blow-up”:** Apologies for the confusion, your intuition is correct and we will clarify in the text. \n\n2. **Comparing convergence rates of GN and GD:** Thank you, we will provide the rates of GD for more clarity. \n   - For GD:   e^{-(mu sigma_min)/4 t} where sigma_min is the smallest eigenvalue of the NTK matrix A_w. \n   - For GN: for instance, when choosing H to be the identity (natural gradient) so that  L_H =\\mu_H=1 and by choosing alpha = 1, the rate is: e^{- t mu }.\n\n   The eigenvalue sigma_min can be arbitrarily small, thus drastically slowing down the convergence rate of GD. This is not the case for GN.   \n\n3. **Pre-training the last layer to near optimality:** \nWe will clarify that the RF baseline we provide is the extreme case where we train the last layer to exact optimality (instead of near optimality). In that case, the gradient vanishes and it is not possible to move away from the RF solution (which is a global one), so that feature learning is not possible. The intermediate case where the last layer is only near optimal can allow changing the inner weights to some extent but this gets worse as the last layer is closer to optimality.  \n   - The failure mode of  GN (diverging dynamics): We have not observed the divergence of the dynamics empirically. \n\n4. **Larger step-sizes** We noticed that larger values do not allow GD to converge to a minimizer (as you noted in Fig 6, you can see that the training loss remains large). Please note that this non-convergence for large step sizes is expected in optimization theory: there is a maximal step size beyond which GD with constant step size cannot converge, usually, this value is found by trial and error, which is what we did as best as our compute power allows.   \n\nWe will include clarifications for: \n   - Overparameterized model: we mean a model that can achieve zero training error on some fixed training set. The condition (A) ensures there exists a parameter value $(v_0,u_0)$ achieving 0 training error (which also explains the sentence in L207). We will include a simple proof to show that. \n   - NTK/mean-field regimes/limits."}, {"review_id": "OV89W5tH19", "reviewer": "Reviewer_2qbp", "summary": "This work theoretically and empirically studied learning dynamics and generalization properties of one-hidden-layer networks trained with the Gauss-Newton algorithm. The main theoretical result is that the authors provided a loss convergence guarantee, which is provably faster than the same networks trained with gradient descent. The authors then empirically studied generalization properties of the GN-trained NN in a student-teacher setting. ", "questions": "1. How can we generalize the results in this work, theoretically or empirically, to approximate second order methods (e.g. natural GD) that people use to train NNs?", "limitations": "See weaknesses.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 1, "strengths": "The theoretical treatment of the GN in Sec. 4 is well appreciated. While I did not check the correctness of the proof, it appears intuitive that GN methods have better convergence. ", "weaknesses": "My primary concern is the relevance of the work. My impression is that GN method is virtually never used in NN training, and a brief literature search turned up only a few papers discussing this. This is simply due to the dramatically higher compute demand compared to GD, as the authors mentioned. The size of the Jacobian is N by N, where N is the number of trainable parameters. So training any modern NN with GN can be prohibitively expensive, since it involves pseudo-inversing the Jacobian at every step. While the author mentions other methods such as natural gradient descent, which are in spirit connected to GN, it is unclear how the theoretical/empirical results in this work generalize to those methods, which may be more relevant to this conference.\n\nMy second concern, which I am less confident of, is that the fact that GN converges faster than GD (at least in terms of the big O analysis) comes as no surprise. That GN converges faster than GD in a nonlinear least-squared problem appears to be well known (https://www.princeton.edu/~aaa/Public/Teaching/ORF363_COS323/F14/ORF363_COS323_F14_Lec9.pdf). I understand that additional technical requirements need to be proven in the one-hidden-layer NN case, but I'm not sure that they are of importance conceptually. \n\nFinally, the empirical part of the paper and the theoretical paper part are pretty unrelated, besides the fact that they both involve one-hidden layer networks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We believe there is a misunderstanding of the relevance of our contribution. The following clarifications should address all your concerns. If you still have other concerns, we are more than happy if you could kindly share them as this would help us improve this work.  \n\n1. **Relevance:**  We believe there is a big misunderstanding here:\n   - **Natural gradient is covered in our work as discussed in L176.**  It corresponds to taking the matrix H to be the identity as discussed in L176. The result of proposition 2 holds for this case as well.  The experimental section corresponds also to a natural gradient on a Gaussian model where the network represents the meanwhile the variance is fixed. In that particular case, GN and natural gradient have the same updates (see ref [31] ). Hence, our work shows that GN and Natural gradient can find solutions that generalize well. We believe this to be highly relevant to the community interested in second-order methods for deep learning. \n\n   - **The computational cost:** Does not require inverting a matrix of the size of the parameters**, only the size of a batch of data. This is thanks to the Woodbury matrix identity that is commonly used for these methods. We use it on a full batch of data and discuss it in L282-284 (note that in the paper N denotes the sample size not the number of parameters). Using smaller batches is scalable and was shown to work well in practice (see ref 11). In our setting, we use a full batch as we focus on the ideal algorithm GN without introducing additional biases or stochasticity in the estimation which can interfere with the interpretation of the results. This approach allows us to understand the ideal performance that a scalable approximation of GN can aim towards.\n\n\n\t\n2. **Novelty/significance of the convergence result:**: \n   - The convergence result you mentioned in the link holds for invertible Hessians and is only local. This is not applicable to over-parameterized networks where the hessian is degenerate by construction. This makes the dynamical system prone to diverge in finite time (blow-up) and is more challenging to deal with. Our result precisely handles this possible blow-up (see section A.2 of the appendix) using non-standard techniques that go beyond convex analysis and that require dynamical systems analysis instead.\n   - **“I'm not sure that they are of importance conceptually”:** The convergence result of GN might seem intuitive when thinking about the invertible hessian case. However, there is a gap between conjecturing the result from this particular case and rigorously proving it in the more challenging scenario of non-convex over-parameterized networks. Our result relies on non-standard optimization techniques which can still be insightful. In particular, the technical tools used come from dynamical system analysis (as discussed in L225-230) and are different from convex analysis often exploited to establish convergence of GN.\n\n3. **Theoretical vs empirical parts:** The motivation for using GN/natural gradient is the faster convergence rate for the training objective. However, these are known for convex objectives with non-degenerate hessian. Two questions stand in the way to understand if GN/natural gradient is useful for learning over-parameterized neural networks:\n   - When can GN reach a global solution for the training objective? The theory part addresses the first point and focuses only on the training objective. The empirical part confirms the global convergence since GN systematically obtains 0 training error faster than GD.   \n   - Does such a solution generalize well to test data? The empirical part studies the generalization of properties of the solution obtained by GN compared to GD. It provides two practical insights:\n       - The commonly used prescription for GN is to use the largest possible step size to achieve faster convergence of the training objective and benefit from the faster rate of GN/natural gradient. We show that this results in poor generalization (similar to random features without learning the hidden layer).\n       - Instead, we show that smaller step sizes for GN allow better feature learning and are thus desirable but they come at a higher computational cost.\n\n\n"}, {"review_id": "Xwfp0GofvO", "reviewer": "Reviewer_oHC4", "summary": "This paper studies applying the Gauss-Newton method to train the neural network to solve a regression problem in the NTK and mean-field regime. The paper establishes a linear convergence rate on Gauss-Newton in this setting. The authors further provide experiments in the teacher-student network setting to study the effect of initialization and step size on convergence and generalization. ", "questions": "Due to the current scale of NeurIPS and the workload on the reviewer, I am unable to go through the appendix in detail. The authors claim their results also holds for the mean-field regime and the proofs in the appendix seems to only address the NTK regime. I encourage the authors to clarify their results. I think for the mean-field regime the authors may have made some assumptions like the initial point is close to the true minimizer which can circumvent the problem of dealing with non-convexity in the mean-field regime. However, this, again, returns back to the classical convex regime. ", "limitations": "The authors discussed the limitations in the conclusion section. ", "rating": 3, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 1, "strengths": "The presentation of this paper is quite good. All the proofs and theorems and experimental results are well-documented. The proofs and experiment results in the appendix are very nicely structured and organized, making it very clear for the reviewer to understand what the authors are doing. ", "weaknesses": "The major weaknesses of this work is on its contribution. First of all, the Gauss-Newton method is a well-known classical method in convex optimization to achieve a faster super-linear convergence rate than first-order methods such as gradient descent which can only achieve linear convergence rate. Therefore, it is natural to apply this method to improve the optimization on training neural network. However, the challenging problem in training neural network is due to its non-convexity. The authors consider the NTK regime and mean-field regime. It is well-known that the optimization problem in the NTK regime is strongly-convex-like. Therefore, the results the authors established in the paper is non-surprising and further, proposition 1 is only able to establish a linear convergence rate whereas Gauss-Newton is able to achieve locally super-linear convergence rate in the classical convex setting, with some additional assumptions needed to be made. Note that if Gauss-Newton is only able to achieve linear convergence, then the advantage of Gauss-Newton is non-existent over gradient descent since gradient descent can also converge linearly and the cost per-step is much smaller than Gauss-Newton. Also, the goal of optimization in deep learning is for the model to generalize, and it is not clear to me how the Gauss-Newton method can affect the model's generalization. \n\nTherefore, currently, I am issuing a rejection for this submission. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We believe the following clarifications should address all your concerns about the strength of our contribution. If you still have concerns, we are more than happy if you could kindly share them as this would help us improve this work.   \n\n1. **Convergence:** The non-convexity and degeneracy of the loss introduces challenges do not appear when studying GN for strongly convex objectives. This is apparent in the proofs which do not rely on classical convexity arguments. Below, we discuss the relevance of the result in terms of linear vs super-linear rate and mean-field limit.\n   - “Linear vs super-linear rate”: The linear rates we provide are general as they include both GN (when the matrix H is the hessian) and natural gradient (when H is the identity). The super-linear rates are only possible when choosing H to be the hessian and are only local. Moreover, Natural gradients are still useful and extensively studied despite their linear rate (see next point). \nImproved conditioning: the linear rate for GN still has much better conditioning than gradient descent which justifies its advantage over gradient descent even if the cost per iteration is higher. Our result highlights precisely this. \n\n   - Mean-field limit: Indeed, the convergence result assumes the initial point is close to optimality. This is stated in Proposition 2 of the main text: L215-216. By mean-field limit, we mean using a scaling of $1/M$ instead of $1/\\sqrt(M)$ where $M$ is the number of particles (L140). Both NTK and mean-field are captured in proposition 2. That being said, we clarify that this does not recover the classical convex regime for two reasons:\n      - The loss is not convex even near the local minimizer: this was shown in Liu et al  2021. Hence, one can’t use convexity arguments in parameter space. \n      - The pre-conditioner in GN can become singular and make the dynamical system diverge (also called finite-time blow-up). This never happens in the strongly convex case where the pre-conditioner has a lower-bounded singular value. This is precisely what makes the analysis more challenging than the strongly convex case. \n\n2. **Generalization of GN:** These are precisely the empirical contributions of the paper: \n   - GN generalizes better with smaller step sizes and exhibits a phenomenon of “hidden learning”: Both train and test losses can remain large when using small step sizes because the linear layer is not well fitted, still the hidden features keep improving. This ultimately yields better generalization after fitting the linear layer alone (Figures 1 and 2).\n\n   - We show that using a larger step size for GN does not generalize well despite being faster at optimizing the training objective: This challenges the common belief of using a larger step size for GN.\n\n   - We show empirically that GN can generalize well, even better than GD. This challenges a common belief that GN does not find solutions that generalize well.\n"}]

Author rebuttal:

Meta-review: This work studies the global convergence and generalization properties of the Gauss-Newton (GN) method when optimizing one-hidden layer networks in the over- parameterized regime. This is a welcome contribution, as the literature focused more on gradient descent. Following the discussion period, the authors went through lots of efforts to clarify the paper.
Decision: Accept (poster)