A Diffusion Theory For Deep Learning Dynamics: Stochastic Gradient Descent Exponentially Favors Flat Minima
This work addresses a foundational problem in deep learning theory by providing a theoretical explanation for SGD's empirical success in finding generalizable solutions, which is significant for researchers and practitioners in machine learning.
The authors developed a density diffusion theory to quantitatively explain how Stochastic Gradient Descent (SGD) selects flat minima over sharp ones in deep learning, proving that SGD favors flat minima exponentially more than sharp minima, while Gradient Descent with white noise only does so polynomially.
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 often 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. To the best of our knowledge, we are the first to theoretically and empirically prove that, benefited from the Hessian-dependent covariance of stochastic gradient noise, 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 the batch size and learning rate. Thus, large-batch training cannot search flat minima efficiently in a realistic computational time.