Universal generalization guarantees for Wasserstein distributionally robust models

Thank you for your reading and positive feedback. Your several remarks helped us improving the paper and we addressed each of them in our revision. Please find below our answers to the important points from your review.

Questions

Why not submit to JMLR? The paper is very rigorous and rather long and technical for an ML conference? You also examine the problem in good detail.

Yes we tried to propose a thorough study. We acknowledge that certain parts of our work are technical. We have made an effort to structure them rigorously and have placed them in the appendices. In the main, we have instead focused on the context, the results, and the key ideas that underpin their derivation. We believe that this material will be valuable to a significant portion of the community, whether they work on theoretical and practical aspects.

Could you provide a simple example in Theorem 3.2, where the optimal coupling is known under (say) Gaussianity assumptions?

This is not evident even for simple examples: the optimal coupling depends on $Q$ which can be any distribution satisfying $W_c^\tau(P,Q) \leq \rho$. In order to gain more interpretable results, approximation results might be established to quantify how much the right hand side is close to the exact risk. Considering existing works on this aspect, see e.g. [1], we believe this would require more structure on $c$. This is a relevant and non trivial question.

I'm a bit confused. What does $\operatorname{argmax}_\Xi f$ mean in (5) a sup norm or something?

This is the set of maximizers of $f$ on $\Xi$, the points attaining the maximum of $f$ on $\Xi$, the ``argmax''. This now appear in the notations.

Why is $\operatorname{min} c(\xi,\zeta) ...$ measurable?

As you say, this can be given by the measurable Maximum Theorem. In our proofs, we did not discuss measurability of every function because it was often satisfied through stronger notions such as continuity or semicontinuity (for instance to prove Lemma 4.1 in appendix D.1).

Each result assumes that the (difficult to interpret) $\rho\_{\text{crit}}$ is "large enough" Can you please provide a general set of conditions ensuring that $\rho\_{\text{crit}}$ can be bounded away from 0.

This is quite abstract, indeed. We interpret it as follows: when $P$ is supported on $\Xi$, $\rho_{\text{crit}} > 0$ if and only if $\mathcal{F}$ contains no constant function. We have a proposition in appendix establishing this; we have added a pointer to it in the main, for more clarity.

In theorem 3.2, why is $\pi^{P,Q} \ll \pi_0$?

This is due to the definition of KL divergence. This is now mentioned as a footnote.

Is it fair to compare, verbally, our results to those of Fournier et al [...] ?

Fournier and Guillin concentration result can indeed be improved when some structure can be leveraged. In this paper, we consider instead the general case, with no restriction on specific distributions; this is the setting of the seminal work of Mohajerin Esfahani and Kuhn (2018), relying on Fournier and Guillin, for the statistical results. In our work, our take is that, in this general setting, we can still gain in sample complexity, not using the structure of the distributions, but rather by leveraging the structure of the WDRO optimization problem itself.

[1] Waiss Azizian, Franck Iutzeler, and Jerome Malick. Regularization for wasserstein distributionally robust optimization. ESAIM: Control, Optimisation and Calculus of Variations, 29:33, 2023b.

We now continue with your minor points. We have addressed all of them in the revision. Please find our answer to the most important ones below.

Minor points

The definition of Wasserstein distance circa (1) is incorrect, $\Xi$ must be Polish, and $c$ must be a power $p \geq 1$ of a metric topologizing $\Xi$ ; what you write is just some transport problem [...]

On the terminology ''Wasserstein distance'' (1): Your are right, writing "Wasserstein distance" in our case is a slight abuse of terminology. WDRO may be studied for general costs (see e.g. [2]), under our Assumption 2.1. Note that we also require $c(\xi,\zeta)$ to be zero if and only if $\xi = \zeta$. As you say, this setting does not imply $W_c$ to be a metric. We now use the terminology "optimal transport cost" as in [2] to avoid any confusion.

Assumption 1 vs. Line 145 You say that is just a measurable space, then later you say its a compact metric space [...]

