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

Wasserstein distance. For $p\geq 1$, the $p$-Wasserstein distance between two probability measures $\mu$ and $\nu$ on $\mathbb{R}^{d}$ is defined as Villani (2009):

where the infimum is taken over all couplings of $X\sim\mu$ and $Y\sim\nu$. In particular, the dual representation for the $1$-Wasserstein distance is given as Villani (2009):

where $\text{Lip}(1)$ consists of the functions $h:\mathbb{R}^{d}\rightarrow\mathbb{R}$ that are $1$-Lipschitz.

Wasserstein algorithmic stability. Algorithmic stability is a crucial concept in learning theory that has led to numerous significant theoretical breakthroughs (Bousquet and Elisseeff, 2002; Hardt et al., 2016). To begin, we will present the definition of algorithmic stability as stated in Hardt et al. (2016):

In this context, we purposefully employ a distinct notation for the loss function $\ell$ (in contrast to $f$) since our theoretical framework necessitates measuring algorithmic stability through a surrogate loss function, which may differ from the original loss $f$. More precisely, our bounds will be based on the 1-Wasserstein distance, hence, we will need the surrogate loss $\ell$ to be a Lipschitz continuous function, as we will detail in (5). On the other hand, for the original loss $f$ we will need some form of convexity (e.g., strongly convex, convex, or dissipative) and we will need the gradient of $f$ to be Lipschitz continuous, in order to derive Wasserstein bounds. Unfortunately, under these assumptions, we cannot further impose $f$ itself to be Lipschitz, hence the need for surrogate losses. Nevertheless, the usage of surrogate losses is common in learning theory, see e.g, Farghly and Rebeschini (2021); Raj et al. (2023b), and we present concrete practical examples in the Appendix.

Now, we present a result from Hardt et al. (2016) that establishes a connection between algorithmic stability and the generalization performance of a randomized algorithm. Prior to presenting the result, we define the empirical and population risks with respect to the loss function $\ell$ as follows:

For a randomized algorithm, if $\nu$ and $\hat{\nu}$ denotes the law of $\mathcal{A}(X)$ and $\mathcal{A}(\hat{X})$ then for a $\mathcal{L}$-Lipschitz surrogate loss function $\ell$, we have the following generalization error guarantee,

The above result can be directly obtained from the combination of the results given in (5), Definition 1, and Theorem 2 (see also Raginsky et al. (2016)).

Next, we recall the Wasserstein perturbation bound for Markov chains from Rudolf and Schweizer (2018). Let $(\theta_{n})_{n=0}^{\infty}$ be a Markov chain with transition kernel $P$ and initial distribution $p_{0}$, i.e., we have almost surely

and $p_{0}(A)=\mathbb{P}(\theta_{0}\in A)$ for any measurable set $A\subseteq\mathbb{R}^{d}$ and $n\in\mathbb{N}$. We assume that $(\hat{\theta}_{n})_{n=0}^{\infty}$ is another Markov chain with transition kernel $\hat{P}$ and initial distribution $\hat{p}_{0}$. We denote by $p_{n}$ the distribution of $\theta_{n}$ and by $\hat{p}_{n}$ the distribution of $\hat{\theta}_{n}$. By $\delta_{\theta}$, we denote the Dirac delta distribution at $\theta$, i.e. the probability measure concentrated at $\theta$. For a measurable set $A\subseteq\mathbb{R}^{d}$, we also use the notation $\delta_{\theta}P(A):=P(\theta,A)$.

Lemma 3 provides a sufficient condition for the distributions $p_{n}$ and $\hat{p}_{n}$ after $n$ iterations to stay close to each other given the initial distributions $p_{0}$. Lemma 3 will provide a key role in helping us derive the main results in Section 3.1 and Section 3.2. Later, in the Appendix, we will state and prove a modification of Lemma 3 (see Lemma 23 in the Appendix) that will be crucial to obtaining the main result in Section 3.3.

To illustrate the proof technique, we start by considering a quadratic loss of the form: $f(\theta,x_{i}):=(a_{i}^{\top}\theta-y_{i})^{2}/2$ where, $x_{i}:=(a_{i},y_{i})$ and $\nabla f(\theta,x_{i})=a_{i}(a_{i}^{\top}\theta-y_{i})$. In this setting, the SGD recursion takes the following form:

The sequence $(H_{k},q_{k})$ are i.i.d. and for every $k$, $(H_{k},q_{k})$ is independent of $\theta_{k-1}$.

Similarly, we can write down the iterates of SGD with a different data set $\hat{X}_{n}:=\{\hat{x}_{1},\ldots,\hat{x}_{n}\}$ with $\hat{x}_{i}=(\hat{a}_{i},\hat{y}_{i})$, where $\hat{X}_{n}$ differs from $X_{n}$ with at most one element:

