Original: S YMMETRY , C ONSERVATION L AWS , AND L EARNING D YNAMICS IN N EURAL N ETWORKS Anonymous authors Paper under double-blind review
Revised: N EURAL M ECHANICS : SYMMETRY AND BROKEN CONSERVATION LAWS IN DEEP LEARNING DYNAMICS

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: Understanding 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 Noether’s theorem in physics.
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 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 expressions 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 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?
Revised: This raises a foundational question: what, if anything, can we quantitatively understand about the learning dynamics of large-scale, non-linear neural network models driven by real-world datasets and optimized via stochastic gradient descent with a ﬁnite batch size, learning rate, and with or without momentum?

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

Original: (2013), inﬁnite width layers Jacot et al.
Revised: (2013), inﬁnite width layers Jacot et al.

Original: (2018), or single hidden layers Saad &
Revised: (2018), or single hidden layers Saad &

Original: Solla (1995).
Revised: Solla (1995).

Original: Many of these works have also ignored the complexity introduced by stochasticity and discretization by only focusing on the learning dynamics under gradient ﬂow.
Revised: Many of these works have also ignored the complexity introduced by stochasticity and discretization by only focusing on the learning dynamics under gradient ﬂow.

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

Original: 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).
Revised: 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 major assumption (see Fig. 1).

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

Original: 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.
Revised: 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 leverage continuous differentiable symmetries in the loss to unify and generalize geometric constraints on neural network gradients and Hessians (section 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 prove that each of these differentiable symmetries has an associated conservation law under the learning dynamics of gradient ﬂow (section 4).We construct
Revised: We prove that each of these differentiable symmetries has an associated conservation law under the learning dynamics of gradient ﬂow (section 4).

Original:  a more realistic continuous model for stochastic gradient descent by modeling weight decay, momentum, stochastic batches, and ﬁnite learning rates (section 5).
Revised: We construct a more realistic continuous model for stochastic gradient descent by modeling weight decay, momentum, stochastic batches and ﬁnite learning rates (section 5).

Original: We show that under this more realistic model the conservation laws of gradient ﬂow are broken, 
Revised: We show that under this more realistic model the conservation laws of gradient ﬂow are broken, yielding simple ODEs governing the dynamics for the previously conserved parameter combinations (section 6).

Original: yielding simple ODEs governing the dynamics for the previously conserved
Revised: We solve these ODEs to derive exact learning dynamics for the parameter combinations, which we validate empirically on VGG-16 trained on Tiny ImageNet with and without batch normalization (section 6).  R ELATED WORK

Original: parameter
Revised: [Figure]

Original: combinations (section 6). We solve these ODEs to derive exact learning dynamics for the parameter combinations, which we validate empirically
Revised: We plot the perparameter dynamics (left) and per-channel squared

Original: on VGG-16  trained on Tiny ImageNet withand without batch normalization (section
Revised: Euclidean norm dynamics (right) for the convolutional layers of a VGG-16 model (with batch normalization) trained on Tiny ImageNet with

Original: R ELATED WORK
Revised: SGD with learning rate η = 0 . 1 , weight decay λ

Original: [Figure]
Revised: = 10 − 4 , and batch size S = 256 .

Original:  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: While the parameter dynamics are noisy and chaotic, the neuron dynamics are smooth and patterned. 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.

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

Original: Saxe et al.
Revised: Saxe et al.

Original: (2013) considered linear neural networks with speciﬁc orthogonal initializations, deriving exact solutions for the learning dynamics under gradient ﬂow.
Revised: (2013; 2019) and Lampinen & Ganguli (2018) considered linear neural networks with speciﬁc orthogonal initializations, deriving exact solutions for the learning dynamics under gradient ﬂow.

Original: The theoretical tractability of linear networks has further enabled analyses on the properties of loss landscapes Kawaguchi (2016), convergence Arora et al.
Revised: The theoretical tractability of linear networks has further enabled analyses on the properties of loss landscapes Kawaguchi (2016), convergence Arora et al.

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

Original: Du & Hu (2019), and implicit acceleration by overparameterization Arora et al. (2018b).
Revised: Du & Hu (2019), and implicit acceleration by overparameterization Arora et al. (2018b).

Original: Saad & Solla (1995)  studieda single hidden layer architecture with non-linearities in a student-teacher setup, deriving a set of complex ODEs describing the learning dynamics.
Revised: Saad & Solla (1995) and Goldt et al. (2019) studied single hidden layer architectures with non-linearities in a studentteacher setup, deriving a set of complex ODEs describing the learning dynamics.

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

Original: (2018b); Mei et al.
Revised: (2018b); Mei et al.

Original: Jacot et al.
Revised: Jacot et al.

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

Original: 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: 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 (though see also Fort et al.

Original: Arora et al.
Revised: Arora et al.

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

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

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

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

Original: 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: We make this fundamental contribution by constructing a framework harnessing the geometry of the loss shaped by symmetry and realistic continuous equations of learning.

Original: Geometry of the loss.
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 simplify, unify, and generalize the literature through the lens 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: (2018a) activations, ﬁnding that gradient geometry plays an important role in constraining the learning 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 has 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: Most recently, Li et al.
Revised: Most recently, Li et al.

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

Original: Equations of learning.
Revised: 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: ; Li et al.
Revised: ; Li et al.

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

Original: ; Chaudhari & Soatto (2018);
Revised: ; Chaudhari & Soatto (2018);

Original: Jastrz˛ebski et al.
Revised: Jastrz˛ebski et al.

Original: ; 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

Original: [Figure] 
Revised: [Figure] Figure

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

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

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

Original: BN  [ ww 2 ] [ xx 2 ] |  where BN is the batch normalization function.
Revised: BN (cid:0) [ ww 2 ] [ xx 2 ] (cid:124) (cid:1) where BN is the batch normalization function.

Original: In (c) we consider the vector ﬁeld associated with a linear path w 2 w 1 x .
Revised: In (c) we consider the vector ﬁeld associated with a linear path w 2 w 1 x .

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: 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 form of implicit regularization, while others have been leveraged algorithmically in applications from network pruning to interpretability.

Original: Remarkably, all these geometric constraints can 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 f ( ✓ ) where ✓ 2 R m .
Revised: Consider a function f ( θ ) where θ ∈ R m .

Original: This function possesses a differentiable symmetry if it is invariant under the differentiable action  of a group G on the parameter vector ✓ , i.e., if ✓ 7! ( ✓ , ↵ ) where ↵ 2 G , then F ( ✓ , ↵ ) =
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 θ (cid:55)→ ψ ( θ, α ) where α ∈ G , then F ( θ, α ) =

Original: f (  ( ✓ , ↵ )) =
Revised: f ( ψ ( θ, α )) =

Original: f ( ✓ ) for all ( ✓ , ↵ ) .
Revised: f ( θ ) for all ( θ, α ) .

Original: The existence of a symmetry enforces a geometric structure on the gradient, r F .
Revised: The existence of a symmetry enforces a geometric structure on the gradient, ∇ F .

Original: Evaluating the gradient at the identity element  of G , so that ( ✓ , ↵ )
Revised: Evaluating the gradient at the identity element I of G (so that for all θ , ψ ( θ, I )

Original: = ✓ , yields the result,
Revised: = θ ) yields the result,

Original: [Equation]
Revised: [Equation]

Original: which implies that the gradient r f is perpendicular to the vector ﬁeld @
Revised: where (cid:104) , (cid:105) denotes the inner product. This equality implies that the gradient ∇ f is perpendicular to the vector ﬁeld ∂ α ψ | α

Original: ↵ that generates the symmetry, for all ✓ .
Revised: = I that generates the symmetry, for all θ .

Original: The symmetry also enforces a geometric structure on the Hessian, H F .
Revised: The symmetry also enforces a geometric structure on the Hessian, H F .

Original: Evaluating the Hessian at the identity elementof G yields the result,
Revised: Evaluating the Hessian at the identity element yields the result,

Original: [Equation]
Revised: [Equation]

Original: which constrains the Hessian H f .
Revised: which constrains the Hessian H f .

Original: See appendix A for the derivation of these properties and other geometric consequences of symmetry.
Revised: See appendix A for the derivation of these properties and other geometric consequences of symmetry.

Original: 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.
Revised: We will now consider the speciﬁc setting of a neural network parameterized by θ ∈ R m , the training loss L ( θ ) , and three families of symmetries (translation, scale, and rescale) that commonly appear in modern network architectures.

Original: Translation symmetry.
Revised: Translation symmetry.

Original: 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 , . .
Revised: Translation symmetry is deﬁned by the group R and action θ (cid:55)→ θ + α 1 A where 1 A is the indicator vector for some subset A of the parameters { θ 1 , . .

Original: , ✓ m } .
Revised: , θ m } .

Original: The loss function
Revised: The loss function

Original: L thus possesses translation symmetry if L ( ✓ ) =
Revised: L thus possesses translation symmetry if L ( θ ) =

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

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

Original: [Equation]
Revised: [Equation]

Original: ) and that the Hessian matrix H = @ 2 ✓ L has the indicator vector in its kernel,
Revised: and that the Hessian matrix H = ∂ 2 θ L has the indicator vector in its kernel,

Original: [Equation]
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 =e zi
Revised: Wx + b be the input to the softmax function such that σ ( z ) i =

Original: P j e zj .
Revised: [Equation] 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 the gradient constraints h @
Revised: this translation, yielding the gradient constraints (cid:104) ∂

Original: L @ w i , 1 i = h @
Revised: L ∂w i , 1 (cid:105) = (cid:104) ∂

Original: L @ b , 1 i = 0 , visualized in Fig.
Revised: L ∂b , 1 (cid:105) = 0 , visualized in Fig.

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

Original: Scale symmetry.
Revised: Scale symmetry.

Original: Scale symmetry is deﬁned by the group GL
Revised: Scale symmetry is deﬁned by the group GL

Original: +

Revised: +

Original: ( R ) and action ✓ 7!
Revised: ( R ) and action θ (cid:55)→ α A (cid:12)

Original: ↵ A  ✓ where ↵
Revised: θ where α A =

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

Original: The loss function possesses scale symmetry if L ( ✓{+θ ) = L ( α A (cid:12) θ ) for all α ∈ GL +
 ( R+} ) 
Revised: The loss function possesses scale symmetry if L ( {+θ ) = L ( α A (cid:12) θ ) for all α ∈ GL +
 ( R+} ) .

Revised: This symmetry immediately implies the loss gradient is everywhere perpendicular to the parameter vector itself ∂ α ψ | α

Original: = 
Revised: = I

Revised: [Equation]

Revised: A

Revised: [Equation]

Revised: =

Revised: [Equation]

Revised: , (5) and relates to the Hessian matrix, where [ ∂ α ∂ θ ψ | α

Original:  L  ( ↵ A ✓ ) for all ↵ 2 GL
Revised: = I ] ∂ θ L = diag ( 1 A ) g = g A , as Hθ A + g A = 0 .

Revised: Batch normalization.

Revised: Batch normalization leads to scale invariance during training.

Revised: Let z = w (cid:124) x

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

Revised: z − E [ z ] √ Var ( z ) where E [ z ] is the sample mean and Var ( z ) is the sample variance given a batch of data.

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.

Revised: Thus, these parameters observe scale symmetry in the loss and their gradients satisfy (cid:10) ∂

Revised: L ∂w , w (cid:11) + (cid:10) ∂

Revised: L ∂b , b (cid:11) = 0 , as has been previously noted by Ioffe & Szegedy (2015); Van Laarhoven (2017), and visualized in Fig.

Revised: 2 for the toy model where w, b ∈ R .

Revised: Rescale symmetry.

Original: +1{+Rescale symmetry is deﬁned by the group GL +
+} ( R ) .
Revised: {+Rescale symmetry is deﬁned by the group GL +
+} ( R ) and action θ (cid:55)→ α A 1 (cid:12) α − 1 A 2 (cid:12) θ where A 1 and A 2 are two disjoint sets of parameters.

Revised: The loss function possesses rescale symmetry

Original: {+if L ( θ ) = L ( α A 1 (cid:12) α − 1 A 2 (cid:12) θ ) for all α ∈ GL +
 ( R ) .+} This symmetry immediately implies the loss gradient is everywhere perpendicular to the  parameter vector itself @ ↵ ↵ = ✓  1 A
Revised: {+if L ( θ ) = L ( α A 1 (cid:12) α − 1 A 2 (cid:12) θ ) for all α ∈ GL +
 ( R ) .+} This symmetry immediately implies the loss gradient is everywhere perpendicular to the sign inverted parameter vector ∂

Original:  = ✓ A , h g, ✓ A i = 0 ,
Revised: α ψ | α = I = θ A 1 − θ A 2 = θ (cid:12)

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

Original: and relates to the Hessian matrix, where @ ✓ L @ ↵ @ ✓ ↵ = g diag
Revised: and relates to the Hessian matrix, where [ ∂ α ∂ θ ψ | α = I ] ∂ θ

Original: [Equation]
Revised: [Equation]

Original: [Equation] ) =
Revised: g A

