Revisiting the Noise Model of Stochastic Gradient Descent
This work addresses the fundamental understanding of SGD noise for machine learning researchers, providing insights into optimization dynamics and potential improvements in training deep neural networks, though it is incremental in refining existing noise models.
The paper tackles the problem of accurately modeling stochastic gradient noise (SGN) in stochastic gradient descent (SGD), showing through empirical evidence that SGN is heavy-tailed and better described by the SαS Lévy distribution rather than Gaussian, and that different parameters in deep neural networks exhibit distinct SGN characteristics. It results in a novel framework using Lévy-driven stochastic differential equations to approximate SGD dynamics near local minima, demonstrating that learning rate decay reduces SGN and affects properties like mean escape time and trapping probability.
The stochastic gradient noise (SGN) is a significant factor in the success of stochastic gradient descent (SGD). Following the central limit theorem, SGN was initially modeled as Gaussian, and lately, it has been suggested that stochastic gradient noise is better characterized using $SαS$ Lévy distribution. This claim was allegedly refuted and rebounded to the previously suggested Gaussian noise model. This paper presents solid, detailed empirical evidence that SGN is heavy-tailed and better depicted by the $SαS$ distribution. Furthermore, we argue that different parameters in a deep neural network (DNN) hold distinct SGN characteristics throughout training. To more accurately approximate the dynamics of SGD near a local minimum, we construct a novel framework in $\mathbb{R}^N$, based on Lévy-driven stochastic differential equation (SDE), where one-dimensional Lévy processes model each parameter in the DNN. Next, we show that SGN jump intensity (frequency and amplitude) depends on the learning rate decay mechanism (LRdecay); furthermore, we demonstrate empirically that the LRdecay effect may stem from the reduction of the SGN and not the decrease in the step size. Based on our analysis, we examine the mean escape time, trapping probability, and more properties of DNNs near local minima. Finally, we prove that the training process will likely exit from the basin in the direction of parameters with heavier tail SGN. We will share our code for reproducibility.