SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathcal {W} _ {1} \left(\nu_ {k}, \hat {\nu} _ {k}\right) \leq \frac {1 - \rho^ {k}}{1 - \rho} \frac {2 \eta D ^ {2}}{n} \max  \left\{1 + \| \theta \|, \frac {1 - \hat {\rho} + \frac {\eta}{b} \mathbb {E} \| \hat {q} _ {1} \|}{1 - \hat {\rho}} \right\}. \tag {3.4}
$$
Proof. The result directly follows from Lemma 3.1, Lemma 3.2, Lemma 3.3 and Lemma 2.1.  $\square$
By a direct application of (2.5), we can obtain a generalization bound for an  $\mathcal{L}$ -Lipschitz surrogate loss function, as follows:
$$
\left| \mathbb {E} _ {\mathcal {A}, X _ {n}} \left[ \hat {R} (\mathcal {A} (X _ {n}), X _ {n}) - R (\mathcal {A} (X _ {n})) \right] \right| \leq \frac {\mathcal {L}}{1 - \rho_ {0}} \frac {2 \eta D ^ {2}}{n} \max \left\{1 + \| \theta \|, \frac {1 - \rho_ {0} + \frac {\eta}{b} \mathbb {E} \| \hat {q} _ {1} \|}{1 - \rho_ {0}} \right\},
$$
where  $\rho_0 = \sup_X\| 1 - \frac{\eta}{b} H_X\|$ ,  $H_{X} = \sum_{i\in \Omega_{k},a_{j}\in X}a_{j}a_{j}^{\top}$  and  $X$  is a random set of  $n$ -data points from the data generating distribution. The generalization bound obtained above does not include the mean square error in the unbounded case but covers a larger class of surrogate loss functions. Because of this incompatibility, a direct comparison is not possible; however, the rate obtained in the equation above has the same dependence on the number of samples that were obtained in the previous works [LY20]. For least squares, there are other works using integral operators that develop generalization bounds for SGD under a capacity condition [LR17, PVRB18]. However, these bounds only hold for the least square loss.
3.2 Strongly convex case
Next, we consider strongly convex losses. In the remainder of the paper, we will always assume that for every  $x \in \mathcal{X}$ ,  $f(\cdot, x)$  is differentiable.
Before proceeding to the stability bound, we first introduce the following assumptions.
Assumption 3.1. There exist constants  $K_{1}, K_{2} > 0$  such that for any  $\theta, \hat{\theta} \in \mathbb{R}^{d}$  and every  $x \in \mathcal{X}$ ,
$$
\left\| \nabla f (\theta , x) - \nabla f (\hat {\theta}, \hat {x}) \right\| \leq K _ {1} \| \theta - \hat {\theta} \| + K _ {2} \| x - \hat {x} \| \left(\| \theta \| + \| \hat {\theta} \| + 1\right). \tag {3.5}
$$
This assumption is a pseudo-Lipschitz-like condition on  $\nabla f$  and is satisfied for various problems such as GLMs [Bac14]. Next, we assume that the loss function  $f$  is strongly convex.
Assumption 3.2. There exists a universal constant  $\mu >0$  such that for any  $\theta_{1},\theta_{2}\in \mathbb{R}^{d}$  and  $x\in \mathcal{X}$
$$
\langle \nabla f (\theta_ {1}, x) - \nabla f (\theta_ {2}, x), \theta_ {1} - \theta_ {2} \rangle \geq \mu \| \theta_ {1} - \theta_ {2} \| ^ {2}.
$$
By using the same recipe as we used for quadratic losses, we obtain the following stability result.
Theorem 3.2. Let  $\theta_0 = \hat{\theta}_0 = \theta$ . Assume that Assumption 3.1 and Assumption 3.2 hold. We also assume that  $\eta < \min \left\{\frac{1}{\mu}, \frac{\mu}{K_1^2 + 64D^2K_2^2}\right\}$ ,  $\sup_{x \in \mathcal{X}} \|x\| \leq D$  for some  $D < \infty$ , and  $\sup_{x \in \mathcal{X}} \| \nabla f(0, x) \| \leq E$  for some  $E < \infty$ . Let  $\nu_k$  and  $\hat{\nu}_k$  denote the distributions of  $\theta_k$  and
$\hat{\theta}_k$  respectively. Then, we have
$$
\begin{array}{l} \mathcal {W} _ {1} \left(\nu_ {k}, \hat {\nu} _ {k}\right) \leq \frac {8 D K _ {2} \left(1 - \left(1 - \frac {\eta \mu}{2}\right) ^ {k}\right)}{n \mu} \left(\frac {2 E}{\mu} + 1\right) \\ \left. \cdot \max  \left\{1 + 2 \| \theta \| ^ {2} + \frac {2 E ^ {2}}{\mu^ {2}}, 2 - \frac {\eta}{\mu} K _ {1} ^ {2} - \frac {5 6 \eta}{\mu} D ^ {2} K _ {2} ^ {2} + \frac {6 4 \eta}{\mu^ {3}} D ^ {2} K _ {2} ^ {2} E ^ {2} \right\}. \right. \tag {3.6} \\ \end{array}
$$
Similarly to the quadratic case, we can now directly obtain a bound on expected generalization error using (2.5). More precisely, for an  $\mathcal{L}$ -Lipschitz surrogate loss function  $\ell$ , we have
$$
\begin{array}{l} \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} \cdot \frac {8 D K _ {2} \left(1 - \left(1 - \frac {\eta \mu}{2}\right) ^ {k}\right)}{n \mu} \left(\frac {2 E}{\mu} + 1\right) \\ \cdot \max  \left\{1 + 2 \| \theta \| ^ {2} + \frac {2 E ^ {2}}{\mu^ {2}}, 2 - \frac {\eta}{\mu} K _ {1} ^ {2} - \frac {5 6 \eta}{\mu} D ^ {2} K _ {2} ^ {2} + \frac {6 4 \eta}{\mu^ {3}} D ^ {2} K _ {2} ^ {2} E ^ {2} \right\}. \\ \end{array}
$$
The bound above has the same dependence on the number of samples as the ones of the previous stability analysis of (projected) SGD for strongly convex functions [HRS16, LLNT17, LY20]. However, we have a worse dependence on the strong convexity parameter  $\mu$ .
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  $\theta_1, \theta_2 \in \mathbb{R}^d$  and  $x \in \mathcal{X}$ ,
$$
\langle \nabla f (\theta_ {1}, x) - \nabla f (\theta_ {2}, x), \theta_ {1} - \theta_ {2} \rangle \geq m \| \theta_ {1} - \theta_ {2} \| ^ {2} - K.
$$
The class of dissipative functions satisfying this assumption are the ones that admit some gradient growth in radial directions outside a compact set. Inside the compact set though, they can have quite general non-convexity patterns. As concrete examples, they include certain one-hidden-layer neural networks [AS23]; they arise in non-convex formulations of classification problems (e.g. in logistic regression with a sigmoid/non-convex link function); they can also arise in robust regression problems, see e.g. [GGZ22]. Also, any function that is strongly convex outside of a ball of radius for some will satisfy this assumption. Consequently, regularized regression problems where the loss is a strongly convex quadratic plus a smooth penalty that grows slower than a quadratic will belong to this class; a concrete example would be smoothed Lasso regression; many other examples are also given in [EHZ22]. Dissipative functions also arise frequently in the sampling and Bayesian learning and global convergence in non-convex optimization literature [RRT17, GGZ22].
Unlike the strongly-convex case, we can no longer obtain a Wasserstein contraction bound using the synchronous coupling technique as we did in the proof of Theorem 3.2. To circumvent this problem, in this setting, we consider a noisy version of SGD, with the following recursion:
$$
\theta_ {k} = \theta_ {k - 1} - \eta \tilde {\nabla} \hat {F} _ {k} \left(\theta_ {k - 1}, X _ {n}\right) + \eta \xi_ {k}, \quad \tilde {\nabla} \hat {F} _ {k} \left(\theta_ {k - 1}, X _ {n}\right) := \frac {1}{b} \sum_ {i \in \Omega_ {k}} \nabla f \left(\theta_ {k - 1}, x _ {i}\right), \tag {3.7}
$$
where  $\xi_{k}$  are additional i.i.d. random vectors in  $\mathbb{R}^d$ , independent of  $\theta_{k-1}$  and  $\Omega_{k}$ , satisfying the following assumption.
Assumption 3.4.  $\xi_{1}$  is random vector on  $\mathbb{R}^d$  with a continuous density  $p(x)$  that is positive everywhere, i.e.  $p(x) > 0$  for any  $x\in \mathbb{R}^d$  and  $\mathbb{E}[\xi_1] = 0$ ,  $\sigma^2 \coloneqq \mathbb{E}\left[||\xi_1||^2\right] < \infty$ .
Note that the SGLD algorithm [WT11] is a special case of this recursion, whilst our noise model can accommodate non-Gaussian distributions with finite second-order moment.
Analogously, let us define the (noisy) SGD recursion with the data set  $\hat{X}_n$  as
$$
\hat {\theta} _ {k} = \hat {\theta} _ {k - 1} - \eta \bar {\nabla} \hat {F} _ {k} \left(\hat {\theta} _ {k - 1}, \hat {X} _ {n}\right) + \eta \xi_ {k},
$$
and let  $p(\theta, \theta_1)$  denote the probability density function of  $\theta_1 = \theta - \frac{\eta}{b} \sum_{i \in \Omega_1} \nabla f(\theta, x_i) + \eta \xi_1$ . Further let  $\theta_*$  be a minimizer of  $\hat{F}(\cdot, X_n)$ . Then, by following the same three-step recipe, we obtain the following stability bound. Here, we do not provide all the constants explicitly for the sake of clarity; the complete theorem statement is given in Theorem E.1 (Appendix E.1).
Theorem 3.3. Let  $\theta_0 = \hat{\theta}_0 = \theta$ . Assume that Assumption 3.1, Assumption 3.3 and Assumption 3.4 hold. We also assume that  $\eta < \min \left\{\frac{1}{m}, \frac{m}{K_1^2 + 64D^2K_2^2}\right\}$  and  $\sup_{x \in \mathcal{X}} \|x\| \leq D$  for some  $D < \infty$  and  $\sup_{x \in \mathcal{X}} \| \nabla f(0,x) \| \leq E$  for some  $E < \infty$ . For any  $\hat{\eta} \in (0,1)$ , define  $M > 0$  so that  $\int_{\|\theta_1 - \theta_*\| \leq M} p(\theta_*,\theta_1)d\theta_1 \geq \sqrt{\hat{\eta}}$  and for any  $R > \frac{2K_0}{m}$  where  $K_0$  is defined in (E.2) so that
$$
\inf  _ {\theta , \theta_ {1} \in \mathbb {R} ^ {d}: V (\theta) \leq R, \| \theta_ {1} - \theta_ {*} \| \leq M} \frac {p \left(\theta , \theta_ {1}\right)}{p \left(\theta_ {*} , \theta_ {1}\right)} \geq \sqrt {\hat {\eta}}. \tag {3.8}
$$
Let  $\nu_{k}$  and  $\hat{\nu}_{k}$  denote the distributions of  $\theta_{k}$  and  $\hat{\theta}_k$  respectively. Then, we have
$$
\mathcal {W} _ {1} \left(\nu_ {k}, \hat {\nu} _ {k}\right) \leq \frac {C _ {1} \left(1 - \bar {\eta} ^ {k}\right)}{2 \sqrt {\psi (1 + \psi)} (1 - \bar {\eta})} \cdot \frac {2 b}{n}, \tag {3.9}
$$
where for any  $\eta_0\in (0,\hat{\eta})$  and  $\gamma_0\in \left(1 - m\eta +\frac{2\eta K_0}{R},1\right)$ $\psi = \frac{\eta_0}{\eta K_0}$  and  $\bar{\eta} = (1 - (\hat{\eta} -\eta_0))\vee \frac{2 + R\psi\gamma_0}{2 + R\psi}$  The constant  $C_1\equiv C_1(\psi ,\eta ,\hat{\eta},R,\gamma_0,K_1,K_2,\sigma^2,D,E)$  is explicitly stated in the proof.
Contrary to our previous results, the proof technique for showing Wasserstein contraction (as in Lemma 3.1) for this theorem relies on verifying the drift condition (Assumption B.1) and the minorization condition (Assumption B.2) as given in [HM11]. Once these conditions are satisfied, we invoke the explicitly computable bounds on the convergence of Markov chains developed in [HM11].
From equation (2.5), we directly obtain the following generalization error bound for  $\mathcal{L}$ -Lipschitz surrogate loss function,
$$
\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} \cdot \frac {C _ {1} (1 - \bar {\eta} ^ {k})}{2 \sqrt {\psi (1 + \psi)} (1 - \bar {\eta})} \cdot \frac {2 b}{n},
$$
where the constants are defined in Theorem  $3.3^{1}$ . The above result can be directly compared with the result in [FR21, Theorem 4.1] that has the same dependence on  $n$  and  $b$ . However, our result is more general in the sense that we do not assume our noise to be Gaussian noise. Note that [LLQ19] can also accommodate non-Gaussian noise; however, their bounds increase with the number of iterations.
Remark 3.4. In Theorem 3.3, we can take  $R = \frac{2K_0}{m} (1 + \epsilon)$  for some fixed  $\epsilon \in (0,1)$  so we can take
$$
\hat {\eta} = \left(\max  _ {M > 0} \left\{\min  \left\{\int_ {\| \theta_ {1} - \theta_ {*} \| \leq M} p \left(\theta_ {*}, \theta_ {1}\right) d \theta_ {1}, \inf  _ {\substack {\theta , \theta_ {1} \in \mathbb {R} ^ {d}: \| \theta_ {1} - \theta_ {*} \| \leq M \\ \| \theta - \theta_ {*} \| ^ {2} \leq \frac {2 K _ {0}}{m} (1 + \epsilon) - 1}} \frac {p \left(\theta , \theta_ {1}\right)}{p \left(\theta_ {*}, \theta_ {1}\right)} \right\} \right\}\right) ^ {2}. \tag {3.10}
$$
Moreover, we can take  $\eta_0 = \frac{\hat{\eta}}{2}$ ,  $\gamma_0 = 1 - \frac{m\eta\epsilon}{2}$ , and  $\psi = \frac{\hat{\eta}}{2\eta K_0}$ , so that
$$
\bar {\eta} = \max  \left\{1 - \frac {\hat {\eta}}{2}, \frac {2 + \frac {(1 + \epsilon) \hat {\eta}}{m} (1 - \frac {m \eta \epsilon}{2})}{2 + \frac {(1 + \epsilon) \hat {\eta}}{m}} \right\} = 1 - \frac {m \eta \epsilon (1 + \epsilon) \hat {\eta}}{4 m + 2 (1 + \epsilon) \hat {\eta}}, \tag {3.11}
$$
provided that  $\eta \leq 1$ . Note that the parameter  $\hat{\eta}$  in (3.10) appears in the upper bound in equation (3.9) that controls the 1-Wasserstein algorithmic stability of the SGD. It is easy to see from equation (3.9) that the smaller  $\hat{\eta}$ , the smaller the 1-Wasserstein bound. By the definition of  $\hat{\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 the largest as possible.
Next, let us provide some explicitly computable lower bounds for  $\hat{\eta}$  in (3.10). This is achievable if we specify further the noise assumption. Under the assumption that  $\xi_{k}$  are i.i.d. Gaussian distributed, we have the following corollary.
<sup>1</sup>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(\varepsilon + d^{3/2} b^{-1/4} \lambda^{-1} \log 1 / \varepsilon)$  after  $K = O(d \varepsilon^{-1} \lambda^{-1} \log 1 / \varepsilon)$  iterations, where  $b$  is the mini-batch size and  $\lambda$  is the uniform spectral gap of the continuous-time Langevin dynamics. Similar results are given for SGLD in [XR17, Theorem 3.6].
Corollary 3.5. Under the assumptions in Theorem 3.3, we further assume the noise  $\xi_{k}$  are i.i.d. Gaussian  $\mathcal{N}(0,\Sigma)$  so that  $\mathbb{E}[\| \xi_1\|^2 ] = tr(\Sigma) = \sigma^2$  . We also assume that  $\Sigma \prec I_d$  . Then, we have
$$
\begin{array}{l} \hat {\eta} \geq \max _ {M \geq \eta \sup _ {x \in \mathcal {X}} \| \nabla f (\theta_ {*}, x) \|} \left\{\min  \left\{\left(1 - \frac {\exp (- \frac {1}{2} (\frac {M}{\eta} - \sup _ {x \in \mathcal {X}} \| \nabla f (\theta_ {*} , x) \|) ^ {2})}{\sqrt {\det (I _ {d} - \Sigma)}}\right) ^ {2}, \right. \right. \\ \exp \left\{- \frac {\left(1 + K _ {1} \eta\right) \left(\frac {2 K _ {0}}{m} (1 + \epsilon) - 1\right) ^ {1 / 2}}{\eta^ {2}} \| \Sigma^ {- 1} \| \right. \\ \left. \cdot \left. \left((1 + K _ {1} \eta) \left(\frac {2 K _ {0}}{m} (1 + \epsilon) - 1\right) ^ {1 / 2} + 2 \left(M + \eta \sup  _ {x \in \mathcal {X}} \| \nabla f \left(\theta_ {*}, x\right) \|\right)\right) \right\} \right\}. \tag {3.12} \\ \end{array}
$$
The above corollary provides an explicit lower bound for  $\hat{\eta}$  (instead of the less transparent inequality constraints in Theorem 3.3), and by combining with Remark 3.4 (see equation (3.11)) leads to an explicit formula for  $\bar{\eta}$  which is essential to characterize the Wasserstein upper bound in (3.9) in Theorem 3.3.
4 Wasserstein Stability of SGD without Geometric Ergodicity
While the Markov chain perturbation theory enabled us to develop stability bounds for the case where we can ensure geometric ergodicity in the Wasserstein sense (i.e., proving contraction bounds), we have observed that such a strong ergodicity notion might not hold for non-strongly convex losses. In this section, we will prove two more stability bounds for SGD, without relying on [RS18], hence without requiring geometric ergodicity. To the best of our knowledge, these are the first uniform-time stability bounds for the considered classes of convex and non-convex problems.
4.1 Non-convex case without additive noise
The stability result we obtained in Theorem 3.3 required us to introduce an additional noise (Assumption 3.4) to be able to invoke Lemma 2.1. We will now show that it is possible to use a more direct approach to obtain 2-Wasserstein algorithmic stability in the non-convex case under Assumption 3.3 without relying on [RS18]. However, we will observe that without geometric ergodicity will have a non-vanishing bias term in the bound. Note that, since  $\mathcal{W}_1(\nu_k,\hat{\nu}_k)\leq \mathcal{W}_p(\nu_k,\hat{\nu}_k)$  for all  $p\geq 1$ , the following bound still yields a generalization bound by (2.5).
Theorem 4.1. Assume  $\theta_0 = \hat{\theta}_0 = \theta$ . We also assume that Assumption 3.1 and Assumption 3.3 hold and  $\eta < \min \left\{\frac{1}{m}, \frac{m}{K_1^2 + 64D^2K_2^2}\right\}$  and  $\sup_{x \in \mathcal{X}} \|x\| \leq D$  for some  $D < \infty$  and  $\sup_{x \in \mathcal{X}} \| \nabla f(0, x)\| \leq E$  for some  $E < \infty$ . Let  $\nu_k$  and  $\hat{\nu}_k$  denote the distributions of  $\theta_k$  and  $\hat{\theta}_k$  respectively. Then, we have
$$
\mathcal {W} _ {2} ^ {2} \left(\nu_ {k}, \hat {\nu} _ {k}\right) \leq \left(1 - (1 - \eta m) ^ {k}\right) \cdot \left(\frac {4 D ^ {2} K _ {2} ^ {2} \eta (8 B + 2)}{b n m} + \frac {4 K _ {2} D (1 + K _ {1} \eta)}{n m} (1 + 5 B) + \frac {2 K}{m}\right), \tag {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 \to \infty$ , and it is small only when  $K$  from the dissipativity condition (Assumption 3.3) is small. In other words, if we consider  $K$  to be the level of non-convexity (e.g.,  $K = 0$  corresponds to strong convexity), as the function becomes 'more non-convex' the persistent term in the bound will get larger. While this persistent term might make the bound vacuous when  $n \to \infty$ , for moderate  $n$  the bound can be still informative as the persistent term might be dominated by the first two terms.
Moreover, discarding the persistent bias term, this bound leads to a generalization bound with rate  $n^{-1/2}$ , rather than  $n^{-1}$  as before. This indicates that it is beneficial to add additional noise  $\xi_k$  in SGD as in Theorem 3.3 in order for the dynamics to be geometrically ergodic that can lead to a sharp bound as  $n \to \infty$ . Finally, we note that as Theorem 4.1 involves 2-Wasserstein distance, it can pave the way for generalization bounds without requiring a surrogate loss. Yet, this is not immediate and would require deriving uniform  $L^2$  bounds for the iterates, e.g., [RRT17].
4.2 Convex case with additional geometric structure
We now present our final stability bound, where we consider relaxing the strong convexity assumption (Assumption 3.2) to the following milder assumption.
Assumption 4.1. There exists universal constants  $\mu >0$  and  $p\in (1,2)$  such that for any  $\theta_1,\theta_2\in \mathbb{R}^d$  and  $x\in \mathcal{X}$ ,  $\langle \nabla f(\theta_1,x) - \nabla f(\theta_2,x),\theta_1 - \theta_2\rangle \geq \mu \| \theta_1 - \theta_2\| ^p$ .
Note that as  $p < 2$ , the function class can be seen as an intermediate class between convex and strongly convex functions, and such a class of functions has been studied in the optimization literature [Dun81, Ber15].
We analogously modify Assumption 3.1 and consider the following assumption.
Assumption 4.2. There exist constants  $K_{1}, K_{2} > 0$  and  $p \in (1,2)$  such that for any  $\theta, \hat{\theta} \in \mathbb{R}^{d}$  and every  $x \in \mathcal{X}$ ,  $\| \nabla f(\theta, x) - \nabla f(\hat{\theta}, \hat{x}) \| \leq K_{1} \| \theta - \hat{\theta} \|^{\frac{p}{2}} + K_{2} \| x - \hat{x} \| (\| \theta \|^{p-1} + \| \hat{\theta} \|^{p-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  $\theta_0 = \hat{\theta}_0 = \theta$ . Suppose Assumption 4.1 and Assumption 4.2 hold (with  $p \in (1,2)$ ) and  $\eta \leq \frac{\mu}{K_1^2 + 2^{p+4}D^2K_2^2}$  and  $\sup_{x \in \mathcal{X}} \|x\| \leq D$  for some  $D < \infty$  and  $\sup_{x \in \mathcal{X}} \|\nabla f(0,x)\| \leq E$  for some  $E < \infty$ . Then  $\nu_k$  and  $\hat{\nu}_k$  converge to the unique stationary distributions  $\nu_\infty$  and  $\hat{\nu}_\infty$  respectively and moreover, we have
$$
\mathcal {W} _ {p} ^ {p} \left(\nu_ {\infty}, \hat {\nu} _ {\infty}\right) \leq \frac {C _ {2}}{b n} + \frac {C _ {3}}{n}, \tag {4.2}
$$
where the constants  $C_2 \equiv C_2(\eta, \mu, K_2, D, E)$  and  $C_3 \equiv C_3(\eta, \mu, K_1, K_2, D, E)$  are explicitly stated in the proof.
While we have relaxed the geometric ergodicity condition for this case, in the proof of Theorem 4.2, we show that the Markov chain  $(\theta_k)_{k\geq 0}$  is simply ergodic, i.e.,  $\lim_{k\to \infty}\mathcal{W}_p(\nu_k,\nu_\infty) = 0$ . Hence, even though we still obtain a time-uniform bound, our bound holds asymptotically in  $k$ , due to the lack of an explicit convergence rate for  $\mathcal{W}_p(\nu_k,\nu_\infty)$ . On the other hand, the lack of strong convexity here results in a generalization bound with rate  $n^{-1 / p}$ , whereas for the strongly convex case, i.e.,  $p = 2$ , we previously obtained a rate of  $n^{-1}$ . This might be an indicator that there might be still room for improvement in terms of the rate, at least for this class of loss functions.
We proved time-uniform Wasserstein-stability bounds for SGD and its noisy versions under different strongly convex, convex, and non-convex classes of functions. By making a connection to Markov chain perturbation results [RS18], we introduced a three-step guideline for proving stability bounds for stochastic optimizers. As this approach required geometric ergodicity, we finally relaxed this condition and proved two other stability bounds for a large class of loss functions.
The main limitation of our approach is that it requires Lipschitz surrogate loss functions, as it is based on the Wasserstein distance. Hence, our natural next step will be to extend our analysis without such a requirement. Finally, due to the theoretical nature of this study, it does not contain any direct potential societal impacts.
Lingjiong Zhu is partially supported by the grants NSF DMS-2053454, NSF DMS-2208303, and a Simons Foundation Collaboration Grant. Mert Gurbuzbalaban's research are supported in part by the grants Office of Naval Research Award Number N00014-21-1-2244, National Science Foundation (NSF) CCF-1814888, NSF DMS-2053485. Anant Raj is supported by the Marie Sklodowska-Curie Fellowship (project NN-OVEROPT 101030817). Umut Şimsekli's research is supported by the French government under management of Agence Nationale de la Recherche as part of the "Investissements d'avir" program, reference ANR-19-P3IA-0001 (PRAIRIE 3IA Institute) and the European Research Council Starting Grant DYNASTY - 101039676.