Large-Scale Methods for Distributionally Robust Optimization

Introduction

The growing role of machine learning in high-stakes decision-making raises the need to train reliable models that perform robustly across subpopulations and environments . Distributionally robust optimization (DRO) shows promise as a way to address this challenge, with recent interest in both the machine learning community and in operations research . Yet while DRO has had substantial impact in operations research, a lack of scalable optimization methods has hindered its adoption in common machine learning practice.

Finally, the penalized $\chi^{2}$ objective replaces the hard constraint (3) with regularization,

We develop sampling-based algorithms for each of the objectives (2)–(4). In Table 1 we summarize their complexities and compare them to previous work. Each entry of the table shows the number of (sub)gradient evaluations to obtain a point with optimality gap $\epsilon$; for reference, recall that for ERM the stochastic subgradient method requires order $\epsilon^{-2}$ evaluations, independent of $d$ and $N$. We discuss related work further in Section 1.1 after outlining our approach.

We establish that $\overline{\mathcal{L}}$ is a useful surrogate for $\mathcal{L}$ by proving uniform bounds on the error $|\mathcal{L}(x;P_{0})-\overline{\mathcal{L}}(x;n)|$. For CVaR (2) we prove a bound scaling as $1/\sqrt{n}$ and extend it to other objectives, including (3), via the Kusuoka representation . Notably, for the penalty version of the $\chi^{2}$ objective (4) we prove a stronger bound scaling as $1/n$.

This analysis implies that, for large enough mini-batch size $n$, an $\frac{\epsilon}{2}$-minimizer of $\overline{\mathcal{L}}$ is also an $\epsilon$-minimizer of $\mathcal{L}$. Further, for CVaR and the $\chi^{2}$ penalized objective, we show that the variance of the gradient estimator decreases as $1/n$, and we use Nesterov acceleration to decrease the required number of (stochastic) gradient steps.

To obtain algorithms with improved oracle complexities, in Section 4 we present a theoretically more efficient multi-level Monte Carlo (MLMC) gradient estimator which is a slight modification of the general technique of Blanchet and Glynn . The resulting estimator is unbiased for $\nabla\overline{\mathcal{L}}(x;n)$ but requires only a logarithmic number of samples in $n$ in expectation. (In contrast, the above-mentioned mini-batch estimator requires $n$ samples). For CVaR and $\chi^{2}$ penalty we control the second moment of the gradient estimator, resulting in complexity bounds scaling with $\epsilon^{-2}$. In Section 5 we prove that these rates are worst-case optimal up to logarithmic factors.

Unfortunately, direct application of the MLMC estimator for the $\chi^{2}$-constrained objective (3) demonstrably fails to achieve a second moment bound. Instead, in Section 6 we optimize its Lagrange dual—the $\chi^{2}$ penalty—with respect to $x$ and Lagrange multiplier $\lambda$. Using a doubling scheme on the $\lambda$ domain, we obtain a complexity guarantee scaling as $\epsilon^{-3}$.

We conclude the paper in Section 8 with some remarks and directions for future research.

Distributionally robust optimization grows from the robust optimization literature in operations research , and the fundamental uncertainty about the data distribution at test time makes its application to machine learning natural. Experiments in the papers show promising results for CVaR (2) and $\chi^{2}$-constrained (3) DRO, while other works highlight the importance of incorporating additional constraints into the uncertainty set definition . Below, we review the prior art on solving these DRO problems at scale.

Similar dual formulations exist for both the constrained and penalized $\chi^{2}$ objectives (3) and (4), and dual SGM provides similar guarantees to CVaR for the penalized $\chi^{2}$ objective (4). For the constrained problem (3), the additional Lagrange multiplier associated with the constraint induce a so-called “perspective transform” , making the method unstable. Indeed, Namkoong and Duchi report that it fails to converge in practice and instead propose a stochastic primal-dual method with convergence rate $(1+\rho N)\epsilon^{-2}$. Their guarantee is optimal in the weak regularization regime where $\rho\lesssim 1/N$ , but is worse than the full-batch method in the setting where $\rho\gtrsim 1$. Hashimoto et al. propose a different scheme alternating between ERM on $x$ and line search over a Lagrange multiplier, but do not provide complexity bounds. Duchi and Namkoong prove that for a sample of size $N^{\prime}\approx\rho^{2}d\epsilon^{-2}$ the empirical objective converges to $\mathcal{L}(x;P_{0})$ uniformly in $x\in\mathcal{X}$; substituting $N^{\prime}$ into the full-batch complexity bound implies a rate of $\rho^{2}d\epsilon^{-4}$. This guarantee is independent of $N$, but features an undesirable dependence on $d$. Ghosh et al. use the mini-batch gradient estimator and gradually increase the batch size to $N$ as optimization progresses; they do not provide convergence rate bounds. We establish concrete rates for fixed batch sizes independent of $N$.

Preliminaries

We collect notation, establish a few assumptions, and provide the most important definitions for the remainder of the paper in this section.

Most of our bounds do not require Assumptions A1 and A2. Moreover, in Appendix B.2 we argue that these assumptions are frequently not restrictive.

Then, for convex $\phi,\psi$ with $\phi(1)=\psi(1)=0$, a constraint radius $\rho\geq 0$, and penalty $\lambda\geq 0$ the general form of the objectives we consider is

The form (5) allows us to redefine the objectives (2)–(4) for general $P_{0}$ (nonuniform and with infinite support):

$\chi^{2}$ constraint. $\mathcal{L}_{\chi^{2}}$ corresponds to $\phi(t)=\chi^{2}(t)\coloneqq\frac{1}{2}(t-1)^{2}$ and $\psi=0$.

$\chi^{2}$ penalty. $\mathcal{L}_{\chi^{2}\textup{-pen}}$ corresponds to $\phi=0$ and $\psi(t)=\chi^{2}(t)=\frac{1}{2}(t-1)^{2}$.

Additionally, define the following smoothed version of the CVaR objective, which we use in Section 3.

In Appendix A we present additional standard formulations and useful properties of these objectives.

denote the loss with respect to the empirical distribution on $s_{1}^{n}$. Averaging the robust objective over random batches of size $n$, we define the surrogate objective

Mini-batch gradient estimators

In this section, we develop and analyze stochastic subgradient methods using the subgradients of the mini-batch loss (6). That is, we estimate $\nabla\mathcal{L}(x;P_{0})$ by sampling a mini-batch $S_{1},\ldots,S_{n}\stackrel{{\scriptstyle\textup{iid}}}{{\sim}}P_{0}$ and computing

To show that the mini-batch gradient estimator is nevertheless effective for minimizing $\mathcal{L}$, we proceed in three steps. First, in Section 3.1 we prove uniform bounds on the bias $\mathcal{L}-\overline{\mathcal{L}}$ that tend to zero with $n$. Second, in Section 3.2 we complement them with $1/n$ variance bounds on $\nabla\mathcal{L}(x;S_{1}^{n})$. Finally, Section 3.3 puts the pieces together: we apply the SGM guarantees to bound the complexity of minimizing $\overline{\mathcal{L}}$ to accuracy $\epsilon/2$, using Nesterov acceleration to exploit our variance bounds, and choose the mini-batch size $n$ large enough to guarantee (via our bias bounds) that the resulting solution is also an $\epsilon$ minimizer of the original objective $\mathcal{L}$.

where the bound (11) holds under Assumption A1.

To show that $\mathcal{L}(x;P_{0})\geq\overline{\mathcal{L}}(x;n)$ for every loss of the form (5), we use Lagrange duality to write

