Original: S YMMETRY , CONSERVATION LAWS , AND LEARNING DYNAMICS IN NEURAL NETWORKS
Revised: S YMMETRY , C ONSERVATION L AWS , AND L EARNING D YNAMICS IN N EURAL N ETWORKS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Predicting the dynamics of neural network parameters during training is one of the key challenges in building a theoretical foundation for deep learning.
Revised: Predicting the dynamics of neural network parameters during training is one of the key challenges in building a theoretical foundation for deep learning.

Original: A central obstacle is that the motion of a network in high-dimensional parameter space undergoes discrete ﬁnite steps along complex stochastic gradients derived from real-world datasets.
Revised: A central obstacle is that the motion of a network in high-dimensional parameter space undergoes discrete ﬁnite steps along complex stochastic gradients derived from real-world datasets.

Original: We circumvent this obstacle through a unifying theoretical framework based on intrinsic symmetries embedded in a network’s architecture that are present for any dataset.
Revised: We circumvent this obstacle through a unifying theoretical framework based on intrinsic symmetries embedded in a network’s architecture that are present for any dataset.

Original: We show that any such symmetry imposes stringent geometric constraints on gradients and Hessians, leading to an associated conservation law in the continuous-time limit of stochastic gradient descent (SGD), akin to
Revised: We show that any such symmetry imposes stringent geometric constraints on gradients and Hessians, leading to an associated conservation law in the continuous-time limit of stochastic gradient descent (SGD), akin to

Original: Noether’s theorem in physics.
Revised: Noether’s theorem in physics.

Original: We further show that ﬁnite learning rates used in practice can actually break these symmetry induced conservation laws.
Revised: We further show that ﬁnite learning rates used in practice can actually break these symmetry induced conservation laws.

Original: We apply tools from ﬁnite difference methods to derive modiﬁed gradient ﬂow , a differential equation that better approximates the numerical trajectory taken by SGD at ﬁnite learning rates.
Revised: We apply tools from ﬁnite difference methods to derive modiﬁed gradient ﬂow , a differential equation that better approximates the numerical trajectory taken by SGD at ﬁnite learning rates.

Original: We combine modiﬁed gradient ﬂow with our framework of symmetries to derive exact integral expressions for the dynamics of certain parameter combinations.
Revised: We combine modiﬁed gradient ﬂow with our framework of symmetries to derive exact integral expressions for the dynamics of certain parameter combinations.

Original: We empirically validate our analytic predictions for learning dynamics on VGG-16 trained on Tiny ImageNet.
Revised: We empirically validate our analytic predictions for learning dynamics on VGG-16 trained on Tiny ImageNet.

Original: Overall, by exploiting symmetry, our work demonstrates that we can analytically describe the learning dynamics of various parameter combinations at ﬁnite learning rates and batch sizes for state of the art architectures trained on any dataset.
Revised: Overall, by exploiting symmetry, our work demonstrates that we can analytically describe the learning dynamics of various parameter combinations at ﬁnite learning rates and batch sizes for state of the art architectures trained on any dataset.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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 artiﬁcial neural networks.
Revised: 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 artiﬁcial neural networks.

Original: This raises a foundational question: what, if anything, can we quantitatively predict about the training dynamics of large scale deep learning models driven by real-world datasets?
Revised: This raises a foundational question: what, if anything, can we quantitatively predict 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 ﬁnite batch size, learning rate, and with or without momentum?

Original: The central challenge to answering this question is managing the complexity of the stochastic gradients and ﬁnite step sizes that drive learning.
Revised: In order to make headway on this extremely difﬁcult question, existing works have made major simplifying assumptions on the network, such as restricting to identity activation functions Saxe et al. (2013), inﬁnite width layers Jacot et al. (2018), or single hidden layers Saad &

Original: In this work, we disentangle the complexity introduced by  the dataset from geometric properties embedded in a network’s architecture that do not depend on data.
Revised: 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 ﬂow. In the present work, we make the ﬁrst 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 simpliﬁed dynamics that can be solved exactly without introducing a single assumption (see Fig. 1).

Original: To do this, 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.
Revised: To ﬁnd 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.

Original: We systematically apply this approach to modern neural networks to derive exact integral expressionsfor the dynamics of parameter combinations during training and verify our predictions empirically on large scale models and datasets. 
Revised: 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 ﬁrst step towards a foundational understanding of neural network learning dynamics that is not based in simplifying assumptions, but rather the simplifying symmetries embedded in a network’s architecture.

Original: Our main contributions are:
Revised: Our main contributions are:

Original: We  unify previously discovered geometric constraints on neural network gradients and Hessians by leveraging intrinsic translation, scale, and inversion symmetries (Sec. 3).
Revised: We leverage continuous differentiable symmetries in the loss to unify and generalize geometric constraints on neural network gradients and Hessians (section 3).

Original: We show that  these geometric constraints are data-independent, thereby constraining SGD noise to low-rank subspaces for any dataset (Sec.5).
Revised: We prove that each of these differentiable symmetries has an associated conservation law under the learning dynamics of gradient ﬂow (section 4). We construct

Original: We show that each of these symmetries has an associated constant of motion , which in the continuous time limit of stochastic gradient descent (SGD) is conserved (Sec.4).
Revised: a

Original: We demonstrate under weight decay that these previously conserved quantities decay exponentially, thereby providing an account of the impact of weight decay on training (Sec.4).
Revised: more realistic continuous model for stochastic gradient descent by modeling weight decay, momentum, stochastic batches, and ﬁnite learning rates (section 5). We show that under this more realistic model the conservation

Original: We apply tools from ﬁnite difference methods to introduce a stochastic differential equation (SDE) with modiﬁed gradients to model SGD even with ﬁnite-learning rates.
Revised: laws of gradient ﬂow are broken, yielding

Original: (Sec.5).
Revised: simple

Original:  We use this SDE to demonstrate that ﬁnite learning rates break continuous-time conversation laws, and ﬁnd exact solutions for the dynamics of quantities that were previously constant (Sec.6).
Revised: ODEs governing the dynamics for the previously conserved parameter combinations (section 6). We solve these ODEs to derive exact learning dynamics for the parameter combinations, which we validate

Original: We validate our exact solutions empirically on VGG-16with and without batch nornalization and trained on Tiny ImageNet (Sec.6).
Revised: empirically on VGG-16 trained on Tiny ImageNet with and without batch normalization (section

Original:  R ELATED WORK
Revised:  R ELATED WORK

Revised: [Figure]

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

Revised: Existing works have made progress towards this goal through major simplifying assumptions on the architecture and learning rule.

Revised: Saxe et al.

Original: Gradient geometry.
Revised: (2013) considered linear neural networks with speciﬁc orthogonal initializations, deriving exact solutions for the learning dynamics under gradient ﬂow. 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).

Revised: Saad & Solla (1995) studied a single hidden layer architecture with non-linearities in a student-teacher setup, deriving a set of complex ODEs describing the learning dynamics.

Revised: Such shallow neural networks have also catalyzed recent major advances in understanding convergence properties of neural networks Du et al.

Revised: (2018b); Mei et al.

Revised: Jacot et al.

Revised: (2018) considered inﬁnitely wide neural networks with non-linearities, demonstrating that the network’s prediction becomes linear in its parameters.

Revised: This setting allows for an insightful mathematical formulation of the network’s learning dynamics as a form of kernel regression where the kernel is deﬁned by the initialization.

