Learning via Wasserstein-Based High Probability Generalisation Bounds

Summary

This paper focuses on studying PAC-Bayes generalization bounds based on the Wasserstein distance. It introduces novel high-probability bounds applicable to both batch learning (i.i.d. setting) and online learning (non-i.i.d. setting). Unlike previous PAC-Bayes bounds, their bounds are not limited to bounded or sub-Gaussian losses. Furthermore, leveraging their theoretical findings, the authors propose a new SRM training strategy and provide empirical evidence to validate its effectiveness.

Strengths

Compared to previous results, the Wasserstein-distance PAC-Bayes bounds presented in this paper require fewer or weaker assumptions. Additionally, unlike previous PAC-Bayes bounds that typically invoke Gaussian distributions, their bounds accommodate discrete Dirac distributions for both the prior and posterior distributions.

Weaknesses

Regarding the Wasserstein regularization part, it might be necessary to empirically compare it with similar styles of regularization, such as weight decay $||w||$ and the 'distance to initialization' $||w-w_0||$. Although I believe that a data-dependent prior will provide a tighter bound, it is uncertain whether the data-dependent prior can significantly outperform a data-independent prior (e.g., $0, w_0$) considering the significantly increased computational time required.

Questions

Major concerns: 1. There might be a small error in Theorem 2: in the proof of Theorem 2 (Section B.2), you let $\lambda_i=\sqrt{\frac{\ln{(K/\delta)}}{|S_i|}}$ in the last step (i.e. Line 636). However, this choice does not yield the desired $\sum_{i=1}^K\sqrt{\frac{|S_i|\ln{(K/\delta)}}{m^2}}$ in your bound. If I understand correctly, you could consider setting $\lambda_i=\sqrt{\frac{2\ln{(K/\delta)}}{|S_i|}}$, which would result in the last term in the bound becoming $\sum_{i=1}^K\sqrt{\frac{2|S_i|\ln{(K/\delta)}}{m^2}}$. This may be the smallest achievable second term in Section B.2. 2. In Line 206-207, you argue that Eq. (3) contains a trade-off parameter $\lambda$, while Eq. (2) does not. However, it is worth noting that a similar parameter, $\epsilon$, is also involved in your practical training objective, namely Eq. (5). Therefore, the argument regarding the presence of $\lambda$ may not be entirely fair. Perhaps removing Line 206-207 could be considered? 3. Have you considered incorporating the empirical variance term into your algorithm? Since you emphasize its reliance on the prior rather than the posterior, I would appreciate it if you could demonstrate some practical implications. 4. Under the same conditions, such as bounded loss or even zero-one loss, is it possible to compare your Wasserstein-based PAC-Bayes bounds with the standard KL-based PAC-Bayes bounds, especially those that also include variance terms as mentioned in [1,2,3]? Minor comments: 1. Line 131: there is a missing period after $S_K$. 2. RHS of the equation on Line 633: $L{|S_i|}\Longrightarrow L_{|S_i|}$. 3. RHS of the equation on Line 635: $\lambda [S_i|]\Longrightarrow \lambda |S_i|$. [1] Seldin, Yevgeny, et al. "PAC-Bayesian inequalities for martingales." IEEE Transactions on Information Theory 58.12 (2012): 7086-7093. [2] Tolstikhin, Ilya O., and Yevgeny Seldin. "PAC-Bayes-empirical-Bernstein inequality." NeurIPS 2013. [3] Mhammedi, Zakaria, Peter Grünwald, and Benjamin Guedj. "PAC-Bayes un-expected Bernstein inequality." NeurIPS 2019.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

4 excellent

Limitations

There is no potential negative societal impact.

Summary

This paper introduces Wasserstein-based PAC-Bayes bounds both for offline (batch learning) and online learning. These bounds are obtained with a clever use of the Kantorovich-Rubinstein duality together with a control of some additional terms appearing in the proof that only depend on the prior. To use the Kantorovich-Rubinstein duality they require the loss to be Lipschitz and to control the additional terms they require the loss to have a bounded second moment, following the techniques from [HG23a]. The extension from offline to online learning comes almost "for free" in the analysis due to the super-martingale strategies developed in [HG23a]. This way, the presented bounds have a set of weak assumptions (Lipschitzess + bounded second moment) which makes them amenable to losses with potentially heavy tails. Another aspect of the bounds is that they allow data-dependent priors. For this, they separate the training dataset $\mathcal{S}$ into $K$ subsets $\mathcal{S}\_i$and dedicate $K$ data-dependent priors $\pi\_i$ that depend on $\mathcal{S} \setminus \mathcal{S}\_i$. This way, they have an average of Wasserstein distances of the posterior and each of these priors weighted accordingly. This way, for a large number of samples the posterior and the priors start to align helping the bound to become tighter. Finally, they employ their bounds to derive regularization terms that are employed to design algorithms both in offline and online learning. They showcase how the usage of these algorithms does not decrease the performance of linear and deep models and that, in some cases, they help these models achieve better performance.

