On Generalization and Regularization via Wasserstein Distributionally Robust Optimization

Stochastic optimization problems of the form

$\displaystyle\inf_{\bm{\beta}\in\mathcal{D}}\rho(Y\cdot\bm{\beta}^{\top}\mathbf{X})$

(1)

naturally arise in many machine learning (ML) and operations research (OR) applications, where $Y$ denotes a binary random variable, taking values from $\{-1,1\}$, and $\mathbf{X}$ is a random vector in ${\mathbb{R}}^{n}$. The function $f_{\bm{\beta}}(\mathbf{X})=\bm{\beta}^{\top}\mathbf{X}$, parameterized by the decision variable $\bm{\beta}$ $\in\mathcal{D}\subseteq\mathbb{R}^{n}$, can generally be interpreted as an affine decision rule. In ML, the binary random variable $Y$ often represents outcomes in classification problems. When $Y$ is constant, the formulation (1) applies to a broad range of regression problems in ML and risk minimization problems in OR. The function $\rho$ represents a measure of risk, mapping a random variable to a real number that quantifies its risk. The notion of risk in this paper is broadly defined as the undesirability of a random variable $Z$; that is, a larger value of $\rho(Z)$ indicates that $Z$ is less preferable.

In this paper, we study the Wasserstein distributionally robust counterpart of (1), specifically:

$\displaystyle\inf_{\bm{\beta}\in\mathcal{D}}\sup_{F\in{\cal B}_{p}(F_{0},\varepsilon)}\rho^{F}(Y\cdot\bm{\beta}^{\top}\mathbf{X}),$

(2)

where ${\cal B}_{p}(F_{0},\varepsilon)$ denotes a type-$p$ Wasserstein ball of radius $\varepsilon$ centred at a reference distribution $F_{0}$ (c.f. Kuhn et al. (2019)). The superscript $F$ in $\rho^{F}$ indicates that the measure depends on a joint distribution $F$ of $(Y,\mathbf{X})$. This problem has been extensively studied and applied in various fields when $\rho^{F}$ takes the form of an expected value function, i.e., $\rho^{F}(Z)=\mathbb{E}^{F}[\ell(Z)]$ for some $\ell:{\mathbb{R}}\rightarrow{\mathbb{R}}$ (see e.g., Kuhn et al. (2019); Shafieezadeh-Abadeh et al. (2019); Gao et al. (2017); Gao (2022)). This setup reveals two key findings that underscore the potential strength of the Wasserstein distributionally robust optimization (DRO) model in achieving strong out-of-sample performance. First, the model enjoys generalization bounds under mild assumptions (Shafieezadeh-Abadeh et al. (2019)). More notably, Gao (2022) shows that when the data-generating distribution satisfies transportation inequalities, $O(N^{-1/2})$-generalization bounds for the type-$p$ Wasserstein ball ${\cal B}_{p}(F_{0},\varepsilon)$ with $p\in[1,2]$ can be established, free from the curse of dimensionality. Second, the model has an exact regularization equivalent to the nominal problem (1) with an added regularizer on the decision variable $\bm{\beta}$, when the Wasserstein ball is of type-1 and the loss function $\ell$ is Lipschitz continuous (Shafieezadeh-Abadeh et al. (2019)), as well as when the ball is of type-2 and the loss function $\ell$ is quadratic (Blanchet et al. (2019)). This connection to the classical regularization scheme, prevalent in machine learning, offers a compelling interpretation of the DRO model and has spurred considerable interest in its application within ML and OR (see e.g., Blanchet and Kang (2021), Blanchet et al. (2019), Chen and Paschalidis (2018), Gao et al. (2017), Gao et al. (2022)).