Revised: Arora et al.

Revised: (2019) extended these results to convolutional networks and Lee et al.

Revised: (2019) demonstrated how this understanding also allows for predictions of parameter dynamics.

Revised: In the present work, we make the ﬁrst step in an orthogonal direction.

Revised: Instead of introducing unrealistic assumptions, we discover restricted combinations of parameters for which we can ﬁnd exact solutions, as shown in Fig.

Revised: We make this fundamental contribution by constructing a framework harnessing the geometry of the loss shaped by symmetry and realistic continuous equations of learning.

Revised: Geometry of the loss.

Original: A wide range of literature has discussed constraints on gradients originating from speciﬁc architectural building blocks of networks.
Revised: A wide range of literature has discussed constraints on gradients originating from speciﬁc architectural building blocks of networks.

Original: For the ﬁrst part of our work, we build on this literature and unify it through the framework of symmetry.
Revised: For the ﬁrst part of our work, we simplify, unify, and generalize the literature through the lens of symmetry.

Original: 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.
Revised: 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.

Original: More recent works have studied implicit regularization originating from linear Arora et al.
Revised: More recent works have studied implicit regularization originating from linear Arora et al.

Original: (2018b) and homogeneous Du et al.
Revised: (2018b) and homogeneous Du et al.

Original: (2018) activations, ﬁnding that gradient geometry plays an important role in biasing the training dynamics.
Revised: (2018a) activations, ﬁnding that gradient geometry plays an important role in constraining the learning dynamics.

Original: A different line of research studying the generalization capacity of networks have noticed similar gradient structures
Revised: A different line of research studying the generalization capacity of networks have noticed similar gradient structures

Original: Liang et al.
Revised: Liang et al.

Original: Beyond theoretical studies, geometric properties of the gradient and Hessian have been applied to optimize Neyshabur et al.
Revised: Beyond theoretical studies, geometric properties of the gradient and Hessian have been applied to optimize Neyshabur et al.

Original: (2015), interpret Bach et al.
Revised: (2015), interpret Bach et al.

Original: (2015), and prune Tanaka et al.
Revised: (2015), and prune Tanaka et al.

Original: (2020) neural networks.
Revised: (2020) neural networks.

Original: Gradient properties introduced by batch Ioffe & Szegedy (2015), weight Salimans & Kingma (2016) and layer Ba et al.
Revised: Gradient properties introduced by batch Ioffe & Szegedy (2015), weight Salimans & Kingma (2016) and layer Ba et al.

Original: (2016) normalization have been intensely studied.
Revised: (2016) normalization have been intensely studied.

Original: Van Laarhoven (2017); Zhang et al.
Revised: Van Laarhoven (2017); Zhang et al.

Original: (2018) showed that normalization layers are scale invariant, but have an implicit role in controlling the effective learning rate.
Revised: (2018) showed that normalization layers are scale invariant, but have an implicit role in controlling the effective learning rate.

Original: Cho & Lee (2017); Hoffer et al.
Revised: Cho & Lee (2017); Hoffer et al.

Original: ; Chiley et al.
Revised: ; Chiley et al.

Original: Li & Arora (2019); Wan et al.
Revised: Li & Arora (2019); Wan et al.

Original: (2020) have leveraged the scale invariance of batch normalization to understand geometric properties of the learning dynamics.
Revised: (2020) have leveraged the scale invariance of batch normalization to understand geometric properties of the learning dynamics.

Original: Modeling SGD.
Revised: 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.

Original: To make experimentally testable predictions on learning dynamics, we introduce a continuous model for stochastic gradient descent (SGD).
Revised: To make experimentally testable predictions on learning dynamics, we introduce a continuous model for stochastic gradient descent (SGD).

Original: There exists a large body of works studying this subject using stochastic differential equations (SDEs) in the continuous-time limit Mandt et al.
Revised: There exists a large body of works studying this subject using stochastic differential equations (SDEs) in the continuous-time limit Mandt et al.

Original: (2015; 2017); Li et al.
Revised: ; Li et al.

Original: (2017); Smith & Le (2017)
Revised: (2017); Smith & Le (2017)

Original: ; Chaudhari & Soatto (2018); Jastrz˛ebski et al.
Revised: ; Chaudhari & Soatto (2018); Jastrz˛ebski et al.

Original: (2017); Zhu et al.
Revised: ; Zhu et al.

Original: (2018); An et al.
Revised: (2018); An et al.

Original: Each of these works involves making speciﬁc assumptions on the loss or the noise in order to derive stationary distributions.
Revised: Each of these works involves making speciﬁc assumptions on the loss or the noise in order to derive stationary distributions.

Original: More careful treatment of stochasticity led to ﬂuctuation dissipation relationships at steady state without such assumptions Yaida (2018).
Revised: More careful treatment of stochasticity led to ﬂuctuation dissipation relationships at steady state without such assumptions Yaida (2018).

Original: In our analysis, we apply a more recent approach Li et al.
Revised: In our analysis, we apply a more recent approach Li et al.

Original: ; Barrett & Dherin (2020), inspired by ﬁnite difference methods, that augments SDE model with higher-order terms to account for the effect of a ﬁnite step size and curvature in the learning trajectory.
Revised: ; Barrett & Dherin (2020), inspired by ﬁnite difference methods, that augments SDE model with higher-order terms to account for the effect of a ﬁnite step size and curvature in the learning trajectory.

Original:  S YMMETRIES IN THE L OSS SHAPE G RADIENT AND H ESSIAN G EOMETRIES
Revised:  S YMMETRIES IN THE L OSS SHAPE G RADIENT AND H ESSIAN G EOMETRIES

Revised: [Figure]

Revised: We visualize the vector ﬁelds associated with simple network components that have translation, scale, and rescale symmetry.

Revised: In (a) we consider the vector ﬁeld associated with a neuron   [ ww 2 ] | x  where  is the softmax function.

Revised: In (b) we consider the vector ﬁeld associated with a neuron

Revised: BN  [ ww 2 ] [ xx 2 ] |  where BN is the batch normalization function.

Original:  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.
Revised: In (c) we consider the vector ﬁeld associated with a linear path w 2 w 1 x . 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.

Original: Some of these constraints have been noticed previously as a from of implicit regularization, while others have been leveraged algorithmically in applications from network pruning to interpretability.
Revised: Some of these constraints have been noticed previously as a from of implicit regularization, while others have been leveraged algorithmically in applications from network pruning to interpretability.

Original: Surprisingly, these geometric constraints canall be understood as consequences of numerous differentiable symmetries in the loss introduced by neural network architectures.
Revised: Remarkably, all these geometric constraints can be understood as consequences of numerous differentiable symmetries in the loss introduced by neural network architectures.

Original: 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.
Revised: 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.

Original: This invariance introduces associated geometric constraints on the gradient and Hessian.
Revised: This invariance introduces associated geometric constraints on the gradient and Hessian.

Original: Consider a function R ( ✓ )  .
Revised: Consider a function f ( ✓ ) where ✓ 2 R m .

Original: This function possesses a  symmetry if it is invariant under the  action of a group G on the parameter vector ✓ : i.e. if ✓ !
Revised: 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!

Original: ( ✓ , ↵ ) where ↵ 2 G , then R ( ✓  ) = R (  ) for all  ↵ 2 G .
Revised: ( ✓ , ↵ ) where ↵ 2 G , then F ( ✓ , ↵ ) = f ( ( ✓ , ↵ )) = f ( ✓ ) for all ( ✓ , ↵ ) .

