C URVATURE - CORRECTED LEARNING DYNAMICS IN DEEP LINEAR NEURAL NETWORKS
Anonymous authors
Paper under double-blind review
A BSTRACT
Deep neural networks exhibit complex learning dynamics due to the highly nonconvex loss landscape, which causes slow convergence and vanishing gradient problems.
Second order approaches, such as natural gradient descent, mitigate such problems by neutralizing the effect of potentially ill-conditioned curvature on the gradient-based updates, yet precise theoretical understanding on how such curvature correction affects the learning dynamics of deep networks has been lacking.
Here, we analyze the dynamics of training deep neural networks under a generalized family of natural gradient methods that applies curvature corrections, and derive precise analytical solutions.
Our analysis reveals that curvature corrected update rules preserve many features of gradient descent, such that the learning trajectory of each singular mode in natural gradient descent follows precisely the same path as gradient descent, while 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 signiﬁcantly distort the learning trajectory into highly diverging dynamics that signiﬁcantly 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 beneﬁt of full curvature correction in terms of convergence speed, with additional beneﬁt of superior numerical stability and neutralizing vanishing/exploding gradient problems, which holds true also in layer-restricted approximations.
I NTRODUCTION
Difﬁculty 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 ﬁrst 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 deﬁne 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 efﬁcient 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 ﬁrst 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 proﬁles (Saxe et al., 2013), and showed the implicit bias for regularization (Du et al., 2018; Arora et al., 2019) and resistance to overﬁtting (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.
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.
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, sacriﬁces stability of parameter update dynamics to explode when gradient vanishes and vice versa.
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.
NGD makes the learning process prone to overﬁtting by simultaneously learning both the signal and the noise dimensions of data, whereas √ NGD partially retains gradient descent’s resistance to overﬁtting by separating the time-scales between the signal and the noise dimensions.
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 efﬁcient and stable learning algorithms.
S ETUP 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 } di =1 that connect the adjacent layers.
The network’s input-output map is ¯ w ≡ (cid:81) di =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:
[Equation]
where E D is the expectation over the dataset D , Σ x ≡ E D [ xx (cid:124) ] is the input correlations, and ¯ w ∗ ≡ E D [ yx (cid:124) ]
Σ − 1 x .
Neglecting the constant term, the loss function is expressed as
[Equation]
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.
Gradient and Hessian
+
We use bold symbols to collectively represent network parameters and derivatives in array form.
For example, ˙ w ≡
[Equation]
and g ≡ (cid:20)
[Equation]
∆ w
[Equation]
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 deﬁnition, produces gradient update:
[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-deﬁnite (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 ˙ w · H + ˙ w =
Tr [ ˙ w (cid:124) 1 w (cid:124) 2 ˙∆
+
˙
w
cid:124) 2 ˙∆ w (cid:124) 1 ] =
Tr [ ˙∆ ˙∆ (cid:124) ] ≥ 0 , where the dot-product denotes a · b ≡ (cid:80) di
=
Tr [ a i b (cid:124) 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.
[Equation]
are the equivalent continuous-time transformations that yield the invariance ˙∆
= ˙¯ w = w 2 mw 1 − w 2 mw 1 = 0 , ∀ m .
These transformations form the null-space of H + , since ˙ w null · H + ˙ w null =
Tr [ ˙∆ ˙∆ (cid:124) ] = 0 , which is orthogonal to gradient, since ˙ w null · g =
Tr [∆ ˙∆ (cid:124) ]
= 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.
˙
[Equation] [Equation]
null
[Equation]
exhibit the following conservation law
[Equation]
This result was previously only known for gradient descent dynamics (Arora et al., 2018b; Du et al., 2018), which is generalized here.
L EARNING 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 modiﬁes the curvature-corrected dynamics (Section 3.3).
3. PARAMETER DYNAMICS
We follow the singular value decomposition (SVD)-based analysis of Saxe et al.
; 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 ∗ .
Under such initialization, the update dynamics of weight matrices simpliﬁes to their singular value dynamics, with their singular vectors remain unchanging.
This simpliﬁed case admits exact analytic solutions, which provide good approximation to general learning dynamics.
Moreover, this aligned singular vector condition is automatically satisﬁed for networks initialized with small random weights Saxe et al.
; 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 by 3 (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 (cid:107) ˙ σ (cid:107) is proportional to the displacement | σ ∆ | and the coupling strength (cid:107) j (cid:107) (cid:107) ˙
σ (cid:107) ∝
| σ ∆ |(cid:107) j (cid:107) , ( (cid:107) ˙
σ (cid:107) 2 ≡ (cid:80) di =1 ˙ σ
[Equation]
which vanishes for networks with small coupling strength and explodes for large coupling strength.
Natural gradient descent (NGD) NGD ﬁnds the minimum-norm update solution ( min ˙ w · ˙ w ) 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 Λ satisﬁes
[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) is
[Equation]
(12) where σ ∆ of SGD eq (6) is replaced by σ Λ = σ ∆ / (cid:107) j (cid:107) 2 , the singular values of Λ (See S.I.).
NGD dynamics eq (12) follows the same hyperbolic paths of SGD eq (7), but with modiﬁed update speed
[Equation]
(13) which inversely scales with (cid:107) j (cid:107) .
Therefore, NGD’s update speed explodes for small coupling strength, reciprocal to SGD’s vanishing speed phenomenon.
[Figure]
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
[Equation]
(14) which interpolates between NGD ( q = 1 ) and SGD ( q → ∞ ).
Eq (14) follows the same hyperbolic paths of SGD eq (7), but with modiﬁed 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 +
[Equation]
+
[Equation]
I
[Equation] [Equation]
≥ 0 ), which yields the regularized (or damped) inverse
[Equation] [Equation]
similar to Levenberg-Marquardt damping (less the ( (cid:15) + 1) term), whose singular mode dynamics is
[Equation]
[Figure]
where the ratio a ≡ (cid:15)/ (cid:107) j (cid:107) describes the effective degree of interpolation between NGD ( a → 0 ) and SGD ( a → ∞ ).
Note that a should be large enough to provide sufﬁcient damping but not too large to nullify the effect of curvature correction which is difﬁcult to simultaneously satisfy across all singular modes with fkxed (cid:15) .
√ NGD can be considered as providing ideally and adaptively tuned regularization ( a = 1 ) for all singular modes, where the regularization is most effective.
3. M AP 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 (
[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 ﬁeld 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 (cid:124)
i − w
[Equation]
i +1 w i +
= 0 .
Note that this balanced weight condition is automatically approximately satisﬁed if the networks are initialized with small random weights.
Under the balanced weight condition, eq (20) simpliﬁes 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:
[Equation]
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 ﬁnite 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 sacriﬁces 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 inﬁnite depth limit.
Its learning curves exhibit polynomial growth near zero, ¯ σ ( t ) ∝ t 1 / (1 − p ) , which takes ﬁnite 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 inﬁnite depth limit.
Its learning curves exhibit sigmoidal shape that take inﬁnite 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]
(22) which approaches the actual depth d eff → d in the SGD limit ( q → ∞ ), and similarly, approaches d eff → q in the limit of inﬁnite 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 sacriﬁces 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. 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): ( d = 2 example) (cid:20)
[Equation]
which nevertheless satisﬁes the NGD constraint (9) if (cid:80) di
=1 η i = η .
H i denotes the block-diagonal
Hessian + term of layer i .
Singular mode dynamics of eq (23) is (with η i =
η/d )
[Equation]
(24) where the layer-restricted factor j
i substitutes the full curvature correction factor (cid:107) j (cid:107) 2 of NGD (12).
This block-diagonalization signiﬁcantly modiﬁes 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 efﬁcient, 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 + η g i /d
1 /q = 0 yields
[Equation]
(25) which conserves σ 2(1 − 1 /q )
i −
[Equation]
) k as constants of motion.
For q = 2 , called √ NGD-d, the singular mode dynamics
[Equation]
(26) follows non-diverging, parallel paths that conserve | σ i
|−| σ k | , with identical parameter update speed as √ NGD’s (cid:107) ˙
σ (cid:107) =
η | σ ∆ | (Fig 1E).
Therefore, √ NGD-d yields neutrally-stable update dynamics that neutralizes the vanishing/exploding update speed problems.
I MPLICIT 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 conﬁgurations 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 difﬁculty 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 (cid:107) w i (cid:107) 2 Frob − (cid:107) w i +1 (cid:107) 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.
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 conﬁrmation of this prediction in S.I.
The conservation of squared difference of norms for homogeneous ReLU networks still requires similar numerically conﬁrmation.
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
[Table]
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 conﬁrms 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 efﬁciently extract low-rank structure of the dataset, such as for ﬁnding matrix factorizations with minimum nuclear norm (Gunasekar et al., 2017; Arora et al., 2019).
It also allows deep networks to avoid overﬁtting via early stopping by ﬁrst learning the signal dimensions of noisy dataset, before the overﬁtting of the noise dimensions occurs, as long as signal-to-noise ratio is sufﬁciently 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 .
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 overﬁtting.
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 (cid:15) =
1 e − 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 ﬁrst learns the three signal modes, well separated from the onset of overﬁtting 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 proﬁle deviates from exponential decay due to the clipping.
In contrast, √ NGD and √ NGD-d allow fast learning while separating the signal dimensions from the noise dimensions, achieving comparable test loss as SGD update, but also with fast earlystopping time comparable to NGD update.
Note that all three update rules achieve the same test loss after overﬁtting is complete.
This is due to the shared learning path for each singular mode across the methods.
C ONCLUSION
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
[Figure]
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, sacriﬁces 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 overﬁtting by simultaneously learning both the signal and the noise dimensions of data, whereas √ NGD partially retains gradient descent’s resistance to overﬁtting 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 efﬁcient and stable learning algorithms.
Supplemental Materials S.I. M OORE -P ENROSE INVERSE SOLUTION : EQ (15,16,17) IN S ECTION 3.
In section 4, we ﬁnd the Moore-Penrose inverse solution that minimizes the update norm ˙ w · ˙ w = (cid:80) i
Tr [ ˙ w i ˙ w (cid:124)
i ] while satisfying the natural gradient constraint: ( d = 2 example)
[Equation]
This constrained optimization problem is described by the following Lagrangian:
[Equation]
where ˙∆ = w 2 ˙ w 1 + ˙ w 2 w 1 , and dot notation denotes inner-product: a · b
≡ Tr [ a (cid:124) b ] .
Optimality condition on ˙ w i yields
[Equation]
which, via change of variables Λ ≡ ( w 2 Λ 1 + Λ 2 w 1 ) /η , reduces to
[Equation]
which can be plugged into the optimality condition on Λ
i
[Equation]
to produce a linear equation for Λ i :
[Equation]
where
[Equation]
Note that if w 2 , w 1 are invertible, it is easy to see that eq (S.I.6)(S.I.7) reduce to exponentially converging dynamics ˙∆ + η ∆ = ˙¯ w + η ( ¯ w − ¯ w ∗ )
= 0 , with the solution of S (Λ) =
0 driving the parameter update eq (S.I.4)(S.I.5).
This result also holds true for the over-complete cases, where the hidden layer width is larger than the minimum of input layer or output layer size.
For the undercomplete cases, i.e. with bottleneck hidden layers, the exponential convergence applies only to the subspace dimensions permitted by the bottleneck, with the other dimensions remain frozen.
S.I. S INGULAR MODE ANALYSIS EQ (18) IN S ECTION 3.
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 = (cid:81) k
(cid:54) =
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 σ i σ S =(S.I.11)
[Equation]
(S.I.12)
An active singular 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),
[Equation]
which plugs into (S.I.10) to produce the result in the main text
[Equation]
S.I. S CALING 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 (cid:107) j (cid:107) :
[Equation]
deﬁned for individual singular modes under the well-aligned singular vector condition.
For the balanced weight case where the coupling strength term is constant across all singular modes, eq (S.I.15) reduces to more generally applicable scaling relationships:
[Equation] [Equation]
17) where α is the overall scale factor for the weight matrices.
Instead of the detailed description of individual singular mode dynamics, eq (S.I.16)(S.I.17) encapsulate the overall scaling law between weights and update speed that can be readily measured without requiring the aligned singular vector condition.
Under SGD ( q → ∞ ), weight update scales as α d − 1 , and double the power for the loss update, exhibiting the vanishing update problem for small α .
Under sNGD ( q = 1 / 2 ), weight update is constant (zero power-law) with respect to α and loss update scales as α d − 1 .
Under NGD ( q = 1 ), loss update is constant (zero power-law) with respect to α , but the weight update inversely scales as α − ( d − 1) , exhibiting the exploding update problem for small α .
Numerical experiments indeed conﬁrm these scaling laws (Figure S1): For the case of balanced weights the predictions hold exactly, and approximately for random weight matrices.
Note that the zero-power law predictions for weight update under sNGD and for loss update under NGD are exact, because they indeed satisfy the balanced condition: the coupling strength term with zero exponent (cid:107) j (cid:107) 0 is indeed constant.
[Figure]
N UMERICAL CONFIRMATION OF CONSERVATION LAW
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 ( w 1 , w 2 , w 3 ).
It also shows the conserved quantities w 1 w (cid:124) 1 − w (cid:124) 2 w 2 , w 2 w (cid:124) 2 − w (cid:124) 3 w 3 , 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.
s