Original: [Equation]
Revised: [Equation]

Original:  A  , 
Revised: g A 2 , as

Original: [Equation]
Revised: [Equation]

Original: Batch normalization.
Revised: Homogeneous activation.

Original: Batch normalization leads to scale invariance during training. Let z 
Revised: For networks with continuous, homogeneous activation functions φ ( z ) =

Original: =
Revised: φ (cid:48)

Original: w | x + b be the input to  a  neuron with batch normalization such that BN ( z ) =  z  E [ z ] p Var ( z )where E [ z ] is the sample mean
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. For example, consider a hidden neuron with ReLU activation φ ( z ) = max { 0 , z } , such that w 2 φ ( w (cid:124) 1 x + b ) is the computational path through this neuron.

Original: and Var ( z ) is the sample variance given a batch of data. Notice that scaling w  and b by anon-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: 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 scale symmetryin the loss and their gradients satisfy 
Revised: Thus, these parameters observe rescale symmetry and their gradients satisfy (cid:68) [Equation] , w [Equation]

Revised: (cid:69) + (cid:10)

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

Revised: − (cid:68)

Original:  0 , as has been previously noted by Ioffe & Szegedy (2015); Van Laarhoven (2017), and visualized in Fig.
Revised: [Equation] 2 , w 2 (cid:69) = 0 , as has been previously noted by Du et al. (2018a) ; Liang et al.

Revised: ; Tanaka et al.

Original:  2 for the toy model where w, b 2  R  .
Revised: (2020), and visualized in Fig. 2 for the toy model where w 1 , w 2 ∈ R and b = 0 .

Original: Rescale symmetry. Rescale symmetry is deﬁned by the group GL
Revised: We now explore how geometric constraints on gradients and Hessians, arising as a consequence of symmetry, impact the learning dynamics given by stochastic gradient descent (SGD).

Original: [-+
-]
Revised: [-+
-]We will consider a model parameterized by θ , a training dataset { x 1 , . . .

Original:  ( R ) and action ✓ 7!
Revised: , x N } of size N , and a training loss L ( θ ) =

Original: ↵ A 1  ↵  1 A 2  ✓ where A 1 and A 2 are two disjoint sets of parameters.
Revised: 1 N (cid:80) Ni =1 (cid:96) ( θ, x i ) with corresponding gradient g ( θ ) = ∂ L ∂θ . [Figure] where the axes are conserved.

Original: The loss function possesses rescale symmetry if L ( ✓ ) = 
Revised: The color represents the value of the conserved quantity, where blue is positive and red is negative, and the black lines are level sets. The gradient descent update with learning rate η is θ ( n +1) = θ ( n ) −

Original: L ( ↵ A 1  ↵  1 A 2  ✓ )  for all ↵ 2 GL
Revised: ηg ( θ ( n ) ) , which is a forward Euler discretization with step size η of the ordinary differential equation (ODE) [Equation] 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). 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: +1 ( R ) . This symmetry immediately implies the loss gradient is everywhere perpendicular to the sign inverted parameter vector @
Revised: Symmetry leads to conservation. Strikingly similar to Noether’s theorem, which describes a fundamental relationship between symmetry and conservation for physical systems governed by Lagrangian dynamics, here we show that symmetries of a network architecture have corresponding conserved quantities through training under gradient ﬂow. Theorem 1.

Original: ↵
Revised: Symmetry and conservation laws in neural networks. Every differentiable symmetry ψ ( α, θ ) of the loss that satisﬁes (cid:104) θ, [ ∂ α ∂ θ ψ | α = I ] g ( θ ) (cid:105) = 0 has the corresponding conservation law,

Original: [Equation]
Revised: [Equation]

Original: ✓
Revised: through learning under gradient ﬂow. To prove Theorem 1, consider projecting the gradient ﬂow learning dynamics in equation (9) onto the generator vector ﬁeld ∂ α ψ | α

Original: [Equation]
Revised: = I associated with a symmetry.

Original: [Equation]
Revised: As shown in section 3, the gradient of the loss g ( θ ) is always perpendicular to the vector ﬁeld ∂ α ψ | α = I . Thus, the projection yields a differential equation, which can be simpliﬁed to equation (10).

Original: and relates to the Hessian matrix, where @ ✓
Revised: See appendix B for a complete derivation. Application of this general theorem to the translation, scale, and rescale symmetries, identiﬁed in section 3, yields the following conservation law of learning,

Original: [Equation]
Revised: [Equation]

Original: Homogeneous activation.
Revised: Each of these equations deﬁne a conserved constant of learning through training. For parameters with translation symmetry, their sum ( (cid:104) θ A ( t ) , 1 (cid:105) ) is conserved, effectively constraining their dynamics to a hyperplane.