The section notation was indeed a bit general for the main text. We modified it to make it more simple and we wrote a general notation part in the appendix for the proofs. In the main text we now define the spaces $\Xi$, $\mathcal{F}$ and the cost $c$ right from the beginning in the notation part. We kept their specificities (such as compactness or continuity) in Assumption 2.1 to facilitate the comparison of our setting with the literature, and make the presentation as much transparent as possible.

Line 69: "This theoretical feature is specific to WDRO and highlights its potential to give more resilient models." This can be much less hand-wavy. Please explain more precisely/mathematically.

This relates to the previous sentence, which already describe the inequality (3). We reorganized both sentences to make it clearer

Line 176: Why jointly Lipschitz? [...]

This is a good remark, there was an omission here. Joint Lipschitz continuity ensures $\mathcal{I}\_{\mathcal{F}}$ is finite; this is now added in the revision. $\mathcal{I}\_{\mathcal{F}} < \infty$ can also be satisfied when assuming $f(\cdot,\xi)$ $L$-Lipschitz for all $\xi \in \Xi$. For simplicity we assume joint Lipschitz continuity since we can not think of a relevant situation where $\Xi$, $\Theta$ would be compact but $f$ not jointly Lipschitz continuous.

Line 223: There are many more references of the use of this type of metric

This is indeed a natural metric, which exists in many other contexts than WDRO.

Assumption 1: Why call (2) jointly continuous [...]

Yes, but this is continuity on the product space $\Xi \times \Xi$. We use the terminology joint continuity to avoid any confusion with continuity with respect to each variable.

[2] Jose Blanchet and Karthyek Murthy. Quantifying distributional model risk via optimal transport. Mathematics of Operations Research, 44(2):565–600, 2019.

Thank you for your positive feedback. Please find our response to your several questions below.

In Section 3.2, the authors discussed how their results on generalization guarantees apply to linear regression and logistic regression. However, more complicated models such as neural networks with ReLU or other smooth activation functions (e.g. GELU) are not discussed.}

We underline that more complicated models, e.g. neural networks with ReLU, are covered by our general analysis, in contrast to the previous ones.

This being said, we agree that obtaining theoretical or empirical insights into the constants $\lambda_{\text{low}}$ and $\rho_{\text{crit}}$ in deep learning contexts would be interesting. This question remains open and deserves a dedicated study.

---

The theoretical results require a lower bound on $n$, while Theorem 3.4 of Azizian et al. (2023a) applies to all $n \geq 1$. The implications of this requirement should be discussed.

This is a good technical point to raise, thank you. Theorems 3.1 and 3.4 of Azizian et al. (2023a) also (implicitly) require such lower bounds. Let us explain why.

Theorems 3.1 and 3.4 from Azizian et al. 2023 require to choose $\rho$ as $\frac{\alpha}{\sqrt{n}}\leq \rho \leq \frac{\rho_{\text{crit}}}{2} - \frac{\beta}{\sqrt{n}}$ (here we replaced their big O notation by some constants $\alpha >0$ and $\beta > 0$). However, this range can be empty with no further condition on $n$. Only for $n$ high enough $\rho$ can be chosen in the interval $[\frac{\alpha}{\sqrt{n}}, \frac{\rho_{\text{crit}}}{2} - \frac{\beta}{\sqrt{n}}]$. This gives the condition $\frac{\alpha}{\sqrt{n}} \leq \frac{\rho_{\text{crit}}}{2} - \frac{\beta}{\sqrt{n}}$ and then the lower bound $n \geq 4 (\alpha + \beta)^2/\rho_{\text{crit}}^2$. In this case, as we did, we may also remove the upper bound on $\rho$ by monotonicity of the robust loss with respect to $\rho$. We added a comments on this point in appendix.

---

What are the practical implications of the generalization guarantees compared to Azizian et al. (2023a)? Can you provide some numerical results analogous to Appendix H of Azizian et al. (2023a)?