Showing the bound (10) requires a fairly different argument. Our proof uses the dual representation of $\mathcal{L}_{\chi^{2}\textup{-pen}}$ as a minimum of an expected risk over a Lagrange multiplier $\eta$ imposing the constraint that $q$ in (6) sums to $1$ (or that $Q$ in (5) integrates to $1$). Using convexity with respect to $\eta$ we relate the value of the risk at $\eta_{n}$ (the minimizer for sample $S_{1}^{n}$) to $\eta^{\star}$ (the population minimizer), which on expectation are $\overline{\mathcal{L}}_{\chi^{2}\textup{-pen}}$ and $\mathcal{L}_{\chi^{2}\textup{-pen}}$, respectively. We then apply Cauchy-Schwartz and bound the variance of $\eta_{n}$ with the Efron-Stein inequality to obtain a $1/n$ bias bound. ∎

2 Variance analysis

With the bias bounds in Proposition 5 established, we analyze the variance of the stochastic gradient estimators $\nabla\mathcal{L}(x;S_{1}^{n})$. More specifically, we prove that the variance of the mini-batch gradient estimator decreases as $1/n$ for penalty-type robust objectives (with $\phi=0$) for which the maximizing $Q$ has bounded $\chi^{2}$ divergence from $P_{0}$, which we call “$\chi^{2}$-bounded objectives” (see Section A.4). Noting that $\mathcal{L}_{\textup{kl-CVaR}}$ (with $\mathcal{L}_{\textup{CVaR}}$ as a special case) and $\mathcal{L}_{\chi^{2}\textup{-pen}}$ are $\chi^{2}$-bounded yields the following.

(Note that the variance bound on $\mathcal{L}_{\textup{kl-CVaR}}$ is independent of $\lambda$ and therefore holds also for $\mathcal{L}_{\textup{CVaR}}$ where $\lambda=0$).

3 Complexity guarantees

We also consider Nesterov’s accelerated gradient method . For $x_{0}=y_{0}=z_{0}\in\mathcal{X}$, a fixed step-size $\eta>0$ and a sequence $\{\theta_{t}\}$, we iterate

We now state the rates of convergence of the iterations (12) and (13) following the analysis in , with a small variation where the stochastic gradient estimates are unbiased for a uniform approximation of the true objective with additive error $\delta$. We provide a short proof in Appendix B.4.

Since our gradient estimator has norm bounded by $G$, SGM allows us to find an $\epsilon$-minimizer of $\overline{\mathcal{L}}$ in $T\asymp{(GR)^{2}}/{\epsilon^{2}}$ steps. Therefore, choosing $n$ large enough in accordance to Proposition 1 guarantees that we find an $\epsilon$-minimizer of $\mathcal{L}$. The accelerated scheme (13) admits convergence guarantees that scale with the gradient estimator variance instead of its second moment, allowing us to leverage Proposition 2 to reduce $T$ to the order of $1/\epsilon$. The accelerated guarantees require the loss $\mathcal{L}$ to have order $1/\epsilon$-Lipschitz gradients—fortunately, this holds for $\mathcal{L}_{\chi^{2}\textup{-pen}}$ and $\mathcal{L}_{\textup{kl-CVaR}}$.

Let Assumption A1 hold. For all $P$, $\nabla\mathcal{L}_{\textup{kl-CVaR}}(x;P)$ and $\nabla\mathcal{L}_{\chi^{2}\textup{-pen}}(x;P)$ are $(\frac{G^{2}}{\lambda}+H)$-Lipschitz in $x$, and $0\leq\mathcal{L}_{\textup{CVaR}}(x;P)-\mathcal{L}_{\textup{kl-CVaR}}(x;P)\leq\lambda\log({1}/{\alpha})$ for all $x$.

See proof in Appendix A.1.6. Thus, to minimize $\mathcal{L}_{\textup{CVaR}}$ we instead minimize $\mathcal{L}_{\textup{kl-CVaR}}$ and choose $\lambda\asymp\epsilon/\log({1}/{\alpha})$ to satisfy the smoothness requirement while incurring order $\epsilon$ approximation error. For $\mathcal{L}_{\chi^{2}\textup{-pen}}$ with $\lambda\geq\epsilon$ we get sufficient smoothness for free.We can also handle the case $\lambda<\epsilon$ by adding a KL-divergence term to $\psi$ for $\mathcal{L}_{\chi^{2}\textup{-pen}}$.

For $\mathcal{L}=\mathcal{L}_{\textup{CVaR}}$, we have nT\lesssim\frac{(GR)^{2}}{\alpha\epsilon^{2}}\left(1+\min\Big{\{}\frac{\alpha G_{\textup{icdf}}\sqrt{\log\frac{1}{\alpha}+\nu}}{GR},\frac{B^{2}\sqrt{\log\frac{1}{\alpha}+\nu}}{GR\epsilon},\frac{B^{2}}{\epsilon^{2}}\Big{\}}\right).

For $\mathcal{L}=\mathcal{L}_{\chi^{2}\textup{-pen}}$ with $\lambda\leq B$, we have nT\lesssim\frac{(GR)^{2}B}{\lambda\epsilon^{2}}\left(1+\min\Big{\{}\frac{B}{GR}\sqrt{\frac{\epsilon(1+\nu)}{\lambda}},\frac{B}{\epsilon}\Big{\}}\right).

For $\mathcal{L}=\mathcal{L}_{\chi^{2}}$, we have $nT\lesssim\frac{(1+\rho)(GR)^{2}B^{2}}{\epsilon^{4}}\log\frac{(1+\rho)B^{2}}{\epsilon^{2}}$.

For any loss of the from (5), we have $nT\lesssim\frac{(GR)^{2}G_{\textup{icdf}}}{\epsilon^{3}}$.

The smoothness parameter $H$ only appears in rates resulting from Nesterov acceleration. Even there, $H$ appears in lower-order terms in $\epsilon$ since $\nu=\frac{H}{G^{2}}\epsilon$. We also note that the final $G_{\textup{icdf}}\epsilon^{-3}$ rate holds even when the uncertainty set is the entire simplex; therefore, when $G_{\textup{icdf}}<\infty$ it is possible to approximately minimize the maximum loss in sublinear time. Theorem 1 achieves the claimed rates of convergence in Table 1 in certain settings. In particular, it recovers the rates for $\mathcal{L}_{\textup{CVaR}}$ and $\mathcal{L}_{\chi^{2}\textup{-pen}}$ (the first and last column of the table) when $\nu\lesssim 1$, $\lambda\gtrsim(B/(GR))^{2}\epsilon$, and $\alpha\lesssim GR/G_{\textup{icdf}}$. In the next section, we show how to attain the claimed optimal rates for $\mathcal{L}_{\textup{CVaR}}$ and $\mathcal{L}_{\chi^{2}\textup{-pen}}$ without conditions, returning to address the rates for the constrained $\chi^{2}$ objective $\mathcal{L}_{\chi^{2}}$ in Section 6.

Multi-level Monte Carlo (MLMC) gradient estimators

In the previous section, we optimized the mini-batch surrogate $\overline{\mathcal{L}}(x;n)$ to the risk $\mathcal{L}(x;P_{0})$, using Proposition 1 to guarantee the surrogate’s fidelity for sufficiently large $n$. The increasing (linear) complexity of computing the estimator $\nabla\mathcal{L}(x;S_{1}^{n})$ as $n$ grows limits the (theoretical) efficiency of the method. To that end, in this section we revisit a multi-level Monte Carlo (MLMC) gradient estimator of Blanchet and Glynn to form an unbiased approximation to $\nabla\overline{\mathcal{L}}(x;n)$ whose sample complexity is logarithmic in $n$. We provide new bounds on the variance of this MLMC estimator, leading immediately to improved (and, as we shall see, optimal) efficiency estimates for stochastic gradient methods using it.

