PaperlayerView exact revision on arXiv

Source-linked AI summary

Large-Scale Methods for Distributionally Robust Optimization

Daniel Levy, Yair Carmon, John C. Duchi, Aaron Sidford

arXiv:2010.05893v2math.OCcs.LGstat.ML

TL;DR

The paper tackles the lack of scalable optimization methods for distributionally robust learning with CVaR and χ2 uncertainty sets. It develops sampling-based and multilevel Monte Carlo gradient methods, proving gradient-evaluation guarantees independent of training-set size and parameter dimension. Experiments on MNIST and ImageNet report 9×–36× fewer gradient computations than full-batch methods, while a remaining gap concerns the χ2-constrained objective.

  • Problem

    DRO can improve robustness across subpopulations and environments, but existing sampling-based gradient estimates are biased and χ2 methods may scale with training-set size or parameter dimension.

  • Method

    The paper develops sampling-based algorithms using biased mini-batch surrogates and an essentially unbiased MLMC gradient estimator for CVaR and χ2 DRO objectives.

  • Results

    Experiments on MNIST and ImageNet show generalization improvements over ERM and 9×–36× fewer gradient computations than full-batch methods.

  • Takeaways & Limitations

    The algorithms provide a foundation for applying stochastic gradient methods to large-scale convex φ-divergence DRO.

  • Takeaways & Limitations

    The χ2-constrained objective remains between an O(ρϵ^-3) upper bound and an Ω(ρϵ^-2) lower bound, so its optimal rate is unresolved.

Abstract

from arXiv · show

We propose and analyze algorithms for distributionally robust optimization of convex losses with conditional value at risk (CVaR) and $χ^2$ divergence uncertainty sets. We prove that our algorithms require a number of gradient evaluations independent of training set size and number of parameters, making them suitable for large-scale applications. For $χ^2$ uncertainty sets these are the first such guarantees in the literature, and for CVaR our guarantees scale linearly in the uncertainty level rather than quadratically as in previous work. We also provide lower bounds proving the worst-case optimality of our algorithms for CVaR and a penalized version of the $χ^2$ problem. Our primary technical contributions are novel bounds on the bias of batch robust risk estimation and the variance of a multilevel Monte Carlo gradient estimator due to [Blanchet & Glynn, 2015]. Experiments on MNIST and ImageNet confirm the theoretical scaling of our algorithms, which are 9--36 times more efficient than full-batch methods.

1 Introduction

The paper addresses the scalability barrier preventing distributionally robust optimization from becoming common in machine learning. It develops sampling-based methods for CVaR and χ2 objectives with guarantees independent of training-set size and parameter dimension.

  • Motivation: DRO trains models against worst-case distributions in an uncertainty set, unlike ERM, but scalable optimization methods remain limited.Existing χ2 methods can require gradient evaluations scaling linearly or worse with training-set size or parameter dimension.
  • Objectives: The paper studies CVaR, χ2-constrained, and penalized-χ2 objectives, extending the framework beyond uniform finite training distributions.The developments apply to non-uniform distributions with infinite support.
  • Approach: Sampling-based algorithms target optimality gaps ϵ while measuring complexity by the number of gradient evaluations, independent of d and N.For reference, stochastic subgradient methods for ERM require order ϵ^-2 evaluations independent of dimension and training-set size.
  • Approach: Mini-batch estimation yields a CVaR bias bound scaling as 1/√n and a stronger 1/n bound for penalized χ2.For sufficiently large batches, an ϵ^2-minimizer of the surrogate is also an ϵ-minimizer of the population objective.
  • Approach: An MLMC gradient estimator is essentially unbiased, uses a logarithmic expected number of samples in n, and supports complexity bounds scaling with ϵ^-2 for CVaR and χ2 penalty.The χ2-constrained objective instead uses its Lagrange dual and achieves a guarantee scaling as ϵ^-3.
  • Experiments: Experiments on MNIST and ImageNet report generalization improvements over ERM and 9×–36× fewer gradient computations than full-batch methods.The mini-batch estimator was practically superior to MLMC because MLMC’s increased variance outweighed its bias advantage.

2 Preliminaries

