Stochastic gradient descent performs variational inference, converges to limit cycles for deep networks

TL;DR

This paper establishes that stochastic gradient descent implicitly performs variational inference by minimizing a steady-state free-energy F(ρ) composed of an average potential Φ(x) and an entropic term, where Φ generally differs from the original loss f unless the gradient noise is isotropic. It shows that deep networks exhibit highly non-isotropic, low-rank diffusion, leading to out-of-equilibrium dynamics where SGD traces deterministic loops or limit cycles rather than converging to traditional critical points. Empirical results across MNIST and CIFAR demonstrate the pronounced anisotropy of the diffusion matrix D(x) and the consequent misalignment between Φ and f, with implications for generalization and architecture search. The work connects continuous-time SGD analysis to Wasserstein gradient flows and Bayesian inference, clarifying when SGD behaves like variational optimization and highlighting how non-equilibrium behavior can enhance generalization in high-dimensional models.

Abstract

Stochastic gradient descent (SGD) is widely believed to perform implicit regularization when used to train deep neural networks, but the precise manner in which this occurs has thus far been elusive. We prove that SGD minimizes an average potential over the posterior distribution of weights along with an entropic regularization term. This potential is however not the original loss function in general. So SGD does perform variational inference, but for a different loss than the one used to compute the gradients. Even more surprisingly, SGD does not even converge in the classical sense: we show that the most likely trajectories of SGD for deep networks do not behave like Brownian motion around critical points. Instead, they resemble closed loops with deterministic components. We prove that such "out-of-equilibrium" behavior is a consequence of highly non-isotropic gradient noise in SGD; the covariance matrix of mini-batch gradients for deep networks has a rank as small as 1% of its dimension. We provide extensive empirical validation of these claims, proven in the appendix.

Figures (4)

• Figure 1: Eigenspectrum of $D(x)$ at three instants during training ($20\%$, $40\%$ and $100\%$ completion, darker is later). The eigenspectrum in \ref{['fig:smallfc_D']} for the fully-connected network has a much smaller rank and much larger variance than the one in \ref{['fig:lenets_D']} which also performs better on MNIST. This indicates that convolutional networks are better conditioned than fully-connected networks in terms of $D(x)$.
• Figure 2: Eigenspectrum of $D(x)$ at three instants during training ($20\%$, $40\%$ and $100\%$ completion, darker is later). The eigenvalues are much larger in magnitude here than those of MNIST in \ref{['fig:mnist_D']}, this suggests a larger gradient diversity for CIFAR-10 and CIFAR-100. The diffusion matrix for CIFAR-100 in \ref{['fig:allcnns_cifar100_D']} has larger eigenvalues and is more non-isotropic and has a much larger rank than that of \ref{['fig:allcnns_cifar10_D']}; this suggests that gradient diversity increases with the number of classes. As \ref{['fig:allcnns_cifar10_D']} and \ref{['fig:allcnns_cifar10_augment_D']} show, augmenting input data increases both the mean and the variance of the eigenvalues while keeping the rank almost constant.
• Figure 3: \ref{['fig:fft_fcnet']} shows the Fast Fourier Transform (FFT) of $x^i_{k+1} - x^i_k$ where $k$ is the number of epochs and $i$ denotes the index of the weight. \ref{['fig:ac_fcnet']} shows the auto-correlation of $x^i_k$ with $99\%$ confidence bands denoted by the dotted red lines. Both \ref{['fig:fft_fcnet', 'fig:ac_fcnet']} show the mean and one standard-deviation over the weight index $i$; the standard deviation is very small which indicates that all the weights have a very similar frequency spectrum. \ref{['fig:fft_fcnet', 'fig:ac_fcnet']} should be compared with the FFT of white noise which should be flat and the auto-correlation of Brownian motion which quickly decays to zero, respectively. \ref{['fig:fft_fcnet', 'fig:fft_ac']} therefore show that trajectories of SGD are not simply Brownian motion. Moreover the gradient at these locations is quite large (\ref{['fig:fgrad_fcnet']}).
• Figure 4: Gradient field for the dynamics in \ref{['eg:double_well']}: line-width is proportional to the magnitude of the gradient $\lVert\nabla f(x)\rVert$, red dots denote the most likely locations of the steady-state $e^{-\Phi}$ while the potential $\Phi$ is plotted as a contour map. The critical points of $f(x)$ and $\Phi(x)$ are the same in \ref{['fig:double_well1']}, namely $(\pm 1, 0)$, because the force $j(x) = 0$. For $\lambda=0.5$ in \ref{['fig:double_well2']}, locations where $\nabla f(x) = 0$ have shifted slightly as predicted by \ref{['thm:most_likely_locations']}. The force field also has a distinctive rotation component, see \ref{['rem:rotation_force_field']}. In \ref{['fig:double_well3']} with a large $\lVert j(x)\rVert$, SGD converges to limit cycles around the saddle point at the origin. This is highly surprising and demonstrates that the solutions obtained by SGD may be very different from local minima.

Theorems & Definitions (24)

• Definition 1: Diffusion matrix $D(x)$
• Lemma 2: Continuous-time SGD
• Theorem 5: SGD performs variational inference
• Lemma 6: Potential equals original loss iff isotropic diffusion
• Lemma 7: Most likely trajectories of SGD are limit cycles
• Corollary 8: Wasserstein gradient flow for isotropic noise
• Remark 9: SGD has an information bottleneck
• Remark 10: ELBO prior conflicts with SGD
• Remark 11: Learning rate should scale linearly with batch-size to generalize well
• Remark 12: Sampling with replacement is better than without replacement