SlowGuess/NeurIPS_2023_OCR huggingface dataset:
[AFGO19] Necdet Serhat Aybat, Alireza Fallah, Mert Gurbuzbalaban, and Asuman Ozdaglar. A universally optimal multistage accelerated stochastic gradient method. In Advances in Neural Information Processing Systems, volume 32, 2019.
[AFGO20] Necdet Serhat Aybat, Alireza Fallah, Mert Gurbuzbalaban, and Asuman Ozdaglar. Robust accelerated gradient methods for smooth strongly convex functions. SIAM Journal on Optimization, 30(1):717-751, 2020.
[AS23] Shunta Akiyama and Taiji Suzuki. Excess risk of two-layer ReLU neural networks in teacher-student settings and its superiority to kernel methods. In International Conference on Learning Representations, 2023.
[Bac14] Francis Bach. Adaptivity of averaged stochastic gradient descent to local strong convexity for logistic regression. Journal of Machine Learning Research, 15(1):595-627, 2014.
[BE02] Olivier Bousquet and André Elisseeff. Stability and generalization. Journal of Machine Learning Research, 2(Mar):499-526, 2002.
[Ber15] Dimitri Bertsekas. Convex Optimization Algorithms. Athena Scientific, 2015.
[BFGT20] Raef Bassily, Vitaly Feldman, Cristóbal Guzmán, and Kunal Talwar. Stability of stochastic gradient descent on nonsmooth convex losses. In Advances in Neural Information Processing Systems, volume 33, pages 4381-4391, 2020.
[BM13] Francis Bach and Eric Moulines. Non-strongly-convex smooth stochastic approximation with convergence rate  $\mathcal{O}(1/n)$ . In Advances in Neural Information Processing Systems, volume 26, 2013.
[CGZ19] Bugra Can, Mert Gurbüzbalaban, and Lingjiong Zhu. Accelerated linear convergence of stochastic momentum methods in Wasserstein distances. In International Conference on Machine Learning, pages 891-901. PMLR, 2019.
[CJY18] Yuansi Chen, Chi Jin, and Bin Yu. Stability and convergence trade-off of iterative optimization algorithms. arXiv preprint arXiv:1804.01619, 2018.
[Dun81] Joseph C Dunn. Global and asymptotic convergence rate estimates for a class of projected gradient processes. SIAM Journal on Control and Optimization, 19(3):368-400, 1981.
[EHZ22] Murat A. Erdogdu, Rasa Hosseinzadeh, and Matthew S. Zhang. Convergence of Langevin Monte Carlo in Chi-squared and Rényi divergence. In Proceedings of the 25th International Conference on Artificial Intelligence and Statistics (AISTATS), volume 151. PMLR, 2022.
$\left[\mathrm{FGO}^{+}22\right]$  Alireza Fallah, Mert Gurbuzbalaban, Asuman Ozdaglar, Umut Şimsekli, and Lingjiong Zhu. Robust distributed accelerated stochastic gradient methods for multi-agent networks. Journal of Machine Learning Research, 23(1):9893-9988, 2022.
[FR21] Tyler Farghly and Patrick Rebeschini. Time-independent generalization bounds for SGLD in non-convex settings. In Advances in Neural Information Processing Systems, volume 34, pages 19836–19846, 2021.
[FV19] Vitaly Feldman and Jan Vondrak. High probability generalization bounds for uniformly stable algorithms with nearly optimal rate. In Conference on Learning Theory, pages 1270-1279. PMLR, 2019.
[GG23] Guillaume Garrigos and Robert M Gower. Handbook of convergence theorems for (stochastic) gradient methods. arXiv preprint arXiv:2301.11235, 2023.
[RRT17] Maxim Raginsky, Alexander Rakhlin, and Matus Telgarsky. Non-convex learning via stochastic gradient Langevin dynamics: A nonasymptotic analysis. In *Conference on Learning Theory*, pages 1674–1703. PMLR, 2017.
[RS18] Daniel Rudolf and Nikolaus Schweizer. Perturbation theory for Markov chains via Wasserstein distance. Bernoulli, 24(4A):2610-2639, 2018.
[RZGS23] Anant Raj, Lingjiong Zhu, Mert Gurbuzbalaban, and Umut Şimsekli. Algorithmic stability of heavy-tailed SGD with general loss functions. In International Conference on Machine Learning, volume 202, pages 28578-28597. PMLR, 2023.
[Vil09] Cédric Villani. Optimal Transport: Old and New. Springer, Berlin, 2009.
[WT11] Max Welling and Yee W Teh. Bayesian learning via stochastic gradient Langevin dynamics. In Proceedings of the 28th International Conference on Machine Learning (ICML-11), pages 681–688, 2011.
[XR17] Aolin Xu and Maxim Raginsky. Information-theoretic analysis of generalization capability of learning algorithms. In Advances in Neural Information Processing Systems, volume 30, 2017.
$\left[\mathrm{ZZB}^{+}22\right]$  Yikai Zhang, Wenjia Zhang, Sammy Bald, Vamsi Pingali, Chao Chen, and Mayank Goswami. Stability of SGD: Tightness analysis and improved bounds. In Uncertainty in Artificial Intelligence, pages 2364-2373. PMLR, 2022.
- In Section A, we provide further details and examples about the usage of surrogate losses.
- In Section B, we provide technical background for the computable bounds for the convergence of Markov chains which will be used to prove the results in Section 3.3 in the main paper.
- In Section C, we provide technical proofs for 1-Wasserstein perturbation results for the quadratic loss in Section 3.1 in the main paper.
- In Section D, we provide technical proofs for 1-Wasserstein perturbation results for the strongly-convex loss in Section 3.2 in the main paper.
- In Section E, we provide technical proofs for 1-Wasserstein perturbation results for the non-convex loss (with additive noise) in Section 3.3 in the main paper.
- In Section F, we provide technical proofs for 2-Wasserstein stability bounds for the nonconvex loss without additive noise in Section 4.1 in the main paper.
- In Section G, we provide technical proofs for  $p$ -Wasserstein stability bounds for the convex loss with additional geometric structure in Section 4.2 in the main paper.
A On the Usage of Surrogate Losses
While the requirement of surrogate losses is a drawback of our framework, nevertheless our setup can cover several practical settings. In this section, we will provide two such examples.
Example 1. We can choose the surrogate loss as the truncated loss, such that:
$$
\ell (\theta , x) = \min  (f (\theta , x), C),
$$
where  $C > 0$  is a chosen constant. This can be seen as a "robust" version of the original loss, which has been widely used in robust optimization and is conceptually similar to adding a projection step to the optimizer.
Example 2. Another natural setup for our framework is the  $\ell_2$ -regularized Lipschitz loss that was also used in [FR21]. As opposed to the previous case, for the sake of this example, let us consider  $\ell$  as the true loss and  $f$  as the surrogate loss. Then, we can choose the pair  $f$  and  $\ell$  as follows:
$$
f (\theta , x) = \ell (\theta , x) + \frac {\mu}{2} \| \theta \| _ {2} ^ {2},
$$
where  $\mu > 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. In other words, we have a loss  $\ell$ ; however, we run the algorithm on the regularized loss  $f$  to have better convergence properties, and finally, we would like to understand if the algorithm generalizes on  $\ell$  or not, and we are typically not interested if the algorithm generalizes well on the regularized loss  $f$ .
Next, we illustrate how a generalization bound for the loss  $f$ , i.e.,  $\left|\mathbb{E}[\hat{F}(\theta) - F(\theta)]\right|$ . For this example, a bound on the quantity can be obtained by building on our analysis. To obtain such a bound, in addition to the bounds that we developed on  $\left|\mathbb{E}[\hat{R}(\theta) - R(\theta)]\right|$ , we would need to estimate the following quantity:
$$
\left| \mathbb {E} _ {\theta , X _ {n}} \left[ \frac {1}{n} \sum_ {i = 1} ^ {n} \left(f \left(\theta , x _ {i}\right) - \ell \left(\theta , x _ {i}\right)\right) \right] \right|.
$$
For illustration purposes, assume that  $\ell$  is convex and Lipschitz in the first parameter. Then,  $f$  is  $\mu$ -strongly convex. Further consider that we initialize SGD from 0, i.e.,  $\theta_0 = 0$  and set the batch
size  $b$  to 1. Denote  $\theta = \theta_{k}$  as the  $k$ -th iterate of SGD when applied on  $\hat{F}(\theta, X_{n})$ , i.e., (1.1). Further define the minimum:
$$
\theta_ {X _ {n}} ^ {\star} = \arg \min  _ {\theta} \hat {F} (\theta , X _ {n}).
$$
We can now analyze the error induced by the surrogate loss as follows:
$$
\begin{array}{l} \left| \mathbb {E} _ {\theta , X _ {n}} \left[ \frac {1}{n} \sum_ {i = 1} ^ {n} \left(f \left(\theta , x _ {i}\right) - \ell \left(\theta , x _ {i}\right)\right) \right] \right| = \frac {\mu}{2} \mathbb {E} _ {\theta , X _ {n}} \| \theta \| ^ {2} = \frac {\mu}{2} \mathbb {E} _ {\theta , X _ {n}} \left\| \theta - \theta_ {X _ {n}} ^ {\star} + \theta_ {X _ {n}} ^ {\star} \right\| ^ {2} \\ \leq \mu \mathbb {E} _ {\theta , X _ {n}} \left\| \theta - \theta_ {X _ {n}} ^ {\star} \right\| ^ {2} + \mu \mathbb {E} _ {X _ {n}} \left\| \theta_ {X _ {n}} ^ {\star} \right\| ^ {2} \\ \leq \mu \mathbb {E} _ {X _ {n}} \left[ (1 - \eta \mu) ^ {k} \left\| \theta_ {X _ {n}} ^ {\star} \right\| ^ {2} + \frac {2 \eta}{\mu} \sigma_ {X _ {n}} \right] + \mu \mathbb {E} _ {X _ {n}} \left\| \theta_ {X _ {n}} ^ {\star} \right\| ^ {2} \\ = \mu \left(\left(1 - \eta \mu\right) ^ {k} + 1\right) \mathbb {E} _ {X _ {n}} \left\| \theta_ {X _ {n}} ^ {\star} \right\| ^ {2} + 2 \eta \mathbb {E} _ {X _ {n}} \left[ \sigma_ {X _ {n}} \right]. \\ \end{array}
$$
Here, the second inequality follows from standard convergence analysis for SGD [GG23, Theorem 5.7] and we define  $\sigma_{X_n}$  as the stochastic gradient noise variance:
$$
\sigma_ {X _ {n}} := \operatorname {V a r} \left[ \nabla f \left(\theta_ {X _ {n}} ^ {\star}, x _ {i}\right) \right],
$$
where for a random vector  $V$  we define  $\mathrm{Var}[V] \coloneqq \mathbb{E} \| V - \mathbb{E}[V] \|^2$ . Hence, we can see that the error induced by the surrogate loss depends on the following factors:
- The regularization parameter  $\mu$ ,
- The expected norm of the minimizers,
- The step-size  $\eta$ ,
- The expected stochastic gradient noise variance.
These terms can be controlled by adjusting  $\mu$  and  $\eta$ .
B.1 Computable bounds for the convergence of Markov chains
Geometric ergodicity and convergence rate of Markov chains has been well studied in the literature [MT93, MT94, HM11]. In this section, we state a result from [HM11] that provides an explicitly computable bound on the Wasserstein contraction for the Markov chains that satisfies a drift condition that relies on the construction of an appropriate Lyapunov function and a minorization condition.
Let  $\mathcal{P}(\theta, \cdot)$  be a Markov transition kernel for a Markov chain  $(\theta_k)$  on  $\mathbb{R}^d$ . For any measurable function  $\varphi: \mathbb{R}^d \to [0, +\infty]$ , we define:
$$
(\mathcal {P} \varphi) (\theta) = \int_ {\mathbb {R} ^ {d}} \varphi (\tilde {\theta}) \mathcal {P} (\theta , d \tilde {\theta}).
$$
Assumption B.1 (Drift Condition). There exists a function  $V: \mathbb{R}^d \to [0, \infty)$  and some constants  $K \geq 0$  and  $\gamma \in (0, 1)$  so that
$$
(\mathcal {P} V) (\theta) \leq \gamma V (\theta) + K,
$$
for all  $\theta \in \mathbb{R}^d$
Assumption B.2 (Minorization Condition). There exists some constant  $\hat{\eta} \in (0,1)$  and a probability measure  $\nu$  so that
$$
\inf_{\theta \in \mathbb{R}^{d}:V(\theta)\leq R}\mathcal{P}(\theta ,\cdot)\geq \hat{\eta}\nu (\cdot),
$$
for some  $R > 2K / (1 - \gamma)$
We define the weighted total variation distance:
$$
d _ {\psi} (\mu_ {1}, \mu_ {2}) = \int_ {\mathbb {R} ^ {d}} (1 + \psi V (\theta)) | \mu_ {1} - \mu_ {2} | (d \theta),
$$
where  $\psi > 0$  and  $V(\theta)$  is the Lyapunov function that satisfies the drift condition (Assumption B.1). It is known that  $d_{\psi}$  has the following alternative expression [HM11]:
$$
d _ {\psi} (\mu_ {1}, \mu_ {2}) = \sup  _ {\varphi : \| \varphi \| _ {\psi} \leq 1} \int_ {\mathbb {R} ^ {d}} \varphi (\theta) (\mu_ {1} - \mu_ {2}) (d \theta),
$$
where  $\| \cdot \|_{\psi}$  is the weighted supremum norm such that for any  $\psi > 0$ :
$$
\| \varphi \| _ {\psi} := \sup  _ {\theta \in \mathbb {R} ^ {d}} \frac {| \varphi (\theta) |}{1 + \psi V (\theta)}.
$$
It is also noted in [HM11] that  $d_{\psi}$  has yet another equivalent expression:
$$
d _ {\psi} \left(\mu_ {1}, \mu_ {2}\right) = \sup  _ {\varphi : \| \left| \left| \varphi \right| \right| _ {\psi} \leq 1} \int_ {\mathbb {R} ^ {d}} \varphi (\theta) \left(\mu_ {1} - \mu_ {2}\right) (d \theta),
$$
$$
| | \varphi | | _ {\psi} := \sup  _ {\theta \neq \tilde {\theta}} \frac {| \varphi (\theta) - \varphi (\tilde {\theta}) |}{2 + \psi V (\theta) + \psi V (\tilde {\theta})}.
$$
Lemma B.1 (Theorem 1.3. [HM11]). If the drift condition (Assumption B.1) and minorization condition (Assumption B.2) hold, then there exists  $\bar{\eta} \in (0,1)$  and  $\psi > 0$  so that
$$
d _ {\psi} \left(\mathcal {P} \mu_ {1}, \mathcal {P} \mu_ {2}\right) \leq \bar {\eta} d _ {\psi} \left(\mu_ {1}, \mu_ {2}\right)
$$
for any probability measures  $\mu_1, \mu_2$  on  $\mathbb{R}^d$ . In particular, for any  $\eta_0 \in (0, \hat{\eta})$  and  $\gamma_0 \in (\gamma + 2K / R, 1)$  one can choose  $\psi = \eta_0 / K$  and  $\bar{\eta} = (1 - (\hat{\eta} - \eta_0)) \vee (2 + R\psi\gamma_0) / (2 + R\psi)$ .
C Proofs of Wasserstein Perturbation Results: Quadratic Case
C.1 Proof of Lemma 3.1
Proof. Let  $P^{k}(\theta, \cdot)$  denote the law of  $\theta_{k}$  starting with  $\theta_{0} = \theta$  and  $P^{k}(\tilde{\theta}, \cdot)$  the law of  $\tilde{\theta}_{k}$ :
$$
\tilde {\theta} _ {k} = \left(I - \frac {\eta}{b} H _ {k}\right) \tilde {\theta} _ {k - 1} + \frac {\eta}{b} q _ {k}, \tag {C.1}
$$
with  $\tilde{\theta}_0 = \tilde{\theta}$ . Note that
$$
\theta_ {k} = \left(I - \frac {\eta}{b} H _ {k}\right) \theta_ {k - 1} + \frac {\eta}{b} q _ {k}, \tag {C.2}
$$
$$
\tilde {\theta} _ {k} = \left(I - \frac {\eta}{b} H _ {k}\right) \tilde {\theta} _ {k - 1} + \frac {\eta}{b} q _ {k}, \tag {C.3}
$$
$$
\begin{array}{l} \mathbb {E} \left\| \theta_ {k} - \tilde {\theta} _ {k} \right\| = \mathbb {E} \left\| \left(I - \frac {\eta}{b} H _ {k}\right) \left(\theta_ {k - 1} - \tilde {\theta} _ {k - 1}\right) \right\| \\ \leq \mathbb {E} \left[ \left\| I - \frac {\eta}{b} H _ {k} \right\| \left\| \theta_ {k - 1} - \tilde {\theta} _ {k - 1} \right\| \right] = \rho \mathbb {E} \left\| \theta_ {k - 1} - \tilde {\theta} _ {k - 1} \right\|. \tag {C.4} \\ \end{array}
$$
By iterating over  $j = k, k - 1, \ldots, 1$ , we conclude that
$$
\mathcal {W} _ {1} \left(P ^ {k} (\theta , \cdot), P ^ {k} (\tilde {\theta}, \cdot)\right) \leq \mathbb {E} \| \theta_ {k} - \tilde {\theta} _ {k} \| \leq \rho^ {n} \| \theta_ {0} - \tilde {\theta} _ {0} \| = \rho^ {k} \| \theta - \tilde {\theta} \|. \tag {C.5}
$$
C.2 Proof of Lemma 3.2
Proof. First, we recall that
$$
\hat {\theta} _ {k} = \left(I - \frac {\eta}{b} \hat {H} _ {k}\right) \hat {\theta} _ {k - 1} + \frac {\eta}{b} \hat {q} _ {k}, \tag {C.6}
$$
where  $\hat{H}_k\coloneqq \sum_{i\in \Omega_k}\hat{a}_i\hat{a}_i^\top$  and  $\hat{q}_k\coloneqq \sum_{i\in \Omega_k}\hat{a}_i\hat{y}_i$ . Therefore, starting with  $\hat{\theta}_0 = \theta$ , we have
$$
\hat {\theta} _ {1} = \left(I - \frac {\eta}{b} \hat {H} _ {1}\right) \theta + \frac {\eta}{b} \hat {q} _ {1}, \tag {C.7}
$$
$$
(\hat {P} \hat {V}) (\theta) = \mathbb {E} \hat {V} (\hat {\theta} _ {1}) = 1 + \mathbb {E} \| \hat {\theta} _ {1} \| \leq 1 + \hat {\rho} \| \theta \| + \frac {\eta}{b} \mathbb {E} \| \hat {q} _ {1} \| = \hat {\rho} \hat {V} (\theta) + 1 - \hat {\rho} + \frac {\eta}{b} \mathbb {E} \| \hat {q} _ {1} \|. \tag {C.8}
$$
C.3 Proof of Lemma 3.3
Proof. Let us recall that
$$
\theta_ {1} = \left(I - \frac {\eta}{b} H _ {1}\right) \theta + \frac {\eta}{b} q _ {1}, \tag {C.9}
$$
$$
\hat {\theta} _ {1} = \left(I - \frac {\eta}{b} \hat {H} _ {1}\right) \theta + \frac {\eta}{b} \hat {q} _ {1}, \tag {C.10}
$$
$$
\mathcal {W} _ {1} \left(\delta_ {\theta} P, \delta_ {\theta} \hat {P}\right) \leq \mathbb {E} \left\| H _ {1} - \hat {H} _ {1} \right\| \frac {\eta}{b} \| \theta \| + \frac {\eta}{b} \mathbb {E} \| q _ {1} - \hat {q} _ {1} \|. \tag {C.11}
$$
Since  $X_{n}$  and  $\hat{X}_{n}$  differ by at most one element and  $\sup_{x\in \mathcal{X}}\| x\| \leq D$  for some  $D < \infty$ , we have  $(H_1,q_1) = (\hat{H}_1,q_1)$  with probability  $\frac{n - b}{n}$  and  $(H_1,q_1)\neq (\hat{H}_1,q_1)$  with probability  $\frac{b}{n}$  and moreover
$$
\mathbb {E} \left\| H _ {1} - \hat {H} _ {1} \right\| \leq \frac {b}{n} \max  _ {1 \leq i \leq n} \left\| a _ {i} a _ {i} ^ {\top} - \hat {a} _ {i} \hat {a} _ {i} ^ {\top} \right\| \leq \frac {b}{n} \max  _ {1 \leq i \leq n} \left(\left\| a _ {i} \right\| ^ {2} + \left\| \hat {a} _ {i} \right\| ^ {2}\right) \leq \frac {2 b D ^ {2}}{n}, \tag {C.12}
$$
$$
\mathbb {E} \left\| q _ {1} - \hat {q} _ {1} \right\| \leq \frac {b}{n} \max  _ {1 \leq i \leq n} \| a _ {i} y _ {i} - \hat {a} _ {i} \hat {y} _ {i} \| \leq \frac {b}{n} \max  _ {1 \leq i \leq n} \left(\| a _ {i} \| \| q _ {i} \| + \| \hat {a} _ {i} \| \| \hat {q} _ {i} \|\right) \leq \frac {2 b D ^ {2}}{n}. \tag {C.13}
$$
$$
\sup  _ {\theta \in \mathbb {R} ^ {d}} \frac {\mathcal {W} _ {1} \left(\delta_ {\theta} P , \delta_ {\theta} \hat {P}\right)}{\hat {V} (\theta)} \leq \sup  _ {\theta \in \mathbb {R} ^ {d}} \frac {\frac {\eta}{b} \left(\| \theta \| + 1\right) ^ {\frac {2 b D ^ {2}}{n}}}{1 + \| \theta \|} = \frac {2 \eta D ^ {2}}{n}. \tag {C.14}
$$
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  $\theta_{k}$  and  $\hat{\theta}_k$ , we need to establish a sequence of technical lemmas. First, we show a 1-Wasserstein contraction rate in the following lemma.
Lemma D.1. Assume that Assumption 3.1 and Assumption 3.2 hold, and further assume that  $\eta < \min \left\{\frac{1}{\mu}, \frac{\mu}{K_1^2}\right\}$ . Then, for any  $n \in \mathbb{N}$ ,
$$
\mathcal {W} _ {1} \left(P ^ {n} (\theta , \cdot), P ^ {n} (\tilde {\theta}, \cdot)\right) \leq \left(1 - \frac {\eta \mu}{2}\right) ^ {n} \| \theta - \tilde {\theta} \|. \tag {D.1}
$$
Proof. Let  $P^k(\theta, \cdot)$  denote the law of  $\theta_k$  starting with  $\theta_0 = \theta$ :
$$
\theta_ {k} = \theta_ {k - 1} - \eta \tilde {\nabla} \hat {F} _ {k} \left(\theta_ {k - 1}, X _ {n}\right), \tag {D.2}
$$
and  $P^{k}(\tilde{\theta},\cdot)$  the law of  $\tilde{\theta}_k$
$$
\tilde {\theta} _ {k} = \tilde {\theta} _ {k - 1} - \eta \tilde {\nabla} \hat {F} _ {k} (\tilde {\theta} _ {k - 1}, X _ {n}), \tag {D.3}
$$
$$
\theta_ {k} = \theta_ {k - 1} - \frac {\eta}{b} \sum_ {i \in \Omega_ {k}} \nabla f \left(\theta_ {k - 1}, x _ {i}\right), \tag {D.4}
$$
$$
\tilde {\theta} _ {k} = \tilde {\theta} _ {k - 1} - \frac {\eta}{b} \sum_ {i \in \Omega_ {k}} \nabla f (\tilde {\theta} _ {k - 1}, x _ {i}). \tag {D.5}
$$
$$
\begin{array}{l} \mathbb {E} \left\| \theta_ {k} - \tilde {\theta} _ {k} \right\| ^ {2} = \mathbb {E} \left\| \theta_ {k - 1} - \tilde {\theta} _ {k - 1} - \frac {\eta}{b} \sum_ {i \in \Omega_ {k}} \left(\nabla f \left(\theta_ {k - 1}, x _ {i}\right) - \nabla f \left(\tilde {\theta} _ {k - 1}, x _ {i}\right)\right) \right\| ^ {2} \\ = \mathbb {E} \left\| \theta_ {k - 1} - \tilde {\theta} _ {k - 1} \right\| ^ {2} + \frac {\eta^ {2}}{b ^ {2}} \mathbb {E} \left\| \sum_ {i \in \Omega_ {k}} \left(\nabla f \left(\theta_ {k - 1}, x _ {i}\right) - \nabla f \left(\tilde {\theta} _ {k - 1}, x _ {i}\right)\right) \right\| ^ {2} \\ \left. - \frac {2 \eta}{b} \mathbb {E} \left\langle \theta_ {k - 1} - \tilde {\theta} _ {k - 1}, \sum_ {i \in \Omega_ {k}} \left(\nabla f \left(\theta_ {k - 1}, x _ {i}\right) - \nabla f \left(\tilde {\theta} _ {k - 1}, x _ {i}\right)\right) \right\rangle . \right. \tag {D.6} \\ \end{array}
$$
By applying Assumption 3.1 and Assumption 3.2, we get
$$
\begin{array}{l} \mathbb {E} \left. \left\| \theta_ {k} - \tilde {\theta} _ {k} \right\| ^ {2} \right. \\ \leq (1 - 2 \eta \mu) \mathbb {E} \left\| \theta_ {k - 1} - \tilde {\theta} _ {k - 1} \right\| ^ {2} + \frac {\eta^ {2}}{b ^ {2}} \mathbb {E} \left[ \left(\sum_ {i \in \Omega_ {k}} \left\| \nabla f (\theta_ {k - 1}, x _ {i}) - \nabla f (\tilde {\theta} _ {k - 1}, x _ {i}) \right\|\right) ^ {2} \right] \\ \leq (1 - 2 \eta \mu) \mathbb {E} \left\| \theta_ {k - 1} - \tilde {\theta} _ {k - 1} \right\| ^ {2} + \eta^ {2} K _ {1} ^ {2} \mathbb {E} \left\| \theta_ {k - 1} - \tilde {\theta} _ {k - 1} \right\| ^ {2} \\ \leq (1 - \eta \mu) \mathbb {E} \left\| \theta_ {k - 1} - \tilde {\theta} _ {k - 1} \right\| ^ {2}, \tag {D.7} \\ \end{array}
$$
provided that  $\eta \leq \frac{\mu}{K_1^2}$ . By iterating over  $k = n, n - 1, \ldots, 1$ , we conclude that
$$
\begin{array}{l} \left(\mathcal {W} _ {1} \left(P ^ {n} (\theta , \cdot), P ^ {n} (\tilde {\theta}, \cdot)\right)\right) ^ {2} \leq \left(\mathcal {W} _ {2} \left(P ^ {n} (\theta , \cdot), P ^ {n} (\tilde {\theta}, \cdot)\right)\right) ^ {2} \\ \leq \mathbb {E} \left\| \theta_ {n} - \tilde {\theta} _ {n} \right\| ^ {2} \leq (1 - \eta \mu) ^ {n} \left\| \theta_ {0} - \tilde {\theta} _ {0} \right\| ^ {2} \leq \left(1 - \frac {\eta \mu}{2}\right) ^ {2 n} \left\| \theta - \tilde {\theta} \right\| ^ {2}. \tag {D.8} \\ \end{array}
$$
Next, we construct a Lyapunov function  $\hat{V}$  and obtain a drift condition for the SGD  $(\hat{\theta}_k)_{k=0}^{\infty}$ .
Lemma D.2. Assume that Assumption 3.1 and Assumption 3.2 hold. Let  $\hat{V}(\theta) := 1 + \|\theta - \hat{\theta}_*\|^2$ , where  $\hat{\theta}_*$  is the minimizer of  $\hat{F}(\theta, \hat{X}_n) := \frac{1}{n} \sum_{i=1}^{n} \nabla f(\theta, \hat{x}_i)$ . Assume that  $\eta < \min \left\{\frac{1}{\mu}, \frac{\mu}{K_1^2 + 64D^2 K_2^2}\right\}$  and  $\sup_{x \in \mathcal{X}} \|x\| \leq D$  for some  $D < \infty$ . Then, we have
$$
(\hat {P} \hat {V}) (\theta) \leq (1 - \eta \mu) \hat {V} (\theta) + 2 \eta \mu - \eta^ {2} K _ {1} ^ {2} - 5 6 \eta^ {2} D ^ {2} K _ {2} ^ {2} + 6 4 \eta^ {2} D ^ {2} K _ {2} ^ {2} \| \hat {\theta} _ {*} \| ^ {2}. \tag {D.9}
$$
$$
\hat {\theta} _ {k} = \hat {\theta} _ {k - 1} - \frac {\eta}{b} \sum_ {i \in \Omega_ {k}} \nabla f (\hat {\theta} _ {k - 1}, \hat {x} _ {i}). \tag {D.10}
$$
Therefore, starting with  $\hat{\theta}_0 = \theta$ , we have
$$
\begin{array}{l} \hat {\theta} _ {1} = \theta - \frac {\eta}{b} \sum_ {i \in \Omega_ {1}} \nabla f (\theta , \hat {x} _ {i}) \\ = \theta - \frac {\eta}{n} \sum_ {i = 1} ^ {n} \nabla f \left(\theta , \hat {x} _ {i}\right) + \eta \left(\frac {1}{n} \sum_ {i = 1} ^ {n} \nabla f \left(\theta , \hat {x} _ {i}\right) - \frac {1}{b} \sum_ {i \in \Omega_ {1}} \nabla f \left(\theta , \hat {x} _ {i}\right)\right). \tag {D.11} \\ \end{array}
$$
$$
\mathbb {E} \left[ \frac {1}{b} \sum_ {i \in \Omega_ {1}} \nabla f (\theta , \hat {x} _ {i}) | \theta \right] = \frac {1}{n} \sum_ {i = 1} ^ {n} \nabla f (\theta , \hat {x} _ {i}). \tag {D.12}
$$
$$
\begin{array}{l} (\hat {P} \hat {V}) (\theta) \\ = \mathbb {E} \hat {V} (\hat {\theta} _ {1}) = 1 + \mathbb {E} \left\| \hat {\theta} _ {1} - \hat {\theta} _ {*} \right\| ^ {2} \\ = 1 + \left\| \theta - \hat {\theta} _ {*} + \frac {\eta}{n} \sum_ {i = 1} ^ {n} \nabla f \left(\theta , \hat {x} _ {i}\right) \right\| ^ {2} + \eta^ {2} \mathbb {E} \left\| \frac {1}{n} \sum_ {i = 1} ^ {n} \nabla f \left(\theta , \hat {x} _ {i}\right) - \frac {1}{b} \sum_ {i \in \Omega_ {1}} \nabla f \left(\theta , \hat {x} _ {i}\right) \right\| ^ {2}. \tag {D.13} \\ \end{array}
$$
We can compute that
$$
\begin{array}{l} \left\| \theta - \hat {\theta} _ {*} + \frac {\eta}{n} \sum_ {i = 1} ^ {n} \nabla f (\theta , \hat {x} _ {i}) \right\| ^ {2} \\ = \left\| \theta - \hat {\theta} _ {*} + \frac {\eta}{n} \sum_ {i = 1} ^ {n} \left(\nabla f \left(\theta , \hat {x} _ {i}\right) - \nabla f \left(\hat {\theta} _ {*}, \hat {x} _ {i}\right)\right) \right\| ^ {2} \\ \leq (1 - 2 \eta \mu) \| \theta - \hat {\theta} _ {*} \| ^ {2} + \frac {\eta^ {2}}{n ^ {2}} \left\| \sum_ {i = 1} ^ {n} \left(\nabla f (\theta , \hat {x} _ {i}) - \nabla f (\hat {\theta} _ {*}, \hat {x} _ {i})\right) \right\| ^ {2} \\ \leq \left(1 - 2 \eta \mu + \eta^ {2} K _ {1} ^ {2}\right) \| \theta - \hat {\theta} _ {*} \| ^ {2}. \tag {D.14} \\ \end{array}
$$
Moreover, we can compute that