PaperlayerView exact revision on arXiv

Source-linked AI summary

Neural Mechanics: Symmetry and Broken Conservation Laws in Deep Learning Dynamics

Daniel Kunin, Javier Sagastuy-Brena, Surya Ganguli, Daniel L. K. Yamins, Hidenori Tanaka

arXiv:2012.04728v2cs.LGcond-mat.dis-nncond-mat.stat-mechq-bio.NCstat.ML

TL;DR

The paper addresses the challenge of understanding neural-network learning dynamics under realistic architectures and stochastic, finite-step SGD. It uses architectural symmetries and a more realistic continuous model to derive exact dynamics for selected parameter combinations, which it reports matching empirical VGG-16 results on Tiny ImageNet.

  • Problem

    The paper asks what can be quantitatively understood about learning dynamics in large nonlinear networks trained on real-world data with finite-batch, finite-rate SGD, since prior work often simplified architectures or focused on gradient flow.

    Source: §1 · p. 1

  • Method

    The framework uses loss symmetries to constrain gradients and Hessians, then combines them with a continuous SGD model incorporating weight decay, momentum, stochasticity, and finite-step effects.

    Source: §1 · p. 2, §5 · p. 6

  • Results

    The derived dynamics match nearly perfectly with empirical results for VGG-16 trained on Tiny ImageNet under typical settings including weight decay, stochastic batches, and finite learning rates.

    Source: §6 · p. 8

  • Takeaways & Limitations

    The framework derives exact dynamics for meaningful combinations of parameters and experimentally verifies them on large neural networks and datasets.

    Source: §7 · p. 9

Abstract

from arXiv · show

Understanding the dynamics of neural network parameters during training is one of the key challenges in building a theoretical foundation for deep learning. A central obstacle is that the motion of a network in high-dimensional parameter space undergoes discrete finite steps along complex stochastic gradients derived from real-world datasets. 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. 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. We further show that finite learning rates used in practice can actually break these symmetry induced conservation laws. We apply tools from finite difference methods to derive modified gradient flow, a differential equation that better approximates the numerical trajectory taken by SGD at finite learning rates. We combine modified gradient flow with our framework of symmetries to derive exact integral expressions for the dynamics of certain parameter combinations. We empirically validate our analytic expressions for learning dynamics on VGG-16 trained on Tiny ImageNet. Overall, by exploiting symmetry, our work demonstrates that we can analytically describe the learning dynamics of various parameter combinations at finite learning rates and batch sizes for state of the art architectures trained on any dataset.

1 INTRODUCTION

The paper asks what can be quantitatively understood about neural-network learning dynamics under realistic architectures and SGD, where prior analyses often simplify models or ignore stochasticity and discretization. It uses architectural symmetries to derive geometric constraints and conservation laws, then models realistic SGD to obtain exact dynamics for selected parameter combinations.

1 INTRODUCTION

  • Prior work often simplified architectures or focused on gradient flow, leaving stochasticity and discretization out of its learning-dynamics analysis.Examples include identity activations, infinite-width layers, and single-hidden-layer networks.Source: §1 · p. 1
  • The paper seeks meaningful, exactly solvable parameter combinations without imposing major assumptions on the model or learning dynamics.Its approach relies on symmetries embedded in network architectures.Source: §1 · pp. 1–2
  • The framework combines symmetry-derived gradient and Hessian constraints, gradient-flow conservation laws, and a realistic SGD model to derive exact dynamics for selected parameter combinations.The expressions are validated on VGG-16 trained on Tiny ImageNet, with and without batch normalization.Source: §1 · p. 2

2 RELATED WORK

Prior analyses of neural-network learning dynamics have gained tractability through assumptions on architecture or learning rules, while related work has studied geometric constraints and continuous models of SGD. This paper takes a complementary route: use architectural symmetries and more realistic continuous equations to solve dynamics for restricted parameter combinations.

