[Paper Review] Rethinking 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 Techniques34 references4 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 convergence, with distance traveled scaling as a power law. Using a continuous-time underdamped Langevin model and Fokker-Planck analysis in linear regression, it identifies modified loss and probability currents—shaped by gradient noise and Hessian—as the key drivers of phase space oscillations and nontrivial diffusion exponents.

ABSTRACT

In this work we explore the limiting dynamics of deep neural networks trained with stochastic gradient descent (SGD). We find empirically that 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

• Understand the persistent, non-equilibrium dynamics in neural network parameters long after training loss has stabilized.
• Identify the mechanisms behind anomalous diffusion in parameter space during the limiting phase of SGD optimization.
• Reveal how hyperparameters, gradient noise structure, and the Hessian matrix jointly shape long-term dynamics.
• Establish a continuous-time model for finite-batch SGD that captures velocity-dependent dynamics and phase space oscillations.
• Validate theoretical predictions in a ResNet-18 trained on ImageNet, linking statistical physics to deep learning dynamics.

Proposed method

• Formulate a continuous-time underdamped Langevin equation to model finite learning rate and batch size SGD.
• Derive exact analytic solutions for parameter and velocity dynamics in linear regression using the Fokker-Planck equation.
• Introduce a modified loss function that implicitly regularizes velocity, distinct from the original training loss.
• Analyze probability currents in phase space to explain persistent oscillatory behavior in parameter trajectories.
• Characterize the interplay between gradient noise structure, Hessian curvature, and the resulting nontrivial diffusion exponent.
• Validate theoretical predictions by analyzing the phase space dynamics of a trained ResNet-18 on ImageNet.

Experimental results

Research questions

• RQ1 What drives the persistent, non-convergent movement of neural network parameters after training loss has stabilized?
• RQ2 How does the interplay between gradient noise, Hessian curvature, and optimization hyperparameters shape long-term parameter dynamics?
• RQ3 What is the role of velocity regularization and probability currents in generating oscillatory behavior in parameter-velocity phase space?
• RQ4 To what extent can the anomalous diffusion exponent in parameter space be explained by a modified loss and Fokker-Planck dynamics?
• RQ5 Can the theoretical framework derived in linear models be extended to explain dynamics in deep nonlinear networks like ResNet-18?

Key findings

• After convergence, deep neural networks continue to exhibit anomalous diffusion in parameter space, with distance traveled scaling as a power law in the number of updates.
• The diffusion exponent is nontrivial and determined by the interplay between the Hessian matrix, gradient noise structure, and optimization hyperparameters.
• The key driver of dynamics is not the original training loss, but a modified loss that implicitly regularizes velocity in parameter space.
• Probability currents in phase space induce persistent oscillations, preventing the system from reaching a simple equilibrium.
• Theoretical predictions from the linear regression model with underdamped Langevin dynamics are quantitatively validated in a ResNet-18 trained on ImageNet.
• The limiting dynamics are governed by a balance between velocity regularization and non-equilibrium currents, rooted in statistical physics principles.