The preliminaries define the paper’s general φ-divergence DRO formulation, assumptions on the loss and domain, and the three principal objectives. They also introduce mini-batch notation and the gradient-evaluation complexity measure.

  • Assumptions: The optimization variable lies in a closed convex domain of diameter at most R, while the loss is bounded, convex, and G-Lipschitz.Additional assumptions separately impose H-Lipschitz gradients or a G_icdf-Lipschitz inverse loss CDF.
  • Assumptions: The paper notes that most bounds do not require the additional smoothness and inverse-CDF assumptions, which are argued to be frequently nonrestrictive.The inverse-CDF condition can yield a general 1/n bias bound independent of uncertainty-set size.
  • DRO formulation: The framework uses convex functions φ and ψ with φ(1)=ψ(1)=0, divergence radius ρ, and penalty λ to define generalized DRO objectives.The formulation covers both constrained and regularized uncertainty-set constructions.
  • Objectives: The three main objectives are CVaR, χ2 constraint, and χ2 penalty, represented through specific choices of φ and ψ.A KL-regularized CVaR variant is also defined for later use.
  • Mini-batch surrogate: The mini-batch surrogate averages the robust objective over random batches of size n, enabling stochastic subgradient methods based on batch losses.The surrogate is introduced to replace direct estimation of the population robust objective.
  • Complexity: Method complexity is measured by the number of ∇ℓ evaluations needed to reach accuracy ϵ, with runtime additionally depending on dimension and single-sample evaluation cost.Solving the batch objective itself takes O(n log n) time in the studied problems.

3 Mini-batch gradient estimators

The mini-batch method replaces the full robust objective with a biased but efficiently estimable surrogate, then controls bias and variance to obtain complexity guarantees.

  • Bias of the batch estimator: The mini-batch gradient is unbiased for the surrogate objective but generally biased for the full robust objective.The analysis attributes this bias to the noncommutation of infimum and expectation.
  • Complexity guarantees: Nesterov acceleration exploits the 1/n variance reduction, decreasing stochastic-gradient steps from order 1/ϵ^2 to order 1/ϵ when applicable.The accelerated guarantee applies to objectives with sufficiently smooth gradients, including χ2 penalty and KL-CVaR.
  • Bias of the batch estimator: CVaR surrogate bias scales as 1/√n, while an inverse-cdf Lipschitz assumption yields a general 1/n bound.For CVaR, only local Lipschitzness around the relevant quantile is required.
  • Bias of the batch estimator: The χ2-penalized objective has a 1/n bias bound obtained by combining its dual representation with Efron-Stein control of the multiplier variance.The proof relates the sample and population minimizers of the dual formulation.
  • Variance of the batch estimator: For χ2-bounded penalty-type objectives, the mini-batch gradient variance decreases as 1/n; this includes CVaR and χ2 penalty.The CVaR variance bound is independent of its penalty parameter.
  • Variance of the batch estimator: The variance bounds fail for the χ2-constrained formulation, where instances exist with gradient variance at least on the order of ρ.This prevents directly extending the batch-estimator guarantee to the constrained formulation.
  • Complexity guarantees: The resulting total evaluation complexity is (GR)^2B/(λϵ^2) for χ2 penalty and (GR)^2G_icdf/ϵ^3 for general losses.The χ2-penalty rate is stated for λ ≤ B, while the general-loss bound uses the inverse-cdf Lipschitz parameter.
  • Complexity guarantees: Under specified parameter conditions, the theorem recovers the claimed CVaR and χ2-penalty rates and can approximately minimize maximum loss in sublinear time.The maximum-loss statement holds when G_icdf is finite, even for the entire simplex uncertainty set.

4 Multi-level Monte Carlo (MLMC) gradient estimators

