Year: 2019

Abstract:
Deep neural networks exhibit complex learning dynamics due to highly nonconvex loss landscape. Second order approaches, such as natural gradient descent, mitigate such problems by neutralizing the effect of potentially ill-conditioned curvature, yet it is largely unknown how the current theory of deep learning generalizes beyond gradient descent to these higher order learning rules. To answer these questions, we derive exact solutions to learning dynamics of deep linear networks under a spectrum of curvature-corrected learning rules. Our analysis reveals that curvature corrected learning preserves a core feature of gradient descent, a conservation law, such that the learning trajectory follows precisely the same path in the underlying manifold as gradient descent, only accelerating the temporal dynamics along the path. We also show that layer-restricted approximations of natural gradient, which are widely used in most second order methods (e.g. K-FAC), can significantly distort the learning trajectory into highly diverging dynamics that significantly differs from true natural gradient, which may lead to undesirable network properties. We also introduce fractional natural gradient that applies partial curvature correction, and show that it provides most of the benefit of full curvature correction in terms of convergence speed, with additional benefit of superior numerical stability and neutralizing vanishing/exploding gradient problems, which holds true also in layer-restricted approximations.

heading: 1 INTRODUCTION
text: Difficulty in training deep neural networks arises from the fact that the network’s input-output map fθ(·) is nonlinearly related to its parameters θ. This causes non-convex loss landscape with proliferation of saddle-points and poorly-conditioned curvature where gradient-based first order optimization methods perform poorly (Martens, 2010; Dauphin et al., 2014). Second order methods, such as natural gradient descent (Amari, 1998), compensate for the effect of curvature by using the distance metric intrinsic to the space of input-output maps to define the update steps (Pascanu & Bengio, 2013; Martens, 2014; Bernacchia et al., 2018), rather than the parameter space. Recent advancements led to approximate implementations of these methods that prove efficient for practical scale applications (Ba et al., 2016; Grosse & Martens, 2016; Martens et al., 2018; Osawa et al., 2019).
Despite their practical effectiveness, however, the exact nature of such curvature-corrected learning process remains largely unknown. Do curvature-corrected learning methods simply accelerate convergences towards the same minimum solutions as gradient descent, or do they impose implicit bias toward qualitatively different solutions?
As a first step toward establishing theoretical understanding of these questions, we analyze the exact learning dynamics of deep linear networks under a spectrum of curvature-corrected update rules. Deep linear networks provide an excellent mathematical framework for developing insightful theoretical understanding of the complex inner workings of deep nonlinear networks (Goodfellow et al., 2016). Despite their simplicity, deep linear networks capture the essential nonlinear relationship between network’s input-output maps and their parameters, and exhibit comparable learning behavior to their nonlinear counterparts that can be exactly solved for rigorous analysis. Indeed, many recent works analyzed the learning trajectories of deep linear networks under gradient descent to compute the convergence rate under various initial conditions (Arora et al., 2018a;b; Bartlett et al., 2019; Du & Hu, 2019), revealed decoupled modes of convergence dynamics to explain the origin of multiple
stage-like loss profiles (Saxe et al., 2013), and showed the implicit bias for regularization (Du et al., 2018; Arora et al., 2019) and resistance to overfitting (Advani & Saxe, 2017; Lampinen & Ganguli, 2018; Poggio et al., 2018). Yet, it is uncertain whether these convergence properties generally apply for update rules beyond gradient descent.
Our contribution The main results are summarized as follows.
1. We derive a generalized conservation law that describes the optimization paths of network parameters under gradient descent as well as curvature-corrected update rules. Consequently, curvature correction only affects the speed of convergence without affecting other qualitative properties of parameter update process.
2. There is a trade-off between map dynamics and parameter dynamics. The full curvature correction effect of natural gradient descent (NGD) completely linearizes the map learning dynamics of deep networks, equivalent to that of shallow networks. Such complete linearization, however, sacrifices stability of parameter update dynamics to explode when gradient vanishes and vice versa.
3. We introduce a regularized version of NGD that partially corrects for the effect of curvature, called √ NGD, which facilitates the parameter update dynamics by eliminating the
vanishing/exploding update problems. This makes the map dynamics slightly nonlinear, but no more so than that of single hidden layer networks under gradient descent.
4. NGD makes the learning process prone to overfitting by simultaneously learning both the signal and the noise dimensions of data, whereas √ NGD partially retains gradient descent’s
resistance to overfitting by separating the time-scales between the signal and the noise dimensions.
5. The widely-used block-diagonal approximation of NGD breaches the aforementioned conservation law, resulting in highly divergent parameter update dynamics, which breaks the weight balance across layers. In contrast, block-diagonalization of √ NGD preserves sta-
bility of parameter update dynamics, yielding efficient and stable learning algorithms.

