Title: CURVATURE-CORRECTED LEARNING DYNAMICS IN DEEP LINEAR NEURAL NETWORKS

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.

1. INTRODUCTION:

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 stability of parameter update dynamics, yielding efficient and stable learning algorithms.

2. SETUP AND NOTATIONS:

Consider a depth d network that consists of an input layer, d -1 hidden layers, an output layer, and weight matrices w ≡ {w i } d i=1 that connect the adjacent layers. The network's input-output map is w ≡

d i=1 w i = w d • • • w 1 , such that f w (x) = wx.

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 E D [ wx -y 2 ] = Tr 1 2 ( w -w * )Σ x ( w -w * ) + const,

where E D is the expectation over the dataset D,

Σ x ≡ E D [xx ]

is the input correlations, and

w * ≡ E D [yx ] Σ -1

x . 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 = w 1 ) 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. 

EQUATION

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

EQUATION

)

This operation is indeed PSD, since ẇ

• H + ẇ = Tr[ ẇ 1 w 2 ∆ + ẇ 2 ∆ w 1 ] = Tr[ ∆ ∆ ] ≥ 0, where the dot-product denotes a • b ≡ d i=1 Tr[a i b 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.

w 1 w 2 → mw 1 w 2 m -1 , ∀m. ẇnull ≡ mw 1 -w 2 m

are the equivalent continuous-time transformations that yield the invariance ∆ = ẇ = w 2 mw 1 -w 2 mw 1 = 0, ∀m.

These transformations form the null-space of H + , since ẇnull

• H + ẇnull = Tr[ ∆ ∆ ] = 0, which is orthogonal to gradient, since ẇnull • g = Tr[∆ ∆ ] = 0. Also orthogonal to the null-space is natural gradient, since ẇnull • H † + g = g • H † + ẇ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 ẇ that are orthogonal to the null-space, i.e.

EQUATION

This result was previously only known for gradient descent dynamics (Arora et al., 2018b; Du et al., 2018) , which is generalized here.

3. LEARNING DYNAMICS:

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).

3.1. PARAMETER DYNAMICS:

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 have their map's singular vectors aligned with those of w * .2 Under 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)

EQUATION

The SVD analysis decomposes eq (5) to individual singular mode dynamics. The dynamics of one singular mode is described by3 (See S.I.)

EQUATION

where σ i , σ * , σ, σ ∆ are the singular values of w i , w * , w, ∆, and j i ≡ ∂ σ/∂σ 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

EQUATION

which is the direct consequence of the conservation law (4). The update speed σ is proportional to the displacement |σ ∆ | and the coupling strength j

EQUATION

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 ẇ • ẇ) subject to the constraint (i.e. Moore-Penrose pseudo-inverse solution)

EQUATION

which can be solved using Lagrange multipliers to yield (See S.I.)

EQUATION

where Λ satisfies

EQUATION

) 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

EQUATION

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. 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 + ẇ + ηg = 0, where q √ H + is a fractional power of Hessian + (q ≥ 1). The singular mode dynamics of q √ NGD is

EQUATION

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

EQUATION

)

Note that for q = 2, termed √ NGD, the update speed becomes independent of the coupling strength

EQUATION

thus eliminating the vanishing/exploding update speed problems of SGD/NGD (See Fig 1C ).

Relation to Regularized NGD Alternative interpolation solves (H + ẇ + ηg) + I( ẇ + ηg) = 0 ( ≥ 0), which yields the regularized (or damped) inverse ẇ = -η( + 1)(

EQUATION

similar to Levenberg-Marquardt damping (less the ( + 1) term), whose singular mode dynamics is 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 .

EQUATION

√ NGD can be considered as providing ideally and adaptively tuned regularization (a = 1) for all singular modes, where the regularization is most effective.

3.2. MAP DYNAMICS:

The parameter update ultimately drives the learning dynamics of the input-output map, via Jacobian

EQUATION

which yields the following map learning dynamics under q √ NGD update ( 14)

EQUATION

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, w i w i -w i+1 w i+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

EQUATION

)

where η ≡ η d 1-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) ∝ t 1/(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:

EQUATION

which approaches the actual depth d eff → d in the SGD limit (q → ∞), and similarly, approaches d eff → 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.

3.3. EFFECT OF LAYER-RESTRICTED APPROXIMATION:

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) 

EQUATION

which nevertheless satisfies the NGD constraint (9) if d i=1 η i = η. H i denotes the block-diagonal Hessian + term of layer i. Singular mode dynamics of eq ( 23) is (with

EQUATION

where the layer-restricted factor j 2 i 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 NGD 7 , 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.

EQUATION

)

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 

EQUATION

4. IMPLICIT BIAS FOR REGULARIZATION:

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 w i 2 Frob -w i+1 2 Frob , 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 σ 2 NGD-d ≥ σ 2 NGD can be shown using Jensen's inequality:

1 d d i=1 1 j 2 i ≥ 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 µ ∈ R N is the input data, y µ ∈ R N is the output, w * x µ is the signal and z µ ∈ R N is the noise. Teacher's input-output map w * ∈ R N ×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. Levenberg-Marquardt 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.

5. CONCLUSION:

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. L shown across different depth d and curvature correction factor q. α is the overall scale factor for the weights. Dotted-lines show predictions for the case of balanced weights, eq (S.I.16)(S.I.17). The numerical results show the result from networks of random weights sampled from Gaussian distribution, (Xavier normalization (Glorot & Bengio, 2010) ), They exhibit the same power-law exponent as the dotted-lines (i.e. same slope), but shifted toward right. Blockdiagonal approximation does not affect the power-law exponents, but it affects the coefficient such that NGD-d exhibits much larger parameter update size than NGD, while exhibiting identical zero power-law for loss update. In contrast, sNGD-d exhibits identical match with sNGD for both parameter update and loss update. See S.I. for detailed setup. Error bars show standard deviation from 10 simulations with random weights.

