Title: Uniform-in-Time Wasserstein Stability Bounds for (Noisy) Stochastic Gradient Descent
Venue: NeurIPS 2023
Accepted: yes
Keywords: Algorithmic stability

SGD

Wasserstein distance
Abstract: 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.

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

Abstract

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 timeuniform 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.

1 Introduction

With the development of modern machine learning applications, understanding the generalization properties of stochastic gradient descent (SGD) has become a major challenge in statistical learning theory. In this context, the main goal is to obtain computable upper-bounds on the population risk associated with the output of the SGD algorithm that is given as follows: F ( θ ) := E x ∼D [ f ( θ, x )] , where x ∈ X denotes a random data point, D is the (unknown) data distribution defined on the data space X , θ denotes the parameter vector, and f : R d ×X → R is an instantaneous loss function.

In a practical setting, directly minimizing F ( θ ) is not typically possible as D is unknown; yet one typically has access to a finite data set X n = { x 1 , . . . , x n } ∈ X n , where we assume each x i is independent and identically distributed (i.i.d.) with the common distribution D . Hence, given X n , one can then attempt to minimize the empirical risk ˆ F ( θ, X n ) := 1 n ∑ n i =1 f ( θ, x i ) as a proxy for F ( θ ) . In this setting, SGD has been one of the most popular optimization algorithms for minimizing

ˆ F ( θ ) and is based on the following recursion:

where η is the step-size, b is the batch-size, Ω k is the minibatch that is chosen randomly from the set { 1 , 2 , . . . , n } , and its cardinality satisfies | Ω k | = b .

One fruitful approach for estimating the population risk attained by SGD, i.e., F ( θ k ) , is based on the following simple decomposition:

where the last term is called the generalization error . Once a computable upper-bound for the generalization error can be obtained, this decomposition directly leads to a computable upper bound for the population risk F ( θ k ) , since ˆ F ( θ k ) can be computed thanks to the availability of X n . Hence, the challenge here is reduced to derive upper-bounds on | ˆ F ( θ k ) -F ( θ k ) | , typically referred to as generalization bounds.

Among many approaches for deriving generalization bounds, algorithmic stability [BE02] has been one of the most fruitful notions that have paved the way to numerous generalization bounds for stochastic optimization algorithms [HRS16, CJY18, MWZZ18, FV19, LY20, ZZB + 22]. In a nutshell, algorithmic stability measures how much the algorithm output differs if we replace one data point in X n with a new sample. More precisely, in the context of SGD, given another data set ˆ X n = { ˆ x 1 , . . . , ˆ x n } = { x 1 , . . . , x i -1 , ˆ x i , x i +1 , . . . x n } ∈ X n that differs from X n by at most one element, we (theoretically) consider running SGD on ˆ X n , i.e.,

and we are interested in the discrepancy between θ k and ˆ θ k in some precise sense (to be formally defined in the next section). The wisdom of algorithmic stability indicates that a smaller discrepancy between θ k and ˆ θ k implies a smaller generalization error.

The last decade has witnessed an increasing number of stability bounds for different algorithms applied on different classes of loss functions. In a pioneering study, [HRS16] proved a variety of stability bounds for SGD, for strongly convex, convex, and non-convex problems. Their analysis showed that, under strong convexity and bounded gradient assumptions, the generalization error of SGD with constant step-size is of order n -1 ; whereas for general convex and non-convex problems, their bounds diverged with the number of iterations (even with a projection step), unless a decreasing step-size is used. In subsequent studies [LY20, KWS23] extended the results of [HRS16], by either relaxing the assumptions or generalizing the setting to more general algorithms. However, their bounds still diverged for constant step-sizes, unless strong convexity is assumed. In a recent study, [BFGT20] proved stability lower-bounds for projected SGD when the loss is convex and non-smooth. Their results showed for general non-smooth loss functions we cannot expect to prove time-uniform (i.e., non-divergent with the number of iterations) stability bounds for SGD, even when a projection step is appended.

In a related line of research, several studies investigated the algorithmic stability of the stochastic gradient Langevin dynamics (SGLD) algorithm [WT11], which is essentially a 'noisy' version of SGD that uses the following recursion: θ k = θ k -1 -η ˜ ∇ ˆ F k ( θ k -1 , X n ) + ξ k , where ( ξ k ) k ≥ 0 is a sequence of i.i.d. Gaussian vectors, independent of θ k -1 and Ω k . The authors of [RRT17, MWZZ18] proved stability bounds for SGLD for non-convex losses, which were then extended to more general (non-Gaussian) noise settings in [LLQ19]. While these bounds hinted at the benefits of additional noise in terms of stability, they still increased with the number of iterations, which limited the impact of their results. More recently, [FR21] proved the first time-uniform stability bounds for SGLD under non-convexity, indicating that, with the presence of additive Gaussian noise, better stability bounds can be achieved. Their time-uniform results were then extended to non-Gaussian, heavy-tailed perturbations in [RBG + 23, RZG¸ S23] for quadratic and a class of non-convex problems.

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. Hence, it is not straightforward to extend the existing techniques to different algorithms with different classes of loss functions. Moreover, currently, it is not clear how the noisy perturbations affect algorithmic stability so that time-uniform bounds can be achieved, and more generally, it is not clear in which circumstances one might hope for time-uniform stability bounds.

In this study, we contribute to this line of research and prove novel time-uniform algorithmic stability bounds for SGD and its noisy versions. Our main contributions are as follows:

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 with a constant step-size. Our approach is based on Markov chain perturbation theory [RS18], which offers a three-step proof technique for deriving stability bounds: (i) showing the optimizer is geometrically ergodic, (ii) obtaining a Lyapunov function for the optimizer and the loss, and (iii) bounding the discrepancy between the Markov transition kernels associated with the chains ( θ k ) k ≥ 0 and ( ˆ θ k ) k ≥ 0 . We illustrate this approach on SGD and show that time-uniform stability bounds can be obtained under a pseudo-Lipschitz-like condition for smooth strongly-convex losses (we recover similar results to the ones of [HRS16]) and a class of non-convex losses (that satisfy a dissipativity condition) when a noisy perturbation with finite variance (not necessarily Gaussian, hence more general than [FR21]) is introduced. Our results shed more light on the role of the additional noise in terms of obtaining time-uniform bounds: in the non-convex case the optimizer might not be geometrically ergodic unless additional noise is introduced, hence the bound cannot be obtained. Moreover, 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 [AFGO20, LRP16, FGO + 22, GRZ22, LGY20, AFGO19].

We then investigate the case where no additional noise is introduced to the SGD recursion and the geometric ergodicity condition does not hold. First, for non-convex losses, we prove a time-uniform stability bound, where the bound converges to a positive number (instead of zero) as n →∞ , and this limit depends on the 'level of non-convexity'. Then, we consider a class of (non-strongly) convex functions and prove stability bounds for the stationary distribution of ( θ k ) k ≥ 0 , which vanish as n increases. To the best of our knowledge, these results are novel, and indicate that the stability bounds do not need to increase with time even under non-convexity and without additional perturbations; yet, they might have a different nature depending on the problem class.

One limitation of our analysis is that it requires Lipschitz surrogate loss functions and does not directly handle the original loss function, due to the use of the Wasserstein distance [RRT + 16]. Yet, surrogate losses have been readily utilized in the recent stability literature (e.g., [RBG + 23, RZG¸ S23]) and we believe that our analysis might illuminate uncovered aspects of SGD even with this requirement. All the proofs are provided in the Appendix.

2 Technical Background

2.1 The Wasserstein distance and Wasserstein algorithmic stability

Wasserstein distance. For p ≥ 1 , the p -Wasserstein distance between two probability measures µ and ν on R d is defined as [Vil09]:

where the infimum is taken over all couplings of X ∼ µ and Y ∼ ν . In particular, the dual representation for the 1 -Wasserstein distance is given as [Vil09]:

where Lip (1) consists of the functions h : R d → 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 [BE02, HRS16]. To begin, we will present the definition of algorithmic stability as stated in [HRS16]:

Definition 2.1 ([HRS16], Definition 2.1) . Let RV ( R d ) denote the set of R d -valued random vectors. For a (surrogate) loss function ℓ : R d × X → R , an algorithm A : ⋃ ∞ n =1 X n → RV ( R d ) is ε -uniformly stable if

where the first supremum is taken over data X, ˆ X ∈ X n that differ by one element, denoted by X ∼ = ˆ X .

