Exact, Tractable Gauss-Newton Optimization in Deep Reversible Architectures Reveal Poor Generalization
Authors: Davide Buffelli
Jamie McGowan
Wangkun Xu
Alexandru Cioba
Da-shan Shiu
Guillaume Hennequin
Alberto Bernacchia
Venue: NeurIPS 2024 poster (2024)
Keywords: Gauss-Newton
Optimization
Deep Learning
Reversible Neural Networks
Primary area: optimization_for_deep_networks

Abstract: Second-order optimization has been shown to accelerate the training of deep neural networks in many applications, often yielding faster progress per iteration on the training loss compared to first-order optimizers. However, the generalization properties of second-order methods are still being debated. Theoretical investigations have proved difficult to carry out outside the tractable settings of heavily simplified model classes - thus, the relevance of existing theories to practical deep learning applications remains unclear. Similarly, empirical studies in large-scale models and real datasets are significantly confounded by the necessity to approximate second-order updates in practice. It is often unclear whether the observed generalization behaviour arises specifically from the second-order nature of the parameter updates, or instead reflects the specific structured (e.g. Kronecker) approximations used or any damping-based interpolation towards first-order updates. Here, we show for the first time that exact Gauss-Newton (GN) updates take on a tractable form in a class of deep reversible architectures that are sufficiently expressive to be meaningfully applied to common benchmark datasets. We exploit this novel setting to study the training and generalization properties of the GN optimizer. We find that exact GN generalizes poorly. In the mini-batch training setting, this manifests as rapidly saturating progress even on the training loss, with parameter updates found to overfit each mini-batch without producing the features that would support generalization to other mini-batches. In contrast to previous work, we show that our experiments run in the feature learning regime, in which the neural tangent kernel (NTK) changes during the course of training. However, changes in the NTK are not associated with any significant change in neural representations, explaining the lack of generalization.

Exact, Tractable Gauss-Newton Optimization in Deep
Reversible Architectures Reveal Poor Generalization
Davide Buffelli" MediaTek Research
Jamie McGowan" MediaTek Research
Wangkun Xut Imperial College London
Alexandru Cioba MediaTek Research
Da-shan Shiu
MediaTek Research
Guillaume Hennequin
Alberto Bernacchia
MediaTek Research & University of Cambridge
MediaTek Research
Abstract
Second-order optimization has been shown accelerate the training of deep neural networks in many applications, often yielding faster progress per iteration on the training loss compared to first-order optimizers. However, the generalization prop- erties of second-order methods are still being debated. Theoretical investigations have proved difficult to carry out outside the tractable settings of heavily simplified model classes thus, the relevance cxisting theories to practical deep learning applications remains unclear. Similarly, empirical studies in large-scale models and real datasets are significantly confounded by the necessity to approximate second- order updates in practice. It is often unclear whether the observed generalization behaviour arises specifically from the second-order nature of the parameter updates, instead reflects the specific structured (e.g. Kronecker) approximations used any hamping-based interpolation towards first-order updates.
Here, we show for the first time that exact Gauss-Newton (GN) updates take on tractable form in a class of deep reversible architectures that are sufficiently expressive to be meaningfully applied to common benchmark datasets. We exploit this novel setting to study the training and generalization properties of the GN optimizer. We find that exact GN generalizes poorly. In the mini-batch training setting, this manifests as rapidly saturating progress even on the training loss, with parameter updates found to overfit each mini-batchatch without producing the features that would support generalization to other experiments run in the "lazy" regime, in which the neural tangent kernel (NTK) changes very little during the course training. This behaviour is associated with having significant changes in neural representations, explaining the lack of generalization.
We show that our
Introduction
Efficient optimization of overparameterized neural networks is a major challenge for deep learning. For large models, training remains one of the main computational and time bottlenecks. Much work has therefore been devoted to the development of neural network optimizers that could accelerate training, enabling researchers and engineers to iterate faster and at lower cost their search for better
"quual Contribution. Correspondence to (davide. buffelli, jamie. www.oacnwnmtteeemmmmm (aggovaa)Jtttkkreeeerrccc.o Work done while MediaTek Research.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