Original: This symmetry implies d R d ↵ = h @ R @ ( ✓ ) @ ( ✓ ) @ ↵ i = 0 .
Revised: The existence of a symmetry enforces a geometric structure on the gradient, r F .

Original: Evaluating this expression at  identity element  ↵ 
Revised: Evaluating the gradient at the identity element of G , so that ( ✓ , ↵ )

Original: =I of G , so that = ✓ then yields the result that the gradient @ R @✓ is perpendicular to the vector ﬁeld @
Revised: = ✓ , yields the result, [Equation] which implies that the gradient r f is perpendicular to the vector ﬁeld @

Original: ( ✓ ) @↵ that generates the symmetry, for all ✓ .
Revised: ↵ that generates the symmetry, for all ✓ .

Original: To make this analysis concrete, we consider three families of symmetries, translation, scale, and inversion, that commonly appear in modern network architectures.
Revised: The symmetry also enforces a geometric structure on the Hessian, H F .

Original: 3. T RANSLATION
Revised: Evaluating the

Revised: Hessian at the identity element of G yields the result,

Original: [Figure]
Revised: [Equation]

Original: 
Revised: which constrains the Hessian H f . See appendix A for the derivation of these properties and other geometric consequences of symmetry. We will now consider the speciﬁc setting of a neural network parameterized by ✓ 2 R m , the training loss L ( ✓ ) , and three families of symmetries (translation, scale, and rescale) that commonly appear in modern network architectures.

Original: [Equation]
Revised: Translation symmetry.

Original: x and action  ✓ 7!
Revised: Translation symmetry is deﬁned by the group R and action ↵ ( ✓ ) = ✓ + ↵ 1 A where 1 A is the indicator vector for some subset A of the parameters { ✓ 1 , . . , ✓ m } . The loss function

Original: ✓ + ↵ possesses translation symmetry if[Equation] L 
Revised: L thus possesses translation symmetry if L ( ✓ )

Original: [Equation]
Revised: =

Original:  ✓ + ↵ 1  ) for all ↵ 2 R .
Revised: L ( ✓ + ↵ 1 A ) for all ↵ 2 R .

Original:  a symmetry 
Revised: Such a symmetry in turn implies the loss gradient @ ✓ L = g is orthogonal to the indicator vector @ ↵ ↵ = 1 A ,

Original: [Equation]
Revised: [Equation]

Original: Note the vector valued function R ( ✓ )  = @ 
Revised: ) and that the Hessian matrix H = @ 2 ✓ L has the indicator vector in its kernel,

Original: L @✓ is also translation invariant. The implied orthogonality condition on the gradient of R then implies a constraint on the Hessian H H 1 ✓ = 0 , (2) where 1 ✓ is the indicator vector for the subset of parameters with translation symmetry.
Revised: [Equation]

Original: Softmax function.
Revised: Softmax function.

Original: Any network using the softmax function gives rise to translation symmetry for the parameters immediately preceding the function.
Revised: Any network using the softmax function gives rise to translation symmetry for the parameters immediately preceding the function.

Original: Let z
Revised: Let z

Original: =
Revised: =

Original: Wx + b be the input to the softmax function such that  ( z ) i =
Revised: Wx + b be the input to the softmax function such that  ( z ) i =

Original: e zi
Revised: e zi

Original: P j e zj .
Revised: P j e zj .

Original: Notice that shifting any column of the weight matrix w i 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.
Revised: Notice that shifting any column of the weight matrix w i 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.

Original: Thus, the loss function is invariant w.r.t.
Revised: Thus, the loss function is invariant w.r.t.

Original: this translation, yielding  gradient constraints h @
Revised: this translation, yielding the gradient constraints h @

Original: L @ w i , 1 i = h @
Revised: L @ w i , 1 i = h @

Original: L @ b , 1 i = 0 .
Revised: L @ b , 1 i = 0 , visualized in Fig.

Original: S CALE
Revised: 2 for the toy model where w i 2 R 2 .

Original: Scale  symmetry is deﬁned by the group GL 1{++
+} ( R ) and action ✓ 7!
Revised: Scale symmetry. Scale symmetry is deﬁned by the group GL {++
+} ( R ) and action ✓ 7!

Original: ↵✓ .
Revised: ↵ A  ✓ where ↵

Revised: A = ↵ 1 A + 1 A c .

Original: The loss function possesses scale symmetry if L ( ✓ )
Revised: The loss function possesses scale symmetry if L ( ✓ )

Original: =
Revised: =

Original: L ( ↵✓ ) for all ↵ 2 GL 1 ( R ) .
Revised: L ( ↵ A  ✓ ) for all ↵ 2 GL +1 ( R ) .

Original: This symmetry immediately implies the loss gradient is everywhere perpendicular to the parameter vector itself:
Revised: This symmetry immediately implies the loss gradient is everywhere perpendicular to the parameter vector itself @ ↵ ↵ = ✓  1 A

Revised: = ✓ A , h g, ✓ A i = 0 ,

Original: [Equation]
Revised: [Equation]

Original: Note the vector function R ( ✓ ) = @  L @✓  ✓ also possesses scale symmetry.
Revised: and relates to the Hessian matrix, where @ ✓ L @ ↵ @ ✓ ↵ = g diag [Equation] [Equation]

Original: The implied orthogonality condition on the gradient of R then implies a constraint on the loss gradient and Hessian
Revised: ) = [Equation] A ,

Original: [Equation]
Revised: [Equation]

Original: Batch normalization.
Revised: Batch normalization.

Original: Batch normalization leads to scale invariance during training.
Revised: Batch normalization leads to scale invariance during training.

Original: Let
Revised: Let

Original: z
Revised: z

Original: =
Revised: =

Original: w | x + b be the input to a neuron with batch normalization such that BN ( z ) =
Revised: w | x + b be the input to a neuron with batch normalization such that BN ( z ) =

Original: z  E [ z ] p Var ( z ) where E [ z ] is the sample mean
Revised: z  E [ z ] p Var ( z ) where E [ z ] is the sample mean

Original: and Var ( z ) is the sample variance given a batch of data.
Revised: and Var ( z ) is the sample variance given a batch of data.

Original: 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.
Revised: 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.

Original: Thus, these parameters observe scale symmetry in the loss and their gradients satisfy
Revised: Thus, these parameters observe scale symmetry in the loss and their gradients satisfy

Original: [Equation]
Revised: [Equation]

Original:  , as has been previously noted by Ioffe & Szegedy (2015); Van Laarhoven (2017).
Revised: 0 , as has been previously noted by Ioffe & Szegedy (2015); Van Laarhoven (2017), and visualized in Fig.

Original: 3. I NVERSION
Revised: 2 for the toy model where w, b 2 R .

Original: Inversion symmetry is deﬁned by the group GL 1{++
+} ( R ) and action( ✓ 1 , ✓2 ) 7!
Revised: Rescale symmetry. Rescale symmetry is deﬁned by the group GL {++
+} ( R ) and action ✓ 7!

Original: ( ↵✓ 1 , ↵  1  ✓  2 ) .
Revised: ↵ A 1  ↵  1 A 2  ✓ where A 1 and A 2 are two disjoint sets of parameters.

Original: Thus the loss function possesses scaling symmetry if L ( ✓1 , ✓ 2 ) =
Revised: The loss function possesses rescale symmetry if L ( ✓ ) =