2 RELATED WORK

  • Earlier exact-dynamics analyses used linear networks with specific initializations, single-hidden-layer nonlinear networks, or infinitely wide networks.These settings supported analyses of loss landscapes, convergence, implicit acceleration, and kernel-regression dynamics.Source: §2 · p. 2
  • The paper’s complementary approach finds exact solutions for restricted parameter combinations using loss geometry shaped by symmetry and realistic continuous learning equations.Figure 1 contrasts noisy, chaotic per-parameter dynamics with smooth, patterned per-channel squared-norm dynamics for VGG-16 on Tiny ImageNet.Source: §2 · p. 2, Figure 1 · p. 2
  • Prior work studied gradient constraints from architectural components, including normalization-related scale invariance and its role in effective learning rates and learning dynamics.Related geometric properties have also been applied to network optimization, interpretation, and pruning.Source: §2 · pp. 2–3
  • Prior continuous-time SGD studies often assumed particular loss or noise properties to derive stationary distributions, while later work examined steady-state fluctuation-dissipation relations without those assumptions.This paper instead uses finite-difference-inspired higher-order terms to model finite step size and trajectory curvature.Source: §2 · p. 3

3 SYMMETRIES IN THE LOSS SHAPE GRADIENT AND HESSIAN GEOMETRIES

Differentiable invariances of a network’s loss constrain its gradient and Hessian through the vector fields generated by parameter transformations. The section applies this framework to translation, scale, and rescale symmetries found in common neural-network components.

3 SYMMETRIES IN THE LOSS SHAPE GRADIENT AND HESSIAN GEOMETRIES

  • A differentiable symmetry is a parameter transformation that leaves the loss unchanged, and its generator is perpendicular to the gradient while also constraining the Hessian.The section considers translation, scale, and rescale symmetries in modern network architectures.Source: §3 · pp. 3–4
  • Translation symmetry shifts a parameter subset uniformly, forcing the gradient to be orthogonal to that subset’s indicator vector.For parameters immediately before softmax, shifting a weight-matrix column or bias by a constant leaves the output unchanged and yields a zero-sum bias-gradient constraint.Source: §3 · p. 4
  • Scale symmetry rescales a parameter subset without changing the loss, making the gradient perpendicular to the corresponding parameter vector and imposing a related Hessian constraint.Batch normalization gives scale invariance because scaling a neuron’s weights and bias cancels through its normalized input.Source: §3 · p. 4
  • Rescale symmetry inversely scales two disjoint parameter subsets, constraining the gradient relative to their sign-inverted parameter difference and relating the Hessian to their gradient difference.For homogeneous activations such as ReLU, scaling incoming parameters and inversely scaling outgoing parameters preserves the computational path.Source: §3 · pp. 4–5

4 SYMMETRY LEADS TO CONSERVATION LAWS UNDER GRADIENT FLOW

Differentiable symmetries of a network loss can impose conserved quantities on parameter dynamics under gradient flow. These quantities constrain trajectories to geometric surfaces, with distinct forms for translation, scale, and rescale symmetries.

4 SYMMETRY LEADS TO CONSERVATION LAWS UNDER GRADIENT FLOW

  • Every differentiable loss symmetry satisfying the stated generator-gradient condition has an associated conserved quantity under gradient flow.The proof projects the flow onto the symmetry generator, using the gradient's perpendicularity to that generator.Source: §4 · p. 5
  • Projecting gradient flow onto the symmetry generator yields a differential equation that simplifies to the conservation relation.The loss gradient is perpendicular to the generator vector field.Source: §4 · p. 5
  • Translation symmetry conserves the parameter sum, scale symmetry the Euclidean norm, and rescale symmetry the difference of squared norms.These quantities constrain the flow to a hyperplane, sphere, and hyperbola, respectively.Source: §4 · pp. 5–6, Figure 3 · p. 5

5 A REALISTIC CONTINUOUS MODEL FOR STOCHASTIC GRADIENT DESCENT

The section builds a more realistic continuous model of SGD by accounting for optimizer and batch effects as well as finite-step discretization. Modified dynamics are intended to better represent discrete training trajectories than ordinary gradient flow.