Strengths

The strengths of this paper are multiple, but the most important are: - The presented theorems are general (Lipschitz and second-moment assumptions only). - The proofs of the presented theorems are exposed in a simple way, even though the final results are useful. - The theorems allow for deterministic predictors, which is one of the biggest drawbacks of PAC-Bayes bounds. - The results are validated empirically.

Weaknesses

There are not many major issues with the paper. My biggest concern is the following: - In the section "Batch algorithm" starting in line 275, you use the equation (5) to guide the training of the posterior, which includes the distance $\lVert \textbf{w} - \textbf{w}_i \rVert_2$ that comes from the Wasserstein distance. However, the employed loss is only Lipschitz with respect to the outputs. How do you justify then the usage of the $\ell$_2 norm of the weights, as it could be that the loss is not Lipschitz with respect to the weights? The same happens later in equations (6) and (7). Other smaller issues are the following: - The experimental results don't seem too strong compared with previous methods such as the one presented in [PORPH+21]. Although they use probabilistic networks, it would have been interesting to observe how the two methods compare to each other with the same network architecture. For instance, with a FCN and Gaussian priors, they show population risk certificates of ~0.02, while the presented here are of ~0.09 and ~0.12. - Some of the writing can be clarified. In particular, I am referring to sentences of the type: "for all $i$, $\mathcal{S}$,$\pi\_i(S,\cdot)$ is $\mathcal{F}\_{i-1}$ measurable". This is not very clear, I would recommend changing all these sentences for something like "the distribution $\pi\_i(S,\cdot)$ is $\mathcal{F}\_{i-1}$ measurable for all $i$ and all $\mathcal{S}$". This small change can improve readability. - I missed some comments on Proposition 5 of [AEM22], which is a Seeger-type bound. How do the presented bounds compare to that? - I am not sure the implementation of OGD is correct, given that it ends up having a larger empirical risk than the regularized version. Maybe I am missing something here, or maybe this showcased some limitations of OGD I was not aware of. I also mentioned it explicitly as a question below. Some typos: - In line 29, it should be $R\_{\mu}(h) - \hat{R}\_{\mathcal{S}}(h)$. - In lines 192 and 193, you forgot to include the hat in $\hat{R}$. - In line 354, it should be $\mathfrak{C}\_{\mathcal{S}}$.

Questions

* How do you justify the usage of the $\ell_2$ norm on the weights in equations (5), (6), and (7) if the employed loss is only Lipschitz with respect to the outputs of the hypotheses $h_{\textbf{w}}$ but not necessarily with respect to the weights themselves? * What is the intuition for the fact that Algorithm 2 has a smaller empirical risk than OGD? It seems weird given the fact that Algorithm 2 is essentially a regularized version of OGD, which is an algorithm for ERM.

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Some of the limitations are mentioned throughout the text, yes.

Summary

This work propose PAC-Bayesian learning with the KL divergence replaced by the Wasserstein distance on the metric space of hypotheses $(\mathcal{H}, d)$. Denoting the data distribution by $\mathcal{D}$, the authors study the following learning problems: 1. Batch learning. Under the assumption that the loss function has bounded 2nd moment with respect to $\mathcal{D}$ and priors over $\mathcal{H}$, the authors derive a generalization bound in terms of $LW$, where $W$ is the Wasserstein distance between posterior and prior distributions over $\mathcal{H}$ and $L$ is the Lipschitz constant of the loss function (with respect to $d$). 2. Online learning. Under the assumption that the loss function has bounded 2nd _conditional_ moment with respect to $\mathcal{D}$ and priors over $\mathcal{H}$, the authors derive a similar generalization bound as above, where the population risk is defined by the sum of conditional expections of the loss, given a filtration adapted to the sequence of data. Both bounds can be applied to heavy-tailed loss functions with bounded 2nd moment. And they can be used to derive new learning problems regularized by $LW$. Experimental results with a linear model and an NN model show that these algorithms have better test risks than ERM and OGD on several dataset, among which online learning with the NN model shows the most improvement.

Strengths