heading: 2 SETUP AND NOTATIONS
text: Consider a depth d network that consists of an input layer, d− 1 hidden layers, an output layer, and weight matrices w ≡ {wi}di=1 that connect the adjacent layers. The network’s input-output map is w̄ ≡ ∏d i=1 wi = wd · · ·w1, such that fw(x) = w̄x. The network learns the input-output statistics of a dataset D = {xµ, yµ}Pµ=1 by minimizing the squared-error loss:
L(w) = 1
2 ED[‖w̄x− y‖2] = Tr
[ 1
2 (w̄ − w̄∗)Σx(w̄ − w̄∗)ᵀ
] + const,
where ED is the expectation over the dataset D, Σx ≡ ED[xxᵀ] is the input correlations, and w̄∗ ≡ ED[yxᵀ] Σ−1x . Neglecting the constant term, the loss function is expressed as
L(w) = Tr [ 1
2 ∆Σx∆
ᵀ ] . (∆ ≡ w̄ − w̄∗) (1)
where ∆ denotes the displacement between w̄ and w̄∗.
Shallow networks (d = 1, w̄ = w1) exhibit linear learning dynamics under gradient descent, whose convergence rates scale with eigenvalues of Σx. In this case, curvature correction has the well-understood effect of normalizing the convergence rates, which is also achievable by simple pre-whitening of input correlation. Instead, we are interested in the less-understood effect of how curvature correction facilitates the complex nonlinear dynamics of deep networks (d ≥ 2). Therefore, we consider pre-whitened input distribution Σx = I to isolate the nonlinear effect of curvature correction, but this condition is not critical for the analysis.
Gradient and Hessian+ We use bold symbols to collectively represent network parameters and derivatives in array form. For example, ẇ ≡ [ ẇ1 ẇ2 ] and g ≡ [ ∂L ∂w1 ∂L ∂w2 ] = [ wᵀ2 ∆ ∆wᵀ1 ] represent the continuous-time weight update and the gradient of a depth d = 2 network. Hessian is fully characterized by its operation on weight update, which, by definition, produces gradient update:
Hẇ = ġ =
[ wᵀ2 ∆̇ + ẇ ᵀ 2 ∆
∆̇wᵀ1 + ∆ẇ ᵀ 1
] . (∆̇ = ˙̄w = w2ẇ1 + ẇ2w1) (2)
However, true Hessian-based methods (e.g. Newton-Raphson method) can converge to any extrema types. To guarantee convergence to (local) minimum solutions, natural gradient methods use positive semi-definite (PSD) approximations of Hessian (e.g. Fisher matrix (Amari, 1998; Heskes, 2000; Martens & Grosse, 2015; Bernacchia et al., 2018), Generalized-Gauss-Newton matrix (Martens, 2014; Botev et al., 2017; Roux et al., 2008)1), which correspond to
H+ẇ =
[ wᵀ2 ∆̇
∆̇wᵀ1
] . (3)
This operation is indeed PSD, since ẇ ·H+ẇ = Tr[ẇᵀ1w ᵀ 2 ∆̇ + ẇ ᵀ 2 ∆̇w ᵀ 1 ] = Tr[∆̇∆̇ ᵀ] ≥ 0, where the dot-product denotes a · b ≡ ∑d i=1 Tr[aib ᵀ i ]. We refer to this operation as Hessian+.
Null-space and Conservation laws Deep linear networks exhibit inherent symmetries that their input-output map w̄ is invariant under transformations that multiply arbitrary matrix m to one layer and its inverse to the next layer, i.e. [ w1 w2 ] → [ mw1 w2m −1 ] , ∀m. ẇnull ≡ [ mw1 −w2m ] are the equivalent
continuous-time transformations that yield the invariance ∆̇ = ˙̄w = w2mw1 − w2mw1 = 0, ∀m.
These transformations form the null-space of H+, since ẇnull ·H+ẇnull = Tr[∆̇∆̇ᵀ] = 0, which is orthogonal to gradient, since ẇnull · g = Tr[∆∆̇ᵀ] = 0. Also orthogonal to the null-space is natural gradient, since ẇnull ·H†+g = g ·H†+ẇnull = 0, where H†+ denotes Moore-Penrose pseudo-inverse. These continuous symmetries imply the following, self-explanatory theorem (Noether’s theorem):
Theorem 1 All update rules ẇ that are orthogonal to the null-space, i.e.
ẇ · ẇnull = d∑ i=1 Tr[(wiẇ ᵀ i − ẇ ᵀ i+1wi+1)mi] = 0, ∀mi
exhibit the following conservation law
d/dt (wiw ᵀ i − w ᵀ i+1wi+1) = 0, ∀i (4)
This result was previously only known for gradient descent dynamics (Arora et al., 2018b; Du et al., 2018), which is generalized here.

