Wasserstein Robust Risk And Lipschitz Losses

Why Transport Cost Controls Risk Sensitivity

A distribution can change by moving probability mass across the sample space. The effect on risk depends on two things: how far that mass moves and how much the loss changes over that distance. Wasserstein ambiguity sets express the first quantity through a transport cost; Lipschitz continuity of the loss controls the second.

This connection gives a direct way to bound risk under distribution shift. If a loss changes slowly relative to the chosen metric, a fixed transport budget can change its expected value only modestly. If the loss is steep, or if the metric makes relevant moves costly or cheap, the same budget has a different implication.

Wasserstein Distance and Robust Risk

Let $(\mathcal X,d)$ be a metric space, and let $P$ and $Q$ be Borel probability distributions on it. Assume both have finite first moment: for some, and hence any, $x_0\in\mathcal X$, $\int d(x,x_0)\,dP(x)<\infty$ and $\int d(x,x_0)\,dQ(x)<\infty$. A coupling of $P$ and $Q$ is a joint distribution $\pi$ whose first marginal is $P$ and whose second marginal is $Q$.

Definition: The 1-Wasserstein distance between $P$ and $Q$, with ground metric $d$, is $$ W_1(P,Q)=\inf_{\pi\in\Pi(P,Q)}\int_{\mathcal X\times\mathcal X} d(x,z)\,d\pi(x,z), $$ where $\Pi(P,Q)$ is the set of all couplings of $P$ and $Q$. For a reference distribution $P$ and radius $\rho\geq 0$, the Wasserstein ambiguity set is $\mathcal B_{W_1}(P,\rho)=\{Q: W_1(P,Q)\leq\rho\}$. The worst-case risk of a loss $\ell:\mathcal X\to\mathbb R$ over this set is $\sup_{Q\in\mathcal B_{W_1}(P,\rho)}E_Q[\ell(X)]$.

The transport cost is an average distance, not a maximum distance: a small amount of mass can travel far if the average cost remains within the budget. For the risk statements below, assume $\ell$ is measurable and has finite expectation under $P$. In particular, a globally Lipschitz loss is integrable under any distribution with finite first moment, because its magnitude grows at most linearly with distance from a fixed point.

The Lipschitz Risk Bound

A loss is $L$-Lipschitz with respect to $d$ if $L\geq 0$ and $|\ell(x)-\ell(z)|\leq Ld(x,z)$ for every $x,z\in\mathcal X$. The constant $L$ measures the largest change in loss per unit of distance. The central result turns this pointwise property into a bound on changes in expected risk.

Theorem: Suppose $P$ has finite first moment, and $\ell:\mathcal X\to\mathbb R$ is $L$-Lipschitz. For every distribution $Q$ with finite first moment, $$ \left|E_Q[\ell(X)]-E_P[\ell(X)]\right|\leq L W_1(P,Q). $$ Consequently, $$ E_P[\ell(X)]\leq \sup_{Q:W_1(P,Q)\leq\rho}E_Q[\ell(X)] \leq E_P[\ell(X)]+L\rho. $$

Proof. Take any coupling $\pi\in\Pi(P,Q)$, and write its coordinates as $X$ and $Z$, so that $X\sim P$ and $Z\sim Q$. The Lipschitz condition gives $$ \left|E_Q[\ell(Z)]-E_P[\ell(X)]\right| =\left|\int(\ell(z)-\ell(x))\,d\pi(x,z)\right| \leq\int|\ell(z)-\ell(x)|\,d\pi(x,z) \leq L\int d(x,z)\,d\pi(x,z). $$ This holds for every coupling. Taking the infimum of the right-hand side over all couplings proves the first inequality. If $W_1(P,Q)\leq\rho$, it follows that $E_Q[\ell(X)]\leq E_P[\ell(X)]+L\rho$. Taking the supremum over such $Q$ gives the upper bound. The lower bound holds because $P$ itself belongs to the ambiguity set: $W_1(P,P)=0\leq\rho$. $\square$

