[Paper Review] The Heavy-Tail Phenomenon in SGD

Mert Gürbüzbalaban, Umut Şimşekli+1 more|arXiv (Cornell University)|Jun 8, 2020

Stochastic Gradient Optimization Techniques56 references38 citations

TL;DR

The paper proves that SGD iterates can converge to heavy-tailed stationary distributions in quadratic settings, with tail heaviness governed by stepsize, batch size, dimension, and curvature, and corroborates with experiments on neural networks.

ABSTRACT

In recent years, various notions of capacity and complexity have been proposed for characterizing the generalization properties of stochastic gradient descent (SGD) in deep learning. Some of the popular notions that correlate well with the performance on unseen data are (i) the `flatness' of the local minimum found by SGD, which is related to the eigenvalues of the Hessian, (ii) the ratio of the stepsize $η$ to the batch-size $b$, which essentially controls the magnitude of the stochastic gradient noise, and (iii) the `tail-index', which measures the heaviness of the tails of the network weights at convergence. In this paper, we argue that these three seemingly unrelated perspectives for generalization are deeply linked to each other. We claim that depending on the structure of the Hessian of the loss at the minimum, and the choices of the algorithm parameters $η$ and $b$, the SGD iterates will converge to a \emph{heavy-tailed} stationary distribution. We rigorously prove this claim in the setting of quadratic optimization: we show that even in a simple linear regression problem with independent and identically distributed data whose distribution has finite moments of all order, the iterates can be heavy-tailed with infinite variance. We further characterize the behavior of the tails with respect to algorithm parameters, the dimension, and the curvature. We then translate our results into insights about the behavior of SGD in deep learning. We support our theory with experiments conducted on synthetic data, fully connected, and convolutional neural networks.

Motivation & Objective

• Motivate how capacity and complexity notions relate to SGD generalization in deep learning.
• Show that SGD iterates can converge to heavy-tailed stationary distributions under certain algorithmic and problem settings.
• Characterize how tail-heaviness depends on stepsize, batch size, dimension, and curvature.
• Provide rigorous results in a linear/quadratic setting and connect them to deep learning observations.
• Support theory with experiments on synthetic data and neural networks.

Proposed method

• Model SGD as an iterated random recursion x_k = Psi_Omega_k(x_{k-1}).
• Approximate SGD near a quadratic minimum by an affine recursion x_k ≈ (I - (eta/b) H_k) x_{k-1} + q_k.
• Apply implicit renewal theory and stochastic matrix recursions to derive the tail-index alpha via h(alpha)=1.
• Show that under Gaussian input the tail-index increases with batch size and decreases with stepsize and variance.
• Establish three regimes for stepsize and convergence based on alpha relative to 2.
• Provide non-asymptotic moment bounds and Wasserstein distance convergence results.

Experimental results

Research questions

• RQ1 Does SGD converge to a stationary distribution with heavy tails under standard quadratic/linear regression settings?
• RQ2 How do stepsize, batch size, dimension, and curvature affect the tail-index of the SGD stationary distribution?
• RQ3 Can we explicitly relate tail heaviness to algorithmic parameters in Gaussian and non-Gaussian data models?
• RQ4 What are the implications of heavy tails for convergence rates and generalization in deep learning?
• RQ5 Do experimental results in neural networks corroborate the theoretical heavy-tail behavior?

Key findings

• SGD iterates can have a heavy-tailed stationary distribution with infinite variance in quadratic/linear regression even with light-tailed data.
• There exists a unique alpha > 0 such that h(alpha) = 1, determining the tail decay; u^T x_infty has a polynomial tail with exponent alpha.
• For Gaussian input, tail heaviness increases with curvature and eta/b ratio, and decreases with batch size b.
• The stationary distribution converges to equilibrium in Wasserstein distance at an exponential rate when rho < 0.
• Three regimes are identified: finite variance (alpha > 2), heavy tails (alpha < 2), and potential divergence (rho >= 0) depending on eta and b.
• When alpha <= 1, certain moments of x_k remain finite; when alpha > 1, higher moments are controlled with explicit bounds.