# Uniform-in-Time Wasserstein Stability Bounds for (Noisy) Stochastic Gradient Descent

arXiv: https://arxiv.org/abs/2305.12056 · HTML: https://mlsanity.com/paper/2305.12056/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 43 references that are in the map.
- First posted on arXiv 20 May 2023.

## Builds on (most cited first, 29 of 29)

- [Train faster, generalize better: Stability of stochastic gradient descent](https://mlsanity.com/paper/1509.01240/) (2015)
- [Information-theoretic analysis of generalization capability of learning algorithms](https://mlsanity.com/paper/1705.07809/) (2017)
- [Non-convex learning via Stochastic Gradient Langevin Dynamics: a nonasymptotic analysis](https://mlsanity.com/paper/1702.03849/) (2017)
- [Stability of Stochastic Gradient Descent on Nonsmooth Convex Losses](https://mlsanity.com/paper/2006.06914/) (2020)
- [Handbook of Convergence Theorems for (Stochastic) Gradient Methods](https://mlsanity.com/paper/2301.11235/) (2023)
- [Fine-Grained Analysis of Stability and Generalization for Stochastic Gradient Descent](https://mlsanity.com/paper/2006.08157/) (2020)
- [An Improved Analysis of Stochastic Gradient Descent with Momentum](https://mlsanity.com/paper/2007.07989/) (2020)
- [Non-strongly-convex smooth stochastic approximation with convergence rate O(1/n)](https://mlsanity.com/paper/1306.2119/) (2013)
- [Analysis and Design of Optimization Algorithms via Integral Quadratic Constraints](https://mlsanity.com/paper/1408.3595/) (2014)
- [High probability generalization bounds for uniformly stable algorithms with nearly optimal rate](https://mlsanity.com/paper/1902.10710/) (2019)
- [Generalization Bounds of SGLD for Non-convex Learning: Two Theoretical Viewpoints](https://mlsanity.com/paper/1707.05947/) (2017)
- [Statistical Optimality of Stochastic Gradient Descent on Hard Learning Problems through Multiple Passes](https://mlsanity.com/paper/1805.10074/) (2018)
- [On Generalization Error Bounds of Noisy Gradient Methods for Non-Convex Learning](https://mlsanity.com/paper/1902.00621/) (2019)
- [Stability and Convergence Trade-off of Iterative Optimization Algorithms](https://mlsanity.com/paper/1804.01619/) (2018)
- [Convergence of Langevin Monte Carlo in Chi-Squared and Renyi Divergence](https://mlsanity.com/paper/2007.11612/) (2020)
- [Algorithmic stability and hypothesis complexity](https://mlsanity.com/paper/1702.08712/) (2017)
- [Stability of SGD: Tightness Analysis and Improved Bounds](https://mlsanity.com/paper/2102.05274/) (2021)
- [Adaptivity of averaged stochastic gradient descent to local strong convexity for logistic regression](https://mlsanity.com/paper/1303.6149/) (2013)
- [Global Convergence of Stochastic Gradient Hamiltonian Monte Carlo for Non-Convex Stochastic Optimization: Non-Asymptotic Performance Bounds and Momentum-Based Acceleration](https://mlsanity.com/paper/1809.04618/) (2018)
- [Time-independent Generalization Bounds for SGLD in Non-convex Settings](https://mlsanity.com/paper/2111.12876/) (2021)
- [Perturbation theory for Markov chains via Wasserstein distance](https://mlsanity.com/paper/1503.04123/) (2015)
- [A Universally Optimal Multistage Accelerated Stochastic Gradient Method](https://mlsanity.com/paper/1901.08022/) (2019)
- [Accelerated Linear Convergence of Stochastic Momentum Methods in Wasserstein Distances](https://mlsanity.com/paper/1901.07445/) (2019)
- [Algorithmic Stability of Heavy-Tailed SGD with General Loss Functions](https://mlsanity.com/paper/2301.11885/) (2023)
- [Robust Accelerated Gradient Methods for Smooth Strongly Convex Functions](https://mlsanity.com/paper/1805.10579/) (2018)
- [Algorithmic Stability of Heavy-Tailed Stochastic Gradient Descent on Least Squares](https://mlsanity.com/paper/2206.01274/) (2022)
- [Excess Risk of Two-Layer ReLU Neural Networks in Teacher-Student Settings and its Superiority to Kernel Methods](https://mlsanity.com/paper/2205.14818/) (2022)
- [Robust Distributed Accelerated Stochastic Gradient Methods for Multi-Agent Networks](https://mlsanity.com/paper/1910.08701/) (2019)
- [Generalization in Supervised Learning Through Riemannian Contraction](https://mlsanity.com/paper/2201.06656/) (2022)

## Built on by (most cited first, 6 of 6)

- [Generalization Error Matters in Decentralized Learning Under Byzantine Attacks](https://mlsanity.com/paper/2407.08632/) (2024)
- [Generalization Bounds for Label Noise Stochastic Gradient Descent](https://mlsanity.com/paper/2311.00274/) (2023)
- [Privacy of SGD under Gaussian or Heavy-Tailed Noise: Guarantees without Gradient Clipping](https://mlsanity.com/paper/2403.02051/) (2024)
- [Generalization Bounds for Markov Algorithms through Entropy Flow Computations](https://mlsanity.com/paper/2502.07584/) (2025)
- [Stability, Complexity and Data-Dependent Worst-Case Generalization Bounds](https://mlsanity.com/paper/2507.06775/) (2025)
- [Aortic root landmark localization with optimal transport loss for heatmap regression](https://mlsanity.com/paper/2407.04921/) (2024)

## Abstract (arXiv)

Algorithmic stability is an important notion that has proven powerful for deriving generalization bounds for practical algorithms. The last decade has witnessed an increasing number of stability bounds for different algorithms applied on different classes of loss functions. While these bounds have illuminated various properties of optimization algorithms, the analysis of each case typically required a different proof technique with significantly different mathematical tools. In this study, we make a novel connection between learning theory and applied probability and introduce a unified guideline for proving Wasserstein stability bounds for stochastic optimization algorithms. We illustrate our approach on stochastic gradient descent (SGD) and we obtain time-uniform stability bounds (i.e., the bound does not increase with the number of iterations) for strongly convex losses and non-convex losses with additive noise, where we recover similar results to the prior art or extend them to more general cases by using a single proof technique. Our approach is flexible and can be generalizable to other popular optimizers, as it mainly requires developing Lyapunov functions, which are often readily available in the literature. It also illustrates that ergodicity is an important component for obtaining time-uniform bounds -- which might not be achieved for convex or non-convex losses unless additional noise is injected to the iterates. Finally, we slightly stretch our analysis technique and prove time-uniform bounds for SGD under convex and non-convex losses (without additional additive noise), which, to our knowledge, is novel.