The MLMC estimator removes the mini-batch method’s linear sample-cost growth by using an unbiased multilevel construction with logarithmic expected sample complexity.

  • Estimator construction: MLMC forms an unbiased approximation to the surrogate gradient using only a logarithmic expected number of samples in the target batch size.The expected sample size is n_0(1 + log_2(n/n_0)).
  • Estimator construction: The estimator samples a truncated geometric level and evaluates gradients on progressively larger sample batches.Its parameters include a minimum sample size n_0 and a target size n = 2^jmax n_0.
  • Estimator construction: A modified level distribution and optional n_0 > 1 distinguish this estimator from Blanchet and Glynn’s proposal.The altered level distribution is described as crucial for this setting.
  • Variance and complexity: The MLMC gradient estimator has a second moment bounded logarithmically in n for CVaR and χ2 penalty, enabling the stated complexity guarantees.The variance analysis uses sign consistency across robust maximizers at adjacent levels.
  • Variance and complexity: The same MLMC bounds extend to KL-CVaR and fail for the χ2-constrained gradient, so that formulation requires a different optimization scheme.The alternative jointly optimizes the χ2 penalty dual over x and λ using MLMC estimators for both derivatives.
  • Variance and complexity: For CVaR, the MLMC method achieves an ϵ-optimality guarantee with complexity controlled by n_0αϵ^2 log^2 n.The theorem sets n proportional to B^2/(αϵ^2) and chooses n_0 no larger than order log n.
  • Variance and complexity: For χ2 penalty, the corresponding dependence replaces α^-1 with 1 + B/λ.This gives the analogous complexity conclusion for the penalized χ2 objective.

5 Lower bounds

The paper proves worst-case lower bounds for its MLMC guarantees, using distinct hard-instance constructions for CVaR and penalized χ2 optimization.

  • Lower-bound framework: The lower bounds use a standard stochastic-oracle model in which each query returns a loss and gradient at a chosen point.Complexity is measured by the number of individual first-order evaluations.
  • Lower-bound conclusions: The resulting lower bounds match the MLMC guarantees up to logarithmic factors for CVaR and penalized χ2.The paper presents these bounds as worst-case optimality results within the stated oracle setting.
  • CVaR lower bound: For CVaR, the lower bound holds even in one dimension and remains valid when the oracle reveals the entire sampled loss function.This makes the CVaR lower bound stronger than the penalized-χ2 lower bound.
  • Penalized-χ2 lower bound: The penalized-χ2 lower bound uses a classical high-dimensional hard instance because the objective is not positively homogeneous in x.The construction makes one sampled loss informative and the other uninformative.
  • CVaR lower bound: The CVaR construction reduces optimization to statistically difficult testing between two perturbed Bernoulli distributions.The two robust objectives have well-separated values at their respective minima.

6 A doubling scheme for minimizing Lχ2

The paper minimizes the constrained χ2 DRO objective by optimizing a sequence of penalized approximations over geometrically spaced parameter intervals. This doubling scheme improves the direct penalized approach while yielding a fixed-constraint guarantee, though its rate remains above a known lower bound.

  • Approximation scheme: The constrained χ2 objective is approximated through its penalized counterpart by restricting the penalty parameter λ to a suitable interval.The method jointly optimizes x and λ using projected stochastic gradients with an MLMC gradient estimator.
  • Doubling scheme: A direct parameter substitution gives complexity scaling as ϵ^-5, worse than the mini-batch rate of ϵ^-4.The paper therefore divides the λ range into geometrically spaced sub-intervals.
  • Doubling scheme: The doubling scheme uses intervals with λ(i+1)/λ(i)=2 and computes an approximately optimal solution on each interval before selecting among them.The number of intervals is logarithmic in B/ϵ, and the resulting procedure obtains an ϵ-accurate solution to the constrained χ2 objective.
  • Guarantee and limitation: Theorem 4 gives a rigorous complexity guarantee for minimizing the constrained χ2 objective with fixed constraint ρ by optimizing the penalization parameter λ.The theorem addresses the fixed-ρ constrained problem rather than only the penalized objective.
  • Guarantee and limitation: The optimal rate for the constrained χ2 problem remains open because the O(ρϵ^-3) upper bound exceeds the Ω(ρϵ^-2) lower bound.The authors note that tuning λ directly by validation may be more practical when ρ is unknown.

7 Experiments

