SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Michael Arbel, Romain Menegaux*, and Pierre Wolinski*†  
Univ. Grenoble Alpes, Inria, CNRS, Grenoble INP, LJK,38000 Grenoble, France  
firstname.lastname@inria.fr
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.
Gauss-Newton (GN) and related methods, like the Natural Gradient (NG) can offer faster convergence rates compared to gradient descent due to the improved dependence on the conditioning of the problem [9, 21, 30, 31]. For these reasons, these methods have attracted attention in machine learning as an alternative to gradient descent when optimizing ill-conditioned problems arising from the use of vastly over-parameterized networks and large training sets [29, 45]. Unfortunately, GN's high computational cost per iteration, which involves solving an expensive linear system, restricts its applicability in large-scale deep learning optimization. Addressing this challenge has been a primary focus, with extensive efforts dedicated to 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, 9, 11, 21, 31, 42].
While substantial effort has been put into making GN scalable, little is 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 knowing that state-of-the-art machine learning methods rely on over-parameterized networks for which infinitely many solutions interpolating the data exist albeit exhibit varying generalization properties that depend on implicit bias of the optimization procedure [7, 39]. To date, the majority of studies on generalization in over-parameterized networks have primarily focused on analyzing the dynamics of gradient descent and gradient flows – their continuous-time limit -, making important progress in understanding the implicit bias of these optimization procedures. Specifically, [15]
†Now affiliated with the LMO, Université Paris-Saclay, Orsay, France.
showed that a gradient flow can run under two regimes which yield solutions with qualitatively different generalization properties: the kernel regime, where the internal features of a network remain essentially constant during optimization, often yields poor generalization, and the feature learning regime, which allows the internal features of a network to adapt to the data, seems behind the impressive 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 a network of width  $M$  by a factor  $1 / \sqrt{M}$ , while the feature learning regime can arise when instead taking the mean-field 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-parameterized networks. The added complexity of these methods makes their analysis intricate, particularly in the stochastic setting where inaccurate estimates of the pre-conditioning matrix can alter the dynamics in a 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 nonlinear over-parameterized models remains understudied in the feature learning regime which is a regime of interest as 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 on 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 a regression problem using over-parameterized one-hidden layer networks in a deterministic (full-batch) setting. By examining an exact, deterministic GN method, rather than an approximate or stochastic version, we eliminate potential regularization effects stemming from approximation and stochasticity that might influence convergence and generalization. This strategy, in turn, enables us 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-time limit of GN's dynamics in the over-parameterized regime that holds under the 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-parameterized networks. Instead, our result relies on a simple condition on the initial 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).
2. The empirical study in Section 5 complements the convergence result in Proposition 3 by further investigating the implicit bias of GN on a 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 initialization. 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.
Convergence of Gauss-Newton'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-Marquardt method [25, 38]. Many such works focus on a continuous-time analysis of the methods as it allows for a simplified study [24, 41]. In this work, we also consider continuous-time dynamics and leave the study of discrete-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 a 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-parameterized regime. [34] studied the convergence of the Levenberg-Marquardt dynamics, a regularized version of Gauss-Newton's method, in a non-convex setting under a cubic growth assumption and showed global albeit sub-linear convergence rates. It 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 a 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 in the feature learning regimes.
The implicit bias of gradient flows. In the over-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 a given initialization and several works studied this problem in the case of linear neural networks [6, 33, 40, 49, 51]. For nonlinear networks, there is still no general characterization of the implicit bias in the regression setting, although [10] recently made important progress in the case of a one-hidden layer network with ReLU 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 a regression problem where the goal is to approximate an unknown real-valued function  $f^{\star}$  defined over a subset  $\mathcal{X}$  of  $\mathbb{R}^d$  from a pair of i.i.d. inputs/outputs  $(x_{n},y_{n})_{1\leq n\leq N}$  where each  $x_{n}$  is a sample from some unknown distribution  $\mathbb{P}$  and  $y_{n} = f^{\star}(x_{n})$ . We assume that  $f^{\star}$  belongs to  $L_{2}(\mathbb{P})$ , the set of square-integrable functions w.r.t.  $\mathbb{P}$  and note that  $f^{\star}$  always belongs to  $L_{2}(\hat{\mathbb{P}})$ , where  $\hat{\mathbb{P}}$  denotes the empirical distribution defined by the samples  $(x_{1},\ldots ,x_{N})$ . For simplicity, we assume the data are non-degenerate, meaning that  $x_{n}\neq x_{n^{\prime}}$  whenever  $n\neq n^{\prime}$ . In this case  $L_{2}(\hat{\mathbb{P}})$  is isomorphic to  $\mathbb{R}^{N}$  (i.e.,  $L_{2}(\hat{\mathbb{P}})\cong \mathbb{R}^{N}$ ) and we can identify any function  $f\in L_2(\hat{\mathbb{P}})$  with the evaluation vector  $(f(x_{1}),\dots ,f(x_{N}))$ . We are interested in approximating  $f^{\star}$  using a one-hidden layer network  $f_{w}$  with parameter vector  $w$  belonging to some, possibly infinite, Hilbert space  $\mathcal{W}$ . The network's parameters are learned by minimizing an objective of the form  $\mathcal{L}(f_w)$  where  $\mathcal{L}$  is a non-negative,  $L$ -smooth, and  $\mu$ -strongly convex function defined over  $L_{2}(\hat{\mathbb{P}})$  and achieving a 0 minimum value at  $f^{\star}$ . A typical example that we consider in Section 5 is the mean-squared error over the training set:
$$
\min  _ {w \in \mathcal {W}} \mathcal {L} \left(f _ {w}\right), \quad \mathcal {L} \left(f _ {w}\right) = \frac {1}{2 N} \sum_ {n = 1} ^ {N} \left(f _ {w} \left(x _ {n}\right) - y _ {n}\right) ^ {2}. \tag {1}
$$
For the convergence results, we will consider both finite-width one-hidden layer networks and their mean-field infinite-width limit. In the experiments, we will restrict to finite-width networks although, we will be using a relatively large number of units  $M$  compared to the size of the training data, therefore approximating the mean-field limit.
Finite-width one-hidden layer networks. Given a non-polynomial point-wise activation function  $\gamma$  and some positive integer  $M$ , these networks take the form:
$$
f _ {w} (x) = \frac {1}{M} \sum_ {i = 1} ^ {M} v _ {i} \gamma \left(u _ {i} ^ {\top} x\right), \quad w = (v, u) \in \mathcal {W} := \mathbb {R} ^ {M} \times \mathbb {R} ^ {M \times d}, \tag {2}
$$
where  $v \in \mathbb{R}^{M}$  is the linear weight vector while  $u \in \mathbb{R}^{M \times d}$  is the hidden weight matrix. Popular choices for  $\gamma$  include ReLU [37] and its smooth approximation  $\mathrm{SiLU}(x) = x(1 + e^{-\beta x})^{-1}$  which can be made arbitrary close to ReLU by increasing  $\beta$  [19]. The above formulation can also account for a bias term provided the input vector  $x$  is augmented with a non-zero constant component.
Mean-field infinite-width limit. We consider some base probability  $\mu$  measure with full support on  $\mathcal{X}$  and finite second moment and denote by  $L_{2}(\mu)$  and  $L_{2}(\mu, \mathbb{R}^{d})$  the set of square integrable functions w.r.t.  $\mu$  and taking values in  $\mathbb{R}$  and  $\mathbb{R}^d$ . Given a non-polynomial point-wise non-linearity, we define infinitely-wide one-hidden layer networks  $f_{w}$  as:
$$
f _ {w} (x) = \int v (c) \gamma \left(u (c) ^ {\top} x\right) \mathrm {d} \mu (c), \quad w := (v, u) \in \mathcal {W} := L _ {2} (\mu) \times L _ {2} (\mu , \mathbb {R} ^ {d}). \tag {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 dynamics 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 vs NTK limits and kernel vs feature learning regimes in Appendix A.
Over-parameterization. We consider an over-parameterized setting where the network has enough parameters to be able to fit the training data exactly, thus achieving a 0 training error. To this end, we define the gram matrix  $G$  which is an  $N \times N$  matrix taking the following forms depending on whether we are using a finite-width network  $(G^{F})$  or infinite-width network  $(G^{I})$ :
$$
G _ {n, n ^ {\prime}} ^ {F} (u) = \frac {1}{M} \sum_ {i = 1} ^ {M} \gamma (u _ {i} ^ {\top} x _ {n}) \gamma (u _ {i} ^ {\top} x _ {n ^ {\prime}}), \quad G _ {n, n ^ {\prime}} ^ {I} (u) = \int \gamma (u (c) ^ {\top} x _ {n}) \gamma (u (c) ^ {\top} x _ {n ^ {\prime}}) \mathrm {d} \mu (c).
$$
We say that a network is over-parameterized if  $G(u_0)$  is invertible for some hidden-layer weight  $u_0$ :
(A) (Over-parameterization)  $G(u_0)$  is invertible for an initial parameter  $w_0 = (v_0, u_0) \in \mathcal{W}$ .
When  $u_{0}$  is sampled randomly from a Gaussian, a common choice to initialize a 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 (i.e.  $x_{i} \nmid x_{j}$  whenever  $i \neq j$  for any  $1 \leq i,j \leq N$ ), a property that holds almost surely if the data distribution  $\mathbb{P}$  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 a Gaussian [28].
3.2 Generalized Gauss-Newton dynamics
To solve (1), we consider optimization dynamics based on a generalized Gauss-Newton method. To this end, we denote by  $J_{w}$  the Jacobian of  $(f_{w}(x_{1}),\dots,f_{w}(x_{N}))$  w.r.t. parameter  $w$  which can be viewed as a linear operator from  $\mathcal{W}$  to  $\mathbb{R}^N$ . Moreover,  $\nabla \mathcal{L}(f_w)$  denotes the vector of size  $N$  representing the gradient of  $\mathcal{L}$  w.r.t. the outputs  $(f_{w}(x_{1}),\dots,f_{w}(x_{N}))$ . We can then introduce the Gauss-Newton vector field  $\Phi : \mathcal{W} \to \mathcal{W}$  defined as:
$$
\Phi (w) := \left(J _ {w} ^ {\top} H _ {w} J _ {w} + \varepsilon (w) I\right) ^ {- 1} J _ {w} ^ {\top} \nabla \mathcal {L} \left(f _ {w}\right),
$$
where  $H_{w}$  is a symmetric positive operator on  $L_{2}(\mathbb{P})$  and  $\varepsilon (w)$  is an optional non-negative (possibly vanishing) damping parameter. Starting from some initial condition  $w_0\in \mathcal{W}$  and for a given positive step-size  $\gamma$ , the Gauss-Newton updates and their continuous-time limit are given by:
$$
w _ {k + 1} = w _ {k} - \gamma \Phi (w _ {k}), \quad \dot {w} _ {t} = - \Phi (w _ {t}). \tag {4}
$$
The continuous-time limit is obtained by taking the step-size  $\gamma$  to 0 and rescaling time appropriately. When  $H_{w}$  is given by the Hessian of  $\mathcal{L}$ , the dynamics in (4) recovers the generalized Gauss-Newton dynamics in continuous time when no damping is used [44] and recovers the continuous-time Levenberg-Maquardt dynamics when a positive damping term is added [35]. When the matrix  $H_{w}$  is the identity  $H_{w} = I$ , the resulting dynamics is tightly related to the natural gradient [2, 30]. More generally, we only require the following assumption on  $H_{w}$ :
(B)  $H_{w}$  is continuously differentiable in  $w$  with eigenvalues in  $[\mu_H, L_H]$  for positive  $L_H, \mu_H$ .
Optional damping. Amongst possible choices for the damping, we focus on  $\varepsilon(w)$  of the form:
$$
\varepsilon (w) = \alpha \sigma^ {2} (w) \tag {5}
$$
where  $\alpha$  is a non-negative number  $\alpha \geq 0$ , and  $\sigma^2(w)$  is the smallest eigenvalue of the Neural Tangent Kernel (NTK)  $A_w := J_w J_w^\top$  which is invertible whenever  $J_w$  is surjective. This choice allows us to study the impact of scalar damping on the convergence rate of the GN dynamics. While computing  $\sigma^2(w)$  is challenging in practice, we emphasize that the damping is only optional and that the case when no damping is used, i.e. when  $\alpha = 0$ , is covered by our analysis.
4.1 Global convergence under no-blow up
We start by showing that the continuous-time dynamics in (4) converge to a global optimum provided that it remains defined at all times.
Proposition 1. Let  $w_0 = (v_0, u_0) \in \mathcal{W}$  be an initial condition so that  $u_0$  satisfies Assumption (A). Under Assumption (B) and assuming that the activation  $\gamma$  is twice-continuously differentiable, there exists a unique solution to the continuous-time dynamics in (4) defined up to some maximal, possibly infinite, positive time  $T$ . If  $T$  is finite, we say that the dynamics blows-up in finite time. If, instead  $T$  is infinite, then  $f_{w_t}$  converges to a global minimizer of  $\mathcal{L}$ . Moreover, denoting by  $\mu$  the strong convexity constant of  $\mathcal{L}$  and defining  $\mu_{GN} := 2\mu / (L_H(1 + \alpha / \mu_H))$ , it holds that:
$$
\mathcal {L} \left(f _ {w _ {t}}\right) \leq \mathcal {L} \left(f _ {w _ {0}}\right) e ^ {- \mu_ {G N} t}. \tag {6}
$$
Proposition 1 is proved in Appendix B.2. It relies on the Cauchy-Lipshitz theorem to show that the dynamics are well-defined up to a positive time  $T$  and uses strong convexity of  $\mathcal{L}$  to obtain the rate in (6). The result requires a smooth non-linearity, which excludes ReLU. However, it holds for any smooth approximation to ReLU, such as SiLU. As we discuss in Section 5, this does not make a large difference in practice as the approximation becomes tighter. The rate in (6) is only useful when the dynamics is defined at all times which is far from obvious as the vector field  $\Phi$  can diverge in a finite time causing the dynamics to explode. Thus, the main challenge is to find conditions ensuring the vector field  $\Phi$  remains well-behaved which is the object of Section 4.2.
Comparison with gradient flow. The rate in (6) only depends on the strong-convexity constant  $\mu$  of  $\mathcal{L}$ , the smoothness constant of  $H$ , and damping strength  $\alpha$ , with the fastest rate achieved when no damping is used, i.e.  $\alpha = 0$ . For instance, when choosing  $H$  to the identity and  $\alpha = 0$ , the linear rate becomes  $\mu_{GN} = 2\mu$ . This is in contrast with a gradient flow of  $w \mapsto \mathcal{L}(f_w)$  for which a linear convergence result, in the kernel regime, follows from [15, Theorem 2.4]:
$$
\mathcal {L} \left(f _ {w _ {t}}\right) \leq \mathcal {L} \left(f _ {w _ {0}}\right) e ^ {\mu_ {G D} t},
$$
with  $\mu_{GF} \coloneqq \mu \sigma(w_0) / 4$  where  $\sigma(w_0)$  is the smallest singular value of the initial NTK matrix  $A_{w_0}$ . In practice,  $\mu_{GF} \ll \mu_{GN}$  since the linear rate  $\mu_{GF}$  for a gradient flow is proportional to  $\sigma(w_0)$ , 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  $w_0 = (v_0, u_0) \in \mathcal{W}$  ensuring the Gauss-Newton dynamics never blows up. We start with a simple result, with a proof in Appendix B.1, showing the existence of a vector  $v^{\star}$  so that  $w^{\star} := (v^{\star}, u_0) \in \mathcal{W}$  interpolates the training data whenever the over-parameterization assumption Assumption (A) holds for a given  $u_0$ .
Proposition 2. Let  $w_0 = (v_0, u_0) \in \mathcal{W}$  be an initial condition so that  $u_0$  satisfies Assumption (A). Define  $\Gamma(u) := (\gamma(u^\top x_n))_{1 \leq n \leq N}$  and  $F^\star := (f^\star(x_1), \dots, f^\star(x_N))$  and let  $v^\star$  be given by:
$$
v ^ {\star} = \Gamma (u _ {0}) G (u _ {0}) F ^ {\star}.
$$
Then the vector  $w^{\star} = (v^{\star}, u_{0}) \in \mathcal{W}$  is a global minimizer of  $w \mapsto \mathcal{L}(f_w)$  and a fixed point of the Gauss-Newton dynamics, i.e.  $\Phi(w^{\star}) = 0$ .