SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lingjiong Zhu $^{1}$ , Mert Gurbuzbalaban $^{2,3}$ , Anant Raj $^{4,5}$ , Umut Şimşekli $^{5}$
1: 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 Urbana-Champaign
5: Inria Paris, CNRS, Ecole Normale Supérieure, PSL Research University
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.
With the development of modern machine learning applications, understanding the generalization properties of stochastic gradient descent (SGD) has become a major challenge in statistical learning theory. In this context, the main goal is to obtain computable upper-bounds on the population risk associated with the output of the SGD algorithm that is given as follows:  $F(\theta) \coloneqq \mathbb{E}_{x \sim \mathcal{D}}[f(\theta, x)]$ , where  $x \in \mathcal{X}$  denotes a random data point,  $\mathcal{D}$  is the (unknown) data distribution defined on the data space  $\mathcal{X}$ ,  $\theta$  denotes the parameter vector, and  $f: \mathbb{R}^d \times \mathcal{X} \to \mathbb{R}$  is an instantaneous loss function.
In a practical setting, directly minimizing  $F(\theta)$  is not typically possible as  $\mathcal{D}$  is unknown; yet one typically has access to a finite data set  $X_{n} = \{x_{1},\ldots ,x_{n}\} \in \mathcal{X}^{n}$ , where we assume each  $x_{i}$  is independent and identically distributed (i.i.d.) with the common distribution  $\mathcal{D}$ . Hence, given  $X_{n}$ , one can then attempt to minimize the empirical risk  $\hat{F} (\theta ,X_n)\coloneqq \frac{1}{n}\sum_{i = 1}^{n}f(\theta ,x_i)$  as a proxy for  $F(\theta)$ . In this setting, SGD has been one of the most popular optimization algorithms for minimizing
$\hat{F} (\theta)$  and is based on the following recursion:
$$
\theta_ {k} = \theta_ {k - 1} - \eta \tilde {\nabla} \hat {F} _ {k} (\theta_ {k - 1}, X _ {n}), \qquad \tilde {\nabla} \hat {F} _ {k} (\theta_ {k - 1}, X _ {n}) := \frac {1}{b} \sum_ {i \in \Omega_ {k}} \nabla f (\theta_ {k - 1}, x _ {i}), \qquad (1. 1)
$$
where  $\eta$  is the step-size,  $b$  is the batch-size,  $\Omega_{k}$  is the minibatch that is chosen randomly from the set  $\{1,2,\dots ,n\}$ , and its cardinality satisfies  $|\Omega_k| = b$ .
One fruitful approach for estimating the population risk attained by SGD, i.e.,  $F(\theta_k)$ , is based on the following simple decomposition:
$$
F \left(\theta_ {k}\right) \leq \hat {F} \left(\theta_ {k}\right) + \left| \hat {F} \left(\theta_ {k}\right) - F \left(\theta_ {k}\right) \right|, \tag {1.2}
$$
where the last term is called the generalization error. Once a computable upper-bound for the generalization error can be obtained, this decomposition directly leads to a computable upper bound for the population risk  $F(\theta_k)$ , since  $\hat{F}(\theta_k)$  can be computed thanks to the availability of  $X_n$ . Hence, the challenge here is reduced to derive upper-bounds on  $|\hat{F}(\theta_k) - F(\theta_k)|$ , typically referred to as generalization bounds.
Among many approaches for deriving generalization bounds, algorithmic stability [BE02] has been one of the most fruitful notions that have paved the way to numerous generalization bounds for stochastic optimization algorithms [HRS16, CJY18, MWZZ18, FV19, LY20,  $\mathsf{Z}\mathsf{B}\mathsf{B}^{+}22$ ]. In a nutshell, algorithmic stability measures how much the algorithm output differs if we replace one data point in  $X_{n}$  with a new sample. More precisely, in the context of SGD, given another data set  $\hat{X}_n = \{\hat{x}_1,\dots ,\hat{x}_n\} = \{x_1,\dots ,x_{i - 1},\hat{x}_i,x_{i + 1},\dots x_n\} \in \mathcal{X}^n$  that differs from  $X_{n}$  by at most one element, we (theoretically) consider running SGD on  $\hat{X}_n$ , i.e.,
$$
\hat {\theta} _ {k} = \hat {\theta} _ {k - 1} - \eta \tilde {\nabla} \hat {F} _ {k} \left(\hat {\theta} _ {k - 1}, \hat {X} _ {n}\right), \quad \tilde {\nabla} \hat {F} _ {k} \left(\hat {\theta} _ {k - 1}, \hat {X} _ {n}\right) := \frac {1}{b} \sum_ {i \in \Omega_ {k}} \nabla f \left(\hat {\theta} _ {k - 1}, \hat {x} _ {i}\right), \tag {1.3}
$$
and we are interested in the discrepancy between  $\theta_{k}$  and  $\hat{\theta}_k$  in some precise sense (to be formally defined in the next section). The wisdom of algorithmic stability indicates that a smaller discrepancy between  $\theta_{k}$  and  $\hat{\theta}_k$  implies a smaller generalization error.
The last decade has witnessed an increasing number of stability bounds for different algorithms applied on different classes of loss functions. In a pioneering study, [HRS16] proved a variety of stability bounds for SGD, for strongly convex, convex, and non-convex problems. Their analysis showed that, under strong convexity and bounded gradient assumptions, the generalization error of SGD with constant step-size is of order  $n^{-1}$ ; whereas for general convex and non-convex problems, their bounds diverged with the number of iterations (even with a projection step), unless a decreasing step-size is used. In subsequent studies [LY20, KWS23] extended the results of [HRS16], by either relaxing the assumptions or generalizing the setting to more general algorithms. However, their bounds still diverged for constant step-sizes, unless strong convexity is assumed. In a recent study, [BFGT20] proved stability lower-bounds for projected SGD when the loss is convex and non-smooth. Their results showed for general non-smooth loss functions we cannot expect to prove time-uniform (i.e., non-divergent with the number of iterations) stability bounds for SGD, even when a projection step is appended.
In a related line of research, several studies investigated the algorithmic stability of the stochastic gradient Langevin dynamics (SGLD) algorithm [WT11], which is essentially a 'noisy' version of SGD that uses the following recursion:  $\theta_{k} = \theta_{k - 1} - \eta \tilde{\nabla}\hat{F}_{k}(\theta_{k - 1},X_{n}) + \xi_{k}$ , where  $(\xi_k)_{k\geq 0}$  is a sequence of i.i.d. Gaussian vectors, independent of  $\theta_{k - 1}$  and  $\Omega_{k}$ . The authors of [RRT17, MWZZ18] proved stability bounds for SGLD for non-convex losses, which were then extended to more general (non-Gaussian) noise settings in [LLQ19]. While these bounds hinted at the benefits of additional noise in terms of stability, they still increased with the number of iterations, which limited the impact of their results. More recently, [FR21] proved the first time-uniform stability bounds for SGLD under non-convexity, indicating that, with the presence of additive Gaussian noise, better stability bounds can be achieved. Their time-uniform results were then extended to non-Gaussian, heavy-tailed perturbations in  $[\mathrm{RBG^{+}23},\mathrm{RZGS}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, it is not clear how the noisy perturbations affect algorithmic stability so that time-uniform bounds can be achieved, and more generally, it is not clear in which circumstances one might hope for time-uniform stability bounds.
In this study, we contribute to this line of research and prove novel time-uniform algorithmic stability bounds for SGD and its noisy versions. Our main contributions are as follows:
- We make a novel connection between learning theory and applied probability, and introduce a unified guideline for proving Wasserstein stability bounds for stochastic optimization algorithms with a constant step-size. Our approach is based on Markov chain perturbation theory [RS18], which offers a three-step proof technique for deriving stability bounds: (i) showing the optimizer is geometrically ergodic, (ii) obtaining a Lyapunov function for the optimizer and the loss, and (iii) bounding the discrepancy between the Markov transition kernels associated with the chains  $(\theta_{k})_{k\geq 0}$  and  $(\hat{\theta}_k)_{k\geq 0}$ . We illustrate this approach on SGD and show that time-uniform stability bounds can be obtained under a pseudo-Lipschitz-like condition for smooth strongly-convex losses (we recover similar results to the ones of [HRS16]) and a class of non-convex losses (that satisfy a dissipativity condition) when a noisy perturbation with finite variance (not necessarily Gaussian, hence more general than [FR21]) is introduced. Our results shed more light on the role of the additional noise in terms of obtaining time-uniform bounds: in the non-convex case the optimizer might not be geometrically ergodic unless additional noise is introduced, hence the bound cannot be obtained. Moreover, our approach is flexible and can be generalizable to other popular optimizers, as it mainly requires developing Lyapunov functions, which are often readily available in the literature [AFGO20, LRP16, FGO+22, GRZ22, LGY20, AFGO19].
- We then investigate the case where no additional noise is introduced to the SGD recursion and the geometric ergodicity condition does not hold. First, for non-convex losses, we prove a time-uniform stability bound, where the bound converges to a positive number (instead of zero) as  $n \to \infty$ , 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  $(\theta_k)_{k \geq 0}$ , which vanish as  $n$  increases. To the best of our knowledge, these results are novel, and indicate that the stability bounds do not need to increase with time even under non-convexity and without additional perturbations; yet, they might have a different nature depending on the problem class.
One limitation of our analysis is that it requires Lipschitz surrogate loss functions and does not directly handle the original loss function, due to the use of the Wasserstein distance  $\left[\mathrm{RRT}^{+}16\right]$ . Yet, surrogate losses have been readily utilized in the recent stability literature (e.g.,  $\left[\mathrm{RBG}^{+}23, \mathrm{RZGS}23\right]$ ) and we believe that our analysis might illuminate uncovered aspects of SGD even with this requirement. All the proofs are provided in the Appendix.
2 Technical Background
2.1 The Wasserstein distance and Wasserstein algorithmic stability
Wasserstein distance. For  $p \geq 1$ , the  $p$ -Wasserstein distance between two probability measures  $\mu$  and  $\nu$  on  $\mathbb{R}^d$  is defined as [Vil09]:
$$
\mathcal {W} _ {p} (\mu , \nu) = \left\{\inf  \mathbb {E} \| X - Y \| ^ {p} \right\} ^ {1 / p}, \tag {2.1}
$$
where the infimum is taken over all couplings of  $X \sim \mu$  and  $Y \sim \nu$ . In particular, the dual representation for the 1-Wasserstein distance is given as [Vil09]:
$$
\mathcal {W} _ {1} (\mu , \nu) = \sup  _ {h \in \operatorname {L i p} (1)} \left| \int_ {\mathbb {R} ^ {d}} h (x) \mu (d x) - \int_ {\mathbb {R} ^ {d}} h (x) \nu (d x) \right|, \tag {2.2}
$$
where  $\operatorname {Lip}(1)$  consists of the functions  $h:\mathbb{R}^d\to \mathbb{R}$  that are 1-Lipschitz.
Wasserstein algorithmic stability. Algorithmic stability is a crucial concept in learning theory that has led to numerous significant theoretical breakthroughs [BE02, HRS16]. To begin, we will present the definition of algorithmic stability as stated in [HRS16]:
Definition 2.1 ([HRS16], Definition 2.1). Let  $\mathcal{RV}(\mathbb{R}^d)$  denote the set of  $\mathbb{R}^d$ -valued random vectors. For a (surrogate) loss function  $\ell : \mathbb{R}^d \times \mathcal{X} \to \mathbb{R}$ , an algorithm  $\mathcal{A} : \bigcup_{n=1}^{\infty} \mathcal{X}^n \to \mathcal{RV}(\mathbb{R}^d)$  is  $\varepsilon$ -uniformly stable if
$$
\sup  _ {X \cong \hat {X}} \sup  _ {z \in \mathcal {X}} \mathbb {E} \left[ \ell (\mathcal {A} (X), z) - \ell (\mathcal {A} (\hat {X}), z) \right] \leq \varepsilon , \tag {2.3}
$$
where the first supremum is taken over data  $X, \hat{X} \in \mathcal{X}^n$  that differ by one element, denoted by  $X \cong \hat{X}$ .
In this context, we purposefully employ a distinct notation for the loss function  $\ell$  (in contrast to  $f$ ) since our theoretical framework necessitates measuring algorithmic stability through a surrogate loss function, which may differ from the original loss  $f$ . More precisely, our bounds will be based on the 1-Wasserstein distance, hence, we will need the surrogate loss  $\ell$  to be a Lipschitz continuous function, as we will detail in (2.2). On the other hand, for the original loss  $f$  we will need some form of convexity (e.g., strongly convex, convex, or dissipative) and we will need the gradient of  $f$  to be Lipschitz continuous, in order to derive Wasserstein bounds. Unfortunately, under these assumptions, we cannot further impose  $f$  itself to be Lipschitz, hence the need for surrogate losses. Nevertheless, the usage of surrogate losses is common in learning theory, see e.g., [FR21, RZGS23], and we present concrete practical examples in the Appendix.
Now, we present a result from [HRS16] that establishes a connection between algorithmic stability and the generalization performance of a randomized algorithm. Prior to presenting the result, we define the empirical and population risks with respect to the loss function  $\ell$  as follows:
$$
\hat {R} (\theta , X _ {n}) := \frac {1}{n} \sum_ {i = 1} ^ {n} \ell (\theta , x _ {i}), \quad R (\theta) := \mathbb {E} _ {x \sim \mathcal {D}} [ \ell (\theta , x) ].
$$
Theorem 2.1 ([HRS16], Theorem 2.2). Suppose that  $\mathcal{A}$  is an  $\varepsilon$ -uniformly stable algorithm, then the expected generalization error is bounded by
$$
\left| \mathbb {E} _ {\mathcal {A}, X _ {n}} \left[ \hat {R} \left(\mathcal {A} \left(X _ {n}\right), X _ {n}\right) - R \left(\mathcal {A} \left(X _ {n}\right)\right) \right] \right| \leq \varepsilon . \tag {2.4}
$$
For a randomized algorithm, if  $\nu$  and  $\hat{\nu}$  denotes the law of  $\mathcal{A}(X)$  and  $\mathcal{A}(\hat{X})$  then for a  $\mathcal{L}$ -Lipschitz surrogate loss function  $\ell$ , we have the following generalization error guarantee,
$$
\left| \mathbb {E} _ {\mathcal {A}, X _ {n}} \left[ \hat {R} (\mathcal {A} (X _ {n}), X _ {n}) - R (\mathcal {A} (X _ {n})) \right] \right| \leq \mathcal {L} \sup  _ {X \cong \hat {X}} \mathcal {W} _ {1} (\nu , \hat {\nu}). \tag {2.5}
$$
The above result can be directly obtained from the combination of the results given in (2.2), Definition 2.1, and Theorem 2.1 (see also  $\mathrm{[RRT^{+}16]}$ ).
2.2 Perturbation theory for Markov chains
Next, we recall the Wasserstein perturbation bound for Markov chains from [RS18]. Let  $(\theta_{n})_{n = 0}^{\infty}$  be a Markov chain with transition kernel  $P$  and initial distribution  $p_0$ , i.e., we have almost surely
$$
\mathbb {P} \left(\theta_ {n} \in A \mid \theta_ {0}, \dots , \theta_ {n - 1}\right) = \mathbb {P} \left(\theta_ {n} \in A \mid \theta_ {n - 1}\right) = P \left(\theta_ {n - 1}, A\right), \tag {2.6}
$$
and  $p_0(A) = \mathbb{P}(\theta_0\in A)$  for any measurable set  $A\subseteq \mathbb{R}^d$  and  $n\in \mathbb{N}$ . We assume that  $(\hat{\theta}_n)_{n = 0}^\infty$  is another Markov chain with transition kernel  $\hat{P}$  and initial distribution  $\hat{p}_0$ . We denote by  $p_n$  the distribution of  $\theta_{n}$  and by  $\hat{p}_n$  the distribution of  $\hat{\theta}_n$ . By  $\delta_{\theta}$ , we denote the Dirac delta distribution at  $\theta$ , i.e. the probability measure concentrated at  $\theta$ . For a measurable set  $A\subseteq \mathbb{R}^d$ , we also use the notation  $\delta_{\theta}P(A)\coloneqq P(\theta ,A)$ .
Lemma 2.1 ([RS18], Theorem 3.1). Assume that there exist some  $\rho \in [0,1)$  and  $C\in (0,\infty)$  such that
$$
\sup  _ {\theta , \tilde {\theta} \in \mathbb {R} ^ {d}: \theta \neq \tilde {\theta}} \frac {\mathcal {W} _ {1} \left(P ^ {n} (\theta , \cdot) , P ^ {n} (\tilde {\theta} , \cdot)\right)}{\| \theta - \tilde {\theta} \|} \leq C \rho^ {n}, \tag {2.7}
$$
for any  $n \in \mathbb{N}$ . Further assume that there exist some  $\delta \in (0,1)$  and  $L \in (0,\infty)$  and a measurable Lyapunov function  $\hat{V} : \mathbb{R}^d \to [1,\infty)$  of  $\hat{P}$  such that for any  $\theta \in \mathbb{R}^d$ :
$$
(\hat {P} \hat {V}) (\theta) \leq \delta \hat {V} (\theta) + L, \tag {2.8}
$$
where  $(\hat{P}\hat{V})(\theta)\coloneqq \int_{\mathbb{R}^d}\hat{V} (\hat{\theta})\hat{P} (\theta ,d\hat{\theta}).$  Then, we have
$$
\mathcal {W} _ {1} \left(p _ {n}, \hat {p} _ {n}\right) \leq C \left(\rho^ {n} \mathcal {W} _ {1} \left(p _ {0}, \hat {p} _ {0}\right) + \left(1 - \rho^ {n}\right) \frac {\gamma \kappa}{1 - \rho}\right), \tag {2.9}
$$
where  $\gamma \coloneqq \sup_{\theta \in \mathbb{R}^d}\frac{\mathcal{W}_1(\delta_\theta P,\delta_\theta\hat{P})}{\hat{V}(\theta)}$ $\kappa \coloneqq \max \left\{\int_{\mathbb{R}^d}\hat{V} (\theta)d\hat{p}_0(\theta),\frac{L}{1 - \delta}\right\} .$
Lemma 2.1 provides a sufficient condition for the distributions  $p_n$  and  $\hat{p}_n$  after  $n$  iterations to stay close to each other given the initial distributions  $p_0$ . Lemma 2.1 will provide a key role in helping us derive the main results in Section 3.1 and Section 3.2. Later, in the Appendix, we will state and prove a modification of Lemma 2.1 (see Lemma E.5 in the Appendix) that will be crucial to obtaining the main result in Section 3.3.
3 Wasserstein Stability of SGD via Markov Chain Perturbations
In this section, we will derive time-uniform Wasserstein stability bounds for SGD by using the perturbation theory presented in [RS18]. Before considering general losses that can be non-convex, we first consider the simpler case of quadratic losses to illustrate our key ideas.
3.1 Warm up: quadratic case
To illustrate the proof technique, we start by considering a quadratic loss of the form:  $f(\theta, x_i) \coloneqq (a_i^\top \theta - y_i)^2 / 2$  where,  $x_i \coloneqq (a_i, y_i)$  and  $\nabla f(\theta, x_i) = a_i (a_i^\top \theta - y_i)$ . In this setting, the SGD recursion takes the following form:
$$
\theta_ {k} = \left(I - \frac {\eta}{b} H _ {k}\right) \theta_ {k - 1} + \frac {\eta}{b} q _ {k}, \quad \text {w h e r e ,} \quad H _ {k} := \sum_ {i \in \Omega_ {k}} a _ {i} a _ {i} ^ {\top}, \quad q _ {k} := \sum_ {i \in \Omega_ {k}} a _ {i} y _ {i}. \tag {3.1}
$$
The sequence  $(H_{k},q_{k})$  are i.i.d. and for every  $k$ ,  $(H_{k},q_{k})$  is independent of  $\theta_{k - 1}$ .
Similarly, we can write down the iterates of SGD with a different data set  $\hat{X}_n \coloneqq \{\hat{x}_1, \dots, \hat{x}_n\}$  with  $\hat{x}_i = (\hat{a}_i, \hat{y}_i)$ , where  $\hat{X}_n$  differs from  $X_n$  with at most one element:
$$
\hat {\theta} _ {k} = \left(I - \frac {\eta}{b} \hat {H} _ {k}\right) \hat {\theta} _ {k - 1} + \frac {\eta}{b} \hat {q} _ {k}, \quad \text {w h e r e} \quad \hat {H} _ {k} := \sum_ {i \in \Omega_ {k}} \hat {a} _ {i} \hat {a} _ {i} ^ {\top}, \quad \hat {q} _ {k} := \sum_ {i \in \Omega_ {k}} \hat {a} _ {i} \hat {y} _ {i}. \tag {3.2}
$$
Our goal is to obtain an algorithmic stability bound, through estimating the 1-Wasserstein distance between the distribution of  $\theta_{k}$  and  $\hat{\theta}_{k}$  and we will now illustrate the three-step proof technique that we described in Section 1. To be able to apply the perturbation theory [RS18], we start by establishing the geometric ergodicity of the Markov process  $(\theta_{k})_{k\geq 0}$  with transition kernel  $P(\theta ,\cdot)$ , given in the following lemma.
Lemma 3.1. Assume that  $\rho \coloneqq \mathbb{E}\left\| I - \frac{\eta}{b} H_1\right\| < 1$ . Then, for any  $k \in \mathbb{N}$ , we have the following inequality:  $\mathcal{W}_1\left(P^k (\theta ,\cdot),P^k (\tilde{\theta},\cdot)\right) \leq \rho^k\| \theta -\tilde{\theta}\|$ .
We note that since  $H_{1} \succeq 0$ , the assumption in Lemma 3.1 can be satisfied under mild assumptions, for example when  $H_{1} \succ 0$  with a positive probability, which is satisfied for  $\eta$  small enough.
In the second step, we construct a Lyapunov function  $\hat{V}$  that satisfies the conditions of Lemma 2.1.
Lemma 3.2. Let  $\hat{V}(\theta) \coloneqq 1 + \| \theta \|$ . Assume that  $\hat{\rho} \coloneqq \mathbb{E}\left\|I - \frac{\eta}{b}\hat{H}_1\right\| < 1$ . Then, we have
$$
(\hat {P} \hat {V}) (\theta) \leq \hat {\rho} \hat {V} (\theta) + 1 - \hat {\rho} + \frac {\eta}{b} \mathbb {E} \| \hat {q} _ {1} \|. \tag {3.3}
$$
In our third and last step, we estimate the perturbation gap based on the Lyapunov function  $\hat{V}$  in the form of (2.7), assuming that the data is bounded. Such bounded data assumptions have been commonly made in the literature [Bac14, BM13].
Lemma 3.3. If  $\sup_{x\in \mathcal{X}}\| x\| \leq D$  for some  $D < \infty$ , then, we have  $\sup_{\theta \in \mathbb{R}^d}\frac{\mathcal{W}_1(\delta_\theta P,\delta_\theta\hat{P})}{\hat{V}(\theta)}\leq \frac{2\eta D^2}{n}$ .
Note that Lemma 2.1 relies on three conditions: the Wasserstein contraction in (2.7), which is obtained through Lemma 3.1, the drift condition for the Lyapunov function in (2.8), which is obtained in Lemma 3.2 and finally the estimate on  $\gamma$  in (2.9) which is about the one-step 1-Wasserstein distance between two semi-groups that in our context are associated with two datasets that differ by at most one element, which is obtained in Lemma 3.3. The only place the neighborhood assumption  $(\sup_{x\in \mathcal{X}}\| x\| \leq D)$  is used is in the expression of  $\gamma$  in equation (2.9). Now, having all the ingredients, we can invoke Lemma 2.1 and we obtain the following result which provides a 1-Wasserstein bound between the distribution of iterates when applied to datasets that differ by one point.
For  $Y \in \bigcup_{n=1}^{\infty} \mathcal{X}^n$  and  $k \geq 0$ , let  $\nu(Y, k)$  denote the law of the  $k$ -th the SGD iterate when  $Y$  is used as the dataset, i.e.,  $\nu(X, k)$  and  $\nu(\hat{X}, k)$  denote the distributions of  $\theta_k$  and  $\hat{\theta}_k$  obtained by the recursions (3.1) and (3.2) respectively. As shorthand notation, set  $\nu_k := \nu(X, k)$  and  $\hat{\nu}_k := \nu(\hat{X}, k)$ .
Theorem 3.1. Assume  $\theta_0 = \hat{\theta}_0 = \theta$ . We also assume that  $\rho \coloneqq \mathbb{E}\left\|I - \frac{\eta}{b}H_1\right\| < 1$  and  $\hat{\rho} \coloneqq \mathbb{E}\left\|I - \frac{\eta}{b}\hat{H}_1\right\| < 1$  and  $\sup_{x \in \mathcal{X}} \|x\| \leq D$  for some  $D < \infty$ . Then, we have