Gram-Gauss-Newton Method: Learning Overparameterized Neural Networks for Regression Problems
Authors: Tianle Cai*
Ruiqi Gao*
Jikai Hou*
Siyu Chen
Dong Wang
Di He
Zhihua Zhang
Liwei Wang
Venue:  (2020)
Keywords: Deep learning
Optimization
Second-order method
Neural Tangent Kernel regression

Abstract: First-order methods such as stochastic gradient descent (SGD) are currently the standard algorithm for training deep neural networks. Second-order methods, despite their better convergence rate, are rarely used in practice due to the pro- hibitive computational cost in calculating the second-order information. In this paper, we propose a novel Gram-Gauss-Newton (GGN) algorithm to train deep neural networks for regression problems with square loss. Our method draws inspiration from the connection between neural network optimization and kernel regression of neural tangent kernel (NTK). Different from typical second-order methods that have heavy computational cost in each iteration, GGN only has minor overhead compared to first-order methods such as SGD. We also give theoretical results to show that for sufficiently wide neural networks, the convergence rate of GGN is quadratic. Furthermore, we provide convergence guarantee for mini-batch GGN algorithm, which is, to our knowledge, the first convergence result for the mini-batch version of a second-order method on overparameterized neural net- works. Preliminary experiments on regression tasks demonstrate that for training standard networks, our GGN algorithm converges much faster and achieves better performance than SGD.

Under review as a conference paper at ICLR 2020
GRAM-GAUSS-NEWTOL METHOD: LEARNING OVER- PARAMETERIZED NEURAL NETWORKS FOR REGRES- SION PROBLEMS
Anonymous authors Paper under double-blle blind review
ABSTRACT
First-order methods such as stochastic gradient descent (SGD) are currently the standard algorithm for training deep neural networks. Second-order methods, despite their better convergence rate, are rarely used in practice due to the pro- hibitive computational cost in calculating the second-order information. In this paper, we propose a novel Gram-Gauss-Vee Gauss-Newton (GGN) algorithm to train deep neural networks for regression problems with square loss. Our method draws inspiration from the connection between neural network optimization and kernel regression of neural tangent kernel (NTK). Different from typical second-order methods that have heavy computational cost in each iteration, GGN only has minor overhead compared first-order methods such SGD. We also give theoretical results to show that for sufficiently wide neural networks, the convergence rate of GGN is quadratic. Furthermore, we provide convergence guarantee for mini-batch GGN algorithm, which is, to our knowledge, the first convergence result for the mini-batch version of second-order method on overparameterized neural net- works. Preliminary experiments on regression tasks demonstrate that for training standard networks, our GGN algorithm converges much faster and achieves better performance than SGD.
INTRODUCTION
First-order methods such as Stochastic Gradient Descent (SGD) are currently the standard choice for training deep neural networks. The merit of first-order methods is obvious: they only calculate the gradient and therefore are computationally efficient. In addition to better computational efficiency, SGD has even more advantages among the first-order methods. At each iteration, SGD computes the mini-batch batch instead of all training data. Such randomness introduced by sampling the mini-bated can lead to better generalization (Hardt al. 2013 Keskar et al.) [016 Masters &
gradient only on a mini-
Luschi 2018 Mou et al. 2017 Zhu et al.) 2018) and better convergence (Ge et al. [015]] Jin et al.
2017a b). which is crucial when the function class is highly overparameterizze deep neural networks.
Recently there is a huge body of works trying to develop more efficient first-order methods beyond
SGD Ducch al. 2011 Kingma & Ba 2014 Luo al. 2019 Liu al. 2019).
methods, despite their better convergence rate, are rarely used to train deep neural networks. At each iteration, the algorithm has to compute second order information, for example, the Hessian or its approximation, which is typically an m by matrix where m is the number of parameters of the neural network. Moreover, the algorithm needs to compute the inverse of this matrix. The computational cost is prohibitive and usually it is not even possible to store such matrix.
Second-order
Recently, there is a series works (Du et al. 2018h a Zou al.) 2018 Allen-Zhu et al.. 201 2018a Oymak & Soltanolkotabi] [2018] Arora et al. 2019h Cao & Gu 2019 Zou & Gu considering the optimization of neural networks based on the idea of neural tangent kernel (NTK) (Jacot et al., 2018 Roughly speaking, the idea of NTK is to linearly approximate the output of network w.r.t. the parameters at a local region, thus resulting in a kernel feature map, which is the gradient of the output w.r.t. the parameters. Jacot et al. (2018) shows that when the width of the network tends to infinity, the NTK tends to be unchanged during gradient flow since the network only needs o o(1) change (goes to the width tends to infinity) of parameter fit the data, the change of kernel

---

Under review as a conference paper at ICLR 2020
is also o( 1). As for finite- finite-width width networks, however, the NTK changes and previous works (Allen-Zht jet al. 2018b; Arora et al., (019a) show that for sufficiently wide neural networks the optimization dynamic of GD/SGD is equivalent to that of using GD/SGD solve NTK kernel regression problem where the kernel is slowly evolving (see Lemma] T in Section 2 for a precise description). A
natural question then arises:
Can we gain acceleration by directly solving kernel regression w.r.f. the NTK at each step?
In this paper, we give a positive answer to this question and reveal the connection between NTK regression and Gauss-Newton method. We propose novel optimization method the Gram-Gause Newton (GGN) method. Instead of doing gradient descent, GGN solves the kernel regression w.r.t. the NTK at each step of the optimization. Following this idea, theoretically prove that for overparameterized networks, GGN enjoys quadratic convergence compared to the linear rate of gradient descent.
Besides theoretical fast convergence rate, GGN is also very efficient in practice. In fact, GGN is implicitly a reformulation of the Gauss-Newter method (see Section 3.1 for details) which is classic second-order algorithm often used for solving nonlinear regression problems with square loss. In the Gauss-Newton method, one uses J as an approximation of the Hessian (see Section 2 for formal description) where J is the Jacobian matrix. However, the original Gauss-Newtor method faces challenges when used for training deep neural networks. Most seriously, the size of the approximate Hessian J J is by m, where m is the number of parameters of the neural network. Moreover, for overparameterized neural networks, for training commonly-used neural networks.
is not invertible, which may make the algorithm intractable
GGN bypasses the difficulty stated above as follows. Instead of using J TJJ J as approximate Hessian and applying Newton-type method, each step of GGN only involves the Gram matrix JJ whose size is n by n where n is the number of data. Furthermore, as already mentioned, to get better generalization performance, it is crucial to use mini-batch to introduce sampling noise when calculat- ing derivatives. Therefore, like SGD, we also use mini-batch in GGN. In this case, the size of the Gram matrix further reduces b by b. where b the batch size. Though conventional wisdom may suggest that applying mini-batch scheme to second-order methods will introduce biased estimation of the accelerated gradient direction, we give the first convergence result for mini-batch second-order method on overparameterized networks. Regarding computational complexity, we show that at each iteration, the overhead of GGN is small compared to SGD: the extra computation of GGN is mainly the matrix product JJJ and the inverse of this matrix whose size is small for mini-batch. Detailed analyses can be found in Section 3.3 We next conduct experiments on two regression tasks to study the effectiveness of the GGN algorithm. We demonstrate that in these two real applications, using practical neural network (c.g.. ResNet-32) with standard width, our proposed GGN algorithm can converge faster and achieve better performance than several baseline algorithms.
1.1 RELATED WORKS
Despite the prevalence of first-order methods for training deep neural networks, there have been continuing efforts in developing practical second-order methods (Becker et al. 1988. Pascanu & Bengio 2013). We summarize some of these works below.
The main approach for these methods is to develop delicate approximations of the second-order information matrix so that the update direction can be computed as efficiently as possible. For example, Botev al. (2017) proposed recursive block-diagonal approximation of the Hessian. The blocks are Kronecker factored and can be efficiently computed and inverted. Grosse and Martens in series of works developed the K-FAC method (Martens & Grosse 2015) Grosse & Martens 2016). The key idea is a Kronecker-factored approximation of the Fisher information matrix, which is used as the second-order matrix in natural gradient methods. These works received considerable attention and have been further improved (Wu et al. 2017:Geeore George et al., 22088) Martens et al. 2018). | Bernacchia et al. (2018) derived an exact expression of the natural gradient update, but only works for linear networks. Different from all these works, our GGN algorithm does not try to approximate the second-or -ooder matrix whose size is inevitably huge. Instead, we present an solution of the updating direction, reducing the computational cost significantly. One exceptional concurrent work Ren & Goldfarb (209)) also aims to use the exact Gauss-Newton update. They focus reducing the complexity of inverting approximate Hessian by MorrmannMoorioon--a Morrssoo----oonouuu
2

---

