Noether’s Learning Dynamics: Role of Symmetry Breaking in Neural Networks

q0 q1 q2 Physical space ( b ) Classical Mechanics q0 q1 qN Parameter space ( a ) Learning Dynamics Figure 1 : Learning dynamics in analogy to classical mechanics . ( a ) Learning dynamics of a neural network in parameter space . ( b ) Motion of a particle in physical space . With the rapid increase in available data and computational power , artificial neural networks have become an essential tool both in science and engineering . To obtain an accurate model , one has to overcome the challenge of tuning millions of randomly initialized parameters q 2 RN to minimize the loss function f ( q ) 2 R , a measure of the discrepancy between the model ’ s predictions from the input data and the provided truth . In minimizing the loss , the learning rule specifies how to iteratively update these parameters based on the recent trajectory and the local geometry of the non-convex loss landscape . We can represent the repeated updates of the millions of parameters during learning as discrete movements of a point in high-dimensional parameter space ( Fig . 1a ) . The architecture shapes the static loss landscape and the learning rule governs the motion of the model through parameter space during training . However , it is the complex interaction of the architecture and learning rule which determines the final state of the trained model . Inspired by tools and ideas from classical mechanics ( Fig . 1b ) , we develop a theoretical framework which unifies the effects of the architecture and learning rule on the geometry of the learning dynamics . * Correspondence to : Hidenori Tanaka < hidenori.tanaka @ ntt-research.com > 2 Work done during an internship at Physics & Informatics Laboratories , NTT Research , Inc. 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Symmetry as a mathematical lens . One of the most powerful mathematical tools for studying the geometry of complex objects is symmetry . A symmetry of an object is a transformation that when applied to the object leaves certain properties unchanged . For example , consider how a circle can be rotated around its center without changing its overall appearance . In otherwords , a circle has a continuous rotational symmetry . In particular , the form of the symmetry transformation constrains the form of the object , the circle is the only 2D shape that has continuous rotational symmetry . Objects with symmetry are not restricted to geometric shapes , but could be functions or even dynamical systems . Recently , the lens of symmetry has been applied in the study of neural networks [ 1 , 2 ] . A recent work [ 1 ] identified an array of symmetries in neural networks to study the geometry of the loss function f ( q ) in the parameter space . For example , batch normalization [ 3 ] defined as BN ( qx ) = ( qx E [ qx ] ) / p Var [ qx ] introduces a scale symmetry in the parameters immediately preceding the normalization layer , BN ( ( 1 + s ) qx ) = BN ( qx ) for any positive scalar s 2 R+ . Geometry of learning dynamics . What does the symmetry of the loss function in parameter space imply about the symmetry of the learning dynamics ? In classical mechanics , a symmetry of the potential function in space usually implies the equivariance of the trajectory , namely , a solution of the Newton ’ s equation of motion ( grey ) remains a solution under the symmetry transformations ( red ) as shown in Fig . 2 ( a ) . This implies that the time evolution of the potential energy is invariant under symmetry transformations at any point in the dynamics . However , this intuition does not hold in modern neural networks . Consider the concrete example of a VGG11 model with a batch normalization layer after each convolutional filter . The scale symmetry of batch normalization implies the loss function is invariant to the scaling transformation on the convolutional filters q , f ( ( 1 + s ) qx ) = f ( qx ) , however the learning dynamics are not . To demonstrate this we train two VGG11 models on Tiny ImageNet with a standard initialization ( grey ) and a scaled initialization ( red ) . Notice , that the scaling does not change the initial loss due to the scale symmetry , however it does effect the learning dynamics of the loss , as shown in the bottom of Fig . 2 ( b ) . In otherwords , even though the loss landscape observes scale symmetry at all steps in training , the learning dynamics do not . Understanding symmetry of the learning dynamics requires a new theoretical perspective . Implicit acceleration and variational perspective on learning dynamics . In physics , variational formulations based on the Lagrangian and Hamiltonian provide a unifying framework to study the role of symmetry in dynamical systems . However , there is still a gap between the tools and concepts in physics and the learning dynamics of neural networks . We identify the root cause of this gap to be the learning rule , namely the fact that the continuous-time limit of gradient descent is a first order differential equation in time , gradient flow dqdt = rf , unlike classical mechanics which is governed by Netwon ’ s second law md 2q dt2 = rf , which is second-order in time . We bridge this gap by realizing that the effect of discretization , is to introduce an implicit acceleration term proportional to the learning rate . Leveraging this connection , we build the mechanics of learning dynamics : Sec . 1 . From Newtonian to Lagrangian mechanics of learning via implicit acceleration . We show how gradient descent with a finite learning rate introduces implicit acceleration making it amendable to the Lagrangian formulation of accelerated optimizers [ 4 ] . Sec . 2 . Unlike in physics , kinetic symmetry breaking ( KSB ) is inherent in learning dynamics . We define kinetic symmetry breaking , where the kinetic energy corresponding to the learning rule explicitly breaks the symmetry of the potential energy corresponding to the loss function . Sec . 3 . We derive Noether ’ s learning dynamics ( NLD ) induced by the KSB . We generalize Noether ’ s theorem in physics to derive NLD , which accounts for damping , the unique symmetries of the loss , and the non-Euclidean metric used in learning rules . Sec . 4 . KSB induces impicit adaptive optimization and enables successful learning . We leverage NLD to exactly solve for the time-evolution of the effective learning rate for networks with normalization layers , and establish an exact analogy between two seemingly unrelated components of modern deep learning : normalization and adaptive optimization . 1 From Newtonian to Lagrangian mechanics of learning . The fundamental law of mechanics is the principle of least action [ 5 ] . We first review how pioneering works by Wibisono , Wilson , and Jordan [ 4 , 6 , 7 ] apply the variational principles of physics to model accelerated methods in optimization . We then show how gradient descent with a finite learning rate introduces implicit acceleration making it amendable to the Lagrangian formulation . Variational formulation of classical mechanics . Just over a century after Newton , Joseph-Louis Lagrange demonstrated how Newton ’ s equations of motion could be derived from a variational framework of least action . Rather than integrating an equation of motion over time to get the path an object takes from point a to point b , Lagrange prescribed that the path taken by the object is the path that minimizes the action , the path integral of a function termed the Lagrangian , over all possible paths from a to b . This principle of least action transformed mechanics from a differential description to a variational description . Consider for example the simple setup of a particle of mass m at position q 2 RN in N dimensional physical space moving according to a potential f ( q ) 2 R and Newton ’ s second law mq̈ = rf , where q̈ = d 2q dt2 is the acceleration . Let the Lagrangian be the function L ( q , q̇ , t ) = 12mq̇ 2 f ( q ) , the difference between the kinetic energy T = 12mq̇ 2 and the potential energy V = f ( q ) of the particle . The principle of least action states that the time evolution of the particle ’ s position q ( t ) is the function that minimizes the action S [ q ] = R t1 t0 L ( q , q̇ , t ) dt . In otherwords , q ( t ) , with boundary conditions q ( t0 ) = a and q ( t1 ) = b , must satisfy the stationary condition S q ( t ) = 0 , which is equivalent to satisfying the Euler-Lagrange equation ( see SI Sec . A for the review ) , d dt ✓ @ L @ q̇ ◆ = @ L @ q . ( 1 ) Indeed , substituting the Lagrangian into the Euler-Lagrange equation , we recover the Newton ’ s equation of motion mq̈ = rf confirming the consistency of this variational formulation . Damping in optimization . In learning dynamics , one of the universal features is damping , a term proportional to the velocity q̇ = dqdt in the equation motion . Intuitively , damping is needed for any optimization , since without damping the learning dynamics would conserve energy and not converge . We can incorporate damping into the previous Lagrangian example by introducing an explicit time-dependence L ( q , q̇ , t ) = e µm t ( 12mq̇ 2 f ( q ) ) , where µ is the damping ( or friction ) coefficient [ 8 ] . Substituting this Lagrangian into the Euler-Lagrange equation and dividing both sides by e µ m t , we obtain the equation of motion of a damped particle mq̈ + µq̇ = rf as intended . Non-Euclidean geometry in optimization . Many optimization algorithms assume an underlying geometry . For example , consider gradient descent with learning rate ⌘ . The update equation , q ( t + ⌘ ) = q ( t ) ⌘rf , is equivalent to solving the minimization problem , q ( t + ⌘ ) = argminq0 ⇥ hrf ( q ( t ) ) , q0i+ ⌘ 1 12 |q 0 q ( t ) |2 ⇤ , where h , i represents inner product . In otherwords , we can interpret each update of gradient descent as a trade-off ( specified by ⌘ ) between minimizing the loss f ( q ) , while remaining close to its current location in Euclidean distance . For non-Euclidean measures of distance we would obtain different optimization algorithms . Mirror descent generalizes this idea by replacing the distance measure by the Bergman divergence Dh ( y , x ) = h ( y ) h ( x ) hrh ( x ) , y xi where h ( x ) is a distance-generating function . Indeed , when the distance generating function is h ( x ) = 12 |x| 2 , then the Bregman divergence simplifies to be the squared Euclidean distance Dh ( y , x ) = 12 |x y| 2 and mirror descent is equivalent to gradient descent . The mirror descent framework encapsulates a broad array of optimization methods such as gradient descent , natural gradient descent , and others [ 9 ] . The Bregman Lagrangian . Integrating building blocks of optimization such as damping and nonEuclidean geometry into the variational framework of mechanics , Wibisono , Wilson , and Jordan introduced the Bregman Lagrangian [ 4 , 6 ] , L ( q , q̇ , t ) = e↵t+ t Dh ( q + e ↵t q̇ , q ) e tf ( q ) . ( 2 ) For various choices of the time-dependent factors ↵t , t , and t the Bregman Lagrangian , under the principle of least action , defines continuous approximations to an array of learning rules such as Nesterov ’ s accelerated gradient methods [ 10 , 11 ] and natural gradient flow [ 9 ] . The flexibility of this framework makes our analysis generalize to a broad array of learning rules , as summarized in table 1 . Implicit acceleration and the Lagrangian formulation of gradient descent . In the continuoustime limit of diminishing learning rate ⌘ ! 0 , gradient descent q ( t+ ⌘ ) = q ( t ) ⌘rf becomes a first order differential equation , gradient flow q̇ = rf . Unlike Nesterov ’ s momentum [ 11 ] , the trajectory of gradient descent with heavy ball momentum as applied in deep learning [ 12 ] , follows gradient flow with the rescaling ( 1 ) q̇ = rf . Here , we apply modified equation analysis [ 1 , 13 ] to incorporate the effect of a finite learning rate and leverage the fact that the discrete steps of gradient descent closely follow the continuous-time trajectory of the second-order differential equation in time . We can informally see this by Taylor expanding an update with respect to a small , but finite learning rate q ( t+ ⌘ ) = q + ⌘q̇ + ⌘ 2 2 q̈ +O ( ⌘ 3 ) and then plugging this expression to the gradient descent updates q ( t+⌘ ) q ( t ) ⌘ = rf . We keep the first order in ⌘ and obtain a second-order differential equation ⌘2 q̈ + q̇ = rf . We call the additional term ⌘ 2 q̈ , implicit acceleration . This dynamics reduces to gradient flow q̇ = rf in the continuous-time limit as expected . ( Note that a complementary approach is to view the effects of discrete step size as implicit gradient regularization on the original loss f ( q ) ! f ( q ) + ⌘4 |rf | 2 [ 1 , 14 , 15 , 16 , 17 ] . ) Performing a similar analysis [ 1 , 13 ] accounting for the effects of the heavy ball momentum with weight decay k yields a second-order differential equation , ⌘ 2 ( 1 + ) q̈ + ( 1 ) q̇ +rf ( q ) + kq = 0 , ( 3 ) where ⌘2 ( 1 + ) q̈ represents the implicit acceleration . Indeed , we can connect this second-order dynamics to the Bregman Lagrangian by considering the factors ↵t = log ⌘ ( 1+ ) 2 , t = log ⌘ ( 1+ ) 2 , t = 2 ( 1 ) ⌘ ( 1+ ) t , and the distance generating function h ( x ) = 12 |x| 2 . The Lagrangian describing the learning dynamics under gradient descent with a finite learning rate ⌘ , heavy ball momentum , the loss function f ( q ) , and weight decay constant k is L ( q , q̇ , t ) = e 2 ( 1 ) ⌘ ( 1+ ) t ⌘ ( 1 + ) 4 |q̇|2 ✓ f ( q ) + k 2 |q|2 ◆ . ( 4 )

Summary: The gradient descent with a momentum can be regarded as the Newtonian equation with damping. The output of the batch normalization is invariant to the scale of the input (i.e., a scale symmetry). When a BN is inserted after a convolution, the whole network is invariant to the weight scale. A typical gradient descent algorithm ignores (i.e., breaks) this symmetry. Thanks to this symmetry breaking, the learning rate is implicitly adapted, and the adaptation is equivalent to RMSProp's.