- This is a strong result in the field of PAC-Bayes learning. Replacing the KL divergence by the Wasserstein distance is logical when the geometry of the underlying metric space is in consideration. - The bound is just as nice as the KL divergence version of the PAC-Bayes bound. And the ability to convert the bound into the algorithm is powerful. - The experimental results show that the bound can be used to improve models' generalization.

Weaknesses

- It seems that the batch size substantially affects the test risks. For example, learning with $\epsilon = 1/m$ on MNIST and TICTACTOE with the linear model, and TICTACTOE on the NN model. Are there any specific reasons that make these models underperform on these datasets? - Alg 2. seemingly performs as well as OGD with the linear model; on the other hand, the former does noticeably better with the NN model. The authors might want to discuss on why the NN model performs well in this case. - A minor comment on the organization: In Section 4.1: Batch algorithm, the authors set $L=1/2$ and proceed to define a specific loss that has the said Lipschitz constant, which is quite confusing to read. I think it is more logical to introduce the loss function first.

Questions

- Though not required, but I think high level ideas of the proofs of Theorem 1 and 3 would be a nice addition. - In Section 4.1: Batch algorithm, has the definition of $d(h,h_i)$ been defined before? I guess the authors have make a convention that $d(h,h_i)=d(w,w_i)$. - Can the authors performs the experiment against ERM with standard regularization as well?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

The authors might want to add some limitation of the bounds. For example, it might be difficult to compute the Wasserstein distance over nonparametric models, such as kNN.

Summary

The paper builds on recent advances in PAC-Bayesian learning and derives new Wasserstein distance-based generalization bounds. Besides batch learning for iid data, a first set of results are derived for online learning (with non iid data). The authors also provide tight bounds for the case of heavy tailed loss with bounded second order moments. Through experimental results the authors argue that these bounds can inspire new learning algorithms that achieve comparable (and in some cases better) empirical performance than classical benchmarks (ERM for batch and OGD for online setting).

Strengths

The introduction and main results are well placed in the context of recent literature in this area. Although this reviewer has not worked on the topic of PAC-Bayes learning, it was possible to understand the authors’ contribution relative to work by Haddouche and Guedj, Catoni and co-authors, and recent paper of Amit et al from NeurIPS 2022. One would have expected more clarity on the use of “supermatingale toolbox” in the main text of the paper, but I suppose this can be addressed using minor revisions of the paper.

Weaknesses

However, there are several points that impede the broader understanding of the work. Firstly, the role and possible limitations imposed by L-Lipschitz assumption should be clarified. In the discussion of Theorem 1, the authors mention that the variance terms are considered with respect to the prior distributions $\pi_{i,\mathcal{S}}$ and not $\rho$. Can you clearly explain why this is a strength of the result (if so)? In comment after equation (2), I agree that Wasserstein distance as regularizer…allows the user of multiple priors. But the issue of selecting K (as opposed to choice of $\lambda$ in classical algorithm). The authors do provide some discussion of role of $K$ but the nuance of its role in regularization is still not very clear (besides, the discussion on tradeoff needs a bit more clarity). Appreciate the interest in online setting, but the contribution/merits of algorithm based on equation (4) are not very clear. Firstly, what is the conceptual novelty here in comparison to online counterpart to equation (3)? Secondly, the particularization to Dirac masses seems like a restriction. The algorithmic aspects need a bit more clarity as well. The computational results seem promising and an assertive discussion on comparison with ERM and OGD benchmarks is needed.

Questions

Please see my comments/suggestion in the Weaknesses section.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

4 excellent

Limitations

Would have liked to see a more robust discussion on the limitations of the proposed approach. A clear explanation about when and to what extent geometry of data space can improve the results obtained through Wasserstein distance based approach would be much appreciated. Also comment on future work to refine results on online setting.

Answer to rebuttal

Thank you for your rebuttal. My major concern is now addressed with the new Lemma, thanks. This only begs the question on how to choose the hyper-parameter $\varepsilon$ with respect to that Lipschitz constant or, in other words, how much the bound is used for regularization. But you already discuss that in the paragraph after (5), so a small further discussion will do. Also, all my minor concerns and questions are clarified. (P.S.: I missed that in [AEM22, Prop. 5] they needed that $m \cdot \Delta\_S^{\mathrm{kl}} \in \mathcal{F}\_S$ for some reason. Then, I thought one could choose $\mathcal{F}\_S$ to be the set of Lipschitz functions. Sorry for that.).

Thank you

I thank the authors for addressing all of my concerns. There is an open problem of finding a suitable batch size left which is interesting in its own right.