Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Global Dynamics of Heavy-Tailed SGDs in Nonconvex Loss Landscape: Characterization and Control

Xingyu Wang, Chang-Han Rhee

TL;DR

The paper addresses why SGD generalizes well in nonconvex, multimodal loss landscapes by modeling heavy-tailed gradient noise combined with gradient clipping. It develops a metastability-based framework showing that, under truncation, the global dynamics of SGD converge to a continuous-time Markov chain that visits only the widest minima, with exit times governed by the width of attraction fields and the tail index. This leads to a principled algorithmic approach: inflate heavy tails and apply gradient clipping to steer training toward flat minima, supported by a rigorous scaling theory and first-exit analyses. Empirical results—ranging from 1D simulations to DL experiments with LeNet, VGG11, and Wide-ResNet architectures—confirm that the truncated heavy-tailed strategy yields flatter minima and better test generalization, even when integrated with adaptive optimizers like Adam. The work provides both a theoretical lens on SGD’s global dynamics and practical guidelines for optimization in deep learning.

Abstract

Stochastic gradient descent (SGD) and its variants enable modern artificial intelligence. However, theoretical understanding lags far behind their empirical success. It is widely believed that SGD has a curious ability to avoid sharp local minima in the loss landscape, which are associated with poor generalization. To unravel this mystery and further enhance such capability of SGDs, it is imperative to go beyond the traditional local convergence analysis and obtain a comprehensive understanding of SGDs' global dynamics. In this paper, we develop a set of technical machinery based on the recent large deviations and metastability analysis in Wang and Rhee (2023) and obtain sharp characterization of the global dynamics of heavy-tailed SGDs. In particular, we reveal a fascinating phenomenon in deep learning: by injecting and then truncating heavy-tailed noises during the training phase, SGD can almost completely avoid sharp minima and achieve better generalization performance for the test data. Simulation and deep learning experiments confirm our theoretical prediction that heavy-tailed SGD with gradient clipping finds local minima with a more flat geometry and achieves better generalization performance.

Global Dynamics of Heavy-Tailed SGDs in Nonconvex Loss Landscape: Characterization and Control

TL;DR

The paper addresses why SGD generalizes well in nonconvex, multimodal loss landscapes by modeling heavy-tailed gradient noise combined with gradient clipping. It develops a metastability-based framework showing that, under truncation, the global dynamics of SGD converge to a continuous-time Markov chain that visits only the widest minima, with exit times governed by the width of attraction fields and the tail index. This leads to a principled algorithmic approach: inflate heavy tails and apply gradient clipping to steer training toward flat minima, supported by a rigorous scaling theory and first-exit analyses. Empirical results—ranging from 1D simulations to DL experiments with LeNet, VGG11, and Wide-ResNet architectures—confirm that the truncated heavy-tailed strategy yields flatter minima and better test generalization, even when integrated with adaptive optimizers like Adam. The work provides both a theoretical lens on SGD’s global dynamics and practical guidelines for optimization in deep learning.

Abstract

Stochastic gradient descent (SGD) and its variants enable modern artificial intelligence. However, theoretical understanding lags far behind their empirical success. It is widely believed that SGD has a curious ability to avoid sharp local minima in the loss landscape, which are associated with poor generalization. To unravel this mystery and further enhance such capability of SGDs, it is imperative to go beyond the traditional local convergence analysis and obtain a comprehensive understanding of SGDs' global dynamics. In this paper, we develop a set of technical machinery based on the recent large deviations and metastability analysis in Wang and Rhee (2023) and obtain sharp characterization of the global dynamics of heavy-tailed SGDs. In particular, we reveal a fascinating phenomenon in deep learning: by injecting and then truncating heavy-tailed noises during the training phase, SGD can almost completely avoid sharp minima and achieve better generalization performance for the test data. Simulation and deep learning experiments confirm our theoretical prediction that heavy-tailed SGD with gradient clipping finds local minima with a more flat geometry and achieves better generalization performance.
Paper Structure (19 sections, 21 theorems, 207 equations, 2 figures, 4 tables)