It remains largely unclear, however, whether the generalization bounds established for the specific case of the expected value function can be achieved in the broader setting of (2) – namely, for a general measure $\rho$ and any type-$p$ Wasserstein ball ${\cal B}_{p}(F_{0},\varepsilon)$ with $p\in[1,\infty]$ – without suffering from the curse of dimensionality. In particular, the result of Gao (2022), aside from the limitations imposed by the transportation inequality assumption on the underlying distribution and the restriction to $p\in[1,2]$, fundamentally relies on the unbiased property of the sample average – a property that does not extend to risk measures such as Conditional Value-at-Risk (CVaR) (Rockafellar and Uryasev (2002)). It is also worth noting that, even in the literature on ML, the question of how to obtain generalization bounds, possibly through the classical regularization scheme, for the nominal problem (1) has largely been left open (c.f. Shalev-Shwartz and Ben-David (2014)). In this paper, we shed new light on the efficacy of Wasserstein DRO in broad out-of-sample tests by showing that the Wasserstein DRO model (2) can circumvent the curse of dimensionality for virtually any decision criterion under affine decision rules. As a crucial insight, we uncover that the Wasserstein DRO remarkably coincides with the recent max-sliced Wasserstein DRO (Olea et al. (2022)) for any decision criterion, even though the ball defined in the former is a far less conservative choice – specifically, it is smaller and tighter than the one defined in the latter. Additionally, we extend our results to general decision rules under the light-tail assumption by leveraging the measure concentration property of a Wasserstein ball.

Another focus of this work is to investigate whether the key equivalence of Wasserstein DRO to classical regularization in ML can be extended to a more general setting in (2). To answer this, we first pay particular attention to the setting where the Wasserstein ball ${\cal B}_{p}(F_{0},\varepsilon)$ is of any type, i.e., $p\in[1,\infty]$, and the measure $\rho$ is an expected function $\rho^{F}(Z)=\mathbb{E}^{F}[\ell(Z)]$ with a loss function $\ell$ having a growth order $p$. We show that for many loss functions $\ell$ arising from practical applications, it is possible to establish an exact relation between the Wasserstein DRO model (2) and the classical regularization scheme in more general terms. Notably, we prove that no loss functions beyond those we identify satisfy this equivalence, offering a definitive answer to the extent to which Wasserstein DRO can be interpreted and solved from a regularization perspective. Our results also reveal the tractability of solving Wasserstein DRO for many non-Lipschitz continuous loss functions $\ell$ and higher-order Wasserstein balls ($p>1$). While Sim et al. (2021) approach such problems via approximation due to reformulation challenges, we show that exact solutions are possible through regularization reformulations. Moreover, we extend this equivalence to settings where $\rho$ goes beyond expected value functions. Previously, this exact relation was known only for variance (Blanchet et al. (2022)) or distortion risk measures (Wozabal (2014)). We expand it to a significantly broader class of measures, encompassing higher moment coherent risk measures (Krokhmal (2007)), criteria emerging more recently from OR and ML (Rockafellar and Uryasev (2013), Rockafellar et al. (2008), Gotoh and Uryasev (2017)), and notably, general decision-theoretical models such as the celebrated rank-dependent expected utility (RDEU) (Quiggin (1982)).

Let $(\Omega,\mathcal{A},\mathbb{P})$ be an atomless probability space. A random vector $\bm{\xi}$ is a measurable mapping from $\Omega$ to $\mathbb{R}^{n+1}$, $n\in\mathbb{N}$. Denote by $F_{\bm{\xi}}$ the distribution of $\bm{\xi}$ under $\mathbb{P}$. For $p\geqslant 1$, let $L^{p}:=L^{p}(\Omega,\mathcal{A},\mathbb{P})$ be the set of all random variables with a finite $p$th moment, and $\mathcal{M}_{p}(\Xi)$ be the set of all distributions on $\Xi\subseteq\mathbb{R}^{n+1}$ with finite $p$th moments in each component. Let $L^{\infty}:=L^{\infty}(\Omega,\mathcal{A},\mathbb{P})$ be the set of all bounded random variables, and $\mathcal{M}_{\infty}(\Xi)$ be the set of all distributions on $\Xi\subseteq\mathbb{R}^{n+1}$ with bounded support. Let $q$ denote the Hölder conjugate of $p$, i.e., $1/p+1/q=1$. We use $\mathrm{ess\mbox{-}sup}Z$ to represent the essential supremum of the random variable $Z$, i.e., $\mathrm{ess\mbox{-}sup}Z=\inf\{t:F_{Z}(t)\geqslant 1\}$. Recall that for any two distributions $F_{1},\,F_{2}\in\mathcal{M}_{p}(\Xi)$, the type-$p$ Wasserstein metric is defined as

$\displaystyle W_{d,p}\left(F_{1},F_{2}\right)=\inf_{\pi\in\Pi(F_{1},F_{2})}\big{(}\mathbb{E}^{\pi}[d(\bm{\xi}_{1},\bm{\xi}_{2})^{p}]\big{)}^{{1}/{p}},$