Experiments on digit recognition and ImageNet evaluate how gradient-estimator batch size affects DRO optimization and generalization. Larger batches substantially reduce work relative to full-batch methods, while DRO improves worst-subgroup metrics with dataset-dependent average-performance trade-offs.

  • Experimental setup: The experiments measure how total work depends on the gradient estimator and mini-batch size across two heterogeneous datasets.The study focuses on digit recognition and ImageNet classification settings.
  • Convergence: Figure 1 reports DRO-objective convergence across five repetitions, with shaded regions showing minimum-to-maximum variability and zooms exposing small-batch bias floors.The objective traces are compared as optimization progresses.
  • Convergence: Batch-estimator bias becomes negligible for nearly every batch size above 10, except for the χ2 objective in the digits experiment.The bias is particularly weak for the penalized χ2 objective.
  • Convergence: The MLMC estimator avoids a bias floor by design but converges more slowly than the mini-batch estimator in these experiments.Its optimal convergence guarantees rely on smoothness and, for CVaR, sufficiently rapid bias-floor decay.
  • Efficiency: 27–36× more work is required by full-batch methods than by batch sizes 50–5K on ImageNet, while digits require 9–16× more work.These comparisons measure reaching within 2% of the optimal value.
  • Generalization: DRO improves worst-subgroup loss by 17.5%–27% on digits and worst-performing-class log loss by 5%–10% on ImageNet, while ImageNet average performance declines.On ImageNet, average loss rises by up to 10% and average accuracy drops by roughly 1 point.

8 Conclusion

The paper establishes stochastic-gradient guarantees for large-scale convex φ-divergence DRO and identifies where those guarantees and practical benefits remain limited. It concludes that CVaR and penalized χ2 rates are resolved, while constrained χ2 optimality and broader machine-learning benefits require further study.

  • Contributions: The work provides rigorous convergence guarantees for large-scale convex φ-divergence DRO using stochastic gradient methods.The conclusion presents these guarantees as a foundation for practical use.
  • Open problems: The constrained χ2 problem remains theoretically unresolved, with an O(ρϵ^-3) upper bound and an Ω(ρϵ^-2) lower bound.The authors state that the upper bound may be optimal only in high-dimensional settings.
  • Open problems: The practical benefits of large-scale φ-divergence DRO require further research and appear larger when distinct subgroups exist and low logarithmic loss matters.These are experimental suggestions rather than a general conclusion across all machine-learning settings.
  • Scope: The paper focuses theoretically and experimentally on convex losses, although its algorithms apply directly to non-convex losses and its bias and variance analyses do not require convexity.The implementation supports arbitrary prediction models through automatic differentiation.
  • Objective formulations: The robust objective can be characterized through φ-divergence uncertainty sets and equivalent dual formulations involving expectations and Lagrange multipliers.For pure-penalty objectives, the divergence-constraint multiplier is unnecessary.
  • Objective formulations: CVaR corresponds to a bounded likelihood-ratio uncertainty set, while χ2 objectives use χ2-divergence-based uncertainty sets.The paper also notes that some results extend to more general φ- or Rényi-divergence balls.
  • Objective formulations: For penalized χ2, sufficiently large λ reduces the objective to empirical risk with variance regularization.The finite-sample formulation uses sorted losses to characterize the associated optimizer.
  • Optimization properties: A subgradient of the robust objective is the expectation under a maximizing distribution Q⋆ of a loss subgradient.Smoothness of the loss supports Lipschitz-gradient results for penalized χ2 and KL-CVaR objectives.

A.3 Stochastic gradient method on the dual objective

The section analyzes stochastic gradient optimization of the dual robust objective, combining bounded dual variables with batch-bias estimates to obtain convergence guarantees under several distributional assumptions.

  • Stochastic gradient method: The dual formulation enables stochastic gradient methods for minimizing the robust objective over model parameters.The analysis derives runtime guarantees where direct references were not found.
  • Dual-variable bounds: For CVaR, restricting the multiplier η to [0, B] follows because the optimal η is the value at risk and lies within the loss range.For Lχ2-pen, the corresponding range is [−λ, B].
  • χ2-bounded objectives: The three studied objectives are χ2-bounded, with constants 1/α − 1 for KL-CVaR, ρ for Lχ2, and B/λ for Lχ2-pen.These bounds control the divergence of worst-case distributions used in the analysis.
  • Worst-case tightness: The CVaR bias lower bound scales as B√((1 − α)/(αn)) for a Bernoulli loss, showing the general upper bounds are worst-case tight.The Bernoulli construction establishes that the rate cannot generally be improved without additional assumptions.
  • Inverse-cdf assumptions: A Lipschitz inverse-cdf assumption yields 1/n bias, but it requires positive loss density throughout the loss range and fails for discrete or separated-bulk distributions.Under a relaxed approximation condition, the bias is bounded by 2Gicdf/(n + 1).