Original: For networks with continuous, homogeneous activation functions 
Revised: For parameters with scale symmetry, their euclidean norm ( | [Equation] A ( t

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

Revised: ) is conserved, effectively constraining their dynamics to a sphere Ioffe & Szegedy (2015); Van Laarhoven (2017).

Original:  ( 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.
Revised: For parameters with rescale symmetry, their difference in squared euclidean norm ( | θ A 1 ( t ) | 2 − | θ A 2 ( t ) | 2 ) is conserved, effectively constraining their dynamics to a hyperbola Du et al. (2018a). In Fig. 3 we visualize the level sets of these conserved quantities for the toy models discussed in Fig.

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

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.

Revised: However, empirically we know these laws are broken, as demonstrated in Fig.

Revised: What causes this discrepancy?

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

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.

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

Revised: Modeling weight decay.

Original: For example, consider a hidden neuron with ReLU activation  ( z )
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 optimizer’s update equation. For stochastic gradient descent, the result leads to the updated continuous model [Equation]

Revised: Modeling momentum.

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

Original: = max { 0 , z } , such that w 2  ( w | 1 x + b ) is the computational path through this neuron.
Revised: 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 [Equation]

Revised: g ( θ ( n ) ) .

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: 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. 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 1 , [Equation]

Revised: Modeling stochasticity.

Revised: Stochastic gradients ˆ g B ( θ ) arise when we consider a batch B of size S drawn uniformly from the indices { 1 , . . .

Revised: , N } forming the unbiased gradient estimate ˆ g B ( θ ) =

Revised: S (cid:80) i ∈B ∇ (cid:96) ( θ, x i ) .

Revised: When the batch size is much smaller than the size of the dataset, S (cid:28) N , then we can model the batch gradient as an average of S i.i.d.

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

Revised: Using the central limit theorem, we assume ˆ g B ( θ ) − g

Revised: ( θ ) is a Gaussian random variable with mean µ = 0 and covariance matrix Σ( θ ) .

Original: Thus, these parameters observe rescale symmetry and their gradients satisfy
Revised: However, because both the batch gradient and true gradient observe the same geometric properties introduced by symmetry, the noise has a special low-rank structure. As we showed in section 3, the gradient of the loss, regardless of the batch, is orthogonal to the generator vector ﬁeld ∂ α ψ | α

Original: D
Revised: = I associated with a symmetry. This implies the stochastic noise must also observe the same property. In order for this relationship to hold for arbitrary noise, then Σ( θ ) ∂ α ψ | α

Original: [Equation]
Revised: = I = 0 . In other words, the differential symmetry inherent in neural network architectures projects the noise introduced by stochastic gradients onto low rank subspaces, leaving the gradient ﬂow dynamics in these directions unchanged 2 . Modeling discretization. The effect of discretization when modeling continuous dynamics is a well studied problem in the numerical analysis of partial differential equations.

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.

Revised: We present two methods based on modiﬁed equation analysis, which modify gradient ﬂow to account for the effect of discretization.

Revised: [Figure]

Revised: [Figure]

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

Revised: Li et al.

Revised: ; Feng et al.

Revised: (2019) and

Revised: most recently Barrett & Dherin (2020), demonstrate that the gradient descent trajectory closely follows the steepest descent path of a modiﬁed loss function (cid:101) L .

Revised: The divergence between these trajectories fundamentally depends on the learning rate η

Revised: and the curvature

Revised: H .

Revised: As derived in Barrett & Dherin (2020), and summarized in appendix D, this divergence is given by the gradient correction

Revised: − 12 Hg , which is the gradient of the squared norm − η 4 |∇L| 2 .

Revised: Thus, the modiﬁed loss is (cid:101)

Revised: L = L + η 4 |∇L| 2 and the modiﬁed gradient ﬂow ODE is

Revised: [Equation]

Revised: See Fig.

Revised: 4 for an illustrative example of this method applied to a quadratic loss in R 2 .

Revised: Modiﬁed ﬂow.

Revised: Rather than modifying gradient ﬂow with higher order “spatial” derivatives of the loss function, here we introduce higher order temporal derivatives.

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

Original: {+Rearranging the update equation for gradient descent, θ t +
+}
Revised: {+Rearranging the update equation for gradient descent, θ t +
+}

Revised: =

Revised: θ t − ηg ( θ t ) , and assuming θ ( t )

Revised: = θ t and θ ( t + η ) =

Revised: θ t +1 , gives the equality − g ( θ t )

Revised: = θ ( t + η ) − θ ( t ) η , which Taylor expanding the right side results in the differential equation − g

Revised: ( θ t ) =

Revised: dθdt +

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

Original:  , w
Revised: 2 + O ( η 2 ) . Notice that in the limit as η → 0 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] [Equation]

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.

Revised:  C OMBINING S YMMETRY AND M ODIFIED G RADIENT F LOW TO D ERIVE E XACT L EARNING D YNAMICS

Revised: As shown in section 4, each symmetry results in a conserved quantity under gradient ﬂow.

Original: E
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. To do this we (i) consider a realistic continuous model for SGD, (ii) project these learning dynamics onto the generator vector ﬁelds ∂ α ψ | α = I 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.

Revised: For simplicity, we ﬁrst consider the continuous model of SGD incorporating weight decay and modiﬁed loss, but not momentum and stochasticity 3 .

Revised: In this setting, the exact dynamics, as fully derived in appendix E, for the parameter combinations tied to the symmetries are,

Revised: [Equation]

Revised: Notice how these equations are equivalent to the conservation laws derived in section 4 when η = λ = 0 .

Revised: Remarkably, even in typical hyperparameter settings (weight decay, stochastic batches, ﬁnite learning rates), these solutions match nearly perfectly with empirical results 4 from modern neural networks (VGG-16) trained on real-world datasets (Tiny ImageNet), as shown in Fig.

Revised: We will now discuss each equation individually.

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

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

Revised: S ca l e

Revised: R e s ca l e

Revised: Translation dynamics.

Revised: For parameters with translation symmetry, equation 18 implies that the sum of these parameters ( (cid:104) θ A ( t ) , 1 (cid:105) ) decays exponentially to zero at a rate proportional to the weight decay.

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

Revised: This is due to the lack of curvature in the gradient ﬁeld for these parameters (as shown in Fig. 2).

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.

Revised: The ﬁrst row in Fig.demonstrates this qualitatively, as all trajectories are smooth exponential functions that converge faster for increasing levels of weight decay.

Revised: Scale dynamics.

Revised: For parameters with scale symmetry, equation 19 implies that 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.

Revised: Compared to the translation dynamics, the scale dynamics do depend on the data through the gradient norms accumulated throughout training.

Revised: Without weight decay λ = 0 , the ﬁrst term stays constant and the second term grows monotonically.

Revised: With weight decay λ > 0 , the ﬁrst term decays monotonically to zero, while the second term can decay or grow, but always stays positive.

Revised: The second row in Fig.

Revised: 5 demonstrates these qualitative relationships.

Revised: Without weight decay the norms increase monotonically as predicted and with weight decay the dynamics are non-monotonic and present more complex behavior.

Revised: To better understand the forces driving these complex dynamics, we can examine the time derivative of equation 19,

Revised: [Equation]

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

Revised: The centripetal effect due to weight decay is a direct consequence of its regularizing inﬂuence, pulling the parameters towards the origin.

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

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.

Revised: In particular, at stationarity, then ddt | θ ( t ) | 2 = 0 , which rearranging equation (21) gives the condition ω ( t ) ≡ (cid:12)(cid:12) dθdt (cid:12)(cid:12) /

Revised: | θ | = (cid:113) 2 λη .

Revised: Consistent with the results of Wan et al.

Revised: (2020), this implies that at stationarity the angular speed ω ( t ) of the weights is constant and governed only by the learning rate η and weight decay constant λ .

Revised: When considering stochasticity explicitly, as explained in appendix F, then these dynamics also depend on the covariance of the gradient noise Σ( θ ) and batch size S .

Revised: Rescale dynamics.

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.

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

Revised: This leads to quite sophisticated, complex dynamics, consistent with the third row in Fig.

Revised: Despite the complexity, our theory, nevertheless, quantitatively matches the empirics.

Revised: 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] [Table]

Revised: +RIyQ1xq7Var1xxq+5cZBW8HCqQq9Erf/n9hGUxV8gkNabruSkGE6pRMMmnJT8zPKVsRAe8a1HRmJtgMr95Ss6s0ydRou1TSObu74kJjY0Zx6HtjCkOzXJtZv5X62YYXQcTodIMuWKLRVEmCSZkFgDpC80ZyrEFyrSwtxI2pJoytDGVbAje8pdXoXVR9Sw/XFbqt3kcRTiBUzgHD66gDvfQgCYwSOEZXuHNyZwX5935WLQWnHzmGP7I+fwB0LWQNw==</latexit>

Revised: Harmonic oscillation with momentum.

Revised: We will now consider a continuous model of SGD with momentum.

Revised: As discussed in appendix G, we consider the continuous model incorporating weight decay, momentum, stochasticity, and modiﬁed ﬂow.

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.

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.

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

Revised: Tiny ImageNet, as in right panel of Fig.

Revised:  C ONCLUSION

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

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

Revised: In this paper, we constructed a unifying theoretical framework harnessing the geometric properties of symmetry and more realistic continuous equations for learning dynamics that model weight decay, momentum, stochasticity, and discretization.

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

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.

Revised: Overall, this work provides a ﬁrst step towards understanding the mechanics of learning in neural networks without unrealistic simplifying assumptions.

Revised: A CKNOWLEDGEMENTS

Revised: We thank Daniel Bear, Lauren Gillespie, Kyogo Kawaguchi, Brett Larsen, Eshed Margalit, Alain

Revised: Studer, Sho Sugiura, Umberto Maria Tomasini and Atsushi Yamamura for helpful discussions.

Revised: This work was funded in part by the IBM-Watson AI Lab.

Revised: D.K. thanks the Stanford Data Science Scholars program for support.

Revised: J.S. thanks the Mexican National Council of Science and Technology (CONACYT) for support.

Revised: S.G. thanks the James S. McDonnell and Simons Foundations, NTT Research, and an NSF CAREER Award for support.

Revised: D.L.K.Y thanks the McDonnell Foundation (Understanding

Revised: Human Cognition Award Grant No. 220020469)

Revised: , the Simons Foundation (Collaboration on the

Revised: Global Brain Grant No. 543061), the Sloan Foundation (Fellowship FG-2018-10963), the National

Revised: Science Foundation (RI 1703161 and CAREER Award 1844724) and the DARPA Machine Common

Revised: Sense program for support and the NVIDIA Corporation for hardware donations

Revised: A S YMMETRY AND G EOMETRY

Revised: Here we derive in detail the geometric properties of the loss landscape introduced by symmetry, as discussed in section 3.

Revised: Consider a function f ( θ ) where θ ∈ R m .

Revised: This function possesses a symmetry if it is invariant under the action ψ of a group G on the parameter vector θ , i.e., if θ (cid:55)→ ψ ( θ, α ) where α ∈ G , then F ( θ, α ) =

Revised: f ( ψ ( θ, α ))

Revised: =

Revised: f ( θ ) for all ( θ, α ) .

Revised: Symmetry enforces a geometric relationship between the gradient ∇ F and Hessian H F of the composition F ( θ, α ) with the gradient ∇ f and Hessian H f of the original function f ( θ ) .

Revised: This relationship can be described by ﬁve constraints on ∇ f and H f .

Revised: Considering these general formulae when f ( θ )