(3)

where $d(\cdot,\cdot):\Xi\times\Xi\to\mathbb{R}_{+}\cup\{\infty\}$ is a metric on $\Xi$. The set $\Pi(F_{1},F_{2})$ denotes the set of all joint distributions of $\bm{\xi}_{1}$ and $\bm{\xi}_{2}$ with marginal distributions $F_{1}$ and $F_{2}$, respectively. In this paper, we consider $\bm{\xi}=(Y,{\bf X})\in\Xi$ with $\Xi:=\{-1,1\}\times\mathbb{R}^{n}\subseteq\mathbb{R}^{n+1}$ and apply the type-$p$ Wasserstein metric (3) with $d(\bm{\xi}_{1},\bm{\xi}_{2})=d((Y_{1},\mathbf{X}_{1}),(Y_{2},\mathbf{X}_{2}))$, defined by the following additively separable form

$\displaystyle d((Y_{1},\mathbf{X}_{1}),(Y_{2},\mathbf{X}_{2}))=\|\mathbf{X}_{1}-\mathbf{X}_{2}\|+\Theta(Y_{1}-Y_{2}),$

(4)

where $\|\cdot\|$ is any given norm on $\mathbb{R}^{n}$ with its dual norm $\|\cdot\|_{*}$ defined by $\|\mathbf{y}\|_{*}=\sup_{\|\mathbf{x}\|\leqslant 1}\mathbf{x}^{\top}\mathbf{y}$. The function $\Theta:\mathbb{R}\rightarrow\{0,\infty\}$ satisfies $\Theta(s)=0$ if $s=0$ and $\Theta(s)=\infty$ otherwise. Thus, the function (4) assigns an infinitely large cost to any discrepancy in $Y$, i.e., $Y_{1}-Y_{2}\neq 0$, and reduces to a general norm on $\mathbf{X}$ when there is no discrepancy in $Y$, i.e., $Y_{1}-Y_{2}=0$. Using this norm, for $F_{0}\in\mathcal{M}_{p}(\Xi)$ and $\varepsilon\geqslant 0$, we define the ball of distributions $\overline{{\cal B}}_{p}(F_{0},\varepsilon)$ as

$\displaystyle\overline{{\cal B}}_{p}(F_{0},\varepsilon)=\left\{F\in\mathcal{M}_{p}(\Xi):W_{d,p}(F,F_{0})\leqslant\varepsilon\right\},$

(5)

which we refer to as the type-$p$ Wasserstein ball throughout this paper.

In the remainder of this paper, we first demonstrate that the DRO problem (2), using the above definition of the Wasserstein ball:

$\displaystyle\inf_{\bm{\beta}\in\mathcal{D}}\sup_{F\in\overline{{\cal B}}_{p}(F_{0},\varepsilon)}\rho^{F}(Y\cdot\bm{\beta}^{\top}\mathbf{X}),$

(6)

can achieve generalization guarantees for any measure $\rho$ and type-$p$ Wasserstein ball. We then elucidate this guarantee for the widely used affine decision rules by exploring the connection between the optimization model (6) and the classical regularization scheme in Section 4. This framework covers various problem types, including:

• •

  Classification:

  $$\inf_{\bm{\beta}\in\mathcal{D}}\sup_{F\in\overline{{\cal B}}_{p}(F_{0},\varepsilon)}\rho^{F}(Y\cdot\bm{\beta}^{\top}\mathbf{X}),$$

• •

  Regression:

  $$\inf_{\bm{\beta}_{r}\in\mathcal{D}}\sup_{F\in\overline{{\cal B}}_{p}(F_{0},\varepsilon)}\rho^{F}((1,-\bm{\beta}_{r})^{\top}\mathbf{X}),$$

• •

  Risk minimization:

  $$\inf_{\bm{\beta}\in\mathcal{D}}\sup_{F\in\overline{{\cal B}}_{p}(F_{0},\varepsilon)}\rho^{F}(\bm{\beta}^{\top}\mathbf{X}),$$

where $F_{0}(\{Y=1\})=1$ in the cases of regression and risk minimization. These formulations accommodate numerous machine learning and operations research applications not yet addressed in the existing literature on Wasserstein DRO. Table 1 highlights examples considered in this paper.

