EXACT ANALYSIS OF CURVATURE CORRECTED LEARNING DYNAMICS IN DEEP LINEAR NETWORKS

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 sta- bility 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 ≡ { 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 . 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 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 .

Summary: 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.