S.I.2 SINGULAR MODE ANALYSIS EQ (18) IN SECTION 3.2:

We follow the SVD-base analysis under the aligned singular vector condition (Saxe et al., 2013) . We introduce σ i , σ, σ ∆ , σ Λ , σ S which represent the singular values of w i , w, ∆, Λ, S(Λ) of one particular singular mode, and j i ≡ ∂ σ ∂σi = k =i σ k In this representation, eq (S.I.4)(S.I.5) reduce to σi = -σ Λ j i (S.I.10) whereas, eq (S.I.8)(S.I.9) reduce to 

S.I.3 SCALING LAWS OF UPDATE DYNAMICS:

The parameter update ( 14) and map learning dynamics (20) of a singular mode exhibit the following scaling relationships with respect to the coupling strength j :

σ /|σ ∆ | ∝ j 1-2/q , σ/σ ∆ ∝ -j 2-2/q , (S.I.15)

:

Fisher matrix and Generalized-Gauss-Newton matrix are equivalent in many cases, including the least squares problem considered here(Pascanu & Bengio,

2013;Martens, 2014).

Given 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 * , L d = 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.4 The 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).

This expression reduces to SGD in the limit → ∞, which differs from the usual regularized inverse ẇ = -η( I + H+) -1 g, which reduces to 0.

Figures and Tables:

Figure1: Learning dynamics of a singular mode of a single hidden layer network (d = 2). The contour lines visualize the manifolds of constant displacement levels σ ∆ ≡ σ 1 σ 2 -σ * . The optimal solution σ ∆ = 0 is shown in black. Tangent space of the manifolds defines the null-space of Hessian + . The vector field visualizes the displacement-normalized update [ σ1 , σ2 ]/|σ ∆ |, whose amplitude is the normalized update speed: σ /|σ ∆ | ∝ j 1-2/q . (A,B,C) SGD, NGD, √ NGD share the same update directions defined by the hyperbolic paths that conserve σ 2 1 -σ 2 2 (red lines), orthogonal to the null-space. But they exhibit different update speed: SGD exhibits vanishing speed problem for small weights, while NGD has the opposite problem. In contrast, √ NGD exhibits constant normalized speed. (D) NGD-d exhibits radially diverging vector field that conserves σ 1 /σ 2 . The learning trajectories of NGD and NGD-d traverse the contour lines with synchronized timing (red dots), (E) √ NGD-d exhibits vector field of constant direction and amplitude that conserves |σ 1 | -|σ 2 |.

Figure2: Learning curves of input-output map and loss profile for various stiffness p levels: p = 0 corresponds to NGD update, 1/2 ≤ p < 1 corresponds to √ NGD update, and 1 ≤ p < 2 corresponds to SGD update for network depth ranging between 2 ≤ d < ∞. Top: Learning curves of map singular modes σ(t) from eq (21). Dashed lines show the mode-strength of dataset σ * . Note that large p increases the stiffness of dynamics, i.e. extreme changes of time-scale: between extreme slow and extreme fast. Half-max points (black circles) are shown to visualize the overall time-scale of learning dynamics, which decreases with mode strength as σ-p * . Bottom: Corresponding loss profiles. Initial conditions: σ(0) = 0 for p < 1, and σ(0) = σ * /100 for p ≥ 1. η = 1.

Figure 3: Curvature correction effect on generalization dynamics: (A) Singular mode strength of input-output correlation of training dataset. Dataset is generated from a rank-3 teacher network with added noise (SNR = 10). (B, C) Training and testing loss profiles of a 3-hidden-layer student network (See text).Note that all methods eventually converge to the identically overfitted solution (horizontal dashed-line). (D) Time-separated learning dynamics of across singular modes. To obtain the individual mode components, we computed the network's input-output correlation matrix, and projected it to the singular vector basis of the data correlation matrix.

Figure S.I.1: Scaling relationships for normalized loss update L/L and weight/parameter update speed ẇ /√L shown across different depth d and curvature correction factor q. α is the overall scale factor for the weights. Dotted-lines show predictions for the case of balanced weights, eq (S.I.16)(S.I.17). The numerical results show the result from networks of random weights sampled from Gaussian distribution, (Xavier normalization(Glorot & Bengio, 2010)), They exhibit the same power-law exponent as the dotted-lines (i.e. same slope), but shifted toward right. Blockdiagonal approximation does not affect the power-law exponents, but it affects the coefficient such that NGD-d exhibits much larger parameter update size than NGD, while exhibiting identical zero power-law for loss update. In contrast, sNGD-d exhibits identical match with sNGD for both parameter update and loss update. See S.I. for detailed setup. Error bars show standard deviation from 10 simulations with random weights.

mode should have at least one non-zero σ i , which according to eq (S.I.11), implies σ S = 0. Therefore, from eq (S.I.12),

Peer review (iclr2020):
In this paper, the authors study the training dynamics of natural gradient descent with linear DNNs. Specifically, they showed that the curvature corrected natural gradient descents preserve the learning trajectory of plain gradient descent and only affect the temporal dynamics. A fractional natural gradient descent method that only applies partial curvature correction is proposed to address the numerical stability issue of vanilla natural gradient descent. The paper is well presented and the derivations of analytical solutions are clear. Although the analysis is limited to the linear case, it does give some insight into the behaviors of curvature corrected gradient descents. However, it would be interesting to see how these algorithms would perform under the no-linear case. Also, will these results hold for other neural network structures? For example, the Hessian of each layer of ResNet will be close to orthogonal and how will that affect the NGD? Overall I think this is an interesting paper with strong results and vote for accepting.