Wasserstein DRO is typically applied in a data-driven setting, where the data-generating distribution $F^{*}$ of random variables $(Y,\mathbf{X})\in\Xi=\{-1,1\}\times{\mathbb{R}}^{n}$ is only partially observed through sample data points $(\widehat{y}_{i},\widehat{\mathbf{x}}_{i})$ for $i=1,...,N$, independently drawn from $F^{*}$. In this context, the empirical distribution $\widehat{F}_{N}$ is used as the reference distribution $F_{0}\in\mathcal{M}_{p}(\Xi)$ in Wasserstein DRO, where $\widehat{F}_{N}:=\frac{1}{N}\sum_{i=1}^{N}\delta_{(\widehat{y}_{i},\widehat{\mathbf{x}}_{i})}$ and $\delta_{\mathbf{x}}$ denotes the point-mass at $\mathbf{x}$. The central question is whether the (Wasserstein) in-sample risk

$$\sup_{F\in\overline{{\cal B}}_{p}(\widehat{F}_{N},\varepsilon_{N})}\rho^{F}(Y\cdot\bm{\beta}^{\top}\mathbf{X}),$$

can provide an upper confidence bound on the out-of-sample risk

$$\rho^{F^{*}}(Y\cdot\bm{\beta}^{\top}\mathbf{X})$$

across all possible decisions $\bm{\beta}\in{\cal D}$, with a radius $\varepsilon_{N}$ that decays at a rate scaling gracefully with the dimensionality of the random vector $(Y,\mathbf{X})$ – in other words, generalization bounds that overcome the curse of dimensionality.

Current methods face significant limitations. While generalization bounds based on the measure concentration properties of type-$p$ Wasserstein balls offer flexibility in accommodating general risk measures $\rho$, they are constrained by the curse of dimensionality inherent in high-dimensional Wasserstein balls (see e.g., Esfahani and Kuhn (2018)). Approaches such as Gao (2022) address this issue by focusing on expected value functions, i.e., $\rho^{F}(Z)=\mathbb{E}^{F}[\ell(Z)]$, leveraging the concentration properties of sample averages. However, extending these methods beyond expected value functions proves challenging, as most risk measures lack analogous concentration properties.

We take a novel perspective to tackle this challenge by focusing on the structural properties of the projection set:

$\displaystyle\overline{\mathcal{B}}_{p|\bm{\beta}}(F_{0},\varepsilon):=\{F_{Y\cdot\bm{\beta}^{\top}\mathbf{X}}:F_{(Y,\mathbf{X})}\in\overline{\mathcal{B}}_{p}(F_{0},\varepsilon)\},$

(7)

where $F_{0}\in\mathcal{M}_{p}(\Xi)$, $\bm{\beta}\in\mathbb{R}^{n}$ and $\varepsilon\geqslant 0$. We uncover that this set coincides with two particularly revealing one-dimensional sets: a one-dimensional Wasserstein ball and the projection set of a high-dimensional max-sliced Wasserstein ball. To formalize this, for $G_{0}\in\mathcal{M}_{p}(\mathbb{R}^{n})$, we first recall the definition of a type-$p$ Wasserstein ball on $\mathbb{R}^{n}$ with the metric $d(\cdot,\cdot)$ being a norm:

$\displaystyle{\cal B}_{p}(G_{0},\varepsilon)=\left\{G\in\mathcal{M}_{p}(\mathbb{R}^{n}):W_{p}(G,G_{0})\leqslant\varepsilon\right\}~{}~{}{\rm with}~{}~{}W_{p}(G,G_{0}):=\inf_{\pi\in\Pi(G,G_{0})}\left(\mathbb{E}^{\pi}[\|\bm{\xi}_{1}-\bm{\xi}_{2}\|^{p}]\right)^{{1}/{p}},$

(8)

and a type-$p$ max-sliced Wasserstein ball (Olea et al. (2022)):

$\displaystyle{\cal B}_{p}^{\rm ms}(G_{0},\varepsilon)=\left\{G\in\mathcal{M}_{p}(\mathbb{R}^{n}):\sup_{\bm{\beta}:\|\bm{\beta}\|_{*}=1}\inf_{\pi\in\Pi(G,G_{0})}\left(\mathbb{E}^{\pi}[|\bm{\beta}^{\top}\bm{\xi}_{1}-\bm{\beta}^{\top}\bm{\xi}_{2}|^{p}]\right)^{1/p}\leqslant\varepsilon\right\}.$

(9)

