Optimal Transport

Motivation

Regularized least squares is the standard tool of high-dimensional linear regression. The square-root LASSO of Belloni, Chernozhukov and Wang (Biometrika, 2011) minimizes $\sqrt{\mathrm{MSE}_n(\beta)} + \lambda\|\beta\|_1$. Unlike the LASSO, its optimal regularization parameter does not depend on the unknown noise level. Regularization is usually justified through sparsity or bias–variance arguments. Blanchet, Kang and Murthy (arXiv:1610.05627, J. Appl. Probab. 56(3), 2019) give a different justification. The square-root LASSO, and every $\ell_p$-penalized square-root least-squares estimator, is exactly a distributionally robust estimator. It minimizes the worst-case expected square loss over all data distributions within a given optimal-transport distance of the empirical distribution.

The rest of the paper builds on this representation: the radius of the transport ball is the regularization parameter, which the paper's Robust Wasserstein Profile function selects by a statistical criterion (mission 3 of this series). The duality theorem underneath, Proposition 1, is due to Blanchet and Murthy (Math. Oper. Res., 2019). Closely related representations for logistic regression appear in Shafieezadeh-Abadeh, Mohajerin Esfahani and Kuhn (NeurIPS 2015), where they are approximate. The cost function introduced in this paper makes them exact.

Setting

The training data are $n \ge 1$ pairs $(X_1, Y_1), \dots, (X_n, Y_n)$ with predictors $X_i \in \mathbb R^d$ and responses $Y_i \in \mathbb R$. No distributional assumption is made; the data are fixed vectors. The empirical distribution is $P_n = \frac1n \sum_{i=1}^n \delta_{(X_i, Y_i)}$. For $\beta \in \mathbb R^d$ the square loss is $l(x, y; \beta) = (y - \beta^T x)^2$ and the mean square error is $\mathrm{MSE}_n(\beta) = \frac1n\sum_{i=1}^n (Y_i - \beta^T X_i)^2$.

A cost function $c$ assigns to two points $z, w$ of $\mathbb R^d \times \mathbb R$ a value $c(z, w) \in [0, \infty]$, the cost of moving a unit of mass from $z$ to $w$. The optimal transport cost between probability measures $P$ and $Q$ is

$$D_c(P, Q) = \inf\Big\{ \mathbb E_\pi[c(U, W)] : \pi \text{ a probability measure on pairs } (U, W),\ \pi_U = P,\ \pi_W = Q \Big\}. \qquad (7)$$

The worst-case expected loss at radius $\delta \ge 0$ is $\sup_{P : D_c(P, P_n) \le \delta} \mathbb E_P[l(X, Y; \beta)]$, and the distributionally robust regression problem (8) minimizes it over $\beta$.

Two costs are used. With $q \in (1, \infty]$:

• the squared $\ell_q$ cost on $\mathbb R^{d+1}$, $c((x, y), (u, v)) = \|(x, y) - (u, v)\|_q^2$ (Proposition 2);
• the cost $N_q^2$, where (14) $N_q((x, y), (u, v)) = \|x - u\|_q$ if $y = v$ and $+\infty$ otherwise. Under this cost the responses cannot be moved, and only the predictors are perturbed (Theorem 1).

The exponent $p$ is the dual of $q$, $1/p + 1/q = 1$, and $\bar\beta = (-\beta, 1)$.

Formalization targets

Goal: Theorem 1 (p. 11)

For the cost $c = N_q^2$, every $\delta \ge 0$ and every $\beta \in \mathbb R^d$,

$$\sup_{P :\, D_c(P, P_n) \le \delta} \mathbb E_P\big[(Y - \beta^T X)^2\big] = \Big(\sqrt{\mathrm{MSE}_n(\beta)} + \sqrt\delta\,\|\beta\|_p\Big)^2 ,$$

and consequently

$$\inf_{\beta \in \mathbb R^d} \sup_{P :\, D_c(P, P_n) \le \delta} \mathbb E_P\big[(Y - \beta^T X)^2\big] = \inf_{\beta \in \mathbb R^d} \Big(\sqrt{\mathrm{MSE}_n(\beta)} + \sqrt\delta\,\|\beta\|_p\Big)^2 .$$

The second identity is the printed theorem; the first is what its proof establishes for each $\beta$. The goal states both.

Milestones

1. Proposition 1 (p. 10): strong duality. For a lower semicontinuous cost vanishing on the diagonal, an upper semicontinuous loss and $\delta > 0$, the worst-case expected loss equals $\min_{\gamma \ge 0} \{\gamma\delta + \frac1n \sum_i \varphi_\gamma(X_i, Y_i)\}$, with $\varphi_\gamma(z) = \sup_u \{l(u) - \gamma c(u, z)\}$ (11).
2. (28) (pp. 28–29): the closed form of $\varphi_\gamma$ for the square loss and the squared $\ell_q$ cost.
3. (29) and the display after it (p. 29): $\inf_{\gamma > b^2} \{\gamma\delta + \frac{\gamma}{\gamma - b^2} M\} = (\sqrt M + b\sqrt\delta)^2$ for $M, b, \delta \ge 0$.
4. Proposition 2 (p. 10): the analogue of the goal for the squared $\ell_q$ cost, with $\|\bar\beta\|_p$ in place of $\|\beta\|_p$ (13).
5. Outline of the proof of Theorem 1, last display (p. 29): the closed form of $\varphi_\gamma$ for the cost $N_q^2$.

Significance