heading: 3 LEARNING DYNAMICS
text: In this section, we analyze the learning dynamics of the network parameters w (Section 3.1) and the update dynamics of the input-output map w̄ (Section 3.2) under a spectrum of curvature-corrected update rules. We then analyze how block-diagonal approximation modifies the curvature-corrected dynamics (Section 3.3).

heading: 3.1 PARAMETER DYNAMICS
text: We follow the singular value decomposition (SVD)-based analysis of Saxe et al. (2013); Advani & Saxe (2017); Lampinen & Ganguli (2018), by considering network weights that are initialized to
1Fisher matrix and Generalized-Gauss-Newton matrix are equivalent in many cases, including the least squares problem considered here (Pascanu & Bengio, 2013; Martens, 2014).
have their map’s singular vectors aligned with those of w̄∗.2Under such initialization, the update dynamics of weight matrices simplifies to their singular value dynamics, with their singular vectors remain unchanging. This simplified case admits exact analytic solutions, which provide good approximation to general learning dynamics. Moreover, this aligned singular vector condition is automatically satisfied for networks initialized with small random weights Saxe et al. (2013); Advani & Saxe (2017).
Steepest gradient descent (SGD) Under SGD update, deep networks’ weight parameters exhibit coupled nonlinear dynamics: (d = 2 example, η: learning rate)
ẇ + ηg =
[ ẇ1 + η w ᵀ 2 ∆
ẇ2 + η∆w ᵀ 1
] = 0. (5)
The SVD analysis decomposes eq (5) to individual singular mode dynamics. The dynamics of one singular mode is described by3(See S.I.)
σ̇i + η σ∆ ji = 0 (σ∆ = σ̄ − σ̄∗, σ̄ = d∏ i=1 σi) (6)
where σi, σ̄∗, σ̄, σ∆ are the singular values of wi, w̄∗, w̄, ∆, and ji ≡ ∂σ̄/∂σi = σ̄/σi denotes the coupling between the input-output map and parameters, i.e. Jacobian. Note that this singular mode dynamics follows the hyperbolic paths
σ2i − σ2k = constant, ∀i, k (7) which is the direct consequence of the conservation law (4). The update speed ‖σ̇‖ is proportional to the displacement |σ∆| and the coupling strength ‖j‖
‖σ̇‖ ∝ |σ∆|‖j‖, (‖σ̇‖2 ≡ ∑d i=1 σ̇ 2 i , ‖j‖2 ≡ ∑d i=1 j 2 i ) (8)
which vanishes for networks with small coupling strength and explodes for large coupling strength.
Natural gradient descent (NGD) NGD finds the minimum-norm update solution (min ẇ · ẇ) subject to the constraint (i.e. Moore-Penrose pseudo-inverse solution)
H+ẇ + ηg =
[ wᵀ2 (∆̇ + η∆)
(∆̇ + η∆)wᵀ1
] = 0, (9)
which can be solved using Lagrange multipliers to yield (See S.I.)[ ẇ1 + η w ᵀ 2 Λ
ẇ2 + ηΛw ᵀ 1
] = 0, (10)
where Λ satisfies wᵀ2S(Λ) = S(Λ)w ᵀ 1 = 0. (S(Λ) ≡ (w2w ᵀ 2 )Λ + Λ(w ᵀ 1w1)−∆) (11)
Remarkably, the only change from SGD update (5) is replacing ∆ with Λ as the main drive of dynamics eq (10), which preserves orthogonality to null-space and hence the conservation law (4) 4. The singular mode dynamics of NGD update eq (10) is5
σ̇i + η σ∆ ji ‖j‖2 = 0, (12)
where σ∆ of SGD eq (6) is replaced by σΛ = σ∆/‖j‖2, the singular values of Λ (See S.I.). NGD dynamics eq (12) follows the same hyperbolic paths of SGD eq (7), but with modified update speed
‖σ̇‖ ∝ |σ∆| ‖j‖
(13)
which inversely scales with ‖j‖. Therefore, NGD’s update speed explodes for small coupling strength, reciprocal to SGD’s vanishing speed phenomenon.
2Given SVD of weight matrices wi = LiDiRᵀi and w̄∗ = L∗D∗R ᵀ ∗, where D are the diagonal singular value matrices and L/R are the left/right singular vector matrices, the aligned singular vector condition assumes R1 = R∗, Ld = L∗ and Ri+1 = Li for all layers 1 ≤ i ≤ d− 1.
3 The dynamics eq (6),(12) apply to all active singular modes. Inactive modes that have σ̄ = 0 stay frozen. The number of active modes is determined by the bottleneck size, i.e. the narrowest width of network.
4The Moore-Penrose pseudo-inverse solution is guaranteed to be orthogonal to the null-space, since nonzero null-space components only increase the solution’s norm without affecting the constraint eq (9).
Fractional Natural Gradient Descent ( q √
NGD) Above results can be generalized to a spectrum of update rules that apply partial curvature corrections, described by q √ H+ẇ + ηg = 0, where q √ H+ is a fractional power of Hessian+ (q ≥ 1). The singular mode dynamics of q √ NGD is
σ̇i + η σ∆ ji ‖j‖2/q = 0, (14)
which interpolates between NGD (q = 1) and SGD (q →∞). Eq (14) follows the same hyperbolic paths of SGD eq (7), but with modified update speed
‖σ̇‖ ∝ |σ∆|‖j‖1−2/q. (15)
Note that for q = 2, termed √ NGD, the update speed becomes independent of the coupling strength
‖σ̇‖ = η |σ∆|, (16) thus eliminating the vanishing/exploding update speed problems of SGD/NGD (See Fig 1C).
Relation to Regularized NGD Alternative interpolation solves (H+ẇ + ηg) + I(ẇ + ηg) = 0 ( ≥ 0), which yields the regularized (or damped) inverse
ẇ = −η( + 1)( I + H+)−1g, 6 (17) similar to Levenberg-Marquardt damping (less the ( + 1) term), whose singular mode dynamics is
σ̇i + η σ∆ ji ‖j‖ ( a‖j‖+ 1 a+ ‖j‖ ) = 0, (a ≡ /‖j‖) (18)
6 This expression reduces to SGD in the limit → ∞, which differs from the usual regularized inverse ẇ = −η( I + H+)−1g, which reduces to 0.
where the ratio a ≡ /‖j‖ describes the effective degree of interpolation between NGD (a → 0) and SGD (a → ∞). Note that a should be large enough to provide sufficient damping, but not too large to nullify the effect of curvature correction, which is difficult to simultaneously satisfy across all singular modes with fkxed . √ NGD can be considered as providing ideally and adaptively tuned regularization (a = 1) for all singular modes, where the regularization is most effective.