It would indeed be great to have numerical illustrations of generalization properties in the general setting we consider. This is however a difficult point that would require a complete and rigorous work. This is out of scope of this theoretical paper. Among the many challenges, efficient practical procedures to compute WDRO models in deep learning are still missing. This is how we conclude the paper insisting on future work.

Thank you very much for your positive feedback, all your comments/remarks, and your two main questions above. It helps us improving our presentation, and especially on the position w.r.t Azizian et al. 2023a, which is the closest work.

In the revision, we now provide a discussion, just below our main assumptions (Assumption 2.1) where we compare our setting to Azizian et al. 2023a in details. This complements the key differences, already mentioned in the related work section (lines 120 to 130).

Let us recall here our improvements, compared to Azizian et al (2023a):

• The setting of Azizian et al (2023a) restricts to smooth functions $f \in \mathcal{F}$ (twice differentiable with uniformly bounded derivatives on a convex sample space). We only require the $f \in \mathcal{F}$ to be continuous on a metric space. In addition to nonsmooth functions, this allows us to consider distributions on sample spaces with discrete and continuous variables (as for e.g. classification tasks).

• Their proof require to take $c$ as the squared norm and the reference distribution $\pi_0(\cdot|\xi)$ as a Gaussian distribution. We consider instead general costs $c$, continuous with respect to a distance on $\Xi$ and an arbitrary reference probability distribution.

For instance, this setting is captured by us but not by Azizian et al. (2023a):

(i) The sample space $\Xi = B(0,R) \times \{0,1\}$ where $R > 0$

(ii) The loss family $\mathcal{F} = \\{ f(\theta, \cdot) \ : \ \theta \in \Theta \\}$ with the cross entropy loss $f(\theta, x,y) = - y \log(h(\theta,x)) - (1 - y) \log(h(\theta,x))$ where $h(\theta, \cdot)$ is a feedforward network with RELU activations and $\Theta$ is compact.

(iii) The cost function: $c((x,y), (x',y')) = \|x-x'\|_{p}^{q} + \kappa \mathbb{1}\_{y \neq y'}$

• Moreover, in their proof, to overcome nonsmoothness of WDRO (which poses concerns for applying concentration results), they require two technical assumptions: a compactness condition (1) and growth conditions around maximizers (2), this is their Assumption 5:

(1) For any $R > 0$, there exists $\Delta > 0$ such that $\forall f \in \mathcal{F}, \; \forall \zeta \in \Xi, \; d\left(\zeta, \operatorname{argmax} f\right) \geq R \implies f(\zeta) - \max f \leq -\Delta.$

(2) There exist $\mu > 0$ and $L > 0$ such that, for all $f \in \mathcal{F}$, $\xi \in \Xi$ and $\xi^*$ a projection of $\xi$ on $\operatorname{argmax} f$, $f(\xi^*) \geq f(\xi) + \frac{\mu}{2} \|\xi - \xi^*\|^2 - \frac{L}{6} \|\xi - \xi^*\|^3.$

We do not rely on these conditions. They are rather strong and difficult to verify since maximizers over the sample space are hard to control in general and both depend on the sample space and the function class geometries. In particular, (1) requires $\mathcal{F}$ to be compact with respect to a distance defined by summing the sup norm and the Hausdorff distance between maximizers sets. Equivalently, this can be seen as the continuity of $f \mapsto \operatorname{argmax} f$ on $\mathcal{F}$ (see our Proposition F.4 in appendix), which is hard to verify.

The main technical difficulties in our proof was thus to get rid Assumption 5 from Azizian et al 2023a and to deal with the nonsmooth aspect of the WDRO objective. To this purpose, we simplified the proof and use nice nonsmooth analysis tools. We highlight this aspect in the sketch of proof, section 4.2 ''Definition of the lower bound" where we present a maximal radius function.

About your smaller question: Indeed, the assumption is vacuous in this case and we may take $\omega = 1$. We decided to keep the constant $\omega$ in the results in order to highlight the dependence of $\lambda_{\text{low}}$ on the hypothesis domain. To make our statements more precise, we added the condition $\omega \geq 1$ in Assumption 3.1.1.