5 A REALISTIC CONTINUOUS MODEL FOR STOCHASTIC GRADIENT DESCENT

  • The model also incorporates explicit L2 weight decay, which is commonly implemented by modifying the optimizer update rather than the training loss.Source: §5 · p. 6
  • Momentum updates parameters using an exponentially moving average of gradients, with damping and momentum coefficients controlling the update.The damping coefficient effectively reduces the learning rate, while the momentum coefficient controls how past gradients contribute to later updates.Source: §5 · p. 6
  • For small batches, stochastic-gradient noise is modeled as Gaussian and its covariance has zero action along symmetry generators.The symmetry geometry confines this noise to low-rank subspaces, leaving gradient-flow dynamics in those symmetry directions unchanged.Source: §5 · p. 6
  • Finite-step discretization is modeled using modified equation analysis, which adjusts continuous dynamics with higher-order derivatives.The temporal-derivative approach recovers gradient flow as the learning rate approaches zero and adds a second-order term for small learning rates.Source: §5 · pp. 6–7
  • The modified loss adds a learning-rate-scaled squared-gradient-norm term to the original loss to model finite-rate gradient descent.The passage gives the modified loss as eL = L + η/4|∇L|².Source: §5 · p. 7
  • Modified continuous dynamics visually track discrete gradient descent and momentum trajectories better than the original continuous dynamics in the quadratic-loss illustration.The figure uses modified loss for gradient dynamics and modified flow for momentum dynamics.Source: §5 · p. 6, Figure 4 · p. 7

6 COMBINING SYMMETRY AND MODIFIED GRADIENT FLOW TO DERIVE EXACT LEARNING DYNAMICS

Combining symmetry geometry with realistic continuous SGD models yields exact dynamics for parameter combinations whose gradient-flow conservation laws are broken. The predicted dynamics match VGG-16 experiments, with distinct behaviors for translation, scale, rescale, and momentum settings.

6 COMBINING SYMMETRY AND MODIFIED GRADIENT FLOW TO DERIVE EXACT LEARNING DYNAMICS

  • The predicted dynamics match VGG-16 training on Tiny ImageNet nearly perfectly despite weight decay, stochastic batches, and finite learning rates.Figure 5 compares empirical trajectories with theoretical predictions for models with and without batch normalization.Source: §6 · p. 8, Figure 5 · p. 8
  • Translation-symmetric parameter sums decay exponentially at a weight-decay-controlled rate, independent of learning rate, dataset, or task.Higher weight decay produces faster convergence in the Figure 5 trajectories.Source: §6 · p. 8, Figure 5 · p. 8
  • Scale-symmetric parameter norms depend on accumulated gradient norms: they grow monotonically without weight decay but can become non-monotonic when weight decay is present.The dynamics reflect competition between weight decay’s inward effect and finite-step discretization’s outward effect.Source: §6 · p. 8
  • At stationarity, the scale-symmetry effects of weight decay and discretization balance, constraining parameters to a high-dimensional sphere.The balance is between the regularizing pull toward the origin and the outward effect of finite learning-rate updates.Source: §6 · p. 9
  • Rescale-symmetric norm differences combine decaying initialization memory with accumulated gradient-norm differences, producing complex dynamics that theory quantitatively matches.For sufficiently large weight decay, the norm difference decays exponentially to zero.Source: §6 · p. 9
  • With momentum, the parameter-combination dynamics follow driven harmonic-oscillator solutions; typical settings are overdamped, while large β yields empirically verified oscillations.Figure 6 compares empirical and theoretical column-sum trajectories across β values.Source: §6 · p. 9

7 CONCLUSION

The paper develops a symmetry-based framework for deriving parameter-combination dynamics without simplifying away realistic training features. It connects architectural symmetries to gradient and Hessian geometry, conservation laws, and continuous models of gradient descent and momentum.