Under review as a conference paper at ICLR 2020
Formula and require subtle implementation tricks to use backpropagation. In contrast, GGN has simpler update rule and better guarantee for neural networks.
In concurrent and independent work, Zhang et al. (2019a) showed that natural gradient method and K-FAC have a linear convergence rate for suffiiently wide networks in full-batch setting. In contrast, method enjoys higher-order (quadratic) convergence rate guarantee for overparameterizzed networks, and we focus on developing a practical and theoretically sound optimization method. We also reveal the relation between our method and NTK kernel regression, so using results based on NTK (Arora et al.) 20195), one can easily give generalization guarantee of our method. Another independent work (Achiam al. 2019) proposed a preconditioned Q-learning algorithm which has similar form of our update rule. Unlike the methods considered in Zhang et al. (2019a); AAchiam et al. 2019) which contain the learning rate that needed be tuned, our derivation of GGN does not introduce a learning rate term (or understood as suggesting that the learning rate can be fixed to be 1 to get good performance which verified in Figure 2 (c))
NEURAL TANGENT KERNEL AND THE CLASSIC GAUSS-NEW METHOD FOR NONLINEAR LEAST SQUARES REGRESSION
Nonlinear least squares regression problem is a general machine learning problem. Given data pairs {x, yo / and class of nonlinear functions f. c.g. neural networks, parameterized by w. the nonlinear least squares regression aims to solve the optimization problem
min L(w) (u(W,,))) W
(I)
In the seminal work (Jacot et al. with infinite width. They showed that optimization on this problem using gradient flow involves a special kernel which is called neural tangent kernel (NTK). The follow-up works further extended the
2018), the authors consider the case when f is a neural network
relation between optimization and NTK which can be concluded in the following lemma:
Lemma 1 (Lemma 3.1 in Arora et al. (2)1a)), see also Dou & Liang] (2019); Mei et al. 2019) )
dw, Consider optimizing problem ① by gradient descent with infinitesimally small learning rate: dr (W.(. where w, is the parameters at time I. Let f, = (/(w,, x;) x=))²= E IR" be the network outputs on all x, time t, and y (y:)== be the desired outputs. Then fe follows the following evolution:
df, dt -ff -))-
(2)
where G2 is an positive semidefinite matrix, iee. the Gram matrix w.r.t. the NTK at time 1, whose (i, entry is (Vw (fwf(ww f(w,, xx), W xz)).
The key idea of Jacot et al. (2018) and its extensions (Du ct al. 2018b a Zou et al. 2018 Allen -hu al., 200 Saj Oymak & Soltanolkoabbi 2018, Lee al. 2019 Yang 2019 Arora et al. 2019b a Cao & Gul 2019 Zou & Gul 2019) is that when the network is sufficiently wide, the Gram matrix at initialization Go is close
fixed positive definite matrix defined by the infinite-width kernel and G is close to Go during training for all t. Under this situation, G, remains invertible, and the above dynamics is then identical to the dynamics of solving kernel regression with gradient flow w.r.t. the current kernel at time t. In fact, |Arora al. (2019a) rigorously proves that fully-traine sufficiently wide ReLU neural network is equivalent to the kernel regression predictor.
As pointed out in Chizat & Bach (0018), the idea of NTK can be summarized as a linear approximation using first order Taylor expansion. We give an example of this idea on the NTK at initialization:
f(w,x;) f(wo, x;) / / x:) (w -oo;
where Vwf x) can then be viewed as an explicit expression of feature map at x, w Wo is the parameter reproducing kernel Hilbert space (RKHS) induced by NTK and f(w, x:) - ʃ (Wo, x;) the target value.
The idea of linear approximation is also used in the classic Causs-Newtor method (Golub) 1965) to obtain an acceleration algorithm for solving nonlinear least squares problem (I. Concretely, at iteration c, Gauss-Newton method takes the following first-order approximation:
(w) f(w,, x,) V f xx) (w wt);
3

---

