Neural Mechanics: Symmetry and Broken Conservation Laws in Deep Learning Dynamics

1 INTRODUCTION . Just like the fundamental laws of classical and quantum mechanics taught us how to control and optimize the physical world for engineering purposes , a better understanding of the laws governing neural network learning dynamics can have a profound impact on the optimization of artificial neural networks . This raises a foundational question : what , if anything , can we quantitatively understand about the learning dynamics of large-scale , non-linear neural network models driven by real-world datasets and optimized via stochastic gradient descent with a finite batch size , learning rate , and with or without momentum ? In order to make headway on this extremely difficult question , existing works have made major simplifying assumptions on the network , such as restricting to identity activation functions Saxe et al . ( 2013 ) , infinite width layers Jacot et al . ( 2018 ) , or single hidden layers Saad & Solla ( 1995 ) . Many of these works have also ignored the complexity introduced by stochasticity and discretization by only focusing on the learning dynamics under gradient flow . In the present work , we make the first step in an orthogonal direction . Rather than introducing unrealistic assumptions on the model or learning dynamics , we uncover restricted , but meaningful , combinations of parameters with simplified dynamics that can be solved exactly without introducing a major assumption ( see Fig . 1 ) . To find the parameter combinations , we use the lens of symmetry to show that if the training loss doesn ’ t change under some transformation of the parameters , then the gradient and Hessian for those parameters have associated geometric constraints . We systematically apply this approach to modern neural networks to derive exact integral expressions and verify our predictions empirically on large scale models and datasets . We believe our work is the first step towards a foundational understanding * Equal contribution . Correspondence to kunin @ stanford.edu & hidenori.tanaka @ ntt-research.com of neural network learning dynamics that is not based in simplifying assumptions , but rather the simplifying symmetries embedded in a network ’ s architecture . Our main contributions are : 1 . We leverage continuous differentiable symmetries in the loss to unify and generalize geometric constraints on neural network gradients and Hessians ( section 3 ) . 2 . We prove that each of these differentiable symmetries has an associated conservation law under the learning dynamics of gradient flow ( section 4 ) . 3 . We construct a more realistic continuous model for stochastic gradient descent by modeling weight decay , momentum , stochastic batches , and finite learning rates ( section 5 ) . 4 . We show that under this more realistic model the conservation laws of gradient flow are broken , yielding simple ODEs governing the dynamics for the previously conserved parameter combinations ( section 6 ) . 5 . We solve these ODEs to derive exact learning dynamics for the parameter combinations , which we validate empirically on VGG-16 trained on Tiny ImageNet with and without batch normalization ( section 6 ) . 2 RELATED WORK The goal of this work is to construct a theoretical framework to better understand the learning dynamics of state-of-the-art neural networks trained on real-world datasets . Existing works have made progress towards this goal through major simplifying assumptions on the architecture and learning rule . Saxe et al . ( 2013 ; 2019 ) and Lampinen & Ganguli ( 2018 ) considered linear neural networks with specific orthogonal initializations , deriving exact solutions for the learning dynamics under gradient flow . The theoretical tractability of linear networks has further enabled analyses on the properties of loss landscapes Kawaguchi ( 2016 ) , convergence Arora et al . ( 2018a ) ; Du & Hu ( 2019 ) , and implicit acceleration by overparameterization Arora et al . ( 2018b ) . Saad & Solla ( 1995 ) and Goldt et al . ( 2019 ) studied single hidden layer architectures with non-linearities in a studentteacher setup , deriving a set of complex ODEs describing the learning dynamics . Such shallow neural networks have also catalyzed recent major advances in understanding convergence properties of neural networks Du et al . ( 2018b ) ; Mei et al . ( 2018 ) . Jacot et al . ( 2018 ) considered infinitely wide neural networks with non-linearities , demonstrating that the network ’ s prediction becomes linear in its parameters . This setting allows for an insightful mathematical formulation of the network ’ s learning dynamics as a form of kernel regression where the kernel is defined by the initialization ( though see also Fort et al . ( 2020 ) ) . Arora et al . ( 2019 ) extended these results to convolutional networks and Lee et al . ( 2019 ) demonstrated how this understanding also allows for predictions of parameter dynamics . In the present work , we make the first step in an orthogonal direction . Instead of introducing unrealistic assumptions , we discover restricted combinations of parameters for which we can find exact solutions , as shown in Fig . 1 . We make this fundamental contribution by constructing a framework harnessing the geometry of the loss shaped by symmetry and realistic continuous equations of learning . Geometry of the loss . A wide range of literature has discussed constraints on gradients originating from specific architectural building blocks of networks . For the first part of our work , we simplify , unify , and generalize the literature through the lens of symmetry . The earliest works understanding the importance of invariances in neural networks come from the loss landscape literature Baldi & Hornik ( 1989 ) and the characterization of critical points in the presence of explicit regularization Kunin et al . ( 2019 ) . More recent works have studied implicit regularization originating from linear Arora et al . ( 2018b ) and homogeneous Du et al . ( 2018a ) activations , finding that gradient geometry plays an important role in constraining the learning dynamics . A different line of research studying the generalization capacity of networks has noticed similar gradient structures Liang et al . ( 2019 ) . Beyond theoretical studies , geometric properties of the gradient and Hessian have been applied to optimize Neyshabur et al . ( 2015 ) , interpret Bach et al . ( 2015 ) , and prune Tanaka et al . ( 2020 ) neural networks . Gradient properties introduced by batch Ioffe & Szegedy ( 2015 ) , weight Salimans & Kingma ( 2016 ) and layer Ba et al . ( 2016 ) normalization have been intensely studied . Van Laarhoven ( 2017 ) ; Zhang et al . ( 2018 ) showed that normalization layers are scale invariant , but have an implicit role in controlling the effective learning rate . Cho & Lee ( 2017 ) ; Hoffer et al . ( 2018 ) ; Chiley et al . ( 2019 ) ; Li & Arora ( 2019 ) ; Wan et al . ( 2020 ) have leveraged the scale invariance of batch normalization to understand geometric properties of the learning dynamics . Most recently , Li et al . ( 2020 ) studied the role of gradient noise in reconciling the empirical dynamics of batch normalization with the theoretical predictions given by continuous models of gradient descent . Equations of learning . To make experimentally testable predictions on learning dynamics , we introduce a continuous model for stochastic gradient descent ( SGD ) . There exists a large body of works studying this subject using stochastic differential equations ( SDEs ) in the continuous-time limit Mandt et al . ( 2015 ; 2017 ) ; Li et al . ( 2017 ) ; Smith & Le ( 2017 ) ; Chaudhari & Soatto ( 2018 ) ; Jastrzębski et al . ( 2017 ) ; Zhu et al . ( 2018 ) ; An et al . ( 2018 ) . Each of these works involves making specific assumptions on the loss or the noise in order to derive stationary distributions . More careful treatment of stochasticity led to fluctuation dissipation relationships at steady state without such assumptions Yaida ( 2018 ) . In our analysis , we apply a more recent approach Li et al . ( 2017 ) ; Barrett & Dherin ( 2020 ) , inspired by finite difference methods , that augments SDE model with higher-order terms to account for the effect of a finite step size and curvature in the learning trajectory . 3 SYMMETRIES IN THE LOSS SHAPE GRADIENT AND HESSIAN GEOMETRIES . While we initialize neural networks randomly , their gradients and Hessians at all points in training , no matter the loss or dataset , obey certain geometric constraints . Some of these constraints have been noticed previously as a form of implicit regularization , while others have been leveraged algorithmically in applications from network pruning to interpretability . Remarkably , all these geometric constraints can be understood as consequences of numerous differentiable symmetries in the loss introduced by neural network architectures . A set of parameters observes a differentiable symmetry in the loss if the loss doesn ’ t change under a certain differentiable transformation of these parameters . This invariance introduces associated geometric constraints on the gradient and Hessian . Consider a function f ( θ ) where θ ∈ Rm . This function possesses a differentiable symmetry if it is invariant under the differentiable action ψ of a group G on the parameter vector θ , i.e. , if θ 7→ ψ ( θ , α ) where α ∈ G , then F ( θ , α ) = f ( ψ ( θ , α ) ) = f ( θ ) for all ( θ , α ) . The existence of a symmetry enforces a geometric structure on the gradient , ∇F . Evaluating the gradient at the identity element I of G ( so that for all θ , ψ ( θ , I ) = θ ) yields the result , 〈∇f , ∂αψ|α=I〉 = 0 , ( 1 ) where 〈 , 〉 denotes the inner product . This equality implies that the gradient ∇f is perpendicular to the vector field ∂αψ|α=I that generates the symmetry , for all θ . The symmetry also enforces a geometric structure on the Hessian , HF . Evaluating the Hessian at the identity element yields the result , Hf [ ∂θψ|α=I ] ∂αψ|α=I + [ ∂θ∂αψ|α=I ] ∇f = 0 , ( 2 ) which constrains the Hessian Hf . See appendix A for the derivation of these properties and other geometric consequences of symmetry . We will now consider the specific setting of a neural network parameterized by θ ∈ Rm , the training loss L ( θ ) , and three families of symmetries ( translation , scale , and rescale ) that commonly appear in modern network architectures . Translation symmetry . Translation symmetry is defined by the group R and action θ 7→ θ + α1A where 1A is the indicator vector for some subsetA of the parameters { θ1 , . . . , θm } . The loss function L thus possesses translation symmetry if L ( θ ) = L ( θ + α1A ) for all α ∈ R. Such a symmetry in turn implies the loss gradient ∂θL = g is orthogonal to the indicator vector ∂αψ|α=I = 1A , 〈g,1A〉 = 0 , ( 3 ) and that the Hessian matrix H = ∂2θL has the indicator vector in its kernel , H1A = 0 . ( 4 ) Softmax function . Any network using the softmax function gives rise to translation symmetry for the parameters immediately preceding the function . Let z = Wx + b be the input to the softmax function such that σ ( z ) i = e zi∑ j e zj . Notice that shifting any column of the weight matrix wi or the bias vector b by a real constant has no effect on the output of the softmax as the shift factors from both the numerator and denominator canceling its effect . Thus , the loss function is invariant w.r.t . this translation , yielding the gradient constraints 〈 ∂L∂wi ,1〉 = 〈 ∂L ∂b ,1〉 = 0 , visualized in Fig . 2 for the toy model where wi ∈ R2 . Scale symmetry . Scale symmetry is defined by the group GL+1 ( R ) and action θ 7→ αA θ where αA = α1A + 1Ac . The loss function possesses scale symmetry if L ( θ ) = L ( αA θ ) for all α ∈ GL+1 ( R ) . This symmetry immediately implies the loss gradient is everywhere perpendicular to the parameter vector itself ∂αψ|α=I = θ 1A = θA , 〈g , θA〉 = 0 , ( 5 ) and relates to the Hessian matrix , where [ ∂α∂θψ|α=I ] ∂θL = diag ( 1A ) g = gA , as HθA + gA = 0 . ( 6 ) Batch normalization . Batch normalization leads to scale invariance during training . Let z = wᵀx+ b be the input to a neuron with batch normalization such that BN ( z ) = z−E [ z ] √ Var ( z ) where E [ z ] is the sample mean and Var ( z ) is the sample variance given a batch of data . Notice that scaling w and b by a non-zero real constant has no effect on the output of the batch normalization as it factors from z , E [ z ] , and Var ( z ) canceling its effect . Thus , these parameters observe scale symmetry in the loss and their gradients satisfy 〈 ∂L ∂w , w 〉 + 〈 ∂L ∂b , b 〉 = 0 , as has been previously noted by Ioffe & Szegedy ( 2015 ) ; Van Laarhoven ( 2017 ) , and visualized in Fig . 2 for the toy model where w , b ∈ R. Rescale symmetry . Rescale symmetry is defined by the group GL+1 ( R ) and action θ 7→ αA1 α−1A2 θ where A1 and A2 are two disjoint sets of parameters . The loss function possesses rescale symmetry if L ( θ ) = L ( αA1 α−1A2 θ ) for all α ∈ GL + 1 ( R ) . This symmetry immediately implies the loss gradient is everywhere perpendicular to the sign inverted parameter vector ∂αψ|α=I = θA1 − θA2 = θ ( 1A1 − 1A2 ) , 〈g , θA1 − θA2〉 = 0 ( 7 ) and relates to the Hessian matrix , where [ ∂α∂θψ|α=I ] ∂θL = diag ( 1A1 − 1A2 ) g = gA1 − gA2 , as H ( θA1 − θA2 ) + gA1 − gA2 = 0 . ( 8 ) Homogeneous activation . For networks with continuous , homogeneous activation functions φ ( z ) = φ′ ( z ) z ( e.g . ReLU , Leaky ReLU , linear ) , this symmetry emerges at every hidden neuron by considering all incoming and outgoing parameters to the neuron . For example , consider a hidden neuron with ReLU activation φ ( z ) = max { 0 , z } , such that w2φ ( wᵀ1x + b ) is the computational path through this neuron . Scaling w1 and b by a real constant and w2 by its inverse has no effect on the computational path as the constants can be passed through the ReLU activation canceling their effects . Thus , these parameters observe rescale symmetry and their gradients satisfy〈 ∂L ∂w1 , w1 〉 + 〈 ∂L ∂b , b 〉 − 〈 ∂L ∂w2 , w2 〉 = 0 , as has been previously noted by Du et al . ( 2018a ) ; Liang et al . ( 2019 ) ; Tanaka et al . ( 2020 ) , and visualized in Fig . 2 for the toy model where w1 , w2 ∈ R and b = 0 .

Summary: This paper studies the dynamics of the parameters while training a neural network via SGD. SGD is not studied directly but three continuous time approximations of SGD are considered: the classical gradient flow, a stochastic differential equation of Langevin type, and a "modified" gradient flow whose derivation has roots in the literature. For each dynamics, the authors show that some invariant properties of the loss function (which are often satisfied in practice) imply some invariant quantities for the dynamics.