The result. Theorem 1 identifies $\ell_p$-penalized square-root least squares with a min–max problem over data distributions. For $q = \infty$, $p = 1$ the minimizers are those of the square-root LASSO with $\lambda = \sqrt\delta$. The regularization parameter therefore acquires a meaning: it is the square root of the transport budget an adversary may spend perturbing the predictors. This is the basis of the paper's choice of $\delta$ by the Robust Wasserstein Profile function (§4), and of the interpretation of regularized estimators as robust to covariate perturbations. Proposition 2 shows that letting the adversary also move the responses changes the penalty to $\|(-\beta, 1)\|_p$, which is why the label-preserving cost $N_q$ is needed for an exact match.

Formalizing it. All results are proved on paper; none is formalized. A complete development gives a machine-checked strong-duality theorem for optimal-transport balls with possibly infinite costs (Proposition 1), two explicit worst-case computations, and corrected boundary cases of the closed forms (28) and the outline display, which print $+\infty$ for all $\gamma \le \|\bar\beta\|_p^2$ although the value can be finite at equality. The corrections do not affect the theorems.

Difficulty

The obvious argument fails in two places. The first is the duality step: the supremum ranges over all Borel probability measures on $\mathbb R^{d+1}$ within transport cost $\delta$, an infinite-dimensional set that is not compact in any convenient topology, with a loss that is unbounded above. Exchanging the supremum with the Lagrange multiplier of the budget constraint is Proposition 1, a theorem in its own right (Blanchet–Murthy), and its attainment claim needs $\delta > 0$.

The second is the cost $N_q$, which is $+\infty$ off $\{y = v\}$, so the standard Wasserstein duality theorems, which assume a finite metric cost, do not apply. The degenerate cases $\beta = 0$, $\mathrm{MSE}_n(\beta) = 0$, $\delta = 0$, where the objective in $\gamma$ does not blow up at both ends, must be covered separately.

Formalization scope

• Spaces. A data point is a pair in (Fin d → ℝ) × ℝ with the product σ-algebra and topology. Proposition 2's cost uses the stacked vector in Fin (d+1) → ℝ (response last, built with Fin.snoc), and $\bar\beta$ is the stacked vector of $(-\beta, 1)$.
• Norms. $\|\cdot\|_q$ and $\|\cdot\|_p$ are the norms of PiLp, with exponents in ℝ≥0∞, so $q = \infty$ (the square-root LASSO case) is included. The exponents are linked by p.HolderConjugate q, and $q \in (1, \infty]$ throughout. Theorem 1 does not print a range for $q$; the range is taken from Proposition 2, which the paper calls essentially the same result.
• Transport cost and worst case. Costs are ℝ≥0∞-valued, and $D_c$ is an infimum over probability couplings with both marginals fixed. Expectations of the nonnegative losses are lower Lebesgue integrals, and the worst case is a supremum in ℝ≥0∞ over all probability measures in the ball. No integrability side condition removes measures from the ball. Identities with a real right-hand side are stated after embedding it with ENNReal.ofReal.
• The empirical distribution is the published definition WassersteinDRO.Regularization.empiricalDistribution, applied to $i \mapsto (X_i, Y_i)$, with $n > 0$.
• $\varphi_\gamma$. A point at infinite cost contributes $-\infty$ for every $\gamma \ge 0$, including $\gamma = 0$, as in the paper's treatment of $N_q$. With the convention $0 \cdot \infty = 0$ instead, Proposition 1's minimum would not be attained for the cost $N_q^2$ at $\beta = 0$.
• Proposition 1 is stated for a nonnegative loss and $\delta > 0$; both are restrictions of the page, recorded in the item.
• Corrections. (28) and the outline display are stated with their corrected boundary cases. The one-dimensional lemma behind (29) is stated as a greatest lower bound over $\gamma > b^2$, including $b = 0$, $M = 0$, $\delta = 0$.

A formalization in which the transport infimum did not fix both marginals, allowed sub-probability couplings, or used a Bochner integral would make the worst case trivially $+\infty$ or $0$. The conventions above rule this out: at $\delta = 0$ the ball is $\{P_n\}$ and both sides of the goal equal $\mathrm{MSE}_n(\beta)$.

The work needs Kantorovich-type duality for lower semicontinuous costs on $\mathbb R^m$ (absent from Mathlib), Hölder's inequality with its equality case for PiLp, and elementary one-variable optimization. The duality theorem and the transport-cost definition are reusable beyond this mission: mission 2 of this series (classification) uses Proposition 1 with the cost $N_q$, $\rho = 1$. Contributions that prove Proposition 1, or its weak-duality half, are particularly welcome.

Selected references

• J. Blanchet, Y. Kang, K. Murthy, Robust Wasserstein Profile Inference and Applications to Machine Learning, J. Appl. Probab. 56(3), 2019; arXiv:1610.05627v4. https://arxiv.org/abs/1610.05627
• J. Blanchet, K. Murthy, Quantifying distributional model risk via optimal transport, Math. Oper. Res. 44(2), 2019. https://doi.org/10.1287/moor.2018.0936
• A. Belloni, V. Chernozhukov, L. Wang, Square-root lasso: pivotal recovery of sparse signals via conic programming, Biometrika 98(4), 2011. https://doi.org/10.1093/biomet/asr043
• S. Shafieezadeh-Abadeh, P. Mohajerin Esfahani, D. Kuhn, Distributionally robust logistic regression, NeurIPS 2015. https://arxiv.org/abs/1509.09259
• C. Villani, Optimal Transport: Old and New, Springer, 2009. https://doi.org/10.1007/978-3-540-71050-9