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 authors identify a "kinetic asymmetry" in the learning dynamics under gradient flow, where the kinetic energy is not invariant to the symmetries of the potential energy. An example is scaling: even when the network function is invariant to scaling or ReLUs or batch norm inputs, the learning dynamics are not. They derive a genealized version of Noether's theorem which takes this asymmetry into account and use it to study the implict adaptive optimization exhibited by batch normalization, showing that it is the same as that of RMSprop.