Uniform-in-Time Wasserstein Stability Bounds for (Noisy) Stochastic Gradient Descent
Authors: Lingjiong Zhu
Mert Gurbuzbalaban
Anant Raj
Umut Simsekli
Venue: NeurIPS 2023 poster (2023)
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
Lingjiong Zhu', Mert Gürbüzbalaban"* Anant Raj45, Umut Şimşekli
I: Dept. of Mathematics, Florida State University
2: Dept. of Management Science and Information Systems, Rutgers Business School
3: Center for Statistics and Machine Learning. Princeton University
4: Coordinated Science Laboratory, University of Illinois Tirbana-Chappaig
5: Inria Paris, CNRS, Ecole Normale Supérieure, PSL. Research University
Abstract
Algorithmic stability is an important notion that has proven powerful for deriving generalization bounds for practical algorithms. The last decade has witnessed 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 unified guideline for proving Wasserstein stability bounds for stochastic optimization algorithms. We illustrate our approach on stochastic gradient descent (SGD) and we obtain time-unifone 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 extend them to more general cases by using 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 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 non-convex losses (without additional additive noise), which, to our knowledge, is novel.
bounds for SGD under convex and
Introduction
With the development of modern machine learning applications, understanding the generalization properties stochastic gradient descent (SGD) has become 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(0) == Eg-D [f (8, .r).. where r E X' denotes a random data point, D is the (unknown) data distribution defined on the data space X, denotes the parameter vector, and f ; Idd X
an instantaneous loss function.
In a practical setting, directly minimizing F(0) is not typically possible as D is unknown; yet one typically has access to finite data set X,, {.1;-..,22 _.,} E 2 where assume each r; is independent and identically distributed (i.i.d.) with the common distribution D. Hence, given X,,. one can then attempt to minimize the empirical risk F.),X.) = F(0). In this setting, SGD has been one of the most popular optimization algorithms for minimizing
Σn /00, vx) as proxy for
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

F(0) and is based on the following recursion:
(0 : 1 E Vf(0k-117.))
(1.1)
where v) is the step-size, is the batch-size, S2g is the minibatch that chosen randomly from the set (1,2....,71), and its cardinality satisfies [Stg] b.
One fruitful approach for estimating the population risk attained by SGD, i.c., F(Ox). is based on the following simple decomposition:
F(O,) S F(O2)) + | F(O+)):
(1.2)
where the last term is called the generalization error. Once a computable upper-bound for the generalization
be obtained, this decomposition directly leads computable upper bound for the population risk F(Ox), since F0x) can be computed thanks the availability of Xn. Hence, the challenge here is reduced to derive upper-bounde on [F(0)) - F(Ox)), typically referred as generalization bounds.
Among many approaches for deriving generalization bounds, algorithmic stability [B00 has been one of the most fruitful notions that have paved the way to numerous generalization bounds for stochastic optimization algorithms [HRS16 CJY18 MW7Z18 FV19. LY20 ZZB + 22). In nutshell, algorithmic stability measures how much the algorithm output differs if we replace one data point X,, with a new sample. More precisely, in the context of SGD, given another data set
X, = . } = (#1..... I;-1, 1,-,,,f;, Hi, I2+1.-... Zn E d that differs from X,, by most one element, we (theoretically) consider running SGD on X, i.e.,
i.e.
0. = - - (0
1 E V (1.3) b SE
and we are interested in the discrepancy between 02 and 0r in some precise sense (to be formally defined in the next section). The wisdom of algorithmic stability indicates that a smaller discrepancy between 0g and 0g 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 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 n= whereas for general convex and non-convex problems, their bounds diverged with the number of iterations (even with a projection step), unless decreasing step-size is used. In subsequent studies [LY20 K KWS33 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 recent study, [FGT220] 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.c., 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 [WTII] which is essentially 'noisy" version of
SGD that uses the following recursion: 0s = Oz-1 v/ Fx (O.-1, X,) + Ex. where (Ex))>0 a is sequence of i.i.d. Gaussian vectors, independent of f and f2g. The authors of [RRTI7] MWZZI8 proved stability bounds for SGLD for non-convex losses, which were then extended to more general (non-Gaud Gonn-Gusiin 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-unifone stability bounds for SGLD under convextty,,,,, indicating that, with the presence of additive Gaussian noise, better stability bounds can be achieved. Their time-unne uniform results were then extended to non- Gaussian,, heavy-tailed perturbations in RBG 23 RZG$23 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, 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 novel connection between learning theory and applied probability, and introduce unified guideline for proving Wasserstein stability bounds for stochastic optimization algorithms with constant step-size. Our approach is based on Markov chain perturbation theory (RSI8), which offers a three-step proof technique for deriving stability bounds: (i) showing the optimizer is geometrically ergodic, (ii) obtaining Lyapunov function for the optimizer and the loss, and (iii) bounding the discrepancy between the Markov transition kernels associated with the chains (x)):20 and (0& k22- We illustrate this approach on SGD and show that time- uume-unifom stability Loaudd---pschtz-iiik condition for smooth strongly-conve losses (we recover similar results to the ones of [HRST6]) and 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 terms of obtaining time-uniform bounds: 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 flexible and can be generalizable to other popular optimizers, it mainly requires developing Lyapunov functions, which are often readily available in the GRZ22 LGY20 AFGO19]
bounds can be obtained under
literature AFGO20 LRPI6 FGO 22
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 0o, 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 (0x):20, 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 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$23) 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 2 the p- W-sssessteinn distance between two probability measures μ and y on IRd is defined as [Vil09]:
Wp(,,v) = -X 1/3
(2.1)
where the infimum is taken over all couplings of X μ and Y In particular, the dual representation for the 1- I-Wasserstein distance is given as [Vi109)
sup 1/e Mx))(dx) - f M(r))(dd))) heLip(1)
(2.2)
where Lip(1) consists of the functions h Rd + 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 HRSI6. To begin, we will present the definition of algorithmic stability as stated in [HRS16]:
3

---

Definition 2.1 (HHRS16), Definition 2.1), Let RV(R") denote the set of Rf random vectors.
For a (surrogate) loss function RRd x X - R, an algorithm A U x + (RV(R)) is e-uniformly stable if
sup sup E [(((X)),) ((A((),,))) SE. x=x ZE
(2.3)
where the first supremum is taken over data X, x,X E x that differ by one element, denoted by X x X.
In this context, we purposefully employ a distinct notation for the loss function & (in contrast 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 I-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 derive Wasserstein bounds. Unfortunately, under these assumptions, we cannot further impose / itself to be Lipschitz, hence the need for surrogate losses. Nevertheless, the usage of surrogate losses common in learning theory, see e.g. [R21 RZG5523), and we present concrete practical examples in the Appendix.
Now, we present 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 the loss function as follows:
X.. 1 Σ ((0,z:), R(0) = EewD[ x)].
Theorem 2.1 (HRST6)) Theorem 2.2). Suppose that A is E-uniformly stable algorithm, then the expected generalization error is bounded by
EAX.. (A(A(Xa).X.) - R(A(X.)) S..
(2.4)
For randomized algorithm, v and denotes the law of A(X) and A(8) then for L-Liischitz surrogate loss function L, we have the following generalization error guarantee.
EA,x, [R(A(X.), X..) R(A(X.)) < L sup (i, xex The above result can be directly obtained from the combination of the results given in 2.2 Defini- tion 2.1 and Theorem 2.1
(2.5)
2.1 (see also (RRT* 16
2.2
Perturbation theory for Markov chains
Next, we recall the Wasserstein perturbation bound for Markov chains from RSIS Let (0,, m be Markov chain with transition kernel P and initial distribution po. i.e., have almost surely
P(O, Ale.. .... .0.1)) P(O,, e Al00--1 = P().-1,A).
(2.6)
and (AA) = I0o E A) for any measurable set A C RR and n E N. We assume that (0x) is another Markov chain with transition kernel 12 and initial distribution Po- We denote by Pr. the distribution of 0,, and by Pn. the distribution of 0,. By dg, denote the Dirac delta distribution 0, ie. the probability measure concentrated at 0. For a measurable set A C Rd, we also use the notation P(A) A).
Lemma 2.1 (RSI8), Theorem 3.1). Assume that there exist some E [0, 1) and C' E (0, oo) such that
Wi (P"
Pm Pm S SCp". 100 01
sup 0.0
(2.7)
for any N. Further assume that there exist some 8 (0,1) 1) and L (0, oo) and measurable Lyapunov function V ddd + [1, 00) of P such that for any E Rdd:
(PV)(0) 5 6V + L.
(2.8)

---

where (PV) =- Sa V(0) d0). Then, we have
Wi(pn, SC (D"Wi 1o) +1 (1- (1-10") 1
7k
(2.9)
where 7 = SuPaERd
W1(bgP.da V(0)
K: max { [I2- V(0)dpo(0), 1-6)
Lemma 2.1 provides a sufficient condition for the distributions pn and Pr. after n iterations to stay close to each other given the initial distributions Po- Lemma 2.1 will provide a key role in helping us derive the main results in Section 3.1 and Section 3.2 Later, 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 RSI88. Before considering general losses that can 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, start by considering quadratic loss of the form: f(0, (a, 0 2 /2 where, z; := (a,, and a0;, a, 0 y.). In this setting, the SGD recursion takes the following form:
0g - -
(1- Hm, Oh-1 + Ja.. where,
Ha = Σ aaa
Th = Σ a.l
(3.1)
iEll
iEf
The sequence (H&, qs) are i.i.d. and for every k (H&, 4s) is independent of 0
Similarly, we can write down the iterates of SGD with different data set Xx, == fd1,...,ff., with f (ās, Dt). where X, differs from Xn with most one element:
!1 - O2 (1 -1 H2x 02-1 7 a where 11 == Σ àzā
2
Σ a.h. Eanh: iE0g
(3.2)
Seft
Our goal is to obtain algorithmic stability bound, through estimating the 1-Wasserstein distance between the distribution of 0% and 0g 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 geometrie ergodicity of the Markov process (0z) with transition kernel P(0, ), given in the following lemma.
Lemma 3.1. Assume that = E H. || < 1. Then, for any l E N, we have the following
inequality: W1 (P* pk (0.-) S p² p/16 ||e
01|
We note that since H1 2 O. the assumption in Lemma 3.1 can be satisfied under mild assumptions, for example when H1 0 with positive probability, which is satisfied for n small enough.
In the second step, we construct a Lyapunov function V that satisfies the conditions of Lemma 2.1
Lemma 3.2. Let (0) = 1 + |0]]. Assume that P = 11|1 प z < 1. Then, we have
(PV)(0) - (0) (0) 1- 2
+ 11
(3.3)
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 [Bacl4 BM13].
Lemma 3.3. If suppex ||x||| < for some D oo, then, we have supgerd
W. P.8g V(0)
DnD
5

---

Note that Lemma 2.1 relies on three conditions: the Wasserstein contraction in (2.7h. which is obtained through Lemma 3.1 the drift condition for the Lyapunov function (2.8). which obtained in Lemma 3.2 and finally the estimate on in (2.9) which is about the one-step 1- 1-Wasserstein distance between two semi-groups that 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>ex ||=]|| & D) is used the expression of y in equation (2.9). Now, having all the ingredients, we can invoke Lemma 2.1 and we obtain the following result which provides I-Wasserstein bound between the distribution iterates when applied to datasets that differ by one point.
For Y E U X and k 2 0, let vYY, k) denote the law of the k-th the SGD iterate when Y is used as the dataset, i.e., (XX, k) and v(X, k) denote the distributions of Gg. and 0g obtained by the recursions (3.1 and (.2) respectively. As shorthand notation, set 2/- = v (xx, k)) and iys == v(X, k)).
Theorem 3.1. Assume 9, 0o 0. We also assume that E | o M1 || 1 and p
E T 2 H and supzex ||=||| - D foorr some D < D0. Then, have
Wi(uk, Dk) - 1 p² 22 max 1 |101|
1
(3.4)
p n
Proof. The result directly follows from Lemma
3.1 Lemma
3.2
Lemma
3.3 and Lemma
2.
By a direct application of 2.5 we can obtain a generalization bound for an L-I surrogate loss function, as follows:
EAXX, [R(A(X.), X.)) - R(A(X,)))
C
2qD² max 1 ||0||, 1 - Pu + 3
Po n
Po
where Po supx || 3 Hxll Hx = EiEn, ajaj and X is 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, 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 capacity condition [R17 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 E X, f') is differentiable.
Before proceeding the stability bound, we first introduce the following assumptions.
Assumption 3.1. There exist constants K1, K2 > O such that for any 0, E RRd and every x E X,
|VS x) (0(,,)) +) < K1 |0 0|| K2]1 4 || || + ||²|| +
(3.5)
This assumption is a pseudo-Lipschizz Lipschitz- condition on f and satisfied for various problems such 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 02,0, 02 E Rd and E X,
(VI - v.(02,)), 0 02) 2 m|00 02112.
By using the same recipe as we used for quadratic losses, we obtain the following stability result.
Theorem 3.2. Let Oo = do = 0. Assume that Assumption 3.1 and Assumption 3.2 hold. We also assume that 1j min 1 K7- }. supaex |>||| D for some D < cc, and
suPxEx (0)xx)) < E for some E < 00. Let uy and iy denote the distributions of 0% and