heading: 3.2 MAP DYNAMICS
text: The parameter update ultimately drives the learning dynamics of the input-output map, via Jacobian
˙̄σ = d∑ i=1 σ̇i ji, (19)
which yields the following map learning dynamics under q √ NGD update (14)
˙̄σ = −η(σ̄ − σ̄∗)‖j‖2(1−1/q). (20) In general, eq (20) does not admit closed-form solutions due to the coupling strength term, with the exception of NGD (q = 1). As shown by the vector field in Figure 1, however, the coupling strength changes in a streotypical manner along the learning trajectories. Therefore, the general characteristics of map dynamics can be appreciated from the representative case of balanced weights: σi = σ̄ 1/d ∀i, or in terms of the conserved quantities, wiwᵀi − w ᵀ i+1wi+1 = 0. Note that this balanced weight condition is automatically approximately satisfied if the networks are initialized with small random weights.
Under the balanced weight condition, eq (20) simplifies to
˙̄σ = −η̄ (σ̄ − σ̄∗) σ̄p ( p ≡ 2(d− 1)(q − 1)
d q ) (21)
where η̄ ≡ η d1−1/q is the depth-calibrated learning rate, and p represents the combined effect of depth and curvature correction that determines the stiffness, or degree of nonlinearity, of map
dynamics. Figure 2 shows the following notable closed-form solutions, as well as p = 2 case: σ̄(t) = σ̄∗(1− e−η̄t) ( p = 0) σ̄(t) = σ̄∗ tanh 2(η̄ √ σ̄∗t/2) ( p = 0.5)
σ̄(t) = σ̄∗
1 + (σ̄∗/σ̄(0) − 1)e−η̄σ̄∗t ( p = 1 )
where zero initial condition σ̄(0) = 0 is assumed for p < 1 cases.
NGD update (q = 1, p = 0) Under NGD update, the map dynamics exhibits fully linearized convergence dynamics with a constant time-scale η−1 for all depth d and data mode-strength σ̄∗. Its learning curves exhibit finite growth rate near zero σ̄(t) ≈ η̄ σ̄∗t, which entails exploding parameter update speed as the coupling strength approaches zero. Therefore, the full curvature correction of NGD sacrifices stability of parameter dynamics in order to perfectly cancel out all nonlinearities of map dynamics.
√ NGD update (q = 2, p = 1 − 1/d) For √ NGD update, the stiffness ranges from p = 0.5 for single hidden layer networks to p→ 1 in infinite depth limit. Its learning curves exhibit polynomial growth near zero, σ̄(t) ∝ t1/(1−p), which takes finite time to escape from zero initial condition. even though the initial growth rate vanishes with the coupling strength. The overall time-scale of learning decreases with mode strength as σ̄−p∗ , such that stronger singular modes (large σ̄∗) learn faster than weaker modes.
SGD update (q → ∞, p = 2 − 2/d) Under SGD update, the stiffness ranges from p = 1 for single hidden layer networks to p → 2 in infinite depth limit. Its learning curves exhibit sigmoidal shape that take infinite time to escape from the saddle point at zero initial condition: the escape time diverges as O(− log σ̄(0)) for p = 1 and O(σ̄1−p(0) ) for p > 1. Also, the increased p causes greater separation of time-scales (η̄σ̄p∗)−1 across singular modes, which results in stage-like transitions over the course of training, with each singular mode making sudden transition from slow learning to rapid convergence (Saxe et al., 2013).
Effective Depth Network depth d and curvature correction q interact in a symmetric manner, which can be intuitively understood by representing stiffness in terms of the corresponding network depth under SGD update, called the effective depth:
deff = dq
d+ q − 1 , (22)
which approaches the actual depth deff → d in the SGD limit (q → ∞), and similarly, approaches deff → q in the limit of infinite depth (d → ∞). Therefore, q √ NGD reduces the network’s effective
depth to be strictly less than q. For √ NGD, this upper-limit is 2, i.e. single hidden layer network.
To summarize, curvature correction lowers the nonlinearity/stiffness of map dynamics of deep networks by reducing their effective depth. The full curvature correction effect of NGD perfectly cancels out all nonlinearities of map dynamics to exhibit linear convergence, equivalent to shallow network learning, but it sacrifices stability of parameter dynamics to explode at the saddle point. In contrast, partial curvature correction of √ NGD directly facilitates the parameter update dynamics, which eliminates the vanishing/exploding update problem, and it makes the map dynamics only slightly nonlinear, but no more so than that of single hidden layer networks under gradient descent.