Here, $\|\cdot\|$ is the norm defined in (4) on $\mathbb{R}^{n}$. Without loss of generality and to be consistent with the definition of classic one-dimensional Wasserstein metric, for $n=1$, i.e., on $\mathbb{R}$, we assume that $\|\cdot\|=|\cdot|$ is the absolute-value norm.

With these definitions in place, we can now present the following equivalence results.

The equivalence established above has profound implications for both Wasserstein DRO and max-sliced Wasserstein DRO. It shows that the two frameworks coincide for any decision criterion $\rho$ under affine decision rules, both reducing to the problem of minimizing worst-case risk over a one-dimensional type-$p$ Wasserstein ball – an inherently more tractable problem, as detailed in Section 4. Notably, this equivalence reveals that, although the Wasserstein distance is a significantly stronger metric than the max-sliced Wasserstein distance, i.e., the type-$p$ Wasserstein ball is a much less conservative choice of an uncertainty set than its max-sliced counterpart (Olea et al. (2022)), Wasserstein DRO can still achieve the same generalization bounds. This insight directly leads to the following generalization bounds for general measures $\rho$, free from the curse of dimensionality.

The decay rate of the radius is approximately of the order $N^{-1/(2p)}$. While this rate is unaffected by dimensionality, its dependency on the type of Wasserstein ball, specifically the parameter $p$, is noteworthy. This dependency, often overlooked in the literature compared to the curse of dimensionality, is potentially concerning, as it suggests that the radius may shrink very slowly for higher-order Wasserstein balls. This issue arises from the use of stronger metrics – specifically, norms raised to higher powers – when defining any ball constructed similarly to the Wasserstein ball, including the max-sliced Wasserstein ball (Olea et al. (2022)). Intuitively, a higher-order Wasserstein ball is smaller than a lower-order one, requiring a larger radius to ensure the data-generating distribution is covered with high probability. In the extreme case of a type-$\infty$ Wasserstein ball, it would never contain a data-generating distribution with unbounded support, regardless of the radius.

We circumvent this issue, which we refer to as the “curse of the order $p$”, by recognizing that ensuring high-probability coverage of the data-generating distribution, while sufficient, is not necessary for developing generalization bounds. In fact, we show that even for the type-$\infty$ Wasserstein ball – where the probability of covering a data-generating distribution with unbounded support is effectively zero – the worst-case risk from the type-$\infty$ ball can still bound the out-of-sample risk of the data-generating distribution with high probability. This guarantee requires only the following mild assumption about the measure $\rho$.

For any measures satisfying Assumption 1, one can leverage the generalization bounds for the type-1 Wasserstein ball to derive similar bounds for the type-$\infty$ Wasserstein ball, with a radius scaled by a constant $\lambda$. These bounds apply to any type-$p$ Wasserstein ball, for $p\in[1,\infty]$, resulting in generalization bounds where the decay rate of the radius is no longer dependent on the order $p$.

The core idea presented in this section can be extended beyond affine decision rules, as discussed in greater detail in Section 6. Here, we further illustrate the generality of Assumption 1 by providing additional examples.

Regularization generally refers to any means applied to avoid overfitting and enhance generalization. In this regard, with the generalization guarantees established in the previous section, the Wasserstein DRO model (6) is well justified as a general regularization model. In this section, we provide further insights into the regularization effect of model (6) on the decision variable $\bm{\beta}$, offering a practically useful alternative interpretation. Specifically, we show that the previously observed equivalence between the Wasserstein DRO model and regularized empirical optimization in ML holds across much broader settings, while precisely identifying the boundary of its validity. This equivalence significantly broadens the range of both Wasserstein and max-sliced Wasserstein DRO problems that can now be solved efficiently through their regularization counterparts. Proposition 1 plays a key role in facilitating this analysis. By (10) in Proposition 1, the Wasserstein DRO model (6) can be recast as

$\displaystyle\inf_{\bm{\beta}\in\mathcal{D}}\sup_{F\in\mathcal{B}_{p}(G_{0},\varepsilon)}\ \rho^{F}(\bm{\beta}^{\top}\mathbf{Z}),$

(12)

where $G_{0}\in\mathcal{M}_{p}(\mathbb{R}^{n})$, and $\mathcal{B}_{p}(G_{0},\varepsilon)$ is a Wasserstein ball on $\mathbb{R}^{n}$ defined by (8).