SlowGuess/NeurIPS_2023_OCR huggingface dataset:
[ACB17] Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein Generative Adversarial Networks. In International Conference on Machine Learning (ICML), 2017.
[AEMM22] Ron Amit, Baruch Epstein, Shay Moran, and Ron Meir. Integral Probability Metrics PAC-Bayes Bounds. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[AG18] Pierre Alquier and Benjamin Guedj. Simpler PAC-Bayesian bounds for hostile data. Machine Learning, 107(5), 2018.
[AM18] Ron Amit and Ron Meir. Meta-Learning by Adjusting Priors Based on Extended PAC-Bayes Theory. In International Conference on Machine Learning (ICML), 2018.
[VC68] Vladimir Vapnik and Alexey Chervonenkis. On the uniform convergence of relative frequencies of events to their probabilities. In Doklady Akademii Nauk USSR, 1968.
[VC74] Vladimir Vapnik and Alexey Chervonenkis. Theory of pattern recognition, 1974.
[Vil09] Cédric Villani. Optimal transport: old and new. Number 338 in Grundlehren der mathematischen Wissenschaften. Springer, 2009.
[WDFC19] Hao Wang, Mario Díaz, José Cândido Silveira Santos Filho, and Flávio P. Calmon. An Information-Theoretic View of Generalization via Wasserstein Distance. In IEEE International Symposium on Information Theory (ISIT), 2019.
[XR17] Aolin Xu and Maxim Raginsky. Information-theoretic analysis of generalization capability of learning algorithms. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
[XRV17] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms, 2017.
[Zin03] Martin Zinkevich. Online Convex Programming and Generalized Infinitesimal Gradient Ascent. In International Conference on Machine Learning (ICML), 2003.
[ZLT18] Jingwei Zhang, Tongliang Liu, and Dacheng Tao. An Optimal Transport View on Generalization. arXiv, abs/1811.03270, 2018.
[ZS21] Vincent Zhuang and Yanan Sui. No-Regret Reinforcement Learning with Heavy-Tailed Rewards. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2021.
$\left[\mathrm{ZVM}^{+}21\right]$  Valentina Zantedeschi, Paul Viallard, Emilie Morvant, Rémi Emonet, Amaury Habrard, Pascal Germain, and Benjamin Guedj. Learning Stochastic Majority Votes by Minimizing a PAC-Bayes Generalization Bound. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
The supplementary material is organized as follows:
(i) We provide more discussion about Theorems 1 and 2 in Appendix A;
(ii) The proofs of Theorems 1 to 4 are presented in Appendix B;
(iii) We present in Appendix C additional information about the experiments.
A Additional insights on Section 3.1
In Appendix A.1, we provide additional discussion about Theorem 1 while Appendix A.2 discuss about the convergence rates for Theorem 2.
A.1 Supplementary discussion about Theorem 1
[HG23b, Corollary 10] proposed PAC-Bayes bounds with Wasserstein distances on a Euclidean predictor space with Gaussian prior and posteriors. The bounds have an explicit convergence rate of  $\mathcal{O}\left(\sqrt{dW_1(\rho,\pi) / m}\right)$  where the predictor space is Euclidean with dimension  $d$ . While our bound does not propose such an explicit convergence rate, it allows us to derive learning algorithms as described in Section 4. A broader discussion about the role of  $K$  is detailed in Theorem 2. Furthermore, our bound holds for any Polish predictor space and does not require Gaussian distributions. Furthermore, our result exploits data-dependent priors and deals with the dimension only through the Wasserstein distance, which can attenuate the impact of the dimension.
A.2 Convergence rates for Theorem 2
In this section, we discuss more deeply the values of  $K$  in Theorem 2. This implies a tradeoff between the number of sets  $K$  and the cardinal of each  $S_{i}$ . The tightness of the bound depends highly on the sets  $S_{1},\ldots ,S_{K}$ .
Full batch setting  $\mathbf{K} = \mathbf{1}$ . When  $S_{1} = S$  with  $K = 1$ , the bound of Theorem 2 becomes, with probability  $1 - \delta$ , for any  $\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 2 L \mathrm {W} (\rho , \pi) + 2 \sqrt {\frac {\ln \frac {1}{\delta}}{m}},
$$
where  $\pi = \pi_1$  is data-free. This bound can be seen as the high-probability (PAC-Bayesian) version of the expected bound of [WDFC19]. Furthermore, in this setting, we are able, through our proof technique, to recover an explicit convergence rate similar to the one of [AEMM22, Theorem 12]. It is stated below.
Corollary 5. For any distribution  $\mu$  on  $\mathcal{Z}$ , for any finite hypothesis space  $\mathcal{H}$  equipped with a distance  $d$ , for any  $L$ -Lipschitz loss function  $\ell: \mathcal{H} \times \mathcal{Z} \to [0,1]$ , for any  $\delta \in (0,1]$ , we have, with probability  $1 - \delta$  over the sample  $\mathcal{S}$ , 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 L \sqrt {\frac {2 \ln \left(^ {4} | \mathcal {H} | ^ {2} / \delta\right)}{m}} \mathrm {W} (\rho , \pi) + 2 \sqrt {\frac {\ln \left(^ {2} / \delta\right)}{m}}
$$
where  $\pi$  is a data-free prior.
Proof. We exploit [AEMM22, Equation 35] to state that with probability at least  $1 - \delta /2$ , for any  $(h,h^{\prime})\in \mathcal{H}^{2}$ :
$$
\left| \frac {1}{m} \sum_ {i = 1} ^ {m} \left[ \ell \left(h ^ {\prime}, \mathbf {z} _ {i}\right) - \ell \left(h, \mathbf {z} _ {i}\right) \right] - \underset {\mathbf {z} \sim \mu} {\mathbb {E}} \left[ \ell \left(h ^ {\prime}, \mathbf {z}\right) - \ell (h, \mathbf {z}) \right] \right| \leq L \sqrt {\frac {2 \ln \left(\frac {4 | \mathcal {H} | ^ {2}}{\delta}\right)}{m}} d \left(h, h ^ {\prime}\right).
$$
So, with high probability, we can exploit the Kantorovich-Rubinstein duality with this new Lipschitz constant: with probability at least  $1 - \delta /2$ :
$$
\underset {h \sim \rho} {\mathbb {E}} \left[ \mathrm {R} _ {\mu} (h) - \hat {\mathrm {R}} _ {\mathcal {S}} (h) \right] \leq L \sqrt {\frac {2 \ln \left(\frac {4 | \mathcal {H} | ^ {2}}{\delta}\right)}{m}} \mathrm {W} (\rho , \pi) + \underset {h \sim \pi} {\mathbb {E}} \frac {1}{m} \left[ \sum_ {i = 1} ^ {m} \mathrm {R} _ {\mu} (h) - \ell (h, \mathbf {z} _ {i}) \right],
$$
To conclude, we control the quantity on the right-hand side the same way as in Theorem 1 and Theorem 2. We then have, with probability at least  $1 - \delta /2$ , for a loss function in [0, 1]:
$$
\frac {1}{m} \sum_ {i = 1} ^ {m} \mathbf {R} _ {\mu} (h) - \ell (h, \mathbf {z} _ {i}) \leq 2 \sqrt {\frac {\ln \frac {K}{\delta}}{m}}.
$$
Taking the union bound concludes the proof.
Mini-batch setting  $K = \sqrt{m}$ . When a tradeoff is desired between the quantity of data we want to infuse in our priors and an explicit convergence rate, a meaningful candidate is when  $K = \sqrt{m}$ . Theorem 2's bound becomes, in this particular case:
$$
\underset {h \sim \rho} {\mathbb {E}} \left[ \mathrm {R} _ {\mu} (h) - \hat {\mathrm {R}} _ {S} (h) \right] \leq \frac {2 L}{\sqrt {m}} \sum_ {i = 1} ^ {\sqrt {m}} \mathrm {W} (\rho , \pi_ {i}) + 2 \sqrt {\frac {\ln \frac {\sqrt {m}}{\delta}}{\sqrt {m}}}. \tag {8}
$$
Towards online learning:  $K = m$ . When  $K = m$ , the sets  $S_{i}$  contain only one example. More precisely, we have for all  $i \in \{1, \dots, m\}$  the set  $S_{i} = \{\mathbf{z}_{i}\}$ . In this case, the bound becomes:
$$
\underset {h \sim \rho} {\mathbb {E}} \left[ \mathrm {R} _ {\mu} (h) - \hat {\mathrm {R}} _ {\mathcal {S}} (h) \right] \leq \frac {2 L}{m} \sum_ {i = 1} ^ {m} \mathrm {W} (\rho , \pi_ {i}) + 2 \sqrt {\ln \frac {m}{\delta}}.
$$
This bound is vacuous since the last term is incompressible, hence the need for a new technique detailed in Section 3.2 to deal with it.
The proof of Theorem 1 is presented in Appendix B.1. Appendices B.2 and B.3 introduce two proofs of Theorem 2. Theorem 3's proof is presented in Appendix B.4. Appendix B.5 provides the proof of Theorem 3.
$$
\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} \mathrm {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}
$$
Proof. For the sake of readability, we identify, for any  $i$ ,  $\pi_i$  and  $\pi_{i,\mathcal{S}}$ .
Step 1: Exploit the Kantorovich duality [Vil09, Remark 6.5]. First of all, note that for a  $L$ -Lipschitz loss function  $\ell : \mathcal{H} \times \mathcal{Z} \to [0,1]$ , we have
$$
\left| \left(\left| \mathcal {S} _ {i} \right| \mathrm {R} _ {\mu} \left(h _ {1}\right) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell \left(h _ {1}, \mathbf {z}\right)\right) - \left(\left| \mathcal {S} _ {i} \right| \mathrm {R} _ {\mu} \left(h _ {2}\right) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell \left(h _ {2}, \mathbf {z}\right)\right) \right| \leq 2 \left| \mathcal {S} _ {i} \right| L d \left(h _ {1}, h _ {2}\right). \tag {9}
$$
Indeed, we can deduce Equation (9) from Jensen inequality, the triangle inequality, and by definition that we have
$$
\begin{array}{l} \left| \left(| \mathcal {S} _ {i} | \mathrm {R} _ {\mu} (h _ {1}) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell (h _ {1}, \mathbf {z})\right) - \left(| \mathcal {S} _ {i} | \mathrm {R} _ {\mu} (h _ {2}) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell (h _ {2}, \mathbf {z})\right) \right| \\ = \left| \left(\sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \mathrm {R} _ {\mu} (h _ {1}) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell (h _ {1}, \mathbf {z})\right) - \left(\sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \mathrm {R} _ {\mu} (h _ {2}) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell (h _ {2}, \mathbf {z})\right) \right| \\ \leq \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \underset {\mathbf {z} ^ {\prime} \sim \mu} {\mathbb {E}} \left[ | \ell \left(h _ {1}, \mathbf {z} ^ {\prime}\right) - \ell \left(h _ {2}, \mathbf {z} ^ {\prime}\right) | + | \ell \left(h _ {2}, \mathbf {z}\right) - \ell \left(h _ {1}, \mathbf {z}\right) | \right] \\ \leq \underset {\mathbf {z} ^ {\prime} \sim \mu} {\mathbb {E}} \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} 2 L d (h _ {1}, h _ {2}) \\ = 2 \left| \mathcal {S} _ {i} \right| L d \left(h _ {1}, h _ {2}\right). \\ \end{array}
$$
We are now able to upper-bound  $\mathbb{E}_{h\sim \rho}[\mathbf{R}_\mu (h) - \hat{\mathbf{R}}_{\mathcal{S}}(h)]$ . Indeed, we have
$$
\begin{array}{l} \underset {h \sim \rho} {\mathbb {E}} \left[ \mathrm {R} _ {\mu} (h) - \hat {\mathrm {R}} _ {\mathcal {S}} (h) \right] = \frac {1}{m} \sum_ {i = 1} ^ {K} \underset {h \sim \rho} {\mathbb {E}} \left[ | \mathcal {S} _ {i} | \mathrm {R} _ {\mu} (h) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell (h, \mathbf {z}) \right] \\ \leq \sum_ {i = 1} ^ {K} \frac {2 \left| \mathcal {S} _ {i} \right| L}{m} W \left(\rho , \pi_ {i}\right) + \sum_ {i = 1} ^ {K} \underset {h \sim \pi_ {i}} {\mathbb {E}} \frac {1}{m} \left[ \left| \mathcal {S} _ {i} \right| R _ {\mu} (h) - \sum_ {\mathbf {z} \in S _ {i}} \ell (h, \mathbf {z}) \right], \tag {10} \\ \end{array}
$$
where the inequality comes from the Kantorovich-Rubinstein duality theorem.
Step 2: Define an adapted supermartingale. For any  $1 \leq i \leq K$ , we fix  $\lambda_i > 0$  and we provide an arbitrary order to the elements of  $S_i := \{\mathbf{z}_{i,1}, \dots, \mathbf{z}_{i,|S_i|}\}$ . Then we define for any  $h$ :
$$
M _ {| \mathcal {S} _ {i} |} (h) := | \mathcal {S} _ {i} | \mathrm {R} _ {\mu} (h) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell (h, \mathbf {z}) = \sum_ {j = 1} ^ {| \mathcal {S} _ {i} |} R _ {\mu} (h) - \ell (h, \mathbf {z} _ {i, j}).
$$
Remark that, because our data are i.i.d.,  $(M_{|S_i|})_{|S_i|\geq 1}$  is a martingale. We then exploit the technique [HG23a] to define a supermartingale. More precisely, we exploit a result from [BT08] cited in Lemma 1.3 of [HG23a] coupled with Lemma 2.2 of [HG23a] to ensure that the process
$$
S M _ {| \mathcal {S} _ {i} |} := \underset {h \sim \pi_ {i}} {\mathbb {E}} \left[ \exp \left(\lambda_ {i} M _ {| \mathcal {S} _ {i} |} (h) - \frac {\lambda_ {i} ^ {2}}{2} \left(\hat {V} _ {| \mathcal {S} _ {i} |} (h) + V _ {| \mathcal {S} _ {i} |} (h)\right)\right) \right],
$$
is a supermartingale, where  $\hat{V}_{|\mathcal{S}_i|}(h) = \sum_{j=1}^{|\mathcal{S}_i|} (\ell(h, \mathbf{z}_{i,j}) - R_\mu(h))^2$  and  $V_{|\mathcal{S}_i|}(h) = \mathbb{E}_{\mathcal{S}_i}\left[\hat{V}_{|\mathcal{S}_i|}(h)\right]$ .
Step 3. Combine steps 1 and 2. We restart from Equation (10) to exploit again the Kantorovich-Rubinstein duality.
$$
\begin{array}{l} \underset {h \sim \rho} {\mathbb {E}} \left[ \mathrm {R} _ {\mu} (h) - \hat {\mathrm {R}} _ {\mathcal {S}} (h) \right] = \frac {1}{m} \sum_ {i = 1} ^ {K} \underset {h \sim \rho} {\mathbb {E}} \left[ | \mathcal {S} _ {i} | \mathrm {R} _ {\mu} (h) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell (h, \mathbf {z}) \right] \\ \leq \sum_ {i = 1} ^ {K} \frac {2 | \mathcal {S} _ {i} | L}{m} W (\rho , \pi_ {i}) + \sum_ {i = 1} ^ {K} \frac {1}{m \lambda_ {i}} \lambda_ {i} \underset {h \sim \pi_ {i}} {\mathbb {E}} \left[ | \mathcal {S} _ {i} | R _ {\mu} (h) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell (h, \mathbf {z}) \right], \\ = \sum_ {i = 1} ^ {K} \frac {2 | \mathcal {S} _ {i} | L}{m} W (\rho , \pi_ {i}) + \sum_ {i = 1} ^ {K} \frac {1}{m \lambda_ {i}} \underset {h \sim \pi_ {i}} {\mathbb {E}} \left[ \lambda_ {i} M _ {| \mathcal {S} _ {i} |} \right], \\ \leq \sum_ {i = 1} ^ {K} \frac {2 | \mathcal {S} _ {i} | L}{m} W (\rho , \pi_ {i}) + \sum_ {i = 1} ^ {K} \frac {1}{m \lambda_ {i}} \ln \left(S M _ {| \mathcal {S} _ {i} |}\right) \\ + \frac {1}{m} \sum_ {i = 1} ^ {K} \mathbb {E} _ {h \sim \pi_ {i}} \left[ \frac {\lambda_ {i}}{2} \left(\hat {V} _ {| \mathcal {S} _ {i} |} (h) + V _ {| \mathcal {S} _ {i} |} (h)\right) \right]. \\ \end{array}
$$
The last line holds thanks to Jensen's inequality. We now apply Ville's inequality (see e.g., Section 1.2 of [HG23a]). We have for any  $i$ :
$$
\underset {\mathcal {S} _ {i} \sim \mu^ {| S _ {i} |}} {\mathbb {P}} \left(\forall | S _ {i} | \geq 1, S M _ {| S _ {i} |} \leq \frac {1}{\delta}\right) \geq 1 - \delta .
$$
Applying an union bound and authorising  $\lambda_{i}$  to be a function of  $|S_{i}|$  (thus the inequality does not hold for all  $|S_{i}|$  simultaneously) finally gives with probability at least  $1 - \delta$ , for all  $\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 \sum_ {i = 1} ^ {K} \frac {2 | \mathcal {S} _ {i} | L}{m} \mathrm {W} (\rho , \pi_ {i}) + \sum_ {i = 1} ^ {K} \frac {\ln \left(\frac {K}{\delta}\right)}{\lambda_ {i} m} + \frac {\lambda_ {i}}{2 m} \underset {h \sim \pi_ {i}} {\mathbb {E}} \left[ \hat {V} _ {| \mathcal {S} _ {i} |} (h) + V _ {| \mathcal {S} _ {i} |} (h) \right].
$$
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 priors). 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})$ :
Proof. For the sake of readability, we identify, for any  $i$ ,  $\pi_i$  and  $\pi_{i,S}$ .
$$
\left| \left(\left| \mathcal {S} _ {i} \right| \mathrm {R} _ {\mu} \left(h _ {1}\right) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell \left(h _ {1}, \mathbf {z}\right)\right) - \left(\left| \mathcal {S} _ {i} \right| \mathrm {R} _ {\mu} \left(h _ {2}\right) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell \left(h _ {2}, \mathbf {z}\right)\right) \right| \leq 2 \left| \mathcal {S} _ {i} \right| L d \left(h _ {1}, h _ {2}\right). \tag {11}
$$
Indeed, we can deduce Equation (11) from Jensen inequality, the triangle inequality, and by definition that we have
$$
\begin{array}{l} \left| \left(\left| \mathcal {S} _ {i} \right| \mathrm {R} _ {\mu} \left(h _ {1}\right) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell \left(h _ {1}, \mathbf {z}\right)\right) - \left(\left| \mathcal {S} _ {i} \right| \mathrm {R} _ {\mu} \left(h _ {2}\right) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell \left(h _ {2}, \mathbf {z}\right)\right) \right| \\ = \left| \left(\sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \mathrm {R} _ {\mu} (h _ {1}) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell (h _ {1}, \mathbf {z})\right) - \left(\sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \mathrm {R} _ {\mu} (h _ {2}) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell (h _ {2}, \mathbf {z})\right) \right| \\ \leq \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \underset {\mathbf {z} ^ {\prime} \sim \mu} {\mathbb {E}} \left[ | \ell \left(h _ {1}, \mathbf {z} ^ {\prime}\right) - \ell \left(h _ {2}, \mathbf {z} ^ {\prime}\right) | + | \ell \left(h _ {2}, \mathbf {z}\right) - \ell \left(h _ {1}, \mathbf {z}\right) | \right] \\ \leq \underset {\mathbf {z} ^ {\prime} \sim \mu} {\mathbb {E}} \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} 2 L d \left(h _ {1}, h _ {2}\right) \\ = 2 \left| S _ {i} \right| L d \left(h _ {1}, h _ {2}\right). \\ \end{array}
$$
We are now able to upper-bound  $\mathbb{E}_{h\sim \rho}[\mathrm{R}_{\mu}(h) - \hat{\mathrm{R}}_{\mathcal{S}}(h)]$ . Indeed, we have
$$
\begin{array}{l} \underset {h \sim \rho} {\mathbb {E}} \left[ \mathrm {R} _ {\mu} (h) - \hat {\mathrm {R}} _ {\mathcal {S}} (h) \right] = \frac {1}{m} \sum_ {i = 1} ^ {K} \underset {h \sim \rho} {\mathbb {E}} \left[ | \mathcal {S} _ {i} | \mathrm {R} _ {\mu} (h) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell (h, \mathbf {z}) \right] \\ \leq \sum_ {i = 1} ^ {K} \frac {2 \left| \mathcal {S} _ {i} \right| L}{m} W \left(\rho , \pi_ {i}\right) + \sum_ {i = 1} ^ {K} \underset {h \sim \pi_ {i}} {\mathbb {E}} \frac {1}{m} \left[ \left| \mathcal {S} _ {i} \right| \mathrm {R} _ {\mu} (h) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell (h, \mathbf {z}) \right], \tag {12} \\ \end{array}
$$
Remark that, because our data are i.i.d.,  $(M_{|\mathcal{S}_i|})_{|\mathcal{S}_i|\geq 1}$  is a martingale. We then exploit the technique [CWR23] to define a supermartingale. More precisely, we exploit [CWR23, Lemma A.2 and Lemma B.1] to ensure that the process
$$
S M _ {\left| \mathcal {S} _ {i} \right|} := \underset {h \sim \pi_ {i}} {\mathbb {E}} \left[ \exp \left(\lambda_ {i} M _ {\left| \mathcal {S} _ {i} \right|} (h) - \frac {\lambda_ {i} ^ {2}}{2} L _ {\left| \mathcal {S} _ {i} \right|} (h)\right) \right],
$$
is a supermartingale, where, because our data are i.i.d.,  $L_{|\mathcal{S}_i|}(h) = \mathbb{E}_{\mathcal{S}}\left[\sum_{j=1}^{|S_i|}\ell(h,\mathbf{z}_{i,j})^2\right] = |\mathcal{S}_i| \mathbb{E}_{z \sim \mu}[\ell(h,z)^2]$ .
Step 3. Combine steps 1 and 2. We restart from Equation (12) to exploit the Kantorovich-Rubinstein duality again.
$$
\begin{array}{l} \underset {h \sim \rho} {\mathbb {E}} \left[ \mathrm {R} _ {\mu} (h) - \hat {\mathrm {R}} _ {\mathcal {S}} (h) \right] = \frac {1}{m} \sum_ {i = 1} ^ {K} \underset {h \sim \rho} {\mathbb {E}} \left[ | \mathcal {S} _ {i} | \mathrm {R} _ {\mu} (h) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell (h, \mathbf {z}) \right] \\ \leq \sum_ {i = 1} ^ {K} \frac {2 | \mathcal {S} _ {i} | L}{m} W (\rho , \pi_ {i}) + \sum_ {i = 1} ^ {K} \frac {1}{m \lambda_ {i}} \lambda_ {i} \underset {h \sim \pi_ {i}} {\mathbb {E}} \left[ | \mathcal {S} _ {i} | R _ {\mu} (h) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell (h, \mathbf {z}) \right], \\ = \sum_ {i = 1} ^ {K} \frac {2 | \mathcal {S} _ {i} | L}{m} W (\rho , \pi_ {i}) + \sum_ {i = 1} ^ {K} \frac {1}{m \lambda_ {i}} \underset {h \sim \pi_ {i}} {\mathbb {E}} \left[ \lambda_ {i} M _ {| \mathcal {S} _ {i} |} \right], \\ \leq \sum_ {i = 1} ^ {K} \frac {2 | \mathcal {S} _ {i} | L}{m} W (\rho , \pi_ {i}) + \sum_ {i = 1} ^ {K} \frac {1}{m \lambda_ {i}} \ln \left(S M _ {| \mathcal {S} _ {i} |}\right) \\ + \frac {1}{m} \sum_ {i = 1} ^ {K} \mathbb {E} _ {h \sim \pi_ {i}} \left[ \frac {\lambda_ {i}}{2} L _ {| \mathcal {S} _ {i} |} (h) \right]. \\ \end{array}
$$
The last line holds thanks to Jensen's inequality. We now apply Ville's inequality (see e.g., section 1.2 of [HG23a]). We have for any  $i$ :
$$
\underset {h \sim \rho} {\mathbb {E}} \left[ \mathrm {R} _ {\mu} (h) - \hat {\mathrm {R}} _ {\mathcal {S}} (h) \right] \leq \sum_ {i = 1} ^ {K} \frac {2 | \mathcal {S} _ {i} | L}{m} \mathrm {W} (\rho , \pi_ {i}) + \sum_ {i = 1} ^ {K} \frac {\ln \left(\frac {K}{\delta}\right)}{\lambda_ {i} m} + \frac {\lambda_ {i}}{2 m} \underset {h \sim \pi_ {i}} {\mathbb {E}} [ L | \mathcal {S} _ {i} | (h) ].
$$
Finally, using the assumption  $\mathbb{E}_{h\sim \pi_i}\mathbb{E}_{z\sim \mu}[\ell (h,z)^2 ]\leq 1$  gives, with probability at least  $1 - \delta$ , for all  $\rho \in \mathcal{M}(\mathcal{H})$ :
$$
\underset {h \sim \rho} {\mathbb {E}} \left[ \mathbf {R} _ {\mu} (h) - \hat {\mathbf {R}} _ {\mathcal {S}} (h) \right] \leq \sum_ {i = 1} ^ {K} \frac {2 | \mathcal {S} _ {i} | L}{m} \mathrm {W} (\rho , \pi_ {i}) + \sum_ {i = 1} ^ {K} \frac {\ln \left(\frac {K}{\delta}\right)}{\lambda_ {i} m} + \frac {\lambda_ {i} [ \mathcal {S} _ {i} ]}{2 m}.
$$
Taking for each  $i$ ,  $\lambda_{i} = \sqrt{\frac{2\ln(K / \delta)}{|S_{i}|}}$  concludes the proof.
B.3 Alternative proof of Theorem 2
We state here a slightly tighter version of Theorem 2 for bounded losses, which relies on an application of McDiarmid's inequality instead of supermartingale techniques. This is useful for the numerical evaluations of our bound.
Theorem 6. We assume our loss  $\ell$  to be in [0,1] and  $L$ -Lipschitz. 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 {| \mathcal {S} _ {i} | \ln \frac {K}{\delta}}{2 m ^ {2}}}
$$
where  $\pi_i$  does not depend on  $S_i$ .
First of all, note that for a  $L$ -Lipschitz loss function  $\ell : \mathcal{H} \times \mathcal{Z} \to [0,1]$ , we have
$$
\left| \left(| \mathcal {S} _ {i} | \mathrm {R} _ {\mu} (h _ {1}) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell (h _ {1}, \mathbf {z})\right) - \left(| \mathcal {S} _ {i} | \mathrm {R} _ {\mu} (h _ {2}) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell (h _ {2}, \mathbf {z})\right) \right| \leq 2 | \mathcal {S} _ {i} | L d (h _ {1}, h _ {2}). \tag {13}
$$
Indeed, we can deduce Equation (13) from Jensen's inequality, the triangle inequality, and by definition that we have
We are now able to upper-bound  $\mathbb{E}_{h\sim \rho}[\mathbf{R}_\mu (h) - \hat{\mathbf{R}}_{\mathcal{S}}(h)]$  . Indeed, we have
$$
\begin{array}{l} \underset {h \sim \rho} {\mathbb {E}} \left[ \mathrm {R} _ {\mu} (h) - \hat {\mathrm {R}} _ {\mathcal {S}} (h) \right] = \frac {1}{m} \sum_ {i = 1} ^ {K} \underset {h \sim \rho} {\mathbb {E}} \left[ | \mathcal {S} _ {i} | \mathrm {R} _ {\mu} (h) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell (h, \mathbf {z}) \right] \\ \leq \sum_ {i = 1} ^ {K} \frac {2 \left| \mathcal {S} _ {i} \right| L}{m} W \left(\rho , \pi_ {i}\right) + \sum_ {i = 1} ^ {K} \underset {h \sim \pi_ {i}} {\mathbb {E}} \frac {1}{m} \left[ \left| \mathcal {S} _ {i} \right| \mathrm {R} _ {\mu} (h) - \sum_ {\mathbf {z} \in \mathcal {S} _ {i}} \ell (h, \mathbf {z}) \right], \tag {14} \\ \end{array}
$$