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: The current work studies the implications of continuous symmetries of a DNN on its weight dynamics. Specifically, they consider linearly realized symmetries (such as a shift/translation to the logits), and use the fact that the loss/hessian/mini-batches, are insensitive to such shift to decouple their dynamics. They do so both in the continuous/vanishing-learning-rate case and for small learning rates and manage to provide accurate quantitative predictions for the dynamics of these quantities.