7 CONCLUSION

  • The framework combines symmetry geometry with continuous learning models for weight decay, momentum, stochasticity, and discretization, and experimentally verifies exact dynamics for meaningful parameter combinations.For VGG-16 with batch normalization trained on Tiny ImageNet, the authors report 12,751 analytically describable parameter combinations.Source: §7 · p. 9

A SYMMETRY AND GEOMETRY

  • Architectural symmetries impose geometric constraints on gradients and Hessians; the framework yields fifteen equations for neural-network training losses.The constraints unify properties associated with translation, scale, and rescale symmetries.Source: A SYMMETRY AND GEOMETRY · p. 13, Table 1 · p. 13

A.1 GRADIENT GEOMETRY

  • At the identity transformation, the gradient is perpendicular to the vector field generating a differentiable symmetry.For neural-network losses, the paper summarizes gradient properties for translation, scale, and rescale symmetries.Source: §A.1 · p. 14, Table 2 · p. 14

A.2 HESSIAN GEOMETRY

  • Symmetry also constrains Hessians and links their geometry to gradient geometry, with properties summarized for translation, scale, and rescale symmetries.The paper notes that these symmetries are not exhaustive and gives broader rescale symmetry examples.Source: §A.2 · pp. 14–15

B CONSERVATION LAWS

  • Each differentiable loss symmetry satisfying the stated condition has an associated conservation law under gradient flow.The condition is stated for translation, scale, and rescale symmetries considered in the paper.Source: B CONSERVATION LAWS · p. 15

C LIMITING DIFFERENTIAL EQUATIONS FOR LEARNING RULES

  • Gradient descent and classical momentum arise as first-order discretizations of their respective limiting differential equations as the learning rate tends to zero.Momentum flow is represented using a first-order ODE and step size η(1 − α).Source: C LIMITING DIFFERENTIAL EQUATIONS FOR… · p. 16, §C.2 · p. 16

C.1 GRADIENT DESCENT

  • Gradient descent with step size η is a forward-Euler discretization of gradient flow.The paper identifies gradient flow as the limiting ODE for gradient descent.Source: §C.1 · p. 16

C.2 CLASSICAL MOMENTUM

  • Classical momentum’s discrete updates correspond to momentum flow, a first-order ODE derived using step size η(1 − α).The PyTorch implementation note specifies that the first momentum-buffer gradient is not damped.Source: §C.2 · p. 16

D MODIFIED EQUATION ANALYSIS

Finite learning rates make discrete gradient-descent trajectories diverge from continuous steepest descent, with the discrepancy depending on learning rate and loss curvature. Modified equation analysis models discretization artifacts to produce continuous dynamics that better approximate discrete steps.

D MODIFIED EQUATION ANALYSIS

  • Finite learning rates cause gradient descent to depart from the continuous steepest-descent path, and the divergence depends on learning rate and loss curvature.The paper poses the need for a continuous model that incorporates both factors.Source: D MODIFIED EQUATION ANALYSIS · p. 17
  • Modified equation analysis represents numerical artifacts as higher-order derivatives modifying the original differential equation, yielding continuous dynamics that better approximate discrete trajectories.The approach is drawn from finite-difference methods for numerical solutions of PDEs.Source: D MODIFIED EQUATION ANALYSIS · p. 17

D.1 MODIFIED LOSS

The paper uses modified equation analysis to adjust continuous learning dynamics for finite-step effects. It presents modified-loss and modified-flow approaches, including a second-order momentum model whose quadratic-loss solutions are damped harmonic oscillators.

D.1 MODIFIED LOSS

  • Modified equation analysis motivates corrections to gradient flow using higher-order terms in the learning rate, producing gradient flow on a modified loss.The first-order correction is g1 = −1/2 Hg, and the modified loss adds a term proportional to the squared gradient norm.Source: §D.1 · p. 17
  • For a positive-definite quadratic loss, the paper compares gradient descent, gradient flow, and modified gradient flow, whose trajectories are specified analytically.The modified loss adds an η-weighted quadratic form involving A^2.Source: §D.1 · p. 17