heading: 3.3 EFFECT OF LAYER-RESTRICTED APPROXIMATION
text: Block-diagonal NGD (NGD-d) In most practical deep learning applications, numerically estimating and inverting Hessian+ becomes prohibitively expensive. Instead, most second-order methods approximate NGD by applying layer-restricted curvature corrections, ignoring the off-blockdiagonal Hessian+ terms across different layers (Martens & Grosse, 2015; Ba et al., 2016; Grosse & Martens, 2016; Martens et al., 2018; Bernacchia et al., 2018): (d = 2 example)[
H1ẇ1 + η1g1 H2ẇ2 + η2g2
] = [ wᵀ2w2ẇ1 + η1w ᵀ 2 ∆
ẇ2w1w ᵀ 1 + η2∆w ᵀ 1
] = 0, (23)
which nevertheless satisfies the NGD constraint (9) if ∑d i=1 ηi = η. Hi denotes the block-diagonal Hessian+ term of layer i. Singular mode dynamics of eq (23) is (with ηi = η/d)
σ̇i + η σ∆ d ji = 0, (24)
where the layer-restricted factor j2i substitutes the full curvature correction factor ‖j‖2 of NGD (12). This block-diagonalization significantly modifies the parameter update dynamics by adding nonzero null space components. Instead of the hyperbolic paths (7), eq (24) follows radially diverging paths that conserve σi/σk as constants of motion. Consequently, NGD-d update exhibits larger parameter update speed than NGD7, and converges to less efficient, large norm solutions that are highly sensitive to initial conditions and perturbations (Fig 1D, red line). Despite the vastly different parameter dynamics, however, NGD-d exhibits identical map learning dynamics as NGD ˙̄σ = −η (σ̄ − σ̄∗) (Fig 1BD, red dots), because the input-output map is invariant under null-space transformations.
Block-diagonal √ NGD ( √
NGD-d) More generally, block-diagonalized fractional NGD H
1/q i ẇi + η gi/d 1/q = 0 yields
σ̇i + η σ∆ d1/q j 1−2/q i = 0, (25)
which conserves σ2(1−1/q)i − σ 2(1−1/q) k as constants of motion. For q = 2, called
√ NGD-d, the
singular mode dynamics
σ̇i + η σ∆√ d sign(ji) = 0, (26)
follows non-diverging, parallel paths that conserve |σi|−|σk|, with identical parameter update speed as √ NGD’s ‖σ̇‖ = η |σ∆| (Fig 1E). Therefore, √
NGD-d yields neutrally-stable update dynamics that neutralizes the vanishing/exploding update speed problems.