performing models. Second-order optimizers, in particular, have been shown to deliver substantially
faster per-iteration progress on the training loss Martens and Grosse 2015
Botev et al. 2017 George
2018, Goldfarb et al. 2020 Bae et al. 2022 Petersen et al., 2023 Garcia et al., 2023 and et al.
much work has been done to scale them to large models via suitable approximations [Ba et al.
2017
Anil et al. [021]. However, the generalization properties of second-order optimizers remain poorly understood. Here, we focus on the training and generalization properties of the Gauss-Newton (GN) method, which in many cases of interest also encompasses natural gradient descent (NGD)
Martens,2022 2020].
Theoretical studies of generalization in GN/NGD have been limited to simplified models, such as linear models [Amari et al., [022] or nonlinear models taken to their NTK limit [Zhang et al., 2019]. When applied to wealdw networks and large datasets, GN/NGD has so far required approximations,
such as truncated conjugate gradient iterations in matrix-free approaches [Martens et al. 2010), or
block-diagonal and Kronecker-factored estimation of the Gauss-Newton / Fisher matrix Martens
and Grosse 2015 Botev et al.) 2017, George et al. 2018 Goldfarb et al. 2020 |ae et al. 2022
Petersen et [0.2 Garcia al. [023]] Those approximations are exact only in the limit of constant
NTK [Karakida and Osawa 2020 J, in which models cannot learn any features (Yang and Hu 2021
Aitchison 2020]. To our knowledge, the only case in which exact and tractable GN updates have been obtained is that of deep linear networks [Bernacchia et al. 2018 Huh 2020], which despite exhibiting non-trivial learning dynamics [Saxe al.] [2013] cannot learn interesting datasets nor yield additional insights into generalization beyond the linear regression setting. Critically, the use of
necessary approximations makes it difficult to understand how much of the observed generalization
(or lack thereof) can be attributed to the GN method itself, or to the various ways in which it has been simplified.
Here, we derive an exact, computationally tractable expression for Gauss-Newton updates in deep reversible networks (Dinh et al.) 2015 Mangalam et al., 2022]. In reversible architectures made of stacked, volume-preserye possible to analytically derive specific form of generalized inverse for the network's 's Jacobian. This generalized inverse enables fast, exact GN updates in the overparameterized regime. We highlight that, in contrast to the work of [hang al. [2019). |Ci et al. [2019]. Rudner al. [2019). Karakida and Osawa [220]. we do not assume constant NTK, instead we only require the NTK to remain non-singular during training [Nguyen et al. [222 Liu al. 2022 Charles and Papailiopoulos 2018 as, for example, in the mean-field limit (Abel et al. 2023 Equipped with this new model, we study for the first time the generalization behaviour of GN in realistic settings. In the stochastic regime, find that GN trains too well, overfitting single mini-batch at the expense of impaired performance not only on the test set, but also on the training set. To understand this severe lack of generalization, we conduct a careful examination of the model's 's neural tangent kernel and show that the NTK remains almost unchanged during training, and that the neural representations that arise from after training are not different from those set by the network's 's initialization. Thus, GN tends to remain in the "lazy" regime (Jacot al.] [2018] Chizat et al., 2019], in which representations remain close to those at initialization, lacking generalization.
coupling layers (which we call RevMLPs), we show that is
In summary:
We show that GN updates computed with any generalized inverse of the model Jacobian results the same dynamics of the loss, provided that the NTK does not become singular during training (Theorem
We derive an exact and tractable generalized inverse of the Jacobian in the case of deep reversible neural networks (Proposition 4.4). The corresponding GN updates have the same complexity as gradient descent.
We study the generalization properties of GN in models up to 147 million parameters on MNIST and CIFAR-10, and we show that neural representations do not change during training, as the model remains in the "lazy" regime.
2
Related Work
Exact vs approximate Gauss-Newton in deep learning. Previous work on second-order optimiza- tion of deep learning models focused on either Natural Gradient Descent (NGD), or Gauss-Newton

---

(GN), Since the two are equivalent in many important cases [Martens 2020 here we do not distin-
guish them and we refer simply GN. The most popular methods for computing Gauss-Newtor
updates assume block-diagona and pronecker-factored pre-conditioning matrices [Martens and
Grosse 2015 Botev et al.
2017 George et al. 2018 Goldfarb et al. 2020 [aae et al. 2022 Petersen
et al.] 2023] Garcia et al. 2023
Such approximations are known to be exact only in deep inear net-
works [Bernacchia et al., 2018 Huh 2020] and in the Neural Tangent Kernel (NTK) limit Karakida
and Osawa 2020). both of which cannot learn features (Yang and Hu 2021 Aitchison 2020 Recent
work focused on exact Gauss-Newton
the feature learning (mean-feld)) regime |Arel et al. 2023]
but they studied only small models applied synthetic data. The work of Cai et al.) [019 studies
exact Gauss-Newtor on real data but only models with one diee-innnsionl outputs. Our work is the first to investigate exact Gauss Newstn in the feature learning regime real data and sizeable neural networks.
Reversible neural networks. Reversible neural networks [Dinh et al. 2015) allow saving memory during training of large models, because they do not require storing activations [Gomez et al.. 2017 MacKay et al. 2018), and achieve near state-of-the-ar performance [Mangalam et al. [022]. Reversible neural networks also feature in normalizing likelihood-based generative models, or normalizing flows (Dinh et al. 2016). In different reversible models, the inverse is either computed
analytically with coupling layers [Kingma and Dhariwal 2018 Chang et al., 2018 Jacobsen et al.
2018 and similar algebraic tricks [Papamakarioo et al. 2017 HHogebooo al. [2099] Finzi et al.
2019 Xiao and Liu 2020 Lu and Huang 2020]. is computed numerically [Behrmann et al. 2219 Song ct al. 2019 Huang et al.) 2020]. is learned [Keller et al., 2021 Teng and Choromanska 2019]. In this work, we use analytical inversion with coupling layers, because of the efficiency of automatic differentiation through the inverse function. Our work is the first to use reversible neural networks to compute Gauss-Newton updates. A previous work made a connection between reversible models and Gauss-Newtor [Meulemans al. 2020], but they studied Target Propagation, a very different optimizer.
Generalization of Gauss-Newton in overparameterized models. The generalization properties of Gauss-Newton are currently debated. While Wilson et al. (2017) shows worst-case scenarios for adaptive methods, [Zhang et al. [2019] suggests that GN has similar generalization properties as gradient descent (GD) in the NTK limit. In overparameterized linear models, GN and GD find the same optimum [Amari et al.) 2021]. however GD transiently achieves better test loss before convergence (Wadia et al. 202 J. The loss dynamics of Gauss-Newton is approximately re- ce-aaammetrrization invariant, and it remains unclear whether a specific parameterizations allows GD to generalize better [Kerekes et al., 2021]. Previous work also suggests a trade-off between training speed and generalization of GN: a good generalization is obtained only when slowing down training, either by damping (Wadia et al. 2021 or by small learning rates [Arbel et al.) 2023). Here study for the first time generalization for exact GN in sizeable neural networks and real data, and we show that GN achieves poor generalization with respect to gradient descent and similar first order optimizers.
3 Background
We provide brief introduction to Gauss-Newton and Generalized Gauss-Newton. Given input and target data pairs (x. E RRd x Rid and parameters O E R?, the loss is a sum over batch ((7,, 3:12 } of n data points
(( = - (vs, f(z,,
0)) ((.
E(f()))
(1)
with a twice differentiable and convex function é (e.g. square loss or cross entropy) and a parame- terized model fx, 0) (e.g. deep neural network). In the second equality of (.. we concatenate the model outputs (2; 0) E Rd for all n data points in single (column) vector f(0) E Rnd, with entries fe+n-(j-1) = ( 0),, and define concisely the loss in function space as L(f(O)). The loss Z(f) a convex and twice differentiable function of the model f, but L()) is usually a non-convex function of the parameters 0, due to the non-linearity of the model (0). Gradient descent optimizes parameters according to:
0 0 O2 VeC
(2)
3

---

where a is the learning rate and VeC is the gradient of the loss with respect to the parameters. In the full-batch setting, B is the full training dataset. In the mini-batch setting (stochastic gradient descent, SGD), batch of data B is drawn at random from the dataset at each iteration, without replacement, until all data covered (one epoch), after which random batches are re-drawn.
Gauss-Newton We review two alternative but equivalent views on Gauss-Newton: the Hessian view and the the functional view, which provide different intuitions into the method. The Hessian view understands Gauss-Newton as a second-orde optimization method, from the point of view of the curvature of the loss. The functional view understands Gauss-Newton model inversion, and is more appropriate in the context of our work.
In the functional view. Newuss corresponds to gradient descent in function space [Zhang et al., 2019 Cai et al. 2019 Bae et al. 2022 Amari 1998. Martens 2020j. By assumption, the loss L is convex funcction of the model outputs f, thus it would be convenient to optimize the model outputs directly. Gradient flow in function space is given by df di -
However, we need to optimize parameters 0 to have model that can be applied to new data. If f(t) f()(t)), we use the chain rule to find the update in parameter space that corresponds to gradient flow in function space,
of do
00 dt -
f(((:))
Given the gradient VrL and the Jacobian J = 30
of defined as Jab 8f. (of shape ndy p). this is 30,
linear system of equations that can be solved for the update do by sseudo-inverting the Jacobian. In de
discrete time, with learning rate a, the update is equal to [Björck] 1996 Ben-Israe 1965
0.+1 = 0 a J VpL
where the superscript + denotes matrix pscudo-inversion We use this update in our work, in either the full-batch mini-batch setting. We note that equation (5) implies equation (3. in the continuous time limit, only if the Jacobian has linearly independent rows (J.+ = ). which also guarantees convergence to global minimum (full-batch). This requires overparameterization p 2 ndy. however, even if the model is underparameterized and does not converge to a global minimum, equation (5) is still equivalent to Gauss-Newton in the Hessian view, shown below.
In the Hessian view, Gauss-Newton corresponds to Newton's 's method with a positive-definite ap- proximation of the Hessian, in the case of square loss [Dennis Jr and Schnabel 1996 Nocedal and Wrigght 1999. Bottou et al., 2018]. The approximation is accurate near a global minimum of the loss. therefore Gauss-Newton inherits the accelerated convergence of Newton's 's method near global minima [Dennis Jr and Schnabel, [996). The Gauss -eanss update, with learning rate a, is equal to
01+1 = 0, -a aa (JTJ)* + VoL
(6)
Matrix pseudo-inverse is used instead of inverse when J is singular (damping is also a popular
choice, see Nocedal and Wright [1999]. It is straightforward prove that equations (6) and (5) are identical, by noting that, since C(0) [(f())). then VoL J'VrL by chain rule, and (J²J)+ Jr J by the properties of matrix pseudo-inverse. The Graan-Gaus-Newtonn update of Cai et al. 2019 also equivalent equation (5). it just requires the formula for the Jacobian eseudo-inverss in the case of linearly independent rows.
Generalized Gauss Newton. Following the Hessian view, Generalized Neuss- (GGN)
was introduced for convex losses that are different from square loss [Ortega and Rheinboldt 2000 Schraudolph] 2002). The Hessian is approximated by the positive semi-definite matrix J HJ, where
H = V.C. C. As in the case of square loss, the approximation is accurate near global minimum. That leads to the following update:
θ = 0, - a (JH VOL
Note that GGN reduces to GN for If = Ind, (square loss). In the functional view, Appendix shows that Generalized Gauss-Newton corresponds to Newton's 's method in function space, provided That the Jacobian has linearly independent rows and E is strongly convex. Furthermore, Appendix B provides some intuition into the convergence of GGN flow.

---

4 Exact and tractable Gauss-Newton
The main hurdle in the GN update of equation (5) is the computation of the Jacobian pseudo- inverse.. For a batch size nt, number of parameters p and output dimension d, that requires ((ndp min (nd. p)) compute and (ndp) memory. In this Section, efficiently for reversible models. For a dense neural network of L layers and dimension d, implying (Ld)) parameters, our GN update requires the same memory as gradient descent and O (Lndd p dn compute, compared to (Lnd)) compute of gradient descent.
show that the GN update can be computed
Our method consists of two steps: first, we replace the Jacobian pseudoinverse with generalized inverse, and show that it has identical convergence properties (Theorem 4.3). Second, we show that a specific generalized inverse can be computed efficiently in reversible neural networks (Proposition 4.4). We present both results in the case of square loss (GN), Results for other convex loss functions (GGN) can be derived following steps similar to Appendix B
Replacing the pseudoinverse with generalized inverse. We show that the Jacobian pseudoin- verse in equation (5) can be replaced by generalized inverse that has the same convergence properties.
similar approach was proposed by |Bernacchia al. (2018), Karakida and Osawa (2020), but it was only valid in the case of, respectively, deep linear networks or constant Neural Tangent Kernel (NTK) limit. Here we provide a more general formulation that holds under less restrictive assumptions, e.g. it holds in the mean field regime [Arbel et al.) [023]. We need the following assumption
Assumption 4.1. Assume (0) has linearly independent rows (is surjective) for all 0 in the domain where GN dynamics takes place.
Note that this implies that the network is overparametrized, i.e. p 2 ndy. While, in practice, this assumption may seem strong, it is only slightly stronger than the following version, employed in Arbel et al. (0033]:
Assumption 4.2. J(0) is surjective at initialization Oo.
In Arber et al. 2023]. the authors argue that since surjectivity of J is an open condition, it holds for neighbourhood of to, and moreover continue to prove that the dynamics of GN well defined up to some exit time T from this neighbourhood. They then continue to give assumptions guaranteeing that this dynamies extends to Do. We directly assume we are in this latter setting.
Theorem 4.3. Under Assumption I so that there is a right inverse J satisfying JJ 1, consider the update in parameter space with respect to the flow induced by an arbitrary right inverse J
θ θ: - aJ VVL..
(8)
Then the loss along these trajectories is the same up O(x), i.e. for any two choices Jr and Júy the corresponding iterates 0!' and o!²) satisfy:
IV veL(r(o S O(a).
Moreover, as the Moore-Penrone pseudo-invers a right inverse under the assumptions, the result applies to J+, and consequently to the dynamics of (5)
The proof is in Appendix C The intuition behind this result becomes clearer once we examine the differential of the loss w.r.t. the function outputs, VrC. Notice that, as E is a convex function, it has a unique stationary point, and hence it is natural to interpret V+(( as the error at 0, especially close to the global minimum. We will therefore adopt the notation
<0)) == Vr
(10)
here and throughout the proofs to refer to the deviation from the global minimum at the current parameter value. A key ingredient of the proof of Theorem 4il/ be to establish that, for trajectories induced by GGN or the update in equation
(8) c(t) == e ((((t) satisfies:
de
c(t)
(11)
dt
This trivially implies that O from any initial condition co. so that the evolution of the weights approaches a stationary point for the loss, and hence its global minimum.
The right inverse of the Jacobian, J is non-unique, and, in general, it is not feasible to compute for large models. However, turns out that in the case of reversible models, we have an analytic expression for J which allows computing exact GN nearly the same cost as SGD.
5

---

Computing GN of reversible deep network. Throughout this Section we employ the following notation: for an arbitrary matrix X of shape (d, n) we write the lowercase boldfont corresponding symbol, e.g. x for the row-wise vectorization of the matrix, i.e. xi1d -Xi.-
We consider networks composed of I. reversible layers, and we denote by X (with the associated vectorization Xe) and by We (and we), respectively, the output and the parameters of layer in matrix and vector forms. The output of the model is the output of the last layer, f= xt.
The Jacobian of the full neural network can be expressed as a block matrix consisting of the Jacobians of different layers. Letting 0 (W1, W2. layers
.. WL) the concatenated vector with parameters of all
HXL
(J1 JL)
(12)
(WI...., WL)
with J/ = 0xL, Since the only way w affects x is through the way it affects xe, by the chain rule, ow.
the layer-wise Jacobian can be written as
0x θxe
Je Oxe Ow:
(13)
can be computed by finding that the neural network is reversible, a right inverse of equation particular, the product of the inverse of the first factor 0x, / with any vector can be computed exactly with a single forward differentiation pass on the neural network's ' inverse. The inverse of the second factor 0 (8 can be also computed low complexity when individual layers are linear in the parameters, even if nonlinear in the input. These observations hold for any reversible neural network, but here we use dense coupling layers as specific realization (see Section 2), which we call RevMLP. The activation X E Rd' for layer é is written in matrix form X₁ E R and is split along the first dimension into two components X, = (x(") xx(2), ) where Xo is the input. Here d is an even integer and both x!") and x(2) have shape (2 x n). The equations for single coupling layer are
First, we note that a right inverse of the full Jacobian in equation right inverses of the individual, layer-wise Jacobians of equation 13 Then we show that, given
12
13 can be computed easily. In
x1" = x + w!Th ( (
(14)
x = x
2 + w! (2) (V, ( x())
(15)
where We (W(")
w(2) are trainable parameters, while V3 = (V. (1) v((l) are non-trainable parameters (also known as inverted bottleneck, see Bachmann et al. (2024). and (.)) is any differen- tiable non-lincar function. In the rest of this paper, we use the term layer and block interchangeably
to refer full coupling layer (i.e., where the output is the concatenation of x and x as defined in Equation (14) and Equation (15). Whereas we explicitly refer to "half" -oupled layers to mean Equation (14) or Equation (15). We also define the reshaping operator: for x, a vector of dn components we write R.d..) (x) for the matrix A size (d x n) satisfying: Ar.j xi+d(i-1)
Propasition 4.4. Assuming a x ). update for the weights of each layer is given by
(VV x(()) have linearly independent columns, the GN
(x(1)
(1)
W/ (t+)(+1)==33) (() 7 R($mm))
(2(2)
x
(16)
Ox1,
A
w/2) (t+1) -1++)) R(gmn)
0x,
+xx??
(2)
R
R( -1
{1
(0) ) 0 (v=0)(
OXL
@w!")
(1)
(17)
The proof is in Appendix D
Computational and Memory Complexity
The terms in the braces in equations (6). 17 are Jacobian-vector products and can be easily computed using automatic mode differentiation at the cost of one (reverse) inference pass, i.e. O(nd)).

---

1203
(a)
(b)
Figure 1: Training loss and accuracy on (a) MNIST and (b) CIFAR-10 in full-batch scenario where each dataset is trimmed to a fixed subset of n = 1024 images. GN converges much faster than Adam and SGD.
130
500
502
1,00
Test Loos
Trila
134
85-
ON
Epoch
(a)
24438
13
2x:1
1923
Train Loos
Tost Loss
14.39
14+30
12x
14.18
39.
AMVV
30
30
30
Spach
Epoch
Spochh
(b)
Figure 2: Training loss, test loss, and test accuracy on (a) MNIST and (b) CIFAR-10 in a mini-batch scenario. GN does exhibit the same properties observed in the full-batch setting. In fact, Adam reaches lower training and test loss.
The last factor requires pseudo-invedd inverting a matrix of size n x d/2, which requires O(nd min (n, d)). Since these operations are required in each layer, the overall cost of the update for the full network is O Ln²d), compared (IIdd) of SGD, while the memory complexity the same.
Experiments
For our experiments we train RevMLPs (equations and (15) with 2 (6) blocks for MNIST [LeCun al., 2010 (CIFAR-10 10; Krizhevsky 22099) ReLU non-linearities at all half- half-coupled lay- ers, and an inverted bottleneck of size 8000 resulting models with 12M (MNIST) and 147M (CIFAR-10) parameters. We train these models to classify images flattened ID vectors, using cross-entropy objective. Note that the chosen size of the inverted bottleneck ensures that the assumptions Proposition 4.4 hold.
At each training iteration, we compute the pseudoinverses in equa- tions (6). 17 using an SVD. For numerical stability truncate the SVD to a % tolerance relative to the largest singular value and absolute tolerance of 10-5 whichever gives the smallest rank our main findings are qualitatively robust to these tolerance levels. Full hyperparameters and additional details are reported in Appendix [ and code is provided with the submission. We report results averaged over random seeds. All experiments are performed on single NVIDIA RTXA6000 GPU.
Figure 3: Percentage change in training loss after each update. GN decreases the loss for the current mini-batch more than SGD and Adam early in train- ing.
7

---

5.1 Full-Batch Setting
We first examine the full-batch setting in which the dataset is a random subset of size 1024 of MNIST CIFAR-IO. We tuned the learning rate for each optimizer by selecting the largest one that did not cause the loss diverge. Figure T significantly faster than Adam and SGD in both datasets, in line with theoretical predictions (Equation
10-4
shows that GN is
10-
ADAMI
5.2 Mini-Bater Setting
Eyoch
(a)
Next, we consider the full MNIST and CIFAR-10 datasets in the
mini-batch setting. We follow standard train and test splits for the
datasets, with a mini-batch size of n = 1024. The learning rate for
all methods is tuned towards the largest progress after epoch that
does not exhibit any training instabilities. In both datasets, GN
makes good initial progress on the training and test losses, but then
struggles to sustain the continued progress which Adam exhibits (Fig.2) 2. This surprising early saturation of the GN training and test losses is most pronounced for the CIFAR dataset, where even SGD eventually overtakes GN (see Fig. i Appendix [Efor a longer training run). In the rest of this Section, we use the CIFAR-10 10 setup to study the possible origins of such weak generalization.
Epwck
(b)
Block
Lool-
Lo00
Overfitting each mini-batch. Based on the full-batch results of Figure in which GN was seen to converge very fast, we postulate that the poor generalization behaviour observed in the mini-batch case may be caused by overfitting this hypothesis, at each iteration, we compute the loss on a single mini-batch before and after applying the update computed on that same minihb The resulting percentage change in mini-batch loss is shown in Figure 3 Compared SGD and Adam, GN leads to a much stronger immediate decrease in loss after each up- date, especially early training. Whilst this difference gradually weakens during the course of training, it subsists for 80 epochs, i.e. until well after GN's overall training and test losses have sat- urated (c.f. 2ig These results suggest that GN might require some form of regularization to prevent aggressive incorporation of each mini-bated into the model's parameters. However, we find that neither smaller learning rates (Appendix []. nor weight decay (Appendix D. nor any of the usual techniques for regularizing the pseudoinverse in Equation (16) (Appendix K) appear to help in
20
Spech
(c)
each mini-batch. To test
Mock 3
1.00
Epaca
(d)
Wark
CSA Studunity
CK
340
79
Epoch
this respect (Figures 12 1 and 14).
(c)
Evolution of the Neural Tangent Kernel. We further hypothe- size that GN's 's poor generalization may be due to lack of feature learning. In a similar fashion to Fort et al. [220], we study the evolution of the neural tangent kernel (NTK) when training with GN compared to SGD and Adam. A changing NTK would sug- gest that the model learns features different from those present at initialization. Figure and Figure 4bb show the rate of change of the NTK between epochs, and the evolution of the NTK similarity with initialization, respectively.
Figure 4: NTK and CKA sim- ilarity evolution across training for GN, Adam and SGD. Top two panels include (a) the rate of change of the NTK and (b) the NTK similarity during train- ing with respect to initialization. Bottom three panels, along the same axis, include the CKA sim- ilarities for the (c) first, (d) mid- dle and (c) last block with re- spect to their initial values.
Overall, the NTK changes very little for SGD and GN, suggesting that SGD and GN operate close to the "lazy" training regime [Jacot] et al. 2018 Chizat et al., 2019). On the other side, Adam causes the NTK to change significantly, i.e. Adam does learn features different from the initial ones.
8

---

Feature Learning with Gauss-Newton. Even if features (i.e., the NTK) change during GN training. it remains unclear whether they are associated with changes in neural representations. We examine the change in neural representations during training by computing the Centered Kernel Alignment (CKA: Koonnlihhhht et al., 2019) measure of similarity between the representations at initialization and those learned at each epoch, across all layers of the model.
Figures 4c 4d and 4e illustrate the evolution of CKA similarities for the last, middle and first block (i.e.. a "Tull" coupling layer as described by equations [4 [3 of 12 layer RevMLP trained on CIFAR-10. Plots for all blocks are provided in Appendix H along with pairwise similarities across optimizers (Figures T0 and IT). Similar to the NTK, neural representations do not change during training with GN. SGD behaves similarly, with little change in CKA. Adam. on the contrary, has changes in neural representations that coincide with changes in NTK. Appendix Gishows that the lack of changes in neural representations for GN cannot be explained by a smaller change parameters, in fact both GN and Adam show changes in weight space, while weights of SGD change little (Figure (
Contrary to the findings of Arbel et al. [2023], is evident that the CKA similarities in Figures 44 4d and Aemmannn higher for longer in earlier blocks for GN, implying that GN is slower than Adam adapting its deeper internal representations. Furthermore, in Appendix F and | we address two suggestions from Arbel et al. [2023] and find that the generalization improvements when using in Figure 7 and low variance in Figure g) do not carry over to deeper networks. In particular, Figure 7 shows that continuing training with GN after initially training with Adam exhibits the same phenomena as training with GN throughout- - albeit at slightly lower loss than can be achieved using only GN.
smaller learning rates and/or different initializations
5.3 Experiments without Inverted Bottleneck
The previous experiments used inverted bottlenecks to ensure "linear independence". i.e., to ensure that the model is adequately overparametrized such that the proposed efficient generalized inverse is valid. In other words, inverted bottlenecks ensure that the scalable GN weight updates (cquations) [6) do implement gradient flow in function space (equation 3), such that our results are not potentially confounded by broken theoretical assumptions. Nevertheless, the proposed GN update can still be applied in the absence of inverted bottlenecks. In Figure 5 we report results on the CIFARIO dataset, following the same experimental procedure of the previous experiments, but removing all inverted bottlenecks. In the full-batch setting, GN still performing much better than Adam and SGD. In the mini-batch setting we observe a very similar trend to what is observed in the previous experiments: GN leads to an early saturation of the loss, which instead does not appear in Adam and SGD.
24619
23 ana Not
SOAM NoRK
12x
Toutn Looe
(CA No
2x19
560 No41
14438
12x
12415"
Avl 10
Spoch
Figure 5: Experiments on CIFAR using model without Inverted Bottleneck (Full-batch the left, mini- batch on center and right). While the theoretical guarantees do not hold in this setting, the results follow the same trend observed in Figure? 2
5.4 Regression Experiments
We further performed some experiments on regression tasks from the UCI regression datase In
more detail, we used the Superconductivity |Hamidieh| (2018) and Wine Aeberhard and Forina [992
datasets, and followed the same experimental procedure used for the classification datasets (i.e.,
use the same RevMLP architecture with inverted bottleneck for all optimizers and select the
highest learning rate that does not cause the loss to diverge). Results are shown in Appendix M and follow the same trend observed in the classification experiments: in the full-batch case GN is

---

significantly faster than SGD and Adam, while in the mini-batch case there is an apparent stagnation of the test and train losses under GN.
Summary and limitations
In this paper we have introduced a new, tractable way of computing exact Gauss-Newton updates in models with millions of parameters. We have used this theory to study the generalization behaviour of GN in realistic task settings. We found that, although GN yields fast convergence in the full batch regime as predicted, it does not perform as well in the stochastic setting where it tends to overfit each mini-batch. We observed that the NTK does not change when training with GN, suggesting that it operates in the "lazy" regime. In line with the above, using the CKA metric, we performed an analysis of the neural representations at the start and end of training showing that they remain very close to each other. This can explain the observed lack of generalization.
Our investigations have relied on a specific formulation of GN based on tractable generalized inverse of the Jacobian in reversible networks. While we proved that this inverse leads to the same training loss dynamics, in the limit of small learning rate, as the standard GN formulation is based on the Moore-Penross pseudoinverse (MPP), one cannot exclude the possibility that those two update rules have different learning and generalization properties for finite learning rates. Indeed, the functional view of GN (Section 3 makes it clear that the standard MPP-based GN update corresponds to the minimum- weight update that guarantees (infinitesimal) steepest descent in function space. Whilst also achieving steepest descent, our generalized inverse does not have the same least-squares interpretation although it could imply another form of regularization which future work could uncover. In any case, these differences are difficult to assess precisely because the full Jacobian of the network (let alone its MPP) simply cannot be computed for large models.
Previous applications of approximate GN to deep models found that damping, or truncating, the pseu- doinverse of the GN matrix (or, equivalently. of the Jacobian) is key not only for good generalization but also for successful training [Martens et al. 2010 Wadia et al. 2021]. Our generalized inverse is based on layer-wise factorization of the Jacobian, teasing apart the sensitivity of the network's output to small changes in layer activations and (ii) the sensitivity of those activations to small changes in parameters. The use of exactly reversible networks allows us invert the former very efficiently, but does not easily accommodate damping or truncation, making it difficult to study their effect on generalization in large scale settings. We speculate that variation coupling layer-based reversible networks could be developed that allows for damping or truncation, potentially improving the generalization behaviour of GN. If this can be done, our framework would then enable very efficient training of large models, effectively achieving the training acceleration of second-order methods at the cost of first-order optimizers, all memory efficient architecture.
Acknowledgments and Disclosure of Funding
The authors would like to thank Emmeran Johnson for proving that our layer-wise right inverse of the Jacobian (Equation (38) is actually the Moore-Penross pseudoinverse.
References
S. Aeberhard and M. Forina.
Wine.
UCI
Machine
Learning Repository, 1992. DOI:
L. Aitchison. Why bigger is not always better: on finite and infinite neural networks. In International Conference on Machine Learning, pages 156-164. PMLR, 2020.
S.-i. Amari. Natural Gradient Works Efficiently Learning. Neural Computation, 10(2):251-21) 1998.
S.-i. Amari, J. Ba, R. Grosse, X. Li, A. Nitanda, T. Suzuki, D. Wu, and J. Xu. When does precondi- tioning help or hurt generalization? ICLR, 2021.
R. Anil, V. Gupta, T. Koren, K. Regan, and Y. Singer. Scalable second order optimization for deep learning. 2.21.
10

---

M. Arbel, R. Menegaux, and P. Wolinski. Rethinking Gauss-Newton for learning over- parameterized models. Advances in Neural Information Processing Systems, 37, 2023.
J. Ba, R. Grosse, and J. Martens. Distributed second-order optimization using kronecker-factored approximations, In International Conference on Learning Representations, 2017.
G. Bachmann, S. Anagnostidis,, and T. Hofmann. Scaling MLPs: tale of inductive bias. Advances in Neural Information Processing Systems, 36, 2024.
J. Bae, P. Vicol, J. Z. HaoChen, and R. B. Grosse. Amortized proximal optimization. Advances in Neural Information Processing Systems, 35:8982- 118982899777 2022.
Behrmann, W. Grathwohl, R. T. Chen, D. Duvenaud, and J.-H. Jacobsen. Invertible residual networks. In International conference on machine learning, pages 573-582. PMLR, 2019.
A. Ben-Israel. A modified Newton-Raphson method for the solution of systems of equations. Israel journal of mathematics, 3:94-98, -9, 1965.
A. Bernacchia, M. Lengyel, and G. Hennequin. Exact natural gradient in deep linear networks and its application to the nonlinear case. Advances in Neural Information Processing Systems, 31, 2018.
Å. Björck. Numerical methods for least squares problems. SIAM, 1996.
A. Botev, H. Ritter, and D. Barber. Practical Gauss-Newton optimisation for deep learning. In International Conference on Machine Learning, pages 557-565, 2017.
Bottou, F. Curtis, and J. Nocedal. Optimization methods for large-scale machine learning. SIAM Review, 60(2):223 311, 2018.
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. 2019.
B. Chang, L. Meng, E. Haber, L. Ruthotto, D. Begert, and E. Holtham. Reversible architectures for arbitrarily deep residual neural networks. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
Z. Charles and D. Papailiopoulos. Stability and generalization of learning algorithms that converge to global optima. In International conference on machine learning, pages 745-754. PMLR, 2018.
L. Chizat, E. Oyallon, and F. Bach. On lazy training in differentiable programming. Advances in Neural Information Processing Systems, 32, 2019.
J. E. Dennis Jr and R. B. Schnabel. Numerical methods for unconstrained optimizution and nonlineur equations. SIAM, 1996.
L. Dinh, D. Krueger, and Y. Bengio. Nice: Non-linear independent components estimation. ICLR 2015 Workshop Track, 2015.
L. Dinh, J. Sohl-Dickstein, and S. Bengio. Density estimation using Real NVP.
International
Conference on Learning Representations, 2016.
M. Finzi, P. Izmailov, W. Maddox, P. Kirichenko, and A. G. Wilson. Invertible convolutional networks. In Workshop on Invertible Neural Nets and Normalizing Flows, International Conference on Machine Learning, volume 2019.
S. Fort, G. K. Dziugaite, M. Paul, S. Kharaghani, D. M. Roy, and S. Ganguli. Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel. Advances in Neural Information Processing Systems, 33:5850-558 1118850-5861, 2020.
J. R. Garcia, F. Freddi, S. Fotiadis, M. Li, S. Vakili, A. Bernacchia, and G. Hennequin. Fisher- Legendre (FishLeg) optimization deep neural networks. In The Eleventh International Confer- ence on Learning Representations, 2023.
T. George, C. Laurent, X. Bouthillier, N. Ballas, and P. Vincent. Fast approximate natural gradient descent in a Kronecker factored eigenbasis. Advances in Neural Information Processing Systems, 31, 2018.
11

---

X. Glorot and Y. Bengio. Understanding the difficulty of training deep feedforward neural networks. AISTATS, 2010.
D. Goldfarb, Y. Ren, and A. Bahamou. Practical quasi-newton methods for training deep neural networks. Advances in Neural Information Processing Systems, 33:2386-2396. 2020.
A. N. Gomez, M. Ren, Urtasun, and R. B. Grosse. The reversible residual network: Back- propagation without storing activations. Advances in neural information processing systems, 30, 2017.
K. Hamidich. Superconductivty Data. UCI Machine Learning Repository, 2018. DOI: http://dbboorg/102444///3/3333
E. Hoogeboom, R. Van Den Berg, and M. Welling. Emerging convolutions for generative normalizing flows. In International conference machine learning, pages 2771 PMLR, 2019.
C.-W. Huang, R. T. Chen, C. Tsirigotis, and A. Courville. Convex potential flows: Universal probability distributions with optimal transport and convex optimization. International Conference on Learning Representations, 2020.
D. Huh. cervature-correatte learning dynamics in deep neural networks. ICML, page 9, 2020.
A. Iserles. A first course in the numerical analysis of differential equations. Cambridge University Press.
J. Jacobsen, A. W. Smeulders, and E. Oyallon. i-revnet: Deep invertible networks. International Conference Learning Representations, 2018.
A. Jacot, Gabriel, and C. Hongler. Neural tangent kernel: Convergence and generalization neural networks. Advances in neural information processing systems, 31, 2018.
R. Karakida and K. Osawa. Understanding approximate Fisher information for fast convergence of natural gradient descent in wide neural networks. Advances in neural information processing systems, 33:10891-10901, 2020.
T. A. Keller, J. W. Peters, P. Jaini, E. Hoogeboom, Forré, and M. Welling. Self normalizing flows. In International Conference on Machine Learning, pages 5378-5387. PMLR, 2021.
A. Kerekes, A. Mészáros, and F. Huszár. Depth Without the Magic: Inductive Bias of Natural Gradient Descent. arXi::2111.11542 2021.
D. P. Kingma and P. Dhariwal. Glow: Generative flow with invertible 1xl convolutions. Advances in neural information processing systems, 31, 2018.
S. Kornblith, M. Norouzi, H. Lee. and G. Hinton. Similarity of neural network representations revisited. In International conference on machine learning, pages 3519-3529, PMLR, 2019.
A. Krizhevsky. Learning multiple layers of features from tiny images. pages 32-33, 2009.
Y. LeCun, C. Cortes, and C. Burges. Mnist handwritten digit database. ATT Labs (Online). Available: http://yann.eecu.coo///ee www.l//w.../eed///iiiiee.e 2, 2010.
C. Liu, L. Zhu, and M. Belkin. Loss landscapes and optimization in over-parameterizee non-linear systems and neural networks, Applied and Computational Harmonic Analysis, 59:85-116. 2022.
Y. Lu and B. Huang. Woodbury transformations for deep generative flows. Advances in Neural Information Processing Systems, 33:5801-558 -5810,,111111- 2020.
M. MacKay, P. Vicol, Ba, and R. B. Grosse. Reversible recurrent neural networks. Advances in Neural Information Processing Systems, 31, 2018.
K. Mangalam, H. Fan, Y. C.-Y. Wu, B. Xiong, C. Feichtenhofer, and J. Malik. Reversible vision transformers. Proceedings of the IEFF/CVF Conference on Computer Vision and Pattern Recognition, pages 10830-10840, 2022.
12

---

J. Martens. New insights and perspectives on the natural gradient method. The Journal of Machine Learning Research, 21(1):5776-5851. 2020.
J. Martens and Grosse. Optimizing neural networks with Kronecker-factored approximate curva- ture. In International conference on machine learning, pages 2408-2417. PMLR, 2015.
J. Martens et al. Deep learning via Hessian-free optimization. In ICML, volume 27, pages 735-742, 2010.
A. Meulemans, F. Carzaniga, J. Suykens, J. Sacramento, and B. F. Grewe. A theoretical framework for target propagation. Advances Neural Information Processing Systems, 33:20024-20036, 2020.
Q. Nguyen, M. Mondelli, and G. F. Montufar. Tight bounds on the smallest eigenvalue of the neural tangent kernel for deep relu networks. International Conference on Machine Learning, pages 8119-88 1129. PMLR, 2021.
J. Nocedal and S. J. Wright. Numerical optimization. Springer, 1999.
J. M. Ortega and W. C. Rheinboldt. Iterative solution of nonlinear equations in several variables. SIAM, 2000.
G. Papamakarios, T. Pavlakou, and I. Murray. Masked autoregressive flow for density estimation. Advances in neural information processing systems, 30, 2017.
A. Paszke, S. Gross, Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Köpf, E. Yang. Z. DeVito. M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala. PyTorch: An imperative style, high- performance deep
learning library, 2019.
F. Petersen, T. Sutter, C. Borgelt, D. Huh, H. Kuehne, Y. Sun, and O. Deussen. ISAAC Newton: Input-based approximate curvature for Newton's 's method. In The Eleventh International Conference on Learning Representations, 2023.
T. G. Rudner, F. Wenzel, Y. W. Teh, and Y. Gal. The natural neural tangent kernel: Neural network training dynamics under natural gradient descent. In 4th workshop Bayesian Deep Learning (NeurIPS 2019), 2019.
A. M. Saxe, L.. McClelland, and S. Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6123 2013.
N. N. Schraudolph. Fast curvature matrix-vector products for second-order gradient descent. Neural computation, 14(7):1723-1733) 2002.
Y. Song, C. Meng, and S. Ermon. Mintnet: Building invertible neural networks with masked
convolutions. Advances in Neural Information Processing Systems, 32, 2019.
Y. Teng and A. Choromanska. Invertible autoencoder for domain adaptation. Computation, 7(2):00, 2019.
N. Wadia, D. Duckworth, S. S. Schoenholz, E. Dyer, and J. Sohl-Dickstein. Whitening and second order optimization both make information the dataset unusable during training, and can reduce or prevent generalization. In International Conference on Machine Learning, pages 10617-10629. PMLR, 2021.
A. C. Wilson, R. Roelofs, M. Stern, Srebro, and B. Recht. The marginal value of adaptive gradient methods in machine learning. Advances in neural information processing systems, 30, 2017.
C. Xiao and L. Liu. Generative flows with matrix exponential. In International Conference on Machine Learning, pages 10452-10461. PMLR, 2020.
G. Yang and E. J. Hu. Feature learning in infinite-wide width neural networks. In International Conference on Machine Learning, pages 11727 PMLR, 2021.
G. Zhang, J. Martens, and R. B. Grosse. Fast convergence of natural gradient descent for over- parameterized neural networks. Advances in Neural Information Processing Systems, 32, 2019.
13

---

Appendix
A Functional view of Generalized Gauss-Newton
Assuming that é is strongly convex, Newton's flow in function space is equal to
df - V
(18)
di
म and pseudo-invert the Jacobian. In discrete time, with learning rate a, the functional view of Generalized Gauss Newton is given by
with / = ViC. Following steps similar to Section we use of = J
at
(19) straightforward to show that equations 19 and 77) are identical when the Jacobian has lin- early independent rows. Under these assumptions, (J² J =-1111 and JTt + f = I.
D++1 = a aa
Furthermore, in Section VoL= J'V,L.
B Convergence of GGN flow
In this Section, we give informal derivation of the continuous timm dynamics of GGN. See Beeeoreea and Rheinboldt [2000), Bottou
[2018] for convergence of GGN in discrete time. We consider the optimization of the eттоr under GGN flow in continuous time. Using the definition of the error
v (equation (TO)), the optimization flow of the error can be derived using the chain rule
de HJJ
d0
(20)
dt
The definition of GGN flow is d0
a
(JH
(21)
dt
Therefore, optimization of the error under GGN flow is equal to
de
HJ(J (
(22)
dt
where we defined the matrix A HJ] JT Using the properties of the matrix pseudoinverse, we note that A is a projection operator, namely A A for any integer power Therefore, eigenvalues of A are either zero or one, implying that the error decreases exponentially of A, while it remains constant in the null space of A. In general, the range and null space of A change during training. We note that the projection is orthogonal with respect to the inner product ef He. Using the same steps as Appendix A if J has linearly independent rows and É is strongly convex, then it is straightforward to show that A = Indy.
the range
C Proof of Theorem 4.3
During the proof we will compare the dynamics of the flow curves two vector fields, namely -J+ (0) and the corresponding -J (O)) The dependence on is assumed throughout and we will write (-J+e)la or just J+ e. The flowlines of these vector fields are given by:
d0 (-J+e) dt
(23)
and
d0
= =-JJec)
(24)
dt
respectively. We will further write plain (tt) for the solutions of equation 23 and @(t) for the solutions of equation (4h.. Moreover, we'll write i to mean f()())) for arbitrary (possibly tensor valued) functions f. to distinguish from (t)))))
14

---

First notice that the right inverses J are not canonical and hence equation 24 represents family of equations and associated flowlines. However, the dynamics of the associated error is the same regardless of the choice, and moreover, the same that of itself. Note that E(((t)) e()) dt which gives:
de
de JJ+
(25)
dt
If J has linearly independent rows have JJ I, therefore de
(26)
dl
If we consider the flow of é <(O(1)), replacing J+ with right inverse J satisfies JJ I, and again we obtain
dē -
(27)
dt
Integrating these equations from the same initial condition ((to) - 0(t)) indeed gives the same error flowlines.
So despite the different dynamics of the 0 and Ô errors propagate identically. can use this to
derive a bound between the errors incurred by the k-th forward Euler iterates of equations (23) and (24). which define the gradient descent equations. Denote by O, the i-th Euler iterate of (t) and by
0, the corresponding iterate of (t). Then, we have:
11: - 0(4)) S a (cl1i--t)= -1)
(28)
K
100 - -()1 0(4))
((((:-tt) -1
(29)
K
where a is the step-size, the number of steps can be computed as |ل=ا and the Lipschitz constants of (-J+e) and (-J+e) respectively, and K and K are constants which depend on the maximum norm of d² (t) and df drz 0(t) across domain, e.g. Iserles Since is at least C² and the dynamics of and take place over bounded domain, e is Lipschitz with constant Le- Then have:
dr
(0.) elO.))l < |E(0,) - e(O(t))) + ((((O(t))) (((())))+ E(O.))
(30)
Le 0, 0t+,)l +0 Lel10: -((t))
(31)
The a
(cl(t--o)- -1)
(32)
K
where K = min(K K, K) and = mnax(L, É). This shows that the dynamics of the gradient descent iterates coincides up to first order a.
D Proof of Proposition 4.4
Proof. We rewrite the layer-wise Jacobian as
Jt = OXL Exi
OxL
a(x)
(1)
x,
(2)
(33)
OxI Ow:
a(x;
(1)
.x,
i(w /
w (2)
0xL.
0xL
(V((2)
x
Ia/2
0x, (1)
8x
(x())
o(v(1) (( + 1a/2
(34)
ow!
T
0xL
0xL
La/2
0
0x
(1))
0x
0x
(1) T
x la/2
(35)
Tty
15

---

where of- o (V and o."' =oV( x for brevity.
Given twar-triangular block matrix
A
0
(36)
is possible define a right inverse as
(-CHBA-
A+
0 C+
)
BA
(37)
Using the above, we prove that a right inverse JØ of Equation is equal to
o P--1
(2)
T+
Ta/2
0x dxe Oxe
J
(+) ol
1//2
Bw
0x (2)
(1)
8-1
(2)
@ 1ay2
(1) o
2 1ap2
OXL
(2)
(38)
It is possible to prove that J defined above in Equation 38 actually corresponds to the Moore-Penrose pseudoinverse of Je (see Appendix Nj. From Equation (34) and Equation (38). have that
OXL. 0xr.
0 (2)
+1 0
T
X
La/2
0
ase dxy
JeJe
(39)
0x1
(1)
0x
(2)
( oe
1
8 Id/2
dx..
Txc (1)
Ox
OxL
T
(40)
0x,
Ox
Oxz
Oxg
(1)
dxL
(2)
dXL
0(x (1) x²²)
OXL Oxe Lan
(41)
0x,
(1)
ÕXL
ox
0XL
((x
Xi
(2))
ÕXL
Ox: OXL
where Equation 40 columns and,
(2) follows from the assumption that 0e-1 and of have linearly independent
(2) Ox
dx (2)
Ow
(1)
- & 1a/2 - ([of) ( ++
2 La/2
Ow (1)
0-1
(2)
T+ = (42)
Ia/2
Additionally, the last equality in Equation 41 holds since X is a bijective function of xe, due to the reversibility of the RevMLP. Finally, following from Equation [2
we note that,
(43)
L
which when substituted into Equation weights of each layer,
along with Equation
38
results in the GN update for the
ex("))
w (t+
R(Ein)
E
(2(2)
(44)
L
Ox,
(11)
w! (t
R((m))
+xx
(xx?²)
-1
(1xxx;1)
-
L
exx
@w!!!)
(45)
16

---

2x107
XMMM
ON
200
Ttum Lowe
Tot
30004
627
382
1000
Epoch
Epoch
Epoch
Figure 6: Training loss, test loss, and test accuracy when training for longer amount of epochs on CIFAR-10 shows that Gauss-Newton is unable to further decrease the training loss, while even plain SGD can reach lower values.
E
Longer training curves
Figure 6 displays the result of training the same RevMLP as in Section 5.2 for 1000 epochs on the CIFAR-10 dataset in mini-batch setting (n 1024). We observe that by continuing the training for longer on CIFAR-10, 10, SGD is able to reach lower values of the training loss when compared to Gauss-Newton. In particular,
highlight that even in 1000 epochs, Gauss- to increase its training performance further than the value it reaches after just 50 epochs. In fact, the results in Figure 6 show that Gauss-Newton tends to increase its training loss after 150 epochs of training.
appears unable
Initialization dependencies for Gauss-Newton
To examine if the poor performance of Gauss- Newton depends on poor initialization point, we first train a model with Adam for epochs, before continuing the training with Gauss-Newton. For comparison, we also train model with Gauss-Newtr for 50 epochs and subsequently continue training with Adam to observe if Gauss-Newtes reaches reaches "bad" local minimum that hard to escape from. These results are provided in Figure 7 and com- pared with their single optimizer counterparts. We choose "good" initialization point Adam trained model at 50 epochs (indicated by the dashed line in Figure 7). which has lower training loss than GN can achieve in the same (or larger) number iterations. One can observe that, even when starting from this "good" ini- tialization, GN eventually saturates value of the loss when compared with the values achievable by continuing training with Adam. We also find that Adam can start from point found by GN and continue training without issues, reaching value of the loss that is lower than the saturation point of GN.
S0P
Triin
300 750
Epoch
Figure 7: Training loss when first using Adam (or GN), and then continuing with GN (or Adam) - the purple dashed line indicates the 50 epochs mark which the optimizers are switched. GN shows early saturation of the loss even when starting from a better intialization point.
higher
In reference to Arbel et al. 2023 we also provide additional results in Figure to show the the initial weight variance chosen. Interestingly, our results
dependency of Gauss-Newton different from those in Arbel et al. (2023) and suggest that choosing a higher variance is preferable. However, all curves exhibit the same phenomena as discussed Section 5 and under-perform with respect to Adam. The default choice for all experiments we report is a T0-3
G Change in weights during training
We analyze the change norm and cosine similarity between the weights at initialization and the end of training when model is trained with different optimizers. Results are shown in Figure We
17

---

GN 10-1
GN 10-2
6N 10-
GN 10-1
Train Loss
Test Loss
155
14-
14
13
122
30
Epoch
Epoch
Figure GN Train and test loss with weights initialized accounding different variances at initialization a.
Daper
Exper
Layer
1.0
Colan Sunbaite
Cheare Sanilattee
80-
39
Specc
Luyer
Obam Stnnkutty
Crans: Staikettt
Coaaso Stmilketty
Coltx Sinbalte
Colte Stailatte
Collse Stmilante
Layer 211
Lasee
Ooota Simlects
Dondon Sundaaiie
500
Epoch
Epoit
Eyech
Figure 9: Cosine similarity with the initial weight initialization across training. ADAM and GN move similarly in weight space indicating consistent behaviour in weight space between the two optimizers. Note that, in this Figure, we use the term "layer" to refer to half-coupled layer in the reversible blocks.
18

---

observe that Gauss Newuss changes the weights to an extent similar to Adam, while SGD show much smaller weight changes, suggesting lazy training regime.
H
Extended Centered Kernel Alignment (CKA) analysis
Block
Black
Kooch
1.022
1.:0
1.05
1,400
3.00
Spect
Dpccn
Dpocs
Block
Block
2lcc 6
0TA Sunbarty
Sundaaty
50
300
40
94
Eporh
Eporh
Spoch
Figure 10: CKA similarity evolution across training for GN, Adam and SGD. GN maintains a high CKA similarity with its initial feature space, very similarly to SGD.
In Figure 10 we present the full spectrum of CKA similarities from initialization across blocks (i.e., full coupling layer) when training on CIFAR-10. We observe that GN behaves very much SGD: the representqations remain very similar to those at initialization. Adam instead tends to change the representations during training, more quickly for later blocks, and more slowly for earlier blocks.
Block
Wodd
C3A SinDurty
03892
Egoch
Epuch
Epoch
Block
Kock
Block 6
16
Epoch
Epocc
Spacc
Figure 11: Pairwise CKA similarity evolution across training between GN and models trained with Adam and SGD.
In addition the CKA evolution results for single optimizers, Figure 11 presents the pairwise similarities between models at each epoch trained with different optimization strategies. These results demonstrate explicitly the close correspondence between SGD and GN learned features for each block.
I Learning rate variations of Gauss-Newton
Figure 12 provides additional training runs of Gauss-Newton with different learning rates. These results indicate that forcing GN to learn slower is not sufficient to reduce the effect of the observed saturation of performance.
19

---

Toxt
Epech
Figure 12: Train loss, test loss and test accuracy (left to right) for RevMLP trained on CIFAR-10 with Gauss-Newton using different learning rates,
Wayit-Dayy
Ui height Decay
2030
Trate
12
AMAM
30
61
30
56
Spach
Epoch
Epocss
Figure 13: Train loss, test loss, and test accuracy (left to right) for a RevMLP trained on CIFAR-10 with Gauss-Newton with weight-decay. Adam is also added for comparison.
Adding Regularization to Gauss-Newton
In this Section we explore wether weight-decay be used improve the performance of Gauss- Newton. We use RevMLP on the full CIFAR- 10 dataset with the same setting presented in Section and add weight decay to the loss during training. We tune the strength of the weight decay using the validation set. Results are shown in Figure 13 We notice that weight decay has a minimal effect of the performance of Gauss-Newton.
K Pseudo-Inverse Regularization
In this Section we explore the effects of different strategies for regularizing the pseudoinverse in the proposed Gauss-Newton update (see equations (T6). (I7). We note that regularization is necessary, as the presence of very small singular values causes numerical instabilities. We compute the pseudoinverse using a singular value decomposition, and we try three different strategies:
Damping: we add constant to all the singular values. In particular we add quantity equal to 1% of the maximum singular value (this quantity damping was tuned by selecting the best performing one over the values 1%, 10%, 0.1%)
we set to zero all the singular values smaller than certain threshold. In particular, we use relative tolerance of % with respect to the largest singular value and an absolute tolerance of 10 (we tune this values in a similar fashion to the previous method). This the strategy used for the results Section 5
Truncation:
Noise: add noise the matrix be pseudoinverted:
then compute the SVD and use all singular values. The noise is sampled from a zero-man eero-mean Gaussian with a standard deviation equal to 10% (this value was selected though tuning procedure above) of the standard deviation of the matrix to be pseudoinverted.
Results are shown in Figure 14 where Adam is also added for comparison. We notice that there is small difference between damping and truncating (with the former performing slightly better), while adding noise does not seem effective. Nevertheless, Gauss-Newton always under-performen when compared to Adam.
L Full Hyperparameters & Experimental Details
In this Section we provide additional details on the hyperparameters and experimental details used for our experiments. Full code to reproduce results is also provided with the submission.
20

---

CUl
2410
ADMA
ADAM
12
14x10
14x39
13AT
Epoch
Figure 14: Train loss, test loss and test accuracy (left to right) for a RevMLP trained on CIFAR-10 with Gauss-Newton using different regularization strategies for the pseudoinverse. Adam is also added for comparison.
Mil
(b)
Figure 15: Train loss and test loss on UCI (a) wine and (b) superconductivity regression datasets. (Full-batch on top, mini-batch on bottom).
Implementation details. Our code is based on the PyTorch framework Paszke al. [2019]. In more detail we use version 2.0 for Linux with CUDA 12.1.
Weight initialization. We use standard Xavier Glorot and Bengio (2010) initialization for the weights, while we initialize the biases to zero.
Sampling the inverted bottleneck. We sample the entries of each inverted bottleneck from zero-centered gaussian with a variance layer dimeensioni
Data augmentations. For the MNIST dataset we do not use any data augmentations. For the CIFAR- 10 dataset we follow the standard practice of applying random crops and resizes. We do not use data augmentations for the regression datasets.
Additional hyperparameters for Adam. We tune the learning rate for each experiment and method, as explained in Section 5] and we use the PyTorch default values for the betas parameters in Adam.
M Regression Results
We report the results for the regression experiments in Figure 15
21

---

N
Proof Liyer-wise Right-inverse is the Moore-Penrose Pseudo-inverse
Consider the Jacobian for layer E:
d tay
A
0
J2 =
du
121
B
C
(46)
(v
Ia/2
Dzr 8h
Ozy th
(47)
(2)
((!)) eo @lap2)
(A B
Denote Je,1
C) For our method, we used
(- -CC BA+
A+
0
with
(48)
C+)
At = oe (ve_)x x
T+
@ La/2,
(49)
dz
(2)
B (1)'
(50)
du?
O = oe (v(()) x x(11)
T+
X L4/2
(51)
Then:
A+A
0 Ctc)
Je Jt,1 = CA BA+A
(52)
We show below that B = BAA+ A, and hence that JØ1 = Jtt (i.e., our right-inverse corresponds to the Moore-Penrose Pseudo-inverse).
We can show that (see Section
N.I
0u (2)
(53)
(1)
1440-11,a-410 4(b-1)
T
Σ (We.a)a. (v(2), j.b
of (v-1) x(1)
(v=l
(54)
kia
kJ
With the above, we first compute (BA+):
(BA+) (
(55)
=
[(B)++
(56)
22:3
(v=7) x(2)
I+
1li2 i4)
(57)
12
t2.52
T Σ oe (v=))x10))
(We,) (v.")) kids
d' (v=))x(1)
σe (v_)x
T+
$132
kid
$2.4
ja.k
(58)
Σ (Wea))in (V(2)) hids
de (v=))x("))
ki:
E(or 0e (V2xx12) )²
T+ (v2)) (v
(59)
32.J4
jn
(We.))..k (v(() k.is
oe (v:T
(oe (v())x12 (V()x(1))
T4
kidi
51.5
(60)
22

---

and finally
(BA* A) -1+tt
(61)
-E (BA*) 7) 11.itt4-)
(62)
in
+ -1 (1 T
Σ
(BA+ +
1(1
(63)
j4.j3
Σ (We.))s.h (v,1))
o' (2
J4.k
(oe (oe(v()
(2)
oe (v_() (2))
)
(64)
Σ (Wea))snk (v,10)
de x(
E6 (VL (V ao [1 x T+)
de
(v(2x221)
(v x )²
(65)
ji.j
janjs
24
Σ (Wea)ink (v:20)
de (V
kii
(oe (v( T oe (v
T+
(v x
01
(66)
(Wea)ink (ve0) kidg de xx))
0e (v[)) x
(2)
(67)
kJ
$1
But 1)ttt
(68)
N.1
Expression for B
dx? (2)
B,+ 11)++
Owl (1)
(440-1),+4(0-)
0(xx²)) (2)
1+160-1)
a(w(") (1)
)a (b)-))
((x )is
((W(")) )abb)
a(WeGoe Vi x )as
a(W!")) ja,b
doe (v x(1))
(Wec)en a(W!")) Jab kjj
(v!") Ad
Σ (We,) de (v.) Aid a(W Jonb
(xx!)) tij ke a(W:
Σ (Wecles of (v(" kij (v:"),
k.t
23

---

e a(w!" (v
= E [wee)aa de (v)) kij (viD)
k,
o(W( )a.b
1 (W!"))
Σ (We.)is (ve" k.s (ve"), h,
(vL) (vE
oi a(W(")a
k,t.s
Σ (Wer))s di x(") li (v.")
(v(2)x²²)
(2)
kia T Σ (We,))n (ve") k,a =a (v[]
bij
of (ve") x(²))
j,b
k.j
24

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: All claims made in the abstract and introduction are reflected in the paper. We show that GN updates computed with any generalized inverse of the model Jacobian results in the same dynamics of the loss, and we introduce a tractable form of the Gauss-Newton update for reversible neural networks in Section 4 We then study its behaviour showing poor generalization in Section
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. It is fine to include aspirational goals as motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss limitations in Section
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged
the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
25

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We provide full assumptions and proofs for all the theoretical results introduced in the paper.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide all details about architecture, datasets, and hyperparameters in the main text and appendix. Code also provided with the submission.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access hosted model (e.g.. in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
26

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Question:
Answer: [Yes]
Justification:
We provide code with instructions to replicate our results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://mips.co
for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https for more details.
7 //nips. htcp//bllcc/uudeeeee/eeeeeeeeeee
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
Yes we provide all information regarding the training details and test details.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification:
Yes our results are averaged over multiple seeds, and we include error bars.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi-
dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
27

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma enror bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: In the experimental Section we provide a description of our computational setting (single NVIDIA RTXA6000 GPU).
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics htttp::///eurppe: httpp//wwww...cpppp/ccccccccaaaa
Answer: [Yes]
Justification:
The paper adheres to the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: the paper has mainly theoretical motivations and outcomes, so there is no direct societal impact of the work performed.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
28

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: the paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: we provide reference for the datasets used in this paper.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswithcone com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
29

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: the paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/moodl as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: the paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: the paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
30


Reviews:
[{"review_id": "bJAY8aynNI", "reviewer": "Reviewer_5mk4", "summary": "Even though the Gauss-Newton method is known as an effective second-order optimization, it suffers from intractability of Jacobian pseudoinverse computation.\nThis paper proposes a fast and efficient optimization method which solves the intractability issue of the Jacobian pseudoinverse in Gauss-Newton optimization method in overparameterized neural networks.\nFirst, the Gauss-Newton optimization is re-interpreted as the functional view, corresponding to the gradient descent in a function space, taking the parameters as input.\nThen, with this perspective, the pseudoinverse can be replaced to the generalized inverse matrix, which yields the equivalent convergence properties, by applying the chain rule of the loss gradient into the functional loss gradient times the parameter derivative, and replace the functional loss gradient into the same form using the generalized inverse.\nThen, with this, the newly proposed exact Gauss-Newton method calculates the right inverse matrix of the Jacobian, with the RevMLP architecture.\n\nAnd the authors showed that this newly proposed Gauss-Newton method, however, deploys an overfitting property, compared to the Adam and SGD methods, showing worse test accuracy compared to the Adam and SGD optimization methods. Finally, this paper suggests several hypothesis, such as the minibatch overfitting and the feature learning on the NTK regime.", "questions": "* Is assuming $\\texttt{RevMLP}$, a reversible neural network, setting feasible to use? It seems that the accuracy with CIFAR-10 dataset has a scale of 70%, which is quite lower than our consensus.\n* I'd like to confirm that the reason of overfitting with respect to each minibatch. What happens if the minibatches are re-shuffled after each iteration? When the test set performance is saturated to minibatches, it will take advantage of the randomness of minibatches when shuffling.\n* With the CKA results, it is expected that the Gauss-Newton learning schemes, which becomes enabled by the reversibility of the model, is more effective when the model gets deeper. Then, what happens with the performance when the reversible model gets deeper or shallower? \n* How is the $V_{\\ell}$ defined? I assumed that it is a class of random matrices, such as random Gaussian or random Fourier transform matrix.\n\n====\n\n* (Line 63) batchatches $\\to$ batches\n* (Line 238) reversed $\\to$ reversible?", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "* Even though this paper contains heavy mathematical details, this paper is clear to follow.\n* The proposed method enabled Gauss-Newton optimization method to be tractable in reversible nonlinear models.\n* With experiments, this showed that even though the final result did not achieve test-set improvements, the initial learning curve is much more steeply learned: this implies that the Gauss-Newton method is well applied.", "weaknesses": "* The method is limited to reversible nonlinear models, which is a strongly restricted class of neural networks. Because of this restriction, the prediction performance is worse than conventional results.\n* The gain of replacing pseudoinverse to a generalized inverse is not clear.\n\nFor further questions, please refer to Questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time reviewing our paper -- we are glad they found it clear to follow.\n\nThe reviewer wondered what is gained by replacing the Moore-Penrose pseudoinverse by a generalized inverse. To the best of our knowledge, there is no known tractable way of computing the Moore-Penrose pseudoinverse for deep nonlinear networks. This not only precludes its use in applications, but also makes it difficult to investigate the properties of GN at scale. Thus, the main “gain” of our generalized inverse is that it offers a computationally tractable expression for a weight update that implements gradient flow in function space (the core principle underlying the GN method, as summarized around Eq. 4).\n\nThe reviewer rightly noted that our generalized inverse applies to reversible network archictures only, which we understand can appear restrictive. Coupling layers (Dinh et al, 2015; the form of reversible blocks that we use here) were originally introduced as a technical work-around for efficient computation of Jacobians and inverses (the same reason we use them here). They have since been very influential in the area of generative modelling (normalizing flows, diffusion models, ...) where they achieve SOTA results, and in other settings for learning flexible isomorphisms. The reversible vision transformer (Mangalam et al., 2022) also achieves near-SOTA results, and Liao et al (NeurIPS 2024) have shown that even large pre-trained LLMs can be made reversible by inserting specific adaptors, leading to memory-efficient fine-tuning. In fact, we believe that extending our efficient GN updates to these other classes of models, in addition to developing new theory to address the overfitting behaviour we have uncovered, is a promising future direction. \n\n> With the CKA results, it is expected that the Gauss-Newton learning schemes, which becomes enabled by the reversibility of the model, is more effective when the model gets deeper. Then, what happens with the performance when the reversible model gets deeper or shallower?\n\nAs the model grows deeper, we find no significant performance increase for GN training, whereas Adam and SGD improve by a small amount towards convergence. As the model gets shallower, we find that performance drops in all methods by a similar amount.\n\nThe reviewer's mention of CKA analysis in this context prompted us to think about alternative ways of studying the emergence of representations in RevMLPs. In the paper as it currently stands, our CKA analysis is performed separately for each half-coupling layer (i.e. separately for Eq. 14 and Eq. 15). We wondered what the CKA similarity would look like at the level of the 'full' coupling layers (i.e. the hidden representation jointly contained in the concatenation of $x_\\ell^{(1)} $ and  $ x_\\ell^{(2)} $ for each block $\\ell$). We computed these CKA similarities (w.r.t. before-training representations at initialization) and found that they remain very high throughout training (Figure 3 of one-page rebuttal PDF). These results show that the large (albeit late) changes in half-coupling-layer representations observed in our original analysis were in fact _coordinated_ between the two half layers in each block, such that the compound representation which they jointly hold does not actually change. This finding is of course completely in line with our main conclusions, and to some degrees even simplify the story (we no longer need to appeal to the fact that the observed changes in (half-layer) representations happen _after_ GN has already reached its final performance). We are grateful to the reviewer for prompting us to think along those lines.\n\n> Is assuming RevMLP, a reversible neural network, setting feasible to use? It seems that the accuracy with CIFAR-10 dataset has a scale of 70%, which is quite lower than our consensus.\n\nThe accuracy we obtained using RevMLPs on CIFAR10 is in fact on par with classical MLPs (~60-70\\% test accuracy without data augmentation). MLPs, whether reversible or not, are just very prone to overfitting on this type of vision problems. Prior to submission, we did experiment with our own reversible version of MLP-Mixer (Tolstikhin et al, 2021; essentially using RevMLPs instead of MLPs in each mixer block), for which we were able to derive an analogous Jacobian generalized inverse. For these rev-MLP-Mixers, performance on CIFAR10 was overall much better ($> 90\\%$ test accuracy with Adam), similar to what was reported in the non-reversible version, showing that the specific reversibility property does not really affect performance. Nevertheless, we had found that GN displayed identical overfitting properties in rev-MLP-Mixer as it did in plain RevMLPs, but training runs took much longer. For this reason we decided to simplify both the exposition and our experiments by focusing our paper on plain RevMLPs.\nWe will add a paragraph of discussion on this in our revision.\n\n> I'd like to confirm that the reason of overfitting with respect to each minibatch. What happens if the minibatches are re-shuffled after each iteration? \n\nIn our experiments in the minibatch setting, we can confirm that the entire dataset is re-shuffled at the beginning of each epoch before getting divided into minibatches; thus, each minibatch is (statistically) unique; what we show is that GN tends to overfit to each minibatch, and this isn't rescued by learning slowly over many randomized minibatches.\n\n> How is the V_t defined? I assumed that it is a class of random matrices, such as random Gaussian or random Fourier transform matrix.\n\nApologies for this omission; we have now added a full description of how these inverted bottleneck weights are drawn: i.i.d. from a normal distribution (as the reviewer suspected) with zero mean and variance $2/d$ (the input to those have dimensions $d/2$, hence the factor of 2).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IRtGXPB58F", "reviewer": "Reviewer_TbBj", "summary": "In this work the authors use reversible neural networks to explore the benefits of exact Gauss-Newton optimization. They provide a theoretical framework for efficient Jacobian pseudoinverses in reversible networks. The authors then provide experiments on MNIST and CIFAR10 comparing SGD, ADAM, and SGD-GN training. They find that on a small (1024) subset of the dataset, full batch training, GN performs well; in the minibatch setting however GN training performs very poorly. In this setting SGD performs poorly in general. They also measure different feature learning metrics and show that GN training performs feature learning after most of the optimization has occured.", "questions": "What happens in the full batch experiments with batch size 2048?\n\nWhat happens if L2 regularization is added to the experiments?", "limitations": "No", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper provides a very clean explanation of Gauss Newton training, and writing it down in terms of the pseudoinverse of the Jacobian is also a nice touch. The use of the reversible neural networks is very clever, and allow for study of the exact dynamics of interest rather than mere approximations. The experiments are basic but very clear, and the paper overall provides some intriguing preliminary results as well as a path towards future studies into GN training.", "weaknesses": "The experiments section could be more full; for example, exploring the batch size dependence more fully (or, extending the full batch examples to more datapoints). I would also be interested to see how GN training performs in a setting where SGD works at least as well as ADAM.\n\nIn addition, it would be helpful if more intuitions about the form of the pseudoinverse of J are brought into the main text.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive appraisal of our work, and for the constructive feedback. Below, we provide answers to their questions and address the weaknesses they have identified. \n\n> The experiments section could be more full; for example, exploring the batch size dependence more fully (or, extending the full batch examples to more datapoints). I would also be interested to see how GN training performs in a setting where SGD works at least as well as ADAM.\n\nWe followed the reviewer's suggestion to extend the full-batch example to 2048 samples, which can be found in Figure 1 in the one-page supplemental PDF. We used the same architecture and experimental procedure as in the main paper, including tuning the learning rate for each optimizer. The overall trend remains unaltered: in the full-batch setting, GN is significantly faster than both Adam and SGD; in the mini-batch setting, GN decreases the training and test loss faster initially, but then causes them to saturate early, while they continue to decrease for Adam (these effects on train/test losses are also reflected in classication accuracies). In general, our experiments have not revealed any important difference in behaviour when changing the batch size. We thank the reviewer for highlighting this aspect, and we will include this new result in the final version of the paper.\n\nThe reviewer will also be interested in our response to Reviewer 9txE, where we show results of additional experiments on two of the UCI regression datasets (Wine Quality and Superconductivity). For these experiments, we used a squared error loss to complement our other classification experiments which used a cross-entropy loss. The results are shown in Figure 4 of our one-page supplemental PDF, and they do corroborate our main conclusions. Interestingly, Adam and SGD perform similarly on the Wine Quality dataset in the mini-batch setting -- a situation which the reviewer wished we had explored --, and even in this setting, GN exhibits similar overfitting characteristics to those reported in our original submission. \n\n> What happens in L2 regularization is added to the experiments?\n\nThat is a great question -- L2 regularization being a very standard way of mitigating overfitting, it had also occurred to us that it could be important in addressing the overfitting behaviour of GN that we describe.  Experiments with L2 regularization can be found in Appendix K of the main submission; following the standard prescription of AdamW (Loshchilov & Hutter, 2017) we implemented L2 regularization as weight decay in all three optimizers. We found that adding L2 regularization has almost no effect on the overfitting behaviour of GN (we have tuned the amount of weight decay over 5 runs with different order of magnitudes). We also tried additional forms of regularization in Appendix L and came to similar conclusions.\n\n> It would be helpful if more intuitions about the form of the pseudoinverse of J are brought into the main text\n\nWe agree that providing higher-level intuitions about the form of our generalized inverse of J would be useful. We had attempted some of that already through our word-only description of the pseudoinversion on lines 234-240, but we will rewrite this to be clearer and more specific. In addition to that, we will mention that, in the overparameterized setting, the Moore-Penrose pseudo-inverse would determine a parameter update with minimum Euclidean norm, among those providing descent in function space. Instead, the update given by our right inverse does not have minimum Euclidean norm, but it may minimize a different type of norm that is currently unknown and will be the subject of future studies.\n\nWe hope that the additional analyses address the reviewer's concerns and that they may consider raising the score of our paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "q4vRcNnVaa", "reviewer": "Reviewer_9txE", "summary": "The paper address the problem of efficient computation of Gauss-Newton (GN) updates in deep neural networks and in particular the question whether GN results in better generalization behavior than SGD. For this, the authors devise a neural architecture based on reversible NNs that incorporates additional linear bottleneck layers and in which by application of Moore-Penrose pseudo-inverse the updates can be derived analytically and computed efficiently through Jacobian-vector products. Consequently, their analysis show that empirically GN struggles to learn useful representations on MNIST and CIFAR-10 with the selected model architecture and shows a significantly different feature learning behavior, i.e., they find a strong change in the NTK and high distances of the CKA wrt the initialization point, than SGD or Adam. The authors provide additional extensive ablations in the appendix and conclude that while GN yields fast convergence in the full batch setting it does not perform well in the stochastic setting in which it tends to overfit each individual batch rather than fitting to the data set.", "questions": "1. Could the authors elaborate on how much the choice of the inverse bottleneck layer and its number of random weights effects the results.\n2. From what I understand most of the derivation in the main text focuses on the squared loss. However, the results in the experimental section focus solely on the cross-entropy loss which would include an additional Hessian term. Could the authors elaborate on this and are there any comparable results for the setting of squared loss settings such as on UCI regression data sets?", "limitations": "The paper adequately addressed the limitations.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper concisely presents analytic and efficiently computable GN updates for a class of reversible neural networks. The analysis is through and well executed and the results are interesting in my opinion.", "weaknesses": "The authors claim that their paper introduces exact updates for the first realistic application of neural networks. It is, however, unclear to me how realistic the constructed model is. Even though the invert bottleneck has been used in prior work (e.g., (Bachmann 2024)) these applications seem inflate the features only mildly while the authors seem to use a rather drastic inverse bottleneck to ensure linear independence. Hence, I am wondering how transferable the results actually are to architecture that are typically used and do not include random feature projections such as those used in the proposed work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thorough and detailed comments. We provide answers to questions and weaknesses below.\n\n> I am wondering how transferable the results actually are to architecture that are typically used and do not include random feature projections such as those used in the proposed work.\n\nThis is a great question. As the reviewer noted, we used inverted bottlenecks to “ensure linear independence”, i.e. to ensure that the model is adequately overparameterized such that the efficient generalized inverse we propose is valid. In other words, inverted bottlenecks ensure that the scalable GN weight updates (Eqs. 16-17) do implement gradient flow in function space (Eq 3; the essence of the GN method), such that our results are not potentially confounded by broken theoretical assumptions. Nevertheless, we agree that our GN updates can still be applied in the absence of inverted bottlenecks: even though they may not enjoy the same theoretical guarantees, they could still lead to similar training behaviour and it is worth investigating this empirically. We did this on the CIFAR10 dataset, with results shown in Figure 2 of the one-page rebuttal PDF. We followed the same experimental procedure as in the paper, only removing all inverted bottlenecks, and we tuned the learning rate for each optimizer. In the full-batch setting, GN is still performing much better than Adam and SGD. In the mini-batch setting we observe a very similar trend to what is shown in our main paper: GN leads to an early saturation of the loss, which instead does not appear in Adam and SGD. We thank the reviewer for this suggestion, and we will include these results in the final version of the paper.\n\n> [...] are there any comparable results for the setting of squared loss settings such as on UCI regression data sets?\n\nWe thank the reviewer for the valuable suggestion to apply our method directly to regression tasks with a squared loss. We have added some new experiments on two of the UCI Regression datasets (Wine Quality & Superconductivity), with results shown in Figure 4 of the one-page rebuttal PDF. These results corroborate our main findings in the cross-entropy / classification setting: in the full-batch case GN is significantly faster than SGD and Adam, while in the mini-batch case  there is an apparent stagnation of the test and train losses under GN.\n\nWe hope that the additional analyses address the reviewer's concerns and that they may consider raising the score of our paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all reviewers for their time reviewing our paper; we are glad the reviewers appreciated our paper's main strengths, found it “clear to follow” with a “very clean explanation of Gauss Newton” and “thorough and well executed” analysis, opening “a path towards future studies into GN training”. We also thank the reviewers for raising interesting questions and making valuable suggestions for additional experiments, which we have done for the most part.

In particular, in the rebuttal PDF, we provide the following new results:
- Mini-batch and full-batch results on CIFAR with a batch size of 2048
- Mini-batch and full-batch results on CIFAR without inverted bottleneck
- CKA results considering the output of each layer (and not of each component of the coupling layer as done in the paper)
- Mini-batch and full-batch results on two regression datasets from the UCI library.

We provide reviewer-specific responses below.

In addition to these, we would like to bring to the reviewers' attention a minor error in our derivation of Proposition 4.4, which we have now corrected and which did not affect our main conclusions. In the original derivation, we had inadvertently omitted a term that arises from the relationship between the two half-coupled layers inside each coupling block. Correcting for this, Proposition 4.4 becomes:

Proposition 4.4: Assuming $\sigma(V_{\ell-1}^{(2)} X_{\ell - 1}^{(2)})$ and $\sigma(V_\ell^{(1)} X_\ell^{(1)})$ have linearly independent columns, 


$W_\ell^{(1)}(t+1)= W_\ell^{(1)}(t)- \frac{\alpha}{L}
 \mathcal{R}^{(\frac{d}2, n)} \left[(\partial\mathbf{x}^1_\ell / \partial x_L) \boldsymbol\epsilon \right]
    \sigma\left(V_{\ell-1}^{(2)} {X_{\ell - 1}^{(2)}}\right)^+ = W_\ell^{(1)}(t)- \frac{\alpha}{L} \Delta^{(1)}$  

$W_\ell^{(2)}(t+1) = W_\ell^{(2)}(t) - \frac{\alpha}{L} \mathcal{R}^{(\frac{d}2, n)}
[(\partial \mathbf{x}^2_{\ell} / \partial x_L) \boldsymbol\epsilon - (\partial \mathbf{x}^2_{\ell} / \partial  \mathbf{w}_{\ell}^{(1)}) \mathcal{R}^{(\frac{d}2, d')^{-1}}  \Delta^{(1)}]$ 

$\sigma \left(V_{\ell}^{(1)} X_\ell^{(1)} \right)^+$ (this term should be at the end of the above equation but openreview has issues rendering it)

We have re-run all our experiments with this modification (which we had spotted shortly after submission) and the new loss curves are almost identical to the old ones; nevertheless we thought we ought to let you know of this change.

Meta-review: The use of second order methods such as Gauss Newton for neural network training has recently received renewed attention. This paper makes two distinct contributions to the study of this class of methods. To enable the study of large-scale Gauss-Newton training without the additional error due to the use of tractable Hessian approximations (like, say, Kronecker approximations), the authors develop a new approach that enables the efficient computation of the exact Hessian pseudoinverses in the special setting of invertible neural networks.
They then use this approach to thoroughly study the training behavior of exact Gauss Newton in this setting.

The reviewers appreciated the cleverness of the invertible Hessian trick and the thoroughness of numerical studies. The main perceived limitation of this work was whether the reversible structure is sufficiently representative for machine learning problems of interest. In other words, the authors found an alternative structure to rule out artifacts due to Hessian approximation, but to what extent are their results artifacts of the nonstandard architecture enabling this?

Nevertheless, there was consensus among the reviewers that this work is worth presenting at NeurIPS, which is why I recommend acceptance as a poster.
Decision: Accept (poster)