Revised: =

Revised: L ( θ ) , the training loss of a neural network, yields ﬁfteen distinct equations describing the geometrical relationships between architectural symmetries and the loss landscapes, some of which have been identiﬁed individually in existing literature (Table 1).

Revised: [Table]

Revised: Ioffe & Szegedy (2015) has motivated the effectiveness of Batch Normalization by its scale invariant property.

Revised: In particular, they have noted that Batch Normalization will stabilize back-propagation because “The scale does not affect the layer Jacobian nor, consequently, the gradient propagation.

Revised: Moreover, larger weights lead to smaller gradients, and Batch

Revised: Normalization will stabilize the parameter growth.”

Revised: Van Laarhoven (2017) has then shown that the role of L 2 regularization when combined with batch Ioffe & Szegedy (2015), weight Salimans & Kingma (2016) or layer Ba et al.

Revised: (2016) normalization is not to regularize the function, but to effectively control the learning rate.

Revised: Zhang et al.

Revised: (2018) has thoroughly studied mechanisms of weight decay regularization and derived various geometric properties of loss landscapes along the way.

Revised: Arora et al.

Revised: (2018c) has theoretically analyzed the automatic tuning property of learning rate in networks with

Revised: Batch Normalization.

Original: E = 0 , as has been previously noted by Du et al. 
Revised: Li et al. (2020) studied the interaction of weight decay and batch normalization in the setting of stochastic gradients. Neyshabur et al. (2015) has identiﬁed that SGD is not rescale equivariant even when network outputs are rescale invariant.

Revised: This is a problem because gradient descent performs very poorly on unbalanced networks due to the lack of equivariance.

Revised: Motivated by the issue, they have introduced a new optimizer, Path-SGD, that is rescale equivariant.

Original:  (2018a) 
Revised: Arora et al. (2018b) proved that the weights of linear artiﬁcial neural networks satisfy strong balancedness property. Du et al. (2018a) studied implicit regularization in networks with homogeneous activation functions.

Revised: To do that, they showed conservation law of parameters with rescale invariance.

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

Revised: (2019) have proposed a new capacity measure to study generalization that respects rescale invariance of networks.

Revised: Along the way, they showed geometric properties of gradients and Hessians for networks with rescale invariance.

Revised: However, their results were restricted to a layer without biases.

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

Original: (2020), and visualized in Fig.
Revised: (2020) have proved gradient properties of parameters with rescale invariance at neuron level including biases. A. G RADIENT G EOMETRY If a function f posses a symmetry, then there exists a geometric constraint on the relationship between the gradients ∇ F and ∇ f at all ( θ, α ) , [Equation]

Original: 2 for the toy model where w 1 , w 2 2 R and b = 0 .
Revised: The top element of the gradient relationship, ∂ θ F , evaluated at any ( θ, α ) , yields the property [Equation] 22) which describes how the symmetry transformation affects the function’s gradients despite leaving the output unchanged.

Revised: The bottom element of the gradient relationship, ∂ α F , evaluated at the identity element of G yields the property

Revised: [Equation]

Revised: which implies the gradient ∇ f is perpendicular to the vector ﬁeld

Revised: ∂ α ψ | α

Revised: = I that generates the symmetry, for all θ .

Revised: In the speciﬁc setting when f ( θ ) =

Revised: L ( θ ) , the training loss of a neural network, these gradient properties are summarized in Table 2 for the translation, scale, and rescale symmetries described in section 3.

Revised: Translation

Original: General rescale symmetry.
Revised: [Table] [Equation] [Table] [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: Notice that the ﬁrst row of Table 2 implies that symmetry transformations affect learning dynamics governed by gradient descent for scale and rescale symmetries, while it does not for translation symmetry.

Revised: These observations are in agreement with Van Laarhoven (2017) who has shown that effective learning rate is inversely proportional to the norm of parameters immediately preceding the batch normalization layers and Neyshabur et al.

Revised: (2015) who have noticed that SGD is not invariant to the rescale symmetry that the network output respects and proposed Path-SGD to ﬁx the discrepancy.

Revised: A. H ESSIAN G EOMETRY

Revised: If a function f posses a symmetry, then there also exists a geometric constraint on the relationship between the Hessian matrices H F and H f

Revised: at all ( θ, α ) ,

Revised: [Equation]

Revised: The ﬁrst diagonal element, ∂ 2 θ F , evaluated at any ( θ, α ) , yields the property

Revised: [Equation]

Revised: 24) which describes how the symmetry transformation affects the function’s Hessian despite leaving the output unchanged.

Original: A more general rescale symmetries can be deﬁned by the group GL
Revised: The off-diagonal elements, ∂ θ ∂ α F = ∂ α ∂ θ F , evaluated at the identity element of G yields the property [Equation] [Equation]

Revised: which implies the geometry of gradient and Hessian are connected through the action of the symmetry.

Revised: Lastly, the second diagonal element, ∂ 2 α F , represents an equality, evaluated at the identity element of

Revised: G yields the property

Revised: [Equation]

Revised: [Equation]

Revised: which combines the geometric relationships in equation 23 and equation 25.

Revised: In the speciﬁc setting when f ( θ )

Original: [-+
-]
Revised: [-+
-]= L ( θ ) , the training loss of a neural network, these Hessian properties are summarized in Table 2 for the translation, scale, and rescale symmetries described in section 3. [Table]

Revised: Scale

Original:  ( R ) and action ✓ 7!
Revised: Rescale [Equation] The translation, scale, and rescale symmetries identiﬁed in section 3 is not an exhaustive list of the symmetries present in neural network architectures. For example, more general rescale symmetries can be deﬁned by the group GL + 1 ( R ) and action θ

Original: ↵ k 1 A 1  ↵  k 2 A 2  ✓ , which occur in networks with quadratic activation functions.
Revised: (cid:55)→ α k 1 A 1 (cid:12) α − k 2 A 2 (cid:12) θ , which occur in networks with quadratic activation functions.

Original: 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.
Revised: A stronger form of rescale symmetry also 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.

Original: (2018a).
Revised: (2018a).

Original: S YMMETRY LEADS TO C ONSERVATION L AWSU NDER G RADIENT F LOW
Revised: Interestingly, some of the gradient and Hessian properties for scale symmetry can also be easily proven as consequences of Euler’s Homogeneous Function Theorem when k = 0 . B C ONSERVATION L AWS

Original: We now explore how geometric constraints on gradients and Hessians arising as a consequence of symmetry, impact network learning dynamics.
Revised: Here we repeat Theorem 1 and provide a detailed derivation. Theorem. Symmetry and conservation laws in neural networks.

Original: In this work we focus on stochastic gradient descent (SGD), the workhorse of modern deep{+the loss that satisﬁes (cid:104) θ, [ ∂ α ∂ θ ψ | α = I ] g ( θ ) (cid:105) = 0 has the corresponding conservation law ddt (cid:104) θ, ∂ α ψ | α = I (cid:105) = 
 through+} learning optimization.
Revised: Every differentiable symmetry ψ ( α, θ ) of {+the loss that satisﬁes (cid:104) θ, [ ∂ α ∂ θ ψ | α = I ] g ( θ ) (cid:105) = 0 has the corresponding conservation law ddt (cid:104) θ, ∂ α ψ | α = I (cid:105) = 
 through+} learning under gradient ﬂow.

Original: We will consider a model parameterized by ✓ , a training dataset { x 1 , . . . , x N } of size N , and a training loss L ( ✓ )  = 1 N
Revised: Proof. Project the gradient ﬂow learning dynamics, dθdt = − g ( θ ) , onto the vector ﬁeld that generates the symmetry ∂ α ψ | α = I , evaluated at the identity element,

Original: [Figure]
Revised: [Equation]

Revised: We can factor the left side of this equation as

Original: [Equation]
Revised: [Equation]

Original: [Figure] corresponding gradient g ( ✓ ) = @
Revised: By assumption, (cid:104) θ, [ ∂ α ∂ θ ψ | α

Original: L
Revised: = I ] g ( θ ) (cid:105) = 0 , implying ddt (cid:104) θ, ∂ α ψ | α = I (cid:105) = 0 . The condition (cid:104) θ, [ ∂ α ∂ θ ψ | α =

Revised: I ] g ( θ ) (cid:105) = 0 holds for the translation, scale, and rescale symmetries we consider in section 3.

Revised: For translation symmetry, ∂ α ∂ θ ψ | α = I = 0 .

Revised: For scale symmetry, ∂ α ∂

Revised: C L IMITING D IFFERENTIAL E QUATIONS FOR L EARNING R ULES

Revised: Here we identify ordinary differential equations whose ﬁrst-order discretization give rise to the gradient descent and classical momentum algorithms.

Revised: These differential equations can be understood as the limiting dynamics for their respective discrete algorithms as the learning rate η → 0 .

Revised: C. G RADIENT D ESCENT

Revised: Gradient descent with learning rate η is given by the update equation

Revised: [Equation]

Revised: and initial condition θ 0 .

Revised: Rearranging the difference between consecutive updates gives

Revised: [Equation]

Revised: This is a discretization with step size η of the ﬁrst order ODE

Revised: [Equation]

Revised: where we used the forward Euler discretization ddt θ k

Revised: = θ k +1 − θ k η .

Revised: This ODE is commonly referred to as gradient ﬂow .

Revised: C. C LASSICAL M OMENTUM

Revised: Classical momentum with learning rate η , damping coefﬁcient α , and momentum parameter β , is given by the update equation

Revised: [Equation]

Revised: and initial conditions v 0 = 0 and some θ 0 .

Revised: The difference between consecutive updates is

Revised: [Equation]

Revised: Rearranging this equation we get

Revised: [Equation]

Revised: This is a discretization with step size η

Revised: (1 − α ) of the ﬁrst order

