A Diffusion Theory For Deep Learning Dynamics: Stochastic Gradient Descent Exponentially Favors Flat Minima

Stochastic Gradient Descent (SGD) and its variants are mainstream methods for training deep networks in practice. SGD is known to find a flat minimum that generalizes well. However, it is mathematically unclear how deep learning can select a flat minimum among so many minima. To answer the question quantitatively, we develop a density diffusion theory (DDT) to reveal how minima selection quantitatively depends on the minima sharpness and the hyperparameters. We empirically verify a key property of stochastic gradient noise (SGN) that the SGN covariance is approximately proportional to the Hessian and inverse to the batch size. To the best of our knowledge, we are the first to prove that, benefited from the Hessian-dependent structure of SGN, SGD favors flat minima exponentially more than sharp minima, while Gradient Descent (GD) with injected white noise favors flat minima only polynomially more than sharp minima. We also reveal that either a small learning rate or large-batch training requires exponentially many iterations to escape from minima in terms of the ratio of batch size and learning rate. Thus, large-batch training cannot search flat minima efficiently in a realistic computational time.