heading: 4 IMPLICIT BIAS FOR REGULARIZATION
text: Recent works have shown that learning dynamics of deep networks under SGD update exhibits implicit bias towards keeping the network well regularized. Here, we consider two such properties, and analyze how they generalize under curvature-corrected update rules.
Weight balance Deep neural networks often exhibit redundant parameterizations, such that configurations of parameters implement the same input-output map. One such redundancy, or symmetry, that concerns both deep linear networks and deep ReLU networks is homogeneity, such that multiplying a layer by a positive scalar c and divide another layer by c, does not change the input-output map. The problem is that c can be arbitrarily large or small, yielding potentially unbounded, yet valid solutions. Such unboundedness poses major theoretical difficulty for convergence analysis of gradient-based local optimization methods (Lee et al., 2016; Shamir, 2018).
Fortunately, SGD update exhibits implicit bias toward automatically balancing the norm of different layer’s weight matrices. The proof directly follows from the following conserved quantity of scalar multiplication symmetry ‖wi‖2Frob − ‖wi+1‖2Frob, which is a relaxed version of the aforementioned conservation law eq (4). Thus, if the weights are initially small, this differences between squared norm will remain small through out learning process, thus establishing balancedness across layers Du et al. (2018). As shown in section 2, curvature-corrected updates (e.g. NGD and √
NGD) retain orthogonality to the null-space of symmetry, and thus comply with the same conservation laws as SGD. We show numerical confirmation of this prediction in S.I. The conservation of squared difference of norms for homogeneous ReLU networks still requires similar numerically confirmation.
In contrast, block-diagonalized methods do not follow the same conservation law. NGD-d conserves the ratio between singular values across layers σi/σk, which does not guarantee balancedness even
7 ‖σ̇‖2NGD-d ≥ ‖σ̇‖2NGD can be shown using Jensen’s inequality: 1d ∑d i=1 1 j2i ≥ d∑d i=1 j 2 i ,
with small initialization. √
NGD-d, however, conserves the absolute difference of singular values across layers |σi| − |σk|, which guarantees balancedness, at least under the condition of aligned singular vectors: That is, the ratio between the singular values would approach close to 1 if they grow from small initial values, while maintaining the small absolute difference. Although this does not constitute a formal proof for general case, √ NGD-d confirms to maintain balancedness across layers in numerical simulations (See S.I.).
Low rank approximation / Generalization dynamics The learning dynamics of the input-output map under SGD update separates the time-scales of learning across singular modes (η̄σ̄p∗)−1, such that the singular modes with stronger data correlation preferentially are learned faster (Saxe et al., 2013). This property yields an implicit regularization property for deep networks to efficiently extract low-rank structure of the dataset, such as for finding matrix factorizations with minimum nuclear norm (Gunasekar et al., 2017; Arora et al., 2019). It also allows deep networks to avoid overfitting via early stopping by first learning the signal dimensions of noisy dataset, before the overfitting of the noise dimensions occurs, as long as signal-to-noise ratio is sufficiently large Advani & Saxe (2017); Lampinen & Ganguli (2018) which yields good generalization performance to unseen data. However, this approach requires the network to be trained from small random weight initialization, where SGD suffers from vanishing gradient problem.
In curvature-corrected cases, the learning speed of map dynamics eq (21) scales as σ̄−p∗ . Under NGD, the map dynamics is perfectly linearized (p = 0), which also removes its ability to separate out the time-scales. This makes NGD update prone to large generalization error due learning the noise dimension simultaneously with the signal. In contrast, √ NGD partially retains time-scale separation in the learning dynamics, while also accelerating the parameter update dynamics near zero weights.
We the generalization property of curvature corrected learning rules with student-teacher task from Lampinen & Ganguli (2018), in which the training and test dataset are generated by a teacher network yµ = w̄∗xµ + zµ, where xµ ∈ RN is the input data, yµ ∈ RN is the output, w̄∗xµ is the signal and zµ ∈ RN is the noise. Teacher’s input-output map w̄∗ ∈ RN×N is assumed to have a low-rank structure (rank 3), and the student is a depth d = 4 network of constant width N = 16, whose weight matrices are initialized to have maximum singular value of 0.05. The number of training dataset {xµ, yµ}Pµ=1 is set to be equal to the effective number of parameters P = N , which makes the learning process most susceptible to overfitting. For numerical calculation of NGD and √
NGD, the Hessian+ block between layer i and k are computed as described in Bernacchia et al. (2018) (eq 42), which are then concatenated to the full Hessian+ matrix and numerically inverted (or sqrt-inverted) via eigen-decomposition. LevenbergMarquardt damping of = 10−5 and update clipping are used for numerical stability of NGD.√
NGD does not require such clipping or damping terms.
Figure 3 shows the result of training. SGD exhibits stage-like transitions, which first learns the three signal modes, well separated from the onset of overfitting of the noise modes begins, which allows effective early stopping scheme. However, it suffers from the long plateaus due to vanishing gradient problem.
NGD (and NGD-d) updates learn all singular modes simultaneously including the noise modes (See Fig 3 D), which leads to high generalization error. Note that NGD’s loss profile deviates from exponential decay due to the clipping. In contrast, √ NGD (and √ NGD-d) allows fast learning while separating the signal dimensions from the noise dimensions, achieving comparable test loss as SGD update, but also with fast early-stopping time comparable to NGD update. Note that all three update rules achieve the same test loss after overfitting is complete. This is due to the shared learning path for each singular mode across the methods.