For a given minimum sample size parameter $n_{0}\geq 1$, we define $\widehat{\mathcal{M}}[\nabla\mathcal{L}]$, the MLMC estimator of $\nabla\mathcal{L}$, via

Our estimator differs from the proposal in two aspects: the distribution of $J$ and the option to set $n_{0}>1$. As we further discuss in Appendix C.3, the former difference is crucial for our setting, while the latter is pratically and theoretically helpful yet not crucial. The following properties of the MLMC estimator are key to our analysis (see Appendix C.1 for proofs).

The estimator $\widehat{\mathcal{M}}[\nabla\mathcal{L}]$ with parameters $n=2^{j_{\max}}n_{0}$ satisfies

For all $x\in\mathcal{X}$, the multi-level Monte Carlo estimator with parameters $n$ and $n_{0}$ satisfies

Further paralleling Proposition 2, we obtain similar bounds on the MLMC estimates of $\mathcal{L}_{\textup{CVaR}}$ and $\mathcal{L}_{\chi^{2}\textup{-pen}}$ (in addition to their gradients), and demonstrate that similar bounds fail to hold for $\nabla\mathcal{L}_{\chi^{2}}$ (Proposition 7 in Appendix C.1). Therefore, directly using the MLMC estimator on $\nabla\mathcal{L}_{\chi^{2}}$ cannot provide guarantees for minimizing $\mathcal{L}_{\chi^{2}}$; instead, in Section 6 we develop a doubling scheme that minimizes the dual objective $\mathcal{L}_{\chi^{2}\textup{-pen}}(x;P_{0})+\lambda\rho$ jointly over $x$ and $\lambda$. This scheme relies on MLMC estimators for both the gradient $\nabla\mathcal{L}_{\chi^{2}\textup{-pen}}$ and the derivative of $\mathcal{L}_{\chi^{2}\textup{-pen}}$ with respect to $\lambda$.

Proposition 4 guarantees that the second moment of our gradient estimators remain bounded by a quantity that depends logarithmically on $n$. For these estimators, Proposition 3 thus directly provides complexity guarantees to minimize $\mathcal{L}_{\textup{CVaR}}$ and $\mathcal{L}_{\chi^{2}\textup{-pen}}$. We also provide a high probability bound on the total complexity of the algorithm using a one-sided Bernstein concentration bound. We state the guarantee below and present a short proof in Appendix C.2.

The same conclusion holds when replacing $\mathcal{L}_{\textup{CVaR}}$ with $\mathcal{L}_{\chi^{2}\textup{-pen}}$ and $\alpha^{-1}$ with $1+B/\lambda$.

Lower bounds

The remaining technical contribution in the paper is to revisit the constrained $\chi^{2}$ objective (3), which is resistant to many of the techniques we have thus far developed. In this section, we leverage duality relationships to approximate the constrained objective (3) via its penalized counterpart (4), $\mathcal{L}_{\chi^{2}\textup{-pen}}$. We adjust notation to make the dependence of $\mathcal{L}_{\chi^{2}\textup{-pen}}^{\lambda}$ on $\lambda$ explicit, and defer all proofs to Appendix E.

Our starting point is the recognition that, by duality (cf. [59, Sec. 3.2]),

for any distribution $P_{0}$. For $0\leq\underline{\lambda}\leq\overline{\lambda}$, we may thus consider the approximation

By restricting $\lambda$ to an appropriate range, we can then approximate $\mathcal{L}_{\chi^{2}}$ by its truncated version, as the next lemma shows.

Our strategy is therefore to jointly minimize $f_{\rho}(x,\lambda)=\mathcal{L}_{\chi^{2}\textup{-pen}}^{\lambda}(x;P_{0})+\lambda\rho$ over both $x\in\mathcal{X}$ and $\lambda\in[\underline{\lambda},\overline{\lambda}]$ (rather than $[0,\infty]$), using the approximation guarantee in Lemma 1 to argue that the restriction of $\lambda$ will have limited effect on the quality of the resulting solution. We iterate the projected stochastic gradient method with the multi-level Monte Carlo (MLMC) gradient estimator (16) via

If we can bound the moments of the MLMC-approximated gradients $\widehat{\mathcal{M}}$, we can then leverage standard stochastic gradient analyses to prove convergence. We use the following bound.

Therefore, we may find an $\epsilon$ approximate minimizer with complexity roughly $B^{3}\overline{\lambda}^{2}/(\underline{\lambda}^{3}\epsilon^{2})$:

Directly substituting $\underline{\lambda}=\frac{\epsilon}{2\rho}$ and $\overline{\lambda}=\frac{B}{\rho}$ results in a guarantee scaling as $\epsilon^{-5}$, which is worse than the mini-batch rate of $\epsilon^{-4}$. To improve on this, we divide $[\frac{\epsilon}{2\rho},\frac{B}{\rho}]$ into $K=\log_{2}\frac{B}{\epsilon}$ sub-intervals $[\lambda^{(i+1)},\lambda^{(i)}]$ satisfying $\lambda^{(i+1)}/\lambda^{(i)}=2$. We then perform the stochastic gradient method (17) on each of these intervals $[\lambda^{(i+1)},\lambda^{(i)}]$ in turn, yielding estimates $\bar{x}^{(i)}$ that are each $\lesssim\epsilon$-suboptimal for the approximate objective $\mathcal{L}_{\chi^{2}[\lambda^{(i+1)},\lambda^{(i)}]}$. Using the bounded ratio $\lambda^{(i+1)}/\lambda^{(i)}=2$, this requires complexity roughly $1/(\lambda^{(i+1)}\epsilon^{2})\lesssim\rho/\epsilon^{3}$, giving the following theorem.

Theorem 4 provides a rigorous guarantee on the complexity of minimizing $\mathcal{L}_{\chi^{2}}$ with a fixed constraint $\rho$ by optimizing the parameter $\lambda$ of $\mathcal{L}_{\chi^{2}\textup{-pen}}^{\lambda}$. In practice, we usually have no prior knowledge of $\rho$, so it will often make sense to directly tune $\lambda$ according to validation criteria rather than a target $\rho$. We also note that Duchi and Namkoong prove a lower bound of order ${\rho}{\epsilon^{-2}}$, which is smaller than our ${\rho}{\epsilon^{-3}}$ rate. Establishing the optimal rate for this problem remains an open question.

Experiments

We test our theoretical predictions with experiments on two datasets. Our main focus is measuring how the total work in solving the DRO problems depends on different gradient estimators. In particular, we quantify the tradeoffs in choosing the mini-batch size $n$ in the estimator $\nabla\mathcal{L}(x;S_{1}^{n})$ of Section 3 and the effect of using the MLMC technique of Section 4. To ensure that we operate in practically meaningful settings, our experiments involve heterogeneous data, and we tune the DRO objective to improve the generalization performance of ERM on the hardest subpopulation. We provide a full account of experiments in Appendix F and summarize them below.

Figure 1 plots the training objective as optimization progresses. In Appendix F.4 we provide expanded figures that also report the robust generalization performance. We find that the benefits of DRO manifest mainly when the metric of interest is continuous (e.g., log loss) as opposed to the 0-1 loss.

