Noise Balance and Stationary Distribution of Stochastic Gradient Descent

TL;DR

This paper investigates how stochastic gradient descent (SGD) navigates neural network loss landscapes, revealing that noise regularizes solutions towards symmetry-related solutions and that deep networks exhibit unique nonlinear phenomena in their stationary distributions.

Contribution

The paper introduces a theory linking SGD noise to loss function symmetries and derives the stationary distribution for deep linear networks, highlighting phenomena unique to deep models.

Findings

SGD noise regularizes towards noise-balanced solutions in symmetric loss functions

Stationary distribution of deep linear networks shows phase transitions and broken ergodicity

Deep networks exhibit nonlinear phenomena not present in shallow models

Abstract

The stochastic gradient descent (SGD) algorithm is the algorithm we use to train neural networks. However, it remains poorly understood how the SGD navigates the highly nonlinear and degenerate loss landscape of a neural network. In this work, we show that the minibatch noise of SGD regularizes the solution towards a noise-balanced solution whenever the loss function contains a rescaling parameter symmetry. Because the difference between a simple diffusion process and SGD dynamics is the most significant when symmetries are present, our theory implies that the loss function symmetries constitute an essential probe of how SGD works. We then apply this result to derive the stationary distribution of stochastic gradient flow for a diagonal linear network with arbitrary depth and width. The stationary distribution exhibits complicated nonlinear phenomena such as phase transitions, broken ergodicity, and fluctuation inversion. These phenomena are shown to exist uniquely in deep networks, implying a fundamental difference between deep and shallow models.