Law of Balance and Stationary Distribution of Stochastic Gradient
  Descent
Liu Ziyin Hongchao Li Masahito Ueda
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 prove that the minibatch noise of SGD regularizes the
solution towards a balanced solution whenever the loss function contains a
rescaling 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.