---

& respectively. Then, we have
Wi(v, 0k) = 8DK2(1 - (1- 2/ )k)
2F
+1
max
{
1 11+21101 +
2E
2 K?
56
D²
+
64n n
D2 D2K3 K3 E²
(3.6)
112
Similarly the quadratic case, we can now directly obtain a bound on expected generalization error using 2.5p. More precisely, for an L-Lipschitz surrogate loss function E, we have
X.)- SC- 8DKa(1 (1
24 )))
2E
E.X.
+1
nH
max 1 21 +
2E
2
7
K
56n
D2 +
D2
64n D2 K3 E²
μs
The bound above has the same dependence on the number of samples as the ones of the previous sta- bility analysis of (projected) SGD for strongly convex functions [HRST6 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 01, 0, Rd and E X,
(VI(01,)) - vS(0,, x), ), 01 02) 2 ml 02l - K.
The class of dissipative functions satisfying this assumption are the ones that admit some gradient growth in radial directions outside compact set. Inside the compact set though, they can have quite general non-convexity patterns. As concrete examples, they include certain onc-hidden-layer neural networks [AS23]: they arise in non-convex formulations of classification problems (c.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 ball of radius for some will satisfy this assumption. Consequently, regularized regression problems where the loss is a strongly convex quadratic plus smooth penalty that grows slower than quadratic will belong to this class; a concrete example would be smoothed Lasso regression; many other examples are also given [EH/22] Dissipative functions also arise frequently in the sampling and Bayesian learning and global convergence in non-convex optimization literature [RRT17] [RRT1]]
Unlike the strongly-conven case, synchronous coupling technique as we did in the proof of Theorem [3.2 To circumvent this problem, in this setting, consider noisy version of SGD, with the following recursion:
longer obtain Wasserstein contraction bound using the
02 = 0. (Oa-1,, + ne
1
Σ V((0R-1.21)
(3.7)
b
iEf
where Ea are additional i.i.d. random vectors in Rd, independent of 0 and 2g. satisfying the following assumption.
Assumption 3.4. &1 is random vector on Rd with a continuous density p(x) that is positive everywhere, iee. p(x) > for any x E IRd and [[g]] = 0, x² = E [1]] < oc.
Note that the SGLD algorithm WTII is 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,, as
Ô, = 0 1 X.) + T
Eien. vf(0,x;;) 7/1- Further let 0, be a minimizer of F(,, Xx). Then, by following the same three-step recipe, we obtain the following stability bound. Here, we do not provide the constants explicitly for the sake of clarity; the complete theorem statement is given in Theorem E.I (Appendix E.II
and let p(0,0,) 01) denote the probability density function of 0

---

Theorem 3.3. Let 0, 0g = 6. Assume that Assumption 3.1 Assumption 3.3 and Assumption 3.4 m
hold. We also assume that 1l < min { m. K7+61D=Kg and sep || S D for some D oo and suPnex x0) I E for some E 0o. For any ij E (0, 1), define M 0 so that Jio,- 100 0. ISM p(O,, d1) 2 vnj and for any R >
2Ko where Ko is defined in (E.2) so that
01) 0,0, 0.0 (0)S - I.SM p(0,, 01)
inf
2 vT.
(3.8)
Let Up and by denote the distributions of 02 and 0g respectively, Then, we have
Ci(1 17/)
2b
WI(Vk, S 2/1(11+ + 15 - 17)
(3.9)
where for any E. (0, f)) and "o (1 (1-ma) + 2n R
1).
nRo and îj = (1- )o)) Reno
2+R
The constani C1 Ci (v),, 9i R, r0 K1, K2, D, E) is explicitly stated in the proof.
Contrary to our previous results, the proof technique for showing Wasserstein contraction (as in Lemma 3.I) for this theorem relies verifying the drift condition (Assumption [B.1] and the minorization condition (Assumption [.2 as given in [НМТ Once these conditions are satisfied, we invoke the explicitly computable bounds on the convergence of Markov chains developed in [HM1T].
From equation 2.5 we directly obtain the following generalization error bound for C-Lipschitz surrogate loss function,
EAXX. [((A(X)). -R(AXXX))))))(((()))))))))))) <C
Ci - R7* 2/0( )1 7) n
26
where the constants are defined in Theorem 3.3 The above result can be directly compared with the result in [FR21] Theorem 4.1] ]] that has the same dependence n and b. However, our result is more general in the sense that we do not assume our noise to 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
2Kg (1 + e) for some fixed E (0, so we can take
I
p(0, 01)
max MS0
min
p(O., O2 d2)0
inf 0,0; ER 01
(3.10)
110 1181-M.11sM
0- 0.115M
p(0., 01)
10-0.11 21 2N0 (1+r)-1
Moreover, we can take 170
2. 70
mne
2
and sb = 2nko' so that
1 2 (1+c) (1
mae
mare (1 e)
1j = max
2
(3.11)
2' 2 (1++)
4m + 2(1 + e)ti'
m
provided that 1) S 1. Note that the parameter à in (3.10) appears in the upper bound in equation that controls the 1- W-sssersteinn algorithmic stability of the SGD. It is easy to see from equation that the smaller ii, the smaller the 1-Wasserstein bound. By the defintion of īJ, the larger íi, the smaller the 1-Wasserstein bound. As a result, woudd like to choose to be as large as possible, and the equation (3100) provides an explicit value that can take, which is already the largest as possible.
3.9
3.9
Next, let us provide some explicitly computable lower bounds for in (.10). This is achievable if we specify further the noise assumption. Under the assumption that & are i.i.d. Gaussian distributed, have the following corollary.
By using the decomposition (1.2), we can obtain excess risk bounds for SGLD by combining our results with [XR17: it was shown that gradient Langevin dynamics has the following optimization error is O(e + 2/22 1/4 X log 1/s) after K Odde X log 1/e) iterations, where & is the mini batch size and is the uniform spectral gap of the continuous-tine time Langevin dynamics. Similar results are given for SGLD in XR17 Theorem 3.6].

---

Moreover, discarding the persistent bias term, this bound leads to generalization bound with rate n 1122, rather than n as before. This indicates that it is beneficial to add additional noise & in SGD as in Theorem 3.3 3 in order for the dynamics to be geometrically ergodic that can lead to a sharp bound as - co. Finally, we note that as Theorem 4.1 involves 2- 2-Wasserstein distance, it can pave the way for generalization bounds without requiring a surrogate loss. Yet, this not immediate and Corollary 3.5. Under the assumptions in Theorem 3.3 we further assume the noise Ex are i.i.d.
Gaussian (0, Σ) so that = o². We also assume that Σ La- Then, we have
min {
exp(- +
supnex [V. 2
(1-
fj 2
max M sunsuppcx VV 20),
V (I8 2)
(2Ko + 1) 1/2
(1 Kin)
|
exp
17
(1+ Kin)
2KK (1 +c) -1)
1/2
+2 +2(M
sup IV x)|l xE
(3.12)
The above corollary provides an explicit lower bound for (instead of the less transparent inequality constraints in Theorem (.3), and by combining with Remark 3.4 (see equation (3.IT) leads to an explicit formula for which is essential to characterize the Wasserstein upper bound in 3..9 in Theorem 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 the Wasserstein sense (i.e., proving contraction bounds), we have observed that such 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 [RST8), 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 (Assump-
3.44 be able to invoke Lemma 2.1 We will now show that possible to use a more direct approach to obtain 2- 2-Wasserstein algorithmic stability in the non-convex case under Assumption 3.3 without relying on [SS8]. However, we will observe that without geometric ergodicity will have
non-vanishing bias term in the bound. Note that, since Wi (vk, Dx) S Wp(u,, D&) for all p 2 1, the following bound still yields generalization bound by 22.
Theorem 4.1. Assume Og 0, = 0. We also assume that Assumption 3.1 and Assump-
tion 3.3 hold and 1 min
m K4+64 and supzex ||=||| S D for some D < and
supzex (.(0, x)ll < E for some E < 00. Let un and in denote the distributions of On and Og respectively. Then, we have
[4D*K3n(BB+ +
2K 4K₂ +Ki)) (1 (1+)) +
W2(uk. in) $ (1 111-1110p mym))) (1-1m))+))
bnm
nm
(4.1)
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 + co, 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 co, for moderate the bound can be still informative as the persistent term might be dominated by the first two terms.
would require deriving uniform _ bounds for the iterates, c.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 E (1,2) such that for any 0,, O2 E Rdd and x E X, (V](0), .²) v f02,, .r), 01 02) 2 p0|00 02 I..
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 Dun8 Ber15].
We analogously modify Assumption 3.1 and consider the following assumption.
Assumption 4,2. There exist constants K1, K2 > 0 and p E (1,2) such that for any 0, 0 E Rd and every x E X, [V x) V f(0, a)ll K 0 0 _ K21læ> 0|[(|001- + 10]pp-1 1).
The next theorem establishes a stability bound for the considered class of convex losses in the stationary regime of SGD.
Theorem 4.2. Let Bu = du Suppose Assumption 4.1 and Assumption 4.2 hold (with p E (1,2)) and 7) S K4+2p-1D2K3 and supzex ||=||| I D for some D < oc and supaex [V, x) S E
for some E' &C Then vs, and İy. converge to the unique stationary distributions Voo and & respectively and moreover, we have
C2 Css bn
S
+
(4.2)
n
where the constants C2 = C2(1 M4 K2, D, E) and C3 = Ca(n.u, K2, D, E) are explicilly stated in the proof.
While we have relaxed the geometric ergodicity condition for this case, in the proof of Theorem 4.2 show that the Markov chain (0s)s20 is simply ergodic, i.e., lima- (p,,,, M/x)) 0. Hence, even though we still obtain a time-uniform bound, our bound holds asymptotically in k: due to the lack of explicit convergence rate for (vk )o On the other hand, the lack of strong convexity here results in generalization bound with rate n -1/p, whereas for the strongly convex case, i.e.,
p 2, we previously obtained rate of n 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.
Conclusion
We proved time-uniform ceasserstein-taabllii bounds for SGD and its noisy versions under different strongly convex, convex, and non-convex classes of functions. By making connection to Markov chain perturbation results [RST&], we introduced a three-step guideline for proving stability bounds for stochastic optimizers. As this approach required geometrie 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 DMS-2053454, NSF DMS-2208303, and Simons Foundation Collaboration Grant. Mert Gürbüzbalaban's 's research are supported in part by the grants Office of Naval Research Award Number N00014-21-1-221
National Science Foundation Anant Raj is supported by the Marie Sklodowska-
(NSF) CCF- 18 14888, NSF Curie Fellowship (project NN-OVEROPT 1010308177 17). Umut Şimşekli's research supported by the French government under management of Agence Nationale de la Recherche as part of the "Investissements d' a'vveii"" program, reference
(PRAIRIE 31A Institute) and the European Research Council Starting Grant DYNASTY - 101039676.
10

---