B.3 Proofs of variance bounds

The variance analysis uses Efron–Stein arguments to bound objective and gradient estimators for χ2-bounded penalty objectives, while constructing a counterexample for constrained χ2 gradients.

  • Penalty objectives: The gradient variance guarantee applies when φ = 0 and ψ is strictly convex, covering pure-penalty objectives such as Lkl-CVaR and Lχ2-pen.The proof relies on monotonicity of ψ*′ and χ2-boundedness.
  • Variance bounds: For C-χ2-bounded objectives, Efron–Stein controls objective and gradient variance by the effect of replacing one sample in a mini-batch.The gradient argument uses the G-Lipschitz property of the loss.
  • Constrained χ2 limitation: The χ2 constraint objective has objective variance at most 2(1 + ρ)B2/n, but no matching O((1 + ρ)G2/n) gradient-variance bound holds in general.The paper states that the φ = 0 requirement is necessary for the gradient result.
  • Lower-bound construction: A three-point distribution and G-Lipschitz loss produce constrained-χ2 gradient behavior that differs across sample events, yielding a worst-case variance lower bound.The construction switches between gradients near G and substantially smaller values with constant probability.

B.5 Proofs of complexity bounds

The complexity proofs select batch sizes to control robust-risk bias and iteration counts to control optimization error, producing rates that depend on objective type and regularity assumptions.

  • Complexity framework: With suitable batch size n and iteration count T, stochastic gradient methods achieve expected robust-risk error at most ϵ with total complexity nT.The resulting bounds are stated separately for CVaR, penalized χ2, and constrained χ2 objectives.
  • Penalized-χ2: For Lχ2-pen with λ ≤ B, the complexity satisfies nT ≲ (GR)2B/(λϵ2).This rate is obtained by combining the bias bound with stochastic-gradient convergence.
  • Lipschitz inverse-cdf: For objectives of the form (5), a Lipschitz inverse-cdf gives complexity nT ≲ (GR)2Gicdf/ϵ3.This rate is independent of the uncertainty-set size under the stated regularity condition.
  • CVaR: The theorem distinguishes CVaR rates for non-smooth, smooth, and smooth-loss/Lipschitz-inverse-cdf settings.The proof chooses n from Proposition 1 and T from Proposition 3 for each case.

C.1 Proofs of moment bounds

The moment analysis establishes multilevel Monte Carlo bounds for χ2-bounded pure-penalty objectives and explains why the same guarantees fail for constrained χ2 gradients.

  • MLMC construction: The MLMC estimator uses a randomized truncated geometric batch level and telescopes across sample sizes to remain unbiased for the base objective.Its expected sample size is n0(1 + log2(n/n0)).
  • Moment bounds: For χ2-bounded pure-penalty losses, the MLMC correction satisfies E∥bDk∥2 ≲ 1/k, producing second-moment guarantees with controlled expected cost.This applies to Lkl-CVaR, Lχ2-pen, and their gradients under the stated assumptions.
  • Constrained-χ2 limitation: For constrained χ2, MLMC gradient estimators lack comparable moment bounds, and worst-case optimization may require T ≍ n(GR)2/ϵ2 iterations.In that regime, a mini-batch estimator with batch size n is no worse in the stated worst-case scaling.
  • Complexity guarantees: Theorem 2 provides MLMC complexity guarantees for CVaR and extends them to Lχ2-pen by replacing α−1 with 1 + B/λ.The construction uses n ≍ B2/(αϵ2) and a minimum batch size n0 bounded logarithmically in n.
  • Estimator design: The geometric parameter p = 1/2 balances second moment and expected sample cost; smaller p makes second moments grow polynomially with n/n0.This differs from the proposal of Blanchet and Glynn because their smoother functionals permit smaller p.

D Lower bound proofs

