SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Paul Viallard*
Inria, CNRS, Ecole Normale Supérieure, PSL Research University, Paris, France paul.viallard@inria.fr
Maxime Haddouche*
Inria, University College London and Université de Lille, France maxime.haddouche@inria.fr
Inria, CNRS, Ecole Normale Supérieure PSL Research University, Paris, France umut.simsekli@inria.fr
Benjamin Guedj
Inria and University College London, France and UK benjamin.guedj@inria.fr
Minimising upper bounds on the population risk or the generalisation gap has been widely used in structural risk minimisation (SRM) – this is in particular at the core of PAC-Bayesian learning. Despite its successes and unfailing surge of interest in recent years, a limitation of the PAC-Bayesian framework is that most bounds involve a Kullback-Leibler (KL) divergence term (or its variations), which might exhibit erratic behavior and fail to capture the underlying geometric structure of the learning problem – hence restricting its use in practical applications. As a remedy, recent studies have attempted to replace the KL divergence in the PAC-Bayesian bounds with the Wasserstein distance. Even though these bounds alleviated the aforementioned issues to a certain extent, they either hold in expectation, are for bounded losses, or are nontrivial to minimize in an SRM framework. In this work, we contribute to this line of research and prove novel Wasserstein distance-based PAC-Bayesian generalisation bounds for both batch learning with independent and identically distributed (i.i.d.) data, and online learning with potentially non-i.i.d. data. Contrary to previous art, our bounds are stronger in the sense that (i) they hold with high probability, (ii) they apply to unbounded (potentially heavy-tailed) losses, and (iii) they lead to estimable training objectives that can be used in SRM. As a result we derive novel Wasserstein-based PAC-Bayesian learning algorithms and we illustrate their empirical advantage on a variety of experiments.
Understanding generalisation is one of the main challenges in statistical learning theory, and even more so in modern machine learning applications. Typically, a learning problem is described by a tuple  $(\mathcal{H},\mathcal{Z},\ell)$  consisting of a hypothesis (or predictor) space  $\mathcal{H}$ , a data space  $\mathcal{Z}$ , and a loss function  $\ell :\mathcal{H}\times \mathcal{Z}\to \mathbb{R}$ . The goal is to estimate the population risk of a given hypothesis  $h$ , defined as  $\mathrm{R}_{\mu}(h) = \mathbb{E}_{\mathbf{z}\sim \mu}[\ell (h,\mathbf{z})]$ , where  $\mu$  denotes the unknown data distribution over  $\mathcal{Z}$ .
As  $\mu$  is not known, in practice, a hypothesis  $h$  is usually built by (approximately) minimising the empirical risk, given by  $\hat{\mathbf{R}}_S(h) = \frac{1}{m}\sum_{i=1}^{m}\ell(h,\mathbf{z}_i)$ , where  $\mathcal{S} = \{\mathbf{z}_i\in \mathcal{Z}\}_{i=1}^m$  is a dataset of  $m$  data points, independent and identically distributed (i.i.d.) from  $\mu$ . We define the generalisation gap of a hypothesis  $h$  as  $\hat{\mathbf{R}}_S(h) - \mathbf{R}_\mu(h)$ .
*The authors contributed equally to this work
Developing upper bounds on the generalisation gap, i.e., generalisation bounds has been a longstanding topic in statistical learning. While a plethora of techniques have been introduced, the PAC-Bayesian framework has gained significant traction over the past two decades to provide non-vacuous generalisation guarantees for complex structures such as neural networks during the training phase (see DR17,  $\mathrm{PORPH}^{+}21$ , PRSS21, among others). In these works, the bounds are also used to derive learning algorithms by minimising the right-hand side of a given bound. Beyond neural networks, the flexibility of PAC-Bayes learning makes it a useful toolbox to derive both theoretical results and practical algorithms in various learning fields such as reinforcement learning [FP10], online learning [HG22], multi-armed bandits [SLST  $^+$  11, SLCB  $^+$  12, SAC23], meta-learning [AM18, FM21, RFJK21, RJFK22, DCL  $^+$  21] to name but a few. The PAC-Bayesian bounds focus on a randomised setting where the hypothesis is drawn from a data-dependent distribution  $\rho \in \mathcal{M}(\mathcal{H})$  where  $\mathcal{M}(\mathcal{H})$  denotes the set of probability distributions defined on  $\mathcal{H}$ . A classical PAC-Bayesian result is [Mau04, Theorem 5] (the so-called McAllester bound), which states that, with probability at least  $1 - \delta$ , for any posterior distribution  $\rho \in \mathcal{M}(\mathcal{H})$ ,
$$
\underset {h \sim \rho} {\mathbb {E}} \left[ \mathrm {R} _ {\mu} (h) - \hat {\mathrm {R}} _ {\mathcal {S}} (h) \right] \leq \sqrt {\frac {\mathrm {K L} (\rho \| \pi) + \ln \frac {2 \sqrt {m}}{\delta}}{2 m}},
$$
where  $\pi \in \mathcal{M}(\mathcal{H})$  is any data-free distribution and KL denotes the Kullback-Leibler divergence. In analogy with Bayesian statistics,  $\pi$  is often called the prior, and  $\rho$  is called the posterior - we refer to [Gue19] for a discussion on these terms.
While PAC-Bayesian bounds remain nowadays of the utmost interest to explain generalisation in various learning problems [LGGL19, MGG19, NGG20, MGW20, BG21, HGRS21, ZVM $^{+}$ 21, BG22b, BG22a, BZG22, CSDG22, LFK $^{+}$ 22, FRKP22, BG23, SAC23, RACA23], they mostly rely on the KL divergence or variants which causes two main limitations: (i) as illustrated in the generative modeling literature, the KL divergence does not incorporate the underlying geometry or topology of the data space  $\mathcal{Z}$ , hence can behave in an erratic way [ACB17], (ii) the KL divergence and its variants require the posterior  $\rho$  to be absolutely continuous with respect to the prior  $\pi$ . However, recent studies [CDE $^{+}$ 21] have shown that, in stochastic optimisation, the distribution of the iterates, which is the natural choice for the posterior, can converge to a singular distribution, which does not admit a density with respect to the Lebesgue measure. Moreover, the structure of the singularity (i.e., the fractal dimension of  $\rho$ ) depends on the data sample  $S$  [CDE $^{+}$ 21]. Hence, in such a case, it would not be possible to find a suitable prior  $\pi$  that can dominate  $\rho$  for almost every  $S \sim \mu^m$ , which will trivially make  $\mathrm{KL}(\rho \| \pi) = +\infty$  and the generalisation bound vacuous.
Some works have focused on replacing the Kullback-Leibler divergence with more general divergences in PAC-Bayes [AG18, OH21, PWG22], although the problems arising from the presence of the KL divergence in the generalisation bounds are actually not specific to PAC-Bayes: information-theoretic bounds [GMGA17, XR17, RZ20] also suffer from similar issues as they are based on a mutual information term, which is the KL divergence between two distributions. In this context, as a remedy to these issues introduced by the KL divergence, [ZLT18, WDFC19, GBTS21, LN22] proved analogous bounds that are based on the Wasserstein distance, which arises from the theory of optimal transport [Mon81]. As the Wasserstein distance inherits the underlying geometry of the data space and does not require absolute continuity, it circumvents the problems introduced by the KL divergence. Yet, these bounds hold only in expectation, i.e., none of these bounds is holding with high probability over the random choice of the learning sample  $S \sim \mu^m$ .
In the context of PAC-Bayesian learning, the recent works [AEMM22, CL21] incorporated Wasserstein distances as a complexity measure and proved generalisation bounds based on the Wasserstein distance. More precisely, [AEMM22] proved a high-probability generic PAC-Bayesian bound for bounded losses depending on an integral probability metric [Mü97], which contains the Wasserstein distance as a special case. On the other hand, [CL21] exploited PAC-Bayesian tools to obtain learning strategies with their associated regret bounds based on the Wasserstein distance for the online learning setting while requiring a finite hypothesis space and do not deal with generalisation.
Contributions. The theoretical understanding of the high-probability generalisation bounds based on the Wasserstein distance is still limited. The aim of this paper is not only to prove generalisation bounds (for different learning settings) based on the optimal transport theory but also to propose new learning algorithms derived from our theoretical results.
(i) Using the supermartingale toolbox introduced in [HG23a, CWR23], we prove in Section 3.1, novel PAC-Bayesian bounds based on the Wasserstein distance for i.i.d. data. While [AEMM22] proposed a McAllester-like bound for bounded losses, we propose a Catoni-like bound (see e.g., ARC16, Theorem 4.1) valid for heavy-tailed losses with bounded order 2 moments. This assumption is less restrictive than assuming subgaussian or bounded losses, which are at the core of many PAC-Bayes results. This assumption also covers distributions beyond subgaussian or subexponential ones (e.g., gamma distributions with a scale smaller than 1, which have an infinite exponential moment).
(ii) We provide in Section 3.2 the first generalisation bounds based on Wasserstein distances for the online PAC-Bayes framework of [HG22]. Our results are, again, Catoni-like bounds and hold for heavy-tailed losses with bounded order 2 moments. Previous work [CL21] already provided online strategies mixing PAC-Bayes and Wasserstein distances. However, their contributions focus on the best deterministic strategy, regularised by a Wasserstein distance, with respect to the deterministic notion of regret. Our results differ significantly as we provide the best-regularised strategy (still in the sense of a Wasserstein term) with respect to the notion of generalisation, which is new.
(iii) As our bounds are linear with respect to Wasserstein terms (contrary to those of AEMM22), they are well suited for optimisation procedures. Thus, we propose the first PAC-Bayesian learning algorithms based on Wasserstein distances instead of KL divergences. For the first time, we design PAC-Bayes algorithms able to output deterministic predictors (instead of distributions over all  $\mathcal{H}$ ) designed from deterministic priors. This is due to the ability of the Wasserstein distance to measure the discrepancy between Dirac distributions. We then instantiate those algorithms in Section 4 on various datasets, paving the way to promising practical developments of PAC-Bayes learning.
To sum up, we highlight two benefits of PAC-Bayes learning with Wasserstein distance. First, it ships with sound theoretical results exploiting the geometry of the predictor space, holding for heavy-tailed losses. Such a weak assumption on the loss extends the usefulness of PAC-Bayes with Wasserstein distances to a wide range of learning problems, encompassing bounded losses. Second, it allows us to consider deterministic algorithms (i.e., sampling from Dirac measures) designed with respect to the notion of generalisation: we showcase their performance in our experiments.
Outline. Section 2 describes our framework and background, Section 3 contains our new theoretical results and Section 4 gathers our experiments. Appendix A gathers supplementary discussion, Appendix B contains all proofs of our claims, and Appendix C provides insights into our practical results as well as additional experiments.
2 Our framework
**Framework.** We consider a Polish predictor space  $\mathcal{H}$  equipped with a distance  $d$  and a  $\sigma$ -algebra  $\Sigma_{\mathcal{H}}$ , a data space  $\mathcal{Z}$ , and a loss function  $\ell: \mathcal{H} \times \mathcal{Z} \to \mathbb{R}$ . In this work, we consider Lipschitz functions with respect to  $d$ . We also associate a filtration  $(\mathcal{F}_i)_{i \geq 1}$  adapted to our data  $(\mathbf{z}_i)_{i = 1,\dots,m}$  and we assume that the dataset  $\mathcal{S}$  follows the distribution  $\mathcal{D}$ . In the PAC-Bayes learning, we construct a data-driven posterior distribution  $\rho \in \mathcal{M}(\mathcal{H})$  with respect to a prior distribution  $\pi$ .
**Definitions.** For all  $i$ , we denote by  $\mathbb{E}_i[\cdot]$  the conditional expectation  $\mathbb{E}[\cdot|\mathcal{F}_i]$ . In this work, we consider data-dependent priors. A stochastic kernel is a mapping  $\pi:\cup_{m=1}^{\infty}\mathcal{Z}^{m}\times\Sigma_{\mathcal{H}}\to[0,1]$  where  $(i)$  for any  $B\in\Sigma_{\mathcal{H}}$ , the function  $\mathcal{S}\mapsto \pi(\mathcal{S},B)$  is measurable,  $(ii)$  for any dataset  $\mathcal{S}$ , the function  $B\mapsto \pi(\mathcal{S},B)$  is a probability measure over  $\mathcal{H}$ .
In what follows, we consider two different learning paradigms: batch learning, where the dataset is directly available, and online learning, where data streams arrive sequentially.
Batch setting. We assume the dataset  $S$  to be i.i.d., so there exists a distribution  $\mu$  over  $\mathcal{Z}$  such that  $\mathcal{D} = \mu^m$ . We then define, for a given  $h \in \mathcal{H}$ , the risk to be  $\mathrm{R}_{\mu} := \mathbb{E}_{\mathbf{z} \sim \mu}[\ell(h, \mathbf{z})]$  and its empirical counterpart  $\hat{\mathrm{R}}_S := \frac{1}{m} \sum_{i=1}^{m} \ell(h, \mathbf{z}_i)$ . Our results aim to bound the expected generalisation gap defined by  $\mathbb{E}_{h \sim \rho}[\mathrm{R}_{\mu}(h) - \hat{\mathrm{R}}_S(h)]$ . We assume that the dataset  $S$  is split into  $K$  disjoint sets  $S_1, \ldots, S_K$ . We consider  $K$  stochastic kernels  $\pi_1, \ldots, \pi_K$  such that for any  $S$ , the distribution  $\pi_i(S, .)$  does not depend on  $S_i$ .
Online setting. We adapt the online PAC-Bayes framework of [HG22]. We assume that we have access to a stream of data  $S = (\mathbf{z}_i)_{i=1,\dots,m}$ , arriving sequentially, with no assumption on  $\mathcal{D}$ . In online PAC-Bayes, the goal is to define a posterior sequence  $(\rho_i)_{i\geq 1}$  from a prior sequence  $(\pi_i)_{i\geq 1}$  which can be data-dependent. We define an online predictive sequence  $(\pi_i)_{i=1\dots m}$  satisfying: (i) for all  $i$  and dataset  $S$ , the distribution  $\pi_i(S,.)$  is  $\mathcal{F}_{i-1}$  measurable and (ii) there exists  $\pi_0$  such that for all  $i\geq 1$ , we have  $\pi_i(S,.)\gg \pi_0$ . This last condition covers, in particular, the case where  $\mathcal{H}$  is an Euclidean space and for any  $i$ , the distribution  $\pi_{i,S}$  is a Dirac mass. All of those measures are uniformly continuous with respect to any Gaussian distribution.
Wasserstein distance. In this paper, we focus on the Wasserstein distance of order 1 (a.k.a., Earth Mover's distance) introduced by [Kan60] in the optimal transport literature. Given a distance  $d: \mathcal{A} \times \mathcal{A} \to \mathbb{R}$  and a Polish space  $(\mathcal{A}, d)$ , for any probability measures  $\alpha$  and  $\beta$  on  $\mathcal{A}$ , the Wasserstein distance is defined by
$$
\mathrm {W} (\alpha , \beta) := \inf  _ {\gamma \in \Gamma (\alpha , \beta)} \left\{\underset {(a, b) \sim \gamma} {\mathbb {E}} d (a, b) \right\}, \tag {1}
$$
where  $\Gamma (\alpha ,\beta)$  is the set of joint probability measures  $\gamma \in \mathcal{M}(\mathcal{A}^2)$  such that the marginals are  $\alpha$  and  $\beta$ . The Wasserstein distance aims to find the probability measure  $\gamma \in \mathcal{M}(\mathcal{A}^2)$  minimising the expected cost  $\mathbb{E}_{(a,b)\sim \gamma}d(a,b)$ . We refer the reader to [Vil09, PC19] for an introduction to optimal transport.
3 Wasserstein-based PAC-Bayesian generalisation bounds
We present novel high-probability PAC-Bayesian bounds involving Wasserstein distances instead of the classical Kullback-Leibler divergence. Our bounds hold for heavy-tailed losses (instead of classical subgaussian and subexponential assumptions), extending the remits of [AEMM22, Theorem 11]. We exploit the supermartingale toolbox, recently introduced in PAC-Bayes framework by [HG23a, CWR23, JJKO23], to derive bounds for both batch learning (Theorems 1 and 2) and online learning (Theorems 3 and 4).
3.1 PAC-Bayes for batch learning with i.i.d. data
In this section, we use the batch setting described in Section 2. We state our first result, holding for heavy-tailed losses admitting order 2 moments. Such an assumption is in line, for instance, with reinforcement learning with heavy-tailed reward (see, e.g., LZ11, LWHZ19, ZS21).
Theorem 1. We assume the loss  $\ell$  to be  $L$ -Lipschitz. Then, for any  $\delta \in (0,1]$ , for any sequence of positive scalar  $(\lambda_i)_{i\in \{1,\dots ,K\}}$ , with probability at least  $1 - \delta$  over the sample  $\mathcal{S}$ , the following holds for the distributions  $\pi_{i,\mathcal{S}}\coloneqq \pi_i(\mathcal{S},.)$  and for any  $\rho \in \mathcal{M}(\mathcal{H})$ :
$$
\begin{array}{l} \underset {h \sim \rho} {\mathbb {E}} \left[ R _ {\mu} (h) - \hat {R} _ {\mathcal {S}} (h) \right] \\ \leq \sum_ {i = 1} ^ {K} \frac {2 | \mathcal {S} _ {i} | L}{m} W (\rho , \pi_ {i, \mathcal {S}}) + \frac {1}{m} \sum_ {i = 1} ^ {K} \frac {\ln \left(\frac {K}{\delta}\right)}{\lambda_ {i}} + \frac {\lambda_ {i}}{2} \left(\underset {h \sim \pi_ {i, \mathcal {S}}} {\mathbb {E}} \left[ \hat {V} _ {| \mathcal {S} _ {i} |} (h) + V _ {| \mathcal {S} _ {i} |} (h) \right]\right), \\ \end{array}
$$
where  $\pi_{i,\mathcal{S}}$  does not depend on  $\mathcal{S}_i$ . Also, for any  $i,|S_i|$ , we have  $\hat{V}_{|\mathcal{S}_i|}(h) = \sum_{\mathbf{z}\in \mathcal{S}_i}(\ell (h,\mathbf{z}) - R_\mu (h))^2$  and  $V_{|\mathcal{S}_i|}(h) = \mathbb{E}_{\mathcal{S}_i}\left[\hat{V}_{|\mathcal{S}_i|}(h)\right]$ .
The proof is deferred to Appendix B.1. While Theorem 1 holds for losses taking values in  $\mathbb{R}$ , many learning problems rely in practice on more constrained losses. This loss can be bounded as in the case of, e.g., supervised learning or the multi-armed bandit problem [Sli19], or simply non-negative as in regression problems involving the quadratic loss (studied, for instance, in Cat16, CG17). Using again the supermartingale toolbox, we prove in Theorem 2 a tighter bound holding for heavy-tailed non-negative losses.
Theorem 2. We assume our loss  $\ell$  to be non-negative and  $L$ -Lipschitz. We also assume that, for any  $1 \leq i \leq K$ , for any dataset  $\mathcal{S}$ , we have  $\mathbb{E}_{h \sim \pi_i(.,\mathcal{S}),z \sim \mu}[\ell(h,z)^2] \leq 1$  (bounded order 2 moments for
prors). Then, for any  $\delta \in (0,1]$ , with probability at least  $1 - \delta$  over the sample  $\mathcal{S}$ , the following holds for the distributions  $\pi_{i,\mathcal{S}}\coloneqq \pi_i(\mathcal{S},.)$  and for any  $\rho \in \mathcal{M}(\mathcal{H})$ :
$$
\underset {h \sim \rho} {\mathbb {E}} \left[ R _ {\mu} (h) - \hat {R} _ {\mathcal {S}} (h) \right] \leq \sum_ {i = 1} ^ {K} \frac {2 | \mathcal {S} _ {i} | L}{m} \mathrm {W} (\rho , \pi_ {i, \mathcal {S}}) + \sum_ {i = 1} ^ {K} \sqrt {\frac {2 | \mathcal {S} _ {i} | \ln \frac {K}{\delta}}{m ^ {2}}},
$$
where  $\pi_{i,\mathcal{S}}$  does not depend on  $S_{i}$ .
Note that when the loss function takes values in  $[0,1]$ , an alternative strategy allows tightening the last term of the bound by a factor  $1/2$ . This result is rigorously stated in Theorem 6 of Appendix B.3.
High-level ideas of the proofs. Theorems 1 and 2 are structured around two tools. First, we exploit the Kantorovich-Rubinstein duality [Vil09, Remark 6.5] to replace the change of measure inequality [Csi75, DV76]; this allows us to consider a Wasserstein distance instead of a KL term. Then, we exploit the supermartingales used in [HG23a, CWR23] alongside Ville's inequality (instead of Markov's one) to obtain a high probability bound holding for heavy-tailed losses. Combining those techniques provides our PAC-Bayesian bounds.
Analysis of our bounds. Our results hold for Lipschitz losses and allow us to consider heavy-tailed losses with bounded order 2 moments. While such an assumption on the loss is more restrictive than in classical PAC-Bayes, allowing heavy-tailed losses is strictly less restrictive. While Theorem 1 is our most general statement, Theorem 2 allows recovering a tighter result (without empirical variance terms) for non-negative heavy-tailed losses. An important point is that the variance terms are considered with respect to the prior distributions  $\pi_{i,\mathcal{S}}$  and not  $\rho$  as in [HG23a, CWR23]. This is crucial as these papers rely on the implicit assumption of order 2 moments, holding uniformly for all  $\rho \in \mathcal{M}(\mathcal{H})$ , while we only require this assumption for the prior distributions  $(\pi_{i,\mathcal{S}})_{i=1,\dots,K}$ . Such an assumption is in line with the PAC-Bayesian literature, which often relies on bounding an averaged quantity with respect to the prior. This strength is a consequence of the Kantorovich-Rubinstein duality. To illustrate this, consider i.i.d. data with distribution  $\mu$  admitting a finite variance bounded by  $V$  and the loss  $\ell(h,z) = |h - z|$  where both  $h$  and  $z$  lie in the real axis. Notice that in this particular case, we can imagine that  $z$  is a data point and  $h$  is a hypothesis outputting the same scalar for all data. To satisfy the assumption of Theorem 2, it is enough, by Cauchy Schwarz, to satisfy  $\mathbb{E}_{h \sim \pi_{i,\mathcal{S}},z \sim \mathcal{S}}[\ell(h,z)^2] \leq \mathbb{E}[h^2] + 2V\mathbb{E}[|h|] + V^2 \leq 1$  for all  $\pi_{i,\mathcal{S}}$ . On the contrary, [HG23a, CWR23] would require this condition to hold for all  $\rho$ , which is more restrictive. Finally, an important point is that our bound allows us to consider Dirac distributions with disjoint support as priors and posteriors. On the contrary, KL divergence forces us to consider a non-Dirac prior for our bound to be non-vacuous. This allows us to retrieve a uniform-convergence bound described in Corollary 7.
Role of data-dependent priors. Theorems 1 and 2 allow the use of prior distributions depending possibly on a fraction of data. Such a dependency is crucial to control our sum of Wasserstein terms as we do not have an explicit convergence rate. For instance, for a fixed  $K$ , consider a compact predictor space  $\mathcal{H}$ , a bounded loss and the Gibbs posterior defined as  $d\rho(h) \propto \exp(-\lambda \hat{\mathbf{R}}_S(h)) dh$  where  $\lambda > 0$ . Also define for any  $i$  and  $\mathcal{S}$ , the distribution  $d\pi_{i,\mathcal{S}}(h) \propto \exp(-\lambda \mathbf{R}_{\mathcal{S} / \mathcal{S}_i}(h)) dh$ . Then, by the law of large numbers, when  $m$  goes to infinity, for any  $h$ , both  $\mathbf{R}_S(h)$  and  $(\mathbf{R}_{\mathcal{S} / \mathcal{S}_i}(h))_{i=1,\dots,m}$  converge to  $\mathbf{R}_\mu(h)$ . This ensures, alongside with the dominated convergence theorem, that for any  $i$ , the Wasserstein distance  $\mathrm{W}(\rho, \pi_{i,\mathcal{S}})$  goes to zero as  $m$  goes to infinity.
Comparison with the literature. [AEMM22, Theorem 11] establishes a PAC-Bayes bound with Wasserstein distance valid for bounded losses being Lipschitz with high probability. While we circumvent the first assumption, the second one is less restrictive than actual Lipschitzness and can also be used in our setting. Also [AEMM22, Theorem 12] proposes an explicit convergence for finite predictor classes. We show in Appendix A that we are also able to recover such a convergence.
Towards new PAC-Bayesian algorithms. From Theorem 2, we derive a new PAC-Bayesian algorithm for Lipschitz non-negative losses:
$$
\underset {\rho \in \mathcal {M} (\mathcal {H})} {\operatorname {a r g m i n}} \mathbb {E} _ {h \sim \rho} \left[ \hat {\mathrm {R}} _ {\mathcal {S}} (h) \right] + \sum_ {i = 1} ^ {K} \frac {2 \left| \mathcal {S} _ {i} \right| L}{m} \mathrm {W} (\rho , \pi_ {i, \mathcal {S}}). \tag {2}
$$
Equation (2) uses Wasserstein distances as regularisers and allows the use of multiple priors. We compare ourselves to the classical PAC-Bayes algorithm derived from [Cat07, Theorem 1.2.6] (which
leads to Gibbs posteriors):
$$
\underset {\rho \in \mathcal {M} (\mathcal {H})} {\operatorname {a r g m i n}} \underset {h \sim \rho} {\mathbb {E}} \left[ \hat {\mathrm {R}} _ {\mathcal {S}} (h) \right] + \frac {\mathrm {K L} (\rho , \pi)}{\lambda}. \tag {3}
$$