Original: L ( ↵✓ 1 , ↵  1 ✓ 2  ) for all ↵ 2 GL 1 ( R ) .
Revised: L ( ↵ A 1  ↵  1 A 2  ✓ ) for all ↵ 2 GL +1 ( R ) .

Original: The gradient of the loss with respect to the parameters satisﬁes the property,
Revised: This symmetry immediately implies the loss gradient is everywhere perpendicular to the sign inverted parameter vector @ ↵

Original: [Equation]
Revised: [Equation]

Original: Note the function R ( ✓ ) = @ L @✓  ( ✓ 1 + ✓ 2 ) also possess inversion symmetry.
Revised: ✓ [Equation] [Equation]

Original: The implied orthogonality conditions on the gradients of R imply the Hessian constraint
Revised: and relates to the Hessian matrix, where @ ✓

Original: [Equation]
Revised: [Equation]

Original: Homogeneous activation.
Revised: Homogeneous activation.

Original: For networks with continuous, homogeneous activation functions ( x ) =
Revised: For networks with continuous, homogeneous activation functions

Original:  0 ( x ) x
Revised: [Equation]

Original:  (e.g. ReLU, Leaky ReLU, linear), this symmetry emerges at every hidden neuron by considering all incoming and outgoing parameters to the neuron.
Revised: ( 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.

Original: For example, consider a hidden neuron with ReLU activation and let w 2 ReLU ( w | 1 x + b ) be the computational path through this neuron.
Revised: For example, consider a hidden neuron with ReLU activation  ( z ) = max { 0 , z } , such that w 2  ( w | 1 x + b ) is the computational path through this neuron.

Original: Scaling w 1 and b by a real constant and w 2 by its inverse has no effect on the computational path as the constants can be passed through the ReLU activation canceling their effects.
Revised: Scaling w 1 and b by a real constant and w 2 by its inverse has no effect on the computational path as the constants can be passed through the ReLU activation canceling their effects.

Original: Thus, these parameters observe inversion symmetry and their gradients satisfy 
Revised: Thus, these parameters observe rescale symmetry and their gradients satisfy D [Equation] [Equation]

Revised: , w

Original: [Equation]
Revised: [Equation]

Original:  , as has been previously noted by Du et al.
Revised: E [Equation] [Equation] E = 0 , as has been previously noted by Du et al.

Revised: (2018a)

Original: ; Liang et al.
Revised: ; Liang et al.

Original: ; Tanaka et al.
Revised: ; Tanaka et al.

Revised: (2020), and visualized in Fig.

Revised: 2 for the toy model where w 1 , w 2 2 R and b = 0 .

Original: General inversion symmetry.
Revised: General rescale symmetry.

Original: A more general inversion symmetries can be deﬁned by the group 
Revised: A more general rescale symmetries can be deﬁned by the group GL

Original: GL 1{++
+} ( R ) and action( ✓1 , ✓ 2 ) 7!
Revised: {++
+} ( R ) and action ✓ 7!

Original: ( ↵ k 1 ✓ 1 , ↵  k 2 ✓ 2 ) , which occur in networks with quadratic activation functions.
Revised: ↵ k 1 A 1  ↵  k 2 A 2  ✓ , which occur in networks with quadratic activation functions.

Original: A stronger form of invertible symmetry occurs for linear networks under the action of the group GL  k ( R ) of k ⇥ k invertible matrices deﬁned by ( ✓ 1 , ✓ 2 ) 7! ( G ✓ 1 , ✓ 2 G  1 ) , as noticed previously by Arora et al.
Revised: A stronger form of rescale symmetry occurs for linear networks under the action of the group GL + k ( R ) of k ⇥ k invertible matrices, as noticed previously by Arora et al.

Original: (2018a);
Revised: (2018a);

Original: Du et al.
Revised: Du et al.

Revised: (2018a).

Original:   C ONSERVED Q UANTITIES C ONSTRAIN THE L EARNING D YNAMICS
Revised:  S YMMETRY LEADS TO C ONSERVATION L AWS U NDER G RADIENT F LOW

Original: We now explore how geometric constraints on gradients and Hessians arising as a consequence of symmetry,derived above, impact network learning dynamics.
Revised: We now explore how geometric constraints on gradients and Hessians arising as a consequence of symmetry, impact network learning dynamics.

Original: In this work we focus on stochastic gradient descent (SGD), the workhorse of modern deep learning optimization.
Revised: In this work we focus on stochastic gradient descent (SGD), the workhorse of modern deep learning optimization.

Original: We will consider a model parameterized by ✓ , a training dataset { x 1 , . . .
Revised: We will consider a model parameterized by ✓ , a training dataset { x 1 , . . .

Original: , x N } of size N ,  a training loss L ( ✓ ) =  N
Revised: , x N } of size N , and a training loss L ( ✓ ) = 1 N

Original: P
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: , x i ) with corresponding gradient g ( ✓ ) = @
Revised: [Figure] corresponding gradient g ( ✓ ) = @

Original: L@✓ , and an L 2 regularization constant   0 .
Revised: L

Original: [Figure]
Revised: [Figure]

Original: The gradient descent update with learning rate ⌘ is ✓ 
Revised: The gradient descent update with learning rate ⌘ is ✓ ( n +1)

Original: [Equation] ) , which is a forward Euler discretization  with step size ⌘ of the ordinary differential equation (ODE)
Revised: = ✓ ( n )  ⌘ g ( ✓ ( n ) ) , which is a forward Euler discretization 1 with step size ⌘ of the ordinary differential equation (ODE)

Original: [Equation]
Revised: [Equation]

Original: In the limit as ⌘ ! 0 , gradient descent exactly matches the dynamics of this ODE, which is commonly referred to as gradient ﬂow.
Revised: In the limit as ⌘ ! 0 , gradient descent exactly matches the dynamics of this ODE, which is commonly referred to as gradient ﬂow Kushner & Yin (2003).

Original: Equipped with a continuous model for the learning dynamics, we now ask how do the dynamics interact with the geometric properties introduced by symmetries?
Revised: Equipped with a continuous model for the learning dynamics, we now ask how do the dynamics interact with the geometric properties introduced by symmetries?

Original: Symmetry leads to conservation.
Revised: Symmetry leads to conservation.

Revised: Strikingly similar to Noether’s theorem, which describes a fundamental relationship between symmetry and conservation for physical systems governed by Lagrangian dynamics, every symmetry of a network architecture has a corresponding conserved quantity through training under gradient ﬂow .

Original: Consider some subset of the parameters ✓ ✓ ⇥ that respects either a translation, scale, or inversion symmetry.
Revised: Consider some subset of the parameters A that respects either a translation, scale, or rescale symmetry.

Original: As shown in section 3, thepartial gradient of the loss with respect to these parameters satisﬁes an equality constraint.
Revised: As shown in section 3, the gradient of the loss g ( ✓ ) is always perpendicular to the vector ﬁeld that generates the symmetry @

Original: For example, in the case of translation symmetry h g ( ✓ ) , 1 i = 0 .
Revised: ↵ .

Original: Projectingboth sides of the gradient ﬂow learning dynamics in equation (7) onto theconstant generator vector ﬁeld1 of translations then yields  [Equation]  [Equation]
Revised: Projecting the gradient ﬂow learning dynamics in equation (27) onto the generator vector ﬁeld yields a differential equation h d ✓ dt , @

Original: h ✓ ,

Original: [Equation]

