SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sijia Zhou<sup>1</sup> Yunwen Lei<sup>2*</sup> Ata Kabán<sup>1</sup>
$^{1}$ School of Computer Science, University of Birmingham, Birmingham B15 2TT, United Kingdom
$^{2}$  Department of Mathematics, The University of Hong Kong, Pokfulam, Hong Kong, China
sxz115@student.bham.ac.uk leiyw@hku.hk a.kaban@bham.ac.uk
We give sharper bounds for uniformly stable randomized algorithms in a PAC-Bayesian framework, which improve the existing results by up to a factor of  $\sqrt{n}$  (ignoring a log factor), where  $n$  is the sample size. The key idea is to bound the moment generating function of the generalization gap using concentration of weakly dependent random variables due to Bousquet et al (2020). We introduce an assumption of sub-exponential stability parameter, which allows a general treatment that we instantiate in two applications: stochastic gradient descent and randomized coordinate descent. Our results eliminate the requirement of strong convexity from previous results, and hold for non-smooth convex problems.
Stochastic optimization methods are the workhorses of many modern machine learning problems. A lot of progress has been made in reducing optimization errors with less training time in stochastic settings [6, 21, 41, 46, 64]. One such method is stochastic gradient descent (SGD), which builds a stochastic gradient estimator iteratively based on a randomly sampled example to approximate the gradient for the next iteration. SGD is appealing for large-scale data analysis due to its cheap computational cost, simplicity and efficiency.
Obtaining models that generalize well is the goal in machine learning and we undertake a theoretical analysis of this problem. The generalization behavior of a model can be quantified by the excess risk, which decomposes into two components: the optimization error and the generalization error (or generalization gap). In this paper, we focus on the analysis of the generalization error of stochastic optimization methods. We combine two useful approaches to bound the generalization error, that is algorithmic stability [7, 16] and PAC-Bayes bounds [9, 26, 34, 53]. Our hope is to retain the benefits of both approaches to derive powerful generalization bounds to better understand the behavior of stochastic optimization methods.
Seminal work by Bousquet et al. [7] provided generalization bounds for uniformly stable algorithms in the deterministic case. Sharper generalization bounds were obtained via a moment bound on the generalization gap and a concentration inequality of weakly dependent random variables [8, 18]. Extending stability-based bounds to randomized algorithms is a challenge. In [16], where the stability-based bounds for the randomized algorithms were considered, the results hold for fixed distributions. To give bounds uniformly for all distributions, we need to turn to PAC-Bayes analysis. A recent work analyzed stable algorithms in the PAC-Bayes framework and derived generalization bounds that hold for any distribution [32]. However, a comparison between [8] and [32] reveals that the error convergence rate in the randomized case [32] is much slower than that for the deterministic case [8]. In more detail, it was shown that  $\beta$ -uniformly stable ( $\beta$  is decreasing w.r.t.  $n$ , where  $n$  is the sample
size) and deterministic algorithms would imply generalization bounds of order  $\widetilde{O}(1/\sqrt{n} + \beta)$  [8, 18] $^2$ , while PAC-Bayes bounds of order  $O(1/\sqrt{n} + \sqrt{n}\beta)$  were developed for randomized algorithms [32]. In [32], the randomness comes from the sampling of hyperparameters such as the index of examples chosen in SGD, and the PAC-Bayes bounds hold for any posterior distribution which may depend on the dataset. It is clear that the PAC-Bayes bounds can be slower than those in [8, 18] by a factor of  $\sqrt{n}$ , which motivates a natural question: can we develop PAC-Bayes bounds for randomized algorithms which match the rate of deterministic algorithms?
This paper explores the above question. We provide sharper PAC-Bayes bounds that hold for uniformly stable randomized algorithms. We adapt a moment bound [8] previously used for stable algorithms in deterministic cases and extend it to the PAC-Bayesian framework, to give bounds that hold for randomized predictors. The PAC-Bayes framework is based on the work of [32]. However, we take a different analysis strategy to control the change in hyperparameters, which is based on an assumption of sub-exponential stability parameter. This general assumption allows us to bound the moment generating function (MGF) of the generalization gap within a high probability domain where our assumption holds (see Appendix A.2). Furthermore, we prove that this assumption holds for both SGD and randomized coordinate descent (RCD). We then illustrate the advantage of our results over existing bounds.
Regarding the convergence rate, our main result improves on the existing PAC-Bayes bounds [32] by a factor of  $\sqrt{n}$  (ignoring a log factor). This improvement holds under weaker conditions under which convergence is not guaranteed in [32]. Our primary technical tool is a moment bound, which we extend to randomized learning algorithms in the PAC-Bayes framework. This allows our bounds to hold for all possible posteriors, not just fixed ones [16] or deterministic algorithms [8, 18]. We need to introduce novel techniques to handle the randomness of the hyperparameter, which is a challenge in the PAC-Bayes analysis (details will be given in Section 3.2).
Regarding assumptions, our result holds without the requirement of hyperparameter stability in previous work [32]. Instead, we introduce a new assumption on the sub-exponential behavior of uniform stability by viewing the uniform stability as a function of the random hyperparameter. Interestingly, it suffices to study the sub-exponential behavior of uniform stability under the prior distribution, which makes the sub-exponential assumption easy to check.
We show that the uniform stability of SGD and RCD have a mixture of sub-Gaussian and subexponential tails that satisfy this assumption. Thus, we remove the strong convexity assumption in the existing analysis of SGD [32] and extend the result to non-smooth problems. Our result also applies to RCD, where the randomness arises from the selected coordinates.
The remainder of the paper is organized as follows. We discuss the basics on PAC-Bayes and stability analysis in Section 2. We present our main results in Section 3, and apply it to SGD and RCD in Section 4. We survey the related work on stability and PAC-Bayesian analysis in Section 5, and conclude the paper in Section 6.
We first introduce some notations. Let  $\mathcal{X}$  and  $\mathcal{Y}$  denote the input and output space respectively and let  $\mathcal{Z} = \mathcal{X} \times \mathcal{Y}$ . We are given a set of training examples of size  $n$ ,  $S = \{z_1 = (x_1, y_1), \ldots, z_n = (x_n, y_n)\}$ , drawn independently and identically distributed (i.i.d.) from an unknown distribution  $\mathcal{D}$  on  $\mathcal{Z}$ . We hope to learn a predictor from a class of hypotheses  $\mathcal{H}$  to predict unseen new data drawn from  $\mathcal{D}$ . Let  $\mathcal{W}$  denote the weight space,  $\mathcal{W} \subseteq \mathbb{R}^d$ , and  $\Theta$  denote a hyperparameter space. A deterministic learning algorithm  $A: \mathcal{Z}^n \times \Theta \to \mathcal{H}$  maps the training examples to a hypothesis  $h_{\mathbf{w}} \in \mathcal{H}$  determined by  $\mathbf{w} \in \mathcal{W}$ . The quality of a hypothesis is measured by a loss function,  $\ell: \mathcal{H} \times \mathcal{Z} \to [0, M]$ .
The goal of any learning algorithm is to produce a predictor that generalizes well. That means that the learned predictor applied on previously unseen input data from the marginal of  $\mathcal{D}$  should have a small expected loss. For any  $\theta \in \Theta$ , the risk of a predictor returned by the algorithm  $A$  is a random variable as a function of  $S$ , defined as
$$
R (A (S; \theta)) = \mathbb {E} _ {z \sim \mathcal {D}} [ \ell (A (S; \theta), z) ]. \tag {2.1}
$$
2We use  $\widetilde{O}$  to hide poly-logarithmic factors.
Since  $\mathcal{D}$  is unknown, we have no access to the true risk. Instead, the empirical risk is often used as an approximation
$$
R _ {S} (A (S; \theta)) = \frac {1}{n} \sum_ {i = 1} ^ {n} \ell (A (S; \theta), z _ {i}). \tag {2.2}
$$
We are interested in how well the empirical risk can estimate the risk. The difference between them is the generalization gap  $G(S; \theta) \triangleq R(A(S; \theta)) - R_S(A(S; \theta))$ . We can upper bound the risk by bounding this difference.
2.1 PAC-Bayes basics
We are interested in randomized algorithms, such as stochastic gradient descent (SGD) and randomized coordinate descent (RCD). In such algorithms, there is an in-built random sampling mechanism that we can think of as a random hyperparameter. In other words, a randomized algorithm may be viewed as a deterministic algorithm with hyperparameters  $\theta$  that follow a distribution. Therefore, for the analysis that follows, we will explicitly define distributions on  $\Theta$ . There has been a lot of interest in SGD because it is often used to train a model and its randomness comes from independent sampling of training instances to estimate gradient directions. In this case, the hyperparameters  $\theta \in \Theta$  form a random sequence  $\theta = (\theta_{1}, \dots, \theta_{T})$ , where every  $\theta_{t} \in [n]$ , for  $t \in [T]$ , is an i.i.d. index of a training point from  $S$ . Here  $[n]$  denotes  $\{1, \dots, n\}$ .
Any distribution over the hyperparameter space  $\Theta$  induces a distribution over the predictors (or weights since we often consider parametric models). In the PAC-Bayes methodology, the common approach is to define a prior distribution on the domain of the weights [13, 22, 34, 53]. However, in the context of randomized algorithms of the kind mentioned above it is more natural to exploit the randomness already present through a distribution on the domain of the hyperparameters.
Hence, we will conduct PAC-Bayesian analysis by focusing on the randomness of the hyperparameters. Let  $\mathbb{P}$  be any prior distribution on the hyperparameter domain  $\Theta$ , chosen before seeing the training data. Given a randomized algorithm, we will only need to prove a condition for the prior  $\mathbb{P}$  on  $\Theta$ , and in return our generalization guarantees will hold with high probability for all posteriors  $\mathbb{Q}$  on  $\Theta$ . In particular, we can assume a simple uniform distribution as a prior for SGD, i.e. SGD with uniform sampling, and derive bounds for data-dependent non-uniform sampling [49, 56, 66]. For example, SGD with importance sampling [66] shows better results than uniform sampling, and there are efforts to learn good sampling distributions [32] in a line of research towards self-certified learning algorithms [13, 44].
The quality of a predictor learned by an algorithm  $A$  from the training sample  $S$ , is defined as the expected risk w.r.t the PAC-Bayes posterior  $\mathbb{Q}$  as
$$
R (S, \mathbb {Q}) = \underset {\theta \sim \mathbb {Q}} {\mathbb {E}} [ R (A (S; \theta)) ]. \tag {2.3}
$$
Its empirical counterpart, the expected empirical risk w.r.t.  $\mathbb{Q}$  is
$$
R _ {S} (S, \mathbb {Q}) = \underset {\theta \sim \mathbb {Q}} {\mathbb {E}} \left[ R _ {S} (A (S; \theta)) \right]. \tag {2.4}
$$
In special cases this expectation can be computed analytically; most often it is approximated either by Monte Carlo sampling or by analytical upper bounds.
PAC-Bayes bounds aim to estimate  $R(S, \mathbb{Q})$  in terms of  $R_S(S, \mathbb{Q})$  and the divergence between  $\mathbb{P}$  and  $\mathbb{Q}$ . A key ingredient that PAC-Bayes bounds rest upon is a change of measure inequality, also known as variational formula. For completeness this is given in Appendix (Lemma A.3).
2.2 Algorithmic stability basics
Another relatively recent framework for generalization analysis is based on the concept of algorithmic stability [7]. The key concept in this framework quantifies how sensitive a learning algorithm is to small perturbations of the training data. There are several notions of algorithmic stability, and the one we use in our work is the following uniform stability [7]. We denote  $S \sim S'$  if they are neighboring datasets, i.e.,  $S$  and  $S'$  differ by at most a single example.
Definition 1 (Uniform Stability). For any  $\theta$ , an algorithm  $A: S \mapsto A(S; \theta)$  is  $\beta_{\theta}$ -uniformly stable w.r.t. a loss function  $\ell$  if  $\forall S \sim S' \in \mathcal{Z}^n, \forall z \in \mathcal{Z}$ ,
$$
\left| \ell (A (S; \theta), z) - \ell (A (S ^ {\prime}; \theta), z) \right| \leq \beta_ {\theta}. \tag {2.5}
$$
Recall that  $A$  is a randomized algorithm whose randomness is exclusively due to a random draw of  $\theta$ . Hence, given a fixed instance of  $\theta$ , the algorithm  $A$  becomes a deterministic algorithm. This simple observation helps us reduce the problem from randomized learning to deterministic learning. Based on this observation, the next section presents sharper generalization bounds for randomized algorithms in comparison to previous results [32].
First we introduce a sub-exponential assumption on the stability parameter. This will allow us to make some key innovations: 1) We will be able to remove one of the assumptions required by [32, Theorem 2] (hyperparameter stability) and only require uniform stability. 2) we will state our result in more general terms, and instantiate it to specific algorithms such as SGD and RCD.
Assumption 1 (Sub-exponential stability). Let  $\mathbb{P}$  be a fixed probability distribution on  $\Theta = \prod_{t=1}^{T} \Theta_t$ . We say that a randomized algorithm with random hyperparameters  $\theta \sim \mathbb{P}$  satisfies sub-exponential stability if, for any fixed instance of  $\theta$  it satisfies  $\beta_\theta$ -uniform stability w.r.t. a loss function  $\ell$ , and there exists  $c \in \mathbb{R}$  such that for any  $\delta \in (0, 1/n]$ , with probability at least  $1 - \delta$  over draws of  $\theta \sim \mathbb{P}$ :
$$
\beta_ {\theta} \leq \mathbb {E} _ {\theta \sim \mathbb {P}} [ \beta_ {\theta} ] + c \log (1 / \delta). \tag {3.1}
$$
In other words, this assumption says that the deviation of  $\beta_{\theta}$  from its mean roughly has a subexponential tail. We will show that  $c$  is dominated by  $\mathbb{E}_{\theta \sim \mathbb{P}}[\beta_{\theta}]$  in typical stochastic optimization algorithms such as SGD and RCD.
Note that in the above assumption, the probability is made w.r.t.  $\theta \sim \mathbb{P}$ , which is a prior distribution independent of  $S$ . We can choose  $\mathbb{P}$  to be a uniform distribution. In this case, we will show that SGD satisfies this assumption under mild conditions on the loss functions with very small  $\mathbb{E}_{\theta \sim \mathbb{P}}[\beta_{\theta}]$  and  $c$ .
Recall that for any fixed instance of  $\theta$ , the trained model of the learning algorithm  $A$  on data set  $S$  with hyperparameters  $\theta$  is a deterministic predictor. Our proof strategy for our main result below is to first fix  $\theta$ , and apply a recent technique of obtaining sharp bounds [8] to the resulting deterministic algorithm. After that, we deal with the randomness of  $\theta$ .
In Theorem 1 we give our main result. This bound is useful when we have small  $\mathbb{E}_{\theta \sim \mathbb{P}}[\beta_{\theta}]$  and  $c$ . We will see two examples in Section 4. We denote  $B \gtrsim B'$  if there exists a universal constant  $c_1 > 0$  such that  $B \geq c_1 B'$ . We use  $B \lesssim B'$  if there exists a universal constant  $c_2 > 0$  such that  $B \leq c_2 B'$ . We use  $B \asymp B'$  if  $B \lesssim B' \lesssim B$ . The proof is given in Appendix A.2.
Theorem 1 (Generalization of sub-exponentially stable randomized algorithms). Consider a learning algorithm  $A(S; \theta)$  that satisfies Assumption 1 w.r.t.  $\mathbb{P}$  and  $c \lesssim \mathbb{E}_{\theta \sim \mathbb{P}}[\beta_{\theta}]$ . Assume  $\ell(A(S; \theta), z) \in [0, M]$ . Then for any  $\delta_1 \in (0, 1)$  the following inequality holds with probability at least  $1 - \delta_1$  uniformly for all  $\mathbb{Q}$
$$
\mathbb {E} _ {\theta \sim \mathbb {Q}} \big [ R (A (S; \theta)) - R _ {S} (A (S; \theta)) \big ] \lesssim \Big (D _ {\mathrm {K L}} (\mathbb {Q} | | \mathbb {P}) + \log (1 / \delta_ {1}) \Big) \max  \left\{\mathbb {E} _ {\theta \sim \mathbb {P}} [ \beta_ {\theta} ] \log^ {2} n, \frac {M}{\sqrt {n}} \right\},
$$
where  $D_{\mathrm{KL}}(\mathbb{Q}\| \mathbb{P})$  means the KL divergence between  $\mathbb{P}$  and  $\mathbb{Q}$ , i.e.,  $D_{\mathrm{KL}}(\mathbb{Q}\| \mathbb{P}) = \mathbb{E}_{\theta \sim \mathbb{Q}}\Big[\log \frac{\mathbb{Q}(\theta)}{\mathbb{P}(\theta)}\Big]$ .
To apply Theorem 1, we only need to check the sub-exponential assumption w.r.t. the prior  $\mathbb{P}$ . Then, we use the PAC-Bayesian analysis to transfer this stability assumption w.r.t. the specific  $\mathbb{P}$  to a bound holding for all posterior distributions.
3.1 Comparison
Next, we compare our result (Theorem 1) with the previous generalization bounds on randomized algorithms due to [32] and observe the advantages that our approach offers.
To blend the stability into the PAC-Bayes framework, where hyperparameters  $\theta$  follow a distribution, [32] defined the following new hyperparameter stability assumption for learning algorithms (Definition 2) and obtained the stability-based PAC-Bayes bound given in Theorem 2.
Definition 2 (Hyperparameter Stability). A learning algorithm  $A$  has uniform hyperparameter stability  $\beta_{\Theta}$  w.r.t. the loss function  $\ell$ , if
$$
\sup  _ {S \in \mathcal {Z} ^ {n}} \sup  _ {z \in \mathcal {Z}} \sup  _ {\theta , \theta^ {\prime} \in \Theta : D _ {H} \left(\theta , \theta^ {\prime}\right) = 1} | \ell (A (S; \theta), z) - \ell (A (S; \theta^ {\prime}), z) | \leq \beta_ {\Theta}, \tag {3.2}
$$
where  $D_H(\mathbf{v},\mathbf{v}^{\prime})\triangleq \sum_{i = 1}^{|\mathbf{v}|}\mathbb{I}[v_i\neq v_i^{\prime}]$  is the Hamming distance and  $\mathbb{I}[\cdot ]$  denotes the indicator function, i.e.,  $\mathbb{I}[E] = 1$  if the event  $E$  holds and 0 otherwise.
Observe that, uniform stability (Definition 1) concerns the stability of an algorithm with respect to a change in the training set. In contrast, the above Definition 2 requires stability w.r.t. a change in the hyperparameters. Moreover, the approach in [32] requires stability w.r.t. both the loss function and the hyperparameters to derive the following PAC-Bayes bound.
Theorem 2 (Theorem 2 of [32]). Let  $A$  be a randomized learning algorithm and  $\ell(A(S; \theta), z) \in [0, M]$ . Assume  $A$  is  $\beta_{\theta}$ -uniformly stable w.r.t. loss functions, and  $\beta_{\Theta}$ -uniformly stable w.r.t. hyperparameters. Consider the prior  $\mathbb{P}$  as a fixed probability distribution defined on  $\Theta = \prod_{t=1}^{T} \Theta_t$ . Then for any  $n \geq 1$ ,  $T \geq 1$ , and  $\delta \in (0, 1)$ , the following inequality holds with probability at least  $1 - \delta$  over draws of a data set,  $S \sim \mathbb{D}^n$ , for every posterior  $\mathbb{Q}$  on  $\Theta$
$$
\mathbb {E} _ {\theta \sim \mathbb {Q}} \big [ R (A (S; \theta)) - R _ {S} (A (S; \theta)) \big ] \lesssim \sqrt {\Big (D _ {\mathrm {K L}} (\mathbb {Q} \| \mathbb {P}) + \log (1 / \delta) \Big) \Big ((M + n \mathbb {E} _ {\theta \sim \mathbb {P}} [ \beta_ {\theta} ]) ^ {2} / n + T \beta_ {\Theta} ^ {2} \Big)}.
$$
We now compare our bound (Theorem 1) with the bound above (Theorem 2).
In terms of the rate, our bound in Theorem 1 improves the previous result of Theorem 2 by up to a factor of  $\sqrt{n}$  (up to poly-logarithmic factors when the divergence is polylogarithmic in  $n$  [32]). To see this, consider the case where in both Theorem 1 and Theorem 2, the term that contains the stability parameter dominates over the KL term. Then, in Theorem 1, for  $\beta_{\theta}$ -uniformly stable randomized algorithms, we have
$$
\mathbb {E} _ {\theta \sim \mathbb {Q}} \big [ R (A (S; \theta)) - R _ {S} (A (S; \theta)) \big ] = \max  \Big \{\widetilde {O} (\mathbb {E} _ {\theta \sim \mathbb {P}} [ \beta_ {\theta} ]), \widetilde {O} (n ^ {- \frac {1}{2}}) \Big \} D _ {\mathrm {K L}} (\mathbb {Q} \| \mathbb {P}).
$$
In contrast, Theorem 2 in [32] gives
$$
\mathbb {E} _ {\theta \sim \mathbb {Q}} [ R (A (S; \theta)) - R _ {S} (A (S; \theta)) ] = \widetilde {O} \left(\sqrt {n} \mathbb {E} _ {\theta \sim \mathbb {P}} [ \beta_ {\theta} ] + \sqrt {T} \beta_ {\Theta}\right) D _ {\mathrm {K L}} ^ {\frac {1}{2}} (\mathbb {Q} \| \mathbb {P}). \tag {3.3}
$$
It is clear that, in the case when  $\widetilde{O} (\sqrt{n}\mathbb{E}_{\theta \sim \mathbb{P}}[\beta_{\theta}])$  dominates the KL divergence then having replaced it with  $\widetilde{O} (\mathbb{E}_{\theta \sim \mathbb{P}}[\beta_{\theta}])$  we improved (3.3) by up to a factor of  $\sqrt{n}$  (ignoring a log term). The prior  $\mathbb{P}$  can be selected freely if it is independent of  $S$ . The posterior  $\mathbb{Q}$  can depend on the data. Indeed, a strength of the PAC-Bayesian analysis is that it applies to any  $\mathbb{Q}$ , and therefore it allows to choose a distribution  $\mathbb{Q}$  in a data-dependent manner. The posterior  $\mathbb{Q}$  can be optimized to control the divergence between  $\mathbb{P}$  and  $\mathbb{Q}$ . Therefore, we typically have a small KL divergence.
Observe also that, in Theorem 1, when  $\mathbb{E}_{\theta \sim \mathbb{P}}[\beta_{\theta}] \lesssim O(n^{-\frac{1}{2}})$ , our generalization bound is  $\widetilde{O}(n^{-\frac{1}{2}})$  while (3.3) implies a vacuous bound  $O(1)$  so generalization is not guaranteed. Furthermore, our bounds require  $\mathbb{E}_{\theta \sim \mathbb{P}}[\beta_{\theta}] \lesssim 1/\sqrt{n}$  to get almost optimal rates  $\widetilde{O}(1/\sqrt{n})$ . As a comparison, the results in Eq. (3.3) require stronger condition  $\mathbb{E}_{\theta \sim \mathbb{P}}[\beta_{\theta}] \lesssim 1/n$  and  $\beta_{\Theta} \leq 1/T$  to get the rate  $O(1/\sqrt{n})$ .
In the unlikely case when the KL divergence is the dominating term, for example,  $D_{\mathrm{KL}}(\mathbb{Q}\| \mathbb{P})\gtrsim O(n)$ , then Theorem 2 achieves better result. But in this circumstance, both Theorems 1 and 2 require  $\beta_{\theta}\lesssim O(n^{-\frac{3}{2}})$  to converge at a rate of  $O(n^{-\frac{1}{2}})$ , which we believe this to be also rather uncommon.
In terms of assumptions, we eliminate the hyperparameter stability assumption in [32] at the price of an additional sub-exponential tail assumption on the uniform stability parameter. We will prove later in Section 4 that this general assumption holds for more stochastic optimization methods and even for non-smooth problems. In conclusion, we achieve better results under weaker conditions.
3.2 Proof sketch and challenge in the analysis
We begin by noting that  $\ell$  is applied to the output of  $A(S;\theta)$ , which depends on the training data and the hyperparameters. Therefore,  $\ell(A(S;\theta))$  could be sensitive to changes in both the dataset and the hyperparameters. To address this issue, Theorem 2 assumes hyperparameter stability, which requires small changes in  $\ell(A(S;\theta))$  when the hyperparameters are perturbed.