Revised: ODE

Original: [Figure]
Revised: [Equation] where we used the forward Euler discretization ddt θ k = θ k +1 − θ k η (1 − α ) and backward Euler discretization ddt θ k = θ k − θ k − 1 η (1 − α ) .

Original: The gradient descent update with learning rate ⌘ is ✓ ( n +1)
Revised: We will refer to this equation as momentum ﬂow .

Original: = ✓ ( n )  ⌘ g ( ✓ ( n ) ) , which is a forward Euler discretization 1 with step size ⌘ of the ordinary differential equation (ODE)
Revised: A more detailed derivation for this ODE under Nesterov variants of classical momentum can be found in Su et al. (2014) and Kovachki & Stuart (2019). D

Original: [Equation]
Revised: M ODIFIED [Figure]

Original: 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).
Revised: We visualize the trajectories given by these dynamics using the initial condition x 0 = [ 1 0 ] (cid:124) (white circle) and a step size η = 0 . As we can see, the modiﬁed continuous trajectory better matches the discrete trajectory. 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.

Revised: The divergence between the discrete and continuous trajectories fundamentally depends on the learning rate and the curvature of the loss.

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: It is thus natural to assume there exists a more realistic continuous model for SGD that incorporates both these terms in a non-trivial way. How can we better model the discrete dynamics of gradient descent with a continuous differential equation? Intuition from ﬁnite difference methods. To answer this question we will take inspiration from tools developed for ﬁnite difference methods.

Revised: Finite difference methods are a class of numerical techniques for approximating derivatives in the analysis of partial differential equations (PDE).

Original: Symmetry leads to conservation.
Revised: These approximations are applied iteratively to construct numerical solutions to a PDE given some initial conditions. However, this discretization process introduces numerical artifacts, which can lead to a signiﬁcant difference between the numerical solution and the true solution for the PDE. Modiﬁed equation analysis is a method for understanding this difference by modeling the numerical artifacts as higher-order spatial or temporal derivatives modifying the original PDE. This approach can be used to construct modiﬁed continuous dynamics that better approximate discrete dynamics, as illustrated in Fig.

Revised: D. M ODIFIED L OSS

Revised: Taking inspiration from modiﬁed equation analysis, Li et al.

Revised: ; Feng et al.

Original: 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 .
Revised: (2019) and most recently Barrett & Dherin (2020), demonstrate that the trajectory given by gradient descent closely follows the steepest descent path of a modiﬁed loss function (cid:101) L , rather than the original loss L . As explained in Barrett & Dherin (2020), assume there exists a modiﬁed vector ﬁeld with corrections g i in powers of the learning rate to the original vector ﬁeld g that the discrete dynamics follow. In other words, rather than considering the dynamics given by gradient ﬂow, ddt θ = − g ( θ ) , we consider the modiﬁed differential equation,

Revised: [Equation]

Revised: Truncating the modiﬁed vector ﬁeld up to the order η and using backward error analysis we can derive that the ﬁrst-order correction g 1 = − 12 Hg , which is the gradient of the squared norm − η 4 |∇L| 2 .

Revised: Thus, the truncated modiﬁed differential equation is simply gradient ﬂow on a modiﬁed loss (cid:101)

Original: Consider some subset of  the parameters A that respects either a translation, scale, or rescale symmetry.
Revised: L = L + η 4 |∇L| 2 . Convex quadratic loss. To illustrate modiﬁed loss, we will consider the trajectories of gradient descent w t , gradient ﬂow ˆ w ( t ) , and modiﬁed gradient ﬂow ˜ w ( t ) on the convex quadratic loss L ( w ) =

Revised: 12 w (cid:124)

Original: As shown in section 3, the gradient of the loss g ( ✓ ) is always perpendicular to the vector ﬁeld that generates the symmetry @{+update formula w t +
 =+}
Revised: Aw , where A (cid:31) 0 is some positive deﬁnite matrix, as shown in Fig. For a ﬁnite learning rate η and initial condition w 0 , gradient descent is given by the {+update formula w t +
 =+}

Original: ↵ . 
Revised: w t − ηAw t . Gradient ﬂow is deﬁned as ddt ˆ w = − A ˆ w , a linear ﬁrst-order differential equation.

Original: Projecting the gradient ﬂow learning dynamics in equation (27) onto the generator vector ﬁeld yields a differential equation h d ✓ dt , @
Revised: For the initial condition w 0 , the resulting initial value problem can be solved exactly giving the gradient ﬂow trajectory ˆ w ( t ) = S − 1 e − Λ t Sw 0 , where A = S − 1 Λ S is the diagonalization of the curvacture matrix. For learning rate η , the modiﬁed loss is (cid:101) L ( w ) = 12 w (cid:124)

Revised: Aw +

Revised: [Equation]

Revised: 4 w (cid:124)

Revised: A 2 w and modiﬁed gradient ﬂow is deﬁned as ddt ˜ w = − A ˜ w − η 2 A 2 ˜ w .