heading: 5 CONCLUSION
text: To summarize our contribution, we derived a generalized conservation law that describes the optimization paths of network parameters under gradient descent as well as curvature-corrected update rules. Consequently, curvature correction only affects the speed of convergence without affecting other qualitative properties of parameter update process.
We revealed a trade-off between map dynamics and parameter dynamics: The full curvature correction effect of natural gradient descent (NGD) completely linearizes the map learning dynamics of deep networks, equivalent to that of shallow networks. Such complete linearization, however, sacrifices stability of parameter update dynamics to explode when gradient vanishes and vice versa. Moreover, we introduced √ NGD that partially corrects for the effect of curvature, which facilitates the parameter update dynamics by eliminating the vanishing/exploding update problems. This makes the map dynamics slightly nonlinear, but no more so than that of single hidden layer networks under gradient descent. Moreover, NGD makes the learning process prone to overfitting by simultaneously learning both the signal and the noise dimensions of data, whereas √ NGD partially retains gradient descent’s resistance to overfitting by separating the time-scales between the signal and the noise dimensions. We also showed that the widely-used block-diagonal approximation of NGD breaches the aforementioned conservation law, resulting in highly divergent parameter update dynamics, which breaks the weight balance across layers. In contrast, block-diagonalization of √
NGD preserves stability of parameter update dynamics, yielding efficient and stable learning algorithms.