D.2 MODIFIED FLOW

  • Modified flow models SGD’s discrete updates by assuming a continuous trajectory through the steps and Taylor-expanding the update relation.As η approaches zero the model recovers gradient flow; for small η it adds a second-order temporal term.Source: §D.2 · p. 18
  • For quadratic loss, modified momentum flow is a second-order differential equation whose solutions are damped harmonic oscillators under the stated initial condition.The equation incorporates momentum and the finite learning rate.Source: §D.2 · p. 18

E DERIVING THE EXACT LEARNING DYNAMICS OF SGD

The section derives continuous learning dynamics for symmetry-related parameter combinations by combining symmetry-imposed gradient and Hessian constraints with a continuous SGD model. It distinguishes translation, scale, and rescale dynamics under finite learning rates and weight decay.

Continuous model

  • The continuous SGD model incorporates weight decay through a modified loss and its corresponding gradient and Hessian.The section also considers stochastic batch gradients and Hessians as replacements for their full-batch counterparts, while noting that integration then requires careful stochastic-calculus treatment.Source: E DERIVING THE EXACT LEARNING DYNAMICS… · p. 18

Translation dynamics

  • Translation symmetry leaves the associated constant of learning unchanged by discretization when combined with the gradient and Hessian constraints.The derivation ignores an O(ηλ^2) term, taking weight decay λ to be as small as learning rate η in practice.Source: E DERIVING THE EXACT LEARNING DYNAMICS… · pp. 18–19

Scale dynamics

  • Finite learning rates make the scale-symmetry norm experience a discretization-driven centrifugal force, while weight decay acts as a centripetal force.The section combines scale-symmetry gradient and Hessian constraints with the learning equation to derive its dynamics.Source: E DERIVING THE EXACT LEARNING DYNAMICS… · p. 19

Rescale dynamics

  • For rescale symmetry, discretization can either increase or decrease the previously conserved difference between squared parameter norms.The section derives a differential equation from the corresponding gradient and Hessian constraints.Source: E DERIVING THE EXACT LEARNING DYNAMICS… · p. 19

F MODELING STOCHASTICITY

The section models stochastic gradients and momentum in continuous learning equations, derives symmetry-specific dynamics, and relates these solutions to discrete SGD updates. It also gives solution methods for exponential growth and harmonic oscillators, including driven cases.

F MODELING STOCHASTICITY

  • Symmetry constrains stochastic-gradient noise to low-rank subspaces, and the resulting dynamics can be represented by a deterministic ODE with an additional variance-forcing term.The noise model assumes batches much smaller than the dataset and Gaussian gradient noise; the authors note this Gaussian assumption is not necessary for the stated symmetry property.Source: F MODELING STOCHASTICITY · p. 20

G DERIVING THE EXACT LEARNING DYNAMICS OF SGD WITH MOMENTUM

  • The momentum model incorporates weight decay, stochasticity, and modified flow; its translation dynamics reduce to a harmonic oscillator under the stated initial condition.The section describes modified flow as a more natural and simpler route for incorporating stochasticity than the modified-loss route.Source: G DERIVING THE EXACT LEARNING DYNAMICS… · p. 21

2. Assuming the initial condition

  • For scale symmetry with momentum, the squared norm follows a driven harmonic-oscillator solution whose form depends on whether damping γ is greater than, equal to, or less than frequency ω.The supplied solution gives separate expressions for the overdamped, critically damped, and underdamped cases.Source: G DERIVING THE EXACT LEARNING DYNAMICS… · p. 21, §2 · p. 22

H.1 EXPONENTIAL GROWTH

  • The exponential-growth appendix solves a first-order linear equation by using an exponential trial solution and then an integrating factor for the inhomogeneous case.The homogeneous solution uses α + λ = 0; the inhomogeneous derivation integrates after multiplying by e^(λt).Source: §H.1 · p. 22

H.2 HARMONIC OSCILLATOR

  • The harmonic-oscillator appendix derives solutions across overdamped, critically damped, and underdamped regimes from the characteristic equation.The displayed general form assumes γ > ω, with separate reductions given for γ = ω and γ < ω.Source: §H.2 · pp. 22–23

