Distributionally Robust Optimization via Ball Oracle Acceleration
Abstract
We develop and analyze algorithms for distributionally robust optimization (DRO) of convex losses.
In particular, we consider group-structured and bounded fdivergence uncertainty sets.
Our approach relies on an accelerated method that queries a ball optimization oracle, i.e., a subroutine that minimizes the objective within a small ball around the query point.
Our main contribution is efﬁcient implementations of this oracle for DRO objectives.
For DRO with N non-smooth loss functions, the resulting algorithms ﬁnd an ✏ -accurate solution with e
O
N ✏  2 / 3 + ✏  ﬁrst-order oracle queries to individual loss functions.
Compared to existing algorithms for this problem, we improve complexity by a factor of up to ✏  4 / 3 .
Introduction
The increasing use of machine learning models in high-stakes applications highlights the importance of reliable performance across changing domains and populations [11, 46, 35].
An emerging body of research addresses this challenge by replacing Empirical Risk Minimization (ERM) with Distributionally Robust Optimization (DRO) [6, 50, 49, 22, 40], with applications in natural language processing [47, 61, 35], reinforcement learning [17, 54] and algorithmic fairness [27, 56].
While
ERM minimizes the average training loss, DRO minimizes the worst-case expected loss over all probability distributions in an uncertainty set U , that is, it minimizes
[Equation]
where ` S ( x ) is the loss a model x 2 X incurs on a sample S and X is a closed convex set with bounded Euclidean diameter.
This work develops new algorithms for DRO, focusing on formulations where U contains distributions supported on N training points, where N is potentially large.
We consider two well-studied DRO variants: (1) Group DRO [57, 31, 49], and (2) f -divergence DRO [19, 6, 22].
Group DRO Machine learning models may rely on spurious correlations (that hold for most training examples but are wrongly linked to the target) and therefore suffer high loss on minority groups where these correlations do not hold [30, 27, 11].
To obtain high performance across all groups, Group
DRO minimizes the worst-case loss over groups.
Given a set U
= { w 1 , . . .
, w M } of M distributions over [ N ] , the Group DRO objective is
[Equation]
[Table]
AGD
[Table]
If we deﬁne the loss of group i as L i ( x ) :=
P
Nj =1 w ij ` j ( x ) then objective (2) is equivalent to max q 2  M
P i
[ M ] q i L i ( x ) with  M := { q 2 R M  0 | ~ 1 T q = 1 } .
Note that, unlike ERM, Group
DRO requires additional supervision in the form of subgroup identities encoded by { w i } .
DRO with f -divergence Another approach to DRO, which requires only as much supervision as
ERM, takes U to be an f -divergence ball around the empirical (training) distribution.
For every convex function f : R + !
R [ { + 1 } such that f (1) = 0 , f (0 / 0)
=
and the f -divergence between distributions q and p over [ N ] is D f ( q, p )
[Equation]
The f -divergence DRO problem corresponds to the uncertainty set U = { q
[Equation]
N
[Equation]
D f
[Equation]
q
[Equation]
1 N
[Equation]

[Equation] [Equation]
i
[Equation]
Several well-studied instances of DRO are a special case of this formulation, with the two most notable examples being conditional value at risk (CVaR) and  2 uncertainty sets.
CVaR at level ↵ corresponds to f ( x ) = { x< 1 ↵ } such that U =
[Equation]
N s.t k q k 1  1 / ( ↵ N ) , and has many applications in ﬁnance such as portfolio optimization and credit risk evaluation [48, 37] as well as in machine learning [47, 39, 20, 60, 17, 54].
The  2 uncertainty set with size ⇢ > 0 corresponds to f ( x ) := 12 ⇢ ( x  1) 2 and the resulting DRO problem is closely linked to variance regularization [21] and has been extensively studied in statistics and machine learning [42, 27, 21, 39, 61].
Complexity notion In this paper, we design improved-complexity methods for solving the convex problems (2) and (3) under the assumption that the loss ` i is convex and Lipschitz for all i .
We measure complexity by the (expected) required number of ` i ( x ) and r ` i ( x ) evaluations to obtain ✏ -suboptimal solution, i.e., return x such that L DRO ( x )  min x ? 2 X L DRO ( x ? )  ✏ with constant probability.
Table 1 summarizes our complexity bounds and compares them to prior art.
Throughout the introduction we assume (for simplicity) a unit domain size ( k x  y k  1 for all x, y 2 X ) and that each loss is 1 -Lipschitz.
Prior art Let us review existing methods that solve Group DRO and f -divergence DRO problems (see Section 1.1 for extended discussion).
For a dataset with N training points, the subgradient method [45] ﬁnds an ✏ approximate solution in ✏  2 iterations.
Computing a single subgradient costs
N functions evaluations (since we need to ﬁnd the maximizing q ).
Therefore, the complexity of this method is O
[Equation]
DRO can also be viewed as a game between a minimizing x -player and a maximizing q -player, which makes it amenable to primal-dual methods [43, 42, 49].
If we further assume that the losses are bounded then, for q 2  m , stochastic mirror descent with local norms obtains a regret bound of
O p m log( m ) /T
(see Appendix A.1).
As a consequence, for Group DRO (where m = M ) the complexity is
[Equation]
, and for f -divergence DRO ( m = N ) the complexity is
[Equation]
Levy et al.
[39] studied  2 -divergence and CVaR DRO problems, and proposed using standard gradient methods with a gradient estimator based on multilevel Monte Carlo (MLMC)
For  2 -divergence with ball of size ⇢ they proved a complexity bound of e
[Equation]
, and for CVaR at level ↵ they established complexity e
O  ↵
[Equation]
✏
[Equation]
.
However, for large uncertainty sets (when ⇢ or ↵ approach N ) their method does not improve over the subgradient method.
Stronger complexity bounds are available under the weak smoothness assumption that each ` i has O
[Equation]
-Lipschitz gradient.
Note that this is a weak assumption since if a function ` is not continuously differentiable, it is possible to approximate ` with additive error at most ✏ / 2 , by its Moreau envelope
[Equation]
o (see [15, Appendix A.1] for more details).
In particular, we can apply Nesterov’s accelerated gradient descent method [44] on an entropy-regularized version of our objective to solve the problem with complexity e
O
[Equation]
N
[Equation]
; see
Appendix A.2 for more details.
Our contribution We propose algorithms that solve the problems (2) and (3) with complexity e
O
[Equation]
.
Compared to previous works, we obtain better convergence rates for DRO with general f -divergence when N  1 and for Group DRO when M  N ✏ 4 / 3 .
When the losses have
O ( ✏  1 ) -Lipschitz gradient, we solve f -divergence DRO with complexity
[Equation]
and, under an even weaker mean-square smoothness assumption ( E j ⇠ w i kr ` j ( x )  r ` j ( y ) k 2 
O ( ✏  2 ) k x  y k 2 for all x, y and i ), we solve Group DRO with complexity
[Equation]
[Equation]
Our complexity bounds are independent of the structure of f and { w i } , allowing us to consider arbitrarily f -divergence balls and support a large number of (potentially overlapping) groups.
Our rates are optimal up to logarithmic factors for the special case of minimizing max i 2 [ N ] ` i ( x ) , which corresponds to Group DRO with N distinct groups and f -divergence DRO with f =
Our approach Our algorithms are based on a technique for acceleration with a ball optimization oracle , introduced by Carmon et al.
[13] and further developed in [15, 3].
Given a function F and a query point x , the ball optimization oracle returns an approximate minimizer of F inside a ball around x with radius r ; the works [13, 15, 3] show how to minimize F using e
O ( r  2 / 3 ) oracle calls.
Our development consists of efﬁciently implementing ball oracles with radius r = e
O ( ✏ ) for the DRO problems (2) and (3), leveraging the small ball constraint to apply stochastic gradient estimators that would have exponential variance and/or cost without it.
Carmon et al.
[15] previously executed this strategy for minimizing the maximum loss, i.e., max q 2  N
P i q i
` i ( x ) , which is a special case of both Group DRO and f -divergence DRO.
However, the ball-oracle implementations of [15] do not directly apply to the DRO problems that we consider; our oracle implementations differ signiﬁcantly and intimately rely on the Group DRO and f -divergence problem structure.
We now brieﬂy review the main differences between our approach and [15], highlighting our key technical innovations along the way.
Since the Group DRO objective is max q 2  M
P i q i L i ( x ) for L i ( x ) =
P j 2 [ N ] w ij ` j ( x ) , one may naively apply the technique of [15] with L i replacing ` i .
However, every step of such a method would involve computing quantities of the form e L i ( x ) / ✏ 0 (for some ✏ 0 = e ⇥ ( ✏ ) ), which can be up to N times more expensive than computing e ` j ( x ) / ✏ 0 for a single j .
To avoid such expensive computation we use MLMC [9] to obtain an unbiased estimate of e L i ( x ) / ✏ with complexity O (1) and appropriately bounded variance.
In the weakly-smooth case we also adapt our estimator to facilitate variance reduction [33, 2].
For f -divergence we consider the well-known dual form [6, 50]
[Equation]
where ⇤ ( v ) = max t  0 { vt  ( t ) } is the Fenchel dual of .
Stochastic gradient methods applied directly on the dual formulation are notoriously unstable (see e.g., [42]).
This is due to the fact that the Fenchel dual ⇤ can be very badly behaved even for standard f -divergences.
We solve this problem by (a) considering a small ball, (b) entropy-regularizing .
Our techniques rely on two technical observations: (i) log( ⇤ ✏ 0 ( · )) is 1 / ✏ 0 -Lipschitz for all convex , and (ii) for 1-Lipschitz losses, y ?
( x ) =
argmin y 2 R
⌥ ( x, y ) satisﬁes | y ?
( x )  y ?
( x 0 ) |  k x  x 0 k for all x, x 0 .
To the best of our knowledge, these observations are new and potentially of independent interest.
Paper organization.
Section 2 provides notation and a concise summary of the ball acceleration framework (largely taken from prior work) on which we build our algorithms.
Sections 3 and present our main contributions in the Group and f -divergence DRO settings, respectively.
Finally,
Section 5 concludes with discussion on the limitations and possible extensions of our work.
1.1 Additional related work
MLMC estimators The multilevel Monte Carlo (MLMC) technique was introduced by Giles [26] and Heinrich [29] in order to reduce the computational cost of Monte Carlo estimation of integrals.
Blanchet and Glynn [9] extended this technique to estimating functions of expectation and proposed several applications, including stochastic optimization [10].
In this work we use their estimator for two distinct purposes: (1) obtaining unbiased Moreau envelope gradient estimates for ball oracle acceleration as proposed by Asi et al.
[3], and (2) estimating the exponential of an expectation for
Group DRO.
Levy et al.
[39] also rely on MLMC for DRO, but quite differently than we do: they directly estimate the DRO objective gradient via MLMC, while we estimate different quantities.
Other DRO methods Several additional works proposed algorithms with theoretical guarantees for f -divergence DRO.
Jin et al.
[32] considered non-convex and smooth losses.
Song et al.
[53] proposed an algorithm for linear models with complexity comparable to the “AGD on softmax” approach (Appendix A.2).
Namkoong and Duchi [42] proposed a primal-dual algorithm that is suitable for small uncertainty  2 sets (with size ⇢ ⌧ 1 N ) and Curi et al.
[20] proposed a primal-dual algorithm specialized for CVaR.
Other works consider DRO with uncertainty sets deﬁned using the
Wasserstein distance [24, 23, 51, 34].
Another relevant line of works proposes reﬁnements for DRO that address some of the challenges in applying it to learning problems [60, 59, 55].
Preliminaries
Notation We write k · k for the Euclidean norm.
We denote by B r ( x 0 ) the Euclidean ball of radius r around x 0 .
We let  n := { q 2 R n  0 | 1 T q = 1 } denote the probability simplex in R n .
For the sequence z m , . .
, z n we use the shorthand z nm .
Using F as a generic placeholder (typically for a loss function ` i ), we make frequent use of the following assumption.
Assumption 1.
The function F : X !
R is convex and G-Lipschitz, i.e., for all x, y 2 X we have | F ( x )  F ( y ) |  G k x  y k .
In addition, the domain X is a closed and convex set, and it has Euclidean diameter at most R .
Throughout, N denotes the number of losses and, in Section 3, M denotes the number of groups.
We use ✏ for our target accuracy and r ✏ := ✏ 0 /G for the ball radius, where ✏ 0 = ✏ / (2 log M ) for
Group-DRO (Section 3) and ✏ 0 = ✏ / (2 log N ) for f -divergence DRO (Section 4).
Complexity model We measure an algorithm’s complexity by its expected number of ` i and r ` i evaluations; bounds on expected evaluation number can be readily converted to more standard probability 1 bounds [see 3, Appendix A.3].
Moreover if X ⇢ R d , d = ⌦ (log N ) , 2 and the time to evaluate ` i and r ` i is O ( d ) , the expected runtime of all the algorithms we consider is at most d times the evaluation complexity.
2.1 Ball oracle acceleration
We now brieﬂy summarize the complexity bounds given by the framework of [13, 15, 3] for accelerated minimization using queries to (inexact) ball optimization oracles, deﬁned as follows.
Deﬁnition 1.
An algorithm is a Ball Regularized Optimization Oracle of radius r ( r -BROO) for function F : X !
R if for query point ¯ x 2 X , regularization parameter  > 0 and desired accuracy  > 0 it returns
[Equation]
Proposition 1.
Let F satisfy Assumption 1, let C F be the complexity of evaluating F exactly, and let C  (  ) bound the complexity of an r -BROO query with  ,  .
Assume that C  (  ) is nonincreasing in  and at most polynomial in 1 /  .
For any ✏ > 0 , Algorithm 1 returns x such that
F ( x )  min x ? 2 X F ( x ? )
 ✏ with probability at least 12 .
For m ✏ = O  log GR 2 ✏ r  and  m =
O  m 2 ✏ ✏ r 4 / 3 R 2 / , the complexity of the algorithm is
[Equation]
Informally, the proposition shows that e
O (( R/r ) 2 / 3 ) BROO calls with  = e ⌦ ( ✏ / ( r 3 / 4 R 2 / 3 )) and accuracy  = e
O ( r ) sufﬁce to ﬁnd an ✏ -accurate solution.
As we show in the sequel, for C  (  )
=
e
O
N + ( G  ) the resulting complexity bound is e
O
N ( GR ✏ ) 2 / 3 + ( GR
[Equation]
) .
The summation over j in bound (5) stems from the use of MLMC to de-bias the BROO output (i.e., make it exact in expectation): compared to the original proposal of Asi et al.
[3], our version of the procedure in
Appendix B slightly alters this MLMC scheme by de-biasing one accurate BROO call instead of averaging many inaccurate de-biased calls, improving our bounds by logarithmic factors.
Group DRO
In this section we develop our BROO implementations for the Group DRO objective (2).
In Section 3.
we describe an “exponentiated group-softmax” function that approximates L g-DRO with additive error at most ✏ / 2 .
We then apply stochastic gradient methods on this function to obtain BROO implementations that yield improved rates for Group DRO via Proposition 1: we ﬁrst consider the non-smooth case in Section 3.2 and then the weakly-smooth case in Section 3.3.
Exponentiated group-softmax
Given a cheap and unbiased stochastic gradient estimator of r L g-DRO , we could use a variant of
SGD and minimize L g-DRO to ✏ -suboptimal solution using O ( ✏  2 ) steps.
However, obtaining an unbiased estimator is challenging due to the maximum operator in L g-DRO .
As a ﬁrst step we use entropy smoothing [7, 8, 5, 4] to replace the maximum in L g-DRO with the softmax operation.
More speciﬁcally, we use the trick from [15] and minimize the “exponentiated softmax” (that has the form of a weighted ﬁnite sum) within a small ball.
For target accuracy ✏ , regularization parameter   0 , center point ¯ x 2 X and ✏ 0 = ✏ / (2 log M ) > 0 , the (regularized) group-softmax function is
[Equation]
We will implement a BROO for L smax , ✏ :=
L smax , ✏ , 0 , which is a uniform approximation of L g-DRO : |L g-DRO ( x )  L smax , ✏ ( x ) |  ✏ / 2 for all x 2 X ; see Appendix C.1 for details.
The (regularized) exponentiated group-softmax is  ✏ ,
[Equation]
x
[Equation]
[Equation] [Equation]
] ¯ p i  i
[Equation]
x ) where
[Equation]
x
[Equation]
In the following lemma we (easily) extend Carmon et al.
[15, Lemma 1] to exponentiated groupsoftmax, showing that  ✏ ,  is well-behaved inside a ball of (appropriately small) radius r around ¯ x and facilitates minimizing L smax , ✏ ,  in that ball; see Appendix C.1 for the proof.
Lemma 1.
Let each ` i satisfy Assumption 1, and consider the restriction of L smax , ✏ ,  (6) and  ✏ ,  (7) to B r (¯ x ) .
Then the functions have the same minimizer x ? 2 B r (¯ x ) and, if   O ( G/r ) and r  O ( ✏ 0 /G ) , then (a)  ✏ ,  is ⌦ (  ) -strongly convex, (b) each  i is O ( G ) -Lipschitz and (c) for every x 2 B r (¯ x ) we have L smax , ✏ ,  ( x )  L smax , ✏ ,  ( x ? )
 O (  ✏ ,  ( x )   ✏ ,  ( x ? )) .
3.2 BROO implementation for Group DRO non-smooth losses
To motivate our BROO implementation, let us review how [15] use the exponentiated softmax in the special case of size-1 groups, i.e., L i =
` i , and explain the difﬁculty that their approach faces when the group structure is introduced.
The BROO implementation in [15] is based on SGD variant with the stochastic gradient estimator
ˆ g ( x )
=
e ( L i ( x )  L i (¯ x )) / ✏ 0 r L i ( x ) where i ⇠
¯ p i .
However, for
Group DRO where L i =
P j 2 [ N ] w ij ` j , the estimator ˆ g ( x ) can be up to N times more expensive to compute.
Approximating ˆ g ( x ) by drawing j, j
⇠ w i and taking e ( ` j ( x )  ` j (¯ x )) / ✏ 0 r ` j 0 ( x ) will result in a biased estimator since E j ⇠ w i e ( ` j ( x )  ` j (¯ x )) /
✏ 0 6 =
e ( L i ( x )  L i (¯ x )) / ✏ 0 .
To address this challenge we propose a new gradient estimator based on the multilevel Monte Carlo (MLMC) method [9].
The MLMC unbiased estimator for  i ( x ) =
✏ 0 e ( L i ( x )  L i (¯ x )) / ✏ 0 , which we denote by c
M [  i ( x )] , is deﬁned as follows:
[Equation] [Equation]
S n
[Equation]
⇠ w
[Equation]
and let
[Equation]
where p j :=
P ( J = j ) =
1 / p
j ⇣ 1  1 p⌘ and, for n 2 2 N , we deﬁne b
D n
[Equation]
and
[Equation]
With the MLMC estimator for  i in hand, we estimate the gradient of  ✏ ,  as follows:
Draw i
⇠
p (
¯
x ) , j
⇠
w i and
set ˆ
g ( x ) = 1 ✏
[Equation]
In the following lemma we summarize the important properties of the MLMC and gradient estimators; see Appendix C.2 for the proof.
Lemma 2.
Let each ` i satisfy Assumption 1, and let r  ✏ 0 G ,   Gr and x 2 B r (¯ x ) .
Then c
M [  i ( x )] and ˆ g ( x ) are unbiased for  i ( x ) and r  ✏ ,  ( x ) , respectively, and have bounded second moments:
[Equation]
.
In addition, the complexity of computing
[Equation]
Due to Lemma 2 and since  ✏ ,  is ⌦ (  ) -strongly convex, we can use the Epoch-SGD algorithm of Hazan and Kale [28] with our gradient estimator (8).
This algorithm has rate of convergence
O
G 2 / (  T )  and our gradient estimator requires additional N function evaluations for precomputing the sampling probabilities { ¯ p i } .
We thus arrive at the following complexity bound.
Theorem 1.
Let each ` j satisfy Assumption 1, let ✏ ,  ,  > 0 and let r ✏ = ✏ / (2 G log M ) .
For any query point ¯ x 2 R d , regularization strength   O ( G/r ✏ ) and accuracy  , EpochSGD [28,
Algorithm 1]) with the gradient estimator (8) outputs a valid r ✏ -BROO response and has complexity
C  (  ) =
O
N + G 2  2  .
Consequently, the complexity of ﬁnding an ✏ -suboptimal minimizer of L g-DRO (2) with probability at least 12 is
[Equation]
We provide the proof for Theorem 1 in Appendix C.3; the ﬁnal complexity bound follows from straightforward calculations which we now brieﬂy outline.
According to Proposition 1, ﬁnding an ✏ 2 -suboptimal solution for L smax , ✏ (and consequently an ✏ -suboptimal solution for L g-DRO ) involves e
O  ( R/r ✏ ) 2 /
BROO calls with accuracy  =
e ⌦  r ✏ 2  J/ and regularization strength    m , where J = min { Geom ( 12 ) , m } .
We may therefore bound the complexity of each such call by
[Equation]
where ( ? ) follows from substituting  m =
e ⌦  ✏ r  4 / 3 ✏
R  2 / and m = e
O (1) .
Multiplying this bound by e
O  ( R/r ✏ ) 2 / yields (up to polylogarithmic factors)
the conclusion of Theorem 1.
3.3 Accelerated variance reduction for mean-square smooth losses
In this section we provide an algorithm with an improved rate of convergence under the following mean-square smoothness assumption.
Assumption 2.
For all x, x 0 2 B r (¯ x ) and i 2 [ M ] , E
[Equation]
Note that assuming L -Lipschitz gradient for each ` i implies Assumption 2, but not the other way around.
To take advantage of Assumption 2, we ﬁrst rewrite the function  ✏ ,  ( x ) in a way that is more amenable to variance reduction:
[Equation]
x, x
[Equation]
, where  i ( x, x
[Equation]
and p i ( x 0 ) :
= e
[Equation]
i ( x
[Equation]
(Note that  i ( x, ¯ x ) =
i ( x ) ).
Given a reference point x 0 , to compute a reduced-variance estimator of r  ✏ ,  ( x ) , we draw i ⇠ p i ( x 0 ) and j ⇠ w i , and set:
[Equation]
where r `  j ( x ) :
= r ` j ( x ) +  ( x  ¯ x ) and c
M [  i ( x, x 0 )] is an MLMC estimator for
i ( x, x 0 ) deﬁned analogously to c
M [  i ( x )] (see details in Appendix C.4).
The estimator (9) is not precisely standard
SVRG [33] since we use c
M [  i ( x, x 0 )] as an estimator for  i ( x, x 0 ) .
Simple calculations show that
E ˆ g x 0 ( x ) =
r  ✏ ,  ( x ) and the following lemma shows that ˆ g satisﬁes a type of variance bound conducive to variance-reduction schemes; see Appendix C.4 for the proof.
Lemma 3.
Let each ` j satisfy Assumptions 1 and 2.
For any   Gr , r = ✏ 0 G and x, x 0 2 B r (¯ x ) , the variance of ˆ g x 0 ( x ) is bounded by Var(ˆ g x 0 ( x ))
 O
