[Paper Review] A Tail-Index Analysis of Stochastic Gradient Noise in Deep Neural Networks

Umut Şimşekli, Levent Sagun+1 more|arXiv (Cornell University)|Jan 17, 2019

Gaussian Processes and Bayesian Inference57 references69 citations

TL;DR

The paper shows stochastic gradient noise in deep nets is heavy-tailed (alpha-stable) and analyzes SGD as a Levy-driven SDE, with experiments confirming non-Gaussian tails and two SGD phases.

ABSTRACT

The gradient noise (GN) in the stochastic gradient descent (SGD) algorithm is often considered to be Gaussian in the large data regime by assuming that the classical central limit theorem (CLT) kicks in. This assumption is often made for mathematical convenience, since it enables SGD to be analyzed as a stochastic differential equation (SDE) driven by a Brownian motion. We argue that the Gaussianity assumption might fail to hold in deep learning settings and hence render the Brownian motion-based analyses inappropriate. Inspired by non-Gaussian natural phenomena, we consider the GN in a more general context and invoke the generalized CLT (GCLT), which suggests that the GN converges to a heavy-tailed $\\alpha$-stable random variable. Accordingly, we propose to analyze SGD as an SDE driven by a L\\'{e}vy motion. Such SDEs can incur `jumps', which force the SDE transition from narrow minima to wider minima, as proven by existing metastability theory. To validate the $\\alpha$-stable assumption, we conduct extensive experiments on common deep learning architectures and show that in all settings, the GN is highly non-Gaussian and admits heavy-tails. We further investigate the tail behavior in varying network architectures and sizes, loss functions, and datasets. Our results open up a different perspective and shed more light on the belief that SGD prefers wide minima.

Motivation & Objective

• Question the Gaussian noise assumption in SGD and its CLT-based SDE analyses.
• Propose and validate an alpha-stable (heavy-tailed) model for stochastic gradient noise.
• Link tail behavior to SGD dynamics and the tendency to find wide minima through metastability theory.
• Empirically characterize how tail index alpha varies with architecture, dataset, and batch size.

Proposed method

• Adopt alpha-stable (SalphaS) noise model for stochastic gradient noise with tail-index alpha.
• Derive a Lévy-driven SDE as the continuous-time limit of SGD when alpha < 2.
• Use a tail-index estimator designed for alpha-stable distributions to estimate alpha from gradient noise samples.
• Conduct extensive experiments across FCN and CNN architectures on MNIST, CIFAR-10, CIFAR-100 with varying depth, width, and minibatch sizes.
• Analyze metastability and first-exit behavior under Levy noise, highlighting jumps and two SGD phases.

Experimental results

Research questions

• RQ1 Is the stochastic gradient noise in deep networks alpha-stable (heavy-tailed) rather than Gaussian?
• RQ2 How does the tail index alpha vary with network size, architecture, dataset, and minibatch size?
• RQ3 What are the implications of alpha-stable noise for SGD dynamics, metastability, and the preference for wide minima?
• RQ4 Do early-iteration dynamics exhibit jumps in alpha that correlate with accuracy improvements?

Key findings

• Stochastic gradient noise is highly non-Gaussian with heavy tails across configurations.
• Increasing minibatch size has little effect on the tail-index alpha.
• Tail-index alpha is influenced by architecture, dataset, and network size, affecting SGD dynamics.
• Two distinct SGD phases observed: alpha rapidly decreases early, followed by a jump and subsequent stabilization of alpha as accuracy improves.
• Two-phase behavior supports metastability theory: jumps occur when alpha is at its lowest value.
• For CIFAR datasets, alpha values are in the range 1.0–1.2 in many configurations, indicating heavy tails.