H.3 DRIVEN HARMONIC OSCILLATOR

  • The driven-oscillator appendix obtains the forced response by Fourier transform and expresses it as a convolution with Green’s function.The solution combines the homogeneous solution with the driven response, with the shown Green’s-function form assuming γ > ω.Source: §H.3 · p. 23

I DERIVING DYNAMICS IN THE DISCRETE SETTING

  • Unrolling discrete SGD yields exact parameter-combination dynamics that closely parallel the continuous equations, replacing exponential decay and integrals with powers and Riemann sums.The discrete equations make the predictions exact, but offer less qualitative insight; with momentum, unrolling is more complicated and obscures the harmonic behavior.Source: I DERIVING DYNAMICS IN THE DISCRETE… · p. 24

J EXPERIMENTAL DETAILS

The experiments use VGG-16 variants trained on Tiny ImageNet and report symmetry counts for the architectures. The theoretical predictions are computed by approximating integral terms with optimizer-buffer Riemann sums.

Dataset

  • All empirical figures use Tiny ImageNet, with 100,000 training images at 64 × 64 resolution across 200 classes.CIFAR-100 was used only in initial experiments.Source: J EXPERIMENTAL DETAILS · p. 25

Model

  • Experiments use standard VGG-16 with adjusted fully connected layers and also a variant adding batch normalization between convolutional layers and activations.For Tiny ImageNet, the adjusted fully connected layer sizes are 2048, 2048, and 200.Source: J EXPERIMENTAL DETAILS · p. 25

Counting number of symmetries

  • The VGG-16 with batch normalization has 4,227 scale symmetries, while the models with and without batch normalization have 8,323 rescale symmetries and 201 translation symmetries.The rescale count depends on input image size and number of classes; the translation count is 200 classes plus a bias term.Source: J EXPERIMENTAL DETAILS · p. 25

Computing the theoretical predictions

  • Theoretical predictions involving exponentially weighted gradient integrals are computed with custom PyTorch optimizers that approximate the integrals by Riemann sums in buffers.The sum is scaled by the learning rate, and figures compare quantities computed from model states with those from stored integral buffers.Source: J EXPERIMENTAL DETAILS · pp. 25–26

J.1 ADDITIONAL EMPIRICS

The additional empirical figures compare parameter combinations from VGG-16 trained on Tiny ImageNet with theoretical predictions. They show both layer-level and per-neuron plots across planar, spherical, and hyperbolic dynamics.

Planar dynamics

  • Planar dynamics compare empirical final-layer column sums with theoretical predictions for SGD and Momentum.The SGD figure uses batch size 256; the Momentum figure uses batch size 128 and varies momentum coefficient β across 0, 0.9, and 0.99.Source: Figure 8 · p. 26, Figure 11 · p. 27

Spherical dynamics

  • Spherical dynamics plot empirical squared norms and theoretical predictions for convolutional layers under SGD and Momentum.The SGD plot includes VGG-16 with batch normalization at batch size 256; the Momentum plot uses batch size 128 and includes batch normalization.Source: Figure 9 · p. 26, Figure 12 · p. 28

Hyperbolic dynamics

  • Hyperbolic dynamics compare empirical differences in consecutive convolutional-layer squared norms with theoretical predictions for SGD and Momentum.The Momentum figure uses batch size 128 and varies β across 0, 0.9, and 0.99; the SGD figure uses batch size 256.Source: Figure 10 · p. 27, Figure 13 · p. 28

Per-neuron dynamics

  • Per-neuron plots show squared norms and consecutive-layer norm differences against theoretical predictions for SGD and Momentum.The spherical plots use batch normalization, while the hyperbolic plots use models without batch normalization; these per-neuron figures use zero weight decay.Source: Figure 14 · p. 29, Figure 15 · p. 29, Figure 16 · p. 30

Related papers

Loading 2012.04728v2…