Original:  i = 0 .
Revised: ↵ i = 0 .

Original: Solving the initial value problem given by this ODE results in the planar relationship [Equation] Analogously, we can show that for scale symmetry we obtain the spherical relationship
Revised: Integrating this equation through time results in the conservation laws,

Original: [Equation]
Revised: [Equation]

Original: and for inversion symmetry we obtain the hyperbolic relationship [Equation] Each of these equations deﬁne a conserved constant of motion through training when  = 0 and an exponentially decaying quantity when  > 0 .
Revised: Each of these equations deﬁne a conserved constant of learning through training. For parameters with translation symmetry, their sum is conserved, effectively constraining their dynamics to a hyperplane.

Original: This result is strikingly similar toNoether’s theorem Noether, a fundamental relationship between symmetry and conservation for physical systems governed by Lagrangian dynamics. L 2 regularization removes degrees of freedom.
Revised: For parameters with scale symmetry, their euclidean norm is conserved, effectively constraining their dynamics to a sphere. For parameters with rescale symmetry, their difference in squared euclidean norm is conserved, effectively constraining their dynamics to a hyperbola.

Original: Despite its widespread adoption as a form of explicit regularization, the actual impact of L 2 regularization on the training dynamics of neural networks is not well understood. Some of the proposed regularizing mechanisms are increasing the effective learning rate, penalizing the input-output Jacobian norm, and reducing the damping coefﬁcient for second-order optimization methods Zhang et al. However, a much clearer picture of L 2 regularization’s impact can be found in the conservation equations (8), (9), and (10).
Revised: In Fig.

Original: When  > 0 the  conserved quantities exponentially decay through training at a rate proportional to the regularization constant. Geometrically, this means the degrees of freedom associated with the network symmetry, the intercepts of the hyperplane, the radius of the sphere, the axes of the hyperbola, diminish exponentially to zero.
Revised: 3 we visualize the level sets of these conserved quantities for the toy models discussed in Fig.

Original:  A  C ONTINUOUS M ODEL FOR S TOCHASTIC G RADIENT D ESCENT
Revised:  A R EALISTIC C ONTINUOUS M ODEL FOR S TOCHASTIC G RADIENT D ESCENT

Original: In section 4 we combined the geometric constraints introduced by symmetries with gradient ﬂow to make exact predictions for simple combinations of parameters during training.
Revised: In section 4 we combined the geometric constraints introduced by symmetries with gradient ﬂow to derive conservation laws for simple combinations of parameters during training.

Original: However, there is something strange about our predictions. For example, in the case of parameters with scale symmetry, such as the weight parameters before a batch normalization layer, we would predict that SGD training with L 2 regularization should result in the norm of the weights exponentially decaying to zero, something we know is not true in practice.
Revised: However, empirically we know these laws are broken, as demonstrated in Fig.

Original: So what could be going wrong?
Revised: What causes this discrepancy?

Original: Gradient ﬂow is too simple of a continuous model for  SGD training.
Revised: Gradient ﬂow is too simple of a continuous model for realistic SGD training.

Original: It fails to incorporate the effect of  stochasticity introduced by random batches and discretization introduced by a ﬁnite learning rate.
Revised: It fails to incorporate the effect of weight decay introduced by explicit regularization, momentum introduced by commonly used hyperparameters, stochasticity introduced by random batches, and discretization introduced by a ﬁnite learning rate.

Original: Here, we construct a more realistic continuous model for stochastic gradient descent. 
Revised: Here, we construct a more realistic continuous model for stochastic gradient descent. Modeling weight decay.

Original: 5. M ODELING S TOCHASTICITY
Revised: Explicit regularization through the addition of an L 2 penalty on the parameters, with regularization constant  , is very common practice when training modern deep learning models. This is generally implemented not by modifying the training loss, but rather directly modifying the optimizers update equation. For stochastic gradient descent, the result leads to the updated continuous model [Equation]

Original: Stochastic gradients ˆ g B (  ✓  ) arise when we consider a minibatch B of size S drawn uniformly from the indices { 1 , . . .
Revised: Modeling momentum. Momentum is a common extension to SGD that uses an exponentially moving average of gradients to update parameters rather than a single gradient evaluation Rumelhart et al. The method introduces two additional hyperparameters, a damping coefﬁcient ↵ and a momentum coefﬁcient  , and applies the two step update equation, ✓ ( n +1) = ✓ ( n )  ⌘ v ( n +1) where v ( n +1) =

Original: , N } forming the unbiased gradient estimate ˆ gB ( ✓  ) 
Revised:  v ( n ) + (1  ↵ ) g ( ✓ ( n ) ) . When ↵ =  = 0 , we regain classic gradient descent. In general, ↵ effectively reduces the learning rate and  controls how past gradients are used in future updates resulting in a form of “inertia” accelerating and smoothing the descent trajectory.

Original: = 1 S
Revised: Rearranging the two-step update equation, we ﬁnd that gradient descent with momentum is a ﬁrst-order discretization with step size ⌘ (1  ↵ ) of the ODE,

Original: [Equation]
Revised: [Equation]

Original: i )  . 
Revised: A more detailed derivation for this ODE can be found in appendix B. Modeling stochasticity. Stochastic gradients ˆ g B ( ✓ ) arise when we consider a batch B of size S drawn uniformly from the indices { 1 , . . . , N } forming the unbiased gradient estimate ˆ g B ( ✓ ) =