This paper contains 19 sections, 21 theorems, 207 equations, 2 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Overview of the paper
  3. Comparison to Related Works
  4. Notations and Problem Settings
  5. Main Results
  6. Characterization of Global Dynamics of Heavy-Tailed SGD
  7. Control of Global Dynamics of Heavy-Tailed SGD
  8. Experiments
  9. Simulation Experiments in $\mathbb{R}^1$
  10. Deep Learning Experiment 1: An Ablation Study
  11. Deep Learning Experiment 2: Adam + Wide Residual Networks
  12. Metastability in Reducible Cases
  13. First Exit Analyses and Related Lemmas
  14. Sample Path Convergence to Jump Processes
  15. Proof of Theorem \ref{['corollary irreducible case']} and Corollary \ref{['theorem: metastability, unclipped']}
  16. ...and 4 more sections

Key Result

Theorem 2

Let Assumptions assumption: geometry, metastability--assumption: regularity condition on b hold. Let $p \in [1,\infty)$, $i_0 \in [K]$, and $\bm x_0 \in I_{i_0}$. As $\eta \downarrow 0$, where $\bm Y^{*|b}_t$ is a continuous-time Markov chain with state space $V^*_b$ (see def: V * b, metastability).

Figures (2)

  • Figure 1.1: (Left) Histograms of the locations visited by SGD when driven by different noise and exploring the multimodal function $f$ plotted in part (e), with the dashed lines indicating the local minima of $f$. Under truncated heavy-tailed noise, SGD hardly ever visits the two narrower minima $m_1$ and $m_3$, and spends almost all its time around the wider minima $m_2$ and $m_4$. (Right) Typical trajectories of different SGD methods when exploring $f$, with the dashed lines indicating the locations of the local minima. Without truncation, heavy-tailed SGDs keep jumping between all local minima of $f$ (see part (a) of the figure). In contrast, when driven by truncated heavy-tailed noise, the global dynamics of SGD resemble those of a continuous-time Markov chain that only visits the wider minima of $f$ (see part (b) of the figure).
  • Figure 3.1: Typical transition graphs under different choices of the truncation threshold $b$, illustrated with a univariate example. (Left) A univariate function $f$ with three attraction fields, where the numbers indicate the distance between each local minimum $m_i$ and the neighboring attraction field to the left or right. Note that in this univariate setting, we have $\mathcal{J}_b(i) = \lceil r(i) / b \rceil$ where $r(i) = \inf\{ |m_i - y|:\ y \notin I_i \}$, and, for each $k \leq \mathcal{J}_b(i)$, we have $\mathcal{G}^{(k)|b}(m_i) = [ m_i - kb, m_i + kb ]$. (Middle) The typical transition graph under $b = 0.5$. In particular, note that $\mathcal{J}^*_b(2) = \lceil 0.6/b \rceil = 2$, and the interval $\mathcal{G}^{( \mathcal{J}_b(2) )|b}(m_2) = [m_2 - 2 b, m_2 + 2b]$ intersects with both $I_1$ and $I_3$ (so the edges $m_2 \to m_1$ and $m_2 \to m_3$ are included in the typical transition graph). The entire graph $\mathcal{G}_b$ is irreducible since all nodes communicate with each other. (Right) The typical transition graph under $b = 0.4$. In this case, note that we still have $\mathcal{J}^*_b(2) = \lceil 0.6/b \rceil = 2$, but now $[m_2 - 2 b, m_2 + 2b]$ does not intersect with $I_3$. As a result, the typical transition graph does not contain the edge $m_2 \to m_3$, leading to two communication classes $G_1 = \{m_1,m_2\},\ G_2 = \{m_3\}$.

Theorems & Definitions (40)

  • Remark 1: Connection to the Relative Width of $I_i$
  • Definition 1: Typical Transition Graph
  • Theorem 2
  • Corollary 3
  • Corollary 4
  • Theorem 1: Metastability of $\bm X^{\eta|b}_t$: Absorbing Case
  • Theorem 2: Metastability of $\bm X^{\eta|b}_t$: Transient Case
  • Theorem 1: Theorem 2.8 of wang2023large
  • Lemma 2: Lemma 4.4 of wang2023large
  • Lemma 3: Lemma 4.5 of wang2023large
  • ...and 30 more