Our analysis in Section 3.1 bounds the suboptimality of solutions resulting from using a mini-batch estimators with batch size $n$, showing it must vanish as $n$ increases. Figure 1 shows that smaller batch sizes indeed converge to suboptimal solutions, and that their suboptimality becomes negligible very quickly: essentially every batch size larger than $10$ provides fairly small bias (with the exception of $\mathcal{L}_{\chi^{2}}$ in the digits experiment). The effect of bias is particularly weak for $\mathcal{L}_{\chi^{2}\textup{-pen}}$, consistent with its superior theoretical guarantees. We note, however, that the suboptimality we see in practice is far smaller than the worst-case bounds in Proposition 1. We investigate this in Appendix F.5, where we show that the bias $\mathcal{L}-\overline{\mathcal{L}}$ is in fact consistent with our theory, but the minimizers of $\mathcal{L}$ and $\overline{\mathcal{L}}$ are more similar than expected a priori.

While the MLMC estimator does not suffer from a bias floor (by design), it is also much slower to converge. This may appear confusing, since the MLMC convergence guarantees are optimal (for $\mathcal{L}_{\textup{CVaR}}$ and $\mathcal{L}_{\chi^{2}\textup{-pen}}$) while the mini-batch estimator achieves the optimal rate only under certain assumptions. Recall, however, that these assumptions are smoothness of the loss (which holds in our experiments) and—for CVaR—sufficiently rapid decay of the bias floor, which we verify empirically.

For batch sizes in the range 50–5K, the traces in Figure 1 look remarkably similar. This is consistent with our theoretical analysis for $\mathcal{L}_{\textup{CVaR}}$ and $\mathcal{L}_{\chi^{2}\textup{-pen}}$, which shows that the variance decreases linearly with the batch size and we may therefore (with Nesterov acceleration) increase the step size proportionally and expect the total work to remain constant. As theory predicts, this learning rate increase is only possible up to a certain batch size (roughly 5K in our experiments), after which larger batches become less efficient. Indeed, to reach within 2% of the optimal value, the full-batch method requires 27–36$\times$ more work than batch sizes 50–5K for ImageNet, and 9–16$\times$ more work for the digits experiment (see Table 5 and 6 for a precise breakdown of the number of epochs required per algorithm for each robust objective).

We also repeat our experiments with the dual SGM and prima-dual methods mentioned in Table 1 and compare them with them our proposed method; see Appendix F.6 for details.

Conclusion

This work provides rigorous convergence guarantees for solving large-scale convex $\phi$-divergence DRO problems with stochastic gradient methods, laying out a foundation for their use in practice; we conclude it by highlighting two directions for further research.

First, while our work resolves the optimal theoretical convergence rates for CVaR and $\chi^{2}$ penalty objectives, the corresponding result for $\chi^{2}$ constraint remains open. In particular, there is a gap between our $O(\rho\epsilon^{-3})$ upper and the $\Omega(\rho\epsilon^{-2})$ lower bound of Duchi and Namkoong . Moreover, combining the uniform convergence results in Duchi and Namkoong with a cutting plane method gives complexity guarantees scaling a roughly as $\rho^{2}d^{2}\epsilon^{-2}$, so the $O(\rho\epsilon^{-3})$ rate can only be optimal in high-dimensional settings.

The authors would like to thank Hongseok Namkoong for discussions and insights, as well as Nimit Sohoni for comments on an earlier draft. DL, YC and JCD were supported by the NSF under CAREER Award CCF-1553086 and HDR 1934578 (the Stanford Data Science Collaboratory) and Office of Naval Research YIP Award N00014-19-2288. YC was supported by the Stanford Graduate Fellowship. AS is supported by a Microsoft Research Faculty Fellowship, NSF CAREER Award CCF-1844855, NSF Grant CCF-1955039, a PayPal research gift, and a Sloan Research Fellowship.

References

Appendix

In this section we collect several basic results which we use in subsequent derivations in the paper: Section A.1 gives several additional characterization of the robust objective $\mathcal{L}$, Section A.2 briefly discusses the computation of $\mathcal{L}$ and its costs, Section A.3 gives a short derivation of the complexity guarantees for “dual SGM” in Table 1, and Section A.4 introduces the notion of losses contained in a $\chi^{2}$ divergence ball. Finally, Section A.5 lists a few standard probabilistic bounds.

Here we give several equivalent characterizations of the robust objective

For $\widehat{P}[s_{1}^{n}]$ uniform on $s_{1},s_{2},\ldots,s_{n}$ (which we abbreviate $s_{1}^{n}$), we write

We can convert the maximization over $r$ in Eq. (20) (or $Q$ in (18)) with minimization over Lagrange multipliers for the constraint that $r$ sums to 1 and the $\phi$-divergence constraint, yielding

where the expectation is over $S\sim P$, i.e. the distribution from which we observe samples. On a finite sample $s_{1}^{n}$ we have

For pure-constraint objectives (with $\psi=0$), $\Upsilon$ simplifies to

For pure-penalty objective (with $\phi=0$) the Lagrange multiplier $\nu$ is unnecessary and we have

Note that $\Upsilon$ is an expectation (i.e., an empirical risk) which means that to minimize $\mathcal{L}(x;P)$ we can, in principle, apply ERM jointly on $x,\eta$ and $\nu$, as we further discuss in Section A.3.

Finally, we note that any $Q^{\star}$ attaining the supremum in (18) is of the form

where $\eta^{\star}$ and $\nu^{\star}$ are optimal Lagrange multipliers in (22) and ${(\nu^{\star}+\lambda\psi)^{*}}^{\prime}$ is a subderivative of $(\nu^{\star}+\lambda\psi)^{*}$. For $\phi=0$ this specializes to

We note that this last expression is a direct consequence of (20), since $\mathcal{R}$ is the set of measures never exceeding $\frac{1}{\alpha}$. On a finite sample $s_{1}^{n}$ this gives the closed-form expression

In words, $\mathcal{L}_{\chi^{2}\textup{-pen}}$ is a sum of a CVaR (at level $F(\eta^{\star})$), a conditional variance regularization term and an outage probability regularization term. This expression simplifies considerably when $\lambda$ is sufficiently large. Specifically, we have,

That is, for sufficiently large $\lambda$ the objective $\mathcal{L}_{\chi^{2}\textup{-pen}}$ is simply the empirical risk with variance regularization (see also ).

An expression for $\mathcal{L}_{\chi^{2}}$ follows via (29)

Similarly, for a sample of size $n$ and a maximizing $q^{\star}$, we have

Therefore, if $Q^{\star}_{x}$ is Lipschitz w.r.t. $x$ in the 1-norm then $\nabla\mathcal{L}$ is Lipschitz as well. This is indeed the case $\mathcal{L}_{\chi^{2}\textup{-pen}}$ and $\mathcal{L}_{\textup{kl-CVaR}}$.

Since entropy is 1-strongly-convex w.r.t. the 1-norm, for $\mathcal{L}_{\textup{kl-CVaR}}$ we have that the penalty $\lambda\psi$ is $\lambda$-strongly-convex w.r.t. the 1-norm and therefore [56, Lemma 2]

establishing that $\nabla\mathcal{L}_{\chi^{2}\textup{-pen}}$ is also $(H+G^{2}/\lambda)$-Lipschitz.

A.2 Computational cost

Since we are interested in large-scale application, we assume that $d\gg\log(nB/\varepsilon)$ and therefore the time to compute the objective and its gradient is $O(nd)$.

For simplicity and stability, our code implements the computation of $q^{\star}$ using bisection over $\eta$ for each of $\mathcal{L}_{\chi^{2}\textup{-pen}},\mathcal{L}_{\chi^{2}}$ and $\mathcal{L}_{\textup{kl-CVaR}}$.

A.3 Stochastic gradient method on the dual objective