[Equation]
Accelerated variance reduction methods for convex functions typically require a stronger variance bound of the form Var(ˆ
g x 0 ( x ))  2 L ( F ( x 0 )  F ( x )  hr F ( x ) , x 0  x i ) for every x
[cf. 1, Lemma 2.4].
The guarantee of Lemma 3 is weaker, but still allows for certain accelerated rates via, e.g., the
Katyusha X algorithm [2].
With it, we obtain the following guarantee.
Theorem 2.
Let each ` j satisfy Assumptions 1 and 2.
Let ✏ > 0 , ✏ 0 =
✏ / (2 log M ) and r ✏ = ✏ 0 /G .
For any query point ¯ x 2 R d , regularization strength   O ( G/r ✏ ) and accuracy  , KatyushaX s [2, Algorithm 2] with the gradient estimator (9) outputs a valid r ✏ -BROO response and has complexity C  (  ) =
O ⇣⇣
N + N 3 /
( G + p ✏ 0 L ) p ✏⌘ log
Gr ✏ ⌘ .
Consequently, the complexity of ﬁnding an ✏ -suboptimal minimizer of L g-DRO (2) with probability at least 12 is
O N
[Equation]
We provide the proof of Theorem 2 in Appendix C.5.
For the special case of Group DRO with a single group satisfying Assumption 2 with L = ⇥
[Equation]
, i.e. minimizing the average loss, we have the lower bound
[Equation]
[62] and for the case of N distinct groups, i.e. minimizing the maximal loss, we have the lower bound e
[Equation]
[15].
This implies that in the weakly mean-square smooth setting the term scaling as N 3 / 4 ✏  1 and the term scaling as N ✏  2 / 3 are unimprovable.
DRO with f -divergence
In this section we develop our BROO implementation for the f -divergence objective (3).
In Section 4.
we reduce the original DRO problem to a regularized form using Lagrange multipliers.
Next, in
Section 4.2 we show that adding negative entropy regularization to the objective produces the stability properties necessary for efﬁcient ball optimization.
In Section 4.3 we describe a BROO implementation for the non-smooth case using a variant of Epoch-SGD [28], and in Section 4.4 we implement the BROO under a weak-smoothness assumption by carefully restarting an accelerated variance reduction method [1].
4.1 The dual problem
We ﬁrst note that (due to Slater’s condition), by Lagrange duality, the objective (3) is equivalent to
L f -
[Equation]
=
[Equation]
Writing ( s ) :
= ⌫ N f ( Ns ) for : R + !
R , we therefore consider objectives of the form
[Equation]
where G is the Lipschitz constant of each loss ` i and for the last equality we use Lagrange duality, with ⇤ ( v ) :
= max t 2 dom ( ) { vt  ( t ) } the Fenchel dual of (for more details see Appendix D.1).
We show that under weak assumptions (introducing logarithmic dependence on bounds on f and the losses) we can solve the constrained problem (3) to accuracy ✏ by computing a polylogarithmic number of O ( ✏ ) -accurate minimizers of (10); see Appendix D.2 for details.
Since the complexity of solving (10) holds for any ⌫ > 0 , and we have a lower bound for the required ⌫ , for the remainder of this section we focus on minimizing L for arbitrary convex .
4.2 Stabilizing the gradient estimator
While minimizing (10) can be viewed as ERM (over x and y ), straightforward application of SGD does not solve it efﬁciently.
To see this, consider the standard gradient estimator formed by sampling i ⇠ Unif ([ N ]) and taking ˆ g x = N ⇤
( ` i ( x )  Gy ) r ` i ( x ) and ˆ g y
[Equation]
For general , this estimator will have unbounded second moments, and therefore SGD using them would lack a convergence guarantee.
As an extreme example, consider = 0 (corresponding to minimizing the maximum loss) whose conjugate function ⇤ ( v ) is 0 for v  0 and 1 for v > 0 , leading to meaningless stochastic gradients.
We obtain bounded gradient estimates in two steps.
First, we ﬁnd a better distribution for i using a reference point ¯ x 2 X with corresponding ¯ y = argmin
y 2 R
⌥
(¯ x, y ) .
Namely, we note that the optimality condition for ¯ y implies that ⇤
( ` i (¯ x )  G ¯ y ) is a pmf over [ N ] .
Therefore, we may sample i ⇠ ⇤
( ` i (¯ x )  G ¯ y ) and estimate the gradient of ⌥ at ( x, y ) using ˆ g x = ⇢ i ( x, y )
r ` i ( x ) and ˆ g y
=
G  1  ⇢ i ( x, y )  , where ⇢ i ( x, y ) =
⇤
( ` i ( x )  Gy ) ⇤
( ` i (¯ x )  G ¯ y ) .
However, for general (and = 0 in particular), the ratio ⇢ i ( x, y ) can be unbounded even when x, y are arbitrarily close to ¯ x, ¯ y .
Our second step ensures that ⇢ i ( x, y ) is bounded around ¯ x, ¯ y by adding a small negative entropy term to , deﬁning ✏ ( q ) :
= ( q ) + ✏ 0 q log q where ✏
[Equation] [Equation]
✏
[Equation]
log N , (11) and
L , ✏ ( x ) =
min y 2 R ⌥ ✏ ( x, y ) with ⌥ ✏ ( x, y ) :=
X i 2 [ N ] ⇤ ✏ ( ` i ( x )  Gy )
+ Gy.
Due to our choice of ✏ 0 , we have |L ( x )  L , ✏ ( x ) |  ✏ / 2 for all x 2 R d , and consequently an ✏ / 2 -accurate minimizer of L , ✏ is also an ✏ -accurate for L (see Lemma 18 in
Appendix D.3).
When =
we have ⇤ ✏ ( v ) =
e ( v  1) / ✏ 0 and therefore the corresponding ⇢ i ( x, y )
=
e ( ` i ( x )  ` i (¯ x )  G ( y  ¯ y )) / ✏ 0 .The
following lemma, which might be of independent interest, shows that the same conclusion holds for any convex .
Lemma 4.
For any convex : R + !
R and ✏ deﬁned in (10) , log
[Equation] [Equation]
is 1 ✏ 0 -Lipschitz.
See proof in Appendix D.4.
Thus, ⇤ ✏ 0 ( v ) / ⇤ ✏ 0 (¯ v ) =
e log ⇤ ✏ 0 ( v )  log ⇤ ✏ 0 (¯ v )  e ( v  ¯ v ) / ✏ 0 and ⇢ i ( x, y )
 e