Under review as a conference paper at ICLR 2020
where W, stands for the parameter at iteration t. We note that this is also the linear expansion for deriving NTK at time t. According to Eq. and 4). to update the parameter, one can instead solve the following problem.
W = argmin
2 [ + Je(w we) -112
(5)
where fe, y have the same meaning as in Lemma 1 and J, (V (W,,(w;, Xwl( Raxm is the Jacobian matrix.
necessary and sufficient condition for to be the solution of Eq. is
(J) J.) (w we) J.
J. (f: y)-
(6)
Below we will denote H, := J. Jr E mm For under- pardtetraraeeee model (i.e., the number of parameters mn is less than the number of data n), H, invertible, and the update rule is
W++ = W. J. y).
This can also be viewed as an approximate Newton's 's method using H, = JJ.J, to approximate the Hessian matrix. In fact, the exact Hessian matrix is
a V w E (f((..))) - -1)) = J, (f((w,))))) - yo V (wr.
(8)
In the case when f is only mildly nonlinear w.r.t. w at data point x,'s, V (w,, x,) ~ 0, and H2 is close the real Hessian. In this situation, the behavior of the Newton method is similar to that of Newton's 's method, and thus can achieve a superlinear convergence rate (Golub) 1965).
3 THE
METHOD
The classic second-order methods using approximate Hessian such as Gauss-Newton method de- scribed in the previous section face obvious difficulties dealing with the intractable approximate Hessian matrix when the regression model is overparameterized neural network.
we develop a Gram-Gauss- Newton (GGN) method which is inspired by NTK kernel regression and does not require the computation of the approximate Hessian. In Section 3.2 we show that for sufficiently wide neural networks, GGN has quadratic convergence rate. In Section 3.3 we show that the additional computational cost (per iteration) of GGN compared to SGD is small.
In Section 3.T,
3.1 THE GRAM-GAUSS GAUSS METHOD FOR OVERPARAMETERIZED NEURAL NETWORKS
We now describe our GGN method to learn overparameterized neural networks for regression problems. As mentioned in the previous sections, for sufficiently wide networks, using gradient descent for solving the regression problem (1) has similar dynamics as using gradient descent for solving NTK kernel regression (Lemma T w.r.t. NTK each step. However, one can also solve the kernel regression problem w.r.t. the NTK at each step immediately using the explicit formula of kernel regression. By explicitly solving instead of using gradient descent to solve NTK kernel regression, can expect the optimization to get accelerated. We propose our Gram-Gauss-Nent (GGN) method to directly solve the NTK kernel regression with Gram matrix G, at each time step t. Note that the feature map of NTK time t, based the approximation in Eq. {4), can be expressed as r + w [we. x), and the linear parameters in RKHS are w w,, also the target is f(w. x;) [(w,, x:). Therefore, the kernel (ridgeless) regression solution (Mohri al. 2018 Liang & Rakhlin 2018) of Jr,s(w w,) = (f(W, x;) - f(W,, x.)) w.r.t. (w - w.) gives the update
W w. G. -
where Jr,s is the matrix of features at iteration t computed the training data set S which is equal to the Jacobian fr, and ys are the vectorized outputs of neural network and the corresponding targets S respectively, and
G Jr,sJT
4

---

Under review as a conference paper at ICLR 2020
is the Gram matrix of the NTK on S.
One may wonder what is the relation between our derivation from NTK kernel regression and the Gauss-Newton method. We point out that for overparameterized models, there are infinitely many solutions of Eq. (5) but our update rule ( essentially uses the minimum norn solution. In other words, the GGN update rule re-derives the Gauss- This somewhat surprising connection is due to the fact that in kernel learning, people usually choose kernel with powerful expressivity, i.e. the dimension of feature space is large or even infinite. However, by the representer theorem (Mohri et al., 2018), the solution of kernel (ridgeless) regression lies in the n-dimensional subspace RKHS and minimizes the RKHS norm. We refer the readers to Chapter 11 Mohri et al. (2018) for details.
method with the minimum norm solution.
As mentioned in Section | the design of learning algorithms should consider not only optimization but also generalization. It has been shown that using mini-batch instead of full batch to compute derivatives crucial for the learned model have good generalization ability (Hardt et al. 2015] Keskar et al. 2016 Masters & Luschi 2018 Mou et al., 2017 Zhu et al., 2018). Therefore, we propose a mini-batch version of GGN. The update rule is the following:
(10) where B, is the mini-batch used at iteration I, Jt,B, and Gt. are the Jacobian and the Gram matrix computed using the data of B, respectively, and B, YB, are the vectorized outputs and the corresponding targets on B, respectively. G.,B, Jr, J is very small matrix when using a typical batch size.
Wi+1 W. - J., G. (G.,F, -Y,),
One difference between Eq. 10p and Eq. 7 is that our update rule only requires to compute the Gram matrix G., and its inverse. Note that the size of G.,,, is equal to the size of the mini-batch and is typically very small. So this also greatly reduces the computational cost.
Using the idea of kernel ridge regression (which can also be viewed as Levenberg-Marde extension (Levenberg 1944) of Gauss-Newton method), we introduce the following variant of GGN:
We+1 W. - JT, (AG,,B, + all (fr, - yB,),
(11)
where À > 0 is another peeeeppareetter controlling the learning process. Our algorithm formally described in Algorithm T
Algorithm (Mini-batch) Gram- Frram-Gauss-Nert Newton Method
Input: Training dataset S. Hyper-parameters À and a.
2: Initialize the network parameter wo. Set / = 0.
3: for each iteration do
4: Fetch mini-batch B, from the dataset.
5: Calculate the Jacobian matrix Jr,
6:
Calculate the Gram matrix G.,,, = -r, JT, B.
7: Update the parameter by W/+1 = We Jr.p (AGr.B, + taI)- - )..)
8:
t=t+1
9:
end for
3.2 CONVERGENCE ANALYSIS FOR OVERPARAMETERIZER NEURAL NETWORKS
In this subsection, we show that for two-layer neural networks, if the width is sufficiently large, then: (I) Full- butl-batch GGN converges with quadratic convergence rate. (2) Mini-batch GGN converges linearly. (For clarity, here we only present proof for two-layer neural networks, but we believe that it is not hard for the conclusion to be extended to deep neural networks using the techniques developed in Du al. (2018)); Zou & Gu (2019))
As we explained through the lens NTK, the result is consequence of the fact that for wide enough neural networks, if the weights are initialized according suitable probability distribution, then with high probability the output of the network is close to linear function w.r.t. the parameters (but nonlinear w.r.t. the input of the network) in neighborhood containing the initialization point and a global optimum. Although the neural networks used in practice are far from that wide, this still motivates us design the GGN algorithm.
5

---

Under review as a conference paper at ICLR 2020
Neural network structure.
We use the following two-layer network
M
f(w.x) 1 a TaWTTTTT VM
a (w,x),
(12)
VM
where x Rd is the input, M is the network width, W [w] W and (+)) is the activation function. Each entry of W is i...d initialized with the standard Gaussian distribution w, ~ N La) and each entry of a is initialized from the uniform distribution
[+1]. Similar to Du et al. (2018b),
we only train the network on parameter W just for the clarity of the proof. We also assume the
activation function (-) E-Lipschitz and 3-smooth, and l and are regarded as ((1) absolute constants.
The key finding. as pointed out in Jacot et al. (2018); Du et al. (2018h a), is that under such initialization, the Gram matrix G has an asymptotic limit, which is, under mild conditions (e.g. input data not being degenerate etc., see Lemma F.2 of Du et al. (2018a)), a positive definite matrix
((x,, x;) = Bw (u.1) [x x (wx;) )o [wxx;]
(13)
Assumption 1 (Least Eigenvalue of the Limit of Gram Matrix). We assume the matrix K defined in 13 above is positive definite, and denote its least eigenvalue as
Ao
Amin K) 0.
Now we are ready to state our theorem of full-batch GGN:
Theorem (Quadratic Convergence of Full-batch GGN on Overparameterized Neural Networks).
Assume Assumption 7 holds. Assume the scale of the data is [xxil2 = O(1). lyal = O(1) for i 3-- the network width
( max
nn n²dd (odlog((66]) )),
M =
XB
then with probability 1 1-6 8 over the random initialization, the full-batch version of GGN whose update rule is given in Eq. 19) satisfies the following:
The Gram matrix Gt,s at each iteration is invertible;
2) The loss converges to zero way that
C [[G++1 - v1l2 /M IIf.-
(14)
for some C that independent of M, which is second-order convergence.
For the mini-batch version of GGN, by the analysis of its NTK limit, the algorithm essentially doing serial subspace correction (Xul (001) on subspaces induced by mini-batch. So mini-batch GGN is similar to the Gauss-Siedel method (Golub & Van Loan 1996) applied to solving systems of linear equations, as shown the proof of the following theorem. Similar to the full batch situation, GGN takes the exact solution of the "kernel regression problem on the subspace" which is faster than just doing gradient step optimize on the subspace. Moreover, we note that existing results of the convergence of SGD on overparameterized networks usually use the idea that when the step size is bounded by a quantity related to smoothness, the SGD can reduced to GD. However, our analysis takes different way from the analysis of GD, thus does not rely on small step size.
In the following, we denote Go Rnxn as the initial Gram matrix. Let n bk, where is the batch size and & is the number of batches, and let
Go.s = Go((i - 1)b + 1 ib, (j - 1)6 - 1 jbb
be the (i, b x b block of Go. We define the iteration matrix
= D -)- "xn
(15)
where
Go.11
0
0
0
0
G
G
0
0
D
:
Go
G
Go
0.

---

Under review as a conference paper at ICLR 2020
represents the block-diagonal and block-lower-triangee parts of Gg. We will show that the conver- gence mini-base batch GGN is highly related to the spectral radius of A. To simplify the proof, we make the following mild assumption on A:
Assumption 2 (Assumption on the Iteration Matrix), Assume the matrix A defined in TS above is diagonalizable. So we choose an arbitary diagonalization of as A P P-PPPQQQ and denote
11 1
==
We note that Assumption2 only for the sake of simplicity. Even if does not hold, an infinitesimally small perturbation can make any matrix diagonalizable, and it will not affect the proof.
Now ready to state the theorem for mini-batch GGN.
Theorem (Convergence of Mini-batch GGN on Overparameterizeed Neural Networks). Assume Assumption and 2 hold. Assume the scule of the data ||x=l|2 = O(1) 1ydl 0(1) for {1. ... n). We use the mini-batch version of GGN whose update rule is given in Eq. 1O and the batch B, is chosen sequentially and cyclically with a fixed batch size and = n/b updates per
epoch. If the network width
M = max (2
112n (1(11
n n2dlog(16n//)
.2
11
X
then with probability 1 over the random initialization, we have the following:
The Gram matrix Gt,B, at each iteration is invertible:
2) The loss converges to zero way that after T epochs, we have
A
T
frk yl2 S pva 1 - 8
(16)
n
Proof sketch for Theorem and 2
Denote J. = =[(W)) and
J J((1 - 8)W, wW W1)dds.
For the full-batch version, we have
|[f++1 yll
|[fe y + Je (voc W(+1) voc(W.))
[f: y - J G, (f: -y)ll2
[[, -J,441 J. G (f: -y)
S [[, JJ,,441lll | 11. G, || ||f - yl2-
(17)
Then control the first term in Eq. 117 in way similar to the following:
J Jr,t11l2 S
|W W S
J
||,
lla [ff yll2.
VM
VM
and if || 1| G. || can be upper bounded, we get our result Eq.
(14.
For the mini-batch version, similarly we have
f41 - = f. y Ju J J., Gr, (fr.B, - YB.)
f, y J. J L.B. G. (fe y),
(18)
where the subscript B, denotes the sub-matrix/vector corresponding to the batch, and G E bxx zero-padded version G to make Eq. (18 hold. Therefore, after one epoch (from (tk + 1)-th to the ((t + 1th update), we have
((-1)k y
II
JI-Je Gilee)) (fix () Ar(fit y)-
(=+++1))-1
7

---

Under review as a conference paper at ICLR 2020
We will see that the matrix A, is close to the matrix A defined in 13 so it boils down to analyzing the spectral properties of A.
For both theorems, we can compute that as M increases, the norm of the update I[ W l does not increase with M, so the update is small compared to the Gaussian initialization where [Woll2 = @(VM). From this we can derive that the matrices J. G etc. remain close to their initialization, which makes bounding their norms possible. The full proof is in the appendix.
C
In conclusion, the accelerated convergence is related to the local linearity and the stability of the Jacobian and Gram matrix. We emphasize that our theorems serve justification of our GGN algorithm, because we expect that GGN works in practice, even under milder situations when M architectures, and that GGN would still perform much better than first-order methods.
motivation than
large as the theorem demands for deep networks with different
3.3 ANALYSIS OF THE PER-ITERATION COMPUTATIONAL COMPLEXITY
We have proved that for sufficiently overparametrized deep neural networks, full-batch GGN has quadratic convergence rate. In this subsection, we analyze the per-iteration computational cost of GGN, and compare it to that of SGD.
For every mini-batch (i.e., iteration), there are two major steps of computation in GGN:
(A). Forward, and then backpropagate for computing the Jacobian matrix J.
(B). Use J to compute the update J (AG + (f-))
We show that the computational complexity of (A) is the same that of SGD with the same batch size; and the computational complexity of (B) is small compared to for typical networks and batch sizes. Thus, the per-iteration computation overhead of GGN is very small compared SGD. Overall, in terms of training time, GGN can be much faster than SGD.
For the computation in step (A), the forward part is just the same that of SGD. For the backward part,for every input data, GGN keeps track of the output's 's derivative for the nodes in the middle of the computational graph. This part is just the same as backpropagation that GGN also, for every inpul data, keeps track of the output's 's derivative for the parameters; while in SGD the derivatives for the parameters are averaged over a batch of data. However, it not difficult to see the computational costs of GGN and SGD are the same.
SGD. What different is
For the computation in step (B), observe that the size of the Jacobian is b x m where b is the batch size and m is the number of parameters. The Gram matrix G.,B, = L.B. in our Gram- Gauss-Newton method is of size x b and only requires O(b²m + b²) for computing Gr n, and matrix inverse. Multiplying the two matrices to f y requires even less computation. Overall, the computational cost in step (B) is small compared that of step (A).
EXPERIMENTS
Given the theoretical findings above, in this section, compare proposed GGN algorithm with several baseline algorithms in real applications. In particular, we mainly study two regression tasks, AFAD-LITE (Nu al. 2016) and RSNA Bone Age (rsn) 2017).
4.1 EXPERIMENTAL SETTING
AFAD-LITE task is to predict the age of human from the facial information. The training data of the AFAD-LITE task contains 60k facial images and the corresponding age for each image. We choose ResNet-32 (He et al.) 2016) as the base model architecture. During training, all input images are resized to 64 64. We study two variants of the ResNet-32 architecture: ResNet-32 with batch normalization layer (referred to ResNetBN), and ResNet-32 with Fixup initialization (Zhang et al.) 2019b) (referred to as ResNetFixup). In both settings, we use SGD as our baseline algorithm. In particular, we follow Qian 1999) to use its momentum variant and set the hyper-parameter Ir=0.003 and momentum=0.99 determined by selecting the best optimization performance using grid search. Since batch normalization is computed over all samples within mini-batch, it is consistent with
8

---

Under review as a conference paper at ICLR 2020
0.5
0.5
ReN S S0D+HRaNNBBN
GON ResNedFixup
0.4
0.
SOD1 ResNeRN
Training Less
Training Loss
0.3
0.3
0.2
0.2
0.1
0.1
0.0
0.0
20
100
Epochs
(a) Loss-time curve on AFAD-LITE
(b) Loss-epoch curve on AFAD-LITE
0.30
0.30
GGN- ceNtRRsNetFlaup
GGN-ResNeel ResNedFixsp
0.25
SGD ResNeeFiop
0.25
SGDI Ned'up
SGD BEDDRReNNeINN
SGDI ResNeLBN
Training Loss
Training Loss
0.0
0.5
1.0
1.5
2.0 Epochs
2.5
3.0
3.5
4.0
Time
(c) Loss-time curve on RSNA Bone Age
(d) Loss- curve on RSNA Bone Age
Figure 1: Training curves of GGN and SGD on two regression tasks.
our assumption in Section 2 that the regression function has the form of f(w, x), which only depends on w and a single input datum x. For this reason, the GGN algorithm does not directly apply to ResNetBN, and we test our proposed algorithm on ResNetFixup only. We set À and = 0.3 for GGN, We follow the common practice to set the batch size to 128 for our proposed method and all baseline algorithms. Mean square loss is used for training.
RSNA Bone Age task is part of the 2017 Pediatric Bone Age Challenge organized by the Radiological Society of North America (RSNA). contains 12,611 labeled images. Each image in this dataset is radiograph of left hand labeled with the corresponding bone age. During training, all input images are resized to 64 * 64. We also choose ResNetBN and ResNetFixup for this experiment, and use ResNetBN and ResNetFixup trained in the first task warm-start initialization. We use Ir= 0.01 and momentum= 0.9 for SGD, and use X = 1 and a = 0.1 for GGN. Batch size is set to 128 in these experiments, and mean square loss is used for training.
4.2 EXPERIMENTAL RESULTS
Convergence. The training loss curves different optimization algorithms for AFAD-LITE and RSNA Bone Age tasks are shown in Figure T On both tasks, our proposed method converges much faster than the baselines. We can see from Figure Ta and Figure that, on the AFAD-LITE task, the loss using our GGN method quickly decreases to nearly zero in 30 epochs. On the contrary, for both baselines using SGD, the loss decays much slower than our method terms of wall clock time and epochs. Similar advantage of GGN can also be observed on the RSNA bone age task.
Generalization performance and different hyper-parameters, We can see that our proposed method trains much faster than other baselines. However, as a machine learning model, generalization performance also needs to be evaluated. Due to space limitation, we only provide the test curve for the RSNA Bone Age task in Figure 2a From the figure, we can see that the test loss of our proposed method also decreases faster than the baseline methods. Furthermore, the loss of our GGN algorithm is lower than those of the baselines. These results show that the GGN algorithm can not only accelerate the whole training process, but also learn better models.

---

Under review as a conference paper at ICLR 2020
0.30
0.20
GON ReeNed/ssup
0.175
1.1.
ResNesPloup
1-12-1
SCID-ReuNoIBN
0.150
0.20
0.125
Training
0.100
0.30
0.075
0.050
0.05
0.025
0.000
2
10
Time
Epochs
(a)
(b)
(c)
on RSNA Bone Age dataset. (a) Test curves of GGN and SGD. (b) Training curves of GGN with different configurations. The optimal is shown for X = 0.25 and X = 4 by grid search. (c) Training loss at 10th epoch of models trained using GGN with different hyper-parameters.
Figure 2: Test performance and ablation study on hyper-parameee
We then study the effect of hyper- perammttsseeters used the GGN algorithm. We try different À and a on the RSNA Bone Age task and report the training loss of all experiments at the 10th epoch. All results are plotted in Figure 2c In the figure, the x-axis is the value a. The gray value of each point corresponds to the loss, the lighter the color, the higher the loss. We can see that the model converges faster when À is close to 1. In GGN, a can be considered as the inverse value of the learning rate in SGD. Empirically, we find that the convergence speed of training loss is not that sensitive to a given a proper X, such as À = 1. Some training loss curves of different hyper-parameter configurations are shown in Figure 2b
and the y-axis is the value of
5 CONCLUSION AND DISCUSSIONS
We propose a novel Gaee--Gausss-eeeee (GGN) method for solving regression problems with square loss using overparameterized neural networks. Despite being a second-order method, the computation overhead of the GGN algorithm at each iteration is small compared to SGD. We also prove that if the neural network sufficiently wide, GGN algorithm enjoys quadratic convergence rate. Experimental results two regression tasks demonstrate that GGN compares favorably to SGD on these data sets with standard network architectures. Our work illustrates that second-order methods have the potential to compete with first-order methods for learning deep neural networks with huge number of parameters.
In this paper, mainly focus the regression task, but method be easily generalized to other tasks such as classification as well. Consider the k-category classification problem, the neural network outputs vector with it entries. Although this will increase the computational complexity of getting the Jacobian whose size increases k times, i.e., J E (bkk)xm each row of J can be still computed parallel, which means the extra cost only comes from parallel computation overhead when calculate in fully parallel setting. While most first-order methods for training neural networks can hardly make use of the computational resource in parallel or distributed settings to accelerate training, GGN method can exploit this ability. For first-order methods, basically extra computational resource can only be used to calculate more gradients at a time by increasing batch size, which harms generalization lot. But for GGN, more resource can be used to refine the gradients and achieve accelerated convergence speed with the help of second-orde information. It is an important future work study the application of GGN to classification problems.
REFERENCES
Rsna pediatric bone age challenge. http://maaaaaaleeesscoooooooooor 2017.
Joshua Achiam, Ethan Knight, and Pieter Abbeel. Towards characterizing divergence in deep q-learning. arXiv preprint arXiv: / 2019.
10

---

Under review as a conference paper at ICLR 2020
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. arXiv preprint 1xXiv:88111.49115 2018a.
Zeyuan Alu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. arXiv preprint arXiv:/811.03900 2018b.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. arXiv preprint arXiv: 1955, 2019a.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterizze two-layer neural networks. arXiv preprint arXiv:1901.0858,, 2019b.
Sue Becker, Yann Le Cun, et al. Improving the convergence of back-propagation learning with second order methods. 1988.
Alberto Bernacchia, Máté Lengyel, and Guillaume Hennequin. Exact natural gradient in deep linear networks and its application to the nonlinear case. In Advances in Neural Information Processing Systems, pp- 15450, 2018.
Aleksandar Botev, Hippolyt Ritter, and David Barber. Practical gauss-newton optimisation for deep learning. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, pp. 557-565. JMLR. org, 2017.
Searning-Voluuee
Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. arXiv preprint arXivv190S.13210, 2019.
Lenaic Chizat and Francis Bach. A note on lazy training in supervised differentiable programming. arXiv preprint arXiv: 1812:0180 2018.
Xialiang Dou and Tengyuan Liang. Training neural networks as learning data-adaptive kernels:
Provable representation and approximation benefits. arXiv preprint arXiv:1901.07114, 2019.
Simon Du, Jason D Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. arXiv preprint arXiv: 18/1 10111111110044. 2018a.
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterizee neural networks. arXiv preprinl arXiv: 180.00044 2018b.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(1u1:2121-2159, 2011.
Rong Ge, Furong Huang, Chi Jin, and Yang Yuan. Escaping from saddle points - online stochastic gradient for tensor decomposition. In Proceedings of The 28th Conference on Learning Theory, pp. 797-777-889 842,2822222 2015.
Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. Fast approximate natural gradient descent in a kronecker factored eigenbasis. Advances in Neural Information Processing Systems, pp. 9550-9560, 2018.
Gene Golub. Numerical methods for solving linear least squares problems. Numerische Mathematik, 703):206-216, 1965.
Gene H. Golub and Charles F. Van Loan. Matrix computations (3rd ed.). Johns Hopkins University Press, Baltimore, MD, USA, 1996. ISBN 08018540 1544884448.
Roger Grosse and James Martens. A kronecker-factored approximate fisher matrix for convolution layers. In International Conference on Machine Learning, pp. 573-582, 2016.
Moritz Hardt, Benjamin Recht, and Yoram Singer. Train faster, generalize better: Stability of stochastic gradient descent. arXiv preprint arXiv:1509.0122 01240, 2015.
Kaiming He. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition. pp. 770-778, 2016.
11

---

Under review as a conference paper at ICLR 2020
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization neural networks. In Advances in neural information processing systems, pp. 8571-8580 2018.
Chi Jin, Rong Ge, Praneeth Netrapalli, Sham M. Kakade, and Michael I. Jordan. How escape saddle points efficiently. In Proceedings the 34th International Conference on Machine Learning, pp. 1724-1732, 2017a.
Chi Jin, Praneeth Netrapalli, and Michael I Jordan. Accelerated gradient descent escapes saddle points faster than gradient descent. arXiv preprint arXiv: 2017b.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv: /609:::660.0000 04836, 2016.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 22.6980 2014.
Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Jaseha Sohl-Dickstein and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. arXiv preprint arXiv: /902.06600 2019.
Kenneth Levenberg. A method for the solution of certain non-linear problems in least squares. Quarterly of applied mathematics, 2(2):164-168, 1944.
Tengyuan Liang and Alexander Rakhlin. Just interpolate: Kernel" ridgeless" regression can generalize. 0038,,0387.000 2018.
arXiv preprint
Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. arXiv preprint 1xXiv:1908.03265 2019.
Liangchen Luo, Yuanhao Xiong, Yan Liu, and Xu Sun. Adaptive gradient methods with dynamic bound of learning rate. arXiv preprint arXiv: 1o.X:v:992200844 2019.
James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In International conference on machine learning, pp. 2408-2417, 2015.
James Martens, Jimmy Ba, and Matt Johnson. Kronecker-factored curvature approximations for recurrent neural networks. 2018.
Dominic Masters and Carlo Luschi. Revisiting small batch training for deep neural networks. arXiv preprint arXiv:1804.07611 0766122,,0.6..... 2018
Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-fre bounds and kernel limit. arXiv preprint arXiv:1902.0601 09001000............ 2019.
Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar. Foundations of machine learning. 2018.
Wenlong Mou, Liwei Wang, Xiyu Zhai, and Kai Zheng. Generalization bounds of sgld for non-convex learning: Two theoretical viewpoints. arXiv preprint 1707.007 2017.
Zhenxing Niu, Mo Zhou, Le Wang, Xinbo Gao, and Gang Hua. Ordinal regression with multiple output cnn for age estimation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 4920-4928, 2016.
Samet Oymak and Mahdi Soltanolkotabi. Overparameterizee nonlinear learning: Gradient descent takes the shortest path? arXiv preprint arXiv:1812. 2018.
Razvan Pascanu and Yoshua Bengio. Revisiting natural gradient for deep networks. arXiv preprint arXiv: 2013.
Ning Qian. On the momentum term in gradient descent learning algorithms. Neural networks, 12(1): 145-151, 1999.
12

---

Under review as a conference paper at ICLR 2020
Yi Ren and Donald Goldfarb. Efficient subsampled gauss-newton newton and natural gradient methods for training neural networks. arXiv preprint arXiv:19060025:: 02353, 2019.
Roman Vershynin. Introduction to the non-asymptotic analysis of random matrices. arXiv preprint arXiv:1011.3027, 2010.
Martin J. Wainwright. High--Dimensionl Statistics: A Non-Asympiotic Viewpoint. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2019. doi:
10.1017/978 108627771.
Yuhuai Wu, Elman Mansimov, Roger B Grosse, Shun Liao, and Jimmy Ba. Scalable trust-region method for deep reinforcement learning using kronecker-fie factoored approximation. In Advances in neural information processing systems, pp. 5279-5288, 2017.
Yuxin Wu and Kaiming He. Group normalization. Proceedings of the European Conference on Computer Vision (ECCV), pp. 3-19, 2018.
Jinchao Xu. The method of subspace corrections. Journal of Computational and Applied Mathematics, 128(1-2):335-66, 2001.
Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint arXiv: 1902:0990 2019.
Guodong Zhang, James Martens, and Roger Grosse. Fast convergence of natural gradient descent for overparameterized neural networks. arXiv preprint
2019a.
Hongyi Zhang. Yann N Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normalization. arXiv preprint Seti::1901.0932.. 2019b.
Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise in stochastic gradient descent: Its behavior of escaping from minima and regularization effects. arXiv preprint arXiv: 1803.00 2018.
Difan Zou and Quanquan Gu. An improved analysis of training over wve-pparameterized deep neural networks. arXiv preprint arXiv: 04688, 2019.
Difan Zou, Yuan Cao. Dongruo Zhou, and Quanquan Gu. Stochastic gradient descent optimizes over-parameterized deep relu networks, arXiv preprint arXiv:1811.08888, 2018
A SOME PREPARATIONS FOR THE PROOF OF THEOREMS
Notations.
We use the following notations for the rest of the sections.
Let [n] = (1, (1,...,,..
Jw.x Rd denotes the gradient 8/(Wx) which is of the same size of W.
JW
The bold Jw or J(W) denotes the Jacobian with regard to all data, with each gradient for W vectorized, i.e.
vec(Jw.x))
Jw
(Md).
(19)
vec(Jw.x
w, denotes the r-th row of W, which is the incoming weights of the r-th neuron.
Wo stands for the parameters at initialization.
dw.x == o' (Wx) E Mx denotes the (entry-wise) derivative of the activation function.
We use ( to denote the inner product, ||:11, to denote the Euclidean norm for vectors or the spectral norm for matrices, and || to denote the Frobenius norm for matrices.
13

---

Under review as a conference paper at ICLR 2020
We easily derive the formula for J as where o the point-wise product. So we can also easily solve G as
Jw
((dw,x o a)x )
(20)
VM
G5 = (Jx,,
Σ x
(21)
M
Our analysis is based on the fact that G stays not too far from its infinite-width limit
((,, X2) = (( (x] xyo' (wo,(w )o (wxx))
which is positive definite matrix with least eigenvalue denoted Åg, and we assume Ào Ào is small data-dependent constant, and without loss of generality we assume Ào S 1, or else we can just take A 1 Amia ( >
The first lemma is about the estimation of relevant norms at initialization.
Lemma (Bounds on Norms at Initialization). If M Ω (dlogg16n)) (6)//)). then with probability at least 1 - 8/2 the following holds
(a). |Woll2 = O(VM).
(b). f(Wo, o, x,) = O(1), e [n].
(c). || II. O(1). for [n].
The lemma is a well-known result concerning the estimation of singular values of Gaussian random matrices (see Corollary 5.35 of Vershynin (2010)). Notice that Wa RMxd is Gaussian random matrix, the Corollary states that with probability 2 //2 one has
Proof. (a).
Woll2 S VM v +t.
By choosing M = max(d, /2 (8/6),, we obtain Wo S 3vM with probability 1 8/4. (b). First, a,, E [M] are Rademacher variables. thereby
with probability 1 - 2e Mt? /2 we have 17 Σm a, I t. This means we take M f (log(16/6)).
M Pr Σ
5
a, - 2
(22)
Next, the vector T WoX; E Mx1 is standard Gaussian vector. Suppose the activation (.)) is (-Lipschitz and is ((1) by our assumption, with the vector fixed, the function
: R - +
V,
a ((xxx v.) = S(W,,(WWW. x;)
VM
has Lipschitz parameter of / ||x ll / O(IVMM). According to the classic result on the concentration of a Lipschitz function over Gaussian variables (see Theorem 2.26 of Wainwright (2019)), have
Mt²
Pr[[o(v,) - Ew. 2exp
2/2
||x:|?
which means if M ((lo(166n8)).
M
1 a To (Wox))
VM a, Www
=0(1)
(23)
VM
holds jointly for all i E [n] with probability 1 - 5/8. Note that
[ww- < [o( + 1 E (111 0(1).
0(1).
(24)
14

---

Under review as a conference paper at ICLR 2020
Plugging in 22 and 24 into 23 1 6/4, for all E |
that long as M [((oo(16n/6)), then with probability
J(Wo: x;)
VM - o(Wox))=-Op
(c). Since a Oip)-sipschizz we have ||dw x. 0l.e = 0 According to (20 we can easily know that |l 1pp S VM [Diag(d)]l
||ll2 | = O(1)
The next lemma is about the least eigenvalue of the Gram matrix G at initialization. It shows that when M is large, Gw, is close to K and has a lower bounded least eigenvalue. It is the same as Lemma 3.1 in Du et al. (20T8b), but here restate it and its proof for the reader's convenience.
Lemma 3 (Bound on the Least Eigenvalue of the Gram Matrix Initialization). If the width
1og(2n/6) then with probability at least 1 - 6/2 over random initialization, we have Ag
M = Ω
Amin (w.) Ao.
Proof. Because is Lipschitz, o'(wx,) (owx)) is bounded by O(1). For every fixed (i,j) pair, at initialization G.y is an average of independent random variables, and by Hoeftding's inequality, applying union bound for all n² of (i,j) pairs, with probability 5/2 at initialization we have
(v /og(2n//)) M
Gij Kiyl 50
and then
n12 (12 log (oloo((n) (en//) (0(
Gw.
K S
K| =0
M
Thus if M = Ω (n² log (2n/6) Ag
we can have |ww, K112 < do and thus Amin (W,,)
2 Ao
Next, in Lemma and 5 we will bound the relevant norms and the least eigenvalue of G inside some scope of W that covers the whole optimization trajectory starting from Wo. Specifically, consider the range
BRR) A (W: | W Woll S R).
where is determined later make sure that the optimization trajectory remains inside (R)). The idea of the whole convergence theorem is that when the width M is large, R is very small compared to its initialization scale: ||Woll, = O(VM). This way, neither the Jacobian nor the Gram matrix changes much during optimization.
Lemma 4 (Bounds on Norms in the Optimization Scope). Suppose the events in Lemma hold. There exists a constant C' > 0 such that if A 2 CRR., we have the following:
(a) For any W BR)), we have
[W O(VM).
(25)
(b) For any W1, W2 B(R), W
W2llF < R', then have
||ww.x JWaxx 1F 0
R'
Vi, and [Jw.. -J Jwall2 00
n
(26)
VM
M
Also, for any W E B(R), we have
|| |Ir 011 and |[ =01l)=
(27)
15

---

Under review as a conference paper at ICLR 2020
Proof. (a). This is straightforward from Lemma a)), the definition of ((R), and M -2(R)).
(b). According to the ((1) smoothnsss of the activation, we have
, dw,,x. l1 011) |Wix; W2xill2 = O(
so we can bound
|| Jwaxille
[Diag(a)(dwixx - dw2,x,x I
VM
1
[Diag(a)]l ||d dwa.x1 llz |xxill2
VM
R'
<
VM
And according 19
we have
[Jw, Jwall2 S Jww, Jwzlle
|Jwix Jwa.x
O
R'
M
Also, taking W1 = W and W₂ = Wa. combining with Lemma c.. we see there exists such
that for M 2 CR have |.w | (1), and naturally [Jw O(yn). Note: The constants hidden in the O(-) notation are irrelevant with C.
The next Lemma deals with the Gram matrix within B(R), ensuring lower bound on the least eigenvalue throughout the optimization process (Du et al. 2018a).
Lemma (Least Eigenvalue in the Optimization Scope). For W E B(R), suppose the events in
Lemma and 3 hold, there exists constant C' such that for M 2 C'n2 we have
X/
Ao
Gw Gwall2 S
4
and thus combined with Lemma we know that Gw remains invertible when W B(R) and satisfies || 2 Ag
Proof. Based on the results in Lemma
b), we have
|w Gwo 2
|| Jw Jwol2
(28)
$ [Jw Jw. Jw.JJWlll J.. (JW - Jw)ll
(29)
nR
s 0
(30)
VM
To make the above less than A. choosing M greater than some
suffices, and the lemma is thus proved. Again the constants hidden in the O--) notation are irrelevant with C².
C' R X
B PROOF OF THEOREM 1
Proof idea. In this section, we use W., E {0. 1..... to represent the parameter W after iterations. For convenience, J,, G fe is short for Jw,: Gw,, fw, respectively. We introduce
Jr,41 ( J((1 s)W, sW Ws+1)ds.
(31)
16

---

Under review as a conference paper at ICLR 2020
For each iteration, if G2 is invertible, we have
f++1 = ft - y + Ja ved W Wi)
= f; - J Ju G yy
(J: J., G, (f. y),
hence
[641 yll2 S [J Je.+il2 112 J 1l G. 112 [I.. - -yll- Then control the first term of the right hand side based Lemma 4 in the following form
(32)
[J, Je,till2 s
Iw W
VM
J_ G (f: - -y>ll2
VM
J
G [. -
VM
and plugging into 332 along with norm bounds on J and G we obtain second-order convergence.
in Lemma and (the constant chosen to make the right hand side of hold). We prove that there exists an easily verify
Formal proof. Let R |W: We+illr for / E {0, J. We take R= θ
A6 )) (with enough constant) that suffices. First that all the requirements for M in Lemma 2-5 can be satisfied. Hence, with probability at least 1 - all the events Lemma 2-5 hold. Under this situation, do induction on show the following:
M Ω max
n²4 x (0666 Ag
(a). W, B(R).
(b). I > 0, then [[f, yll2 S
n/2
[. - y
XBMM
As long as (b) is true for all t, then choosing M large enough to make sure the series {l[f: y converges to zero, we obtain the second-order convergence property.
For = 0, (a) and (b) hold by definition. Suppose the proposition holds for t = 0.... .T. Then for
I = 0, ... T, G2 is invertible. Recall that the update rule is veo (W++1) = vecl (W)) -J. (f.
y), we have
R = w, W IF
|vve W.) vec( (el+11l
[]] ll ||, 112 |lfe y
0 Vn
l y1l2
(Lemma and
(33)
Ao
According to Lemma b) and the assumption that the target label yı O(),, have [[o -12 (n)). When T > 1, the decay ratio at the first step is bounded as
[f yll [fo y
so
[
yll2 s 0o
n
=tr,
X
M
and taking M = 2 with enough constant make is constant less than 1, in which the ordond-order convergence property (b) will ensure a faster ratio of decay at each subsequent step in |f, y - Combining 333 we have
x
T R
8 =0
5
ff.
y1l2
50
n
p--l
n2 Ao
< R.
(34)
Ao
Ag
1=0
1=0
t=1
17

---

Under review as a conference paper at ICLR 2020
Therefore, WT+1 B(R), and (34) also holds when T 0, (a) is true.
Since (R) convex, this means we also have s &W++ (1 - WT+ E B(R) for E [0, 1]. Hence can bound the difference of the Jacobian as
JT Jrll2 -
(1- -
0
sRt
ds
(Lemma 4b)
M
0
n2
||[T
y112
(Using 33
(35)
ADVM
M
So we use the bound of (32 and obtain
[[T+1 - y S Jr Jr
112
|
GT yyl
3/2
||f
(Using (35) and Lemma 4
V M
which proves (b). This concludes our proof.
C PROOF OF THEOREM 2
In this section, we give the proof the convergence of our mini-batch GGN algorithm (Theorem 2
Some additional notations. For the convenience of our proof, will slightly different notations than that in Section 3.2 Let n. bk, where b is the batch size, A: is the number of batches, equivalently, number of updates in each epoch. For epoch
21,2..... and batch i E [k]. we will use Wt E [R to denote the parameters before the i-th update in epoch t. Let Jus J(W,s) 1s) Ra G G(W,s)
R=** and fli f(W..) Rxx (All of Wo. W1, WI1. Go. etc. can represent initial values in the proof.) Also, let Wt = Ja (1+1)1 etc. the i-th batch, we use the data (x1, 31 )1 (11===-(10000 1))+11 To specify the matrix related to a batch, we use the following indexing method: For i', i" [k]],
G. G. Jr,, f, fei, and let W W((++1) W (
Jus Ju((t' 1)b+ 1: i'b, 6, 1 Md) (Mx()))
which is the rows of Jes that represent the Jacobian of the i-th batch. Similarly, we have
G..e Ju, Roxb
fu, = f.,(('' 1)6 1: ib) 1x1
yæ ((((1116+111b) E bx
Similar to 311 in the proof of Theorem
I we make use of the notation
J Ju = L/' J(s is + (1 - s) (sss)))ds
(36)
and similarly define J J((,5+1),5 E R.bxMM etc. In addition, we make use of the matrix
G Obx
G 1G.
0bx(k-910
gdxn
in the formula of the update.
Similar the proof of Theorem the idea of our proof that due the fact that the change of W is small compared to the scale of its initialization, the matrix J and G remains stable. Informally, this makes the algorithm close to solving kernel regression iteratively by batches, which is equivalent to solving system of linear equations
JJ - f(Wo) y,
(37)
18

---

Under review as a conference paper at ICLR 2020
for variables R=x] means solving
(where vee( W) vec(Wo) J r), using the Gauss-Siedel method, which
(oki) (1 (i - 1)b)
JJ
plow (( 11-1 1: :ibb))
- (r(Wo - (1))6-1))-1110
plold)((i + 1: kb))
for the i-th batch in every epoch. Therefore, it is natural that the matrix A = L. (D-L)-+ in 15 introduced. We will show later that the real update follows
f++
(fe -y) y)
for some A, ~ A.
In order prove the theorem, need some additional lemmas follows.
Lemma 6 (Formula for the Update), If the Gram matrix each step is invertible, then:
(a) The update of W is
vec W Wx(++1) = vec W (x) J (fu, - ys)
(38)
-voc(W.t) - JJ, Gr yy)
(39)
(b) The formula for I is
Gu)(ff (((+1) - y= (I, Ju G. yy)
(40)
(c) The update of satisfies
f((+t)) - --U-DD L,) -,) (f. - y),
(41)
where
Ga.
O
0
G
J
0
D.
L
0
0
Gek
J J n.1
J(22,) 12,
0
and
J((1.)), J
G.,
J((23). 12.2
21,1
J(22.)) 12,2
G
J ti
Un Diag
Jn G sbxab
G+ 11041,22),442 2 (041)441
G((42))(42)(22)
0
J J001142) 1041).41
31(42.44)) (7442)
G (k-
Or, if we denote
Un (E Ribxb))
0
Un
Un(E ((
Dis La (E Ribxib D Ln
DH Lis
D - Lule Rlk ijxx (6x-
then we have
U Dr L
0
f(+41) -
(fr -)-
(42)
(D (D): La)
I(c-
19

---

Under review as a conference paper at ICLR 2020
Specifically, we have
f++1 y y - L,)) -y =: A.(f: y),
(43)
where
J 11,1
Ga,1
J J((2,3), J 12
J Jick ce 1+1, 1 bk
0
J 21223))2 12,2
222
J ((k,)+1), thek
U, -
0
1( & bk.k
GER.RR
Proof. (a) For (38), this is exactly the update formula (TO) for the i-th batch where the Jacobian and Gram matrix are Jui and G respectively. Note that
G
Caa y.) = [0xx
G (bx(x-))b) (fes - y),
we then obtain
33
from 38
(b). We have
(fi(1-1) - yy)
= ((44) - fui
- vec(W(.))
Jig Gu.lfe- y)
(By formula (39)
we obtain (40.
(c). Based on
40
we know that
f(++ - y II (L. Jw Gry) - y);
Y'=í
where the index goes in decreasing order from left to right. So in order to prove (41 we only need to prove that
a - Jup we Gev) (D: L() Utis
(44)
Y=i
which we will prove by induction on i. For i = 0 it is trivial that D, L, Un by definition.
Suppose (44) holds for then
II (L, J T
(D, - Lt)
-=i+1
(I, J Ja Mi+ 11 Gu
Un J (0 G-1 (04))00)04))) Un((1i+ (i 1)b,1:
J J( (041)411
J J104
(041).41
-U+
[bx(()) Lxb 0px(k-
J J (11))441
=Ua++))
which proves
41
Note that by the definition, have U Dr
Lio then obtain (42 by
U
D.
Lu
0
-1
U, - -1
U 15
D
Lu
D
L
Ün
(D
L -1
Un
DD L,,)
L (D.) -2
(Drs Lu) -
(n.
L,s)
(D&
(D Led)
-1
1
20

---

Under review as a conference paper at ICLR 2020
that the iteration matrix A, is close the matrix A defined in 15 The convergence of the algorithm much related to the eigenvalues of A. In the next two lemmas, we bound the spectral radius of A and provide an auxiliary result the convergence on perturbed matrices based on the spectral radius.
By 45
Lemma 7. Suppose the least eigenvalue of the initial Gram matrix Go satisfies Amin (Go) 2 3
4 Ao.
( (" X
(n" log (2n/8) according to Lemma 3) Also assume
(which is true with probability 5 if
o le = ((1) for any E [n]] (which true with probability 1 6, according to Lemmal 2). Then the spectral radius A, or equivalently, maximum norm of eigenvalues of A, denoted p(A), satisfies
(A) 11-11
(45)
Proof. For any eigenvalue À E C of A, it is an eigenvalue of T so there exists a corresponding unit vector v E Caxl such that λv = A A (D ) Lv, which means
X(D Jv - =.
(46)
Since Go = - is positive definite with eigenvalues at least 1 Ao we have v (D-L- L L 2 Ao which means
-D--L
3 d 2-2 2 =>0.
(47)
where d= v'Dv, 1 = v*I,, and V = Let v, = v(ii - 1)b +1 : ib) E Rx be the i-th
batch of the vector v, and v, 101x(1 1)6: v. 11x(x=(x) 19b] E [_ax1 is not hard see that
d= v Dv
v Go
Σ Gov 2
l.
3
4 Ao
(48)
1=1
==1
Also, since by our assumption each entry of L (or of Go) is at most 0(1), have
| v'
|اا|| | O()).
(49)
Now we use take an inner product of with 46) and get
0 v'(v LTv Lv) X(d 1) +
(50)
and therefore
(By solving 50
d
d
1112
(d
l) - 1)
did-2 (ed)2R((1)
5
(Using 47
48
Ω
+1
1 Ω
This concludes the proof for this lemma.
Lemma 8. Denote p(A) Po < 1. Let A be diagonalized as A = =P-IQP and |P]|, P
(see Assumption 2 Suppose we have A Aoll2 S for 1 < T, then
II A
S sulp + 18)TT
21

---

Under review as a conference paper at ICLR 2020
Proof. We have
T
T
A
II (A + (A, - Aou)
i=1
II P (Q++((A Ao)P-1)P
i=1
II (Q (Q)P(A-AO)PP)) P
T
Plll
II (IQl12 + |Pl|2 |A. Aol1 [Pll)a) |P-1/2
julp n
In addition to the bounds used in Theorem 1 we provide the next lemma with useful bounds of the norms and eigenvalues of relevant matrices in the optimization scope
- (RR) = (W [WWWllle- - R).
Lemma' (Relevant Bounds for the Matrices in the Optimization Scope), We assume the events in Lemma and hold, and let M Cp²n"R" H²
for some large enough constant C. then:
Ag
For any (t. i) pair, we assume Wes E B(R) all i' E [k] when t' s and i' E when t' = in the following propositions.
(a). minver", ,|M[;= min.ER".|M, (I(D
Llvll2 2 Ao
(b). Suppose up to f, W, is in ((R) (which means for i' E [k + 1), and t' E [t- 1). Wa E B(R))
Then ||A -l12 L-p(A)
2μ
Proof. (a). Because for ||v||2 1,
[ID L)vll2 2 v v(D L)v
1 v
Dv + v VD - LD-LLL"JJ
2 0+
Amin (Go)
2
3
IV
Ap-
(By the positive-definttenss of Go and Lemma
(b). By Lemma 4 we know that within BRR), the size of the Jacobian || ll w.r.t. data x; is O(1),
and the difference || Jw.. |p within O(R) is bounded by O R to each entry of D, L. U, etc., including those (1+1)) terms by the convexity of B(R), and we that each entry of these matrices has size at most ((1) and varies inside O Therefore we get
VMM this can be applied
R
range.
VM
J = =(n).
DD, Lell = O(n).
and
|U, - U1l 0 Rn
VM
[D. L.) (D L)lla = Rn
VM
22

---

Under review as a conference paper at ICLR 2020
By our chosen M, O
Rn VM
Ag can definitely hold, so along with (a) we have
-DD -1- =0 ( [D, - -.) = (Nt)
Hence,
|| All =
- - -)) 1 1|2
(U, - UU - L.) + U((: - -L); - L,,) - (D - L)) - (0) 112
Rm
O 1
1
Rn
1
(0
0
0
VM
Ao
Ao
VM
Ao
n.? R
=0
X
and along with Lemma 7
order to have
||A,
A S
I-p(A) 2μ
all we need
n² R
X2
Ω
so choosing some M
Cy² Ag
suffices.
With all the preparations, now we are ready prove Theorem [ same as what we did in the formal proof of Theorem [ where we then used induction on t and now we use induction on the pair (t, i). The key, is still selecting some that in each step Wa: remains in (R).. Combined with the previous Lemmas, in ((R) we have A, being close to A, and then convergence is guaranteed.
The logic of the proof is just the
Formal proof of Theorem
2
Let Ru
W (We(++) Wull
We take
n2
R= θ
X
in the range (R) (where the constant is chosen to make the right hand side of (52 hold). We prove that there exists an
M = max (9) 122n 1
n 2d moddlog(66n//
Ω
X
with large enough constant that suffices. First we can easily verify that all the requirements for M in Lemma 2-9 (most importantly, Lemma() can be satisfied. Hence, with probability at least - 8 all the events in Lemma 2-9 hold. Under this situation, we do induction on (f, (1 E (1,2, ... }, E [k)] by dictionary order) to show that:
Wis E B(R).
For (1, = (1, it holds by definition.
Suppose the proposition holds up to (t, Then since W E B(R), by Lemma 4 we know that
Amin( (G.)) > de This naturally gives Amin (Gr,ss) > 2
Ag which means G is invertible.
2
Similar to the argument Ja, J1,,41)) D.s, Li. Un, etc., whose index of (t',"") only contains i' with i' S i, is of scale
the proof of Lemma
we know that each entry of
O(1) and varies at most O n VM
which is result of Lemma Then know
[Ju,ll = 0(Vb) = O(vn).
D - I((e-1) =0(n)
D - E O(n),
23

---

Under review as a conference paper at ICLR 2020
etc. We then also know
(Det-1) - L(((-))) - (D, - Le) 17010 0
nR
VM
Since we can also have II((D, L₂) (D
(1: F S
aR
and given our
VM
choice of M and R the right hand side is less than 8 Aa. along with Lemma (
(which says
min
LDD-Lyyly 2 8 Au
VER" 1
and naturally, since this matrix is block-lower-triangular,
||D - L)): 1010 2
min VER /Vl2
holds) we know that D, Le is invertible and has the bound
(D, - -b,) 1/2 $ 4
Based on the update rule 38 we have
Rss = | W (((t)))l
||vec[ voc( Wl11)llz
25.55 (fes. - yi) l
( G 0 (0
0
|fe: -o-yy yll1
= 1
Au
No
vT
U1i-1)(D - Le(a-1))-1
[[f y/l2
Ao
1D) (D:((-1) - (D L((--1))-1
(
(By formula (42)
vn
00
Ú (D - L((a-1))
Au
(D (Dec-1))
(D(L-))(((ET)) - L(((-1)) -1
(p) Ilfe -yllz
3/2
|ff: - -
(51)
A
3/2
II i=1
A
||f] -yl1
(By Eq. 433
X
1/2
3/2
(1-1)
(1
(l( - 22
(Lemma and 9 (b).)
X
n²
pn?
2
(t-1)
00
(Lemma
n²
where Eq. 51 uses the following bounds proved above:
Ú Ú16-11(Deni-1) L) 11 < 1i [D - L =0 (
- (Di(i-1) - - -11 11, s (Da--) - En(a-))
(De(--1) L(((-1)) 0
0
n2
No
[1--11001 = =(n)).
24

---

Under review as a conference paper at ICLR 2020
Since this also hold for previous (t, pairs, have
ukn
A
(1'-1)
Σ
Rev
S
1
52
n1
uknn
n
0
A
n1
0
S
R
(52)
X
which is the reason why we need to take R = +
This means that Wt: E B(R) holds. And by throughout the optimization process.
induction, we have proved that W remains in ((R)
The last thing to do is to bound f; y. By the same logic from above, have
l[f yll2
II A
[f
yll2
AB
(t-1)
$ uvnn
1 1 52
n
which proves theorem.
D ADDITIONAL EXPERIMENTAL RESULTS
In this section, give test performance curve of AFAD-LITE dataset in Figure under the same setting with Section
In addition, provide more baseline results, e.g. Adam (Kingma & Ba 2014) and K-FAC (Martens) & Grosse 2015) on RSNA Bone Age dataset. Since we find that, as another alternation of BN, Group Normalization (GN) (Wu & He 2018) can largely improve the performance of Adam, also implement the GN layer for our GGN method. We use grid search to obtain approximately best
for every experiment. All experiments are performed with batch size 128, input size 64*64 and weight decay 10-1 We set the number of groups to for ResNerGN. Other hyper-parameters are listed below.
SGDD+ResNetBN: learning rate 0.01, momentum 0.9.
Adam+ResNetBNN learning rate 0.001.
SGD+ResNoGGN: rate 0.002, momentum 0.9.
Adam+ResNerGNN: leannigg rate 0.0005.
K-FAC+ResNet: learning rate 0.02, momentum 0.9, 0.1, update frequency 100.
GGN+ResNetGN: 1, a 0.075.
The convergence results summarized in Figure Note make comparison clearer, logarithmic scale for training curves.
25

---

Under review as a conference paper at ICLR 2020
0.50
0.50
GON- ReNN-ReNNe
0.45
SGD PResNetFiop
0.45
SGDI RemNed kuup
SODHRRNNNRR
0.40
0,40
Text
0.35
0.35
0.30
0.30
0.25
20
40
60
80
100
Epoclss
Time
Figure 3: Test performance on AFAD-LITE dataset.
10:1
101
Training Loss
1014
10-4
1077
20ml
5GD RosNedBNN
50D RosNNitNN
10
Adamn Reshuctn
Adamm Reathecth
500 Res-Res
Adim- Re divenE
K04C Raaihie
RPAC-Rishe Rishes
6Gh- RisNNION
CGN NustACON
Time
Epochs
(a) Training loss-time curve on RSNA
(b) Training loss-epoch curve on RSNA
0.30
0.30
S5D+ReeNettX
0.25
SOD-RRSNNNGN
0.25
5CD+ ReskietSN
0.20
Test Loss
0.20
Test
0.15
0.15
0.10
0.30
8.30
8.30
2.5
10.0
12.5
25.0
17.5
20.0
Time
Epochs
(c) Testing loss- time curve on RSNA
(d) Testing loss-epoch curve on RSNA
Figure 4: Training and testing curves of GGN and other baselines on RSNA Bone Age dataset.
26


Reviews:
[{"review_id": "rygJfehGqB", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper presents a second order optimization algorithm, along with convergence proof of the algorithm in both batch and minibatch setting. The effectiveness of the method is demonstrated on two regression tasks. My overall assessment is that the method is still quite limited and the method itself is not novel, but I am willing to change my score to accept if my concerns have been addressed.\n\n(1) The method is not novel. The same algorithm was proposed and applied to the RL setting [1]. \n(2) The method is still quite limited to 1-output function scenario, where the NTK matrix is easy to compute. This limitation though is not mentioned in the paper. I hope the author should have a discussion on this and admit this limitation.\n(3) Also due to (2), the experiments shown in the paper are on toy data and hence lack of strong empirical support.\n(4) The method doesn't scale up to large batch size.\n(5) In the theoretical section, the paper states \n\"However, to our knowledge, no convergence result considering large learning rate (e.g. has the same scale with the\nupdate of GGN) has been proposed.\"\nThis is not true. Here are some papers: [2,3,4]\n(6) Lack of some second order optimization baselines, e.g., KFAC.\n\nMisc:\n(1) For section 3.3, first of all, (B) costs at least half of (A) as it requires a backward pass. \n(2) For section 3.3, the authors write:\n\" What is different is that GGN also, for every input data, keeps track of the output’s derivative for the parameters; while in\nSGD the derivatives for the parameters are averaged over a batch of data.\"\nIs there a simple way of implementing/computing the gradient for *every* input data on GPU? How is that compared to computing the average? I wish to see more evidence of showing they're the same as authors claimed.\n\n[1] Towards Characterizing Divergence in Deep Q-Learning.\n[2] The Power of Interpolation: Understanding the Effectiveness of SGD in Modern Over-parametrized Learning.\n[3] Fast and Faster Convergence of SGD for Over-Parameterized Models (and an Accelerated Perceptron).\n[4] Fast Convergence of Stochastic Gradient Descent under a Strong Growth Condition", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "rJxaGCqN9S", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Authors propose minimizing neural network using kernel ridge regression. (Formula 9 and Algorithm 1). Main difference of this method is compared to Gauss-Newton, is that it uses JJ' as curvature, which has dimensions b-by-by (batch size b), instead of J'J as curvature, which has dimensions m-by-m (number of parameters m).\n\nWhen b is much smaller than m, this matrix is tractable to represent exactly. Related approach is taken by KKT (see Figure 1 of https://arxiv.org/pdf/1806.02958.pdf) which also replaces J'J with more tractable JJ'.\n\nThere is a long history of authors trying to extend second order methods to deep learning and and finding that curvature estimated on a small batch is extremely noisy, requiring large batches (see papers by Nocedal's group). Authors propose a method that estimates curvature from small batches. Given the history of failures in small-batch curvature estimation, the bar is high to show that small-batch curvature estimation works.\n\nBulk of the paper is dedicated to theoretical convergence and connections between concepts. Since the focus of the paper is on a new optimization method for deep learnning, I feel like convergence proofs can be moved to Appendix, and more of the paper should focus on practical aspects of the method. Also the connections to other concepts (ie, tangent kernel) are not essential to the paper and could be better left over for a tutorial paper.\n\nI'm not convinced that their method works well enough to have practical impact.\n\n- Their method seems to be limited to neural network with one output (ie, univariate regression task). This is a serious limitation and paper should highlight this more on this, given that vast majority of applications and benchmarks involve more than output variable.\n\n- Practical implementation details are skimmed over. Section 3.3 brings up that to compute Jacobian, one needs to keep track of the output derivative on per-example basis. How is this accomplished? Modern frameworks like PyTorch and TensorFlow don't give an easy way to compute per-example derivatives efficiently.\n\n- Experiments are performed on two tasks that are not well known in the literature. The choice is somewhat understandable given that their method performs for univariate regression, but also this makes it hard to evaluate whether the method works. SGD vs Gram-Gauss evaluation use parameter settings which are not comparable, so it's impossible to tell whether the improvement are due to better choice of hyper-parameters.\n\n\nThe changes needed to make this paper acceptable are extensive, and I would recommed a reject.\n\nI would recommend authors attempt the following changes for future submission:\n\n1. Make it work for multivariate regression. There's a conversion technique to represent multivariate regression in the same form as univariate regression (see Section 2.4 of \"Rao, Toutenberg\" Linear Models). Essentially it comes down concatenating o output Jacobians (o output classes) along the batch dimension.\n\n2. Use this to evaluate the method on standard benchmarks like MNIST and CIFAR and show that it doesn't cause a significant worsening in quality. Given that similar approach (KKT paper) found bigger improvement on RNN task, an RNN task may be useful.\n\n3. Give more details on implementation. How was Jacobian calculation implemented? Which framework? How was the per-example computation made tractable? Making small-scale experiments reproducible through anonymous github submission would also help ", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly."}, "extra_text": {}}, {"review_id": "SkgoyRG6KH", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Post-rebuttal: I've read author's response and other reviews. As pointed out by other reviewers, the proposed algorithm is restricted to single-output regression and the claim \"accelerate convergence without much computational overhead\" might not be true in general multi-output regression tasks. I believe the lack of multi-output regression experiments makes the paper a bit weak, therefore I changed my score to 3 and vote for rejection.\n\nThat being said, I do find the algorithm interesting and the theoretical results impressive. I encourage the authors to include experiments on multi-output regression tasks (or tone down the claim about computational overhead) and resubmit the paper.\n\n------------------------------------------------------------------------------------------------------------------------------------------------------------------\nBased on recent progress on the connection between neural network training and kernel regression of neural tangent kernel, this paper proposes a Gram-Gauss-Newton (GGN) algorithm to train deep neural networks for regression problems with square loss. For overparameterized shallow networks, the authors proved global convergence of the proposed algorithm in both full-batch and mini-batch setting. To my knowledge, the proof of global convergence in the mini-batch setting is novel and might be of independent interest for other work.\n\nOverall, this paper is well-written and easy to follow. It's interesting to see that the proposed algorithm can achieve quadratic convergence while most previous papers only get linear convergence. \nGiven that, I'd like to give a score of 6 and I'm willing to increase my score if the authors can resolve my concerns below.\n\nConcerns:\n- For the algorithm, if I understand correctly, it's actually same as natural gradient descent with generalized inverse. I think the authors should make the connection clear. I would like to see more discussions with natural gradient descent or Newton methods in the next revision.\n- The authors claim that the proposed GGN algorithm only has minor computational overhead compared to first-order methods. I doubt if it's true in general. In section 3.3, the authors argue that computing individual Jacobian matrices for every example in the minibatch has roughly the same computation as the backpropagation in SGD. As far as I know, it's not true in practice. In addition, the inverse of the Gram matrix can also be expensive when the output dimension (the dimension of y) is large.\n\nMinor Comments:\n- In the paper, the theoretical results are based on the assumption of smooth activation function. I wonder if it is possible to include the case of ReLU activation as it's the most popular activation function in deep learning.\n- I don't have a good understanding about why mini-batch version would converge after reading the paper. To me, second-order methods with mini-batch estimation of the preconditioner would lead to biased gradient estimation. Could you comment on that?", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "BJxhNaz2qH", "reviewer": "AnonReviewer5", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: The authors propose the Gram-Gauss-Newton method for training neural networks. Their method draws inspiration from the connection between the neural network optimization and kernel regression of neural tangent kernel. \n\nTheir method is described in Algorithm 1, but to summarize it, they use the Gauss-Newton method to train neural networks, and prove quadratic convergence for the full-batch training. They also have a mini-batch version of GGN, the practical version, and this is proven to have linear convergence.\n\nThe authors also provide experiments that includes the usual loss v. epoch, but also loss v. wallclock time (which is nice when proposing second-order-like methods where extra computations are necessary), and a test error v. epoch (which is again nice for second-order methods as explained below).\n\nStrengths: The paper has nice proofs of the theorems, and they show a method with quadratic convergence (but full-batch training) without having to invert the full Jacobian matrix whose size depends on the number of parameters, but rather inverting the Gram matrix, whose size depends on the number of training data.\n\nDue to the seeming extra computational cost of the method, (the method requires computing the full Jacobian matrix which depends on the number of neural network weights) I am grateful that they provided comparisons with wallclock time to SGD.\n\nAnd there is this notion that second-order methods have been shown to not generalize as well as first-order methods, and thus it was nice to see that they had an experiment where they tested generalization.\n\nThe background information was also nice to read.\n\nWeaknesses: They do not compare it with other methods optimization methods, such as Adam (a first-order method) or natural gradient (a second-order method), and I would have thus liked to have seen comparisons to these.\n\nI would have also liked to see a test loss v. time/epoch for the AFAD-LITE task as well (they only have it for the RSNA Bone Age task), at least provided in the appendix if there was not enough space.\n\nIn the references, there are numerous citations of the arXiv versions of papers, but I suggest the authors replace them with the conference/journal versions if those papers were accepted in conferences/journals (and I spot some that were).\n\nOther comments: (i) In the first sentence of 3.3., I think one should replace “GGN has quadratic convergence rate” with “full-batch GGN has quadratic convergence rate,” as in the subsequent sections you are discussing mini-batch GGN.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "BJeshEY55r", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors propose a scalable second order method for optimization using a quadratic loss. The method is inspired by the Neural Tangent kernel approach, which also allows them to provide global convergence rates for GD and batch SGD. The algorithm has a computational complexity that is linear in the number of parameters and requires to solve a system of the size of the minibatch.  They also show experimentally the advantage of using their proposed methods over SGD. \n\n\t- The paper is generally easy to read except section 3.1 which could be clearer when establishing the connexion between the proposed algorithm and NTK.\n\n\t- The proposed algorithm seems to be literally a regularized Gauss-Newton with Woodbury matrix inversion lemma applied to equation (7). Additional simplifications occur due to the pre-multiplication by the jacobian and give (9). However, this is not clear in the paper, instead section 3.1, is a bit vague about the derivation of (9).\n\t- In terms of theory, the proofs of thm 1 and 2 seem sound. They rely essentially on the convergence results established for NTK in [Jacot2018, Chizat2018]. The main novelty is that the authors provide faster rates for the Gauss-Newton pre-conditioner which leads to second-order convergence. The second theoretical contribution is to extend the proof to batched gradient descent. Both are somehow expected, although the second one is more technical.\n\t- However, the convergence rates provided for batched gradient descent (thm 2) rely on a rather unrealistic assumption: the size of the network should grow as n^18 where n is the sample size. This makes the result less appealing as in practice this is highly unlikely to be the case.\n\t-  The convergence analysis for the NTK dynamics, which is essential in the proof, relies on a particular scaling 1/sqrt(M) of the function with the number of parameters. In [Chizat2018], it is discussed that although it leads to convergence in the training loss, generalization can be bad. Is there any reason to think in this case, things would be different?\n\n\t- Experiments: Experiments were done on two datasets to solve a regression task. They show that training loss decreases indeed faster than SGD and finds better solutions. A more fair comparison would be against other second-order optimizers like KFAC.\n\t- How was the learning rate chosen for the other methods? Was the same lr used?\n\t- The authors say that the algorithm has the same cost of one backward pass, could they be more specific about the implementation?\n\t- What are the test results for the second dataset? Could they be reported somewhere (in the appendix?)\n        - Both tasks are univariate regression, can the method be applied successfully in a multivariate setting? \nI don't see how the proposed method is different from exactly doing regularized gauss newton, so to me the algorithm is not novel in itself. Besides the method seems to require a quadratic loss function which limits its application.\n\n\n----------------------------------------------------------------------------------\nRevision:\n\n\n I've read the author's response and other reviews. I think the paper will be stronger if extended to more general cases (multivariate output + more general losses), thus I encourage the authors to resubmit the paper with stronger experiments.\n\n\n\n\n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The article considers Gauss-Newton as a scalable second order alternative to train neural networks, and gives theoretical convergence rates and some experiments. The second order convergence results rely on the NTK and very wide networks. The reviewers pointed out that the method is of course not new, and suggested that comparison not only with SGD but also with methods such as Adam, natural gradients, KFAC, would be important, as well as additional experiments with other types of losses for classification problems and multidimensional outputs. The revision added preliminary experiments comparing with Adam and KFAC. Overall, I think that the article makes an interesting and relevant case that Gauss-Newton can be a competitive alternative for parameter optimization in neural networks. However, the experimental section could still be improved significantly. Therefore, I am recommending that the paper is not accepted at this time but revised to include more extensive experiments. 

Decision: Reject