References
[AFG019] Necdet Serhat Aybat, Alireza Fallah, Mert Gurbuzbalaban, and Asuman Ozdaglar. A universally optimal multistage accelerated stochastic gradient method. 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)7177-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 local strong con- vexity 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-522 2002.
[Ber15] Dimitri Bertsekas. Convex Optimization Algorithms. Athena Scientific, 2015.
[BFGT20] Raef Bassily, Vitaly Feldman, Cristóbal Guzmán, and Kunal Talwar. Stability of stochas- tic gradient descent on nonsmooth convex losses. In Advances in Neural Information Processing Systems, volume 33, pages 1381-4399, 2020.
[BM13] Francis Bach and Eric Moulines. Non-strongly-comwe smooth stochastic approximation with convergence rate ((1/n). In Advances Neural Information Processing Systems, volume 26, 2013.
[CGZ19] Bugra Can, Mert Gürbüzbalaban, and Lingjiong Zhu. Accelerated linear convergence of stochastic momentum methods 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:/804.016/9/ 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):368- 400. 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üizbalaban, Asuman Ozdaglar, Umut Şimşekli, and Lingjiong Zhu. Robust distributed accelerated stochastic gradient methods for multi multi-agent net- works. Journal of Machine Learning Research, 23(1)99893-9988 2022.
[FR21] Tyler Farghly and Patrick Rebeschini. Time-independent generalization bounds for SGLD non-convex settings. Advances in Neural Information Processing Systems, volume 34, pages 19836-18846, 2021.
[FV19] Vitaly Feldman and Jan Vondrak. High probability generalization bounds for uniformly stable algorithms with nearly optimal rate. In Conference 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 stochas- tic gradient Hamiltonian Monte Carlo for nonconvex stochastic optimization: Nonasymp- totic performance bounds and momentum- acceleration. Operations Research, 70(5):2931 -947,, 2022.
[GRZ22] Mert Gürrüüzalaaan,, Andrzej Ruszczyúiski, and Landi Zhu. stochastic subgradient method for distributionally robust non-conve, and non-smooth learning. Journal of Optimization Theory and Applications, 1944)11014-141, 2022.
[HM11] Martin Hairer and Jonathan C. Mattingly. Yet another look 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 dynam- ical 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 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 1902.00692..0611,, 2019.
[LR17] Junhong Lin and Lorenzo Rosasco. Optimal rates for multi-pase stochastic gradient methods. Journal of Machine Learning Research, 18(1):3375-342,, 2017.
[LRP16] Laurent Lessard, Benjamin Recht, and Andrew Packard. Analysis and design of opti- mization 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. Com- munications and Control Engineering Series. Springer- Verlag, London, 1993.
[MT94] Sean Meyn and Richard I.. 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 Şimşekli. Algorithmic stability of heavy-tailed stochastic gradient descent on least squares. In International Conference on Algorithmic Learning Theory, volume 201, pages 1292- 1342. PMLR, 2023.
[RRT+16] Maxim Raginsky, Alexander Rakhlin, Matthew Tsao, Yihong Wu, and Aolin Xu. Information- -hheoretic analysis stability and bias of learning algorithms. 2016 IEEE Information Theory Workshop (ITW), pages 26-30. IEEE, 2016.
12

---

[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. Bernoulii, 24(4A):2610-2633 2018.
[RZGŞ23] Anant Raj. Lingjiong Zhu, Mert Gürbüzbalaban, and Umut Şimşekli. Algorithmic stability of heavy-tailed SGD with general loss functions. In International Conference on Machine Learning, volume 202, pages 28578-28597. PMLR, 2023.
[Vi109] 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 (ICMI.-I1), pages 681-688, 2011.
[XR1]] Aolin Xu and Maxim Raginsky. aeformaaiinntheorrtti analysis of generalization capa- bility of learning algorithms. In Advances in Neural Information Processing Systems, volume 30, 2017.
[ZZB+22] 22] Yikai Zhang, Wenjia Zhang, Sammy Bald, Vamsi Pingali, Chao Chen, and Mayank Goswami. Stability of SGD: Tightness analysis and improved bounds. In Uncertainty Artificial Intelligence, pages 2364-2373. PMLR, 2022.
13

---