( ` i ( x )  ` i (¯ x )  G ( y  ¯ y )) / ✏ 0 continues to hold.
Therefore, if | y  ¯ y |  ✏ 0 /G =
r ✏ and x 2 B r ✏ (¯ x ) (so that | ` i ( x )  ` i (¯ x ) |  ✏ 0 if ` i satisﬁes Assumption 1), we have the bound ⇢ i ( x, y )  e 2 .
It remains to show that we may indeed restrict y to be within distance r ✏ from ¯ y .
To this end, we make the following observation which plays a key part in our analysis and might also be of independent interest (see proof in Appendix D.4).
Lemma 5.
For G > 0 , ` ( x ) =
( ` 1 ( x ) , . . .
, ` N ( x )) and y ?
( x ) =
argmin y 2 R
⌥
✏ ( x, y ) , we have | y ?
( x )  y ?
( x 0 ) |  1 G k ` ( x )  ` ( x 0 ) k 1 for all x, x 0 2 X .
Moreover, if each ` i is G -Lipschitz, we have | y ?
( x )  y ?
( x 0 ) |  k x  x 0 k .
Lemma 5 implies that x ? , y ?
= argmin x 2 B r ✏ (¯ x ) ,y 2 R ⌥ ✏ ( x, y ) satisfy | y ?
¯ y |  k x ?
¯ x k  r ✏ .
Therefore, when minimizing ⌥ ✏ (or any regularized version of it) inside the ball B r ✏ (¯ x ) , we may restrict y to [¯ y  r ✏ , ¯ y + r ✏ ] without loss of generality.
We also note that Lemma 5 holds for all values of ✏ and is therefore valid even without entropy regularization (as long as ⇤ is strongly convex y ? is unique, and if y ? is not unique then we can still choose y ?
such that the bound of this lemma holds).
4.3 BROO implementation for f -divergence DRO with non-smooth losses
By the discussion above, to implement a BROO for L , ✏ ( x ) (with radius r ✏ = ✏ 0 /G , regularization  , and query ¯ x 2 X ) it sufﬁces to minimize ⌥ ✏ ,  ( x, y ) :
= ⌥ ✏ ( x, y )
+  2 k x  ¯ x k 2 over x 2 B r ✏ (¯ x ) and y
[¯ y  r ✏ , ¯ y + r ✏ ] , where ¯ y = argmin
y 2 R
⌥
✏ (¯ x, y ) .
To that end we estimate the gradient of ⌥ ✏ ,  ( x, y ) as follows.
Letting ¯ p i =
⇤ ✏ 0 ( ` i (¯ x )  G ¯ y ) (making ¯ p a pmf by optimality of ¯ y ), we sample i ⇠ ¯ p and set ˆ g x ( x, y ) =
⇤ ✏ 0 ( ` i ( x )  Gy ) ¯ p
i r ` i ( x, y ) and
ˆ g y ( x, y ) =
G
[Equation]
[Equation]
p
[Equation]
Lemma 4 implies the following bounds on our gradient estimator; see proof in Appendix D.5.
Lemma 6.
Let each ` i be G -Lipschitz, let ¯ x 2 X
and ¯ y = argmin
y 2 R
⌥
✏ ,  (¯ x, y ) .
Let r ✏ =
✏ 0 G , then for all x 2 B r ✏ (¯ x ) and
y 2 [¯ y  r ✏ , ¯ y + r ✏ ] , the gradient estimators ˆ g x
and ˆ g y satisfy the following:
E i ⇠ ¯ p i
[ˆ g x ( x, y )]
= r x ⌥ ✏ ( x, y ) and E i ⇠ ¯ p i
[ˆ g y ( x, y )]
= r y ⌥ ✏ ( x, y ) .
E i ⇠ ¯ p i k
ˆ g x ( x, y ) k 2  e 4 G 2 and E i ⇠ ¯ p i | ˆ g y ( x, y ) | 2  e 4 G 2 .
To implement the BROO using our gradient estimator we develop a variant of the Epoch-SGD algorithm of Hazan and Kale [28] (Algorithm 3 in Appendix D.5).
Similarly to Epoch-SGD, we apply standard SGD on ⌥ ✏ ,  (with gradient estimator (13)) in “epochs” whose length doubles in every repetition.
Our algorithm differs slightly in how each epoch is initialized.
Standard Epoch-SGD initializes with the average of the previous epoch’s iterates, and strong convexity shows that the suboptimality and distance to the optimum shrink by a constant factor after every epoch.
However, since ⌥ ✏ ,  is strongly convex only in x and not in y , we cannot directly use this scheme.
Instead, we set the initial y variable to be argmin y ⌥
✏ ,  ( x 0 , y ) , where x 0 is the initial x variable still deﬁned as the previous epoch’s average; this initialization has complexity N , but we only preform it a logarithmic number of times.
Using our initialization scheme and Lemma 5, we recover the original Epoch-SGD contraction argument, yielding the following complexity bound (see proof in Appendix D.5).
Theorem 3.
Let each ` i satisfy Assumption 1.
Let ✏ ,  ,  > 0 , and r ✏ =
✏ / (2 G log N ) .
For any query point ¯ x 2 R d , regularization strength   O ( G/r ✏ ) and accuracy  < r ✏ / 2 , Algorithm 3 outputs a valid r ✏ -BROO response for L , ✏ and has complexity C  (  ) =
O
[Equation]
Consequently, the complexity of ﬁnding an ✏ -suboptimal minimizer of L (10) with probability at least 12 is
[Equation]
4.4 Accelerated variance reduction for smooth losses
In this section, we take advantage of the following smoothness assumption.
Assumption 3.
For every i 2 [ N ] the loss ` i is L -smooth, i.e., has L -Lipschitz gradient.
Let us rewrite ⌥ ✏ ,  in a form that is more amenable to variance reduction techniques: ⌥ ✏ ,  ( x, y ) =
X i 2 [ N ] ¯ p i  i
[Equation]
x, y ) where  i ( x, y ) :=
[Equation]
and, as before ¯ p i =
⇤ ✏ ( ` i (¯ x )  G ¯ y ) for some ball center ¯ x 2 X
and ¯ y = argmin
y 2 R
⌥
✏ (¯ x, y ) .
In the following lemma, we bound the smoothness of the functions  i , deferring the proof to Appendix D.6.
Lemma 7.
For any i 2 [ N ] , let ` i be G -Lipschitz and L -smooth, let r ✏ = ✏ 0 G and  = O
Gr ✏  .
The restriction of  i to x 2 B r ✏ (¯ x ) and y 2 [¯ y  r ✏ , ¯ y + r ✏ ] is O ( G ) -Lipschitz
and O
[Equation]
[Equation]
smooth.
Since ⌥ ✏ ,  is a ﬁnite sum of smooth functions, we can obtain reduced-variance gradient estimates by the standard SVRG technique [33].
For any reference point x 0 , y 0 , the estimator is
[Equation]
where r is with respect to the vector [ x, y ] .
Similar to non-smooth case, obtaining an efﬁcient
BROO implementation is complicated by the fact that ⌥ ✏ ,  is strongly-convex in x but not in y .
Our solution is also similar: we propose a restart scheme and minimize over y exactly between restarts (Algorithm 4 in Appendix D.6), that gives the following complexity bound (see proof in
Appendix D.6).
Theorem 4.
Let each ` i satisfy Assumptions 1 and 3, let ✏ ,  ,  > 0 , and r ✏ =
✏ 2 G log N .
For any query point ¯ x 2 R d , regularization strength   O ( Gr ✏ ) and accuracy  , Algorithm 4 outputs a valid r ✏ -BROO response for L , ✏ and has complexity C  (  ) =
O
[Equation]
Consequently, the complexity of ﬁnding an ✏ -suboptimal minimizer of L (10) with probability at least 12 is
[Equation]
N
[Equation]
Discussion
Limitations While our work indicates that the ball optimization approach offers signiﬁcant complexity gains for DRO, we note that turning the algorithms we propose into practical DRO methods faces several challenges.
A main challenge is the costly bisection procedure common to all MonteiroSvaiter-type acceleration schemes [41, 25, 13, 52].
Fortunately, very recently, two works [16, 36] (the former partially motivated by our paper) have shown how to remove the bisection from MonteiroSvaiter schemes, signiﬁcantly improving the practical potential of the methods we propose.
However, another practical limitation of our approach is the need to tune many parameters that are not known in advance, such as those relating to the ball radius r ✏ and smoothing level ✏ 0 , as well as step sizes and number of iterations of oracle implementations; a more adaptive setting for these parameters is likely important.
Extensions First, it would be interesting to extend our approach to DRO objectives max q 2 U
P i 2 [ N ] q i ` i ( x ) with uncertainty set U that is an arbitrary subset of the simplex.
While the subgradient method, the primal-dual method (Appendix A.1), and “AGD on softmax” (Appendix A.2) all apply to any U ✓  N , our methods strongly rely on the structure of U induced by
Group- f -divergence DRO, and extending them to unstructured U ’s seems challenging.
Second, it would be interesting to generalize our results in the “opposite” direction of getting better complexity bounds for problems with additional structure.
For Group-DRO our bounds are essentially optimal when the number of groups M = ⌦ ( N ) , but are suboptimal
when M =
O
We leave it as a question for further research if it is possible to obtain a stronger bound such as e
O
N + M ✏  2 / 3 + ✏  , which recovers our result for M = N but improves on it for smaller values of M .
Taking CVaR at level ↵ as a special case of f -divergence DRO, our bounds are optimal when ↵ is close to 1 /N but suboptimal for larger value of ↵ ; it would be interesting to obtain bounds such as
[Equation]
A third possible extension of our research is DRO in the non-convex setting.
For this purpose, it might be possible to use the technique of Carmon et al.
[12] for turning accelerated convex optimization algorithms to improved-complexity methods for smooth non-convex optimization.
Acknowledgments
We thank Marc Teboulle for detailed and helpful feedback.
This work was supported by Israeli
Science Foundation (ISF) grant no. 2486/21, the Len Blavatnik and the Blavatnik Family foundation and the Adelis Foundation.