Our goal is to obtain an algorithmic stability bound, through estimating the $1$-Wasserstein distance between the distribution of $\theta_{k}$ and $\hat{\theta}_{k}$ and we will now illustrate the three-step proof technique that we described in Section 1. To be able to apply the perturbation theory (Rudolf and Schweizer, 2018), we start by establishing the geometric ergodicity of the Markov process $(\theta_{k})_{k\geq 0}$ with transition kernel $P(\theta,\cdot)$, given in the following lemma.

We note that since $H_{1}\succeq 0$, the assumption in Lemma 4 can be satisfied under mild assumptions, for example when $H_{1}\succ 0$ with a positive probability, which is satisfied for $\eta$ small enough.

In the second step, we construct a Lyapunov function $\hat{V}$ that satisfies the conditions of Lemma 3.

In our third and last step, we estimate the perturbation gap based on the Lyapunov function $\hat{V}$ in the form of (10), assuming that the data is bounded. Such bounded data assumptions have been commonly made in the literature (Bach, 2014; Bach and Moulines, 2013).

Note that Lemma 3 relies on three conditions: the Wasserstein contraction in (10), which is obtained through Lemma 4, the drift condition for the Lyapunov function in (11), which is obtained in Lemma 5 and finally the estimate on $\gamma$ in (12) which is about the one-step $1$-Wasserstein distance between two semi-groups that in our context are associated with two datasets that differ by at most one element, which is obtained in Lemma 6. The only place the neighborhood assumption ($\sup_{x\in\mathcal{X}}\|x\|\leq D$) is used is in the expression of $\gamma$ in equation (12). Now, having all the ingredients, we can invoke Lemma 3 and we obtain the following result which provides a 1-Wasserstein bound between the distribution of iterates when applied to datasets that differ by one point.

For $Y\in\bigcup_{n=1}^{\infty}\mathcal{X}^{n}$ and $k\geq 0$, let $\nu(Y,k)$ denote the law of the $k$-th the SGD iterate when $Y$ is used as the dataset, i.e., $\nu(X,k)$ and $\nu(\hat{X},k)$ denote the distributions of $\theta_{k}$ and $\hat{\theta}_{k}$ obtained by the recursions (13) and (14) respectively. As shorthand notation, set $\nu_{k}:=\nu(X,k)$ and $\hat{\nu}_{k}:=\nu(\hat{X},k)$.

Next, we consider strongly convex losses. In the remainder of the paper, we will always assume that for every $x\in\mathcal{X}$, $f(\cdot,x)$ is differentiable.

Before proceeding to the stability bound, we first introduce the following assumptions.

This assumption is a pseudo-Lipschitz-like condition on $\nabla f$ and is satisfied for various problems such as GLMs (Bach, 2014). Next, we assume that the loss function $f$ is strongly convex.

By using the same recipe as we used for quadratic losses, we obtain the following stability result.

Similarly to the quadratic case, we can now directly obtain a bound on expected generalization error using (8). More precisely, for an $\mathcal{L}$-Lipschitz surrogate loss function $\ell$, we have

The bound above has the same dependence on the number of samples as the ones of the previous stability analysis of (projected) SGD for strongly convex functions (Hardt et al., 2016; Liu et al., 2017; Lei and Ying, 2020). However, we have a worse dependence on the strong convexity parameter $\mu$.

Finally, we consider a class of non-convex loss functions. We assume that the loss function satisfies the following dissipativity condition.

The class of dissipative functions satisfying this assumption are the ones that admit some gradient growth in radial directions outside a compact set. Inside the compact set though, they can have quite general non-convexity patterns. As concrete examples, they include certain one-hidden-layer neural networks (Akiyama and Suzuki, 2023); they arise in non-convex formulations of classification problems (e.g. in logistic regression with a sigmoid/non-convex link function); they can also arise in robust regression problems, see e.g. Gao et al. (2022). Also, any function that is strongly convex outside of a ball of radius for some will satisfy this assumption. Consequently, regularized regression problems where the loss is a strongly convex quadratic plus a smooth penalty that grows slower than a quadratic will belong to this class; a concrete example would be smoothed Lasso regression; many other examples are also given in Erdogdu et al. (2022). Dissipative functions also arise frequently in the sampling and Bayesian learning and global convergence in non-convex optimization literature (Raginsky et al., 2017; Gao et al., 2022).

Unlike the strongly-convex case, we can no longer obtain a Wasserstein contraction bound using the synchronous coupling technique as we did in the proof of Theorem 8. To circumvent this problem, in this setting, we consider a noisy version of SGD, with the following recursion:

where $\xi_{k}$ are additional i.i.d. random vectors in $\mathbb{R}^{d}$, independent of $\theta_{k-1}$ and $\Omega_{k}$, satisfying the following assumption.

Note that the SGLD algorithm (Welling and Teh, 2011) is a special case of this recursion, whilst our noise model can accommodate non-Gaussian distributions with finite second-order moment.