The paper proves lower bounds for CVaR and penalized-χ2 distributionally robust optimization, showing that dimension-free algorithms require substantial oracle complexity. The CVaR and penalized-χ2 arguments use different hard-instance constructions.

  • Lower-bound statements: Theorem 3 gives minimax lower bounds for CVaR and penalized-χ2 objectives over convex G-Lipschitz losses.The penalized-χ2 result applies in dimension d_ϵ constrained by a term proportional to (GR)^2ϵ^-2 log(GR/ϵ).
  • CVaR lower bound: The CVaR lower bound reduces stochastic optimization to hypothesis testing through a standard Le Cam argument.The construction uses linear losses ℓ(x,s)=x·s over X=[−R,R].
  • CVaR lower bound: CVaR’s positive homogeneity yields a closed-form reduction that identifies distinct minimizing points for the constructed distributions.The objective satisfies L_CVaR(x;P_v)=|x|·L_CVaR(sign(x);P_v).
  • Penalized-χ2 lower bound: The penalized-χ2 lower bound instead uses a high-dimensional hard instance because the objective is not positively homogeneous in x.The construction makes the robust loss equivalent, up to an additive constant, to a hard convex optimization function.
  • Penalized-χ2 lower bound: Dimension-free optimization of the hard instance requires Ω(ϵ^-2) deterministic-oracle interactions, transferring to the penalized-χ2 problem through informative samples.The reduction uses a Bernoulli sampling process in which only S=1 reveals information about the hard function.

E Doubling schemes proofs

The paper analyzes batch and multilevel methods for penalized-χ2 optimization, including parameter-doubling schemes and their empirical behavior. Experiments show batch-size-dependent bias and work savings over full-batch optimization, while dual SGM becomes less competitive as uncertainty grows.

  • Subproblem guarantees: The penalized-χ2 subproblem is solved over bounded λ intervals using stochastic-gradient iterations with controlled bias and high-probability error.The analysis chooses n0, n, T, γ_x, and γ_λ to obtain an ϵ-accurate solution.
  • Doubling scheme: Theorem 4 analyzes a doubling scheme over geometrically spaced penalty parameters λ(i), combining approximate minimizers across intervals.The scheme uses K=⌈log_2(2B/ϵ)⌉ intervals and selects an iterate with an ϵ-accurate robust objective value.
  • Generalization: DRO improves worst-subgroup performance relative to ERM, with 17.5%–27% lower worst-subgroup loss on Digits and 5.6%–9% lower loss on ImageNet.The paper notes that logarithmic loss is operationally meaningful because it measures calibration.
  • Scope and comparison: Methods targeting known subgroups will likely perform better than subgroup-agnostic DRO, although DRO requires less subgroup annotation.DRO still uses a subgroup validation set to tune the uncertainty-set size.
  • Optimization behavior: Mini-batch methods converge faster than full-batch methods, while MLMC methods are theoretically optimal but empirically slower to converge.The experiments also find that mini-batch bias is much smaller than the theoretical bound and is negligible for batch sizes around 50.
  • Bias behavior: The bias between population and mini-batch robust objectives is large at small batch sizes and generally decays as 1/n.The χ2-constrained Digits experiment is an exception to the observed 1/n decay.

F.6.2 Comparison with primal-dual methods

The experiments compare primal-dual and primal-only methods for constrained-χ2 and CVaR objectives, finding that relative efficiency depends strongly on dataset scale and the number of misclassified examples.

  • Method description: ϵ^-2 iterations suffice for the described primal-dual update under appropriate γx and γq.The passage states this guarantee for the update analyzed with local-norm regret methods.
  • Results: Digits favors primal-dual methods, despite their guarantees carrying an additional factor of N = 60.6K relative to primal-only methods.The analysis attributes this behavior partly to the small number of incorrectly classified training points, fewer than 100.
  • Results: ImageNet reverses the trend: the primal-only method significantly outperforms the primal-dual method when more than 100K examples are misclassified.The experiment uses N = 1.2M for ImageNet, making the dependence on the effective number of misclassified examples consequential.
  • Results: Replicating every dataset element m times leaves the primal-only method unchanged but is expected to increase primal-dual complexity as the norm of ˜gq grows by m.This thought experiment illustrates the differing scalability with dataset size.
  • Limitations: Primal-dual tuning is considerably harder because it requires two learning rates, and overly large γq can produce a suboptimal converged point.Thus, reducing the learning rate until divergence disappears is not a reliable tuning procedure for this method.
Loading 2010.05893v2…