In this context, we purposefully employ a distinct notation for the loss function ℓ (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 ℓ to be a Lipschitz continuous function, as we will detail in (2.2). 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, [FR21, RZG¸ S23], and we present concrete practical examples in the Appendix.

Now, we present a result from [HRS16] 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 ℓ as follows:

Theorem 2.1 ([HRS16], Theorem 2.2) . Suppose that A is an ε -uniformly stable algorithm, then the expected generalization error is bounded by

For a randomized algorithm, if ν and ˆ ν denotes the law of A ( X ) and A ( ˆ X ) then for a L -Lipschitz surrogate loss function ℓ , we have the following generalization error guarantee,

The above result can be directly obtained from the combination of the results given in (2.2), Definition 2.1, and Theorem 2.1 (see also [RRT + 16]).

2.2 Perturbation theory for Markov chains

Next, we recall the Wasserstein perturbation bound for Markov chains from [RS18]. Let ( θ n ) ∞ n =0 be a Markov chain with transition kernel P and initial distribution p 0 , i.e., we have almost surely

and p 0 ( A ) = P ( θ 0 ∈ A ) for any measurable set A ⊆ R d and n ∈ N . We assume that ( ˆ θ n ) ∞ n =0 is another Markov chain with transition kernel ˆ P and initial distribution ˆ p 0 . We denote by p n the distribution of θ n and by ˆ p n the distribution of ˆ θ n . By δ θ , we denote the Dirac delta distribution at θ , i.e. the probability measure concentrated at θ . For a measurable set A ⊆ R d , we also use the notation δ θ P ( A ) := P ( θ, A ) .

Lemma 2.1 ([RS18], Theorem 3.1) . Assume that there exist some ρ ∈ [0 , 1) and C ∈ (0 , ∞ ) such that

̸

for any n ∈ N . Further assume that there exist some δ ∈ (0 , 1) and L ∈ (0 , ∞ ) and a measurable Lyapunov function ˆ V : R d → [1 , ∞ ) of ˆ P such that for any θ ∈ R d :

where ( ˆ P ˆ V )( θ ) := ∫ R d ˆ V ( ˆ θ ) ˆ P ( θ, d ˆ θ ) . Then, we have

Lemma 2.1 provides a sufficient condition for the distributions p n and ˆ p n after n iterations to stay close to each other given the initial distributions p 0 . Lemma 2.1 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 2.1 (see Lemma E.5 in the Appendix) that will be crucial to obtaining the main result in Section 3.3.

3 Wasserstein Stability of SGD via Markov Chain Perturbations

In this section, we will derive time-uniform Wasserstein stability bounds for SGD by using the perturbation theory presented in [RS18]. Before considering general losses that can be non-convex, we first consider the simpler case of quadratic losses to illustrate our key ideas.

3.1 Warm up: quadratic case

To illustrate the proof technique, we start by considering a quadratic loss of the form: f ( θ, x i ) := ( a ⊤ i θ -y i ) 2 / 2 where, x i := ( a i , y i ) and ∇ f ( θ, x i ) = a i ( a ⊤ i θ -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 θ k -1 .

Similarly, we can write down the iterates of SGD with a different data set ˆ X n := { ˆ x 1 , . . . , ˆ x n } with ˆ x i = (ˆ a i , ˆ y i ) , where ˆ 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 θ k and ˆ θ 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 [RS18], we start by establishing the geometric ergodicity of the Markov process ( θ k ) k ≥ 0 with transition kernel P ( θ, · ) , given in the following lemma.

Lemma 3.1. Assume that ρ := E ∥ ∥ I -η b H 1 ∥ ∥ < 1 . Then, for any k ∈ N , we have the following inequality: W 1 ( P k ( θ, · ) , P k ( ˜ θ, · )) ≤ ρ k ∥ θ -˜ θ ∥ .

We note that since H 1 ⪰ 0 , the assumption in Lemma 3.1 can be satisfied under mild assumptions, for example when H 1 ≻ 0 with a positive probability, which is satisfied for η small enough.

In the second step, we construct a Lyapunov function ˆ V that satisfies the conditions of Lemma 2.1.

In our third and last step, we estimate the perturbation gap based on the Lyapunov function ˆ V in the form of (2.7), assuming that the data is bounded. Such bounded data assumptions have been commonly made in the literature [Bac14, BM13].

Lemma 3.3. If sup x ∈X ∥ x ∥ ≤ D for some D < ∞ , then, we have sup θ ∈ R d W 1 ( δ θ P,δ θ ˆ P ) ˆ V ( θ ) ≤ 2 ηD 2 n .

Note that Lemma 2.1 relies on three conditions: the Wasserstein contraction in (2.7), which is obtained through Lemma 3.1, the drift condition for the Lyapunov function in (2.8), which is obtained in Lemma 3.2 and finally the estimate on γ in (2.9) 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 3.3. The only place the neighborhood assumption ( sup x ∈X ∥ x ∥ ≤ D ) is used is in the expression of γ in equation (2.9). Now, having all the ingredients, we can invoke Lemma 2.1 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 ∈ ⋃ ∞ n =1 X n and k ≥ 0 , let ν ( Y, k ) denote the law of the k -th the SGD iterate when Y is used as the dataset, i.e., ν ( X,k ) and ν ( ˆ X,k ) denote the distributions of θ k and ˆ θ k obtained by the recursions (3.1) and (3.2) respectively. As shorthand notation, set ν k := ν ( X,k ) and ˆ ν k := ν ( ˆ X,k ) .

Theorem 3.1. Assume θ 0 = ˆ θ 0 = θ . We also assume that ρ := E ∥ ∥ I -η b H 1 ∥ ∥ < 1 and ˆ ρ := E ∥ ∥ ∥ I -η b ˆ H 1 ∥ ∥ ∥ < 1 and sup x ∈X ∥ x ∥ ≤ D for some D < ∞ . Then, we have

Proof. The result directly follows from Lemma 3.1, Lemma 3.2, Lemma 3.3 and Lemma 2.1.

By a direct application of (2.5), we can obtain a generalization bound for an L -Lipschitz surrogate loss function, as follows:

where ρ 0 = sup X ∥ 1 -η b H X ∥ , H X = ∑ i ∈ Ω k ,a j ∈ X a j a ⊤ j and X is a random set of n -data points from the data generating distribution. The generalization bound obtained above does not include the mean square error in the unbounded case but covers a larger class of surrogate loss functions. Because of this incompatibility, a direct comparison is not possible; however, the rate obtained in the equation above has the same dependence on the number of samples that were obtained in the previous works [LY20]. For least squares, there are other works using integral operators that develop generalization bounds for SGD under a capacity condition [LR17, PVRB18]. However, these bounds only hold for the least square loss.

3.2 Strongly convex case

Next, we consider strongly convex losses. In the remainder of the paper, we will always assume that for every x ∈ X , f ( · , x ) is differentiable.

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

Assumption 3.1. There exist constants K 1 , K 2 > 0 such that for any θ, ˆ θ ∈ R d and every x ∈ X ,

This assumption is a pseudo-Lipschitz-like condition on ∇ f and is satisfied for various problems such as GLMs [Bac14]. Next, we assume that the loss function f is strongly convex.

Assumption 3.2. There exists a universal constant µ > 0 such that for any θ 1 , θ 2 ∈ R d and x ∈ X ,

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

Theorem 3.2. Let θ 0 = ˆ θ 0 = θ . Assume that Assumption 3.1 and Assumption 3.2 hold. We also assume that η < min { 1 µ , µ K 2 1 +64 D 2 K 2 2 } , sup x ∈X ∥ x ∥ ≤ D for some D < ∞ , and sup x ∈X ∥∇ f (0 , x ) ∥ ≤ E for some E < ∞ . Let ν k and ˆ ν k denote the distributions of θ k and

ˆ θ k respectively. Then, we have

Similarly to the quadratic case, we can now directly obtain a bound on expected generalization error using (2.5). More precisely, for an L -Lipschitz surrogate loss function ℓ , 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 [HRS16, LLNT17, LY20]. However, we have a worse dependence on the strong convexity parameter µ .

3.3 Non-convex case with additive noise

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

Assumption 3.3. There exist constants m> 0 and K > 0 such that for any θ 1 , θ 2 ∈ R d and x ∈ X ,

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 [AS23]; 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. [GGZ22]. 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 [EHZ22]. Dissipative functions also arise frequently in the sampling and Bayesian learning and global convergence in non-convex optimization literature [RRT17, GGZ22].

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 3.2. To circumvent this problem, in this setting, we consider a noisy version of SGD, with the following recursion:

where ξ k are additional i.i.d. random vectors in R d , independent of θ k -1 and Ω k , satisfying the following assumption.

Assumption 3.4. ξ 1 is random vector on R d with a continuous density p ( x ) that is positive everywhere, i.e. p ( x ) > 0 for any x ∈ R d and E [ ξ 1 ] = 0 , σ 2 := E [ ∥ ξ 1 ∥ 2 ] < ∞ .

Note that the SGLD algorithm [WT11] 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 ˆ X n as

and let p ( θ, θ 1 ) denote the probability density function of θ 1 = θ -η b ∑ i ∈ Ω 1 ∇ f ( θ, x i ) + ηξ 1 . Further let θ ∗ be a minimizer of ˆ F ( · , 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 E.1 (Appendix E.1).

Theorem 3.3. Let θ 0 = ˆ θ 0 = θ . Assume that Assumption 3.1, Assumption 3.3 and Assumption 3.4 hold. We also assume that η < min { 1 m , m K 2 1 +64 D 2 K 2 2 } and sup x ∈X ∥ x ∥ ≤ D for some D < ∞ and sup x ∈X ∥∇ f (0 , x ) ∥ ≤ E for some E < ∞ . For any ˆ η ∈ (0 , 1) , define M > 0 so that ∫ ∥ θ 1 -θ ∗ ∥≤ M p ( θ ∗ , θ 1 ) dθ 1 ≥ √ ˆ η and for any R > 2 K 0 m where K 0 is defined in (E.2) so that

Let ν k and ˆ ν k denote the distributions of θ k and ˆ θ k respectively. Then, we have

where for any η 0 ∈ (0 , ˆ η ) and γ 0 ∈ ( 1 -mη + 2 ηK 0 R , 1 ) , ψ = η 0 ηK 0 and ¯ η = (1 -(ˆ η -η 0 )) ∨ 2+ Rψγ 0 2+ Rψ . The constant C 1 ≡ C 1 ( ψ, η, ˆ η, R, γ 0 , K 1 , K 2 , σ 2 , D, E ) is explicitly stated in the proof.

Contrary to our previous results, the proof technique for showing Wasserstein contraction (as in Lemma 3.1) for this theorem relies on verifying the drift condition (Assumption B.1) and the minorization condition (Assumption B.2) as given in [HM11]. Once these conditions are satisfied, we invoke the explicitly computable bounds on the convergence of Markov chains developed in [HM11].

From equation (2.5), we directly obtain the following generalization error bound for L -Lipschitz surrogate loss function,

where the constants are defined in Theorem 3.3 1 . The above result can be directly compared with the result in [FR21, 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 [LLQ19] can also accommodate non-Gaussian noise; however, their bounds increase with the number of iterations.

Remark 3.4. In Theorem 3.3, we can take R = 2 K 0 m (1 + ϵ ) for some fixed ϵ ∈ (0 , 1) so we can take

Moreover, we can take η 0 = ˆ η 2 , γ 0 = 1 -mηϵ 2 , and ψ = ˆ η 2 ηK 0 , so that

provided that η ≤ 1 . Note that the parameter ˆ η in (3.10) appears in the upper bound in equation (3.9) that controls the 1 -Wasserstein algorithmic stability of the SGD. It is easy to see from equation (3.9) that the smaller ¯ η , the smaller the 1 -Wasserstein bound. By the defintion of ¯ η , the larger ˆ η , the smaller the 1 -Wasserstein bound. As a result, we would like to choose ˆ η to be as large as possible, and the equation (3.10) provides an explicit value that ˆ η can take, which is already the largest as possible.

Next, let us provide some explicitly computable lower bounds for ˆ η in (3.10). This is achievable if we specify further the noise assumption. Under the assumption that ξ k are i.i.d. Gaussian distributed, we have the following corollary.

Corollary 3.5. Under the assumptions in Theorem 3.3, we further assume the noise ξ k are i.i.d. Gaussian N (0 , Σ) so that E [ ∥ ξ 1 ∥ 2 ] = tr (Σ) = σ 2 . We also assume that Σ ≺ I d . Then, we have

The above corollary provides an explicit lower bound for ˆ η (instead of the less transparent inequality constraints in Theorem 3.3), and by combining with Remark 3.4 (see equation (3.11)) leads to an explicit formula for ¯ η which is essential to characterize the Wasserstein upper bound in (3.9) in Theorem 3.3.

4 Wasserstein Stability of SGD without Geometric Ergodicity

While the Markov chain perturbation theory enabled us to develop stability bounds for the case where we can ensure geometric ergodicity in the Wasserstein sense (i.e., proving contraction bounds), we have observed that such a strong ergodicity notion might not hold for non-strongly convex losses. In this section, we will prove two more stability bounds for SGD, without relying on [RS18], hence without requiring geometric ergodicity. To the best of our knowledge, these are the first uniform-time stability bounds for the considered classes of convex and non-convex problems.

4.1 Non-convex case without additive noise

The stability result we obtained in Theorem 3.3 required us to introduce an additional noise (Assumption 3.4) to be able to invoke Lemma 2.1. 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.3 without relying on [RS18]. However, we will observe that without geometric ergodicity will have a non-vanishing bias term in the bound. Note that, since W 1 ( ν k , ˆ ν k ) ≤ W p ( ν k , ˆ ν k ) for all p ≥ 1 , the following bound still yields a generalization bound by (2.5).

Theorem 4.1. Assume θ 0 = ˆ θ 0 = θ . We also assume that Assumption 3.1 and Assumption 3.3 hold and η < min { 1 m , m K 2 1 +64 D 2 K 2 2 } and sup x ∈X ∥ x ∥ ≤ D for some D < ∞ and sup x ∈X ∥∇ f (0 , x ) ∥ ≤ E for some E < ∞ . Let ν k and ˆ ν k denote the distributions of θ k and ˆ θ k respectively. Then, we have

where the constant B is explicitly defined in the proof.

While the bound (4.1) does not increase with the number of iterations, it is easy to see that it does not vanish as n →∞ , and it is small only when K from the dissipativity condition (Assumption 3.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 →∞ , 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 ξ k in SGD as in Theorem 3.3 in order for the dynamics to be geometrically ergodic that can lead to a sharp bound as n →∞ . Finally, we note that as Theorem 4.1 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., [RRT17].

4.2 Convex case with additional geometric structure

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

Assumption 4.1. There exists universal constants µ > 0 and p ∈ (1 , 2) such that for any θ 1 , θ 2 ∈ R d and x ∈ X , ⟨∇ f ( θ 1 , x ) -∇ f ( θ 2 , x ) , θ 1 -θ 2 ⟩ ≥ µ ∥ θ 1 -θ 2 ∥ p .

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 [Dun81, Ber15].

We analogously modify Assumption 3.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.

Theorem 4.2. Let θ 0 = ˆ θ 0 = θ . Suppose Assumption 4.1 and Assumption 4.2 hold (with p ∈ (1 , 2) ) and η ≤ µ K 2 1 +2 p +4 D 2 K 2 2 and sup x ∈X ∥ x ∥ ≤ D for some D < ∞ and sup x ∈X ∥∇ f (0 , x ) ∥ ≤ E for some E < ∞ . Then ν k and ˆ ν k converge to the unique stationary distributions ν ∞ and ˆ ν ∞ respectively and moreover, we have

where the constants C 2 ≡ C 2 ( η, µ, K 2 , D, E ) and C 3 ≡ C 3 ( η, µ, K 1 , K 2 , D, E ) are explicitly stated in the proof.

While we have relaxed the geometric ergodicity condition for this case, in the proof of Theorem 4.2, we show that the Markov chain ( θ k ) k ≥ 0 is simply ergodic, i.e., lim k →∞ W p ( ν k , ν ∞ ) = 0 . Hence, even though we still obtain a time-uniform bound, our bound holds asymptotically in k , due to the lack of an explicit convergence rate for W p ( ν k , ν ∞ ) . On the other hand, the lack of strong convexity here results in a generalization bound with rate n -1 /p , whereas for the strongly convex case, i.e., p = 2 , we previously obtained a rate of n -1 . This might be an indicator that there might be still room for improvement in terms of the rate, at least for this class of loss functions.

5 Conclusion

We proved time-uniform Wasserstein-stability bounds for SGD and its noisy versions under different strongly convex, convex, and non-convex classes of functions. By making a connection to Markov chain perturbation results [RS18], we introduced a three-step guideline for proving stability bounds for stochastic optimizers. As this approach required geometric ergodicity, we finally relaxed this condition and proved two other stability bounds for a large class of loss functions.

The main limitation of our approach is that it requires Lipschitz surrogate loss functions, as it is based on the Wasserstein distance. Hence, our natural next step will be to extend our analysis without such a requirement. Finally, due to the theoretical nature of this study, it does not contain any direct potential societal impacts.

Acknowledgments

Lingjiong Zhu is partially supported by the grants NSF DMS-2053454, NSF DMS-2208303, and a Simons Foundation Collaboration Grant. Mert Gürbüzbalaban's research are supported in part by the grants Office of Naval Research Award Number N00014-21-1-2244, National Science Foundation (NSF) CCF-1814888, NSF DMS-2053485. Anant Raj is supported by the a Marie SklodowskaCurie Fellowship (project NN-OVEROPT 101030817). Umut ¸ Sim¸ sekli's research is supported by the French government under management of Agence Nationale de la Recherche as part of the 'Investissements d'avenir' program, reference ANR-19-P3IA-0001 (PRAIRIE 3IA Institute) and the European Research Council Starting Grant DYNASTY - 101039676.

References

[AFGO19] Necdet Serhat Aybat, Alireza Fallah, Mert Gurbuzbalaban, and Asuman Ozdaglar. A universally optimal multistage accelerated stochastic gradient method. In Advances in Neural Information Processing Systems , volume 32, 2019.

[AFGO20] Necdet Serhat Aybat, Alireza Fallah, Mert Gurbuzbalaban, and Asuman Ozdaglar. Robust accelerated gradient methods for smooth strongly convex functions. SIAM Journal on Optimization , 30(1):717-751, 2020.

[AS23] Shunta Akiyama and Taiji Suzuki. Excess risk of two-layer ReLU neural networks in teacher-student settings and its superiority to kernel methods. In International Conference on Learning Representations , 2023.

[Bac14] Francis Bach. Adaptivity of averaged stochastic gradient descent to local strong convexity for logistic regression. Journal of Machine Learning Research , 15(1):595-627, 2014.

[BE02] Olivier Bousquet and André Elisseeff. Stability and generalization. Journal of Machine Learning Research , 2(Mar):499-526, 2002.

[Ber15] Dimitri Bertsekas. Convex Optimization Algorithms . Athena Scientific, 2015.

[BFGT20] Raef Bassily, Vitaly Feldman, Cristóbal Guzmán, and Kunal Talwar. Stability of stochastic gradient descent on nonsmooth convex losses. In Advances in Neural Information Processing Systems , volume 33, pages 4381-4391, 2020.

[BM13] Francis Bach and Eric Moulines. Non-strongly-convex smooth stochastic approximation with convergence rate O (1 /n ) . In Advances in Neural Information Processing Systems , volume 26, 2013.

[CGZ19] Bugra Can, Mert Gürbüzbalaban, and Lingjiong Zhu. Accelerated linear convergence of stochastic momentum methods in Wasserstein distances. In International Conference on Machine Learning , pages 891-901. PMLR, 2019.

[CJY18] Yuansi Chen, Chi Jin, and Bin Yu. Stability and convergence trade-off of iterative optimization algorithms. arXiv preprint arXiv:1804.01619 , 2018.

[Dun81] Joseph C Dunn. Global and asymptotic convergence rate estimates for a class of projected gradient processes. SIAM Journal on Control and Optimization , 19(3):368400, 1981.

[EHZ22] Murat A. Erdogdu, Rasa Hosseinzadeh, and Matthew S. Zhang. Convergence of Langevin Monte Carlo in Chi-squred and Rényi divergence. In Proceedings of the 25th International Conference on Artificial Intelligence and Statistics (AISTATS) , volume 151. PMLR, 2022.

[FGO + 22] Alireza Fallah, Mert Gürbüzbalaban, Asuman Ozdaglar, Umut ¸ Sim¸ sekli, and Lingjiong Zhu. Robust distributed accelerated stochastic gradient methods for multi-agent networks. Journal of Machine Learning Research , 23(1):9893-9988, 2022.

[FR21] Tyler Farghly and Patrick Rebeschini. Time-independent generalization bounds for SGLD in non-convex settings. In Advances in Neural Information Processing Systems , volume 34, pages 19836-19846, 2021.

[FV19] Vitaly Feldman and Jan Vondrak. High probability generalization bounds for uniformly stable algorithms with nearly optimal rate. In Conference on Learning Theory , pages 1270-1279. PMLR, 2019.

[GG23] Guillaume Garrigos and Robert M Gower. Handbook of convergence theorems for (stochastic) gradient methods. arXiv preprint arXiv:2301.11235 , 2023.

[GGZ22] Xuefeng Gao, Mert Gürbüzbalaban, and Lingjiong Zhu. Global convergence of stochastic gradient Hamiltonian Monte Carlo for nonconvex stochastic optimization: Nonasymptotic performance bounds and momentum-based acceleration. Operations Research , 70(5):2931-2947, 2022.

[GRZ22] Mert Gürbüzbalaban, Andrzej Ruszczy´ nski, and Landi Zhu. A stochastic subgradient method for distributionally robust non-convex and non-smooth learning. Journal of Optimization Theory and Applications , 194(3):1014-1041, 2022.

[HM11] Martin Hairer and Jonathan C. Mattingly. Yet another look at Harris' ergodic theorem for Markov chains. In Seminar on Stochastic Analysis, Random Fields and Applications VI , pages 109-118, Basel, 2011.

[HRS16] Moritz Hardt, Ben Recht, and Yoram Singer. Train faster, generalize better: Stability of stochastic gradient descent. In International Conference on Machine Learning , pages 1225-1234. PMLR, 2016.

[KWS23] Leo Kozachkov, Patrick M Wensing, and Jean-Jacques Slotine. Generalization as dynamical robustness-The role of Riemannian contraction in supervised learning. Transactions on Machine Learning Research , 4:1-25, 2023.

[LGY20] Yanli Liu, Yuan Gao, and Wotao Yin. An improved analysis of stochastic gradient descent with momentum. In Advances in Neural Information Processing Systems , volume 33, pages 18261-18271, 2020.

[LLNT17] Tongliang Liu, Gábor Lugosi, Gergely Neu, and Dacheng Tao. Algorithmic stability and hypothesis complexity. In International Conference on Machine Learning , pages 2159-2167. PMLR, 2017.

[LLQ19] Jian Li, Xuanyuan Luo, and Mingda Qiao. On generalization error bounds of noisy gradient methods for non-convex learning. arXiv preprint arXiv:1902.00621 , 2019.

[LR17] Junhong Lin and Lorenzo Rosasco. Optimal rates for multi-pass stochastic gradient methods. Journal of Machine Learning Research , 18(1):3375-3421, 2017.

[LRP16] Laurent Lessard, Benjamin Recht, and Andrew Packard. Analysis and design of optimization algorithms via integral quadratic constraints. SIAM Journal on Optimization , 26(1):57-95, 2016.

[LY20] Yunwen Lei and Yiming Ying. Fine-grained analysis of stability and generalization for stochastic gradient descent. In International Conference on Machine Learning , volume 119, pages 5809-5819. PMLR, 2020.

[MT93] Sean P. Meyn and Richard L. Tweedie. Markov Chains and Stochastic Stability . Communications and Control Engineering Series. Springer-Verlag, London, 1993.

[MT94] Sean P. Meyn and Richard L. Tweedie. Computable bounds for geometric convergence rates of Markov chains. Annals of Applied Probability , 4(4):981-1011, 1994.

[MWZZ18] Wenlong Mou, Liwei Wang, Xiyu Zhai, and Kai Zheng. Generalization bounds of SGLD for non-convex learning: Two theoretical viewpoints. In Conference on Learning Theory , pages 605-638. PMLR, 2018.

[PVRB18] Loucas Pillaud-Vivien, Alessandro Rudi, and Francis Bach. Statistical optimality of stochastic gradient descent on hard learning problems through multiple passes. In Advances in Neural Information Processing Systems , volume 31, 2018.

[RBG + 23] Anant Raj, Melih Barsbey, Mert Gürbüzbalaban, Lingjiong Zhu, and Umut ¸ Sim¸ sekli. Algorithmic stability of heavy-tailed stochastic gradient descent on least squares. In International Conference on Algorithmic Learning Theory , volume 201, pages 12921342. PMLR, 2023.

[RRT + 16] Maxim Raginsky, Alexander Rakhlin, Matthew Tsao, Yihong Wu, and Aolin Xu. Information-theoretic analysis of stability and bias of learning algorithms. In 2016 IEEE Information Theory Workshop (ITW) , pages 26-30. IEEE, 2016.

[RRT17] Maxim Raginsky, Alexander Rakhlin, and Matus Telgarsky. Non-convex learning via stochastic gradient Langevin dynamics: A nonasymptotic analysis. In Conference on Learning Theory , pages 1674-1703. PMLR, 2017.

[RS18] Daniel Rudolf and Nikolaus Schweizer. Perturbation theory for Markov chains via Wasserstein distance. Bernoulli , 24(4A):2610-2639, 2018.

[RZG¸ S23] Anant Raj, Lingjiong Zhu, Mert Gürbüzbalaban, and Umut ¸ Sim¸ sekli. Algorithmic stability of heavy-tailed SGD with general loss functions. In International Conference on Machine Learning , volume 202, pages 28578-28597. PMLR, 2023.

[Vil09] Cédric Villani. Optimal Transport: Old and New . Springer, Berlin, 2009.

[WT11] Max Welling and Yee W Teh. Bayesian learning via stochastic gradient Langevin dynamics. In Proceedings of the 28th International Conference on Machine Learning (ICML-11) , pages 681-688, 2011.

[XR17] Aolin Xu and Maxim Raginsky. Information-theoretic analysis of generalization capability of learning algorithms. In Advances in Neural Information Processing Systems , volume 30, 2017.

[ZZB + 22] Yikai Zhang, Wenjia Zhang, Sammy Bald, Vamsi Pingali, Chao Chen, and Mayank Goswami. Stability of SGD: Tightness analysis and improved bounds. In Uncertainty in Artificial Intelligence , pages 2364-2373. PMLR, 2022.

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

APPENDIX

The Appendix is organized as follows:

In Section A, we provide further details and examples about the usage of surrogate losses.

In Section B, we provide technical background for the computable bounds for the convergence of Markov chains which will be used to prove the results in Section 3.3 in the main paper.

In Section C, we provide technical proofs for 1-Wasserstein perturbation results for the quadratic loss in Section 3.1 in the main paper.

In Section D, we provide technical proofs for 1-Wasserstein perturbation results for the strongly-convex loss in Section 3.2 in the main paper.

In Section E, we provide technical proofs for 1-Wasserstein perturbation results for the non-convex loss (with additive noise) in Section 3.3 in the main paper.

In Section F, we provide technical proofs for 2 -Wasserstein stability bounds for the nonconvex loss without additive noise in Section 4.1 in the main paper.

In Section G, we provide technical proofs for p -Wasserstein stability bounds for the convex loss with additional geometric structure in Section 4.2 in the main paper.

A On the Usage of Surrogate Losses

While the requirement of surrogate losses is a drawback of our framework, nevertheless our setup can cover several practical settings. In this section, we will provide two such examples.

Example 1. We can choose the surrogate loss as the truncated loss , such that:

where C > 0 is a chosen constant. This can be seen as a 'robust' version of the original loss, which has been widely used in robust optimization and is conceptually similar to adding a projection step to the optimizer.

Example 2. Another natural setup for our framework is the ℓ 2 -regularized Lipschitz loss that was also used in [FR21]. As opposed to the previous case, for the sake of this example, let us consider ℓ as the true loss and f as the surrogate loss. Then, we can choose the pair f and ℓ as follows:

where µ > 0 . Intuitively, this setting means that, we have a true loss ℓ which can be Lipschitz, but in the optimization framework we consider a regularized version of the loss. In other words, we have a loss ℓ ; however, we run the algorithm on the regularized loss f to have better convergence properties, and finally, we would like to understand if the algorithm generalizes on ℓ or not, and we are typically not interested if the algorithm generalizes well on the regularized loss f .

Next, we illustrate how a generalization bound for the loss f , i.e., ∣ ∣ ∣ E [ ˆ F ( θ ) -F ( θ )] ∣ ∣ ∣ . For this example, a bound on the quantity can be obtained by building on our analysis. To obtain such a bound, in addition to the bounds that we developed on ∣ ∣ ∣ E [ ˆ R ( θ ) -R ( θ )] ∣ ∣ ∣ , we would need to estimate the following quantity:

For illustration purposes, assume that ℓ is convex and Lipschitz in the first parameter. Then, f is µ -strongly convex. Further consider that we initialize SGD from 0 , i.e., θ 0 = 0 and set the batch

size b to 1. Denote θ = θ k as the k -th iterate of SGD when applied on ˆ F ( θ, X n ) , i.e., (1.1). Further define the minimum:

We can now analyze the error induced by the surrogate loss as follows:

Here, the second inequality follows from standard convergence analysis for SGD [GG23, Theorem 5.7] and we define σ X n as the stochastic gradient noise variance:

where for a random vector V we define Var[ V ] := E ∥ V -E [ V ] ∥ 2 . Hence, we can see that the error induced by the surrogate loss depends on the following factors:

The regularization parameter µ ,

The expected norm of the minimizers,

The step-size η ,

The expected stochastic gradient noise variance.

These terms can be controlled by adjusting µ and η .

B Technical Background

B.1 Computable bounds for the convergence of Markov chains

Geometric ergodicity and convergence rate of Markov chains has been well studied in the literature [MT93, MT94, HM11]. In this section, we state a result from [HM11] that provides an explicitly computable bound on the Wasserstein contraction for the Markov chains that satisfies a drift condition that relies on the construction of an appropriate Lyapunov function and a minorization condition.

Let P ( θ, · ) be a Markov transition kernel for a Markov chain ( θ k ) on R d . For any measurable function φ : R d → [0 , + ∞ ] , we define:

Assumption B.1 (Drift Condition) . There exists a function V : R d → [0 , ∞ ) and some constants K ≥ 0 and γ ∈ (0 , 1) so that

for all θ ∈ R d .

Assumption B.2 (Minorization Condition) . There exists some constant ˆ η ∈ (0 , 1) and a probability measure ν so that

for some R > 2 K/ (1 -γ ) .

We define the weighted total variation distance:

where ψ > 0 and V ( θ ) is the Lyapunov function that satisfies the drift condition (Assumption B.1). It is known that d ψ has the following alternative expression [HM11]:

where ∥ · ∥ ψ is the weighted supremum norm such that for any ψ > 0 :

It is also noted in [HM11] that d ψ has yet another equivalent expression:

̸

Lemma B.1 (Theorem 1.3. [HM11]) . If the drift condition (Assumption B.1) and minorization condition (Assumption B.2) hold, then there exists ¯ η ∈ (0 , 1) and ψ > 0 so that

for any probability measures µ 1 , µ 2 on R d . In particular, for any η 0 ∈ (0 , ˆ η ) and γ 0 ∈ ( γ +2 K/R, 1) one can choose ψ = η 0 /K and ¯ η = (1 -(ˆ η -η 0 )) ∨ (2 + Rψγ 0 ) / (2 + Rψ ) .

C Proofs of Wasserstein Perturbation Results: Quadratic Case

C.1 Proof of Lemma 3.1

Proof. Let P k ( θ, · ) denote the law of θ k starting with θ 0 = θ and P k ( ˜ θ, · ) the law of ˜ θ k :

where with ˜ θ 0 = ˜ θ . Note that

which implies that

By iterating over j = k, k -1 , . . . 1 , we conclude that

This completes the proof.

C.2 Proof of Lemma 3.2

Proof. First, we recall that

where ˆ H k := ∑ i ∈ Ω k ˆ a i ˆ a ⊤ i and ˆ q k := ∑ i ∈ Ω k ˆ a i ˆ y i . Therefore, starting with ˆ θ 0 = θ , we have

which implies that

This completes the proof.

C.3 Proof of Lemma 3.3

Proof. Let us recall that

which implies that

̸

Since X n and ˆ X n differ by at most one element and sup x ∈X ∥ x ∥ ≤ D for some D < ∞ , we have ( H 1 , q 1 ) = ( ˆ H 1 , q 1 ) with probability n -b n and ( H 1 , q 1 ) = ( ˆ H 1 , q 1 ) with probability b n and moreover

and

Hence, we conclude that

This completes the proof.

D Proofs of Wasserstein Perturbation Results: Strongly Convex Case

In order to obtain the algorithmic stability bound, that is a 1 -Wasserstein distance between the distribution of θ k and ˆ θ k , we need to establish a sequence of technical lemmas. First, we show a 1 -Wasserstein contraction rate in the following lemma.

Lemma D.1. Assume that Assumption 3.1 and Assumption 3.2 hold, and further assume that η < min { 1 µ , µ K 2 1 } . Then, for any n ∈ N ,

Proof. Let P k ( θ, · ) denote the law of θ k starting with θ 0 = θ :

and P k ( ˜ θ, · ) the law of ˜ θ k :

with ˜ θ 0 = ˜ θ . Note that

Therefore, we have

By applying Assumption 3.1 and Assumption 3.2, we get

provided that η ≤ µ K 2 1 . By iterating over k = n, n -1 , . . . 1 , we conclude that

This completes the proof.

Next, we construct a Lyapunov function ˆ V and obtain a drift condition for the SGD ( ˆ θ k ) ∞ k =0 .

LemmaD.2. Assume that Assumption 3.1 and Assumption 3.2 hold. Let ˆ V ( θ ) := 1+ ∥ θ -ˆ θ ∗ ∥ 2 , where ˆ θ ∗ is the minimizer of ˆ F ( θ, ˆ X n ) := 1 n ∑ n i =1 ∇ f ( θ, ˆ x i ) . Assume that η < min { 1 µ , µ K 2 1 +64 D 2 K 2 2 } and sup x ∈X ∥ x ∥ ≤ D for some D < ∞ . Then, we have

Proof. First, we recall that

Therefore, starting with ˆ θ 0 = θ , we have

Moreover, we have

This implies that

We can compute that

Moreover, we can compute that

Hence, we conclude that

provided that η ≤ K 2 +64 D 2 K 2 . This completes the proof.

µ 1 2

Next, we estimate the perturbation gap based on the Lyapunov function ˆ V .

Lemma D.3. Assume that Assumption 3.1 holds. Assume that sup x ∈X ∥ x ∥ ≤ D for some D < ∞ . Then, we have

where ˆ θ ∗ is the minimizer of ˆ F ( θ, ˆ X n ) := 1 n ∑ n i =1 ∇ f ( θ, ˆ x i ) .

Proof. Let us recall that

which implies that

̸

Since X n and ˆ X n differ by at most one element and sup x ∈X ∥ x ∥ ≤ D for some D < ∞ , we have x i = ˆ x i for any i ∈ Ω 1 with probability n -b n and x i = ˆ x i for exactly one i ∈ Ω 1 with probability b n and therefore

Hence, we conclude that

This completes the proof.

Next, let us provide a technical lemma that upper bounds the norm of θ ∗ and θ ∗ , which are the minimizers of ˆ F ( θ, X n ) := 1 n ∑ n i =1 ∇ f ( θ, x i ) and ˆ F ( θ, ˆ X n ) := 1 n ∑ n i =1 ∇ f ( θ, ˆ x i ) respectively. Lemma D.4. Under Assumption 3.2, we have ∥ θ ∗ ∥ ≤ 1 µ sup x ∈X ∥∇ f (0 , x ) ∥ and ∥ ˆ θ ∗ ∥ ≤ 1 µ sup x ∈X ∥∇ f (0 , x ) ∥ .

Proof. Since f ( θ, x ) is µ -strongly convex in θ for every x ∈ X , we have

which implies that

which yields that ∥ θ ∗ ∥ ≤ 1 µ sup x ∈X ∥∇ f (0 , x ) ∥ . Similarly, one can show that ∥ ˆ θ ∗ ∥ ≤ 1 µ sup x ∈X ∥∇ f (0 , x ) ∥ . This completes the proof.

D.1 Proof of Theorem 3.2

Proof. By applying Lemma D.1, Lemma D.2, Lemma D.3 and Lemma 2.1, we obtain

ˆ

where ˆ θ ∗ is the minimizer of ˆ F ( θ, ˆ X n ) := 1 n ∑ n i =1 ∇ f ( θ, ˆ x i ) . Finally, ∥ θ -ˆ θ ∗ ∥ 2 ≤ 2 ∥ θ ∥ 2 +2 ∥ ˆ θ ∗ ∥ 2 and by applying Lemma D.4, we complete the proof.

E Proofs of Wasserstein Perturbation Bounds: Non-Convex Case

Lemma E.1. Assume that Assumption 3.1, Assumption 3.3 and Assumption 3.4 hold. For any ˆ η ∈ (0 , 1) . Define M > 0 so that ∫ ∥ θ 1 ∥≤ M p ( θ ∗ , θ 1 ) dθ 1 ≥ √ ˆ η and any R > 0 so that R > 2 K 0 m and

where

Then, for any n ∈ N ,

for any θ, ˜ θ in R d , where for any η 0 ∈ (0 , ˆ η ) and γ 0 ∈ ( 1 -mη + 2 ηK 0 R , 1 ) one can choose ψ = η 0 ηK 0 and ¯ η = (1 -(ˆ η -η 0 )) ∨ 2+ Rψγ 0 2+ Rψ and d ψ is the weighted total variation distance defined in Section B.1.

Proof. Our proof relies on a computable bound on the Wasserstein contraction for the Markov chains by [HM11] that satisfies a drift condition (Assumption B.1) that relies on the construction of an appropriate Lyapunov function and a minorization condition (Assumption B.2).

By applying Lemma E.3, we can immediately show that the following drift condition holds. Let V ( θ ) := 1 + ∥ θ -θ ∗ ∥ 2 , where θ ∗ is the minimizer of 1 n ∑ n i =1 ∇ f ( θ, x i ) . Assume that η < min { 1 m , m K 1 +64 D 2 K 2 2 } and sup x ∈X ∥ x ∥ ≤ D for some D < ∞ . Then, we have

where

Thus, the drift condition (Assumption B.1) holds.

Next, let us show that the minorization condition (Assumption B.2) also holds. Let us recall that

We denote p ( θ, θ 1 ) the probability density function of θ 1 with the emphasis on the dependence on the initial point θ . Then, to check that the minorization condition (Assumption B.2) holds, it suffices to show that there exists some constant ˆ η ∈ (0 , 1)

for some R > 2 ηK 0 1 -(1 -mη ) = 2 K 0 m , where q ( θ 1 ) is the density function of a probability distribution function on R d , and (E.7) follows from Lemma E.2. Hence, by Lemma B.1, we have

for any θ, ˜ θ in R d , where for any η 0 ∈ (0 , ˆ η ) and γ 0 ∈ ( 1 -mη + 2 ηK 0 R , 1 ) one can choose ψ = η 0 ηK 0 and ¯ η = (1 -(ˆ η -η 0 )) ∨ 2+ Rψγ 0 2+ Rψ .

Finally, by the Kantorovich-Rubinstein duality for the Wasserstein metric, we get for any two probability measures µ 1 , µ 2 on R d :

Hence, we conclude that

This completes the proof.

The proof of Lemma E.1 relies on the following technical lemma, which is a reformulation of Lemma 35 in [CGZ19] that helps establish the minorization condition (Assumption B.2).

Lemma E.2. For any ˆ η ∈ (0 , 1) and M > 0 so that ∫ ∥ θ 1 -θ ∗ ∥≤ M p ( θ ∗ , θ 1 ) dθ 1 ≥ √ ˆ η and any R > 0 so that

Then, we have where

Proof. The proof is an adaptation of the proof of Lemma 35 in [CGZ19]. Let us take:

Then, it is clear that q ( θ 1 ) is a probability density function on R d . It follows that (E.7) automatically holds for ∥ θ 1 -θ ∗ ∥ > M . Thus, we only need to show that (E.7) holds for ∥ θ 1 -θ ∗ ∥ ≤ M . Since ξ 1 has a continuous density, p ( θ, θ 1 ) is continuous in both θ and θ 1 . Fix M , by continuity of p ( θ, θ 1 ) in both θ and θ 1 , there exists some η ′ ∈ (0 , 1) such that uniformly in ∥ θ 1 -θ ∗ ∥ ≤ M ,

where we can take

In particular, for any fixed ˆ η , we can take M > 0 such that

and with fixed η and M , we take R > 0 such that uniformly in ∥ θ 1 -θ ∗ ∥ ≤ M ,

This completes the proof.

Next, we construct a Lyapunov function ˆ V and obtain a drift condition for the SGD ( ˆ θ k ) ∞ k =0 .

Lemma E.3. Assume that Assumption 3.1, Assumption 3.3 and Assumption 3.4 hold. Let ˆ V ( θ ) := 1 + ∥ θ -ˆ θ ∗ ∥ 2 , where ˆ θ ∗ is the minimizer of ˆ F ( θ, ˆ X n ) := 1 n ∑ n i =1 ∇ f ( θ, ˆ x i ) . Assume that η < min { 1 m , m K 2 1 +64 D 2 K 2 2 } and sup x ∈X ∥ x ∥ ≤ D for some D < ∞ . Then, we have

Proof. First, we recall that

Therefore, starting with ˆ θ 0 = θ , we have

Moreover, we have

This implies that

We can compute that

Moreover, by following the same arguments as in the proof of Lemma D.2, we have

Hence, we conclude that

provided that K 1 +64 D K 2

Next, we estimate the perturbation gap based on the Lyapunov function ˆ V .

Lemma E.4. Assume that Assumption 3.1 and Assumption 3.4 hold. Assume that sup x ∈X ∥ x ∥ ≤ D for some D < ∞ . Then, we have

where ˆ θ ∗ is the minimizer of 1 n ∑ n i =1 ∇ f ( θ, ˆ x i ) and θ ∗ is the minimizer of 1 n ∑ n i =1 ∇ f ( θ, x i ) and d ψ is the weighted total variation distance defined in Section B.1.

Proof. Let us recall that which implies that

where E ξ 1 denotes expectations w.r.t. ξ 1 only, and we can further compute that

Therefore, we have

Similarly, we have

̸

Since X n and ˆ X n differ by at most one element and sup x ∈X ∥ x ∥ ≤ D for some D < ∞ , we have x i = ˆ x i for any i ∈ Ω 1 with probability n -b n and x i = ˆ x i for exactly one i ∈ Ω 1 with probability b n . Therefore, we have

Hence, we conclude that

This completes the proof.

It is worth noting that the Wasserstein contraction bound we obtained in Lemma E.1 in the non-convex case differs from the one we obtained in the strongly-convex case (Lemma D.1) in the sense that the right hand side of (E.3) is no longer ∥ θ -˜ θ ∥ so that Lemma 2.1 is not directly applicable. Instead, in the following, we will provide a modification of Lemma 2.1, which will be used in proving Theorem 3.3 in this paper. The definitions of the notations used in the following lemma can be found in Section 2.2.

Lemma E.5. Assume that there exist some ρ ∈ [0 , 1) and C ∈ (0 , ∞ ) such that

̸

for any n ∈ N . Further assume that there exist some δ ∈ (0 , 1) and L ∈ (0 , ∞ ) and a measurable Lyapunov function ˆ V : R d → [1 , ∞ ) of ˆ P such that for any θ ∈ R d :

Then, we have where

Proof. The proof is based on the modification of the proof of Lemma 2.1 (Theorem 3.1 in [RS18]). By induction we have

We have

Moreover, for any i = 0 , 1 , 2 , . . . ,

so that we obtain d ψ (˜ p i P, ˜ p i ˜ P ) ≤ γκ . Therefore, we have

By the triangle inequality of the Wasserstein distance, we have

This completes the proof.

Next, let us provide a technical lemma that upper bounds the norm of θ ∗ and ˆ θ ∗ , which are the minimizers of ˆ F ( θ, X n ) := 1 n ∑ n i =1 ∇ f ( θ, x i ) and ˆ F ( θ, ˆ X n ) := 1 n ∑ n i =1 ∇ f ( θ, ˆ x i ) respectively.

Lemma E.6. Under Assumption 3.3, we have

Proof. By Assumption 3.3, we have

which implies that

which yields that

Similarly, one can show that

This completes the proof.

E.1 Proof of Theorem 3.3

Before going to the proof, let us restate the full version of Theorem 3.3 that we provide below.

Theorem E.1 (Complete Theorem 3.3 ) . Let θ 0 = ˆ θ 0 = θ . Assume that Assumption 3.1, Assumption 3.3 and Assumption 3.4 hold. We also assume that η < min { 1 m , m K 2 1 +64 D 2 K 2 2 } and sup x ∈X ∥ x ∥ ≤ D for some D < ∞ and sup x ∈X ∥∇ f (0 , x ) ∥ ≤ E for some E < ∞ . For any ˆ η ∈ (0 , 1) . Define M > 0 so that ∫ ∥ θ 1 -θ ∗ ∥≤ M p ( θ ∗ , θ 1 ) dθ 1 ≥ √ ˆ η and any R > 2 K 0 m where K 0 is defined in (E.2) so that

Let ν k and ˆ ν k denote the distributions of θ k and ˆ θ k respectively. Then, we have W 1 ( ν k , ˆ ν k )

where for any η 0 ∈ (0 , ˆ η ) and γ 0 ∈ ( 1 -mη + 2 ηK 0 R , 1 ) one can choose ψ = η 0 ηK 0 and ¯ η = (1 -(ˆ η -η 0 )) ∨ 2+ Rψγ 0 2+ Rψ and d ψ is the weighted total variation distance defined in Section B.1.

Proof. By applying Lemma E.1, Lemma E.3, Lemma E.4 and Lemma E.5, which is a modification of Lemma 2.1, we obtain

where for any η 0 ∈ (0 , ˆ η ) and γ 0 ∈ ( 1 -mη + 2 ηK 0 R , 1 ) one can choose ψ = η 0 ηK 0 and ¯ η = (1 -(ˆ η -η 0 )) ∨ 2+ Rψγ 0 2+ Rψ and d ψ is the weighted total variation distance defined in Section B.1.

Finally, let us notice that ∥ θ ∗ -ˆ θ ∗ ∥ 2 ≤ 2 ∥ θ ∗ ∥ 2 +2 ∥ ˆ θ ∗ ∥ 2 and ∥ θ -ˆ θ ∗ ∥ 2 ≤ 2 ∥ θ ∥ 2 +2 ∥ ˆ θ ∗ ∥ 2 and for every x ∈ X ,

∥∇ f ( θ ∗ , x ) ∥ 2 ≤ 2 ∥∇ f (0 , x ) ∥ 2 +2 ∥∇ f ( θ ∗ , x ) -∇ f (0 , x ) ∥ 2 ≤ 2 ∥∇ f (0 , x ) ∥ 2 +2 K 2 1 ∥ θ ∗ ∥ 2 . (E.52) By applying Lemma E.6, we complete the proof.

E.2 Proof of Corollary 3.5

Proof. Under our assumptions, the noise ξ k are i.i.d. Gaussian N (0 , Σ) so that E [ ∥ ξ 1 ∥ 2 ] = tr (Σ) = σ 2 . Moreoever, we have Σ ≺ I d . Then p ( θ ∗ , θ 1 ) is the probability density function of

Therefore,

Notice that for any Ω 1 ,

Thus, we have

Since ξ ∼ N (0 , Σ) and Σ ≺ I d , for any γ ≤ 1 2 , we have

By Chebychev's inequality, letting γ = 1 2 , for any M ≥ η sup x ∈X ∥∇ f ( θ ∗ , x ) ∥ , we get

Next, for any θ, θ 1 ∈ R d such that ∥ θ 1 -θ ∗ ∥ ≤ M and ∥ θ -θ ∗ ∥ 2 ≤ 2 K 0 m (1 + ϵ ) -1 , we have

where E Ω 1 denotes the expectation w.r.t. Ω 1 , and p Ω 1 denotes the probability density function conditional on Ω 1 . For any given Ω 1 , we can compute that

where

Therefore, for any θ, θ 1 ∈ R d such that ∥ θ 1 -θ ∗ ∥ ≤ M and ∥ θ -θ ∗ ∥ 2 ≤ 2 K 0 m (1 + ϵ ) -1 , we have

We can further compute that

and

Hence, we have

Since it holds for every Ω 1 , we have

Hence, we conclude that

This completes the proof.

F Proofs of Non-Convex Case without Additive Noise

F.1 Proof of Theorem 4.1

Proof. Let us recall that θ 0 = ˆ θ 0 = θ and for any k ∈ N ,

Thus it follows that

where

This implies that

By Assumption 3.1 and Assumption 3.3, we have

provided that η ≤ m K 2 1 .

Since X n and ˆ X n differ by at most one element and sup x ∈X ∥ x ∥ ≤ D for some D < ∞ , we have x i = ˆ x i for any i ∈ Ω k with probability n -b n and x i = ˆ x i for exactly one i ∈ Ω k with probability b n and therefore

̸

and moreover,

where we used the inequality that

for any a, b ∈ R . Therefore, we have

In Lemma E.3, we showed that under the assumption η < min { 1 m , m K 2 1 +64 D 2 K 2 2 } and sup x ∈X ∥ x ∥ ≤ D for some D < ∞ , we have that for every k ∈ N ,

where ˆ V ( θ ) := 1 + ∥ θ -ˆ θ ∗ ∥ 2 , where ˆ θ ∗ is the minimizer of ˆ F ( θ, ˆ X n ) := 1 n ∑ n i =1 ∇ f ( θ, ˆ x i ) . This implies that

so that

where we applied Lemma E.6. Similarly, we can show that

Since θ 0 = ˆ θ 0 = θ , it follows from (F.10), (F.13) and (F.14) that

provided that η < min { m K 2 1 , 1 m , m K 2 1 +64 D 2 K 2 2 } . This completes the proof.

G Proofs of Convex Case with Additional Geometric Structure

In the following technical lemma, we show that the p -th moment of θ k and ˆ θ k can be bounded in the following sense.

Lemma G.1. Let θ 0 = ˆ θ 0 = θ . Suppose Assumption 4.1 and Assumption 4.2 hold and η ≤ µ K 2 1 +2 p +4 D 2 K 2 2 . Then, we have

where ˆ θ ∗ is the minimizer of 1 n ∑ n i =1 ∇ f ( θ, ˆ x i ) and θ ∗ is the minimizer of 1 n ∑ n i =1 ∇ f ( θ, x i ) .

Proof. First, we recall that

Therefore, we have

Moreover, we have

This implies that

We can compute that

Moreover, we can compute that

Hence, by applying (G.7) and (G.8) to (G.6), we conclude that

and hence

Similarly, we can show that

This completes the proof.

Next, let us provide a technical lemma that upper bounds the norm of θ ∗ and ˆ θ ∗ , which are the minimizers of ˆ F ( θ, X n ) := 1 n ∑ n i =1 ∇ f ( θ, x i ) and ˆ F ( θ, ˆ X n ) := 1 n ∑ n i =1 ∇ f ( θ, ˆ x i ) respectively.

Lemma G.2. Under Assumption 4.1, we have

and

Proof. Under Assumption 4.1, we have

where p ∈ (1 , 2) , which implies that

which yields that

Similarly, one can show that

This completes the proof.

Now, we are able to state the main result for the Wasserstein algorithmic stability.

Theorem G.1. Let θ 0 = ˆ θ 0 = θ . Suppose Assumption 4.1 and Assumption 4.2 hold and η ≤ µ K 2 1 +2 p +4 D 2 K 2 2 and sup x ∈X ∥ x ∥ ≤ D for some D < ∞ and sup x ∈X ∥∇ f (0 , x ) ∥ ≤ E for some E < ∞ . Let ν k and ˆ ν k denote the distributions of θ k and ˆ θ k respectively. Then, we have

Proof. Let us recall that θ 0 = ˆ θ 0 = θ and for any k ∈ N ,

Thus it follows that

where

This implies that

By Assumption 4.2 and Assumption 4.1, we have

Since X n and ˆ X n differ by at most one element and sup x ∈X ∥ x ∥ ≤ D for some D < ∞ , we have x i = ˆ x i for any i ∈ Ω k with probability n -b n and x i = ˆ x i for exactly one i ∈ Ω k with probability b n and therefore

̸

and moreover,

Notice that for any x, y ≥ 0 and p ∈ (1 , 2) , we have xy p -1 ≤ x p + y p , y ≤ y p +1 , x ≤ x p +1 and y p -1 ≤ y p +1 , which implies that

Therefore, by applying (G.24) to (G.23), we have

Hence, by applying (G.21), (G.22), (G.25) into (G.20), we conclude that

provided that η ≤ µ K 2 1 . This, together with θ 0 = ˆ θ 0 = θ , implies that

In Lemma G.1, we showed that under the assumption η ≤ µ K 2 1 +2 p +4 D 2 K 2 2 , we have

Hence, by plugging (G.28) and (G.29) into (G.27), we conclude that

Moreover, ∥ θ -ˆ θ ∗ ∥ 2 ≤ 2 ∥ θ ∥ 2 +2 ∥ ˆ θ ∗ ∥ 2 , ∥ θ -θ ∗ ∥ 2 ≤ 2 ∥ θ ∥ 2 +2 ∥ θ ∗ ∥ 2 and by applying Lemma G.2, we obtain

Finally, by the definition of p -Wasserstein distance, we have

This completes the proof.

G.1 Proof of Theorem 4.2

Proof. First, let us show that the sequences θ k and ˆ θ k are ergodic.

First, let us show that the limit θ ∞ , if exists, is unique. Consider two sequences θ (1) k and θ (2) k starting at θ (1) 0 and θ (2) 0 respectively with two limits θ (1) ∞ and θ (2) ∞ . For any k ∈ N ,

By Assumption 4.2 and Assumption 4.1, we have

provided that η ≤ µ K 2 1 . This implies that

provided that η ≤ µ K 2 1 . Thus, ∥ ∥ ∥ θ (2) k -θ (1) k ∥ ∥ ∥ 2 ≤ ∥ ∥ ∥ θ (2) k -1 -θ (1) k -1 ∥ ∥ ∥ 2 for any k ∈ N . Suppose ∥ ∥ ∥ θ (2) j -1 -θ (1) j -1 ∥ ∥ ∥ 2 ≥ 1 for every j = 1 , 2 , . . . , k , then we have

such that

Since p ∈ (1 , 2) ,

for any k ≥ k 0 +1 , which implies that θ (2) k -θ (1) k → 0 as k →∞ so that θ (2) ∞ = θ (1) ∞ . Next, let us show that for any sequence θ k , it converges to a limit. It follows from (G.39) that

for any k ≥ k 0 +1 . Let θ (1) 0 be a fixed initial value in R d , and let θ (2) 0 = θ (1) 1 which is random yet takes only finitely many values given θ (1) 0 so that k 0 is bounded. Therefore, it follows from (G.40) that

where ν k denotes the distribution of θ k , which implies that

Thus, ( ν k ) is a Cauchy sequence in P 2 ( R d ) equipped with metric W 2 and hence there exists some ν ∞ such that W 2 ( ν k , ν ∞ ) → 0 as k →∞ .

Hence, we showed that the sequence θ k is ergodic. Similarly, we can show that the sequence ˆ θ k is ergodic.

Finally, by ergodic theorem and Fatou's lemma, we have

We can then apply (G.31) from the proof of Theorem G.1 to obtain:

where

This completes the proof.

Peer reviews:
comments: SUMMARY:
This paper provides new stability bounds for SGD: they rely on perturbation theory for Markov Chains and on the ergodicity of SGD.

STRENGTHS:
- Well written, well presented paper.
- New connection between the theory for Markov Chain and generalization of learning algorithms.
- First (to my knowledge) uniform in time stability bounds.
- Extensions and limits are properly discussed.

WEAKNESSES:
N/A

QUESTIONS:
- On the fact that most of the theory requires $\ell$ to be Lispchitz-continuous: would it be because this work is only considering the 1-Waserstein distance? With the 2-Wasserstein, maybe this technique would be able to handle $L$-smooth loss, what do you think? In particular, could you leverage Thm. 4.2 in that case?

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the algorithmic stability of SGD in order to bound its expected generalization error when using Lipschitz losses. To do so, they consider Wasserstein stability instead of the standard uniform stability, which is defined using the dual representation of the Wasserstein distance and the representation of the generalization error from [HRS16].\n\nTo bound the Wasserstein stability the paper employs Markov chain perturbation theory.  The main idea to do so is to follow three steps: (i) showing that the optimizer is geometrically ergodic, which essentially boils down to showing that the Wasserstein distance of the transition kernels can be bounded from above by the distance of the inputs of said kernels; (ii) obtaining a Lyapunov function for the optimizer and the loss, which gives a sense of stability to the system (at least intuitively, as it is a necessary and sufficient condition for stability for certain classes of ODEs); and (iii) bounding the Wasserstein distance between the transition kernels of the optimizer with the same initialization but using two neighbouring datasets.\n\nUsing this three-step process they start giving a bound for the quadratic case with rate $O(\\eta / n)$. Similarly, under strong convexity and some pseudo-Lipschitzness assumption on the gradients they show a bound with rate $O(1/n)$ with an appropriate choice of the learning rate. To deal with non-convex cases, they show that if the loss dissipates with Assumption 3.3 (which seems like a relaxed version of strong convexity), then adding a zero-mean noise with bounded second moment is sufficient to obtain a bound with rate $O(b/n)$, where $b$ is the batch size. This bound of course includes SGLD, but is not restricted to this case.\n\nThe addition of the noise in the non-convex case was to make sure that the ergodicity condition in step (i) was provable. The paper also studies what happens when noise is not included. To deal with this, still assuming that the loss dissipates with Assumption 3.3 they can bound the squared Wasserstein distance of order 2. This way, they can obtain bounds with rate $O(1/\\sqrt{bn} + c)$, where $c$ is some constant. Therefore, the exclusion of noise in the analysis generates an extra bias term.\n\nAll the bounds up to here are uniform in time, that is, they hold for every step of the optimization process.\n\nFortunately, assuming that the loss $f$ is $\\mu$ strongly-ish convex, in the sense that $\\langle \\nabla f(\\theta, x) - \\nabla f(\\vartheta, x) , \\theta - \\vartheta \\rangle \\geq \\mu \\lVert \\theta - \\vartheta \\rVert^p$ for some $p \\in (1,2)$; and that the gradients Lipschitz-ish in the sense that $\\lVert \\nabla f (\\theta, x) - \\nabla (\\vartheta, y) \\rVert \\leq K\\_1 \\lVert \\theta - \\vartheta \\rVert^{p / 2} + K\\_2 \\lVert x - y \\rVert ( \\lVert \\theta \\rVert^{p-1} + \\lVert \\vartheta \\rVert^{p-1} + 1)$; then they can show asymptotic generalization bounds with rate $O(1/n^{1/p})$.\n\nSTRENGTHS:\nThe paper is very well written and \"easy\" to follow despite its technical difficulty, especially the proofs. I have to congratulate the authors for making difficult proofs look simple with their explanations. \n\nI believe that studying the stability of algorithms in terms of the stability of the final transition kernels in terms of the Wasserstein distance is novel (although similar bounds on the generalization error are given in [RGBTS21, Theorem 2] in a different context). This is certainly interesting as not only gives some new results but also new strategies to tackle the problem. These new strategies can be useful to build upon for other researchers in the future.\n\nThe bounds for the strongly-convex functions and those for non-convex losses with additive noise pass the \"eye test\", as they achieve the expected rates with $n$, although some of the other parameters are a bit unclear (see weaknesses).\n\nThe bounds for non-convex losses with additive noise with bounded variance are interesting, as they hold for noises other than Gaussian. This can be helpful either (i) to construct SGLD-like algorithms that work better in practice or (ii) to study more general situations in which one could show that the gradient noise in GD actually behaves like an additive noise with bounded variance (see e.g. [WHXHBZ20]).\n\n**Additional references**\n\nBorja Rodríguez-Gálvez, Germán Bassi, Ragnar Thobaben, and Mikael Skoglund. \"Tighter expected generalization error bounds via Wasserstein distance\". NeurIPS 2021.\n\nJingfeng Wu, Wenqing Hu, Haoyi Xiong, Jun Huan, Vladimir Braverman, and Zhanxing Zhu. \"On the Noisy Gradient Descent that Generalizes as SGD\". ICML 2020.\n\nWEAKNESSES:\n* A criticism of the paper could be that the results from Section 3 recover the rate of previous results with respect to $n$, while having sometimes worse rate with other parameters. This is actually not very important in my opinion, as I value more the new techniques and the fact that this new analysis can be employed in different settings to achieve the desired rate.\n\n* Something more concerning to me are some of the assumptions and parameters in the results. These are sometimes hard for me to interpret and contextualize. For example:\n  * Assumption 3.1. The authors explain that some problems such as GLMs satisfy this condition [Bac14]. Which other problems satisfy the condition? For example, the Lipschitz constant of gradients in neural networks can be very large (or even infinite in some situations), yielding bounds that depend on it essentially vacuous. Is this similar to this pseudo-Lipschitzness condition? How does it behave in familiar scenarios? This can be accompanied by a simplification of Theorem 3.2 in some familiar settings.\n* Assumption 3.3. This assumption seems to me as essentially a strong-convexity assumption with some slack. While this makes the problem more general than a strongly convex problem, it is unclear to me how much more general it is. I would have liked more intuition and explanation around this assumption other than just saying that is common in stochastic analysis [RRT17, GGZZ22] and that has been considered in neural network settings [AS22]. Could you give some examples of familiar settings where this assumption holds, and give an intuitive explanation of how much more general than strong convexity this is?\n* Theorem 3.3. Some of the elements such as the parameter $M$ and $\\bar{\\eta}$ are hard to into context and in general the bound is hard to interpret. Even though the authors make an effort to clarify that in Remark 3.4 and Corollary 3.5, it is still difficult to understand well the result in terms of these parameters.\n* Theorem 4.1. How big is the bias term in certain familiar scenarios? Could it be that his term is small enough for us not to care in some settings? \n* Assumptions 4.1 and 4.2. Could you please give some examples of losses and settings where these assumptions are satisfied? This could help contextualize the results.\n\nWhile I already like the paper and its ideas as it is, without a better understanding of these issues, it is hard for me to assess the actual contribution of the presented results. I am happy to increase my score if they are addressed.\n\nQUESTIONS:\nMy questions mainly concern the assumptions employed throughout the text and some of the parameters appearing in the results. These questions are mainly outlined in the weaknesses, but to summarize:\n* Could you give more intuitions on the assumptions that are employed (e.g. how much more general Assumption 3.3. is with respect to strong convexity)?\n* Could you give examples of losses and situations where these assumptions are satisfied so we can better contextualize the results?\n* Similarly, could you give some examples of the parameters (or some order of magnitude empirically maybe) that appear in the results?\n\nLIMITATIONS:\nYes, the authors discuss some of the limitations of the paper through the text. \n\nSome of the limitations regarding other assumptions could be clarified in the text.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors studied an approach to obtain generalization bounds on SGD algorithm under different class of objective functions such as \"quadratic strongly convex\", \"smooth strongly convex\" and one subclass of non-necessarily convex functions described in Assumption 3.3.\n\nInterestingly they convey the message that, as in the strongly convex settings, they can provide generalization bounds in the non convex setting as soon as they add some additive noise to the SGD update.\n\nTheir approach is based on stability analysis, together with convergence properties of Markov chains.\n\nSTRENGTHS:\n- Mixing stability and Markov chain perturbation theory is a very interesting approach that I had never seen before. The technique might be reused in other settings.\n- The trick of adding noise to create generalization and make it work thanks to Markov chain theory is quite beautiful.\n\nWEAKNESSES:\n- Authors claim \"non convex\" everywhere without more precision. I first thought of conditions like hypoconvexity. Some people study smooth optimization without any other assumption. Here the setting is very specific and authors should be more precise about it in introduction. One cannot claim solving problems in \"non-convex\" setting just as soon as its class contains one non-convex function.\n- It is actually quite natural to believe that adding noise reducing generalization error as the dependency in data is reduced (same argument for differential privacy). But we can do it with an algorithm that simply returns a constant independently of the training data. The interesting thing to do is to reduce generalization error while preserving good optimization properties. In this paper, the authors are degrading optimization performance as a constant step-size eta leads to convergence to neighborhood of the solution, and adding noise leads to a larger size of this neighborhood. This brings two questions:\n    - Can we generalize this analysis to decreasing step size as the ergocity would not be geometric anymore?\n    - In case of constant step-size, can authors discuss the values of « optimization error + generalization error » terms depending on the additive noise variance?\n\n- l.291: « This indicates that it is necessary to add additional noise » -> No! While this section can serve as a good intuition of why additive noise what essential, this is not a clear proof of this necessity. Authors showed they found some weaker upper bound without additive noise. They did not show one cannot find a tighter one.\n\n- Notation suggestions:\n    - The word « algorithm » is never defined before being used in Definition 2.5. Since in definition 2.1, $X$ and $\\hat{X}$ are not stochastic (instead we optimize on their values), it is even more important to stress here that an algorithm is indeed stochastic and that the expectation here is taken over this source of randomness conditioned by the values of $X$, $\\hat{X}$ and $z$. However, the notation $\\mathcal{A}: \\bigcup_{n=1}^{\\infty}\\mathcal{X}^n \\rightarrow \\mathbb{R}^d$ means that $\\mathcal{A}$ is deterministic. So either $\\mathcal{A}$ returns a random variable, or it takes an additional random variable as input.\n    - In 2.5, I find weird the fact to associate $\\nu$ and $\\hat{\\nu}$ to $X$ and $\\hat{X}$ as if those two datasets were particular, while we optimize over all the possible datasets. So $\\nu$ and $\\hat{\\nu}$ vary also. I would rather defined a systematic mapping $\\nu: X \\rightarrow \\nu_X$ so that this distribution is defined for all $X$. Or the variables over which the supremum is computed should be $\\nu$ and $\\hat{\\nu}$ that are neighbors to each other. In the latter case, this should come with the definition that two distributions are neighbors iif there exist two neighbor datasets they are generated from.\n\nMinor / typos:\n\nl.133: « $R(\\omega)$ » -> $R(\\theta)$\n\nl.229: I guess $\\xi_k$ is also independent with $\\theta_{k-1}$ and $\\Omega_k$? Maybe specify it.\n\nl252: « our result in » -> our result is\n\nl237: « Thereom » -> Theorem\n\nQUESTIONS:\n- Just to be sure I understood the proof correctly, can authors confirm the only place the neighborhood assumption between datasets is used is in the expression of $\\gamma$ in equation 2.9.  If this is the case, I think this little intuition should be explicit in the paper.\n- I am surprised to see in lemma 3.3 that the bound on $\\gamma$ does not depend on the batch size. Indeed, there is a certain probability that the batch does not contain the data that differs from one dataset to the other, in which case the step $P$ and $\\hat{P}$ are identical. Then, the error is 0. The average error should be the error in the case that the different data is in the batch, times this probability, the latter being proportional to the batch size.\n- The dependency on $\\sigma^2$ is unclear. In Th D1, the bound seems increasing in $\\sigma^2$, then removing the additive noise seems hurtless. However, looking closer, one notices that $\\psi$ depends on $K_0$ that itself depends on $\\sigma^2$. Why cannot we still consider $\\sigma^2=0$? The only thing we have to verify is that $\\psi$ will not explode, i.e. $K_0>0$. But $K_0$'s expression is a complicated mix between many other variables. Hence it is hard to understand what happens when $\\sigma^2=0$. Can you elucidate this question?\n\nLIMITATIONS:\nNot much except maybe it applies to algorithms that contracts, i.e. SGD with constant step-size that cannot converge to the optimum.  It would be nice to handle decreasing step-sizes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the generalization bounds of (noisy) SGD via Wasserstein stability. The paper presents a unified guideline to derive the Wasserstein stability for stochastic optimization with a constant step size, which allows to derive stability bounds with a three-step proof technique: showing the optimizer is geometrically ergodic, obtaining a Lyapunov function for the optimizer and the loss, and bounding the discrepancy between the Markov transition kernels associated with the chains. With this guideline, the paper derives stability bounds for SGD with strongly convex problems, nonconvex problems and a class between convex and strongly convex functions. The paper also develops stability bounds for noisy SGD in a nonconvex case.\n\nSTRENGTHS:\nThe paper presents a new perspective to derive stability bounds by techniques in applied probability. This connection between machine learning theory and applied probability is interesting and can have potential applications to various optimization algorithms.\n\nThe paper develops time-uniform stability bounds, meaning that the stability bounds would not increase to infinity as the number of iterations goes to infinity. The analysis also applies to general additive noise, which extends the existing analysis developed for Gaussian noise.\n\nThe paper also develops stability bounds for the standard SGD in the nonconvex case, although a dissipativity assumption is required.\n\nWEAKNESSES:\nThe analysis for SGD in Section 3.3 is a bit complicated. For example, the choice of $\\hat{\\eta}$ in Eq (3.10) is too complex. The result is not quite intuitive and it is not clear how we can use this bound to explain the behavior of SGD.\n\nThe bound in Theorem 3.2 is of the order larger than $O(1/(n\\mu^5))$. This bound has a crude dependency on $\\mu$. Since $\\mu$ is often very small in practice, the bound is worse than the existing bound of the order $O(1/(n\\mu))$ (HRS16).\n\nThe bound in Theorem 4.1 is vacuous due to the term $K/m$. From this stability bound, we cannot get meaningful generalization bounds.\n\nQUESTIONS:\nLemma 2.1 requires to build a Lyapunov function. The paper choose the function $V(\\theta)=1+|\\theta-\\theta_*|^2$. Can you explain why this is a good choice of Lyapunov function? What is the intuition behind it? What is the basic principle to build a Lyapunov function?\n\nThe paper considers noisy SGD for nonconvex case. Can you explain intuitively why noise is needed in this case? What is the benefit of adding noise in the nonconvex case?\n\n**Minor comments**:\n\nEq (2.6): should $\\mathbb{P}(\\theta_{n-1},A)$ be $P(\\theta_{n-1},A)$?\n\nEq (2.8): the meaning of $\\hat{P}\\hat{V}$ is not given. The meaning is given in the appendix.\n\nBelow Assumption A2: \"that that\"\n\nEq (D.21): $\\theta - \\hat{\\theta}_*+...$ \n\n should be $ \\theta - \\hat{\\theta}_*-...$. \n\nThe same change should be also made in Eq (D.22)\n\nLIMITATIONS:\nI do not see negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper derives Wasserstein stability bounds for a variety of cases for a \"surrogate loss\" under convex/non-convex settings.\n\nSTRENGTHS:\nThe problem is interesting. The bounds in the non-convex case can be useful. The paper is well-written.\n\nWEAKNESSES:\nI am not sure how novel the convex part is, also noted by authors. This part seems to remove the projection step, however, also done for a surrogate loss, so I am not sure if these are comparable. Also for the non-convex case, it is hard to see the relationship w.r.t. existing results due to the notion of surrogate loss.\n\nQUESTIONS:\nI have the following questions.\n\nOne thing I did not understand is the notion of the surrogate loss (as noted). The relationship between the original cost function $f$ and $\\ell$ is never clarified (I checked the appendix and no definition of this). It seems that the whole paper is written for $\\ell$, this is quite confusing. Why would authors not assume $f$ is $L$-Lipschitz and derive everything for $f$ under this restrictive assumption? This should be made clear as there is no way for a reader to understand what is going on in this paper, fundamentally, if the notion of the surrogate and its relationship to $f$ is not made clear at the beginning of this work. Note also that in Definition 2.1, it reads like the paper cited (HRS16) defines the stability for a surrogate loss, whereas there is no mention of surrogate loss in HRS16.\n\nUpon checking a few related work which uses surrogate loss notion, it seems that the convergence in function value may not happen here. This of course makes the comparison impossible with relevant results in the literature.\n\nDoes any surrogate loss work for these results? Do the results hold for a particular selection of losses, e.g., sometimes depends on $p$ as in earlier works? Please clarify.\n\nLIMITATIONS:\nTheory work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-24"}