Uniform-in-Time Wasserstein Stability Bounds for (Noisy) Stochastic Gradient Descent
APPENDIX
The Appendix is organized follows:
Section
we provide further details and examples about the usage surrogate losses.
In Section B
we provide technical background for the computable bounds for the conver-
gence of Markov chains which will be used to prove the results in Section 3.3 in the main paper.
provide technical proofs for I-Wasserstein perturbation results for the quadratic loss in Section 3.1 in the main paper.
In Section
In Section D we provide technical proofs for I-Wasserstein perturbation results for the strongly-convex loss in Section 3.2 in the main paper.
In Section we provide technical proofs for 1- I-Wasserstein perturbation results for the non-convex loss (with additive noise) in Section 3.3 in the main paper.
In Section P provide technical proof's for 2-sssees stability bounds for the non- convex loss without additive noise in Section in the main paper.
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.
In Section G
A On the Usage of Surrogate Losses
While the requirement of surrogate losses drawback of our framework, nevertheless our setup can cover several practical settings. In this section, we will provide two such examples.
Example
We can choose the surrogate loss as the truncated loss, such that:
e(0, x) m(,,f(0,z), C*),
where C' > 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. also used [FR2T]. As opposed to the previous case, for the sake of this example, let us consider & choose the pair and as follows:
Another natural setup for framework is the éz-regularized Lipschitz loss that was
as the true loss and f as the surrogate loss. Then,
1 2 11112
f(0, x) = C x)
where y > O. Intuitively, this setting means that, we have a true loss which can be Lipschitz, but in the optimization framework we consider regularized version of the loss. In other words, have loss é; however, we run the algorithm on the regularized loss f to have better convergence properties, and finally, would like understand if the algorithm generalizes not interested if the algorithm generalizes well on the regularized loss f.
or not, and we typically
Next, we illustrate how a generalization bound for the loss f. i.c., ((F(() F(0) For this example, a bound on the quantity can be obtained by building our analysis. To obtain such bound, in addition to the bounds that we developed on [(A(() R(0)], we would need to estimate the following quantity:
Ea,X. -f( (0,,z)))
For illustration purposes, assume that is convex and Lipschitz the first parameter. Then, f μ-strongly convex. Further consider that we initialize SGD from 0, i... 9 0 and set the batch
14

---

size b to 1. Denote & = 02 as the k-th iterate of SGD when applied on x.,, i.e. define the minimum:
1.1
Further
Øx arg argmin F (0, Xn).
We can now analyze the error induced by the surrogate loss as follows:
n
(f (0, ((0,z,))
M "E, 18, x, 111111 2
1 Es Eg xx, 110 0x + 0x,
nB nm 110 8x 1? + pEx, 110% 11
(1- mm) 11 1|?
27 xx
nEx, n 0 1|
/
((1 m/) + 1) Ex.
|
1|
2nyEx, [ox..]-
Here, the second inequality follows from standard convergence analysis for SGD 5.7] and we define xx, as the stochastic gradient noise variance:
GG23 Theorem
X., == Var v (6x,. xx)]
where for a random vector V we define Var[V] E E[V]I|" Hence, we can see that the error induced by the surrogate loss depends on the following factors:
The regularization parameter μ4,
The expected norm of the minimizers,
The step-size 7),
The expected stochastic gradient noise variance.
These terms can be controlled by adjusting ja and 17.
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 [M In this section, we state result from [HMII) that provides an explicitly computable bound on the Wasserstein contraction for the Markov chains that satisfies a drift condition that relies the construction of an appropriate Lyapunov function and minorization condition.
Let P(0,) .) be a Markov transition kernel for Markov chain (0%) on Rd. For any measurable function 40 RRd
[0, + -20, we define:
(Pt) Jgd (00)p(0,,0)
Assumption B.1 (Drift Condition). There exists a function Rd + [0, 20) and some constants K 0 and y (0, 1) so that
(PV) SV((0) + K,
for all 0 Rdd Assumption B.2 (Minorization Condition). There exists some constant E (0, 1) and probability measure v so that
inf P(0,-) 2 în((), DERE.V(OSRR
for some R > 2K/(1-7)
We define the weighted total variation distance:
de (1n, 122 = / (1 (V+))()))lr - (z)(d)),
15

---

where vo > 0 and (0)) is the Lyapunov function that satisfies the drift condition (Assumption B.1 It is known that di has the following alternative expression [НМI]:
d = sup
(0))( p2)
eelllll Rd
where || Il is the weighted supremum norm such that for any &i >
[s)
I| sup 1 1+ (V)
PER4
is also noted in [MI that dy has yet another equivalent expression:
/ ((0)(u #2)(d0),
sup (plu, *:lleellos
Rd
where
I|lell == sup (x( - (001
040 2 + VV + (()
Lemma B.1 (Theorem 1.3. (НМIJ)) If the drift condition (Assumption B.1 and minorization
condition (Assumption B.2 hold, then there exists vj E (0, and +b > 0 so that das ( I (1d P2)
for any probability measures 11, on Rd. In particular, for any no E (0. f) and To E (Y+2K/R, one can choose do T// K and 1j (1 (n) no)) (2 + Ru)roo/ + Rob).
C Proofs of Wasserstein Perturbation Results: Quadratic Case
C.1 Proof of Lemma
Proof. Let pk (8,-) denote the law of 0g starting with 0g f and pk the law of Gs:
0:
T
"1 b
H
b 9k.
(C.1)
with 0y = 0. Note that
0
[
11 H 0k-1 +
b qk.
(C.2)
b
[
!
Hb)
0k-1
"
b 9k.
(C.3)
b
which implies that
E 0x
== 0=|=
Ōz
HR:
(0k- - 0z 01-1)11
< "1 H4 0k-1 - 0x-1 = =pF 0k
0x 02-1|1
(C.4)
b
By iterating over = - 11 ... 1, we conclude that
Wi (t p* (=(0))) El 0zl p² p01100 0o1l ,² 110
011
(C.5)
This completes the proof.
C.2 Proof of Lemma 3.2
Proof. First, we recall that
11 02 = (1- (1- 11x 0--1 f
(C.6)
where H. Eierr, a.a and dr = -ienn â.hi- Therefore, starting with O = 0. we have
0 11 H1
0 + b a1
(C.7)
which implies that (PV)(0) (0) 5 + ||
7 E (0) +1 p+ "
3 El ||.
(C.8)
This completes the proof.
16

---

C.3 Proof of Lemma 3.3
Proof. Let recall that
1 H1
01 =
0 +
b 4
(C.9)
b
01 b
H1
0 +
b n
(C.10)
which implies that
Wi 6g oo.)
S Hi H
7
01 + E |1 - qqill-
(C.11)
Since Xn, and Xn, differ by at most one element and sup ZE || - D for some D < oc, we have (1) (f11.0 91) with probability n-b and (H1, (H) * (H1. 91) with probability & and moreover n
n
E ||| H1 111
b max (Ila; [lā.1l?) -
2hD
max |a:a à, S n isi5h
(C.12)
T2 SS
n
and
21ll - b
b max (H|a||]lab|| ||@=|ll@@ell)
2bD (C.13)
E Bl11
max
|aath à.D:ll Isisn
isisn
n
Hence, we conclude that
W1 Wi
2 [110]] + 1) 26D2
2n
sup
sup BERd
(C.14)
AERd
V(0)
1+ |0
71
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 0g and Og. need to establish sequence of technical lemmas. First, show 1- l-Wasserstein contraction rate in the following lemma.
Lemma D.1. Assume that Assumption 3.7 and Assumption
3.2 hold, and further assume that
n < min { # Then, for any n N.
Wi (0,)) Pm (0.))) S 5 nu
110 - Ò||
(D.1)
Proof. Let (f,,.) denote the law of lg starting with On = A:
02 = 0 T/ F Xn),
(D.2)
and p'((0.)) -) the law of 0g:
02 0x-1 4 Ps (0k-1, Xn)
(D.3)
with 0u = 6. Note that
02 0-~1
V vf(0k-1,2s).
(D.4)
b
iEfu
0g 0-1
1
Σ
v (0k-1:211)
(D.5)
b
iEf
Therefore, we have
E
O2
0x
2
= E
02
0p-1
Σ
Vf(00-1, vf (0
b
V5(0k-112;) - v. (0x-1:23) 1
112
= E
0 Ōx-1
+
E
Σ
2n E
0k-1 Õn-1,
(v[(0--1,2)) v. v (0k-1.zi)) (D.6)
-b
iE
17

---

By applying Assumption 3.1 and Assumption 3.2
we get
E 8x
0x
$ (1 2npie
0k-1 8
+ 112
E
Σ
ff(Ok-1,
V
vf (0k-1, 24 ||]
62
2
0kk1 0z ||
$ (1 2n/ilE
Oh-1 xx-1
+ 2) K
S (1 T/1)BB (2-1 Oz
(D.7)
provided that 17 S #े.. By iterating over k= = 1,................. 1, conclude that
(m (O,) (0)))) - (ws (0)(0,.), pm (0.
E2 0n 0n 1
s (1 m1)" |00 16--00 -06l² $ 5(1=
7744
2n
0
0
2
(D.8)
This completes the proof.
Next, we construct a Lyapunov function V and obtain a drift condition for the SGD (0x )==
Lemma D.2. Assume that Assumption 3.1 and Assumption 3.2 hold. Ler (0) == 1 1+| -110--0,1 11? where
0, the minimizer of ((, : E. (f(0, 2)). Assume that n < min : k7+
and supnex |a|||| S D for some D oo. Then, we have
(PV)(0) S (1 ( + 20p u²² D2 D2
(D.9)
Proof. First, we recall that
03 = -
11 Σ Vf @0
(D.10)
b
iEs
Therefore, starting with 0u 6, we have
0 0
7
V ff 24)
b iE
1 Σ vv(
= 0
( +1|
v (0 2b)
(D.11)
1m1
SR
Moreover, we have
E
V f(( (f(0 2
1
V ff 2.).
(D.12)
b
This implies that
(P))
= =EV(O)) = 1 + E 8 8, ||
2
n
=1+
0, + 17
(0, 24))
+ +1²EE
v f(0 22) - Σ vf(0
(D.13)
7
iE
18

---

We can compute that
0
0.
+
11
V
(5(0
n 1=1
V5((0 - Vi (0., ft: )
0 0. 7
|
$ (1 - 20 - -0.13+
Σ
f() vs (8... 2;
n²
2=1
< (1 2n 12 0..
(D.14)
Moreover, compute that
n
1 Σ
E
((0,
V 2:)
=E
Σ
V [5 kt) - v5(0, -((0,.)))
b VEN2
b
iE
2
K₂ 1 Σ - 81,210111111
E
Σ
n
iE
<3 +12
S SDD2 (810-6.11* (810-00.1++88..11 +10.11+++1)
(D.15)
Hence, we conclude that
(P)) < (1 - 2np + k2² + 64n/ 64n D2K3)KKK) K3) 0 0. 22 1| 11 21+88/222K3 (8]10_11++ +1)
(1 2np n K + 64n D K3) ((0)
- 2nj 722 K - 56n D2 K2 + 64n K3
$ (1 - m/11 (0) 2np - 1/² K 56 D2 K2 + 64n D2 1²
(D.16)
provided that 17 S K? K7-61D=K3 This completes the proof.
Next, we estimate the perturbation gap based on the Lyapunov function V.
Lemma D.3, Assume that Assumption 3.1 holds. Assume that supzex || S D for some D < co. Then, we have
(6g
4DK29
sup
V(0)
(2|10_ |l + 1),
(D.17)
BER
where 0. is the minimizer of F(0, 8n) = 1/2 E (/(0
Proof. Let recall that
0 0 Σ (S(00 20),
(D.18)
ief
0, " T)())
(D.19)
19

---

which implies that
W1
5)P)
b
E
VS(0,0) - VS(0,
1 E
VV(00x))) vS(O,&))
b
LiEl1
!
E
Σ
K2 |,, 2.1l (2110) + +1
(D.20)
b
LiEO,
Since X,, and Xn differ by at most one element and supzex || I D for some D < oc, we have I; 2; for any i E St, with probability n- and x; + f; for exactly one E 22 with probability & and therefore
n
wi (ooP, (o0P,60P) P)
n b n <
2DK20 (21101 + 1).
+1)=
(D.21)
Hence, we conclude that
Wi (6.P66 V(0)
4DK27
|00 +
sup BER4
1 sup BER
1+ ||0 - 0
4DK29 2||0 - 0.11 + 21 +
sup BeRd
1+ ||0 11
4DK2n 2| - (2 01
sup BERd ni
1 1+ 0.112
4DK2T) [210.11+11
(D.22)
n
This completes the proof.
Next, let us provide a technical lemma that upper bounds the norm of 8, and 0., which are the minimizers of F(0, X,) 1/2 E. Vf mn) and (0, Xx) = 1/2 Σ (0(0, ..) respectively.
Lemma D.4. Under Assumption 3.2 we have ||0_|| I 1 supzex (f( x)]l and ||+.11 S
suPREx [VV
Proof. Since f(0, x) is stsoonngly convex for every x E X, we have
(VF(0, - (O., X,),0 - 0.) = 20)))x
(V5(0, Vf(0,, - -0.) 2 u|10.113;
n
(D.23)
which implies that 0110.1 2 < sup [0) ||0.
(D.24)
xEx
which yields that ||0.1|
supaex [VV x)]] Similarly, one show that ||²_]| <
suPRex [0, This completes the proof.
D.I Proof of Theorem 3.2
Proof. By applying Lemma D.1 Lemma D.2] Lemma D.3 and Lemma 2.1 we obtain
8DK2(1 - (1 = )k)
(x)
(2100_ +1)
nj
max 1 |00 0, ||
56 D² +
64n D2 D²K3310.13 (D.25)
2
K?
20

---

where 0, is the minimizer of :- := Vf(0,ū;). Finally, ||0- 1|0-0.11 22 < 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
f E (0,1). Define M > so that J
I p 0.. 01)d0, Z vn and any so that R 2Kg and
p(.
inf
2 vn.
(E.1)
0.0)ER::(0))SR..0
M&M p(0., 01)
where
Ko == 2m - nkk - K3 64nD2 10.12 2 2K + no²
(E.2)
Then, for any E N.
Wi
dd 6g).
(E.3)
2/
(1+1++++0 vb)
for any 0, 0 R. where for any no E (0, f) and %0 E (1 (1-m)++
2nKo . one can choose
and 1j (1 (n no)) V 2+Reno and dæ is the weighted total variation distance defined 2+
n.o
in Section
B.
Proof. Our proof relies computable bound the Wasserstein contraction for the Markov chains by [HM11 that satisfies a drift condition (Assumption B.T) that relies on the construction of an appropriate Lyapunov function and minorization condition (Assumption [B.2).
By applying Lemma
we can immediately show that the following drift condition holds. Let
((0) 1 + 10 min
where 0, is the minimizer of Ei (0,,0x)). Assume that v <
m Ki -1 and sup,ex ||>||| S D for some D < co. Then, have
(PV) S mn) V (0) + 1Ko,
(E.4)
where
Ko = 2m 2m nk 56n K3 64n 10.11 2K + 1102.
(E.5)
Thus, the drift condition (Assumption B.1 holds. Next, let show that the minorization condition (Assumption B.2 also holds. Let recall that
17 b 0, = 0- -
Σ vf(0,x)) - nEI.
(E.6)
We denote p(0, 0,) ) the probability density function of 0 with the emphasis on the dependence on initial point 0. Then, to check that the minorization condition (Assumption B.2 holds, it suffices to show that there exists some constant ij E (0, 1)
inf
p(, 02)) ())(0,),
for any 0, Rdd
(E.7)
(ERA:V(0)
Kn 2Kg where q(1)) is the density function a probability distribution for some 1-(1-10 ma)
function IR.. and (E.7) follows from Lemma E.2 Hence, by Lemma B.1 we have
do Pm (" (O) P" S 17th
for any 8, Ø Rd, where for any no E n) and %0 E (1- mn +
kn n 1)
one can choose
noo and (1 (n 70)) v 24
2+R
21

---

Finally, by the Seliirrovicc-Rubiiseee duality for the Wasserstein metric, we get for any two probability measures M1, f2 on F.
Wi(ni, sup
(1)) (2)(d Φ is 1-Lipschitz
Rd
sup
(Ø(0) ((0 p2)(d0) o is 1-Lipschitz
Rd
110
0- 01l1111 - p2l(d0)
2/ w1 + v) / (1 + vV())) p21(d0)
de (12,
(E.8)
2 (1++))
Hence, we conclude that
Wi
+1 (b) 7"du(50:00 6z).
2/
This completes the proof.
The proof of Lemma E. relies on the following technical lemma, which is a reformulation of Lemma 35 in [CGZI9] that helps establish the minorization condition (Assumption B.2)
Lemma E.2. For any r) E (0, 1) and M > 0 so that Jie. (10.-00111.0 P(O., 0d) z vi and any R > 0
so that
p(0,0)
inf (0) -0.1EMM p(0,, 01)
2 VT.
(E.9)
Then, we have
inf VER)V(
p(0,0)) 2 :)g(01),
for any 01 Rd.
(E.10)
where
1101-0.11:00
q(O1) = p(),,01)
(E.11)
SM p(O., 02 )d0;
Proof.
The proof is an adaptation of the proof of Lemma 35 in [CGZ19. Let us take:
Jje, l1e 1. 1
((0)) 01)
(E.12)
p p(0,, 0g )d0
Then, is clear that (01) is a probability density function Rdd It follows that (E.7) automatically holds for ||0 0. > M. Thus, we only need to show that (E.7) holds for ||0 8.1 I M. Since E has continuous density, p(0, 01) is continuous in both 0 and 0,. Fix M, by continuity of p(0,0)) 01) in both 0 and 01, there exists some '/' E (0, 1) such that uniformly in ||0= - 0.1l $ M,
inf BER(VVI0)R
p(0, 01) 2 v'p(0.,0 01) = (0)
(E.13)
where we can take
i 1== 1/'
J 1 p(0.,
(B.14)
In partieular, for any fixed ij, we can take M > 0 such that
Jo p(O., 01 )d0 2 Vh.
(F.15)
and with fixed y and M, we take R > 0 such that uniformly in ||0, 0.1 s M,
inf (ER(:V(0)
0(0,0,) 2 Vip(O..,0,)
(E.
This completes the proof.
22

---

Next, we construct a Lyapunov function V and obtain a drift condition for the SGD (0x) )0x))==
Lemma E.3. Assume that Assumption 3.1 Assumption 3.3 and Assumption 3.4 hold. Let V(0)= ==
1 + 0 0,1 where O. is the minimizer of F(0, X
Ei=1 (5 5.). Assume that TJ
min 1 m K7+6AD=K3 and supzex ||=||| D for some D co. Then, we have
(PV) S (1 (L-
56 D2 K3 + +4n D3 +1/ (E.17)
Proof.
First, we recall that
8. 6x_
Σ iEfl
Vf (Be-11.)) + ne..
(E.18)
b
Therefore, starting with 0g 0, we have
0 = -
V [00 ..) + 1
b iEN
7 0
ff 2.) + 17
1 EVV((0.))
-
1 Σ v)(),a)) )
+nE..
(E.19)
n
ign
Moreover, we have
V
ff(0, 2,) 0
V Σ
f( 2.).
(E.20)
b E0
71
This implies that
(P)
(V(O,) = 1 E. 11 0.] 1
1+ 10- 0. 2 (0, .3.)
V/(0, - 1 Σ
+ E
V 5:)
+1202
(F.21)
We can compute that
0
0
17
V f (6, 21))
i=
n
0 0. 17
Σ
Vi (0. 2x) Vs V5((,,
n i=1
< (1 2mm) |0 10-0. 0.1
2nk
n
Σ
n
(VS(0,2))-V (o.. 2
n²
< (1 2mm + n |000 0. 2nK.
(E.22)
Moreover, by following the same arguments as in the proof of Lemma D.2 we have
n 1 Σ 2
(8|10 à. 1² + =8100112 (...1++1).
+1). *+810_1++1)
E 2.)
5 EEA
SDD2 K3
(E.23)
n1
23

---

Hence, we conclude that
(P)) - (1 2mm + k2² 641) | 6.1"
+1+8 (8100,1 112 + +1) + 2n + 1/2
=(1 1-2mm- + K + 64nr (
+ 2 - k² 56n/ KK 2 + 1/²0²²
$ (1 mn) (00) + 20mm+-mm--- - K 56n +
1/002 (F.24)
provided that 7) S KT 6 This completes the proof.
Next, we estimate the perturbation gap based on the Lyapunov function V. Lemma E.4. Assume that Assumption 3.7 and Assumption 3.4 hold. Assume that suprex | for some D < oc. Then, we have
D
(6g
sup BER4
((0)
2b
vb (4 + 8172
max
n1
IV f(0,, 10, x)]13 } (E.25)
1 1+ v. (1+ (4 K - 0.12 - 40²² sup
|
ZEX
where 0, is the minimizer of Ei= (0,, and 0, is the minimizer of Σ vf) and dy is the weighted total variation distance defined in Section B.1
Proof. Let recall that
01 0
!
Σ
vf(0,2;))
nk
(E.26)
b
iER
0 = 8
"
f 2,) - nEI
(E.27)
b
iEQ
which implies that
[V(i)) 1 Ee, [100 0- 13] =1 172
0 0 7 Σ Vf
x;)
(E.28)
iEn
where E,, denotes expectations w.r.t. only, and can further compute that
0, 7 " SE
(0,
21²²
$
0.112
Σ
ff
6²
ienh
2 [Vf(0,z;)]|
211²
< 21110
0.1124
62
VEn
2
211²
< 21
0
[V5(0,x;)
62
Ven.
2n²
2110
1
bk1 0.11 + b sup || Vf x|]
62
#E
< 2110
0.1 2 + 4n
K 1
11 + 4n²
sup x
|| V f(e,, 1|
(E.29)
24

---

Therefore, we have
(V(0)) s 1 1 2110 0.11 4v2
0.11 + 41/
sup xEt
(E.30)
Similarly, have
3, [V(())) s =1+ + 0.12 K? 1.11 + 41
sup xEA
vf
(E.31)
Since Xn and Xn differ by at most one element and supzex || D for some D < oc, we have I₁ 2; for any i Si with probability n-b and # b. for exactly one E Ω with probability & Therefore, we have
de (oP.
2b
1
+ v (1+1 (2+ 4 0.11 + 4n² sup =EX
[VS(O,, (E.32)
Hence, we conclude that
dp( Po)P)
sup BER
V(0)
2b
1 1p 11+ (2 0.11 + 4uf² supzex [V.00,)))11))
sup
n BCR4
1 100 1.
1 + v (1 + 12000 + (4 8 k 0.113)
2b
sup
n
OER
1 110 0.11
$ (14+ (14+8) - 0.112 41j² supzex
[V[ w
1 10 0.112
2b
< max
244+ 877KK))
n
1 1 vb
1 111200² 1200 + (4+ 8,22
0.11 + 41)² sup || [V 20),
(E.33)
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 the strongly-convas case (Lemma [.I] the sense that the right hand side of E.3 no longer | f|| so that Lemma 2.1 is not directly applicable. Instead, in the following. we will provide a modification of Lemma 2. which will be used in proving Theorem 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 p E (0, and C E (0, oo) such that
(P Pm (0. CCp" (6e. 67)
sup 0
(E.34)
for any E N. Further assume that there exist some ó E (0, 1) and L E (0, oo) and measurable Lyapunov function V Rd + [1, oc) of P such that for any 0 E Rdd
(()0)) +
(E.35)
Then, we have
(w(p Pn) 5C (ddd((po Po) + (1 -p") 7K
(E.36)
1
where
dy (8oP,
::= max { v(0)dpo
L
7:= sup RER
((0)
1-
(E.37)
25

---

Proof. The proof is based on the modification of the proof of Lemma 2.1 (Theorem 3.1 in (RS88).
By induction we have
Pr. Pn -p0 po) p P (P-P) pm-tn n M.N.
(E.38)
We have
dø (p.P.PiP) & { de (osP, dpi(0) 5 //c ((0)
(E.39)
Moreover, for any i 0, 1,1,,,....
J ((0)dp:10) = J/e pd (O)(0)) S 8o L (1--0') 1 max po((V)
L
(B.40)
- 8
so that we obtain (p. P. I yK. Therefore, we have ppm S Cp" (D.P.r.P) S Cpp" cpm
Wi Ppm M.
(E.41)
By the triangle inequality of the Wasserstein distance, we have
n-1
Wi(pn, Pa) - Wi (" -00
+
Σ
Mi (p:Ppa-i-1 -p:P
==
n=1
Cpp" (Po. Po) + E
p
pa-i-lyr.
(E.42)
Y=U
This completes the proof.
Next, let us provide a technical lemma that upper bounds the norm of 8, and ô., which are the minimizers of F(0, X,,) == à E Vf(0, x;) and F(0, X,) == : Σm Vf(0, 2,) respectively.
Lemma E.6. Under Assumption 3.3 we have
suprex [VS(0,)))
supzex Vf(0,x) + 4mk
|0+ v
(E.43)
2m
|0+]] supxex V5(0,x))
suP
VS x) 2 4mK
(E.44)
2m
Proof. By Assumption 3.3
we have
- 0.)
1
(f(00,;;)..)
n
V)(00xr VI(0,, -0,)
1=1
ml 0.|00.1 1| K,
(B.45)
which implies that which yields that
m|
ml00.12 K < sup
[0 x| 10011
(E.46)
00. suPxex [V](0x))]
sup [S(0 + 4mK
(E.47)
2m
Similarly, one show that
16.1 supxex [V](0,,))] + v sup [] 4mK
(E.48)
2m
This completes the proof.
26

---

E.1
Proof of Theorem 3.3
Before going the proof, let us restate the full version of Theorem 3.3 that we provide below.
Theorem E.1 (Complete Theorem 3.3). Let Bu 0g 8. Assume that Assumption 3.1 As-
sumption 3.3 and Assumption 3.4 hold. We also assume that ry < min { m² K3+61D=K3 and
SuPzEx x < D for some D < oc and suP [V, I E for some E < Co. For any
(0,1) 1). Define M 0 so that S/o, - P(O.. f1)d01 2 vij and any R. 2Kg, where Ko is
defined in (E.2 so that
p(0,01)
inf -0 p(0., 81)
Vj.
(E.49)
Let un and Un denote the distributions of 02 and 02 respectively. Then, have
W1 W1(us P2)
1- o7t
2/ w11 +1))11 -7)
2b
2/(4+ + 87))
max
2
E E + 4mk
1 1h 1+ + 16(1+220 221
2m
E
E + 4mk
+ 4n²
2E +KK
2m
+ V F2 + 4mK
max
1 202 +2
2m
2
11 K 2
567
2 K2 +
64n
D2 K3
E V + 4mk
+
2K
17
0
m
m
m
2m
111
m
(E.50)
where for any E (0 and E (1- 1m-mm + 2nk
.1) one can choose v
To
and
K
(1 ()))) V
2+ Rehoo and dy is the weighted total variation distance defined in Section B.1 R+R
Proof. By applying Lemma H.I Lemma E.3
Lemma
E.4 and Lemma E.5 which is modification
of Lemma 2. we obtain
Wi Pk)
1 1- 17
2/ (01 + 1)(1 17)
2b
u/4+ 87)2K))
max
n
+ 1
1 1 1220² (4+
-0.12+ +4/2
sup xEx
V =m)l²
max 1+ 1+-10-000 0.11, 11,
2 2
56
D2
64n
D K 0. +
2K
17 02
m
m
m
m
m
(E.51)
where for any To E (0, and to E (1 (1- mn
+
Kg 1) one can choose e n
no
and
7 =
Koo
(1- (m)) v
R+R and de is the weighted total variation distance defined in Section B.1 24Re
27

---

Finally, let us notice that ||0_ - 0.11 - 21 + 12 and ||0 -0--1,1 6.1 - 2| + 21 and for every x E X
(E.52)
By applying Lemma
E.6 we complete the proof.
E.2 Proof of Corollary 3.5
Proof. Under our assumptions, the noise Eu are i..d. Gaussian No0 Σ) so that B||k[]] = tr(2) a². Moreoever, we have 2 < La. Then p(0.., 01) the probability density function of
0.
Σ
V ff(0., + me1.
(E.53)
Therefore,
(1-7
Ji 0. 1.
Σ Σ-
p(0,,01)d0, =
x;) + S M
(E.54)
iE
Notice that for any 21
Vi((,, 2t))
S
17 sup
(E.55)
b iEA
DE
Thus, have
J p(0.. 02)08 = =1-
Σ EE
vf(0., x2) + +
M
2 1 - P
[17Ea
M 17 sup #EX
V f(0t,
M
1 P
|K&l||
sup xE
[f((
(E.56)
n
Since & ~ NN(0, Σ) and E < Ia. for any 7 S 2' we have
E -"ll)))
det (la 27
By Chebychev's inequality, letting = 11, for any M 2 nsuuuee | (f(,,,r)) we get
1 L1M
M
Vf 2))| 2
p(O., 01)00 2 det (1a - E)
exp
sup z
11.1-5M
17
Next, for any 0, 01 E Rdd such that |10 0-11 s M and ||0 0.112
< 2Kg (1 + + 1, we have
01) En, [pn, (0,01)
(F.57)
p(0., 01) En, [a, (0., 01)]
where En, denotes the expectation w.r.t. 12, and po, denotes the probability density function conditional on Ωr For any given 2, we can compute that
po, (0,0,)
Pn: (02, 01)
cxxp {
2n2 (01 - (1B)) (81- T (())) + 2n (0) - -(B8) m(O.)). (E.58)
where
(0)) 0 " Σ v1().)))
((0.) 0, - Σ (0(0,,))).
(E.59)
b SEA
ien
28

---

Therefore, for any 8, 01 R such that ||01 - 0xll M and ||0 0=|1 v 2kg (1 + E) 1, we have
Pn, (0,01)
1
[la(0) - m(0-)]l || || (110, - u()))l + 110 - m((W)1D}
2 exp
pn, (0., 01)
2xp2
2 exp
2712 llx(0) m(0-))l
(((0)) 2110 ((0w))))
(E.60)
can further compute that
|0 m(0,))l - 110, - 0.1l+ +
Σ
Vf(0.,, z;))
M + 7 sup [V)(
(E.61)
ics
TEX
and
11m(0)) - =(00)1 & 110 0.11
7 Σ [Vf(0,2;) - Vf((.,,;))
b iEs2
< (1 + Kin) (im)11--6 6,1l
(2Ko (1 + E) - 1/2
S (1 + Kin)
(E.62)
Hence, we have
(2Kg (1 + c) - 1) 1/2 1)1/2
Pn, (0,0g) 01)
(1+ Knn)
|
exp
Pn. (0,,0) 01)
2772
(2K5 (1 +) -1)
1/2
(1 (1+Kin) Kin)
+2
M
1l sup ZEx
[Vf x)]l
m
(E.63)
Since it holds for every Pi, have
p(0,01)
(1 Kin)
2Ko
(1 + -1) 1/2
2 exp
m
|
p(0.,01)
2172
2Ko (1 + c) +))
1/2
(1+ (1+ K2n)
+2
M + 1 sup / ( x)||
NEX
(E.64)
Hence, we conclude that
- _1/M M sup // [V [(0.,.,)) 2)
exp
n
max
min
M2nsupsex [0....))
Vdet (1a E)
(1+ Kin) ( (2K0 (1 + c) 1)
||>1|||
exp
2n²
2Kg
1/2
(1 (1+Kim) Kin)
(1 +c) -1)
+ M n sup [V, x)]|
m
xEX
(E.65)
This completes the proof.
29

---

F Proofs of Non-Convex Case without Additive Noise
E1 Proof of Theorem 4.1
Proof. Let us recall that 0n 0 = and for any le E N.
02 Ox-1
Σ
(F.1)
b
iELL
0 0
" Σ (.((0k-1.
(F.2)
iEn
Thus it follows that
0. - 02
Ox-1 - 0
Σ
Vf(0k-11 vs (0.-1.
- 7 E
(F.3)
b
iEll
where
V(((n=1, - Vs (0k-1,21) (x))
Ex =
(F.4)
This implies that
00x - 0u 0a1l²
xk 0x-1-
7
Σ
(v - (86-1122) ))
+ l1Eall
"2
=
b
02
iEs
+ 2 Oh-1 - 0z-1
7
([((0=),,3)) - Vʃ (00-,11))))) E (F..)
b
By Assumption 3.1 and Assumption 3.3 we have
Ok-1 Ōn
7
V 1(00-11,2)) - Vs (0(00-1.02)
b
$ (1 2)mm)
| - 0x-1
||2
+ +2nkK +
122 6²
(bKi | 8a-1 1)
2
$ (1 mmm)
8/-1|
+ Kn
(F.6)
provided that 17 S
m K
Since X,, and Xn differ by at most one element and supsex ||=||| < D for some D < oc, we have
J; .; for any i E Sr with probability n-b and vi * iss for exactly one E St with probability b and therefore
E E1E&|1
S b n
(K22D (2100=111 1ll + (21(+==1+11))+ +))
4D²
(SElI0A 1112 112 +2), 2).
(F.7)
n
and moreover,
E M 8k-1- 17 Σ (VF(0k- - (6 12; (v) 3) EEn
iE
b E [1 (111)) | - 0 K22D (2110-111 + 1)]
nt
2K2 (1 + Kin)E
[(106-111 + 110 -106- 11) (216x-1l1 + 1)
2K 2K2Dn Dn (1 + Kin) (1+
3
-
El 1= 112 + 7
1810x=111
(F.8)
where we used the inequality that
(a + b)(b + 1) = 26² 2ab + a
+
11
302
3
12
(F.9)
2
30

---

for any a. a, E R. Therefore, we have
4D2
E 0k -bll" S (1 mmm) 0 -0a11
+ +n2K
88 (8E1|0-111 + +2)
4K2Dy (1 + Kin)
3
+ 2
112
BBl0 10m0 112)
(F.10)
n
In Lemma E.3 we showed that under the assumption 17 < min {
m K7+61DeK3 j and
supxex | for some D < oc, we have that for every A: E N,
(0x) < (1 1pm (1 + 2mm - 1 k? 561/ D2 K3 + 64v 2e./6//2/''' D²K211011" 2nK,
(F.11)
where (00) = 1 110 0- where 0, is the minimizer of (() = 1 n Σ Vf(0, This
implies that
E V (k))
2pm - K2 (66))D2233 D" + 2
$ (1 >m
V (0o)
1 1 mm))
$ + 110 0.1
+ 2
71 KT-
56n D² +
64n D² 12
2K
(F.12)
m
m
m
m
so that
El|0, 1|
< 2E
0g 0.
11
217 K7
112m D² +
128m D2 D²K210,11
4K
2 0 0,
+ 4
+ 2/12
V
E 4mK
21 K?
112m D²K3
S 4|10113 + 4
+ 4
2m
2
E + C+VE" E2 + 4mk 2
128
E + V22 + 4mk
4K
D2 K2
+ 2
B,
m
2m
m
2m
(F.13)
where we applied Lemma E.6 Similarly, we can show that
I8 0x
B.
(F.14)
Since Pu 0u 6, it follows from F.10 F.13) and F.144) that
W (vk, kx)
< E f 02| 1|
TAD*K2
4K2D
(1+
3
7
2K
< (1 (1 -m))*)
(8B 2)
(1 Kin)
B
B
+
bnm
nm
2
2
m
(F.15)
{ m Kft m. Kf
provided that 17 < min
m 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 0g and 0g can be bounded in the following sense.
31

---

Lemma G.1. Let 8o = 80- = 8. Suppose Assumption 4.1 and Assumption 4.2
hold and S
Kf+20- Then, we have
110 - 0,11
81
E 110
- 0.11² -
+
D²K3
(G.1)
k:
kmp
0.
1
E 0 - 0.1% kmj
+
8m
D²
2 22-11 +5).
(G.2)
k
P
where 0, is the minimizer of 1. Ei (f( .;)) and 0, is the minimizer of 1 E vf(0,z;).
Proof.
First, we recall that
02=0 00-0011-
17 Σ Vs (0
(G.3)
b
iEfL
Therefore, we have
Ō = 0-
V ( (0._1,
Σ5
Vs (h- -
Σ vʃ (
+1
1=1
SEl2
(G.4)
Moreover, we have
Σ vf (0._) 10--1
1
Vf (0:-1:2;)
E
(G.5)
lEfl
This implies that
2
E
0.1
E.
0 0 + Σ
Vf (0
1=1
Σ vs (0k_1, ti:) /
?
12
+ 2
Σ
vf (Ba-1,
-
(G.6)
i=1
We compute that
0 - 0,
VS (0A (0A-1; =
+
i=1
0 1 0, + 7 Σ
v (0 1, 2..
;.:))=/ -v/ (0., ))
n
2
1)2
<
0 0.1"
2711
0 0 - -0,1"
Σ
vs (0n-1, -vs
vs (0., d
n
i=1
<
02-1 0.12
| 2711
0 00-1- - 0.1%
1 k?
0 0.1 ||
(G.7)
32

---

Moreover, we can compute that
1 Σ vf (0a-1, /
?
- -
E
(0--1
1=1
=ES2p
vf 0k ( 2, - vf (f(0A-1, 2i:)
221)
=
E
iE
j=1
n
E
K2 n
Σ
رلا - 83 112 1)
iE0x
< 4D K3B
(2110>-1|| + )
s <8D2 K3
4B || +1
s 8D8 x3
[10--111) + 5)
x K3
22
E
0_ 1 + 2P 110.11 +5).
(G.8)
Hence, by applying
G.7 and G.8 to G.6). we conclude that
E 0.
1a6-0,11
0.
- E
0
0.
2nyBB
0 - 0.1" + K2RKE 0 10.-1-8.1 =6.l"
2 E [6--1 - 0.1" + 2 110 +
817 D²K3
$ E 8/- 0,1
- n
0 8.11" + 8772 D2 (20+1 11 + +5) (G.9)
provided that 17 S
K7 20 pot This implies that
E 0 0. 0.1ps
E 00-1
0. 1|
E
0_
2
81 D² 20tt110_1=+5)) (G.10)
1//4
and hence
100 0-
2 E
0x 0,
2
E
1
0.1"²
+
81
12a
(2 120 +5)
kapp
0.1 1|
0
+
811
D² (20+2 +5)
(G.11)
kipps
1
Similarly, we can show that
k
110 - 0.1
D² 8n
E
111
+
12p+1118, +
(G.12)
A:
kmp
1=1
This completes the proof.
Next, let provide technical lemma that upper bounds the norm of 0, and d., which the minimizers of F(0, X.) := 1 Σ a0)) and F( X.) Σ (5(0, ..) respectively.
Lemma G.2. Under Assumption
1 we have
||o.1| <
sup
V (f(0,))
μ ZE
and
1
||0,1| <
-
sup
\/
33

---

Proof. Under Assumption 4.1
we have
n
VF(0, VF(0,, 0.)
2f(0,
=
n1
-V(0, 00. 2 pl10.11P
11
(G.13)
where p E (1, which implies that
u||0.11P < sup =E
[V](0, ||0.1|,
(G.14)
which yields that Similarly, one can show that
||.|| 5 1F- ZEi
sup x
|10.1| S - sup [f(0,,)) pF ZEt
This completes the proof.
Now, we are able to state the main result for the Wasserstein algorithmic stability.
Theorem G.1. Let On 0g 0. Suppose Assumption 4.1 and Assumption 4.2 hold and s
T2 and supzex |>|| S D for some D) < oo and supzex |I (0, x) for some E < cc. Let un and Dy denote the distributions of 02 and O2 respectively. Then, we have
( 1-11
k
KD2K3n
81 D2K3 22 (2/ +5) 20 F 2p+2
bry
kmj
1
4D2 122 2F+2 bn4
(2p+2 (E/P) + +10)
4DK₂
811 D²K? (2/1((//)=t +5) 202 + pFF
(1
K1n) 33 2--1
mp
knp
4DK₂
20 + 2(E/n))=
81 D² (B))
(1+ Kin) 7
+
nu
knp
4DK₂
(1 + K1n) (10 (00 + +5).
(G.15)
ny
Proof. Let recall that Og 0, and for any N,
0x P--1
1)
(f(0)-1,2))
V
(G.16)
b
iEfl
0 = 0-- - 17
Σ
VS (0)
(G.17)
b
tEfu
Thus it follows that
Oy 0- 0x-1 0k-1
17 Σ (Vf(0- - Vs (0n-1, 7 En.
(G.18)
Veft
34

---

where
== = Σ
V/
(h (0 -v/ (0)-1,
v.
(G.19)
SEfs
This implies that
0x 0x| 11
8x-1 0k-1
Σ
Vf(0k-1, - -V (8..1.222) ))
n/²
||
=
b
6²
En
+
2
0z-1 On
7 Σ
(f((0=1)) - vf -V (0a-11
b
se
(G.20)
By Assumption 4.2 and Assumption 4.1 we have
8g 1 a Σ
(f(( - Vf (Pa-1.2:)
SE
00-1 0z ||
|
20p
0 x--1 0x-1 1|
1
bKi |0x-1 0n-1
11²)
(8) 2
122
|
0% ||²
27p4
6 - 0
p
+ k
k
0 - 0x 01-111"
$ 0x-1 8k
|
77/4
8x-1
p
(G.21)
provided that 1) S 127
Since X, and X, differ by at most one element and sup ||>|||| S D for some D < oc, we have I; = è for any i E 24 with probability n-b and x4 * i; for exactly one E 24 with probability and therefore
E 1| = EE (
(K-2D(210 + =1)
4D Kōb
(SE 1100-1126001 + +2)
n
4D2 K2b
SE [10 + 10)
12.
4D2 K3b
(2 10 - 1111"" + 2 11, + 10).
(G.22)
n
and moreover,
E
0x-1-
Σ
(Vf(0--1,I)) - Vf (0n-1, E
b
17 b
[110--1 0m-1
+ +Kin - K22D (2110--110-1 + 1)
-1
E
n
2D 22 E [(11 Kin)
|00-1- - 0 ||| K17)) @210m_11 +
n
2D
(1 + KB (10x-11l + 100--11 + 1 ) (210 ا +
(G.23)
n
Notice that for any x, 2 0 and p (1, 2). we have zy"-1 < r + v" y S 5d + 1. xx 1 and y < yd 11, which implies that
(x+y+1) 1) + 1) 2xy" 1 22 + 2y
+ 15 3z + 77 +5. (G.24)
35

---

Therefore, by applying
G.2 to G.23), we have
OK-1- - 0k-1 7 Σ
(.((k=-1,2)) - -VI (0n-1.2)) 7
iEfl
2DK21 (1 +1+ Kin)
3E 3B 1 11 + 7E| 0x 111 +
2DK21 (1 + Kin)
3 3..2 20 B1 0-1 - 0.11 +3 13-22-1110.11
+ 7 2. 0.11 + 2100 20 +5).
(G.25)
Hence, by applying (..21)
G.22
G.25
into G.20 we conclude that
E
0k
0x|
- E 0v-- - 02-1/
1| 2
nNB 0x- - 0m_1|"
4D² 1D?
(20+21)
0 - 0.1" + 20,t2110111"" + 10)
+
bn
4DK27
(1 + K17)
(3. 3..2 2P 2P-1E116k-1 - -011 + +3-20-1110.
n
+ +7.2p-1B11
6.1 +7 --7-20-1110- 20-- +5) 5) (G.26)
provided that 17 K7-
This, together with 0g - = 0, implies that
0,-1| ||"
E
0:-1
h:
41 K2
2D+2
00-1 - 0. 1|
+ 2 10.1" 10
bmp
4DK₂
(1 + Kin)
3. 3-20-1
20-1 1--1
B|10;-1 - 0,11 +3 - 100
K
1//4
7. 2-
1
Ell0; 0.10 +7 07. 2n 27- + 6). (G.27)
k
In Lemma G.1
we showed that under the assumption 1) S x7-20140233 we have
110 - -0.1
8n K2 (2p+1 5),
E 1100-11 - 0.11² <
+
(G.28)
kny
i=1
0.
E 0 118-1-88 0.11"= 5 kmp
+
811
D² K2 2 +5)
(G.29)
k
36

---

Hence, by plugging G.28 and G.29 into G.27 we conclude that
0--1| p
E
0
3 1=1
0 0. 1
4D2
81
12tt110.11 +6)
20+2
D2 K2
20 10
bmp
kmj
1
4DK2 np
[10-6,,) 0.11
8n
(1 + Kin) 3-
D2 D2 110.11 +5)
kap
p
(110-0.1 0_ 0.11"
4DK2
(1 Kin) -7-2--1
+
8n
10=+110,110 + +6) D2 D²K3
np
4DK2
+ 1/
(1 + K(n) (3 110.11 + 7- 2p $100 +
(G.30)
Moreover, 110-0.11= < 2 110 0.11 2| 2110. and by applying Lemma G.2
we obtain
00-1 - 1
4D2 K3n
2202 (FFF/))
8n D² (20 (2+t((//))== +5)
22+2
+
bm
knp
+ (E/p) -=م + 10
2DK2 np
202 + (2(E/M) 221
81 D²K' (20+ (E/N)
(1+ + Kin) 2-2
+
kija
1
4DK2
(1 + K1n) 7. 7.2P-1
202 + 2(E/u))
8q
D2 K2
(F/) + +6) 6))
np
knps
4DK₂
(1 + Kin) (10-22 (E/ +5).
(G.31)
11/1
Finally, by the definition of p- Wasserstein distance, we have
1 k
k
S
E 11000- - Σ
(G.32)
h 1=1
This completes the proof.
G.1 Proof of Theorem 4.2
Proof. First, let us show that the sequences 6g and 0x are ergodic.
First, let us show that the limit 6xo, if exists, is unique. Consider two sequences o( and e(2) starting at o²") and o(2) respectively with two limits
0 and 0 For any k E N,
0 (1)
(1) R-
(G.33)
b
NEA
0 (2)
!
Σ
Vf
(0(0)
(G.34)
b
IEA
37

---

By Assumption 4.2 and Assumption 4.1 we have
|0 0% 1 - (
Vf
( k-
00-1: 2i)
Vf (8 (0
6
2
|| 20p
1 0/ || (2) k-
+ q
bKi |0 - 6 1²)
6
2 20p
(2) 0
|
+ 1イ kn
(0)))
0 ||
0
m
K-1
04
(G.35)
provided that 17 S
H. R7
This implies that
|0
0//1
2 - k-1
((1)
12
100000 0 -
17/4
(G.36)
k-1
provided that
" K?
Thus, [0(2)
((1)
k
2
k-
000 1| for any to E N. Suppose
2
2 for every = 1,2,.... ki then we have
(((2) - 1
|080) 08¹
knp,
(G.37)
such that
(663) 8 (2)
10 1000 000
2
5 1,
for any 2 ko =
(G.38)
Since E (1,2),
1000 1/ S (1 - n/) 10 - 2
(G.39)
for any 1: > kn + 1, which implies that 0(2) ofth + a as k
that 0 = 0%
Next, let show that for any sequence 0g, it converges to limit. follows from GG39 that
< (1 (1- n/))= Akl | ((D) 1|
(G.40)
for any Ko + 1. Let " be a fixed initial value in Rd. and let o 00D) which random yet takes only finitely many values given o( so that ko is bounded. Therefore, it follows from (G.40) that
W8 M%) S (1 =
(1- n/4 n/o [ho]
|0/20| -
0(1) ko 17]
(G.41)
where vs denotes the distribution of 0g, which implies that
W2() < <c.
(G.42)
Thus, (Vk) is a Cauchy sequence in P2 (Rd) equipped with metric W2 and hence there exists some Vo such that (2 ) 0 k co.
Hence, showed that the sequence 0g is ergodic. Similarly, we can show that the sequence 0g is ergodic.
Finally, by ergodic theorem and Fatou's 's lemma, we have
k E 0, 0.-1|". (G.43)
E 00 0:0 |2-
lim k-000
0,-1
0
||*
lim sup 10
1=1
i=1
38

---

We can then apply (G.3l) from the proof of Theorem G. to obtain:
C2 C3
Wp Poo) -
(G.44)
bnp n
where
C2 = 4D2K3n
88
D K3 2 (E/1)=) +5)) - 20t2 ( + 10). (G.45)
22+2
C3 = 32D 3
(1 + K(n) 10 20 (2pt)(E/ +5)
122
4DK₂
+ Kin)
(10 (10.2" 20-1
(F/ 1011((1//))))) - +5)
(G.46)
This completes the proof.
39


Reviews:
[{"review_id": "8uIvu3EI1L", "reviewer": "Reviewer_tckZ", "summary": "The paper derives Wasserstein stability bounds for a variety of cases for a \"surrogate loss\" under convex/non-convex settings.", "questions": "I 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.", "limitations": "Theory work.", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The problem is interesting. The bounds in the non-convex case can be useful. The paper is well-written.", "weaknesses": "I 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time invested in our paper. \n\nAs far as we can understand, the reviewer has two main concerns: 1) novelty of the strongly convex part and 2) the use of the surrogate loss function. Below, we clarify both of these points and we hope that the reviewer could reconsider their score based on our explanations.\n\n* Novelty of the strongly convex part:\n\nAs we mentioned in the paper (and also as the reviewer acknowledged as well), the main novelty of the results in the strongly convex part is **not** the bound itself we derived, but **how** we obtained it. Here, our novelty is the introduction of the proof technique, which we believe is novel and very flexible: can be applied strongly convex / non-convex (dissipative) losses, and it can be rather easily extended to other optimization algorithms like SGD-momentum – will hopefully form a fertile ground for further articles. We hope that the reviewer could reposition our contributions from this perspective. \n\n* The use of the surrogate losses\n\nWe agree that the requirement of surrogate losses is a drawback of our framework. However, based on the reviewer’s concerns, we suspect that the use of surrogate losses might seem more daunting than it actually is. We will now clarify why we need surrogate losses, and we will provide an example, which we believe is a very natural case and has been used in prior studies. \n\n*The need for surrogate losses:* As our bounds are based on 1-Wasserstein distance, we need the surrogate loss $\\ell$ to be a Lipschitz continuous function. On the other hand, for the original loss $f$ we need some sort of convexity (e.g., strongly convex, convex, or dissipative) and we need the gradient of $f$ to be Lipschitz continuous. Unfortunately, under these assumptions, we cannot further impose $f$ **itself** to be Lipschitz because there is no function that satisfies these assumptions. Hence, the requirement for the surrogate losses. \n\n*Example for surrogate losses:* \n\nExample 1: We can choose the surrogate loss as the *truncated loss*, such that:\n\n$$\\ell(\\theta,x) = \\min(f(\\theta,x) , C)  $$\n\nwhere $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 (only) conceptually linked to adding a projection step to the optimizer. \n\nExample 2: Another natural setup for our framework is the $l_2$ regularized Lipschitz loss that was also used in [FR2021] (see Section 3 in their paper). As opposed to the previous case, for the sake of this example let us consider $\\ell$ as the true loss and $f$ as the surrogate loss. \n\nThen, we can choose the pair $f$ and $\\ell$ as follows:\n\n$$ f(\\theta,x) = \\ell(\\theta,x) + \\lambda ||\\theta||_2^2 $$ \n\nwhere $\\lambda >0$. Intuitively, this setting means that, we have a true loss $\\ell$ which can be Lipschitz, but in the optimization framework we consider a regularized version of the loss – which has also been considered in various studies, the closest to us being [FR2021].\n\nWe sincerely appreciate that the reviewer has gone over the literature for surrogate losses. We agree this part is a little bit overlooked by assuming the reader might be familiar with the concepts. In the next version, we will provide further explanation and examples about surrogate losses, as we noted above. We will also rephrase the part where we introduced the concept of algorithmic stability, as the reviewer suggested. \n\nIf there are further questions, we remain at your disposal.\n\n–\n\nReferences:\n\n[FR2021] Tyler Farghly and Patrick Rebeschini. Time-independent generalization bounds for 364 SGLD in non-convex settings. In Advances in Neural Information Processing Systems, 365 volume 34, pages 19836–19846, 2021."}, {"review_id": "U88Z35ddt4", "reviewer": "Reviewer_yPZ3", "summary": "This 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.", "questions": "Lemma 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)", "limitations": "I do not see negative societal impact.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The 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.", "weaknesses": "The 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their careful reading of our paper. Below are responses to your questions:\n\n**Weaknesses:** \n\n* We will add an intuitive explanation behind formula (3.10) and its implications for SGD. The parameter $\\hat{\\eta}$ appears in the definition of $\\bar{\\eta}$ that 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 $\\bar{\\eta}$, the smaller the 1-Wasserstein bound. By the definition of $\\bar{\\eta}$, the larger $\\hat{\\eta}$, the smaller the 1-Wasserstein bound. As a result, we would like to choose $\\hat{\\eta}$ to be as large as possible, and the equation (3.10) provides an explicit value that $\\hat{\\eta}$ can take (which is already as large as possible from our theory). We will add more discussions and intuitive explanations in the revised version of the paper.\n\n* Our bound in Theorem 3.2 for the strongly convex case $O(\\frac{1}{n\\mu^5})$ is indeed worse than the $O(\\frac{1}{n\\mu})$ bound given in [HRS16] in terms of its dependency to the strong convexity constant $\\mu$. That being said, as we mentioned in the paper, the main novelty of the results for the strongly convex case is *not* the bound itself we derived, but *how* we obtained it. Here, our novelty is the introduction of the proof technique, which we believe is very novel and flexible: can be applied strongly convex / non-convex (dissipative) losses, and it can be rather easily extended to other optimization algorithms like SGD-momentum – will hopefully form a fertile ground for further articles.\n\n* The bound in Theorem 4.1 is indeed not tight due to the persistent $K/m$ term that does not vanish to zero as the number of samples $n$ increases. However, the bound is **not vacuous** unless $n$ is very large: depending on the value of $K$ (the measure of nonconvexity), the bound can be still informative for a nonasymptotic $n$. On the other hand, the existing algorithmic stability bounds (cited in the paper) for similar non-convex problems typically increase with the number of iterations and become infinite and actually vacuous for any range of $n$ as the number of iterations increases. Compared to those bounds, our bound does not increase with the number of iterations, which we believe is a significant improvement.\n \n\n**Questions:** \n\n* Often, the Lyapunov function is required to be bounded from below by a positive constant, and some mathematics literature for simplicity requires $V\\geq 1$. In order to create a positive function that is greater than $1$, the simplest function one can come up with is a quadratic plus $1$. The quadratic choice for the Lyapunov function is quite standard for SGD and related algorithms and is usually the first one to try. On the other hand, in order to use the strong-convexity and dissipativity conditions, it is easier to work with $|\\theta-\\theta_{\\ast}|^{2}$ instead of $|\\theta|^{2}$. The basic principle of building a Lyapunov function is to show a drift condition can be satisfied so that the expected value of the Lyapunov function can be uniformly bounded in time so that it will not grow as time increases. That being said, there is no principled guideline to design a Lyapunov function for an arbitrary optimization algorithm and objective function; however, for popular algorithms and problems classes several Lyapunov functions have already been developed as we mentioned in Line 100. Note that in optimization, to show convergence, we usually need the Lyapunov function to go to zero as iteration number increases; here because we do not need to show convergence to a point, it suffices that Lyapunov function over the iterations stays as a constant. \n\n* The paper considers noisy SGD for the nonconvex case. The noise is useful in this case because in order to apply Wasserstein perturbation theory [RS18]: we need contraction in Wasserstein distance (geometric ergodicity). Informally, in the nonconvex case, the Markov chain associated with SGD might not be geometrically ergodic if the only source of randomness comes from minibatching. Adding noise in this case enables the Markov chain to access all the regions of the state-space with non-zero probability and makes the chain geometrically ergodic. \n\nFormally, in the strongly-convex case, one can easily use the synchronous coupling method to obtain Wasserstein contraction such that no additional noise is needed. However, in the non-convex case, under a more general dissipativity condition, Wasserstein contraction can not be obtained by using the synchronous coupling method. Instead, in order to obtain Wasserstein contraction (geometric ergodicity), we apply the theory developed in [HM11] which says that as long as the Markov chain satisfies a drift condition (Assumption A.1) that relies on the construction of an appropriate Lyapunov function and a minorization condition (Assumption A.2), the Markov chain would converge in some weighted total variation distance that can imply convergence in 1-Wasserstein. The dissipativity condition helps us to show that we can indeed construct a Lyapunov function that satisfies Assumption A.1. However, a minorization condition (Assumption A.2.) requires a certain mixing property that requires some continuity-type assumptions on the noise structure, and this is where we need the additional noise (since the noise from the mini-batch has discrete nature and is not sufficient to obtain mixing property, i.e the minorization property.) \n\n\n** Minor comments **\n\nWe will fix these typographical errors, and revise along your suggestions. Thanks for the feedback."}, {"review_id": "RrfmJ8wKq3", "reviewer": "Reviewer_RR88", "summary": "In 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\n", "questions": "- 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", "limitations": "Not 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.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- 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.", "weaknesses": "- 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their careful reading of our paper. Below are responses to your questions:\n\n**Weaknesses**:\n\n* As the referee pointed out, our analysis applies to non-convex losses that satisfy some conditions (a dissipativity condition and a pseudo-Lipschitz gradient growth condition). As suggested, we will clarify this in our introduction further, emphasizing the assumptions we make. \n\n\n\n* Our approach mainly relies on the Wasserstein perturbation method developed in [RS18], which requires that the Markov chain is time-homogeneous. When the step-size is decreasing, the associated Markov chain is no longer time-homogeneous and thus [RS18] is no longer applicable. Therefore, the regime with decreasing step-size is out of the scope of the current paper. However, we do believe that it is possible to extend the analysis of [RS18] to deal with time-inhomogeneous Markov chains. But that will require us to write a paper that first extends the analysis of [RS18] and then apply it to the context of decreasing step-size. That will be left as a future research direction which will be interesting to pursue.\n\n* In [RRT17], the authors provided analysis for the optimization error depending on the additive noise variance, later on these results were extended to different settings. We can directly combine our results with their results to obtain an optimization error + generalization error bound. We will mention this in the next version of the paper. \n\n* In line 291, what we meant to say was that for our proof technique to work well, adding noise was beneficial. We absolutely agree with the referee that this does not mean that additive noise is a must for an algorithm to perform/generalize better.\n\n* Thanks for the suggestions on the notations. The notation $\\mathcal{A}$ does not mean that it is deterministic, even though it is a function of the dataset, which is deterministic. The extra randomness comes from the randomness in the SGD, not the dataset. We will make this more transparent in the revised version of the paper.\nIndeed, the notation $\\nu_{X}$, $\\nu_{\\hat{X}}$ might be better than the notation $\\nu$ and $\\hat{\\nu}$ that we currently use. However, since the upper bound we eventually obtain does not rely on the particular dependence on $X$ and $\\hat{X}$, but only on the fact that they differ by only one element and data points are bounded with radius $D$, so that there are factors $1/n$ and $D$ that appear in the bound, independent of $X$ and $\\hat{X}$, our notations as they currently are will not cause a misunderstanding. But we do agree that more formally, it is better to use the notation $\\nu_{X}$, $\\nu_{\\hat{X}}$. \n\nWe will fix the issues mentioned in the minor comments, indeed these are typographical errors.\n\n\n**Questions:**\n\n* Thanks for this excellent observation. Indeed, the only place the neighborhood assumption between datasets is used is in the expression of $\\gamma$ in equation (2.9). Lemma 2.1 ([RS18], Theorem 3.1) relies on three conditions, the Wasserstein construction in (2.7), the drift condition for the Lyapunov function in (2.8) and finally the estimate on $\\gamma$ in equation (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. As a result, the only place the neighborhood assumption between datasets is used is in the expression of $\\gamma$ in equation (2.9). We will add some discussions to make this little intuition more explicit in the revised version of the paper.\n\n* Thanks for the comment. Indeed, the probability you mentioned is proportional to the batch-size $b$. However, do not forget that there is the term $\\eta/b$ in (B.9) and (B.10) in the proof of Lemma 3.3., and here $b$ comes from the definition of the SGD, where the factor $1/b$ helps the stochastic gradient become an unbiased estimator of the full gradient. Combining these two effects, you will obtain a bound on $\\gamma$ which is independent of the batch-size $b$.\n\n* Indeed, it is very natural to ask if one can let $\\sigma^{2}=0$ so that one can remove the additional noise added to the SGD in the case of non-convex loss functions. Unfortunately, the answer is no. Note that for the sake of notational convenience, we let $\\sigma^2=\\mathbb{E}\\Vert\\xi_1\\Vert^{2}$ because the term $\\mathbb{E}\\Vert\\xi_1\\Vert^{2}$ appears many times in the bounds that we derived and it is convenient to use the notation $\\sigma^{2}$. On the other hand, the probability density function $p(x)$ of $\\xi_{1}$ implicitly depends on $\\sigma^{2}$. If you let $\\sigma^{2}\\rightarrow 0$, $p(x)$ will converge to the Dirac delta distribution concentrated at $x=0$. When $\\sigma^{2}\\rightarrow 0$, there will be no mixing since $p(x)$ will converge to the Dirac delta distribution concentrated at $x=0$ such that in equation (3.8), $\\hat{\\eta}$ will have to go to zero. If $\\hat{\\eta}$ goes to zero, by the definition of $\\bar{\\eta}$, we will have $\\bar{\\eta}$ goes to $1$, which will make the upper bound in equation (3.9) explode. Therefore, to conclude, we cannot let $\\sigma^{2}=0$ because otherwise, the upper bound in equation (3.9) will diverge and become a trivial upper bound."}, {"review_id": "CSVt0fCPrg", "reviewer": "Reviewer_Wt2h", "summary": "This 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", "questions": "My 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\n", "limitations": "Yes, 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.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The 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.", "weaknesses": "* 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\n\n\n\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your time invested in our paper. Below are our responses to your questions:\n\n**Weaknesses**:\n\n1. Weaker result with respect to other parameters: We get an optimal rate with respect to the number of samples. But dependence on other parameters might not be optimal  in some results, for instance dependency to the strong convexity constant. However, our result is more general with respect to existing SGD stability results. Our bound on expected generalization error applies to an entire class of Lipschitz loss functions, not limited to a single one. However, the existing stability results for SGD that we are aware of hold typically for a single loss function at a time and do not apply to a class of functions. Also, as the referee mentioned, the value of our analysis technique is that it can be employed in different settings.\n\n2. Assumption 3.1 : Assumption 3.1 is more general than the smoothness (Lipschitzness of the gradients of $f$) assumption. In fact, for $K_2 =0$, the assumption becomes a smoothness assumption on $f$ in parameter $\\theta$. Substituting $K_2=0$ in our results provides us the bound for smooth $f$. For example, regularized logistic regression with bounded data will fall into the $K_2 = 0$ case. We will add a remark about this during the revision.  We conclude that, allowing $K_2 >0$ in Assumption 3.1. allows us to cover a larger class. Therefore, we interpret Assumption 3.1 as a more general version of smoothness assumption. As a concrete example, consider the loss $f(x,\\theta) = | \\theta^T x - b|^p$ with $p\\in (1,2)$ which arises in robust regression. The gradient of $f$ with respect to $\\theta$ is Holder with Holder constant $p-1 \\in (0,1)$, and is not Lipschitz but it will be pseudo-Lipschitz. \n\nWe agree with the referee that Lipschitz constants in neural networks can be very large (assuming the norms of the weights are bounded), in this case our constants $K_1$ and $K_2$ in Assumption 3.1 will also be unfortunately large. However, if weight regularization is used; i.e. if a quadratic penalty term is added to the training loss then by a Lagrangian reformulation, this will be equivalent to training with norm constraints on the weights, and this will allow us to control the Lipschitz constants.\n\n3. Assumption 3.3: The class of dissipative functions we consider 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 (arXiv:2205.14818), they arise in non-convex formulations of classification problems; for instance in logistic regression with a sigmoid/non-convex link function. They can also arise in robust regression problems, see e.g. [X Gao, M Gürbüzbalaban, L Zhu - Operations Research, 2022]. Also, any function $f$ that is strongly convex outside of a ball of radius $R$ for some $R >0 $ 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 arxiv:2007.11612v3). Dissipative functions also arise frequently in the sampling and Bayesian learning literature, for instance in the study stochastic gradient Langevin dynamics (SGLD), this assumption is used to show that a  stationary distribution exists (see e.g. arxiv:1702.03849)  \n\n4. Theorem 3.3 : Thanks for the comment. In the revised version, we will add further explanations to discuss these parameters and how they affect the main results of the paper. For example, the parameter $\\bar{\\eta}$ 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 $\\bar{\\eta}$, the smaller the 1-Wasserstein bound. By the definition of $\\bar{\\eta}$, the larger $\\hat{\\eta}$, the smaller the 1-Wasserstein bound. As a result, we would like to choose $\\hat{\\eta}$ to be as large as possible, and the equation (3.10) provides an explicit value that $\\hat{\\eta}$ can take (which is already as large as possible from our theory). The parameter $M$ can be interpreted in a similar way. We will add more discussions and intuitive explanations in the revised version of the paper.\n\n5. Theorem 4.1 :   If the loss is close to being strongly convex, then the bias term $K/m$ will be small. For example, consider a loss which is $\\mu$-strongly convex outside a ball of radius $R$, and inside this ball of radius $R$, it can be non-convex with a bounded Hessian. As the size of the ball radius $R$ gets smaller, we have $K\\to 0$ and $\\mu\\to m$ so that $K/m$ will go to zero as $R$ to zero. A simple toy example to visualize this would be the double-well example (see Example 1 in [X Gao, M Gürbüzbalaban, L Zhu - Operations Research, 2022]). That being said, we admit that such losses are not common in the statistical learning literature and by simply adding Gaussian noise to the iterations, we can analyze more general non-convex losses as we detail in Section 3.3.\n\n6. Assumptions 4.1, 4.2 : Even though Assumptions 4.1 and 4.2 have been used in prior theoretical studies as we cited in the paper, after your question, we now realized that unfortunately they did not provide any concrete examples that satisfy these assumptions. Nevertheless, as an example, we suspect that the loss f(\\theta, x) = |a^T \\theta - b|^p satisfies the assumption where $p \\in (1,2)$ and $x=(a,b)$ is the input-output pair. We could not formally prove this in the limited rebuttal period; however, we will try our best to come up with concrete examples for these assumptions. We agree with you that such examples would play an important role in terms of appreciation of the results, and we would like to thank you for bringing this to our attention."}, {"review_id": "sWAF5v8aBD", "reviewer": "Reviewer_9dAF", "summary": "This paper provides new stability bounds for SGD: they rely on perturbation theory for Markov Chains and on the ergodicity of SGD.", "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", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- Well written, well presented paper.\n- New connection between the theory for Markov Chain and generalization of learning algorithms.\n- First (to my knowledge) uniform in time stability bounds.\n- Extensions and limits are properly discussed.", "weaknesses": "N/A", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for a careful reading of our paper and finding that our paper is well-written with some new results on the connections between the theory for Markov Chain and generalization of learning algorithms.\n\nThanks for the excellent question. In our paper, most of the theory is derived using the tool of the Wasserstein perturbation method developed in [RS18], and using this particular tool, 1-Wasserstein algorithmic stability bound is obtained. In order to obtain the generalization error bound from the 1-Wasserstein algorithmic stability bound, one has to assume that the surrogate loss is Lipschitz [RRT+16]. However, some of the results in this paper do not rely on the Wasserstein perturbation method [RS18], such as Theorem 4.2. For these results, we are able to obtain algorithmic stability in $p$-Wasserstein for some $p>1$. In particular, with the 2-Wasserstein, we are indeed able to handle smooth loss functions if we also establish uniform L^2 bounds following the approach in [RRT17]; see also [FR21]. We will add these discussions in the revised version of the paper.\n"}]

Author rebuttal:

Meta-review: In this paper, the authors have introduced Wasserstein stability bounds for stochastic optimization algorithms, utilizing tools derived from the Wasserstein perturbation bound for Markov chains. The outcomes not only recover some existing uniform stability bounds for SGD in the literature but also extend to encompass other nonconvex scenarios under certain assumptions. Reviewers  acknowledge the intriguing link forged between learning theory and applied probability and the technical innovation. This approach can find potential applications in other optimization algorithms.

While concerns were initially raised regarding the discrepancy between the surrogate loss function $f$ and the original loss function $\ell$, extensive discussions led to a consensus that there exist avenues (e.g.by adding L2 regularization terms) to mitigate this discrepancy.

Taking into account the collective recommendations from the reviewers, I recommend accepting the paper. I would like to encourage the authors to thoroughly address the raised comments, particularly by incorporating discussions about the relationship between $f$ and $\ell$ and bolstering this with practical examples from machine learning.
Decision: Accept (poster)