heading: 4 NUMERICAL CONFIRMATION OF CONSERVATION LAW
text: See Figure S2. Figure S2 plots the learning trajectory of a 3 layer network, and shows the elements of the weight matrices evolving over time (w1, w2, w3). It also shows the conserved quantities w1w ᵀ 1 − w ᵀ 2w2, w2w ᵀ 2 − w ᵀ 3w3, which indeed remain constant for SGD, NGD and √ NGD, while
it blows up for NGD-d. √
NGD also violates the conservation law, but the weights remain balanced over time.

Review summaries:
Authors analyse curvature corrected optimization methods in the context of deep learning. They build their analysis on Saxe et.al.s work. They show that curvature corrected methods preserve properties of SGD. They also show the disadvantages of layer restricted approximations. They show the importance of time scales in optimization. The paper looks to deep learning from a dynamical systems perspective and hence their experiments are fitting to this framework.

In this manuscript, the authors analyze the dynamics of training deep linear neural networks under a generalized family of natural gradient methods that apply curvature corrections. They first show that the learning trajectory (direction of singular mode dynamics) in natural gradient descent follows the same path as gradient descent, while only accelerating the temporal dynamics along the path. Moreover, the authors show that the learning trajectory in layer-restricted approximations of natural gradient descent can significantly differ from the true natural gradient. Also, the authors proposed a fractional natural gradient that applies partial curvature correction which in addition to faster convergence, neutralizes vanishing/exploding gradient problems.