Revised: For the initial condition w 0 , the resulting initial value problem can be solved exactly giving the modiﬁed gradient ﬂow trajectory ˜ w (

Revised: [Equation]

Original: ↵ i = 0 .
Revised: = [Equation] e − [Equation]

Revised: D. M ODIFIED F LOW

Revised: Rather than modify gradient ﬂow with higher order “spatial” derivatives of the loss function, here we introduce higher order temporal derivatives.

Revised: We start by assuming the existence of a continuous trajectory θ ( t ) that weaves through the discrete steps taken by SGD and then identify the differential equation that generates the trajectory.

Original: Integrating this equation through time results in the conservation laws,
Revised: Rearranging the update equation for SGD, θ t +1 = θ t − ηg B ( θ t ) , and assuming θ ( t ) = θ t and θ ( t + η ) =

Revised: θ t +1 , gives the equality − g B ( θ t )

Revised: = θ ( t + η ) − θ ( t ) η , which

Revised: Taylor expanding the right side results in the differential equation − g B ( θ t ) = dθdt + η 2 d 2 θ dt 2 + O ( η 2 ) .

Original: [Equation]
Revised: Notice that in the limit as η → 0 we regain gradient ﬂow. For small η , η (cid:28) 1 , we obtain a modiﬁed version of gradient ﬂow with an additional second-order term. This approach of modifying ﬁrst-order differential equations with higher order temporal derivatives was applied by Kovachki & Stuart (2019) to modify momentum ﬂow, capturing the harmonic motion of momentum. Convex quadratic loss.

Revised: To illustrate modiﬁed ﬂow, we will consider the trajectories of momentum w t , momentum ﬂow ˆ w ( t ) , and modiﬁed momentum ﬂow ˜ w ( t ) on the convex quadratic loss L ( w )

Original: Each of these equations deﬁne a conserved constant  of learning through training.
Revised: = 12 w (cid:124) Aw , where A (cid:31) 0 is some positive deﬁnite matrix, as shown in Fig. For a ﬁnite learning rate η , dampening α = 0 , momentum constant β , and initial conditions v 0 = 0 , w 0 , then momentum is given by the pair of recursive update equations v t +1 = βv t + Aw t and w t +1 =

Original: For parameters with translation symmetry, their sum is conserved, effectively constraining their dynamics to a hyperplane.
Revised: w t − ηv t +1 . Momentum ﬂow is deﬁned as (1 − β ) ddt ˆ w = − A ˆ w , a linear ﬁrst-order differential equation. For a given initialization w 0 , the resulting initial value problem can be solved exactly as in the case of gradient ﬂow. Modiﬁed momentum ﬂow is deﬁned as η 2 (1 + β ) d 2 dt 2 ˜ w + (1 − β ) ddt ˜ w =

Revised: −

Revised: A ˜ w , a linear second-order differential equation.

Revised: For a given initialization w 0 and the assumed initial condition ddt ˜ w (0) = 0 , then the resulting initial value problem can be solved exactly as a system of damped harmonic oscillators.

Revised: E D ERIVING THE E XACT L EARNING D YNAMICS OF SGD

Revised: We consider a continuous model of SGD without momentum incorporating weight decay (equation 14) and modiﬁed loss (equation 16), such that

Revised: [Equation]

Revised: where L λ =

Revised: L + λ 2 | θ | 2 is the regularized loss.

Revised: The gradient of the modiﬁed loss is,

Revised: [Equation]

Original: For parameters with scale symmetry, their euclidean norm isconserved, effectively constraining their dynamics to a sphere.
Revised: where we used ∇L λ = g + λθ and H λ = H + λI . Thus, the equation of learning we consider is

Original: For parameters with rescale symmetry, their difference in squared euclidean norm is conserved, effectively constraining their dynamics to a hyperbola. In Fig.
Revised: [Equation] To incorporate the effect of stochasticity (equation 31) in this equation of learning we could replace the full-batch gradient g and Hessian H with their stochastic batch counterparts ˆ g B and ˆ H B respectively.

Original: 3 we visualize the level sets of these conserved quantities for the toy models discussed in Fig.
Revised: However, careful treatment of these terms using stochastic calculus is needed when integrating the resulting stochastic differential equation, which we discuss at the end of this section.

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

Original: 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.
Revised: In the case of parameters with translation symmetry, the effect of discretization essentially leaves the dynamics for the constant of learning unchanged. Combining the geometric properties of gradient ( (cid:104) g, 1 A (cid:105) = 0 ) and Hessian ( H 1 A = 0 ) introduced by translation symmetry with the equation of learning (equation 27) gives the differential equation, [Equation] where we used the simpliﬁcation,

Original: However, empirically we know these laws are broken, as demonstrated in Fig.
Revised: [Equation] and ignored the O ( ηλ 2 ) term as we set the weight decay constant λ to be as small as the learning rate η in practice.

Original: What causes this discrepancy?
Revised: The solution to this differential equation is

Original: Gradient ﬂow is too simple of a continuous model for realistic SGD training.
Revised: [Equation] Scale dynamics.

Original: It fails to incorporate the effect ofweight decay introduced by explicit regularization, momentum introduced by commonly used hyperparameters, stochasticity introduced by random batches, and discretization introduced by a ﬁnite learning rate.
Revised: In the case of parameters with scale symmetry, the effect of discretization does distort the original geometry. 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.

Original: Here, we construct a more realistic continuous model for stochastic gradient descent. Modeling weight decay. 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.
Revised: Combining the geometric constraints on the gradient ( (cid:104) g, θ A (cid:105) = 0 ) and Hessian ( Hθ

Original: For stochastic gradient descent, the result leads to the updated continuous model
Revised: A = − g A ) introduced by scale symmetry with the equation of learning (equation 27) gives the following differential equation, [Equation] where we used the simpliﬁcations (cid:104) dθdt ,

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

Original: 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.
Revised: and ignored the O ( ηλ 2 ) term. The solution to this differential equation is [Equation]

Original: The method introduces two additional hyperparameters, a damping coefﬁcient ↵ and a momentum coefﬁcient  , and applies the two step update equation, ✓ (n +1)
Revised: Rescale dynamics. In the case of parameters with rescale symmetry, the effect of discretization also distorts the original geometry. However, unlike in the sphere, the force originating from discretization can both increase or decrease the previously conserved quantity ( | θ A 1 ( t ) | 2 − | θ A 2 ( t ) | 2 ). Combining the geometric properties of gradient (

Original: = ✓ ( n )  ⌘ v ( n +1) where v ( n +1) =
Revised: [Equation]

Original:  v ( n ) + (1  ↵ ) g
Revised: and Hessian [Equation] g

Original: ( ✓ ( n ) ) . When ↵ =  = 0 , we regain classic gradient descent.
Revised: A 2 = 0 ) introduced by rescale symmetry with the equation of learning (equation 27) gives the following differential equation, [Equation]

Original: 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. 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,
Revised: where we used the simpliﬁcation,

Original: [Equation]
Revised: [Equation]

Original: A more detailed derivation for this ODE can be found in appendix B.
Revised: and ignored the O ( ηλ 2 ) term. This is the same differential equation as in equation (29), just with a different forcing term. Thus, the solution to this differential equation is

Original: Modeling stochasticity.
Revised: [Equation] F M ODELING S TOCHASTICITY

Original: Stochastic gradients ˆ g B ( ✓ ) arise when we consider a batch B of size S drawn uniformly from the indices { 1 , . . .
Revised: Stochastic gradients ˆ g B ( θ ) arise when we consider a batch B of size S drawn uniformly from the indices { 1 , . . .

Original: , N } forming the unbiased gradient estimate ˆ g B ( ✓ ) =  S 
Revised: , N } forming the unbiased gradient estimate ˆ g B ( θ ) = 1 S (cid:80)

Original: P i 2 B r ` ( ✓ , x i ) .
Revised: i ∈B ∇ (cid:96) ( θ, x i ) .

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: When the batch size is much smaller than the size of the dataset, S (cid:28) 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 ( θ ) (cid:124) .

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: = ✓ ( n )  ⌘ g ( ✓ ( n ) )
Revised: = θ ( n ) − ηg ( θ ( n ) )

Original: + ⌘ p S G ( ✓ ) ⇠ , where ⇠ is a standard normal random variable.
Revised: + η √ 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[Equation]
Revised: This update is an Euler-Maruyama discretization with step size η of the stochastic differential equation

Original: [Equation]
Revised: [Equation]

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

Original: Many stochastic differential equation approximations for SGD similar to equation (15) have been derived  in  previous works.
Revised: Equation (31) has been derived as a model for SGD in many previous works Mandt et al.

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 ( θ ) (cid:124) , 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. (2015), and ⌃
Revised: C where C is some constant matrix Mandt et al. (2015), and Σ( θ )

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: As we showed in section 3, the gradient of the loss, regardless of the batch, is orthogonal to the generator vector ﬁeld@
Revised: As we showed in section 3, the gradient of the loss, regardless of the batch, is orthogonal to the generator vector ﬁeld

Original: ↵ associated with a symmetry.
Revised: ∂ α ψ | α = I associated with a symmetry.

Original: This implies the stochastic noise must also observe the same property, h 1 p S G ( ✓ ) ⇠ , @
Revised: This implies the stochastic noise must also observe the same property, (cid:104)− 1 √ S G ( θ ) ξ, ∂ α ψ | α

Original: ↵ i = 0 .
Revised: = I (cid:105) = 0 .

Original: In order for this relationship to hold for arbitrary noise 2 ⇠ , then G ( ✓ )  | @ ↵ = 0 .
Revised: In order for this relationship to hold for arbitrary noise 6 ξ , then G ( θ ) (cid:124) ∂ α ψ | α = I = 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 onto low rank subspaces.

Original: [Figure] Modeling discretization. The effect of discretization when modeling continuous dynamics is a well studied problem in the numerical analysis of partial differential equations. 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.
Revised: Scale Symmetry with Stochasticity.

Original: We present two methods based on modiﬁed equation analysis, which modify gradient ﬂow to account for the effect of discretization.
Revised: In the previous section, we performed our analysis using continuous-time model with deterministic gradients to facilitate calculations, and replaced them with stochastic batch gradients upon discretization when we evaluated the results empirically.

Original: Modiﬁed loss. 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 bygradient ﬂow. Li et al.; Feng et al.
Revised: Instead, we can also directly model the stochastic noise arising from batch gradient with an Itô process to perform analogous analysis in continuous-time as pointed out by Li et al.

Original: (2019) and
Revised: (2021) in a recent follow-up work.

Original: most recently Barrett & Dherin (2020), demonstrate that the gradient descent trajectory closely follows the steepest descent path of a modiﬁed loss function e
Revised: First, we assume that the time evolution of network parameters θ ( t ) during SGD training is described by the Itô process dθ

Original: L . The diver- gence between these trajectories fundamentally depends on the learning rate ⌘ and the curvature
Revised: = µdt + (cid:112)

Original: H . 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| . Thus, the modiﬁed loss is e L = L +
Revised: ηS G ( θ ) dW t , where

Original: [Equation]
Revised: [Equation]

Original: 4 |r L| 2 and the modiﬁed gradient ﬂow ODE is
Revised: Assuming a set of parameters A respect scale symmetry, then applying Itô’s lemma to the function | θ A ( t ) | 2 yields the Itô process,

Original: [Equation]
Revised: [Equation]

Original: See Fig. 4 for an illustrative example of this method applied to a quadratic loss in R 2 . Modiﬁed ﬂow. Rather than modifying gradient ﬂow with higher order “spatial” derivatives of the loss function, here we introduce higher order temporal derivatives.
Revised: Notice this is a deterministic ODE equivalent to the previously derived ODE with an additional forcing term accounting for the variance of the noise.

Original: 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 + ⌘ ) = ✓ t +1 , gives the equality  g ( ✓ t ) =
Revised: We can perform an analogous analysis for the case of translation and rescale symmetry.

Original: ✓ ( 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 ⌘ ! 
 we regain gradient ﬂow.-]
Revised: We can also consider the effect of stochasticity without the complexity of stochastic calculus by considering the dynamics in the [-limit as ⌘ ! 
 we regain gradient ﬂow.-]discrete setting.

Original: For small ⌘ , we obtain a modiﬁed version of gradient ﬂow with an additional second-order term, [Equation] 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.
Revised: As explained in appendix I, this is possible for the case without momentum, but becomes much more complicated once we consider momentum as well.

Original: C OMBINING S YMMETRY AND M ODIFIED GRADIENT F LOW TO D ERIVE E XACT L EARNING D YNAMICS 
Revised: G D ERIVING THE E XACT L EARNING D YNAMICS OF SGD WITH M OMENTUM

Original: As shown in section 4, each symmetry results in a conserved quantity under gradient ﬂow. We now study how weight decay, momentum, stochastic gradients, and ﬁnite learning rates all interact to break these conservation laws. Remarkably, even when using amore realistic continuous model for stochastic gradient descent, as discussed in section 5, we can derive exact learning dynamics for the previously conserved quantities.
Revised: We consider a continuous model of SGD with momentum incorporating weight decay (equation 14), momentum (equation 15), stochasticity (equation 31), and modiﬁed ﬂow (equation 17). As discussed in appendix C, we can model the effect of momentum by considering the forward Euler discretization θ k + − θ k η (1 − α ) and the backward Euler discretization − β θ k − θ k − 1 η (1 − α ) .

Original: To do this we (i) consider a realistic continuous model for
Revised: These terms introduce the numerical artifacts η (1 − α )

Original: 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. We ﬁrst consider the continuous model of SGD without momentum incorporating weight decay, stochasticity, and modiﬁed loss.
Revised: d 2 dt 2 θ and η (1 − α )β

Original: In this setting, the exact dynamics, as fully derived in appendix D, for the parameter combinations tied to the symmetries are,
Revised: d 2 dt 2 θ respectively, as explained by the modiﬁed ﬂow analysis in appendix D. Incorporating these elements with weight decay gives the equation of learning,

Original: [Equation]
Revised: [Equation]

Original: Notice how these equations are equivalent to the conservation laws derived in section 4 when ⌘ =  = 0 . 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.
Revised: Incorporating the effect of stochasticity into this equation of learning gives the Langevin equation,

Original: [Equation]
Revised: [Equation]

Original: [Table] +QUFiUH5erD/Fp8YZ4CiR5giNC/f3RE5ipSZxaDpjokdqsTYz/6v1Mx1d+zkTaaZB0PlDUcaxTvAsCzxgEqjmEwOESmZ2xXREJKHaJFYzIbiLX16GzkXDNXx/WW/elHFU0TE6QWfIRVeoie5QC7URRTl6Rq/ozXqyXqx362PeWrHKmUP0R9bnDwRFlN0=</latexit> T r a n s l a t i o n S ca l e
Revised: where W t is a standard Weiner process.

Original: R e s ca l e
Revised: Compared to the modiﬁed loss route (described in the previous section), it is much more natural and simple to account for the effect of stochasticity with modiﬁed ﬂow.

Original: Translation dynamics.
Revised: Translation dynamics.

Original: 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.
Revised: Combining the geometric properties of gradient ( (cid:104) g, 1 A (cid:105) = 0 ), Hessian ( H A = 0 ), and stochasticity ( G ( θ ) (cid:124) 1 A = 0 ) introduced by translation symmetry with the Langevin equation of learning (equation 32) gives the differential equation,

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

Original: This isdue to the lack of curvature in the gradient ﬁeld for these parameters (as shown in Fig. 2).
Revised: This is the differential equation for a harmonic oscillator with γ = 1 − β η (1 − α )(1+ β ) and ω = (cid:113) 2 λ η (1 − α )(1+ β ) .

Original: 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.
Revised: Assuming the initial condition ddt (cid:104) θ ( t ) , 1 A (cid:105) (cid:12)(cid:12) t =0 = 0 , then the general solution (as derived in appendix H) is

Original:  Scale dynamics.
Revised: [Equation] Scale dynamics.

Original: For parameters with scale symmetry, equation 19 implies that 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.
Revised: Combining the geometric constraints on the gradient ( (cid:104) g, θ A (cid:105) = 0 ), Hessian ( Hθ A = − g A ) [Equation]

Original: Compared to the translation dynamics, the scale dynamics do depend on the data through the gradient norms accumulated throughout training.
Revised: and stochasticity ( G ( θ A ) (cid:124) θ

Original: Without weight decay  = 0 , the ﬁrst term stays constant and the second term grows monotonically.
Revised: A = 0 ) introduced by scale symmetry with the Langevin equation of learning (equation 32) gives the following differential equation 7 ,

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

Original: To better understand the forces driving these complex dynamics, we can examine the time derivative of equation 19,
Revised: This is the differential equation for a driven harmonic oscillator with γ = 1 − β η [Equation] 1 − α )

Original: [Equation]
Revised: [Equation]

Original: 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 ).
Revised: 1+ β ) , ω = (cid:113) 4 λ η (1 − α )(1+ β ) , and f ( t ) = 2 (cid:12)(cid:12) dθ A dt (cid:12)(cid:12) 2 . Assuming the initial condition ddt | θ A | 2 (cid:12)(cid:12) t =0 = 0 , then the general solution (derived in appendix H) is [Equation]

Original: Thecentripetal 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. 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: The

Original: In particular, at stationarity, then ddt | ✓ ( t ) | 2 = 0 , which rearranging equation (21) gives the condition !
Revised: [Equation] [Equation] = 1 2 d 2 dt 2 | θ A | 2 − | dθ A dt | 2 .

Original:  ( t ) ⌘
Revised: where | θ A ( t ) | 2 h is the solution to the homogeneous harmonic oscillator [Equation] Rescale dynamics. Combining the geometric properties of gradient (cid:104) g, θ A 1 − θ A 2 (cid:105) = 0 , Hessian ( H ( θ

Revised: A 1 − θ A 2 ) + g A 1 − g

Revised: A 2 = 0 ), and stochasticity ( G ( θ A 1 ) (cid:124)

Revised: θ A 1 − G ( θ A 2 ) (cid:124)

Revised: θ

Revised: A 2 = 0 ) introduced by rescale symmetry with the Langevin equation of learning (equation 32) gives the differential equation,

Revised: [Equation]

Revised: dθ A 2 dt (cid:12)(cid:12)(cid:12)(cid:12) 2 (cid:33) .

Revised: This is the same harmonic oscillator given by equation (35) with the different forcing term f ( t ) = 2 (cid:16) | dθ A 1 dt | 2 − | dθ A 2 dt | 2 (cid:17) .

Revised: The general solution is given by equation (36) and (37) replacing | θ | 2 and | dθdt | 2 by | θ A 1 | 2 − | θ A 2 | 2 and | dθ A 1 dt | 2 − | dθ A 2 dt | 2 respectively.

Revised: H G ENERAL SOLUTIONS FOR ODE S H. E XPONENTIAL G ROWTH

Revised: Here we will solve for the general solution of the homogenous ﬁrst-order linear differential equation,

Revised: [Equation]

Revised: Assume a solution of the form x ( t ) =

Revised: e αt .

Revised: Plugging this in gives the auxiliary equation α + λ = 0 .

Revised: Thus, the general solution to the differential equation with initial condition x (0) is

Revised: [Equation]

Revised: Now we will solve the inhomogenous differential equation,

Revised: [Equation]

Revised: Multiply both sides by e λt and factor the left hand side using the product rule such that the differential equation simplies to ddt (cid:0) e λtx ( t ) (cid:1) =

Revised: e λt f ( t ) .

Revised: Integrate this equation and using the fundamental theorem of calculus rearrange to get the solution

Revised: [Equation]

Revised: H. H ARMONIC O SCILLATOR

Revised: Here we will solve the general solution for a harmonic oscillator,

Revised: [Equation]

Revised: ( t ) = 0 .

Revised: Assume a solution of the form x ( t ) =

Revised: e αt .

Revised: Plugging this in gives the auxiliary equation α 2 + 2 γα + ω 2 = 0 with solutions α ± = − γ ± (cid:112) γ 2 − ω 2 .

Revised: Thus, the general solution to the oscillator equation with initial conditions x (0) and dxdt (0)

Revised: = 0 is

Revised: [Equation]

Revised: where C 1 , C 2 are constants

Revised: [Equation]

Revised: Using hyperbolic functions the solution simpliﬁes as

Revised: [Equation]

Revised: The form of this general solution implicitly assumes γ > ω , the overdamped setting.

Revised: When γ = ω , the critically damped setting

Revised: [Equation]

Revised: then the solution reduces to

Revised: [Equation]

Revised: where C 1 = x (0) and C 2 = γx (0) .

Revised: When γ < ω , the underdamped setting, then the solution reduces to

Revised: [Equation]

Revised: where C 1 = x (0) and C 2 = γ √ ω 2 − γ 2 x (0) .

Revised: H.

Revised: D RIVEN H ARMONIC O SCILLATOR

Revised: Here we will solve the general solution for a driven harmonic oscillator,

Revised: [Equation]

Revised: First notice that if x h ( t ) is a solution to the homogenous harmonic oscillator and x d ( t ) a speciﬁc solution to the driven harmonic oscillator, then

Revised: [Equation]

Revised: t ) , is the general solution to the driven harmonic oscillator.

Revised: We will use the Fourier transform to solve for x d ( t ) .

Revised: Let ˆ x d ( τ ) =

Revised: ( ∇ (cid:101)

Revised: L x d )( τ ) and ˆ f ( τ ) =

Revised: ( ∇ (cid:101) L f )( τ ) be the Fourier transforms of x d ( t ) and f ( t ) respectively.

Revised: Applying the Fourier transform to the driven harmonic oscillator equation and rearranging gives

Revised: [Equation]

Revised: which implies by the inverse Fourier transform that x d ( t ) is the convolution,

Revised: [Equation]

Revised: where G ( t ) is Green’s function (the driven solution x d ( t ) for the dirac delta forcing function δ 0 ),

Revised: [Equation]

Revised: which again using hyperbolic functions simpliﬁes as

Revised: [Equation]

Revised: This form of Green’s function is again implicitly assuming γ > ω .

Revised: When γ = ω , the function simpliﬁes to

Original: [Equation]
Revised: [Equation]

Revised: and when γ < ω , the function simpliﬁes to

Original: [Equation]
Revised: [Equation]

Original: q 2 ⌘ .
Revised: Noticing that both G and f are only supported on [0 , ∞ ) , their convolution can be simpliﬁed and the general solution for the driven harmonic oscillator is [Equation] t ) [Equation]

Original: Consistent with the results of Wan et al.
Revised: In section 4 we identiﬁed certain parameter combinations associated with network symmetries that are conserved under gradient ﬂow. However, as we explained in section 5, these conservation laws are not observed empirically. To remedy this discrepancy we constructed more realistic continuous models for SGD, incorporating weight decay, momentum, stochasticity, and ﬁnite learning rates.

Original: (2020), this implies that at stationarity the angular speed !
Revised: In section 6 we derived the exact dynamics for the parameter combinations under this more realistic setting, demonstrating near perfect alignment with the empirical dynamics. What would happen if we instead derived the dynamics for the parameter combinations directly in the discrete setting of SGD?

Original: ( t ) of the weights is constant andgoverned only by the  learning rate ⌘ and weight decay constant  .
Revised: Here, we will identify these discrete dynamics and discuss the relationship between the discrete equations and the continuous solutions. Gradient descent with learning rate η and weight decay constant λ is given by the update equation θ ( n +1) = [Equation] 1 − ηλ ) θ

Original: Rescale dynamics.
Revised: [Equation]

Original: For parameters with rescale symmetry, equation20 is the sum ofan exponentially decaying memory of the difference in norms at initialization and an exponentially weighted integral of difference in gradient norms accumulated through training.
Revised: n ) − ηg ( θ ( n ) ) , and initial condition θ (0) . Using this update equation the sum of the parameters after n + 1 steps can be “unrolled” as, [Equation] Similarly, the squared Euclidean norm of the parameters after n + 1 steps can be “unrolled” as, | θ

Original: 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.
Revised: [Equation] n + [Equation]

Original: This leads to quite sophisticated, complex dynamics, consistent with the third row in Fig.
Revised: Combining these unrolled equations, with the gradient properties of symmetry discussed in section 3, gives the discrete dynamics for the parameter combinations, Translation: [Equation] Scale:

Revised: [Equation]

Revised: Rescale:

Original: Despite the complexity, our theory, nevertheless, quantitatively matches the empirics.
Revised: [Equation] [Table] Notice the striking similarity between the continuous equations (18), (19), (20) presented in sectionand the discrete equations (39), (40), (41). The exponential function with decay rate λ from the continuous solutions are replaced by a power of the base (1 − ηλ ) in the discrete setting.

Revised: The integral of exponentially weighted gradient norms from the continuous solutions are replaced by a

Revised: Riemann sum of power weighted gradient norms in the discrete setting.

Revised: This is further conﬁrmation that the continuous solutions we derived, and the modiﬁed gradient ﬂow equation of learning used, well approximate the actual empirics.

Original: 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.
Revised: While the equations derived in the discrete setting remove any uncertainty about the exactness of the theoretical predictions, they provide limited qualitative understanding for the empirical learning dynamics. This is especially true if we consider the learning dynamics with momentum. In this setting, the process of “unrolling” is much more complicated and the harmonic nature of the empirics, easily derived in the continuous setting, is hidden in the discrete algebra. J E XPERIMENTAL D ETAILS

Revised: An open source version of our code, used to generate all the ﬁgures in this paper, is available at github.com/danielkunin/neural-mechanics.

Revised: Dataset.

Revised: While we ran some initial experiments on Cifar-100, the dataset used in all the empirical ﬁgures in this documents was Tiny Imagenet.

Revised: It is used for image categorization an consists of 100,000 training images at a resolution of 64 × 64 spanning 200 classes.

Revised: Model.

Revised: We use standard VGG-16 models for all out experiments with the following modiﬁcations:

Revised: • The last three fully connected layers at the end have been adjusted for an input at the Tiny ImageNet resolution ( 64 × 64 ) and thus consist of 2048 , 2048 , and 200 layers respectively.

Revised: • In addition to the standard arrangement of conv layers for the VGG-16, we consider a variant where we add a batch normalization layer between every convolutional layer and its activation function.

Revised: Training hyperparameters.

Revised: Certain hyperparameters were varied during training.

Revised: Below we outline the combinations explored.

Revised: All models were initialized using Kaiming Normal, and no learning rate drops or warmup were used.

Revised: [Figure]

Revised: [Table]

Revised: Counting number of symmetries.

Revised: Here we explain how to count the number of symmetries a

Revised: VGG-16 model contains.

Revised: • Scale symmetries appear once per channel at every layer preceding a batch normalization layer.

Revised: For our VGG-16 model with batch norm: 2 · 64+2 · 128+3 · 256+3 · 512+3 · 512+3 = 4 , 227 . • Rescale symmetries appear once per channel where there are aﬁne transforms between layers as well as once per input neuron to a fully connected layer.

Revised: Note that the sizes of the fully connected layers depend on input image size and the number of classes.

Revised: For our VGG-16 model with and without batchnorm on Tiny ImageNet, this is: (2 · 64 + 2 · 128 + 3 · 256 + 3 · 512 + 3 · 512 + 3) + (2048 + 1024 + 1024) = 8 , 323 .

Revised: • Translation symmetries appear once per input value to the softmax function, which for the case of classiﬁcation is always equal to the number of classes, plus the bias term.

Revised: For our case this is 200 + 1 = 201 .

Revised: [Table]

Revised: Computing the theoretical predictions.

Revised: Some of the expressions shown in equations (18), (19), and (20) in section 6 are not trivial to compute as they involve an integral of an exponentially weighted gradient term.

Revised: To tackle this problem, we wrote custom optimizers in PyTorch with additional buffers to approximate the integral via a Riemann sum.

Revised: At every update step, the argument in the integral term was computed from the batch gradients, scaled appropriately, and was accumulated in the buffer.

Revised: Note that the above sum needs to be scaled by the learning rate, which is the coarseness of the grid of this Riemann sum.

Revised: Checkpoints of the model and optimizer states were stored at involve equations model states the hand side the equations from the integral buffers stored in the optimizer states as explained above.

Revised: These two quantities are referred to the and are depicted with

Revised: J. A DDITIONAL E MPIRICS

Revised: In this work we made exact predictions for the dynamics of combinations of parameters during training with SGD.

Revised: Importantly, these predictions were made at the neuron level, but could be aggregated for each layer.

Revised: Here we plot our predictions at both layer and neuron levels for VGGmodels trained on Tiny ImageNet.

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

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: +Dy8eFPHqf/HmvzFt96CtLwQe3plhJm8QC66N6347uZXVtfWN/GZha3tnd6+4f1DXUaIY1lgkItUMqEbBJdYMNwKbsUIaBgIbweh2Wm88otI8kg9mHKMf0oHkfc6osVanjYaSa+K5nfTMm3SLJbfszkSWwcugBJmq3eJXuxexJERpmKBatzw3Nn5KleFM4KTQTjTGlI3oAFsWJQ1R++ns6gk5sU6P9CNlnzRk5v6eSGmo9TgMbGdIzVAv1qbmf7VWYvpXfsplnBiUbL6onwhiIjKNgPS4QmbE2AJlittbCRtSRZmxQRVsCN7il5ehfl72LN9flCo3WRx5OIJjOAUPLqECd1CFGjBQ8Ayv8OY8OS/Ou/Mxb8052cwh/JHz+QN5Q5Eu</latexit>

Revised: [Table]

Revised: [Figure]

Revised: h W i , 1 i

Original: [Table] 
Revised: [Table] + [Table]

Revised: </latexit>  =<latexit sha

Revised: [Table]

Revised: ="

Revised: [Table]

Revised: =

Revised: [Table]

Revised:

Original: [Equation]
Revised: [Table]  = 0 . [Table] [Figure]

Original: classiﬁer
Revised: [Table]

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

Revised:

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

Original: +RIyQ1xq7Var1xxq+5cZBW8HCqQq9Erf/n9hGUxV8gkNabruSkGE6pRMMmnJT8zPKVsRAe8a1HRmJtgMr95Ss6s0ydRou1TSObu74kJjY0Zx6HtjCkOzXJtZv5X62YYXQcTodIMuWKLRVEmCSZkFgDpC80ZyrEFyrSwtxI2pJoytDGVbAje8pdXoXVR9Sw/XFbqt3kcRTiBUzgHD66gDvfQgCYwSOEZXuHNyZwX5935WLQWnHzmGP7I+fwB0LWQNw==</latexit>
Revised: //9korq2vrG+XNytb2zu5edf+gbVSmCW0RxZXuRNhQziRtWWY57aSaYhFx+hCNbqb+wxPVhil5b8cpDQUeSJYwgq2THnvcRWN85ferNb/uz4CWSVCQGhRo9qtfvViRTFBpCcfGdAM/tWGOtWWE00mllxmaYjLCA9p1VGJBTZjPFp6gE6fEKFHaPWnRTP09kWNhzFhELimwHZpFbyr+53Uzm1yGOZNpZqkk84+SjCOr0PR6FDNNieVjRzDRzO2KyBBrTKzrqOJKCBZPXibts3rg+N15rXFd1FGGIziGUwjgAhpwC01oAQEBz/AKb572Xrx372MeLXnFzCH8gff5AzGjj/4=</latexit> [Figure] [Table] [Table]

Original: Harmonic oscillation with momentum.
Revised:  = [Table] < [Table]

Original: We will now consider a continuous model of SGD with momentum.
Revised: /AKb572Xrx372MeLXnFzCH8gff5AzGjj/4=</latexit> <latexit sha1_base64="9bRYE12pYR/egHV9+vPHBDgRrtE=">AAAB8HicbVDLSgMxFL1TX7W+qi7dBIvgqsyIoBuh6MZlBVsr7VAymUwbmseQZIQy9CvcuFDErZ/jzr8xbWehrQcCh3POJfeeKOXMWN//9korq2vrG+XNytb2zu5edf+gbVSmCW0RxZXuRNhQziRtWWY57aSaYhFx+hCNbqb+wxPVhil5b8cpDQUeSJYwgq2THnvcRWN85ferNb/uz4CWSVCQGhRo9qtfvViRTFBpCcfGdAM/tWGOtWWE00mllxmaYjLCA9p1VGJBTZjPFp6gE6fEKFHaPWnRTP09kWNhzFhELimwHZpFbyr+53Uzm1yGOZNpZqkk84+SjCOr0PR6FDNNieVjRzDRzO2KyBBrTKzrqOJKCBZPXibts3rg+N15rXFd1FGGIziGUwjgAhpwC01oAQEBz/AKb572Xrx372MeLXnFzCH8gff5AzGjj/4=</latexit> [Table] [Figure] [Table]

Original: As discussed in appendix E, we consider the continuous model incorporating weight decay, momentum, stochasticity, and modiﬁed ﬂow.
Revised: [Figure] [Table] [Figure] [Table]

Original: 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.
Revised: [Figure] [Figure] [Table] [Equation]

Original: 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.
Revised: [Table] [Equation] [Table] [Equation]

Original: 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.
Revised: [Table] [Equation] [Table] [Figure]

Original: C ONCLUSION
Revised: [Table] W

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

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

Original: 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.
Revised: [Table] +fwB0LWQNw==</latexit> theory conv. 2-1 conv. 3conv. 4-3 conv.

Original: 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.
Revised: 5conv. 6-5 conv. 7conv. 8-7 conv. 9conv.

Original: 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.
Revised: 10-9 conv. We plot the per- difference between the squared Euclidean consecutive convolutional VGG 16 model (without

Original: Overall, this work provides the ﬁrst solid step towards the mechanics of learning in neural networks.
Revised: ) trained on Tiny with Momentum with