Logarithmic landscape and power-law escape rate of SGD

Takashi Mori, Liu Ziyin, Kangqiao Liu, Masahito Ueda

Keywords: stochastic gradient descent, noise structure, escape rate, flat minima, statistical physics

Abstract: Stochastic gradient descent (SGD) undergoes complicated multiplicative noise for the mean-square loss. We use this property of the SGD noise to derive a stochastic differential equation (SDE) with simpler additive noise by performing a random time change. In the SDE, the loss gradient is replaced by the logarithmized loss gradient. By using this formalism, we obtain the escape rate formula from a local minimum, which is determined not by the loss barrier height $\Delta L=L(\theta^s)-L(\theta^*)$ between a minimum $\theta^*$ and a saddle $\theta^s$ but by the logarithmized loss barrier height $\Delta\log L=\log[L(\theta^s)/L(\theta^*)]$. Our escape-rate formula strongly depends on the typical magnitude $h^*$ and the number $n$ of the outlier eigenvalues of the Hessian. This result explains an empirical fact that SGD prefers flat minima with low effective dimensions, which gives an insight into implicit biases of SGD.

One-sentence Summary: We have derived the Langevin equation in the logarithmized loss landscape for SGD with the mean-square loss, which yields a power-law escape rate from local minima.