SGD as Equilibrium Thermodynamics: A Large Deviations Analysis of the Long-Run Distribution

TL;DR

Why does Stochastic Gradient Descent (SGD) find "good" minima in the messy, high-dimensional landscapes of deep learning? This paper provides a definitive answer by treating SGD iterates through the lens of Large Deviation Theory. It proves that SGD behaves like a physical system reaching a Boltzmann-Gibbs equilibrium, where the step-size is the temperature and the final destination depends heavily on the "flatness" of the landscape and the statistics of the noise.

The Problem: The Mystery of the "Long Run"

Standard optimization theory tells us SGD converges to a region with small gradients. However, non-convex landscapes are littered with local minima, maxima, and an exponential number of saddle points. We know SGD "prefers" flat minima, but the mathematical "why" has always been slippery.

Most researchers use Diffusion Approximations (SDEs), but these have a fatal flaw: the time it takes for an SDE to reach its steady state grows exponentially as the step-size $\eta$ decreases. By the time the SDE is relevant, the discretization error between the actual SGD and the SDE has usually exploded.

Methodology: Action Functionals and Quasi-Potentials

Instead of relying on SDEs, the authors implement the Freidlin-Wentzell Theory for discrete-time Markov processes.

1. The Least Action Principle

The authors define an "Action Functional" ${\mathcal{S}}_T[\gamma ]$. Think of this as the "cost" of SGD following a specific trajectory $\gamma$.

• Paths that follow the gradient flow $-ablaf$ have zero cost.
• Paths that fight the flow (to escape a local minimum) have high costs.

2. The Quasi-Potential

The probability of jumping from one local minimum $K_i$ to another $K_j$ is determined by the Quasi-Potential $B_{ij}$, which is the minimum "action" needed to traverse between them.

Figure 1: On the left, the loss landscape. In the middle, the "most likely" transitions between components. On the right, the empirical long-run distribution of SGD matches the theoretical wireframe prediction.

Core Insights: The Hierarchy of Convergence

The paper establishes four fundamental theorems:

1. Exponential Concentration: Critical regions are visited exponentially more often than anywhere else.
2. Boltzmann-Gibbs Analogy: The probability of being at a component $K_i$ is ${\mu }_{\infty }(U_i)\propto \exp (-E_i/\eta )$.
3. The Dominance of Minima: Saddle points and maxima are "exponentially neglected" compared to local minima.
4. Ground State Selection: SGD settles at the global minimum of the energy functional, which might not be the global minimum of the loss function.

Impact of Noise and Flatness

A profound takeaway (Section 4.2) is the role of Flatness. Under certain noise models, the energy levels are determined more by the reciprocal of the Hessian eigenvalues (flatness) than by the objective value $f(x)$ (depth).

$E_j\ge 2\min \left\{\sum \frac{g^{\lambda }(x)-g_i^{\lambda }}{\lambda {\sigma }^2}\right\}$

As a minimum becomes "flatter" (smaller eigenvalues $\lambda$), its energy level effectively decreases, making it exponentially more likely for SGD to settle there even if it’s a "shallow" minimum relative to other sharp peaks.

Critical Analysis & Conclusion

Takeaways

• Noise is a Feature: The noise in SGD isn't just something to be "tamed"; it's the engine that drives the system toward stable, flat, and likely more generalizable "ground states."
• Temperature Control: The step-size $\eta$ is a literal temperature. Decaying $\eta$ is analogous to Simulated Annealing, locking the model into its current energy state.

Limitations

The analysis assumes a constant step-size, whereas modern training utilizes complex schedules and vanishing step-sizes. Furthermore, while the theory is general, calculating the exact "Energy Levels" for a billion-parameter LLM remains computationally intractable.

Final Outlook

This work bridges the gap between optimization and statistical physics. It moves us away from asking "What is the gradient?" toward "What is the action cost of escaping this basin?"—a paradigm shift that may eventually explain why deep learning works at all.