The result controls sensitivity, not necessarily the exact worst-case risk. It guarantees that the robust risk lies between the reference risk and that risk plus $L\rho$. Whether the upper bound is attained depends on the available directions for moving probability mass and on how the loss changes along those directions.

Worked Example: A Discrete Distribution With an Attained Bound

Let $\mathcal X=\{0,1,2\}$ with $d(x,z)=|x-z|$. Let $P$ assign probability $1/2$ to 0 and $1/2$ to 2, and take $\ell(x)=\min\{x,1\}$. Find the reference risk and the worst-case risk over the Wasserstein ball of radius $\rho=1/2$.

The loss values at 0, 1, and 2 are 0, 1, and 1. The function is 1-Lipschitz: its change is at most the distance between any two points. Thus $$ E_P[\ell(X)]=\tfrac12(0)+\tfrac12(1)=\tfrac12, \qquad \sup_{W_1(P,Q)\leq 1/2}E_Q[\ell(X)]\leq\tfrac12+1\cdot\tfrac12=1. $$

Now let $Q$ put probability $1/2$ at 1 and $1/2$ at 2. Moving the mass at 0 to 1 and leaving the mass at 2 in place gives a coupling with cost $(1/2)(1)+(1/2)(0)=1/2$, so $W_1(P,Q)\leq1/2$. In fact the distance equals $1/2$: to supply the $1/2$ mass that $Q$ puts at 1, the cheapest available source mass is at 0, at distance 1, so at least $1/2$ transport cost is necessary.

Under $Q$, the expected loss is $(1/2)(1)+(1/2)(1)=1$. This reaches the upper bound, so the worst-case risk is exactly 1. Here the allowed movement transfers probability to a point where the loss increases at the full Lipschitz rate.

When the Bound Is Exact—and When It Is Not

A simple continuous example shows that the bound can be exact even when distributions have unbounded support. Let the sample space be the real line with its usual distance and consider the loss $\ell(x)=x$. This loss is 1-Lipschitz. If $P\sim N(0,1)$, its expected loss is zero. Translating $P$ to the right increases the expected loss by exactly the translation distance, and costs exactly that distance in $W_1$.

Worked Example: A Continuous Normal Distribution

Let $P=N(0,1)$, let $\ell(x)=x$, and consider a Wasserstein radius $\rho\geq0$. The reference risk is $E_P[X]=0$. The Lipschitz risk bound gives a worst-case risk no greater than $\rho$.

For the reverse inequality, let $Q=N(\rho,1)$. If $X\sim N(0,1)$, the pair $(X,X+\rho)$ is a coupling of $P$ and $Q$ with transport cost $E[|(X+\rho)-X|]=\rho$. Hence $W_1(P,Q)\leq\rho$, so $Q$ is in the ambiguity set. Its expected loss is $E_Q[X]=\rho$. Therefore the worst-case risk is at least $\rho$, and combining the two inequalities gives $$ \sup_{Q:W_1(P,Q)\leq\rho} E_Q[X]=\rho. $$

This example also verifies that the distance of the translated distributions is exactly $\rho$: the risk bound already proved, applied to the 1-Lipschitz function $x\mapsto x$, implies $E_Q[X]-E_P[X]\leq W_1(P,Q)$, so $W_1(P,Q)\geq\rho$. Together with the coupling upper bound, $W_1(P,Q)=\rho$.

In contrast, a Lipschitz bound need not be tight. It limits the loss change by the steepest rate possible anywhere, even if that rate is not achieved along a useful transport direction from the reference distribution. The next example shows how to calculate a bound while keeping that distinction clear.

Worked Example: A Logistic Loss With a Conservative Bound

Let $\mathcal X=[0,1]$ with $d(x,z)=|x-z|$. For a fixed positive label, suppose the score is $g(x)=2x$, giving logistic loss $\ell(x)=\log(1+\exp(-2x))$. Take $P$ to be a point mass at $x=0.5$, and set $\rho=0.1$.

