[Paper Review] The Limiting Dynamics of SGD: Modified Loss, Phase Space Oscillations, and Anomalous Diffusion

Daniel Kunin, Javier Sagastuy-Breña+5 more|arXiv (Cornell University)|Jul 19, 2021

Stochastic Gradient Optimization Techniques4 citations

TL;DR

This paper reveals that stochastic gradient descent (SGD) continues to drive deep neural networks through anomalous diffusion in parameter space long after performance convergence, due to a modified loss function and non-equilibrium probability currents arising from optimization hyperparameters and gradient noise. Using a continuous-time underdamped Langevin model, the authors derive exact dynamics in linear regression and validate predictions in ResNet-18 on ImageNet, showing precise quantitative agreement between theory and experiment for the anomalous diffusion exponent and instantaneous speed.

ABSTRACT

In this work we explore the limiting dynamics of deep neural networks trained with stochastic gradient descent (SGD). As observed previously, long after performance has converged, networks continue to move through parameter space by a process of anomalous diffusion in which distance travelled grows as a power law in the number of gradient updates with a nontrivial exponent. We reveal an intricate interaction between the hyperparameters of optimization, the structure in the gradient noise, and the Hessian matrix at the end of training that explains this anomalous diffusion. To build this understanding, we first derive a continuous-time model for SGD with finite learning rates and batch sizes as an underdamped Langevin equation. We study this equation in the setting of linear regression, where we can derive exact, analytic expressions for the phase space dynamics of the parameters and their instantaneous velocities from initialization to stationarity. Using the Fokker-Planck equation, we show that the key ingredient driving these dynamics is not the original training loss, but rather the combination of a modified loss, which implicitly regularizes the velocity, and probability currents, which cause oscillations in phase space. We identify qualitative and quantitative predictions of this theory in the dynamics of a ResNet-18 model trained on ImageNet. Through the lens of statistical physics, we uncover a mechanistic origin for the anomalous limiting dynamics of deep neural networks trained with SGD.

Motivation & Objective

• To explain the persistent, non-equilibrium parameter space dynamics of deep neural networks after training convergence, which manifest as anomalous diffusion.
• To identify the mechanistic origin of these dynamics by linking optimization hyperparameters, gradient noise structure, and the Hessian matrix at convergence.
• To develop a continuous-time model of SGD as an underdamped Langevin equation that captures finite learning rate and batch size effects.
• To derive analytically the stationary distribution of parameter dynamics in linear regression and show it corresponds to a modified loss with broken detailed balance.
• To validate theoretical predictions on real-world models like ResNet-18 on ImageNet, demonstrating quantitative agreement with empirical anomalous diffusion.

Proposed method

• Derive a continuous-time underdamped Langevin equation to model SGD with finite learning rates and batch-size-induced noise, treating the optimization process as a stochastic dynamical system.
• Analyze the dynamics in linear regression using the Fokker-Planck equation to show that the stationary distribution is a Gibbs measure over a modified loss function, not the original training loss.
• Prove that the modified loss breaks detailed balance, leading to non-zero probability currents that drive oscillatory motion in phase space (parameters and velocities).
• Express the phase space dynamics as a sum of damped harmonic oscillators in the eigenbasis of the data covariance matrix, enabling exact analytic solutions for moments.
• Use the Fokker-Planck equation to derive the stationary distribution and show that the system exhibits non-equilibrium steady-state behavior with persistent motion.
• Validate theoretical predictions on a ResNet-18 model fine-tuned on ImageNet, tracking parameter trajectories and fitting power laws to global displacement to extract the anomalous diffusion exponent.

Experimental results

Research questions

• RQ1 What drives the persistent, non-equilibrium motion of neural network parameters in parameter space after training convergence?
• RQ2 How do optimization hyperparameters—learning rate, batch size, and momentum—quantitatively influence the anomalous diffusion exponent and instantaneous speed in the limiting dynamics?
• RQ3 What is the role of the Hessian matrix and gradient noise structure in shaping the phase space dynamics of SGD?
• RQ4 Does the stationary distribution of parameter dynamics correspond to the original training loss, or is it governed by a modified, implicit regularization term?
• RQ5 Can the theoretical framework based on underdamped Langevin dynamics quantitatively predict empirical anomalous diffusion behavior in real models like ResNet-18 on ImageNet?

Key findings

• The limiting dynamics of SGD are governed not by the original training loss, but by a modified loss function that implicitly regularizes velocity, leading to non-equilibrium steady states.
• Non-zero probability currents in phase space—arising from broken detailed balance—drive oscillatory, non-diffusive motion, explaining the persistent parameter movement after convergence.
• In linear regression, the phase space dynamics are analytically solvable as a sum of damped harmonic oscillators in the eigenbasis of the data covariance matrix.
• The anomalous diffusion exponent and instantaneous speed in ResNet-18 trained on ImageNet are quantitatively predicted by the theoretical model with exact agreement to empirical measurements.
• Despite highly anisotropic original loss landscapes, the limiting dynamics of ResNet-18 exhibit isotropic behavior in parameter space, consistent with the model’s predictions.
• The fitted anomalous diffusion exponent increases with trajectory length, suggesting a sum of power laws with a dominant term emerging over time, consistent with the theoretical framework.