Source-linked AI summary
Regularization via Mass Transportation
Soroosh Shafieezadeh-Abadeh, Daniel Kuhn, Peyman Mohajerin Esfahani
TL;DR
The paper addresses overfitting and the mismatch between empirical and out-of-sample loss when the data-generating distribution is unknown. It minimizes worst-case expected loss over Wasserstein balls around the empirical distribution, yielding tractable regularized models and new generalization guarantees. It also connects classical regularization to transportation costs and validates the theoretical out-of-sample guarantees experimentally.
Problem
Unknown input-output distributions make ideal out-of-sample-risk minimization impossible, while empirical-loss minimization can overfit scarce training data.
Method
The paper minimizes worst-case expected loss over Wasserstein balls around the empirical distribution for regression and classification, with tractable reformulations, kernelization, and neural-network approximations.
Results
The proposed models provide new generalization bounds under minimal assumptions, while classical regularized learning emerges as a special case of the distributionally robust framework.
Takeaways & Limitations
Transportation-based distributional robustness supplies a unified probabilistic account of regularization and tractable tools for controlling out-of-sample loss, error, and risk.
Takeaways & Limitations
Exact neural-network Lipschitz moduli are NP-hard to compute, so the framework uses an efficiently computable upper bound instead.
Abstract
from arXiv · showhide
The goal of regression and classification methods in supervised learning is to minimize the empirical risk, that is, the expectation of some loss function quantifying the prediction error under the empirical distribution. When facing scarce training data, overfitting is typically mitigated by adding regularization terms to the objective that penalize hypothesis complexity. In this paper we introduce new regularization techniques using ideas from distributionally robust optimization, and we give new probabilistic interpretations to existing techniques. Specifically, we propose to minimize the worst-case expected loss, where the worst case is taken over the ball of all (continuous or discrete) distributions that have a bounded transportation distance from the (discrete) empirical distribution. By choosing the radius of this ball judiciously, we can guarantee that the worst-case expected loss provides an upper confidence bound on the loss on test data, thus offering new generalization bounds. We prove that the resulting regularized learning problems are tractable and can be tractably kernelized for many popular loss functions. We validate our theoretical out-of-sample guarantees through simulated and empirical experiments.
1 Introduction
The paper develops distributionally robust learning models that regularize regression and classification by minimizing worst-case expected loss over Wasserstein balls around the empirical distribution. It establishes tractable formulations, probabilistic interpretations of regularization, and new generalization and error-estimation guarantees.
- Motivation: Finite-sample empirical loss can produce fragile, overfit estimators because the unknown input-output distribution is observed only through training samples.Regularization is the standard remedy, typically adding a hypothesis-complexity penalty to empirical loss.
- Distributionally Robust Formulation: Wasserstein ambiguity sets contain distributions transformable into the empirical distribution at bounded transportation cost, enabling worst-case expected loss optimization.Their statistical convergence properties make properly scaled Wasserstein balls natural confidence regions for the data-generating distribution.
- Tractability: The proposed models hedge against all input-output distributions in a Wasserstein ball and admit tractable convex reformulations for many common losses and linear hypotheses.Lifted variants are kernelizable for nonlinear hypotheses, while neural-network models can be approximated with convex regularization and stochastic proximal gradient descent.
- Probabilistic Interpretation: Classical regularized learning models arise when the cost of moving probability mass along the output space tends to infinity.The input-space transportation cost determines the regularization function, while the Wasserstein-ball radius determines its weight.
- Generalization Bounds: The models provide generalization bounds under minimal assumptions, including an improved dimension-independent square-root convergence law when the hypothesis space is bounded.The naïve Wasserstein confidence bound suffers from the curse of dimensionality and converges slowly for high-dimensional inputs.
- Additional Guarantees: The framework also yields tractable confidence intervals for fixed-hypothesis error or risk and convex programs for computing worst-case distributions for stress tests.Robust and distributionally robust learning are equivalent under dispersion or separability conditions, whereas the distributionally robust models remain efficiently solvable without those conditions.
2 Problem Statement
Supervised learning seeks hypotheses that predict unknown outputs from finite data, but empirical-risk minimization can overfit and cannot directly optimize unknown out-of-sample error. The paper proposes replacing this objective with a worst-case expected loss over distributions near the empirical distribution, yielding tractable models and confidence guarantees.
- Classical Statistical Learning: Supervised learning infers an unknown target function from finitely many independent input-output samples drawn from an unknown distribution.Regression uses continuous outputs, while classification uses categorical outputs.
- Classical Statistical Learning: Learning algorithms select hypotheses by minimizing empirical average loss, with linear models representing hypotheses through weight vectors.The empirical distribution is uniform over the training data, and different losses recover familiar regression and classification problems.
- Regularization: Classical regularization adds a hypothesis-complexity penalty to empirical loss, but common probabilistic interpretations depend on arbitrary prior assumptions.For example, L2- and L1-regularization can be linked to Gaussian and Laplacian priors, respectively.
- The Generalization Problem: Finite-data learning can overfit: a hypothesis with low in-sample error may still have high out-of-sample error on future inputs.The out-of-sample error is the desired quantity, but the unknown distribution prevents its direct computation.
- A New Perspective on Regularization: The proposed approach minimizes expected loss over distributions close to the empirical distribution, with the worst-case value serving as an upper confidence bound on out-of-sample error.The resulting models are computationally tractable, can reduce to regularized learning, and admit kernelizable conservative approximations under stated conditions.
3 Tractable Reformulations
The paper develops tractable reformulations of distributionally robust learning over linear, kernelized, and neural-network hypotheses. These reformulations connect worst-case transportation-based losses to robust optimization and classical regularization under stated assumptions.
- 3.1 Distributionally Robust Linear Regression: Distributionally robust linear regression admits finite convex reformulations for suitable convex loss functions and uncertainty sets.The framework covers piecewise-affine or Lipschitz losses under convex, closed input and output spaces and norm-induced transportation metrics.
- 3.1 Distributionally Robust Linear Regression: For κ = ∞, linear distributionally robust regression reduces to classical regularized learning with dual-norm penalty ∥w∥∗ and weight determined by the Wasserstein radius.The resulting regularization depends on the input-space transportation cost, while infinite output-movement cost restricts distributional changes to inputs.
- 3.1 Distributionally Robust Linear Regression: The worst-case expectations in regression and classification can be represented using discrete distributions that are feasible and asymptotically optimal as γ ↓0.The regression and classification results construct such extremal distributions through optimization-based characterizations.
- 3.3 Nonlinear Hypotheses: Reproducing Kernel Hilbert Spaces: Lifted distributionally robust regression and classification problems over RKHS hypotheses are equivalent to tractable kernelized formulations under symmetric positive definite kernels.The resulting optimization requires computational effort similar to linear learning despite the potentially infinite-dimensional hypothesis space.
- 3.4 Nonlinear Hypotheses: Neural Networks: Neural-network distributionally robust learning has conservative regularized empirical-loss upper bounds, but the resulting weight optimization is generally non-convex.A convex regularization upper bound is available when κ = ∞, while exact network Lipschitz computation is NP-hard and is replaced by an easily computed upper bound.
4 Generalization Bounds
The paper derives Wasserstein-based confidence and generalization bounds for distributionally robust learning. Basic bounds suffer from dimensionality dependence, while restricting the hypothesis space yields improved dimension-independent rates and supports RKHS hypotheses.
- Basic Generalization Bound: Wasserstein measure concentration turns ambiguity balls into confidence sets for the unknown data-generating distribution.The construction requires a light-tailed input-output distribution.
- Basic Generalization Bound: The basic generalization bound applies to optimal distributionally robust hypotheses when the Wasserstein radius exceeds ρN(η).The guarantee holds for N ≥ 1, n ≠ 1, η ∈ (0, 1], and ρ ≥ ρN(η).
- Assumptions and Limitations: The light-tail assumption is restrictive but unavoidable for the stated measure-concentration result, though bounded-support, Gaussian, and exponential cases satisfy it.This assumption constrains the distributional scope of the basic analysis.
- Asymptotic Consistency: Under a suitable vanishing-radius schedule, the distributionally robust solution converges almost surely to the ideal solution minimizing out-of-sample error.The radius tends to zero as N increases, and ηN must converge to zero at a carefully chosen rate.
- Curse of Dimensionality: The radius ρN(η) decreases as sample size grows but suffers a fundamental curse of dimensionality, requiring 2^n more samples to halve it.For large N, the radius is proportional to N^−1/(n+1), with improvements possible only for special distributions such as finitely supported ones.
- Extension to Nonlinear Hypotheses: The framework extends directly to RKHS hypotheses, and its generalization bound remains valid even when the hypothesis space has infinite VC-dimension.The extension requires a calm kernel satisfying the stated growth condition.
- Improved Generalization Bound: The improved bound is dimension-independent in its sample-size requirement: halving ρ′N(η) requires increasing N by a factor of 4.The result depends on the hypothesis-space restriction and light-tail assumptions.
- Improved Generalization Bound: With bounded hypothesis-space assumptions, the improved bound changes the radius rate to O([log(η^−1) + n log(N))/N]^(1/2), essentially a square-root law.The restriction controls uniform differences between empirical and true expected loss.
5 Error and Risk Estimation
The paper uses Wasserstein ambiguity sets to estimate pessimistic and optimistic prediction errors or classification risks. These estimates are tractable for linear hypotheses, extend to RKHS models, and yield confidence intervals for data-dependent hypotheses.
- Problem Setup: The analysis estimates out-of-sample prediction error in regression and misclassification risk in classification over a Wasserstein ball around the empirical distribution.The linear setting uses mean absolute error for regression and misclassification probability for classification.
- Regression: Worst-case and best-case regression errors are defined as the supremum and infimum of expected absolute loss over the Wasserstein ambiguity ball.Theorem 5.1 provides estimates for both quantities.
- Classification: Worst-case and best-case classification risks are defined analogously using the misclassification indicator over the same ambiguity ball.Theorem 5.2 provides the corresponding risk estimates.
- Tractability: For a fixed hypothesis, the error and risk estimation problems are tractable linear programs that can be solved efficiently.The formulations use auxiliary variables and norm constraints associated with transportation costs.
- Confidence Intervals: Setting the radius to ρN(η/2) gives confidence 1 − η intervals for true error and risk, even for hypotheses depending on the training data.The true regression-error interval can be calculated analytically from the reformulation.
- Nonlinear Hypotheses: The estimation results extend to nonlinear hypotheses in an RKHS by lifting inputs into feature space and representing nonlinear hypotheses as linear feature-space hypotheses.The lifted formulations preserve tractability.
6 Numerical Results
The experiments evaluate distributionally robust classifiers across handwritten-digit, UCI, object-recognition, and generalization-bound settings. Results consistently support the method’s out-of-sample performance, learned-parameter behavior, and confidence-bound interpretation.
- 6.2 Regularization with Learned Parameters: The DRSVM was evaluated on three binary MNIST digit-pair problems, varying Wasserstein radius, label-flipping cost, and kernel function.Linear hypotheses were used initially, with results averaged over 100 independent trials.
- 6.2 Regularization with Learned Parameters: Cross-validation learned all free parameters, and the DRSVM consistently outperformed the RSVM across polynomial, Laplacian, and Gaussian kernels.The polynomial kernel with degree one subsumes the non-kernelized DRSVM.
- 6.2 Regularization with Learned Parameters: The DRSVM outperformed the RSVM in AUC for all input norms and datasets, matching the RSVM’s improvement over classical hinge-loss minimization.AUC was used because the considered datasets were imbalanced, making CCR potentially misleading.
- 6.3 Multi-Label Classification: On PASCAL VOC 2007, classifiers were assessed using AP and mAP across 20 binary object categories.Experiments included separate binary classifiers and fine-tuning the last layer of pretrained ALEXNET.
- 6.3 Multi-Label Classification: Fine-tuning the last layer of ALEXNET may improve performance, while the distributionally robust spectral norm regularizer consistently outperformed almost all alternatives.The classifier was optimized with stochastic proximal gradient descent and data augmentation.
- 6.4 Generalization Bounds: The generalization-bound experiment found that increasing ρ widens the confidence interval until it covers the true risk; at ρ ≈ 0.009, the interval [0.008, 0.162] contained it with probability 95%.At ρ = 0, the bounds coincide with empirical training risk, which is lower than test risk under overfitting.
A.1 Proofs of Section 3
The proofs reformulate worst-case expectations over Wasserstein balls using robust optimization, duality, and support-function arguments. These reformulations establish tractable expressions for several loss functions and connect existing regularizers to distributional robustness.
- Robust reformulation: Wasserstein-ball worst-case expectations can be re-expressed as classical robust optimization problems.This follows from a preparatory robust-reformulation lemma and couplings whose empirical marginal is uniform over training samples.
- Convex losses: Convex Lipschitz losses reduce to finite expressions governed by the loss Lipschitz modulus and the dual norm of the predictor.The key lemma yields the loss at the nominal point when lip(L)∥β∥*≤γ and +∞ otherwise.
- Tractable reformulations: Piecewise-linear, Huber, ε-insensitive, and pinball losses admit tractable reformulations through epigraphs, conjugates, support functions, and strong duality.The proofs introduce auxiliary variables for max terms and invoke Slater-based duality where needed.
- Duality: Strong duality and support-function identities simplify the minimax formulation into a finite optimization problem under the stated convexity and feasibility conditions.The derivation uses Slater’s condition, dual norms, and the fact that training samples lie in the support set.
- Worst-case attainment: Discrete distributions constructed near an extremal point belong to the Wasserstein ball, and their expected losses converge to the worst-case value.The convergence argument establishes the claimed equality for the worst-case expectation.
- Piecewise-linear losses: For convex piecewise-linear losses, the proof compares equivalent optimization formulations and identifies restrictions caused by requiring perturbations to be independent of the loss piece.The comparison relies on conditional distributions associated with the uniform empirical marginal.
A.2 Proofs of Section 4
The proofs derive concentration guarantees for Wasserstein distances and worst-case risks under light-tail assumptions. They extend scalar concentration to linear-growth functions and use covering arguments to obtain uniform bounds over predictors.
- Concentration: A preparatory lemma establishes sub-Gaussian tails for sample averages of linear-growth functions under the light-tail assumption.The resulting bound holds for every δ ∈[0, 1] and N ∈N, irrespective of the particular linear-growth function.
- Pushforward measures: The scalar concentration result transfers to pushforward distributions of linear-growth functions, including empirical Wasserstein distances.The pushforward inherits the light-tail constants, and the result applies to functions with asymptotic growth rate at most 1.
- Scope: Theorem 4.6 is proved explicitly for classification with Ξ = R^n × {−1, 1}; the regression proof is omitted as a similar modification.This is the proof’s stated scope boundary.
- Uniform bounds: Uniform guarantees over the hypothesis space are obtained by covering predictors with a finite ∆-net and controlling loss variation through Lipschitz continuity.The covering-number argument uses |W∆| < (Ω/∆)^n − 1 and bounds the loss change via lip(L) and the input norm.
- Uniform deviation: The proof combines concentration for the loss and input norm with probability subadditivity to bound the uniform deviation event.The auxiliary parameter δ is chosen so the resulting probability is at most η.
- Conclusion: The final estimate follows after substituting the covering-number relation and the definition of ρ′, yielding the claimed probability bound.The proof concludes that the probability in the target event does not exceed η.
A.3 Proofs of Section 5
The proofs derive tractable formulations for absolute-value and classification losses using earlier Wasserstein reformulations, dual norms, and linear programming duality. They also obtain a corresponding best-case risk expression.
- Absolute-value loss: The absolute-value loss is treated as the ε-insensitive loss with ε = 0, so the corresponding formulation follows directly from the earlier corollary.The proof fixes w and sets Ξ = R^(n+1).
- Dual formulation: The subordinate minimization problem is converted into linear constraints involving t_i, r_i, v_i, q_i, and the dual norm of (w, −1).Strong linear programming duality and the duality theorem justify the conversion.
- Classification constraints: The classification formulation imposes −1 ≤ v_i ≤ 1 together with v_i∥(w, −1)∥* ≤ λ and −v_i∥(w, −1)∥* ≤ λ.These constraints arise after substituting the dual representation into the reformulated problem.
- Classification loss: Indicator functions for classification errors are represented as pointwise maxima of extended real-valued concave functions.This transforms the classification formulation into a tractable reformulation of the worst-case problem.
- Best-case risk: Strong linear programming duality yields the stated best-case risk expression after substituting the dual constraints.The best-case derivation is analogous, with details omitted in the proof.