Differentiating gives $\ell'(x)=-2/(1+\exp(2x))$, so $|\ell'(x)|\leq2$ on $[0,1]$. The mean value theorem therefore makes $\ell$ 2-Lipschitz. The reference risk is $$ E_P[\ell(X)]=\ell(0.5)=\log(1+e^{-1})\approx0.3133. $$ The robust risk is consequently at most $0.3133+2(0.1)=0.5133$.

For a feasible alternative, let $Q$ be a point mass at $0.4$. Its Wasserstein distance from $P$ is $|0.5-0.4|=0.1$, and its risk is $$ E_Q[\ell(X)]=\ell(0.4)=\log(1+e^{-0.8})\approx0.3711. $$ Thus the worst-case risk is at least $0.3711$ and at most $0.5133$. The bound is valid but is not shown to be exact; the global Lipschitz constant 2 overstates the loss change along this particular move.

How the Metric and Loss Geometry Matter

The bound $L\rho$ is meaningful only after specifying the ground metric. Rescaling distances rescales the numerical Lipschitz constant and changes the interpretation of the transport budget. For example, if $d'(x,z)=c\,d(x,z)$ for $c>0$, then $W_1$ under $d'$ is $c$ times $W_1$ under $d$, while a loss that is $L$-Lipschitz under $d$ is $L/c$-Lipschitz under $d'$. The product of the corresponding constant and distance budget is consistent when the budget is rescaled as well.

For a vector input $x\in\mathbb R^p$ and a linear loss $\ell(x)=a^\top x$, the Euclidean metric gives Lipschitz constant $\|a\|_2$, because the Cauchy–Schwarz inequality yields $$ |a^\top x-a^\top z|=|a^\top(x-z)|\leq\|a\|_2\|x-z\|_2. $$ Thus the sensitivity bound is determined by the direction and magnitude of $a$, as well as by the transport radius. A different metric can assign different costs to movement in different coordinates, changing which distributional perturbations fit within the same budget.

More generally, when a loss is built by composing a prediction function with a loss function, its sensitivity depends on both pieces. A stable prediction function and a slowly changing loss can yield a small Lipschitz constant; a steep loss or a rapidly changing score can yield a larger one. The constant should therefore be justified for the actual loss and metric, rather than selected only for convenient arithmetic.

Common Pitfall

Assuming the bound is an equality. From Lipschitz continuity one may conclude that the worst-case risk is at most $E_P[\ell(X)]+L\rho$; one may not conclude it equals that quantity without showing that a feasible distribution attains the increase. In the logistic example, the calculations establish a lower bound of about 0.3711 and an upper bound of about 0.5133, not equality with either one.
Using a Lipschitz constant without naming the metric. The inequality $|\ell(x)-\ell(z)|\leq Ld(x,z)$ depends on $d$. A constant derived for Euclidean distance cannot automatically be used with a different transport cost.
Treating the transport radius as a maximum displacement. $W_1$ limits the average cost under a coupling. It does not require every unit of mass to move by at most $\rho$.

Key takeaway: For an $L$-Lipschitz loss, every distribution within Wasserstein distance $\rho$ of $P$ has risk at most $L\rho$ above the reference risk. The ground metric determines what movement costs, and the loss geometry determines how much that movement can matter.

Check Your Understanding

Use the definitions and results above to answer the following questions.

Why does finite first moment, together with a Lipschitz loss, ensure that the loss has finite expectation?
If a loss is 3-Lipschitz and the Wasserstein radius is 0.2, what upper bound does the theorem give on the increase in risk?
In the discrete example, which movement of probability mass attains the robust-risk upper bound?
Why does the translated normal distribution in the continuous example have Wasserstein distance exactly $\rho$, rather than merely at most $\rho$?
What extra argument would be needed to claim that a Lipschitz risk bound is the exact worst-case risk?