Original:  When the batch size is much smaller than the size of the dataset, S ⌧ N , then we can model the batch gradient as an average of S i.i.d.
Revised: S P i 2 B r ` ( ✓ , x i ) . When the batch size is much smaller than the size of the dataset, S ⌧ N , then we can model the batch gradient as an average of S i.i.d.

Original: samples from a noisy version of the true gradient g ( ✓ ) .
Revised: samples from a noisy version of the true gradient g ( ✓ ) .

Original: Using the central limit theorem, we assume ˆ g B ( ✓ )  g ( ✓ ) is a Gaussian random variable with mean µ = 0 and covariance matrix ⌃ =
Revised: Using the central limit theorem, we assume ˆ g B ( ✓ )  g ( ✓ ) is a Gaussian random variable with mean µ = 0 and covariance matrix ⌃ =

Original: 1 S G ( ✓ ) G ( ✓ ) | .
Revised: 1 S G ( ✓ ) G ( ✓ ) | .

Original: Under this assumption, the stochastic gradient update can be written as ✓ ( n +1)=
Revised: Under this assumption, the stochastic gradient update can be written as ✓ ( n +1)

Original: (1  ⌘ ) ✓ ( n )  ⌘ g ( ✓ ( n ) )
Revised: = ✓ ( n )  ⌘ g ( ✓ ( n ) )

Original: + ⌘ p S G ( ✓ ) ⇠ , where ⇠ is a standard normal random variable.
Revised: + ⌘ p S G ( ✓ ) ⇠ , where ⇠ is a standard normal random variable.

Original: This update is an Euler-Maruyama discretization with step size ⌘ of the stochastic differential equation 
Revised: This update is an Euler-Maruyama discretization with step size ⌘ of the stochastic differential equation [Equation]

Original: [Equation]
Revised: [Equation]

Original: where W t is a standard Wiener process.
Revised: where W t is a standard Wiener process.

Original: The kernel of the covariance matrix.

Original: Many stochastic differential equation approximations for SGD similar to equation (11) have been derived in previous works.
Revised: Many stochastic differential equation approximations for SGD similar to equation (15) have been derived in previous works.

Original: In order to simplify the analysis, many of these works have then made additional assumptions on the covariance matrix ⌃ =
Revised: In order to simplify the analysis, many of these works have then made additional assumptions on the covariance matrix ⌃ =

Original: G ( ✓ ) G ( ✓ ) | , such as ⌃ =
Revised: G ( ✓ ) G ( ✓ ) | , such as ⌃ =

Original: H where H is the Hessian matrix Jastrz˛ebski et al.
Revised: H where H is the Hessian matrix Jastrz˛ebski et al.

Original: (2017), ⌃ =
Revised: (2017), ⌃ =

Original: C where C is some constant matrix Mandt et al.
Revised: C where C is some constant matrix Mandt et al.

Original: (2015), and ⌃
Revised: (2015), and ⌃

Original: =
Revised: =

Original: I where I is the identity matrix Chaudhari & Soatto (2018).
Revised: I where I is the identity matrix Chaudhari & Soatto (2018).

Original: However, without any additional assumptions, the differential symmetries intrinsic to neural network architectures add fundamental constraints on ⌃ .
Revised: However, without any additional assumptions, the differential symmetries intrinsic to neural network architectures add fundamental constraints on ⌃ .

Original: Again consider some subset of the parameters ✓ ✓ ⇥ that respects either a translation, scale, or inversion symmetry. As we showed in section 3 the partial gradients of the loss with respect to these parameters have special geometric properties, independent of the dataset. For example, in the case oftranslation symmetry, both the mini-batch and the full-batch gradients are orthogonal to the all-ones vector, h ˆ g B ( ✓ ) , 1 i =
Revised: As we showed in section 3, the gradient of the loss, regardless of the batch, is orthogonal to the generator vector ﬁeld @

Original: h g ( ✓ ) , 1 i = 0 . This implies the stochastic noise must also observe the same properties, h 1 p S G ( ✓ ) ⇠ , 1 i = 0 .
Revised: ↵ associated with a symmetry. This implies the stochastic noise must also observe the same property, h 1 p S G ( ✓ ) ⇠ , @ ↵ i = 0 .

Original: In order for this relationship to hold for arbitrary noise  ⇠ , this implies G ( ✓ ) T 1 = 0 .
Revised: In order for this relationship to hold for arbitrary noise 2 ⇠ , then G ( ✓ ) | @ ↵ = 0 .

Original: Analogously, we have for scale symmetry G ( ✓ ) | ✓ = 0 and for the inversion symmetry G ( ✓ 1 ) | ✓ 1  G ( ✓ 2 ) | ✓ 2 = 0 .

Original: In other words, the differential symmetry inherent in neural network architectures projects the noise introduced by stochastic gradients to low rank subspaces.
Revised: In other words, the differential symmetry inherent in neural network architectures projects the noise introduced by stochastic gradients to low rank subspaces.

Original: 5. M ODELING D ISCRETIZATION

Original: While equation (11) is certainly a more realistic model for SGD than gradient ﬂow, its major limitation is that it still assumes an inﬁnitesimal step size when in fact the learning rate is ﬁnite. Gradient descent always moves in the direction of steepest descent on a loss function, however, due to the ﬁnite nature of the learning rate, it fails to remain on the continuous steepest descent path. The divergence between the discrete and continuous trajectories fundamentally depends on the learning rate and the curvature of the loss. It is thus natural to assume that a realistic continuous model for SGD must incorporate both these terms in a non-trivial way.
Revised: [Figure]

Original: Modiﬁed gradient ﬂow.
Revised: Modeling discretization.

Original: The effect of discretization when modeling continuous dynamics is a well studied problem in the numerical analysis of partial differential equations.
Revised: The effect of discretization when modeling continuous dynamics is a well studied problem in the numerical analysis of partial differential equations.

Original: One tool commonly used in this setting, is modiﬁed equation analysis Warming & Hyett (1974), which asks how can we better model discrete steps with a continuous differential equation?
Revised: One tool commonly used in this setting, is modiﬁed equation analysis Warming & Hyett (1974), which determines how to better model discrete steps with a continuous differential equation by introducing higher order spatial or temporal derivatives. We present two methods based on modiﬁed equation analysis, which modify gradient ﬂow to account for the effect of discretization. Modiﬁed loss.

Revised: Gradient descent always moves in the direction of steepest descent on a loss function L at each step, however, due to the ﬁnite nature of the learning rate, it fails to remain on the continuous steepest descent path given by gradient ﬂow.

Original: Li et al.
Revised: Li et al.

Original: ; Feng et al.
Revised: ; Feng et al.

Original: (2019) and
Revised: (2019) and

Original: more recently Barrett & Dherin (2020), applied this approach to the dynamics of stochastic gradient descent, demonstrating that the gradient descent trajectory closely follows the steepest descent path of a modiﬁed loss function e 
Revised: most recently Barrett & Dherin (2020), demonstrate that the gradient descent trajectory closely follows the steepest descent path of a modiﬁed loss function e L . The diver- gence between these trajectories fundamentally depends on the learning rate ⌘ and the curvature

Revised: H .

Revised: As derived in Barrett & Dherin (2020), and summarized in appendix C, this divergence is given by the gradient correction  12 Hg , which is the gradient of the norm  ⌘ 4 | r L| .

Original: L , rather than the original loss  L .
Revised: Thus, the modiﬁed loss is e L = L +

Original: Using backward error analysis they analytically calculate this modiﬁed loss as e
Revised: [Equation] 4 | r L| 2 and the modiﬁed gradient ﬂow ODE is

Original: [Equation]
Revised: [Equation]

Original: Here, we apply 1 this method  to modify the forcing term F ( ✓ ) =
Revised: See Fig. 4 for an illustrative example of this method applied to a quadratic loss in R 2 . Modiﬁed ﬂow.

Original:  g ( ✓  )  ✓ in equation (11), as
Revised: Rather than modifying gradient ﬂow with higher order “spatial” derivatives of the loss function, here we introduce higher order temporal derivatives. We start by assuming the existence of a continuous trajectory ✓ ( t ) that weaves through the discrete steps taken by gradient descent and then identify the differential equation that generates the trajectory. Rearranging the update equation for gradient descent, ✓ t +1 = ✓ t  ⌘ g ( ✓ t ) , and assuming ✓ ( t ) = ✓ t and ✓ ( t + ⌘ ) =

Original: [Equation]{+✓ t +1 , gives the equality  g ( ✓ t ) = ✓ ( t + ⌘ )  ✓ ( t ) ⌘ , which Taylor expanding the right side results in the differential equation  g ( ✓ t ) = d ✓ dt + ⌘ 2 d 2 ✓ dt 2 + O ( ⌘ 2 ) . Notice that in the limit as ⌘ ! 
+}
Revised: {+✓ t +1 , gives the equality  g ( ✓ t ) = ✓ ( t + ⌘ )  ✓ ( t ) ⌘ , which Taylor expanding the right side results in the differential equation  g ( ✓ t ) = d ✓ dt + ⌘ 2 d 2 ✓ dt 2 + O ( ⌘ 2 ) . Notice that in the limit as ⌘ ! 
+}

Original: Using this forcing term results in the ﬁrst-order, stochastic differential equation we term modiﬁed  gradient ﬂow ,
Revised: we regain gradient ﬂow. For small ⌘ , we obtain a modiﬁed version of gradient ﬂow with an additional second-order term,

Original: [Equation]
Revised: [Equation]

Original: In appendix B we considers an illustrative example providing intuition for modiﬁed gradient ﬂow.
Revised: This approach to modifying ﬁrst-order differential equation with higher order temporal derivatives was applied by Kovachki & Stuart (2019) to construct a more realistic continuous model for momentum, as illustrated in Fig.

Original:  A NALYZING T RAINING D YNAMICS THROUGH THE L ENS OF S YMMETRY 
Revised:  C OMBINING S YMMETRY AND M ODIFIED G RADIENT F LOW TO D ERIVE E XACT L EARNING D YNAMICS

Original: As shown in section 4, each symmetry implies a conserved quantity under gradient ﬂow with  = 0 .
Revised: As shown in section 4, each symmetry results in a conserved quantity under gradient ﬂow.

Original: We now employ modiﬁed gradient ﬂow to examine how weight decay, gradient stochasticity and ﬁnite learning rates all interact to break the conservation law of gradient ﬂow, and we predict the dynamics of these previously conserved quantities.
Revised: We now study how weight decay, momentum, stochastic gradients, and ﬁnite learning rates all interact to break these conservation laws. Remarkably, even when using a more realistic continuous model for stochastic gradient descent, as discussed in section 5, we can derive exact learning dynamics for the previously conserved quantities.

Original: To do so, we  project the learning dynamicsgiven by modiﬁed gradient ﬂow equation (13) onto  generator vector ﬁelds  associated with each of thesymmetries while harnessing symmetry induced geometric constraints on the gradients and Hessians 2 .
Revised: To do this we (i) consider a realistic continuous model for SGD, (ii) project these learning dynamics onto the generator vector ﬁelds @ ↵ associated with each symmetry, (iii) harness the geometric constraints from section 3 to derive simpliﬁed ODEs, and (iv) solve these ODEs to obtain exact dynamics for the previously conserved quantities.

Original: In the case of parameters with translation symmetry, the effect of discretization essentially leaves the dynamics for the constant of motion unchanged. Combining the geometric properties of graident ( h g, 1 i = 0 ), Hessian ( H 1 = 0 ), andstochasticity ( G ( ✓ ) T 1 = 0 ) introduced by translation symmetry with modiﬁed gradient ﬂow gives the following differential equation, [Equation]
Revised: We ﬁrst consider the continuous model of SGD without momentum incorporating weight decay, stochasticity, and modiﬁed loss.

Original: The solution of this differential equation is simply [Equation] ✓
Revised: In this setting, the exact dynamics, as fully derived in appendix D, for the parameter combinations tied to the symmetries are,

Original: [Equation]Notice that this is equivalent to equation (8), which we derived using gradient ﬂow. Thus, even when considering the effect of a ﬁnite step size, the dynamics for the parameters with translation symmetry are unchanged. This is due to the absence of curvature in the constraining hyperplane geometry.
Revised: [Equation]

Original: In the case of parameters with scale symmetry, the effect of discretization does distort the original geometry.
Revised: Notice how these equations are equivalent to the conservation laws derived in section 4 when ⌘ =  = 0 .

Original: A ﬁnite learning rate leads to a centrifugal force that monotonically increases the previously conserved quantity ( | ✓ | 2 ), while weight decay acts as a centripetal force decreasing the quantity. Competition between these two forces leads to interesting dynamics with surprising consequences
Revised: Remarkably, even in typical hyperparameter settings (weight decay, stochastic batches, ﬁnite learning rates), these solutions match nearly perfectly with empirical results from modern neural networks (VGG-16) trained on real-world datasets (Tiny ImageNet), as shown in Fig. We will now discuss each equation individually. [Equation]

Original: [Table]
Revised: [Table]

Original: /J1IzymFa8zIFKaFjMoS+QUFiUH5erD/Fp8YZ4CiR5giNC/f3RE5ipSZxaDpjokdqsTYz/6v1Mx1d+zkTaaZB0PlDUcaxTvAsCzxgEqjmEwOESmZ2xXREJKHaJFYzIbiLX16GzkXDNXx/WW/elHFU0TE6QWfIRVeoie5QC7URRTl6Rq/ozXqyXqx362PeWrHKmUP0R9bnDwRFlN0=</latexit>
Revised: +QUFiUH5erD/Fp8YZ4CiR5giNC/f3RE5ipSZxaDpjokdqsTYz/6v1Mx1d+zkTaaZB0PlDUcaxTvAsCzxgEqjmEwOESmZ2xXREJKHaJFYzIbiLX16GzkXDNXx/WW/elHFU0TE6QWfIRVeoie5QC7URRTl6Rq/ozXqyXqx362PeWrHKmUP0R9bnDwRFlN0=</latexit> T r a n s l a t i o n

Original: [Figure]
Revised: S ca l e

Original: for the stationary state of these parameters. Combining the geometric constraints on the gradient ( hg, ✓ i = 0 ), Hessian ( H ✓ =  g ), and stochasticity ( G ( ✓ ) T ✓ = 0 ) introduced by scale symmetry with modiﬁed gradient ﬂow gives the following differential equation,
Revised: R e s ca l e Translation dynamics. For parameters with translation symmetry, equation 18 implies that the sum of these parameters ( h ✓ A ( t ) , 1 i ) decays exponentially to zero at a rate proportional to the weight decay.

Original: [Equation]
Revised: Equation 18 does not directly depend on the learning rate ⌘ nor any information of the dataset or task.

Original: By solving this ﬁrst order inhomogeneous differential equation, we get [Equation] [Equation]
Revised: This is due to the lack of curvature in the gradient ﬁeld for these parameters (as shown in Fig. 2).

Original: Note that the zeroth order term in ⌘ is equivalent to gradient ﬂow equation (9), but here it is corrected by an additional non-negative integral term that is ﬁrst order in ⌘ . This integral term captures the spiraling behavior of gradient descent due to the curvature of the constraining spherical geometry and a ﬁnite learning rate. We can make this exact bytaking the time derivative of the equation (15) giving the temporal dynamics for the constant of motion,
Revised: This implies that at initialization we can deterministically predict the trajectory for the parameter sum as simple exponential functions with a rate deﬁned by the weight decay. The ﬁrst row in Fig.demonstrates this qualitatively, as all trajectories are smooth exponential functions that converge faster for increasing levels of weight decay.

Original: [Equation]
Revised: Scale dynamics.

Original: It is clear from equation (16) that[-even for  = 
-] the time derivative isnon-negative implying the previosly conserved quantity is monotonically increasing.
Revised: For parameters with scale symmetry, equation 19 implies that[-even for  = 
-] the norm for these parameters ( | ✓ A | 2 ) is the sum of an expo- nentially decaying memory of the norm at initialization and an exponentially weighted integral of gradient norms accumulated through training.

Original: More generally for    0 there is a competition between the centripetal effect of weight decay and the centrifugal effect of discretization.
Revised: Compared to the translation dynamics, the scale dynamics do depend on the data through the gradient norms accumulated throughout training. Without weight decay  = 0 , the ﬁrst term stays constant and the second term grows monotonically. With weight decay  > 0 , the ﬁrst term decays monotonically to zero, while the second term can decay or grow, but always stays positive. The second row in Fig.

Original: At the stationary state of the dynamics, these forces will balance leading the dynamics of these parameters to be constrained to the surface of a high-dimensional sphere.
Revised: 5 demonstrates these qualitative relationships. Without weight decay the norms increases monotonically as predicted and with weight decay the dynamics are non-monotonic and present more complex behavior. To better understand the forces driving these complex dynamics, we can examine the time derivative of equation 19,

Original: Constant angular speed at the stationary state.
Revised: [Equation] From this equation we see that there is a competition between a centripetal effect due to weight decay (  2  | ✓ A ( t ) | 2 ) and a centrifugal effect due to discretization ( ⌘ | g A | 2 ). The centripetal effect due to weight decay is a direct consequence of its regularizing inﬂuence, pulling the parameters towards the origin. The centrifugal effect due to discretization originates from the spherical geometry of the gradient ﬁeld in parameter space – because scale symmetry implies the gradient is always orthogonal to the parameter itself, each discrete update with a ﬁnite learning rate effectively pushes the parameters away from the origin.

Original: When the dynamics of | ✓ ( t ) | 2 have converged to some stationary state, then d dt | ✓ ( t ) | 2 = 0 , which rearranging equation (16) gives the condition ![Equation]
Revised: At the stationary state of the dynamics, these forces will balance leading the dynamics of these parameters to be constrained to the surface of a high-dimensional sphere. In particular, at stationarity, then ddt | ✓ ( t ) | 2 = 0 , which rearranging equation (21) gives the condition !

Original:  ) ⌘
Revised: ( t ) ⌘

Original: [Equation] 
Revised: [Equation] [Equation]

Original:  ⌘ .
Revised: q 2 ⌘ .

Original: Consistent with the results of Wan et al.
Revised: Consistent with the results of Wan et al.

Original: (2020), this implies that at stationarity the angular speed !
Revised: (2020), this implies that at stationarity the angular speed !

Original: ( t ) of the weights is constant and governed only by the learning rate ⌘ and weight decay constant  . 
Revised: ( t ) of the weights is constant and governed only by the learning rate ⌘ and weight decay constant  . Rescale dynamics.

Original: In the case of parameters with inversion symmetry,  the effect ofdiscretization also distorts the original geometry. However, unlike in  the sphere, the force originating from discretization can both increase or decrease the previosly conserved quantity ( | ✓ 1 ( t ) | 2  | ✓ 2 ( t ) | 2 ).
Revised: For parameters with rescale symmetry, equation 20 is the sum of an exponentially decaying memory of the difference in norms at initialization and an exponentially weighted integral of difference in gradient norms accumulated through training. Similar to the scale dynamics, the rescale dynamics do depend on the data through the gradient norms, however unlike the scale dynamics we have no guarantee that the integral term is always positive.

Original: Combining the geometric properties of gradient ( h g 1 , ✓ 1 i  h g 2 , ✓ 2 i = 0 ), Hessian ( H ✓ 1  H ✓ 2 + g ✓ 1  g ✓ 2 = 0 ), and stochasticity ( G ( ✓ 1 ) | ✓ 1  G ( ✓ 2 ) | ✓ 2 = 0 ) introduced by inversion symmetry with modiﬁed gradient ﬂow gives the following differential equation,
Revised: This leads to quite sophisticated, complex dynamics, consistent with the third row in Fig. Despite the complexity, our theory, nevertheless, quantitatively matches the empirics. The only apparent pattern from the empirics is that for large enough weight decay, the regularization dominates any complexity introduced by the gradient norms and the difference in parameter norms decays exponentially to zero.

Original: [Equation]
Revised: [Equation]

Original: [Table][Figure] [Table]
Revised: [Table]

Original: [Figure]
Revised: [Equation]

Original: Ou
Revised: classiﬁer

Original: [Figure]
Revised: [Equation]

Original: /JHz

Original: [Figure]
Revised: [Table]

Original: QN5Q5Eu</latexit>
Revised: +RIyQ1xq7Var1xxq+5cZBW8HCqQq9Erf/n9hGUxV8gkNabruSkGE6pRMMmnJT8zPKVsRAe8a1HRmJtgMr95Ss6s0ydRou1TSObu74kJjY0Zx6HtjCkOzXJtZv5X62YYXQcTodIMuWKLRVEmCSZkFgDpC80ZyrEFyrSwtxI2pJoytDGVbAje8pdXoXVR9Sw/XFbqt3kcRTiBUzgHD66gDvfQgCYwSOEZXuHNyZwX5935WLQWnHzmGP7I+fwB0LWQNw==</latexit> Harmonic oscillation with momentum.

Original: [Figure]
Revised: We will now consider a continuous model of SGD with momentum.

Original: [Table] [Figure] This is the same differential equation as in equation (14), just with a different forcing term.
Revised: As discussed in appendix E, we consider the continuous model incorporating weight decay, momentum, stochasticity, and modiﬁed ﬂow.

Original: Thus, the solution for this differential equation is 
Revised: Under this model, the solutions we obtain take the form of driven harmonic oscillators where the driving force is given by the gradient norms, the friction is deﬁned by the momentum constant, the spring coefﬁcient is deﬁned by the regularization rate, and the mass is deﬁned by the the learning rate and momentum constant.

Original: [Equation]
Revised: For most standard hyperparameter choices, these solutions are in the overdamped setting and align well with the ﬁrst-order solutions for SGD without momentum up to a time rescaling, as shown in the left and middle panel of Fig.

Original: Again the leading order terms are equivalent to equation (10), which we derived using gradient ﬂow, with an additional integral term. However,unlike for scale symmetry, this integral terms can be both negative or positive.
Revised: However, for large values of beta we can push the solution into the underdamped regime where we would expect harmonic oscillation and indeed, we can empirically verify our predictions, even at scale for VGG-16 trained on Tiny ImageNet, as in right panel of Fig.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: Despite being the central guiding principle in the exploration of the physical world Anderson (1972);
Revised: Despite being the central guiding principle in the exploration of the physical world Anderson (1972);

Original: Gross (1996), symmetry has been underutilized in understanding the mechanics of neural networks.
Revised: Gross (1996), symmetry has been underutilized in understanding the mechanics of neural networks.

Original: In this paper, we developed a unifying theoretical framework that combines the power of symmetry and modiﬁed equation analysis to derive exact equations for  combinations of parameters during training, which we experimentally veriﬁed on large scale neural networks and datasets.As shown in ﬁgures 3, 4, and 5, our theoretical predictions align remarkably well with the empirics.
Revised: In this paper, we constructed a unifying theoretical framework harnessing the geometric properties of symmetry and realistic continuous equations for learning dynamics modeling weight decay, momentum, stochasticity, and discretization. We use this framework to derive exact dynamics for meaningful combinations of parameters, which we experimentally veriﬁed on large scale neural networks and datasets.

Original: While, these plots are aggregated at the layer-level, our predictions are actually on combinations of parameters at the neuron-level, as shown in appendix A. Remarkably, even for models with millions of parameters, we can make exact predictions on thousands of  parameter combinations through the course of training with large-scale datasets. This work providesa strong theoretical foundation for principled exploration of (i) architecture design guided by the principles of symmetry, (ii) optimization inspired by modiﬁed gradient ﬂow, and (iii) the synergistic interaction of implicit and explicit regularization.
Revised: For example, in the case of a VGG-16 model with batch normalization trained on Tiny-ImageNet (one of the model/dataset combinations we considered in section 6) there are 12 , 751 distinct parameter combinations whose dynamics we can analytically describe. Overall, this work provides the ﬁrst solid step towards the mechanics of learning in neural networks.