Analogously, let us define the (noisy) SGD recursion with the data set $\hat{X}_{n}$ as

and let $p(\theta,\theta_{1})$ denote the probability density function of $\theta_{1}=\theta-\frac{\eta}{b}\sum_{i\in\Omega_{1}}\nabla f(\theta,x_{i})+\eta\xi_{1}$. Further let $\theta_{*}$ be a minimizer of $\hat{F}(\cdot,X_{n})$. Then, by following the same three-step recipe, we obtain the following stability bound. Here, we do not provide all the constants explicitly for the sake of clarity; the complete theorem statement is given in Theorem 25 (Appendix E.1).

Contrary to our previous results, the proof technique for showing Wasserstein contraction (as in Lemma 4) for this theorem relies on verifying the drift condition (Assumption 7) and the minorization condition (Assumption 8) as given in Hairer and Mattingly (2011). Once these conditions are satisfied, we invoke the explicitly computable bounds on the convergence of Markov chains developed in Hairer and Mattingly (2011).

From equation (8), we directly obtain the following generalization error bound for $\mathcal{L}$-Lipschitz surrogate loss function,

where the constants are defined in Theorem 9¹¹1By using the decomposition (2), we can obtain excess risk bounds for SGLD by combining our results with Xu and Raginsky (2017): it was shown that gradient Langevin dynamics has the following optimization error is $O(\varepsilon+d^{3/2}b^{-1/4}\lambda^{-1}\log 1/\varepsilon)$ after $K=O(d\varepsilon^{-1}\lambda^{-1}\log 1/\varepsilon)$ iterations, where $b$ is the mini-batch size and $\lambda$ is the uniform spectral gap of the continuous-time Langevin dynamics. Similar results are given for SGLD in (Xu and Raginsky, 2017, Theorem 3.6). . The above result can be directly compared with the result in (Farghly and Rebeschini, 2021, Theorem 4.1) that has the same dependence on $n$ and $b$. However, our result is more general in the sense that we do not assume our noise to be Gaussian noise. Note that Li et al. (2019) can also accommodate non-Gaussian noise; however, their bounds increase with the number of iterations.

Next, let us provide some explicitly computable lower bounds for $\hat{\eta}$ in (22). This is achievable if we specify further the noise assumption. Under the assumption that $\xi_{k}$ are i.i.d. Gaussian distributed, we have the following corollary.

The above corollary provides an explicit lower bound for $\hat{\eta}$ (instead of the less transparent inequality constraints in Theorem 9), and by combining with Remark 10 (see equation (23)) leads to an explicit formula for $\bar{\eta}$ which is essential to characterize the Wasserstein upper bound in (21) in Theorem 9.

The stability result we obtained in Theorem 9 required us to introduce an additional noise (Assumption 4) to be able to invoke Lemma 3. We will now show that it is possible to use a more direct approach to obtain 2-Wasserstein algorithmic stability in the non-convex case under Assumption 3 without relying on Rudolf and Schweizer (2018). However, we will observe that without geometric ergodicity will have a non-vanishing bias term in the bound. Note that, since $\mathcal{W}_{1}(\nu_{k},\hat{\nu}_{k})\leq\mathcal{W}_{p}(\nu_{k},\hat{\nu}_{k})$ for all $p\geq 1$, the following bound still yields a generalization bound by (8).

While the bound (25) does not increase with the number of iterations, it is easy to see that it does not vanish as $n\rightarrow\infty$, and it is small only when $K$ from the dissipativity condition (Assumption 3) is small. In other words, if we consider $K$ to be the level of non-convexity (e.g., $K=0$ corresponds to strong convexity), as the function becomes ‘more non-convex’ the persistent term in the bound will get larger. While this persistent term might make the bound vacuous when $n\to\infty$, for moderate $n$ the bound can be still informative as the persistent term might be dominated by the first two terms.

Moreover, discarding the persistent bias term, this bound leads to a generalization bound with rate $n^{-1/2}$, rather than $n^{-1}$ as before. This indicates that it is beneficial to add additional noise $\xi_{k}$ in SGD as in Theorem 9 in order for the dynamics to be geometrically ergodic that can lead to a sharp bound as $n\rightarrow\infty$. Finally, we note that as Theorem 12 involves 2-Wasserstein distance, it can pave the way for generalization bounds without requiring a surrogate loss. Yet, this is not immediate and would require deriving uniform $L^{2}$ bounds for the iterates, e.g., Raginsky et al. (2017).

We now present our final stability bound, where we consider relaxing the strong convexity assumption (Assumption 2) to the following milder assumption.

Note that as $p<2$, the function class can be seen as an intermediate class between convex and strongly convex functions, and such a class of functions has been studied in the optimization literature (Dunn, 1981; Bertsekas, 2015).

We analogously modify Assumption 1 and consider the following assumption.

The next theorem establishes a stability bound for the considered class of convex losses in the stationary regime of SGD.