Source-linked AI summary
Distributionally Robust Optimization and Generalization in Kernel Methods
Matthew Staib, Stefanie Jegelka
TL;DR
Existing φ-divergence and Wasserstein uncertainty sets have limitations for machine-learning DRO, motivating alternatives that preserve robustness and support generalization. The paper studies MMD DRO, shows its connection to Hilbert-norm regularization, derives Gaussian-kernel generalization results and a new regularizer, and develops a tractable approximation extending variance regularization.
Problem
Existing φ-divergence and Wasserstein uncertainty sets have drawbacks for DRO, including support restrictions, computational expense, and nontrivial assumptions.
Method
The paper defines DRO uncertainty sets using MMD, analyzes the adversary through Hilbert-space mean embeddings, and develops an empirical approximation beyond kernel methods.
Results
MMD DRO is roughly equivalent to regularizing by the Hilbert norm of the loss, yields Gaussian-kernel generalization bounds matching standard ones up to a small constant, and generalizes variance regularization.
Takeaways & Limitations
MMD DRO provides structural connections among robustness, regularization, and generalization while extending the scope of variance-based regularization.
Takeaways & Limitations
The adversary analysis relaxes distribution constraints, and Wasserstein-related comparisons remain tied to assumptions on ground metrics and losses.
Abstract
from arXiv · showhide
Distributionally robust optimization (DRO) has attracted attention in machine learning due to its connections to regularization, generalization, and robustness. Existing work has considered uncertainty sets based on phi-divergences and Wasserstein distances, each of which have drawbacks. In this paper, we study DRO with uncertainty sets measured via maximum mean discrepancy (MMD). We show that MMD DRO is roughly equivalent to regularization by the Hilbert norm and, as a byproduct, reveal deep connections to classic results in statistical learning. In particular, we obtain an alternative proof of a generalization bound for Gaussian kernel ridge regression via a DRO lense. The proof also suggests a new regularizer. Our results apply beyond kernel methods: we derive a generically applicable approximation of MMD DRO, and show that it generalizes recent work on variance-based regularization.
1 Introduction
The paper introduces MMD-based DRO to address limitations of φ-divergence and Wasserstein uncertainty sets, establishing links between robustness, regularization, and generalization. It derives structural results, a Gaussian-kernel generalization analysis, a new regularizer, and a tractable approximation extending variance regularization.
- φ-divergence balls around empirical distributions typically exclude the population distribution, while Wasserstein approaches are more computationally expensive and assumption-dependent.
- MMD DRO defines uncertainty sets using maximum mean discrepancy, a kernel-based distance between distributions.
- MMD DRO is roughly equivalent to regularizing by the Hilbert norm of the loss rather than the model.
- For Gaussian kernel ridge regression, a bound on the loss Hilbert norm yields generalization bounds matching standard ones up to a small constant.
- The generalization proof suggests a new regularizer for Gaussian kernel ridge regression.
- The paper derives a computationally tractable approximation of MMD DRO that applies broadly and generalizes variance-based regularization.
2 Background and Related Work
This section reviews DRO’s distributional uncertainty framework, its generalization principle, and the dominant φ-divergence and Wasserstein formulations. It then situates MMD among kernel-based distribution distances and contrasts its ability to vary higher-order moments with moment-based approaches.
- DRO minimizes worst-case expected loss over distributions in an uncertainty set centered on the empirical distribution.
- If the uncertainty set contains the population distribution with high probability, the DRO objective provides a high-probability bound on population loss.
- The choice of distributional discrepancy controls the uncertainty radius and the tractability of the DRO problem.
- φ-divergence DRO can equal variance regularization, but empirical-support restrictions typically prevent its uncertainty set from containing the population distribution.
- Wasserstein uncertainty sets can contain continuous distributions, but their analysis commonly relies on Euclidean or Mahalanobis ground metrics and nontrivial assumptions.
- MMD uses kernel embeddings to define a distance between distributions, and its uncertainty sets can allow higher-order moments to vary according to the kernel.
3 Generalization bounds via MMD DRO
MMD DRO yields a high-probability population-risk bound because empirical-to-population MMD shrinks at O(n^-1/2), while the DRO adversary reduces to Hilbert-norm regularization under an RKHS assumption.
- 3 Generalization bounds via MMD DRO: MMD DRO uses an uncertainty radius tied to empirical accuracy, and MMD between the empirical sample and population shrinks at O(n^-1/2).This rate determines how large the uncertainty set must be to contain the population distribution with high probability.
- 3 Generalization bounds via MMD DRO: With k(x, x) ≤ M, a radius chosen from the MMD concentration bound gives a high-probability upper bound on population risk.The bound holds with probability 1 − δ and is obtained through the DRO generalization principle.
- 3.1 Bounding the DRO adversary’s problem: The worst-case embedding moves from µP in the direction of ℓf, identifying the adversary’s optimal direction in the Hilbert space.The displayed solution specifies the embedding shift proportional to the loss function.
- 3.1 Bounding the DRO adversary’s problem: For ℓf ∈ H, optimizing over mean embeddings converts the DRO adversary into empirical risk plus ϵ∥ℓf∥H.The equality follows by optimizing the loss expectation over a Hilbert-space ball around the population mean embedding.
- 3.1 Bounding the DRO adversary’s problem: MMD DRO therefore produces an empirical-risk-plus-loss-norm bound, unlike Wasserstein results that penalize gradients and may include higher-order terms.The Wasserstein comparison also notes slower radius shrinkage O(n^-1/d) in high dimension.
4 Connections to kernel ridge regression
The kernel-ridge analysis makes the loss-norm regularizer tractable for Gaussian kernels and recovers standard generalization bounds while motivating an additional complexity penalty.
- 4 Connections to kernel ridge regression: The proposed objective penalizes ∥ℓf∥H rather than the usual kernel-method penalty ∥f∥H, but this composition norm is not generally computationally tractable.The difficulty arises because f and ℓf may not belong to the same RKHS.
- 4.1 Bounding norms of products: For Gaussian-kernel ridge regression, the squared loss is ℓf(x) = (f(x) − h(x))^2, so bounding its RKHS norm reduces to bounding norms of products.The analysis specializes to Gaussian kernels because product norms require bandwidth-sensitive tools.
- 4.1 Bounding norms of products: For f, g ∈ Hσ, the product-norm theorem bounds ∥fg∥ at a reduced bandwidth using trace representations and positive-semidefinite matrix inequalities.The proof represents f, g, and fg through matrix traces before applying tr(XY) ≤ tr(X)tr(Y).
- 4.2 Implications: kernel ridge regression: The resulting DRO generalization theorem controls population risk for Gaussian kernel ridge regression under separate bounds on f, h, f^2, and h^2.The separate bounds provide finer control of function-class complexity than typical norm-based bounds.
- 4.2 Implications: kernel ridge regression: The DRO-based theorem recovers standard Gaussian-kernel ridge-regression generalization results up to a universal constant.The paper further notes that the DRO bound can motivate regularizing with both ∥f∥σ and the additional loss-related norm.
- 4.2 Implications: kernel ridge regression: The Representer Theorem does not generally reduce optimization to finite kernel expansions when the objective includes ∥f^2∥2, although this norm can be evaluated and differentiated.For a kernel expansion, the squared norm is expressed using a trace involving the Gram matrix.
5 Approximation and connections to variance regularization
The paper approximates MMD DRO by restricting adversarial distributions to the empirical sample support, yielding a tractable weight-based problem. This approximation connects MMD DRO to variance regularization and can be stronger with an appropriate kernel, though it does not guarantee out-of-sample performance.
- Connections to variance regularization: The resulting regularizer connects MMD DRO to variance regularization and, with the right kernel, makes MMD DRO a stronger regularizer than variance.
- Approximation: The approximation restricts adversarial distributions Q to the empirical support, representing them with nonnegative simplex weights.This restriction makes the adversary weaker than the unrestricted MMD ball.
- Approximation: The support restriction prevents guarantees of out-of-sample performance because the approximating distributions generally have different support.
- Approximation: For sufficiently small uncertainty radius ϵ, the nonnegativity constraints are inactive and the adversary’s problem has a closed-form solution.
6 Experiments
Experiments compare the proposed tighter regularizer with the standard kernel ridge regression regularizer in easy and hard synthetic settings. With suitable λ, the proposed regularizer performs better and is less sensitive when λ is too low.
- Experimental setup: The experiment compares the proposed ∥f∥²_{σ/2} regularizer with ∥f∥²_σ on a synthetic Gaussian-kernel problem.The study uses easy and hard noise regimes and sweeps λ.
- Results: With λ tuned properly, the tighter ∥f∥²_{σ/2} regularizer yields better performance in both easy and hard settings.
- Results: The proposed ∥f∥²_{σ/2} regularizer is less sensitive to λ because performance decays slowly when λ is too low.
A Proofs of main structural results
The proof derives a candidate solution using weak duality and establishes equality by proving strong duality. Completing the square identifies the optimizing distribution and yields the matching upper and lower bounds.
- Duality argument: The proof first uses weak duality to derive a candidate solution, then uses that candidate to establish strong duality.
- Candidate solution: Completing the square isolates the only term depending on µ_Q and identifies an optimizer by making the nonpositive squared-norm term zero.
- Candidate solution: Optimizing the dual variable produces an upper bound, with the balancing choice determined by the terms involving ϵ and ∥f∥_H.
- Strong duality: Substituting the candidate µ*_Q into the original supremum supplies the matching lower bound, proving equality.
B Gaussian kernel bounds
The appendix develops algebraic and kernel identities used to establish Gaussian-kernel norm bounds. These identities rewrite composite squared-distance expressions and connect them to trace and inner-product representations.
- Vector identities: The appendix introduces vector identities that rewrite sums of squared distances using norms, inner products, and cross terms.
- Vector identities: For the defined quantities S and T, algebraic expansion shows that S = T.
- Vector identities: Expanding the first three terms of the transformed expression cancels cross terms and yields the norm representation required for T.
- Gaussian-kernel bounds: The Gaussian-kernel proof uses these identities to rewrite exponential expressions in terms of four pairwise squared distances.
- Gaussian-kernel bounds: The appendix then applies the resulting identities and kernel-matrix relations to prove Proposition 4.1 and derive trace forms for Hilbert-space norms.
- Gaussian-kernel bounds: The derivation relates Gaussian-kernel expansions to inner products and trace expressions, including formulas for ∥g∥²_σ and ∥fg∥²_2.
- Gaussian-kernel bounds: The appendix uses cyclicity of the trace and previously established identities to complete the norm calculations.
B.1 Trace inequality
The proof uses the trace inner product and Cauchy–Schwarz inequality, with positive semidefiniteness ensuring the relevant eigenvalue terms are nonnegative.
- The trace inner product is defined as ⟨X, Y⟩ = tr(XᵀY) = tr(XY) because X is symmetric.
- Cauchy–Schwarz bounds the trace expression in the proof.
- Nonnegative eigenvalues justify the inequality, and the same reasoning applies to any positive semidefinite matrix, including Y.
B.2 Extensions of Proposition 4.1
The section develops corollaries and extensions of Proposition 4.1, including a tractable finite-dimensional computation of the product norm using shared kernel expansions and a kernel matrix.
- The extensions aim to make ∥fg∥σ tractable to compute.
- The functions f and g share a finite kernel expansion over the same inputs but may have different coefficients.
- The kernel matrix K is formed from pairwise kernel evaluations of the shared expansion points.
- The proof represents the associated matrices through vectors whose inner products reproduce K, then uses trace identities and ZᵀZ = K.
C Proofs for Section 5
These proofs reduce the Section 5 optimization problems using convexity, balancing, matrix identities, strong duality, and concave quadratic maximization.
- Dropping the wi ≥ 0 constraint rewrites problem (23), while its objective becomes a convex quadratic in λ.
- For the supremum problem, the argument uses a simplification, Slater’s condition, strong duality, and concave quadratic maximization.
- The optimal inner value is substituted back into the overall problem formulation.
- The unconstrained optimizer is λ* = (1ᵀK⁻¹ℓ)/(1ᵀK⁻¹1), after which the remaining positive terms are balanced.
- The resulting balanced objective yields the overall optimal value for the problem.
- The proof of Lemma 5.2 applies the matrix inversion lemma to derive the needed expressions involving K⁻¹.