Here we discuss the convergence guarantees for a simple stochastic gradient method using the dual expression (22) for $\mathcal{L}(x;P_{0})$ in order to minimize it over $x$. While several works consider such methods (see Section 1.1), we could not find direct reference for their runtime guarantees, and we therefore briefly derive it below.

Focusing on objectives with $\phi=0$ (as in (25)), and writing $\gamma_{x}$ and $\gamma_{\eta}$ for step sizes, we write the iterations on $x$ and the Lagrange multiplier $\eta$ as

Where $S_{1},S_{2},\ldots$ are drawn iid from $P_{0}$.

Let $\epsilon\in(0,B)$. For CVaR and a suitable choice of $\gamma_{x},\gamma_{\eta}$ the average iterate $\bar{x}_{T}=\frac{1}{T}\sum_{t\leq T}x_{t}$ satisfies

Similarly, for $\chi^{2}$ penalty we have

By Proposition 3, the expected sub-optimality of $\bar{x}_{T}$ is $\lesssim(\Gamma_{x}R+\Gamma_{\eta}(\overline{\eta}-\underline{\eta})/\sqrt{T}$, where $\Gamma_{x}^{2}$ (respectively $\Gamma_{\eta}$) is an upper bound on the second moment of $\nabla\Upsilon(x,\eta;S)$ (respectively $\frac{\partial}{\partial\eta}\Upsilon(x,\eta;S)$). For $\mathcal{L}_{\textup{CVaR}}$ we have $\Gamma_{x}\leq G/\alpha$, $\Gamma_{\eta}\leq 1/\alpha$ and $\overline{\eta}-\underline{\eta}=B$. For $\mathcal{L}_{\chi^{2}\textup{-pen}}$ we have $\Gamma_{x}\leq G(1+B/\lambda)$, $\Gamma_{\eta}=1+B/\lambda$ and $\overline{\eta}-\underline{\eta}=B+\lambda$. The result follows from substituting $T\asymp\left(\Gamma_{x}^{2}R^{2}+\Gamma_{\eta}^{2}(\overline{\eta}-\underline{\eta})^{2}\right)\epsilon^{-2}$. ∎

A number of our results hold for general subclass of the objective (18) with the following property.

The three objectives we focus on are $\chi^{2}$-bounded.

The objectives $\mathcal{L}_{\textup{kl-CVaR}},\mathcal{L}_{\chi^{2}}$ and $\mathcal{L}_{\chi^{2}\textup{-pen}}$ are $\chi^{2}$-bounded with constants $C=\frac{1}{\alpha}-1$, $C=\rho$ and $C=B/\lambda$, respectively.

That $\mathcal{L}_{\chi^{2}}$ is $\rho$-$\chi^{2}$-bounded is obvious from definition. For $\mathcal{L}_{\chi^{2}\textup{-pen}}$ we have

A.5 General results

We conclude this section of the appendix by stating three general results that aid our analysis. First, we give a lemma stating that a binomial random variable with parameters $n$ and $\alpha$ has a constant probability of being at least $\sqrt{\alpha(1-\alpha)n}$ below its mean.

Second, we state the Efron-Stein inequality in vector form, which follows from applying the standard scalar bound element-wise.

Third, we give a general lemma on the variance of sampling without replacement, which we specialize to the simplex for later use.

Let $p\in\Delta^{k}$ and let $\mathcal{I}$ be a random subset of $[k]$ of size $k/2$. Then

Let us denote $q=p-\tfrac{1}{k}\boldsymbol{1}$. We have

Appendix B Proofs from Section 3

This section completes the proof and discussion of the results in Section 3. First, in Section B.1, we prove the bias bounds in Proposition 1 and argue their tightness in the worst case. Section B.2 provides additional discussion of the smoothness and Lipschitz inverse-cdf assumptions sometimes used in this section. Then, in Section B.3 we bound the variance of the mini-batch estimators for $\chi^{2}$-bounded penalty objectives and their gradient, obtaining Proposition 2 as a corollary. We also argue that similar bounds do not hold for the $\chi^{2}$ constraint objective. In Section B.4 we review the standard convergence guarantees for stochastic gradients iterations with and without Nesterov acceleration, and in Section B.5 we combine all these ingredients to prove Theorem 1.

We first show that the bound $\mathcal{L}\geq\overline{\mathcal{L}}$ holds for any loss of the form (18) and then proceed to show each of the bounds (8)–(11). We remark here that the bound (9) actually holds for any $\rho$-$\chi^{2}$-bounded objective (Definition 1).

where the inequality follows from exchanging the expectation and the infimum.

where $f_{\mathsf{Beta}(a,b)}$ is the density function of the Beta random variable of parameters $a,b$. Substituting back, we have

Recalling Eq. (27) for $\mathcal{L}_{\textup{CVaR}}(x;P_{0})$, and recalling that $F^{-1}(u)\in[0,B]$ for all $u$ by assumption, we bound the bias as

To conclude, it suffices to bound the tail probability of the Beta random variables. We have [see, e.g., 50, Ex. 5 in Sec. 4.6]

and the multiplicative Chernoff bound [40, Theorem 4.3] gives

Substituting into (40) and using $\mathcal{L}_{\textup{CVaR}}(x;P_{0})\leq B$ when $\alpha n\leq 9$ gives the final bound

We start with the expression (20) specialized for the $\mathcal{L}_{\chi^{2}}$,

The restriction of $\mathcal{R}$ to non-increasing functions is “free” since $F^{-1}$ is non-decreasing. Our strategy is to relate $F^{-1}$ to CVaR and then apply the corresponding bias bounds (8)—this type of transformation is closely related to the Kusuoka representation of coherent risk measures . Specifically, note that

Therefore, for any $r\in\mathcal{R}$ integration by parts gives

Taking a supremum over $r\in\mathcal{R}$, we conclude that

It remains to bound the quantity $E$, which we do via the the Cauchy-Schwarz inequality and the definition of $\mathcal{R}$, which gives

for all $r\in\mathcal{R}$, giving the required bound.

The bound (42) hold for any loss (18) and not just $\mathcal{L}_{\chi^{2}}$. Moreover, the final bound using Cauchy-Schwarz is equally valid for any $\rho$-$\chi^{2}$-bounded uncertainty set. In particular, consider the Cressie-Read uncertainty sets corresponding to $k$-norm the constraint $\|r\|_{k}^{2}\leq 1+2\rho$. For $k>2$ they satisfy $\|r\|_{2}^{2}\leq 1+2\rho$ and our bias bounds holds (using Hölder’s inequality instead of Cauchy-Schwarz removes the logarithmic factor). For $k\in(1,2)$ Hölder’s inequality gives bounds decaying as $n^{-(k-1)/k}$.

Starting with CVaR, we return to the expression (38) for the bias and note that

holds for all $u$, because when $u<1-\alpha$ we have that $1_{\{u\geq 1-\alpha\}}-\mathcal{I}_{\alpha}(u)\leq 0$ and so we increase the LHS by replacing $F^{-1}$ with an under-estimate, while for $u\geq 1-\alpha$ we have $1-\mathcal{I}_{\alpha}(u)\geq 0$ due to (39) and we increase the LHS be replacing it with an $F^{-1}$ with an over-estimate. Substituting into (38) and calculating gives

Above, $(i)$ uses the fact that $\frac{1}{\alpha}\mathcal{I}_{\alpha}$ is a convex combination of densities to deduce that

and $\eta_{n}$ is the unique solution to $\frac{1}{n}\sum_{i=1}^{n}(Z_{i}-\eta)_{+}=\lambda$. Convexity of $\Upsilon$ w.r.t. $\eta$ gives us

To handle the second variance we use the Efron-Stein inequality (Lemma 5). Let $I$ be uniformly distributed on $[n]$, and define

Combining (47), (45) and (44) gives the result for $\lambda\leq B$.

In the edge case that $\lambda\geq B$, Eq. (A.1.4) gives us that

We note that in this case may easily form an unbiased estimator of $\mathcal{L}_{\chi^{2}\textup{-pen}}$ by using the standard unbiased variance estimator. ∎

As before, we treat each case separately.

where the final bound follows from the Berry-Esseen theorem (see Lemma 4).

The $\chi^{2}$ divergence between two Bernoulli random variables is

Therefore, for any $p\in(0,1)$, the element in $\mathcal{U}^{\rho}_{\chi^{2}}(\mathsf{Bernoulli}(p))$ that maximizes $Q(1)$ is $Q=\mathsf{Bernoulli}(q)$ with $q=\min\left\{1,p+\sqrt{2\rho}\sqrt{p(1-p)}\right\}$. Set $p_{0}=\frac{1}{1+2\rho}$ and note that the function $f(p)=p+\sqrt{2\rho}\sqrt{p(1-p)}=p+\sqrt{(1-p_{0})/p_{0}}\sqrt{p(1-p)}$ satisfies $f(p_{0})=1$ and $f^{\prime}(p)\geq\frac{1}{2p_{0}}$ for all $p\leq p_{0}$. Therefore, we have

for all $p\in(0,1)$, with equality at $p=p_{0}$. In particular, setting $P_{0}=\mathsf{Bernoulli}(p_{0})$ implies $\mathcal{L}_{\chi^{2}}(x;P_{0})=B$ and for a sample $S_{1}^{n}\sim P_{0}^{n}$ with $\hat{p}=\frac{1}{n}\sum_{i\in[n]}S_{i}$ we have $\mathcal{L}_{\chi^{2}}(x;S_{1}^{n})\leq B\left(1-\frac{p_{0}-\hat{p}}{2p_{0}}\right)$. Therefore

where $(\star)$ follows from the CVaR case and for the final equality we substitute the definition of $p_{0}$.

with equality at $p=\lambda/B$. Taking $p_{0}=\lambda/B$ and and for a sample $S_{1}^{n}\sim P_{0}^{n}$ letting $\hat{p}=\frac{1}{n}\sum_{i\in[n]}S_{i}$, we have

by Berry-Esseen and Chernoff, and the result follows by substituting $p_{0}=\lambda/B$. ∎

B.2 Discussion of additional assumptions

The guarantees for the accelerated gradient iterations (13), detailed in Appendix B.4, require the objective function be smooth, i.e., have Lipschitz gradient. However, the degree of smoothness need not be high: as Nesterov and subsequent work observed, even if $\nabla\mathcal{L}$ is order $G^{2}/\epsilon$ Lipschitz, acceleration allows finding an $\epsilon$ accurate solution in roughly $GR/\epsilon$ steps (a quadratic improvement over the SGM rate), as long as the gradient variance is itself of order $\epsilon$; the accelerated rates in Theorem 1 stem from this fact.

B.3 Proofs of variance bounds

We give a more general statement of the variance bound using the notion of $C$-$\chi^{2}$-bounded objectives (Definition 1); Proposition 2 follows immediately from Claim 4.

If in addition $\phi=0$ and $\psi$ is strictly convex, we have

We first show the bound on the objective variance. By the the Efron-Stein inequality (see Lemma 5), we have

Next, to show the bound on the gradient variance we invoke Efron-Stein elementwise to obtain

By the expression (34) for $\nabla\mathcal{L}$ we have

The remainder of the proof is identical to that of the objective variance bound, except with $2G$ replacing $B$. ∎

(Note that we may assume without loss of generality that $\rho\gtrsim 1/n$, so that $P(S=0)=1-p_{1}-p_{2}>0$, since for $\rho=0$ we already have a standard $1/n$ lower bound on the variance). We set the loss values to be

The source of high variance in this construction is that, for a sample $S_{1}^{n}$, the maximizing $q^{\star}$ behaves very differently when $S_{i}=2$ for some $i$ and when $S_{i}\neq 2$ for all $i$. In the former case, we show that $q^{\star}$ puts significant mass on samples with $S_{i}\neq 1$, so $\nabla\mathcal{L}(0;S_{1}^{n})<G(1-c)$ for some $c\gtrsim\rho/(1+\rho)$. In the latter case, we show that with constant probability $q^{\star}$ places mass only on samples with $S_{i}=1$, and so $\nabla\mathcal{L}(0;S_{1}^{n})=G$. Since either scenario occurs with constant probability, the variance bound follows.

We bound $\nabla\mathcal{L}_{\chi^{2}}$ conditional on each event in turn.

To bound the gradient under event $\mathfrak{E}_{b}(S_{1}^{n})$, assume that without loss of generality that $S_{1}=2$ is the unique sample with that value. We consider separately the cases $q^{\star}_{1}>2/3$ and $q^{\star}_{1}\leq 2/3$. In the former, we clearly have $\nabla\mathcal{L}_{\chi^{2}}(0;S_{1}^{n})\leq-q^{\star}_{1}G+(1-q^{\star}_{1})G<-G/3$. In the latter case, we recall Eq. (33) showing that $q^{\star}$ is of the form

Assuming that $S_{1}=2$, we have that the total weight under $q^{\star}$ of samples with gradient $-G$ is

which implies $\nabla\mathcal{L}_{\chi^{2}}(0;S_{1}^{n})\leq-\frac{7\rho}{4(1+2\rho)}G+(1-\frac{7\rho}{4(1+2\rho)})G\leq G(1-\frac{7\rho}{2(1+2\rho)})$. We conclude that

B.4 Convergence rates of stochastic gradient methods

We state below the classical convergence rates for standard and accelerated stochastic gradient methods, under a somewhat non-standard assumption that the stochastic gradient estimates are unbiased for a uniform approximation of the objective function with additive error $\delta$.

gives us the rates (14) and (15) but for $\overline{F}$ rather than $F$. That is, it guarantees that SGM finds $\bar{x}_{T}$ such that

To remove the bars, we use $0\leq F(x)-\overline{F}(x)\leq\delta$ to write

B.5 Proofs of complexity bounds

To prove each bound in the theorem we choose $n$ large enough via one of the bounds in Proposition 1 and then choose $T$ to guarantee $\epsilon$-accurate solution via Proposition 3. For a (potentially random) point $\bar{x}\in\mathcal{X}$ and robust risk $\mathcal{L}$, we define the shorthand

We summarize our choices of $n$ and $T$ for different robust objectives, under different assumptions in Table 3. In the statement of the theorem, we sometimes upper bound $a\vee b\coloneqq\max\{a,b\}$ by $a+b$ for readability, and state the tighter rates here.

In that case, setting $n\asymp\frac{B^{2}}{\alpha\epsilon^{2}}$ guarantees that the bias is smaller than $\epsilon$ and setting $T\asymp\frac{(GR)^{2}}{\epsilon^{2}}$ yields that $\mathsf{err}(\bar{x}_{T};\mathcal{L}_{\textup{CVaR}})\lesssim\epsilon$.

$\overline{\mathcal{L}}_{\textup{kl-CVaR}}$ being $(\tfrac{G^{2}\log(1/\alpha)}{\epsilon}+H)$-smooth, the final iterate of the sequence (13) achieves

To make sure that the second and third terms are smaller than $\epsilon$, we set $T=\frac{(GR)^{2}}{\alpha n\epsilon^{2}}\vee\frac{GR}{\epsilon}\sqrt{\log\tfrac{1}{\alpha}+\nu}$. To guarantee small bias, we set $n\asymp\frac{B^{2}}{\alpha\epsilon^{2}}$; the resulting complexity is

We once again set $T=\frac{(GR)^{2}}{\alpha n\epsilon^{2}}\vee\frac{GR}{\epsilon}\sqrt{\log\tfrac{1}{\alpha}+\nu}$, and choosing $n\asymp\frac{G_{\textup{icdf}}}{\epsilon}$ yields the result.

and setting $n\asymp\frac{B^{2}}{\lambda\epsilon}$ and $T\asymp\frac{(GR)^{2}}{\epsilon^{2}}$ yields the fist rate.

First, noting that $\lambda\geq\epsilon$ guarantees that $R^{2}(\frac{G^{2}}{\lambda}+H)\leq\frac{(GR)^{2}(1+\nu)}{\epsilon}$. Furthermore, we simplify the variance term since $B/\lambda\geq 1$. We thus set $T\asymp\frac{GR}{\epsilon}\sqrt{1+\nu}\vee\frac{(GR)^{2}B}{\lambda n\epsilon^{2}}$ and choose $n\asymp\frac{B^{2}}{\lambda\epsilon^{2}}$. This yields the final result

This case is straightforward—without any bound on the variance in the worst-case, we turn to the basic SGM guarantee (14); we have

We set $T\asymp\frac{(GR)^{2}}{\epsilon^{2}}$ and $n\asymp\frac{(1+2\rho)B^{2}}{\epsilon^{2}}\log((1+2\rho)B^{2}\epsilon^{-2})$. We then have

The sequence of iterates (12) yield error

and setting $n\asymp\frac{G_{\textup{icdf}}}{\epsilon},T\asymp\frac{(GR)^{2}}{\epsilon^{2}}$ concludes the proof of the theorem. ∎

Appendix C Proofs of Section 4

We now provide additional discussion of the multilevel Monte Carlo estimator for general functions $\mathsf{F}$, whose form we restate here for convenience

Section C.1 provides upper bounds on the moments of $\widehat{\mathcal{M}}$ for estimating $\mathcal{L}_{\textup{kl-CVaR}}$, $\mathcal{L}_{\chi^{2}\textup{-pen}}$, and their gradients, proving Claim 2 and Proposition 4. In that section we also prove that similar second moment bounds do not always hold for $\nabla\mathcal{L}_{\chi^{2}}$. In Section C.2 we prove the complexity guarantees in Theorem 2, and we conclude in Section C.3 with a comparison of some of our design choices to the original proposal of Blanchet and Glynn .

For any function $\mathsf{F}$, the estimator $\widehat{\mathcal{M}}[\mathsf{F}]$ with parameters $n=2^{j_{\max}}n_{0}$ satisfies

We have the following bound on the second moment of the estimator,

Let $\mathcal{L}$ be an objective of the form (18) with $\phi=0$ and strictly convex $\psi$. If $\mathcal{L}$ is $C$-$\chi^{2}$-bounded, we have that for all $x\in\mathcal{X}$, the multi-level Monte Carlo estimator with parameters $n$ and $n_{0}$ satisfies

where $(i)$ is due to Lemma 6 and $(ii)$ follows from the assumption that $\mathcal{L}$ is $C$-$\chi^{2}$-bounded. Substituting into (56), we have

This concludes the argument for the gradient.

Finally, we know that $q^{\prime}$ is optimal for $S_{1}^{k/2}$ and so

where $(i)$ follows from Lemma 6 and the $C$-$\chi^{2}$-boundedness of $\mathcal{L}$. Substituting into (56) yields the desired bound on $\widehat{\mathcal{M}}[\mathcal{L}]$. ∎

Having established the gradient estimator upper bounds for pure-penalty objectives, we demonstrate that similar bounds do not extend to the case of $\chi^{2}$ constraint.

We reuse the construction and notation in the proof of Proposition 6 and so do not repeat it. For a sample $S_{1}^{n}$, we consider the event $\mathfrak{E}_{a}(S_{1}^{n/2})$ where $S_{i}\neq 2$ for all $i\leq n/2$ and there are at least $np_{1}/2$ samples with value 1. We argue in the proof of Proposition 6 (Eq. (52)) that under this event we have

We also consider the event $\mathfrak{E}_{b}(S_{1}^{n})$ that there is exactly one sample with value 2 and less the $np_{1}$ samples with value 1. As per the proof of Proposition 6 (Eq. (53)) we have

Moreover, by the arguments in the proof of Proposition 6, we have

Therefore, since $\widehat{\mathcal{D}}_{n}=\nabla\mathcal{L}_{\chi^{2}}(0;S_{1}^{n})-\frac{1}{2}\nabla\mathcal{L}_{\chi^{2}}(0;S_{1}^{n/2})-\frac{1}{2}\nabla\mathcal{L}_{\chi^{2}}(0;S_{1}^{n/2})$, we have

Since the number $T$ of SGM iterations must be proportional to the second moment of the gradient estimator, Proposition 7 tells us that in the worst case we might have to set $T\asymp n(GR)^{2}/\epsilon^{2}$, in which case we might as well use a mini-batch estimator with batch size $n$ and run $(GR)^{2}/\epsilon^{2}$ SGM steps.

C.2 Proof of complexity bounds

The convergence guarantee of Proposition 3 and the second moment bound of Proposition 4 directly give that iterates of the form (12) with the MLMC gradient estimator guarantees a regret smaller than $\epsilon$ for $n\asymp\frac{B^{2}}{\alpha\epsilon^{2}}$, $1\lesssim n_{0}\lesssim\frac{\log n}{\alpha}$ and $T\asymp\frac{(GR)^{2}}{n_{0}\alpha\epsilon^{2}}\log^{2}{n}$. However, since the multilevel estimator randomizes the batch size, it remains to show that the number of samples concentrates below the claimed bound. Let $K_{t}=n_{0}2^{J_{t}}$ be the batch size at time $t$, and note that

Therefore, since $K_{1}^{T}$ are iid, a one-sided Bernstein bound [66, Prop. 2.14] implies that

Solving in $\delta$ for the RHS to be equal to $\frac{1}{n}$ yields $\delta=\frac{n\log n}{3}(1+\sqrt{1+216\tfrac{Tn_{0}}{n\log n}})$. We replace $n,n_{0}$ and $T$ by their values and conclude the proof. ∎

C.3 Comparison with Blanchet and Glynn [7]

There are two differences between our MLMC estimator and the proposal of Blanchet and Glynn . First, we take $J$ to be a truncated $\mathsf{Geo}(1/2)$ random variable while they suggest $J\sim\mathsf{Geo}(2^{-3/2})$ without truncation—as we further discuss below, this modification is crucial for ensuring a useful second moment bound in our setting. The second difference is that we allow for a minimum sample size $n_{0}>1$ as opposed to $n_{0}=1$ in . This modification is somewhat less important, as $n_{0}=1$ suffices for optimal gradient complexity, but choosing slightly larger $n_{0}$ is helpful in practice and can provably reduce the sequential depth of SGM by logarithmic factors.

Appendix D Lower bound proofs

This section proves our lower bounds, which we restate for ease of reference.

Since our proofs for CVaR and $\chi^{2}$ penalty are quite different, we present them separately in Theorems 3a and 3b, respectively.

To prove the CVaR lower bound we use the following standard Le Cam reduction from stochastic optimization to hypothesis testing.

[17, Chapter 5] Let $\mathcal{P}$ be a set of distributions and $P_{-1},P_{1}\in\mathcal{P}$ and define

Armed with Lemma 7, we state and prove the lower bound for CVaR.

Let us first assume that $\alpha\leq\frac{1}{2}$. For $\delta\leq\min\{\alpha,1-2\alpha\}$, $\mu>0$ and $v\in\{-1,1\}$ we consider the distributions $P_{v}$ such that for $S_{v}\sim P_{v}$ we have

It therefore suffices to compute $\mathcal{L}_{\textup{CVaR}}(\pm 1;P_{\pm 1})$. A quick calculation yields

We thus have a closed-form expression for the CVaR objective: for $P_{1}$ we have

which clearly attains its minimum at $x=0$ where it has value . Choosing $\mu$ such that

gives $\mathcal{L}_{\textup{CVaR}}(1;S_{-1})=-G\mu$ and

which attains its minimum at $x=R$ where it has value $-GR\mu$. We therefore have that

where that last transition used $\delta\leq\alpha$ and $\alpha\leq 1/2$.

and the result follows from substituting $n\leq\frac{(GR)^{2}}{2048\alpha\epsilon^{2}}$. When $\alpha\geq 1/2$ the result follows from the standard lower bound for stochastic convex optimization (e.g. [17, Thm. 5.2.10]). ∎

Computation of $\mathcal{L}_{\chi^{2}\textup{-pen}}(x;P_{\pm})$ for the CVaR lower bound construction (57) shows that the argument does not easily transfer to the penalized-$\chi^{2}$ objective because—as opposed to constrained-$\chi^{2}$ and CVaR—it is not positive homogeneous in $x$.

We construct the hard instance for $\mathcal{L}_{\chi^{2}\textup{-pen}}$ based on the standard hard instance for non-stochastic convex optimization, whose properties are as follows.

In other words, any “dimension-free” algorithm needs to interact $\Omega(\epsilon^{-2})$ times with the deterministic oracle to obtain an $\epsilon$-suboptimal point. With this result, we prove our lower bound for optimizing $\mathcal{L}_{\chi^{2}\textup{-pen}}$.

for $c>0$ independent of $G,R,\lambda$ and $\epsilon$.

(If $\lambda>GR$ the result follows from the standard $(GR)^{2}/\epsilon^{2}$ lower bound for convex optimization). Expressing the resulting objective $\mathcal{L}_{\chi^{2}\textup{-pen}}$ with the dual form (29) gives

since $\eta^{\star}=f(x)-GR\geq-GR$. We get that minimizing $\mathcal{L}_{\chi^{2}\textup{-pen}}$ is equivalent to optimizing $f$.

Appendix E Doubling schemes proofs

We now complete the proofs of the claims in Section 6.

When in addition we have $\lambda^{\star}\geq\epsilon/(2\rho)$ then clearly $\min_{x\in\mathcal{X}}\mathcal{L}_{\chi^{2}[\frac{\epsilon}{2\rho},\frac{B}{\rho}]}(x;P_{0})=\mathcal{L}_{\chi^{2}[0,\infty]}(x;P_{0})=\min_{x^{\prime}\in\mathcal{X}}\mathcal{L}_{\chi^{2}}(x^{\prime};P_{0})$. Otherwise, if $\lambda^{\star}<\epsilon/(2\rho)\eqqcolon\lambda_{\epsilon}$ we may write

Recall the definition (55) of the MLMC estimator of a general $\mathsf{F}$ and the expression (56) for its second moment. Suppose that $\mathsf{F}(\cdot)=\mathsf{F}_{1}(\cdot)+\mathsf{F}_{2}(\cdot)+c$, where $c$ is a constant. Then

We apply this observation to \widehat{\mathcal{M}}\big{[}\tfrac{\partial}{\partial\lambda}\mathcal{L}_{\chi^{2}\textup{-pen}}^{\lambda}(x;\cdot)+\rho\big{]} by noting that

By exactly the same argument that proves the gradient second moment bound in Proposition 4’. The result then follows by substituting into (60). ∎

We take $n\asymp\frac{B}{\underline{\lambda}\epsilon}$ to guarantee bias below $\epsilon/2$ by Proposition 1, and we take $n_{0}\asymp\frac{B}{\underline{\lambda}}\log n$ to guarantee that

Let $\bar{\lambda}_{T}=\sum_{t\leq T}\lambda_{t}$ be the average of the $\lambda$ iterates in (17). By appropriate choice of $\eta$ and $\eta^{\prime}$ we guarantee (via Proposition 3) that

By Lemma 3, finding an $\epsilon$ approximate solution in the interval $[\lambda^{(i+1)},\lambda^{(i)}]$ requires

gradient computations, where we have used $\lambda^{(i)}/\lambda^{(i+1)}\leq 2$, $\lambda^{(i)}\leq\frac{B}{\rho}$, and $\lambda^{(i+1)}\geq\frac{\epsilon}{2\rho}2^{K-i}$. Summing over $i$ (and applying a union bound) gives the claimed guarantee. Since the minimizer of $f_{\rho}(x,\lambda)$ over $x\in\mathcal{X}$ and $\lambda\in[\frac{\epsilon}{2\rho},\frac{B}{\rho}]$ is equivalent is identical to its minimizer in one of the intervals $[\lambda^{(i+1)},\lambda^{(i)}]$ for $i\leq K$, the result follows from Lemma 1. ∎

Appendix F Experiments

In this section we give a detailed description of our experiments. We begin with a description of the problems we study (Section F.1) followed by our hyperparameter settings (Section F.2) and brief remarks about our PyTorch implementation (Section F.3). Then, in Sections F.4 and F.5 we present and discuss our results in detail, including speed-up factors over full-batch optimization, a study of the generalization impacts of the DRO objective, and direct empirical evaluation of the bias $\mathcal{L}-\overline{\mathcal{L}}$ which we bound in Proposition 1.

The ImageNet dataset comprises of $1.2\cdot 10^{6}$ training images and $5\cdot 10^{4}$ test images with $1000$ different classes. We featurize the dataset using a pre-trained ResNet-50 (trained on ImageNet itself with an ERM objective). We use those features as the input to a linear classifier, with regularized multi-class logarithmic loss as in the previous experiment. As the robust generalization metric, we report the average loss and accuracy on the 10 classes with highest test loss.

F.2 Hyperparameter tuning

We tune our stepsizes with a coarse-to-fine strategy. More precisely, for each stepsize in $\{10^{i}\}_{-5\leq i\leq 0}$, we perform a single run of the experiment, and pick the best two stepsizes in terms of the final training value. For these two stepsizes, we evaluate $\frac{\eta}{2},2\eta$ and select the stepsize that gives the best value of the training loss. For this final stepsize, we repeat the experiments with $5$ different seeds (affecting weight initialization and mini batch samples but not the dataset structure) and report the minimum and maximum across seeds at each iteration. We select all the stepsizes in our experiments using this strategy, except for batch size $n=10$ in ImageNet where we extrapolated the stepsize from other batch sizes. Table 4 summarizes our step size choices—for batch sizes up to 5K we see a clear linear relationship between the batch size and optimal step size.

F.3 PyTorch Integration

Figure 2 illustrates our integration of DRO into PyTorch. Users simply define the robust loss they wish to use (in the example $\mathcal{L}_{\chi^{2}}$ with $\rho=1$) and feed the loss for the examples in the batch to the robust layer. While our current implementation only supports the robust objectives we analyze—namely, CVaR, KL-regularized CVaR, constrained-$\chi^{2}$ and penalized-$\chi^{2}$—it is easy to extend to other choices of $\phi$ and $\psi$.