Title: Drago: Primal-Dual Coupled Variance Reduction for Faster Distributionally Robust Optimization
Venue: NeurIPS 2024
Accepted: yes
Keywords: Distributionally Robust Optimization

Stochastic Optimization

Convex Optimization

Saddle Point
Abstract: We consider the penalized distributionally robust optimization (DRO) problem with a closed, convex uncertainty set, a setting that encompasses learning using $f$-DRO and spectral/$L$-risk minimization. We present Drago, a stochastic primal-dual algorithm which combines cyclic and randomized components with a carefully regularized primal update to achieve dual variance reduction. Owing to its design, Drago enjoys a state-of-the-art linear convergence rate on strongly convex-strongly concave DRO problems witha fine-grained dependency on primal and dual condition numbers. The theoretical results are supported with numerical benchmarks on regression and classification tasks.

DRAGO: Primal-Dual Coupled Variance Reduction for Faster Distributionally Robust Optimization

Abstract

We consider the penalized distributionally robust optimization (DRO) problem with a closed, convex uncertainty set, a setting that encompasses learning using f -DRO and spectral/ L -risk minimization. We present DRAGO, a stochastic primal-dual algorithm that combines cyclic and randomized components with a carefully regularized primal update to achieve dual variance reduction. Owing to its design, DRAGO enjoys a state-of-the-art linear convergence rate on strongly convex-strongly concave DRO problems with a fine-grained dependency on primal and dual condition numbers. The theoretical results are supported by numerical benchmarks on regression and classification tasks.

1 Introduction

Contemporary machine learning research is increasingly exploring the phenomenon of distribution shift, in which predictive models encounter different data-generating distributions in training versus deployment [Wiles et al., 2022]. A popular approach to learn under potential distribution shift is distributionally robust optimization (DRO) of an empirical risk-type objective

where ℓ i : R d → R denotes the loss on training instance i ∈ [ n ] := { 1 , . . . , n } , and q = ( q 1 , . . . , q n ) ∈ Q is a vector of n weights for each example. The feasible set Q , often called the uncertainty set , is a collection of possible instance-level reweightings arising from distributional shifts between train and evaluation data, and is often chosen as a ball about the uniform vector 1 /n = (1 /n, . . . , 1 /n ) in f -divergence [Namkoong and Duchi, 2016, Carmon and Hausler, 2022, Levy et al., 2020] or a spectral/ L -risk-based uncertainty set [Mehta et al., 2023].

We consider here the penalized version of (1), stated as

where µ, ν ≥ 0 are regularization parameters and D ( q ∥ 1 /n ) denotes some statistical divergence (such as the Kullback-Leibler (KL) or χ 2 -divergence) between the original weights 1 /n and shifted weights q . For clarity, we focus on the cases of µ, ν > 0 , but also describe the modifications to the methods, results, and proofs for cases in which µ = 0 or ν = 0 , in Appx. C.4. See Fig. 1 for intuition on the relationship between the uncertainty set, divergence D , and hyperparameter ν .

Standard (1) and penalized (2) DRO objectives have seen an outpour of recent use in reinforcement learning and control [Lotidis et al., 2023, Yang et al., 2023, Wang et al., 2023a, Yu et al., 2023, Kallus et al., 2022, Liu et al., 2022] as well as creative applications in robotics [Sharma et al., 2020], language modeling [Liu et al., 2021], sparse neural network training [Sapkota et al., 2023]

1.0

1.0

1.0

0.8

0.8 -

0.8

0.6 -

0.6

+

0.6 +

+10

0.4

0.4 -

0.2

0

0.2 +

+ 0°0

0.0

+0°0

0.0

0.0

0.0

0.2

0.2

0.2

0.4

0.4

0.4

0.6

0.6

0.6

0.8

0.8

0.8

1.0

1.0

0.8

0.6

0.4

0.2

0.0

1.0

1.0

0.8

0.6

0.4

0.2

0.0

1.0

1.0

0.8

0.6

0.4

0.2

0.0

and defense against model extraction [Wang et al., 2023b]. However, even in a classical supervised learning setup, current optimization algorithms for DRO have limitations in both theory and practice.

For context, we consider the large-scale setting in which the sample size n is high, and the training loss in each example is accessed through a collection of n primal first-order oracles { ( ℓ i , ∇ ℓ i ) } n i =1 . Quantitatively, we measure the performance of algorithms by runtime or global complexity of elementary operations to reach within ε of the minimum of ¯ L ( w ) = max q ∈Q L ( w,q ) , while qualitatively, we consider the types of uncertainty sets that can be handled by the algorithm and convergence analysis. Under standard assumptions, ¯ L is differentiable with gradient computed via

̸

In the learning setting, we are interested in stochastic algorithms that can approximate this gradient with b < n calls to the oracles. For uniformly randomly sampled i ∈ [ n ] , nq ⋆ i ( w ) ∇ ℓ i ( w ) + µw is an unbiased estimator of ∇ ¯ L ( w ) . However, computing q ⋆ i ( w ) depends on the first step in (3) which itself requires calling all n oracles (see (2)), i.e., it is no different than the cost of full batch gradient descent. A direct minibatch stochastic gradient descent approach would approximate L ( w, · ) in the first step of (3) with only b calls to generate approximate weights ˆ q ( w ) . Because ˆ q ( w ) = q ⋆ ( w ) in general for b < n , these methods have non-vanishing bias, i.e., do not converge [Levy et al., 2020].

This motivated research on DRO-specific stochastic algorithms with theoretical convergence guarantees under particular assumptions (see Tab. 1 and Appx. B for details) [Namkoong and Duchi, 2016, Levy et al., 2020, Carmon and Hausler, 2022]. Although we highlight the dependence on sample size n and suboptimality ε , the dependence on all constants is given in Tab. 1. For f -divergence-based uncertainty sets in the standard oracle framework, several methods achieve a O ( ε -2 ) complexity. Levy et al. [2020] do so by proving uniform bias bounds, so that if b scales as O ( ε -2 ) , the convergence guarantee is achieved. However, if the required batch size b exceeds the training set size n , then the method reduces to the sub-gradient method, as we can see in Tab. 1. These sublinear rates typically stem from two causes. The first is the adoption of a 'fully stochastic' perspective on the oracles, wherein each oracle's output is treated as an independent random sample drawn from a probability distribution. The second is the non-smoothness of the objective, as we shall see below.

Variance reduction techniques, on the other hand, exploit the fact that the optimization algorithm takes multiple passes through the same dataset, and achieve linear rates of the form O (( n + κ ℓ ) log( ε -1 )) in empirical risk minimization when the objective is both smooth (i.e., has Lipschitz continuous gradient) and strongly convex and κ ℓ is an associated condition number [Johnson and Zhang, 2013, Defazio et al., 2014]. Assuming access to stronger oracles involving constrained minimization and applying a variance reduction scheme, Carmon and Hausler [2022] achieve O ( nε -2 / 3 + n 3 / 4 ε -1 ) for f -divergences as well, but do not obtain linear convergence due to the second type of cause: the objective ∇ ¯ L is non-smooth when ν = 0 .

Recently, Mehta et al. [2024] handled the ν > 0 case for spectral risk uncertainty sets, and their variance-reduced algorithm achieves a linear O (( n + κ Q κ ℓ ) ln(1 /ε )) convergence guarantee (where κ Q ≥ 1 measures the 'size' of the uncertainty set), but only with a lower bound of order Ω( n ) on the problem parameter ν . The challenge of this problem, considered from a general optimization viewpoint beyond DRO, stems from the non-bilinearity of the coupled term ∑ n i =1 q i ℓ i ( w ) and the constraint that ∑ n i =1 q i = 1 for probability vectors. If the coupled term was bilinear (i.e., of the form q ⊤ Aw for A ∈ R n × d ) and the constraints applied separately to each q i , then dual decomposition techniques could be used. Qualitatively, algorithms and analyses often rely on particular uncertainty sets; for example, Kumar et al. [2024] use duality arguments specific to the Kullback-Leibler uncertainty set to create a primal-only minimization problem. See Appx. B for a detailed discussion of related work from the ML and the optimization lenses. Given the interest from both communities, we address whether a stochastic DRO algorithm can simultaneously 1) achieve a linear convergence rate for any ν > 0 and 2) apply to many common uncertainty sets.

Contributions Wepropose DRAGO, a minibatch primal-dual algorithm for the penalized DRO problem (2) that achieves ε -suboptimality in

iterations, where b ∈ { 1 , . . . , n } is the minibatch size, κ Q = nq max := n max q ∈Q ,i ∈ [ n ] q i measures the size of the uncertainty set, and G and L are the Lipschitz continuity parameters of ℓ i and ∇ ℓ i , respectively. For commonly used parameters of uncertainty sets, nq max is bounded above by an absolute constant independent in n (see Prop. 3), so for d < n and b = n/d , we maintain an O ( n ) per-iteration complexity (the dual dimensionality) while reducing the number of iterations to O (( d + κ Q L/µ + d √ nG 2 / ( µν ))) ln (1 /ε )) . Theoretically, the complexity bound we achieve in (4) is the best one among current penalized DRO algorithms, delineating a clear dependence on smoothness constants of the coupled term and strong convexity constants of the individual terms in (2). Practically, DRAGO has a single hyperparameter and operates on any closed, convex uncertainty set for which the map l ↦→ arg max q ∈Q { q ⊤ l -νD ( q ∥ 1 /n ) } is efficiently computable. DRAGO is also of general conceptual interest as a stochastic variance-reduced primal-dual algorithm for min-max problems. It delicately combines randomized and cyclic updates, which effectively address the varying dimensions of the two problems (see Sec. 2). The theoretical guarantees of the algorithm are explained in Sec. 3. Numerical performance benchmarks are shown in Sec. 4.

2 The DRAGO Algorithm

We present here the D istributionally R obust A nnular G radient O ptimizer (DRAGO). While similar in spirit to a primal-dual proximal gradient method with a stochastic flavor, there are several innovations that allow the algorithm to achieve its superior complexity guarantee. These include using 1) minibatch stochastic gradient estimates to improve the trade-off between the per-iteration complexity and required number of iterations (especially when n ≫ d ), 2) a combination of randomized and cyclically updated components in the primal and dual gradient estimates, and 3) a novel regularization term in the primal update which reduces variance in the gradient estimate (i.e., coupled variance reduction). Here, we describe the algorithm in a manner that helps elucidate the upcoming theoretical analysis (Sec. 3). On the other hand, in Appx. D, we present an alternate description of DRAGO that is amenable to direct implementation in code.

Notation & Terminology Let ψ : R n → R ∪ { + ∞} be a proper, convex function such that Q ⊆ dom( ψ ) := { q ∈ R n : ψ ( q ) < + ∞} . Let ψ have a non-empty subdifferential for each q ∈ Q , and denote by ∇ ψ a map from q ∈ Q to an arbitrary but consistently chosen subgradient in ∂ψ ( q ) . We denote the Bregman divergence generated by ψ as ∆ ψ ( q, ¯ q ) = ψ ( q ) -ψ (¯ q ) -⟨∇ ψ (¯ q ) , q -¯ q ⟩ . We employ the Bregman divergence in this way for purely technical reasons, and for common cases this version will not be invoked. Finding a minimizer of (2) is equivalent to finding a saddle-point ( w ⋆ , q ⋆ ) ∈ W × Q which satisfies

Method

Assumptions

Uncertainty Set

Runtime / Global Complexity (Big-

˜

O

)

Sub-Gradient Method

ℓ

i

is

G

-Lipschitz

∥

w

-

w

′

∥

2

R

(if included)

µ >

0

(if included)

∇

ℓ

i

is

L

-Lipschitz and

ν >

0

(if included)

Support Constrained

Support Constrained

Support Constrained

nd

·

(

GR

)

2

ε

-

2

nd

·

G

2

µ

-

1

ε

-

1

nd

·

µ

-

1

(

L

+

µ

+

nG

2

/ν

)

log(1

/ε

)

L

CVaR-SGD

†

L

χ

2

†

L

χ

2

-

pen-SGD

†

[Levy et al., 2020]

ℓ

i

is

G

-Lipschitz and in

[0

, B

]

∥

w

-

w

′

∥

2

R

for all

w,w

′

∈

W

ν >

0

(if included)

θ

-CVaR

ρ

-ball in

χ

2

-divergence

χ

2

-divergence penalty

min

{

n, B

2

θ

-

1

ε

-

2

}

d

·

(

GR

)

2

ε

-

2

min

{

n,

(1 +

ρ

)

B

2

ε

-

2

}

d

·

(

GR

)

2

ε

-

2

min

{

n, B

2

ν

-

1

ε

-

1

}

d

·

(

GR

)

2

ε

-

2

BROO

∗

BROO

∗

[Carmon and Hausler, 2022]

ℓ

i

is

G

-Lipschitz

∥

w

-

w

′

∥

2

R

for all

w,w

′

∈

W

∇

ℓ

i

is

L

-Lipschitz (if included)

1

-ball in

f

-divergence

1

-ball in

f

-divergence

nd

·

(

GR

)

2

/

3

ε

-

2

/

3

+

d

(

GR

)

2

ε

-

2

nd

·

(

GR

)

2

/

3

ε

-

2

/

3

+

n

3

/

4

d

(

GRε

-

1

+

L

1

/

2

Rε

-

1

/

2

)

LSVRG

LSVRG

[Mehta et al., 2023]

ℓ

i

is

G

-Lipschitz

∇

ℓ

i

is

L

-Lipschitz

µ >

0

,

ν >

0

,

κ

:= (

L

+

µ

)

/µ

Spectral Risk Measures (

ν

small)

Spectral Risk Measures (

ν

Ω(

nG

2

/µ

)

)

None

(

n

+

κ

Q

κ

)

d

log(1

/ε

)

Prospect

Prospect

[Mehta et al., 2024]

ℓ

i

is

G

-Lipschitz

∇

ℓ

i

is

L

-Lipschitz

µ >

0

,

ν >

0

,

κ

:= (

L

+

µ

)

/µ, δ

=

G

2

/

(

µν

)

Spectral Risk Measures (

ν

small)

Spectral Risk Measures (

ν

Ω(

nG

2

/µ

)

)

n

(

n

+

d

{

nδ

+

κnq

, n

3

δ

2

κ

2

, n

3

δ

3

}

log(1

/ε

)

(

n

+

κ

Q

κ

) (

n

+

d

)log(1

/ε

)

DRAGO (Ours)

ℓ

i

is

G

-Lipschitz

∇

ℓ

i

is

L

-Lipschitz

µ >

0

,

ν >

0

,

b

:=

Support Constrained

n

(

d

+

κ

Q

L/µ

+

d

√

nG

2

/

(

µν

)

)

log(1

/ε

)

We denote the Jacobian of ℓ := ( ℓ 1 , . . . , ℓ n ) at w as ∇ ℓ ( w ) ∈ R n × d . We refer to the gradient of the coupled term q ⊤ ℓ ( w ) of (2) with respect to w and q (that is, ∇ ℓ ( w ) ⊤ q and ℓ ( w ) ) as the primal and dual gradients, respectively. In both the definition of the algorithm as well as the analysis, we will consider a sequence of positive constants ( a t ) t ≥ 1 with the additional values a 0 = 0 . The partial sums of this sequence will be denoted A t = ∑ t τ =0 a τ . Here, t represents the iteration counter while the convergence rate will be proportional to A -1 t (see Sec. 3), so we wish for the sequence to grow as fast as possible. As mentioned in Sec. 1, d is the primal dimension, n is the dual dimension as well as the sample size, and b denotes a batch size that divides n for ease of presentation.

Algorithm Description Wespecify the algorithm by recursively defining a sequence of primal-dual iterates { ( w t , q t ) t ≥ 1 } that achieve a particular convergence guarantee. First, we may assume that 0 ∈ W without loss of generality, so fix w 0 = 0 and q 0 = 1 /n . For any t ≥ 1 , we introduce v P t -1 and v D t as to-be-specified stochastic gradient estimates of the quantities ∇ ℓ ( w t -1 ) ⊤ q t -1 and ℓ ( w t ) , respectively. We choose to update w t before q t , so that v P t -1 may depend only on the history { ( w τ , q τ } t -1 τ =0 , whereas v D t may depend additionally on w t as well. Letting ( C t ) t ≥ 1 and ( c t ) t ≥ 1 be to-be-specified sequences of positive constants with C 0 = c 0 = 0 , we employ primal-dual proximal approach guided by the updates

For the primal update, despite the non-standard term (6), w t can be computed easily in closed form when W = R d . Otherwise, retrieving w t relies on computing an ℓ 2 -norm projection onto W . On the other hand, the maximization (7) can often be solved exactly or to high accuracy using methods specific to each uncertainty set; we describe these in detail in Appx. D.2. The randomized and cyclic coordinate-wise components mentioned above are contained within the upcoming formulas for v P t -1 and v D t . After specifying these vectors, the constants ( C t , c t ) will be chosen in the analysis to recover the final algorithm.

Next, we describe the computation of v P t -1 and v D t , which will rely on quantities that are stored by the algorithm along its iterations. We store three tables of values ˆ ℓ t ∈ R n , ˆ q t ∈ R n , and ˆ g ∈ R n × d , which are approximations of ℓ ( w t ) , q t , and ∇ ℓ ( w t ) , respectively. Before explaining how these tables are updated, consider the data indices { 1 , . . . , n } to be partitioned into n/b blocks, written ( B 1 , . . . , B n/b ) for B K = ( K -b +1 , . . . , Kb ) . On each iteration t , we randomly sample independent block indices I t , J t ∼ Unif[ n/b ] and define

As for the tables of approximations, we update them on each iteration without suffering the O ( nd ) computational cost of querying every first-order oracle ( ℓ 1 , ∇ ℓ 1 ) , . . . , ( ℓ n , ∇ ℓ n ) . We set ( ˆ ℓ 0 , ˆ q 0 , ˆ g 0 ) = ( ℓ ( w 0 ) , q 0 , ∇ ℓ ( w 0 )) , and for iteration t ≥ 1 update block K t := ( n/b ) mod t +1 via

While q t in particular is always known by the algorithm, we use the approximation ˆ q t which is possibly dual infeasible for every t ≥ 1 , but will approach q t as t grows large. Using this approximation is essential to controlling the per-iteration time complexity, as described below. In addition, the design of (8) and (9) is grounded in the long line of work on incremental methods (both deterministic and randomized). The table of past gradients updated cyclically resembles IAG [Blatt et al., 2007], whereas the randomized component resembles methods such as SAGA [Defazio et al., 2014] and stochastic PDHG [Chambolle et al., 2018]. The full algorithm description is given in Algorithm 1. In the next section, we show that ( a t , C t , c t ) t ≥ 0 can be determined by a single hyperparameter.

Computational Complexity Wealso discuss the per-iteration time complexity and the global space complexity of Algorithm 1, whereas the number of required iterations for ε -suboptimality is given in the next section. We see that the per-iteration time complexity is O ( n + bd ) , as we query b first-order oracles in both the primal and dual updates and all other operations occur on n -length or d -length vectors. While we need to employ O ( nd ) operations to compute ˆ g ⊤ t -1 ˆ q t -1 in (8) when t = 1 , this quantity can be maintained with O ( bd ) operations in every subsequent iteration as only b rows of ˆ q t and ˆ g t are editted in each iteration. For space complexity, we may consider storing the entire gradient table ˆ g t in memory, resulting in an O ( nd ) complexity. However, due to our time complexity calculation, we may reduce the space complexity to O ( n + bd ) by storing only w t -1 , . . . , w t -n/b and recomputing the relevant values of ˆ g t in (8) in every iteration. Therefore, the use of block-cyclic updates in the historical tables may significantly reduce the space complexity as compared to randomized updates (as in Defazio et al. [2014]).

3 Theoretical Analysis

We provide the theoretical convergence rate and global complexity of DRAGO along with technical highlights of the proof which may be of independent interest.

Convergence Analysis We measure suboptimality using the primal-dual gap :

where the saddle point ( w ⋆ , q ⋆ ) exists under Asm. 1. Note that γ t ≥ 0 , as it is the sum of nonnegative quantities

To state the main result, define E t as the conditional expectation over ( I t , J t ) given ( w t -1 , q t -1 ) and E 1 as the marginal expectation over the optimization trajectory. Consider the following assumptions.

Algorithm 1 D istributionally R obust A nnular G radient O ptimizer (DRAGO)

Input: Sequence of constants ( a t , C t , c t ) t ≥ 0 , block size b ∈ { 1 , . . . , n } , number of iterations T . Initialize w 0 = 0 d , q 0 = 1 /n , ˆ ℓ 0 = ℓ ( w 0 ) , ˆ g 0 = ℓ ( w 0 ) , and ˆ q 0 = q 0 .

Sample blocks I t and J t uniformly on [ n/b ] , and define K t = t mod ( n/b ) + 1 .

Primal Update:

Dual Update:

Update table ( ˆ ℓ t,k , ˆ g t,k ) to ( ℓ k ( w t ) , ∇ ℓ k ( w t )) if k ∈ B K t or ( ˆ ℓ t -1 ,k , ˆ g t -1 ,k ) if k / ∈ B K t .

end for

Update table ˆ q t,k to q t,k if k ∈ B K t or ˆ q t -1 ,k if k / ∈ B K t .

Assumption 1. Let ℓ 1 , . . . , ℓ n be G -Lipschitz continuous and L -smooth, in that for all i ∈ [ n ] ,

i.e., each ℓ i is G -Lipschitz continuous and L -smooth with respect to ∥·∥ 2 . Let µ > 0 and ν > 0 , and let q ↦→ D ( q ∥ 1 n /n ) be 1 -strongly convex with respect to ∥·∥ 2 . Finally, Q is closed, convex, and contains 1 /n .

Regarding Asm. 1, an example of a loss function satisfying both Lipschitzness and smoothness is given by the Huber loss used in robust statistics [Huber, 1981]. Another setting in which both assumptions are satisfied is when the domain W is compact, as smoothness will imply Lipschitz continuity. Compactness is a common assumption when pursuing statistical guarantees such as uniform convergence. As for the assumption of strong convexity, we describe modifications of the algorithm when µ = 0 or ν = 0 in Appx. C.4 along with corresponding changes in the analysis.

Theorem 2. For a constant α > 0 , define the sequence

along with its partial sum A t = ∑ t τ =1 a τ . Under Asm. 1, there is an absolute constant C such that using the parameter

the iterates of Algorithm 1 satisfy:

We can compute a point ( w T , q T ) achieving an expected gap no more than ε with bigO complexity

By dividing the result of Thm. 2 by A T , we see that both the expected gap and expected distanceto-optimum in the primal and dual sequences decay geometrically in T . By plugging in b = n/d we get the following runtime for Algorithm 1:

Note in particular the individual dependence on the condition numbers L/µ and nG 2 / ( µν ) , as opposed to the max-over-min-type condition numbers such as those achievable by generic primaldual or variational inequality methods (see Appx. B.3). The proof of Thm. 2 is provided in Appx. C, along with a high-level overview in Appx. C.1. Our analysis relies on controlling a t γ t by bounding a t L ( w ⋆ , q t ) above and bounding a t L ( w t , q ⋆ ) below. Key technical steps occur in the lower bound. We first apply that w ↦→ a t q ⊤ t ℓ ( w ) is convex to produce a linear underestimator of the function supported at w t , and use that w t is the minimizer of a strongly convex function, so

Note that the terms above will be negated when combining the upper and lower bounds. The labeled terms are highly relevant in the analysis, with the second being non-standard. By expanding the definition of v P t and adding and subtracting w t -1 , we write

When choosing the learning rate η correctly, the first two terms will telescope in expectation (see Lem. 11). The third term, after applying Young's inequality, requires controlling 1 b ∑ i ∈ I t ∥∇ ℓ i ( w t -1 ) q t -1 -ˆ g t -2 ,i ˆ q t -2 ,i ∥ 2 2 in expectation, which we dub the 'primal noise bound' (Lem. 5). When we combine the upper and lower bounds, we get a similar inner product term a t 〈 q ⋆ -q t , ℓ ( w t ) -v D t 〉 , and mirroring the arguments above, we encounter the term 1 b ∑ j ∈ J t ( ℓ j ( w t ) -ˆ ℓ t -1 ,j ) 2 2 which also requires a 'dual noise bound' (Lem. 6). Without loss of generality, assume that the blocks are ordered such that

By computing the conditional expectation E t [ · ] := E [ ·| w t -1 ] , we have that

where the second line follows from (14) and the third from G -Lipschitzness. This will telescope with the second term introduced in (13), showing the importance of the regularization. The argument follows similarly even when (14) does not hold, as the blocks can simply be 'renamed' to achieve the final bound. While the proof is technical, this core idea guides the analysis and the algorithm design.

4 Experiments

In this section, we provide numerical benchmarks to measure DRAGO against baselines in terms of evaluations of each component { ( ℓ i , ∇ ℓ i ) } n i =1 and wall clock time. We consider regression and classification tasks. Letting ( x i , y i ) denote a feature-label pair, we have that each ℓ i represents the

0

50000

100000

10

-

6

10

-

3

10

0

Suboptimality

yacht

0

50000

100000

10

-

7

10

-

4

10

-

1

energy

0

50000

100000

10

-

6

10

-

3

10

0

concrete

0

1

2

10

-

6

10

-

3

10

0

yacht

0

1

2

10

-

6

10

-

3

10

0

energy

0

1

2

10

-

6

10

-

3

10

0

concrete

0

50000

100000

First-Order Oracle Evaluations

10

-

4

10

-

2

10

0

Suboptimality

acsincome

0

50000

100000

First-Order Oracle Evaluations

10

-

6

10

-

3

10

0

kin8nm

0

50000

100000

First-Order Oracle Evaluations

10

-

7

10

-

4

10

-

1

power

0

5

Wall-Clock Time (Seconds)

10

-

4

10

-

2

10

0

acsincome

0

5

Wall-Clock Time (Seconds)

10

-

6

10

-

3

10

0

kin8nm

0

5

Wall-Clock Time (Seconds)

10

-

7

10

-

4

10

-

1

power

SGD

LSVRG

Drago (

b

= 16

)

Drago (

b

=

n/d

)

Drago (

b

= 1

)

squared error loss or multinomial cross-entropy loss, given by

respectively. In the latter case, we denote w = ( w 1 , . . . , w C ) ∈ R C × d , indicating multiclass classification with label set Y := { 1 , . . . , C } . We show results for the conditional value-at-risk (CVaR) [Rockafellar and Royset, 2013] in this section, exploring the effect of sample size n , dimensionality d , and regularization parameter ν on optimization performance. The parameter ν in particular has interpretations both as a conditioning device (as it is inversely related to the smoothness constant of w ↦→ max q ∈Q L ( w,q ) and as a robustness parameter, as it controls the essential size of the uncertainty set. Detailed experimental settings are contained in Appx. E, including additional experiments with the χ 2 -divergence ball uncertainty set. The code to reproduce these experiments can be found at https://github.com/ronakdm/drago.

We compare against baselines that can be used on the CVaR uncertainty set: distributionally robust stochastic gradient descent (SGD) [Levy et al., 2020] and LSVRG [Mehta et al., 2023]. For SGD, we use a batch size of 64 and for LSVRG we use the default epoch length of n . For DRAGO, we investigate the variants in which b is set to 1 and b = n/d a priori, as well as cases when b is a tuned hyperparameter. On the y -axis, we plot the primal gap

where we approximate L ( w ⋆ , q ⋆ ) by running LBFGS [Nocedal and Wright, 1999] on the primal objective until convergence. On the x -axis, we display either the exact number of calls to the firstorder oracles of the form ( ℓ i , ∇ ℓ i ) or the wall clock time in seconds. We fix µ = 1 but vary ν to study its role as a conditioning parameter, which is especially important as prior work establishes different convergence rates for different values of ν (see Tab. 1).

4.1 Regression with Large Block Sizes

In this experiment, we consider six regression datasets, named yacht ( n = 244 , d = 6 ) [Tsanas and Xifara, 2012], energy ( n = 614 , d = 8 ) [Baressi Segota et al., 2020], concrete ( n = 824 , d = 8 ) [Yeh, 2006], acsincome ( n = 4000 , d = 202 ) [Ding et al., 2021], kin8nm ( n = 6553 , d = 8 ) [Akujuobi and Zhang, 2017], and power ( n = 7654 , d = 4 ) [T¨ ufekci, 2014]. In each case, there is a univariate, real-valued output. Notice that most datasets, besides acsincome, are low-dimensional as compared to their sample size. Thus, the default block size n/d becomes relatively large, imposing an expensive per-iteration cost in terms of oracle queries. However, when the block size is high, the stochastic gradient estimates in each iteration have lower variance and the table components are updated more frequently, which could improve convergence in principle. The main question of this

0

100

200

300

First-Order Oracle Evaluations (K)

10

-

5

10

-

3

10

-

1

Suboptimality

ν

=

1.0

0

100

200

300

First-Order Oracle Evaluations (K)

10

-

5

10

-

3

10

-

1

ν

=

0.01

0

100

200

300

First-Order Oracle Evaluations (K)

10

-

5

10

-

3

10

-

1

ν

=

0.001

0

2

4

Wall-Clock Time (Seconds)

10

-

2

10

-

1

10

0

ν

=

1.0

0

2

4

Wall-Clock Time (Seconds)

10

-

3

10

-

2

10

-

1

10

0

ν

=

0.01

0

2

4

Wall-Clock Time (Seconds)

10

-

4

10

-

3

10

-

2

10

-

1

10

0

ν

=

0.001

section is whether DRAGO efficiently manages this trade-off via the block size parameter b . Results for gradient evaluations and wall clock time are on the left and right panels of Fig. 2, respectively.

Results The DRAGO variant for b = n/d is not included on the left plot, as the number of queries (almost 2,000 in the case of power) penalizes its performance heavily. Still, the same variant performs best or near best on all datasets in terms of wall clock time (right plot). Thus, if the computation of the queries is inexpensive enough, DRAGO can achieve the lowest suboptimality within a fixed time budget. This is most striking in the case of kin8nm, in which DRAGO achieves 10 -7 primal gap within 1 second, versus LSVRG which is only able to reach within 10 -2 of the minimum in the same amount of time. We also experiment with tuning b to reach a balance between the cost of queries and distance to optimum in the left plot with the b = 16 variant. In the datasets with n ≤ 1 , 000 , DRAGO can match the performance of baselines with only b = 1 , whereas in the larger datasets, a batch size of 16 is needed to be comparable.

4.2 Text Classification Under Ill-Conditioning

We consider a natural language processing example using the emotion dataset [Saravia et al., 2018], which is a classification task consisting of six sentiment categories: sadness, anger, love, fear, joy, and surprise. To featurize the text, we fine-tune a pre-trained BERT network [Devlin et al., 2019] on a held-out set of 8,000 training examples to learn a vectorial representation. We then use a disjoint subset of 8,000 training points and apply PCA to reduce them to 45-dimensional vectors. Because of the six classes this results in d = 270 parameters to learn. To study the effect of dual regularization, we consider ν ∈ { 1 . 0 , 0 . 01 , 0 . 001 } . As ν decreases, the dual solution may shift further from uniformity, and potentially increase the distributional robustness of the learned minimizer. However, the objective can also become poorly conditioned, introducing a key trade-off between optimization and statistical considerations when selecting ν . The results are in Fig. 3.

Results The run time required for LSVRG to make 500K gradient evaluations is too large to be considered. We also observe that LSVRG is vulnerable to ill-conditioned objectives, as it is outperformed by SGD for smaller values of ν in terms of wall clock time. Within 4 seconds, DRAGO can achieve close to a 10 -5 primal suboptimality gap while the gaps of SGD and LSVRG are 2 to 3 orders of magnitude larger in the same amount of time. We hypothesize that because the dual variables in LSVRG are updated once every n iterations, the primal gradient estimates may accrue excessive bias. DRAGO with b = n/d , making ∼ 30 individual first-order queries per iteration, is performant in terms of oracle queries and wall clock time even as ν drops by 3 orders of magnitude.

5 Conclusion

We proposed DRAGO, a stochastic primal-dual algorithm for solving a host of distributionally robust optimization (DRO) problems. The method achieves linear convergence without placing conditions on the dual regularizer, and its empirical performance remains strong across varying settings of the sample size n , dimension d , and the dual regularization parameter ν . The method combines ideas of variance reduction, minibatching, and cyclic coordinate-style updates even though the dual feasible set (a.k.a. the uncertainty set) is non-separable. Opportunities for future work include extensions to non-convex settings and applications to min-max problems beyond distributional robustness, such as missing data imputation and fully composite optimization.

Acknowledgements This work was supported by NSF DMS-2023166, CCF-2019844, DMS2134012, NIH, IARPA 2022-22072200003, U. S. Office of Naval Research under award number N00014-22-1-2348. Part of this work was done while R. Mehta and Z. Harchaoui were visiting the Simons Institute for the Theory of Computing.

Broader Impact Distributionally robust optimization (DRO) within machine learning is heavily motivated by problems in artificial intelligence (AI) safety, such as mitigating catastrophic performance of models on minority groups or end-users. While this work is focused on theoretical and algorithmic aspects of the problem, we intend to increase accessibility and scalability for downstream applications as well.

References

Akujuobi and X. Zhang. Delve: A Dataset-Driven Scholarly Search and Analysis System. SIGKDD Explor. Newsl. , 19, 2017.

Alacaoglu and Y. Malitsky. Stochastic variance reduction for variational inequality methods. In Conference on Learning Theory , pages 778-816. PMLR, 2022.

Alacaoglu, V. Cevher, and S. J. Wright. On the complexity of a practical primal-dual coordinate method. arXiv preprint arXiv:2201.07684 , 2022.

Baressi Segota, N. Andelic, J. Kudlacek, and R. Cep. Artificial Neural Network for Predicting Values of Residuary Resistance per Unit Weight of Displacement. Journal of Maritime & Transportation Science , 57, 2020.

Blanchet, Y. Kang, and K. Murthy. Robust Wasserstein Profile Inference and Applications to Machine Learning. Journal of Applied Probability , 56, 2019.

Blatt, A. O. Hero, and H. Gauchman. A Convergent Incremental Gradient Method with a Constant Step Size. SIAM Journal on Optimization , 2007.

Cai, C. Song, S. Wright, and J. Diakonikolas. Cyclic Block Coordinate Descent With Variance Reduction for Composite Nonconvex Optimization. In ICML , 2023.

Cai, A. Alacaoglu, and J. Diakonikolas. Variance reduced Halpern iteration for finite-sum monotone inclusions. In ICLR , 2024.

Carmon and D. Hausler. Distributionally Robust Optimization via Ball Oracle Acceleration. In NeurIPS , 2022.

Chambolle, M. J. Ehrhardt, P. Richt´ arik, and C.-B. Schonlieb. Stochastic Primal-Dual Hybrid Gradient Algorithm with Arbitrary Sampling and Imaging Applications. SIAM Journal on Optimization , 28(4):2783-2808, 2018.

Condat. Fast projection onto the simplex and the ℓ 1 -ball. Mathematical Programming , 2016.

Defazio, F. Bach, and S. Lacoste-Julien. SAGA: A Fast Incremental Gradient Method With Support for Non-Strongly Convex Composite Objectives. NeurIPS , 27, 2014.

Devlin, M.-W. Chang, K. Lee, and K. Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies , 2019.

Ding, M. Hardt, J. Miller, and L. Schmidt. Retiring Adult: New Datasets for Fair Machine Learning. In NeurIPS , volume 34. Curran Associates, Inc., 2021.

S. Du, G. Gidel, M. I. Jordan, and C. J. Li. Optimal Extragradient-Based Bilinearly-Coupled Saddle-Point Optimization, 2022.

He, Z. Harchaoui, Y. Wang, and L. Song. Point Process Estimation with Mirror Prox Algorithms. Applied Mathematics & Optimization , 2020.

J. Huber. Robust statistics . Wiley New York, 1981.

Johnson and T. Zhang. Accelerating Stochastic Gradient Descent using Predictive Variance Reduction. In NeurIPS , volume 26, 2013.

Kallus, X. Mao, K. Wang, and Z. Zhou. Doubly Robust Distributionally Robust Off-Policy Evaluation and Learning. In ICML , 2022.

Kovalev, A. Gasnikov, and P. Richt´ arik. Accelerated Primal-Dual Gradient Method for Smooth and Convex-Concave Saddle-Point Problems with Bilinear Coupling. In NeurIPS , 2022.

Kuhn, P. M. Esfahani, V. A. Nguyen, and S. Shafieezadeh-Abadeh. Wasserstein Distributionally Robust optimization: Theory and Applications in Machine Learning. Operations research & management science in the age of analytics , 2019.

Kumar, K. A. Majmundar, D. M. Nagaraj, and A. Suggala. Stochastic Re-weighted Gradient Descent via Distributionally Robust Optimization. TMLR , 2024.

Levy, Y. Carmon, J. Duchi, and A. Sidford. Large-Scale Methods for Distributionally Robust Optimization. In NeurIPS , 2020.

J. Li, A. Yuan, G. Gidel, Q. Gu, and M. I. Jordan. Nesterov meets optimism: rate-optimal separable minimax optimization. In ICML , 2023.

Lin, C. Jin, and M. I. Jordan. Near-Optimal Algorithms for Minimax Optimization. In COLT , 2020.

Z. Liu, B. Haghgoo, A. S. Chen, A. Raghunathan, P. W. Koh, S. Sagawa, P. Liang, and C. Finn. Just Train Twice: Improving Group Robustness without Training Group Information. In ICML , 2021.

Liu, Q. Bai, J. Blanchet, P. Dong, W. Xu, Z. Zhou, and Z. Zhou. Distributionally Robust Q -Learning. In ICML , 2022.

Lotidis, N. Bambos, J. Blanchet, and J. Li. Wasserstein Distributionally Robust Linear-Quadratic Estimation under Martingale Constraints. In AISTATS , 2023.

Mehta, V. Roulet, K. Pillutla, L. Liu, and Z. Harchaoui. Stochastic Optimization for Spectral Risk Measures. In AISTATS , 2023.

Mehta, V. Roulet, K. Pillutla, and Z. Harchaoui. Distributionally Robust Optimization with Bias and Variance Reduction. In ICLR , 2024.

Namkoong and J. C. Duchi. Stochastic Gradient Methods for Distributionally Robust Optimization with f-divergences. In NeurIPS , 2016.

Namkoong and J. C. Duchi. Variance-based Regularization with Convex Objectives. In NeurIPS , 2017.

Nocedal and S. J. Wright. Numerical Optimization . Springer, 1999.

Palaniappan and F. Bach. Stochastic Variance Reduction Methods for Saddle-Point Problems. NeurIPS , 29, 2016.

T. Rockafellar and J. O. Royset. Superquantiles and Their Applications to Risk, Random Variables, and Regression. In Theory Driven by Influential Applications . Informs, 2013.

Sapkota, D. Wang, Z. Tao, and Q. Yu. Distributionally Robust Ensemble of Lottery Tickets Towards Calibrated Sparse Network Training. In NeurIPS , 2023.

Saravia, H.-C. T. Liu, Y.-H. Huang, J. Wu, and Y.-S. Chen. CARER: Contextualized Affect Representations for Emotion Recognition. In EMNLP , 2018.

D. Sharma, M. Toubeh, L. Zhou, and P. Tokekar. Risk-Aware Planning and Assignment for Ground Vehicles using Uncertain Perception from Aerial Vehicles. In 2020 IEEE/RSJ IROS , 2020.

Song, S. J. Wright, and J. Diakonikolas. Variance reduction via primal-dual accelerated dual averaging for nonsmooth convex finite-sums. In ICML , 2021.

Song, C. Y. Lin, S. Wright, and J. Diakonikolas. Coordinate Linear Variance Reduction for Generalized Linear Programming. In NeurIPS , 2022.

Tsanas and A. Xifara. Accurate Quantitative Estimation of Energy Performance of Residential Buildings Using Statistical Machine Learning Tools. Energy and Buildings , 49, 2012.

T¨ ufekci. Prediction of Full Load Electrical Power Output of a Base Load Operated Combined Cycle Power Plant using Machine Learning Methods. International Journal of Electrical Power &Energy Systems , 60, 2014.

Wang, N. Si, J. Blanchet, and Z. Zhou. A Finite Sample Complexity Bound for Distributionally Robust Q-learning. In AISTATS , 2023a.

Wang, L. Shen, T. Liu, T. Duan, Y. Zhu, D. Zhan, D. Doermann, and M. Gao. Defending against Data-Free Model Extraction by Distributionally Robust Defensive Training. In NeurIPS , 2023b.

Wiles, S. Gowal, F. Stimberg, S.-A. Rebuffi, I. Ktena, K. D. Dvijotham, and A. T. Cemgil. A fine-grained analysis on distribution shift. In ICLR , 2022.

Xie, L. Luo, Y. Lian, and Z. Zhang. Lower Complexity Bounds for Finite-Sum Convex-Concave Minimax Optimization Problems. In ICML , 2020.

Yang, S. Zhang, N. Kiyavash, and N. He. A Catalyst Framework for Minimax Optimization. In NeurIPS , 2020.

Yang, Y. Guo, P. Xu, A. Liu, and A. Anandkumar. Distributionally Robust Policy Gradient for Offline Contextual Bandits. In AISTATS , 2023.

Yeh. Analysis of Strength of Concrete Using Design of Experiments and Neural Networks. Journal of Materials in Civil Engineering , 18, 2006.

Yu, T. Lin, E. V . Mazumdar, and M. Jordan. Fast Distributionally Robust Learning with VarianceReduced Min-Max Optimization. In AISTATS , 2022.

Yu, L. Dai, S. Xu, S. Gao, and C. P. Ho. Fast Bellman Updates for Wasserstein Distributionally Robust MDPs. In NeurIPS , 2023.

Zhang, M. Hong, and S. Zhang. On lower iteration complexity bounds for the convex concave saddle point problems. Mathematical Programming , 194, 2022.

Appendix

Appx. A contains all notation introduced throughout the paper. Appx. B discusses convergence rate comparisons in detail with contemporary work. The full convergence analysis of DRAGO is given in Appx. C. A description of the algorithm amenable to implementation is given in Appx. D, whereas experimental settings are described in detail within Appx. E.

Table of Contents

A Notation

B Comparisons to Existing Literature

In this appendix, we compare our work to existing literature along two axes: 1) distributional robust optimization (DRO), and 2) primal-dual algorithms for saddle-point problems. In the first category, we are primarily concerned with questions of practical and statistical interest, such as which uncertainty sets can be used, how the size of the uncertainty set affects the convergence rate, and what assumptions are needed on the distribution of losses. In the second category, we discuss computational complexity under various assumptions such as smoothness and strong convexity of the objective.

B.1 Directly Using Gradient Descent

In the penalized case, we add that the objective w ↦→ max q ∈Q L ( w,q ) is ( L + µ + nG 2 ν ) -smooth when the losses ℓ i are G -Lipschitz continuous and L -smooth. For this reason, we may consider simply applying full-batch gradient descent to this objective, which is included in our comparisons. To see why this smoothness condition holds, define h ( l ) := max q ∈Q { q ⊤ l -νD ( q ∥ 1 /n ) } , so that when q ↦→ D ( q ∥ 1 /n ) is 1 -strongly convex with respect to ∥·∥ 2 2 , it holds that ∇ h is (1 /ν ) -Lipschitz continuous with respect to ∥·∥ 2 2 , and that ∇ h ( l ) is non-negative and sums to one (i.e. is a probability mass function on [ n ] ). Then, by the chain rule, for any w 1 , w 2 ∈ W , we have that

Thus, when referring to the gradient descent on w ↦→ max q ∈Q L ( w,q ) = h ( ℓ ( w )) + µ 2 ∥ w ∥ 2 2 , we reference the smoothness constant ( L + µ + nG 2 ν ) .

B.2 Distributionally Robust Optimization (DRO)

Examples of DRO Problems Our problem

accommodates several settings of interest across machine learning. For example, f -DRO with parameter ρ [Namkoong and Duchi, 2016, Carmon and Hausler, 2022] results by defining the uncertainty set and penalty as

where D f ( q ∥ p ) = ∑ n i =1 p i f ( q i /p i ) denotes an f -divergence generated by f (which is always well-defined when p i = 1 /n ). Common examples include the Kullback-Leibler (KL) divergence, generated by f KL ( x ) = -x log x and the χ 2 -divergence generated by f χ 2 ( x ) = ( x -1) 2 . Spectral risk measures [Mehta et al., 2023] are parametrized by an n -length vector σ = ( σ 1 , . . . , σ n ) such that 0 ≤ σ 1 ≤ . . . ≤ σ n and ∑ n i =1 σ i = 1 . The penalty in that setting may also be in the form of an f -divergence [Mehta et al., 2024], so that

where conv ( · ) denotes the convex hull. The most notable example of such an uncertainty set is the θ -conditional value-at-risk, or CVaR [Rockafellar and Royset, 2013], wherein the largest θn values of σ are set to 1 / ( nθ ) and the remaining are set to zero (with a fractional component if θn is not an integer). Finally, Wasserstein-DRO [Kuhn et al., 2019, Blanchet et al., 2019, Yu et al., 2022] with parameter δ typically sets the uncertainty set to be a δ -ball { Q : W c ( Q,P n ) ≤ δ } in Wasserstein distance, where Q is a probability measure, P n is the empirical measure of the training data, and

W c is the Wasserstein distance associated to cost function c . This differs from f -DRO and spectral risk minimization because in the latter settings, the 'shifted' distribution Q is assumed to remain on the same n atoms as before, so that it may simply be specified by a probability vector q ∈ [0 , 1] n (resp. P n by 1 /n ). However, as shown by Yu et al. [2022], Wasserstein-DRO can be reformulated into a problem of the form (16) if the following conditions are satisfied: 1) the loss is of a generalized linear model ℓ i ( w ) = Ψ( ⟨ w,x i ⟩ , y i ) with feature-label pair ( x i , y i ) and discrepancy function Ψ , 2) the cost function c (( x, y ) , ( x ′ , y ′ )) is of the form ∥ x -x ′ ∥ 2 + β | y -y ′ | for β > 0 , and 3) the function Ψ is Lipschitz continuous with known constant.

Comparison of DRO Approaches The performance of classical and recent algorithms designed for the problems above is detailed in Tab. 3. The rightmost column displays the oracle complexity , meaning the number of queries to individual/component loss functions { ( ℓ i , ∇ ℓ i ) } as a function of the desired suboptimality ε . The desiderata in the large-scale setting is to decouple the contribution of the sample size n and the condition number in smooth, strongly convex settings ( µ > 0 and L < ∞ ) or the quantity 1 /ε in non-smooth or non-strongly convex settings. For readability, we encode dependence on n in red and dependence on 1 /ε in blue within Tab. 3. In certain cases, such as in the sub-gradient method, the dependence on n is understood as part of a smoothness constant. Similarly, because q max is of order 1 /n in the best case and 1 in the worst case, we interpret κ Q to play the role of a condition number that measures the size of the uncertainty set. That being said, κ Q is upper bounded by a constant independent of n in common cases. We collect examples of these results below.

Proposition 3. In (16) if Q is chosen to be the CVaR uncertainty set with tail probability θ , then κ Q ≤ 1 θ . If Q is chosen to be the χ 2 -DRO uncertainty set with radius ρ , then κ Q ≤ √ 1 + ρ .

Proof. For the θ -CVAR, we have that q i ∈ [0 , 1 / ( nθ )] for all i ∈ [ n ] . Thus,

For the χ 2 -DRO uncertainty set, we have by direct computation that

This implies that for any i ∈ [ n ] , we have that q 2 i ≤ (1 + ρ ) /n 2 which implies that

the result as desired.

The closest comparison to our setting is that of LSVRG [Mehta et al., 2023] and Prospect [Mehta et al., 2024]. Including DRAGO, these methods all achieve linear convergence on (16), under the assumption of strongly convex regularization in the primal objective. However, both LSVRG and Prospect demand a stringent lower bound of Ω( nG 2 /µ ) on the dual regularization parameter ν to achieve this rate, which essentially matches ours in this regime (as √ nG 2 / ( µν ) reduces to a constant). When this condition does not hold, however, LSVRG does not have a convergence guarantee while Prospect underperforms against the sub-gradient method. Furthermore, while Propsect has a single hyperparameter, LSVRG must search over both a learning rate and an epoch length. DRAGO, on the other hand, is the only method that achieves unconditional linear convergence and fully captures the dependence on the dual regularization parameter ν , which is a smoothness measure of the objective max q L ( · , q ) .

Moving to methods that are not linearly convergent, Levy et al. [2020] consider a variant of minibatch SGD that solves the maximization problem within the mini-batch itself. In Tab. 3, we include a term min { n, b ( ε ) } , where b ( ε ) denotes a required batch size. This is because Levy et al. [2020] measures complexities in terms of calls to first-order oracles for a loss summed across an arbitrary batch size. Thus, our comparison relies on multiplying the required batch size with the number of iterations required for a desired suboptimality (see Levy et al. [2020, Theorem 2], for example). In the setting in which there is a fixed training set of size n , we incorporate this requirement on the batch size in the complexity bound, to account for the case in which the theoretically required batch size grows so large that it exceeds the full batch size itself. Indeed, as commented by Carmon

and Hausler [2022], when the uncertainty set is large (i.e. θ is small for CVaR and ρ is large for χ 2 -divergence), the method may underperform against the sub-gradient method. This is true of methods such as in Mehta et al. [2024] and ours that depend on κ Q ≤ n . This indicates that across DRO settings, increased uncertainty set sizes can bring the performance of incremental methods arbitrarily close to that of the full-batch sub-gradient method. Finally, notice that DRAGO also has a dependence on the batch size b in Tab. 3. While it would appear that b = 1 would minimize oracle complexity, we describe in the next section how b can be chosen to minimize global complexity, which includes the cost of each oracle call as a function of the primal dimension d . Because of the diversity of settings, a coarser measure such as first-order oracle evaluations may be sufficient to compare methods along the DRO axis; we use the finer-grained global complexity to compare methods along the axis of saddle-point algorithms.

B.3 Primal-Dual Saddle Point Algorithms

For context, both (1) and (2) are coupled min-max problems with primal-dual variables ( w,q ) . Observe that the coupled term in the objective (depending on both w and q ) is not bilinear in general. Such objectives have received much less attention in the optimization literature on primal-dual methods than their bilinearly-coupled counterparts. For the bilinear setting, methods such as stochastic Chambolle-Pock-style algorithms [He et al., 2020, Song et al., 2021] use stochastic variance-reduced updates, and in particular, employ coordinate-style updates of a single dual variable per iteration. To illustrate the advantage of these updates, notice that the primal and dual dimensions are d and n , respectively. When n is significantly larger than d , and assuming that each ( ℓ i , ∇ ℓ i ) call comes at an O ( d ) cost, updating the dual variables at O ( n ) cost every iteration becomes the primary computational bottleneck. Coordinate-style updates can eliminate this dependence, but apply only when the dual feasible set Q decomposes as Q 1 × . . . ×Q n , i.e., is separable .

There is a long line of work on variance-reduced algorithms for general stochastic optimization problems [Johnson and Zhang, 2013, Defazio et al., 2014, Palaniappan and Bach, 2016, Cai et al., 2023] and structured and/or bilinearly coupled min-max problems [Yang et al., 2020, Song et al., 2021, Du et al., 2022, Kovalev et al., 2022, Alacaoglu et al., 2022]. However, few of these results can be directly applied to (2) due to non-bilinearity and non-separability of the dual feasible set. These issues motivated work on reformulating common DRO problems as bilinearly coupled minmax problems in Song et al. [2022]. However, the guarantees obtained in Song et al. [2022] are of the order O ( ε -1 ) .

Another viewpoint is that (2) can be written as a monotone variational inequality (VI) problem for the operator F ( w,q ) = ( ∇ w ℓ ( w ) ⊤ q + µw, -ℓ ( w ) + ν ∇ q D ( q ∥ 1 /n )) ∈ R d + n , where we assumed 1 -smoothness of D for ease of presentation. Under our assumptions, this operator is min { µ, ν } -strongly monotone and max { nG 2 , L } -Lipschitz continuous, and VI algorithms such as mirrorprox, Popov's method, and variance-reduced methods like Alacaoglu and Malitsky [2022], Cai et al. [2024] can be used. However, the max-over-min dependence on the individual Lipschitz constants and strong convexity constants is unfavorable compared to the individual condition numbers observed in (4). If F is called directly, the global complexity will be n 2 because q is n -dimensional along with n in the condition number. A finite sum approach could improve dependence on the number of oracle calls, but the global complexity would still be n 2 due to the non-separability of the dual feasible set and again the n -dimensional dual variables.

To summarize the points above and in order to make comparisons among methods for general minmax problems, we first collect aspects of (16) that make it a highly specialized problem in this regard. The major points include:

The objective has a finite-sum structure, in that it can be written as ∑ n i =1 f i ( w,q ) , where f i ( w,q ) := q i ℓ i ( w ) -νD ( q ∥ 1 /n ) + µ 2 ∥ w ∥ 2 2 .

The dimension of the dual variable q is equal to n , the number of functions in the sum (i.e. the sample size), and could be much larger than d .

The dual regularizer q ↦→ D ( q ∥ 1 /n ) is not necessarily smooth. This encompasses common statistical divergences such as the Kullback Leibler (KL).

The dual feasible set Q is non-separable, as any feasible dual iterate must sum to 1 .

For discussions in which linear convergence is a pre-requisite, it is typically assumed that ( w,q ) ↦→ L ( w,q ) is a smooth map, and finer-grained results depend on the individual Lipschitz continuity parameters of w ↦→∇ w L ( w,q ) , q ↦→∇ q L ( w,q ) , w ↦→∇ q L ( w,q ) , and q ↦→∇ w L ( w,q ) . To make the regularized DRO setting of (16) comparable to classical and contemporary saddle-point methods, we make the additional assumption in this section that the (rescaled) χ 2 -divergence penalty is used, so that D ( q ∥ 1 /n ) = 1 2 ∥ q -1 /n ∥ 2 2 and we may compute the smoothness constants as

In making this assumption, however, we emphasize that our results do not depend on (17) , (18) , and (19) being true . We assume here that calls to the oracles ( ℓ i , ∇ w ℓ i ) cost O ( d ) operations while calls to ∇ q L cost O ( n ) operations, making the total O ( n + d ) . This per-iteration complexity is a subtlety of DRO which is essential to recognize when comparing methods. With DRAGO, we also have a batch size parameter b , setting the complexity to O ( n + bd ) . By tuning b , we may achieve improvements in terms of global complexity over standard primal-dual methods. Results comparing linearly convergent methods in terms of the global complexity of elementary operations are given in Tab. 4. The table contains a 'half-life' column, which is the constant τ multiplied with log(1 /ε ) when describing the number of iterations needed to reach ε -suboptimality as τ log(1 /ε ) . Before comparisons, observe that the optimal batch size for DRAGO is b = n/d , in the sense that the global number of arithmetic operations is of order

This is an improvement over setting b = 1 , in which case the same number is

Next, the most comparable and recent setting is that of Li et al. [2023]. Notice that DRAGO is able to improve by a factor of d on the L µ term, as long as κ Q = O (1) .

While less comparable, we mention known lower bounds for completeness. In terms of lower bounds, since (16) enjoys a particular structure, such as decomposability into so-called marginal terms νD ( q ∥ 1 /n ) and µ 2 ∥ w ∥ 2 2 and a coupled term q ⊤ ℓ ( w ) , we are not necessarily constrained by the more general lower bounds of Zhang et al. [2022] and Xie et al. [2020]. For example, the proximal incremental first-order (PIFO) model of Xie et al. [2020] assumes that we observe first-order information from a single component in the sum; in DRAGO, using a batch size of n/d is required to achieve the desired improvement. Zhang et al. [2022], on the other hand, do not treat the finite sum class of problems considered here. Furthermore, we still list the per-iteration cost as O ( n + d ) because a single PIFO call to the dual gradient is of size n .

C Convergence Analysis of DRAGO

This section provides the convergence analysis for DRAGO. We first recall the quantities of interest and provide an alternate description of the algorithm that is useful for understanding the analysis. A high-level overview is given in Appx. C.1, and the remaining subsections comprise steps of the proof.

C.1 Overview

See Tab. 2 for a reference on notation used throughout the proof, which will also be introduced as it appears. Define q 0 = 1 /n and recall the objective

By strong convexity of w ↦→ max q L ( w,q ) with respect to ∥·∥ 2 and strong concavity of q ↦→ min w L ( w,q ) with respect to ∥·∥ 2 , a primal-dual solution pair ( w ⋆ , q ⋆ ) is guaranteed to exist and is unique, and thus we define

In other words, we have that max q ∈Q L ( w ⋆ , q ) = L ( w ⋆ , q ⋆ ) = min w ∈W L ( w,q ⋆ ) . We proceed to describe the algorithm, the optimality criterion, and the proof outline.

Algorithm Description First, we describe two sequences of parameters that are used to weigh various terms in the primal and dual updates. The parameters are set in the analysis, and the version of the algorithm in Appx. D is a description with these values plugged in. However, we keep them as variables in this section to better describe the logic of the proof.

Specifically, let ( a t ) t ≥ 1 be a sequence of positive numbers and define a 0 = 0 in addition. Denote A t = ∑ T s =1 a s . These will become the averaging sequence that will aggregate successive gaps γ t (see (27)) into the return value ∑ T t =1 a t γ t , which will be upper bounded by a constant in T (in expectation). We also define another similar sequence ( c t ) t ≥ 1 , and define C t = A t -( n/b -1) c t for batch size b . Weassume for simplicity that b divides n . When all constants are set, the algorithm will reduce to that given in Algorithm 1. We have one hyperparameter α > 0 , which may be interpreted as a learning rate.

Using initial values w 0 = 0 and q 0 = 1 /n , initialize the tables ˆ ℓ 0 = ℓ ( w 0 ) ∈ R n , ˆ g 0 = ∇ ℓ ( w 0 ) ∈ R n × d , and ˆ q 0 = q 0 ∈ R n . In addition, partition the [ n ] sample indices into M := n/b blocks of size b , or B 1 , . . . , B M with B K := (( K -1) b +1 , . . . , Kb ) for K ∈ [ M ] . We can set the averaging sequence according to the following scheme:

The initial value a 2 = 4 α is a slight modification for theoretical convenience, and the algorithm operates exactly as in Appx. D in practice. In order to retrieve the Appx. D version, we simply replace the condition above with a t = (1 + α ) a t -1 for t ≥ 2 .

Consider iterate t ∈ { 1 , . . . , T } . We sample a random block I t uniformly from [ M ] and compute the primal update.

We see that C t -1 + c t -1 ( n/b -1) = A t -1 , so the inner objective of the update (23) is A t µ -strongly convex (as the one in (26) is A t ν -strongly concave). Note that when n/b < 1 , we simply treat the method as not including the additional regularization term c t -1 µ 2 ∑ t -2 s = t -n/b ∥ w -w s ∨ 0 ∥ 2 2 . Proceeding, we then modify the loss and gradient table. The loss update has to occur between the primal and dual updates to achieve control over the variation in the dual update (see Appx. C.2). Define K t = t mod M +1 as the (deterministic) block to be updated, and set

Define e j to be the j -th standard basis vector. Then, sample a random block J t uniformly from [ M ] compute

Notice the change in indices between (21) and (25), which accounts for the update in the loss table that occurs in between. Finally, we must update the remaining table. Set

We define the random variable H t := { ( I s , J s ) } t s =1 as the history of blocks selected at all times up to and including t , and define E t [ · ] to be the conditional expectation operator given H t -1 . In other words, E t integrates the randomness { ( I s , J s ) } T s = t . Accordingly E 1 [ · ] is the marginal expectation of the entire random process. We may now describe the optimality criterion.

Proof Outline Construct the gap function

and aim to bound E 1 [ ∑ t s =1 γ s ] by a constant. Throughout the proof, we will use a free parameter α , with update rule

Because we will search for α down to an absolute constant, we will often swap ( 1 + α 4 ) for (1 + α ) for readability. We assume the right-hand side of (28) holds in Step 1 and Step 2 . We then select α to satisfy this condition (and all others) in Step 3 below. The proof occurs in five steps total.

Lower bound the dual suboptimality a t L ( w ⋆ , q t ) .

Upper bound the primal suboptimality a t L ( w t , q ⋆ ) , and combine both to derive a bound on the gap function for t ≥ 2 .

Derive all conditions on the learning rate constant α and batch size b .

Bound γ 1 and sum γ t over t for a T -step progress bound.

Bound the remaining non-telescoping terms to complete the analysis.

We begin with a section of technical lemmas that will not only be useful in various areas of the proof but also capture the main insights that allow the method to achieve the given rate. Given these lemmas, the main proof occurs in Appx. C.3 and otherwise follows standard structure.

C.2 Technical Lemmas

This section contains a number of lemmas that describe common structures in the analysis of quantities in the primal and dual. Lem. 4 bounds cross terms that arise when there are inner products between the primal-dual iterates and their gradient estimates. Lem. 5 and Lem. 6 are respectively the primal and dual noise bounds, constructed to control the variation of the terms δ P t and δ D t appearing in (21). Finally, Lem. 10 exploits the cyclic style updates of the ˆ q t table to bound the term ∥ ˆ q t -1 -q t -2 ∥ 2 2 which is used in the primal noise bound.

Cross Term Bound Both of the estimates of the gradient of the coupled term q ⊤ ℓ ( w ) with respect to the primal and dual variables share a similar structure (see (22) and compare to (29)). They are designed to achieve a particular form of telescoping, with a remaining squared term that can be controlled by case-specific techniques. This can be observed within Lem. 4. In the sequel, we will refer to a sequence of random vectors ( u t ) t ≥ 1 as adapted to H t , where H t = { ( I s , J s ) } t s =1 is the history of random blocks. This simply means that u t , when conditioned on H t -1 , is only a function of the current random block ( I t , J t ) . Similarly, conditioned on H t -1 , we have that u t -1 is not random. In the language of probability theory, { σ ( H t ) } t ≥ 1 forms a filtration and u t is σ ( H t ) -measurable, but this terminology is not necessary for understanding the results.

Lemma4 (Cross Term Bound) . Let ( x t ) t ≥ 1 , ( y t ) t ≥ 1 , and (ˆ y t ) t ≥ 1 denote random sequences of R m -valued vectors, and let x ⋆ ∈ R m be a vector. Denote by I t be the index of a uniform random block [ M ] .

Finally, let ( x t , y t , ˆ y t ) t ≥ 1 adapted to H t (as defined above). Then, for any positive constant γ > 0 ,

Proof. By plugging in the value of v t and using x ⋆ -x t = x ⋆ -x t -1 + x t -1 -x t , we have that

where the final step follows from Young's inequality with parameter γ .

In the primal case, we have that v t = v P t , y t = ∇ ℓ ( w t ) ⊤ q t , ˆ y t = ˆ g ⊤ t ˆ q t , and x t = w t . In the dual case, we have that v t = v D t , y t = ℓ ( w t ) , ˆ y t = ˆ ℓ t +1 , and x t = q t . The next few lemmas control the third term appearing in Lem. 4 for the specific case of the primal and dual sequences.

Noise Term Bounds Next, we proceed to control the δ P t and δ D t by way of Lem. 5 and Lem. 6. As discussed in Sec. 3, a key step in the convergence proof is establishing control over these terms. Define π ( t, i ) to satisfy ˆ q t,i = q π ( t,i ) ,i and ˆ g t,i = ∇ ℓ i ( w π ( t,i ) ) , that is, the time index of the last update of table element i on or before time t . This notation is used to write the table values such as ˆ q t in terms of past values of the iterates (e.g., q t ).

Lemma 5 (Primal Noise Bound) . When t ≥ 2 , we have that

Proof. By definition, we have that

where we use that I t is drawn uniformly over n/b . Continuing again with the term above, we have

where we use that ∥∇ ℓ i ( w ⋆ ) ∥ 2 2 ≤ G 2 because every ℓ i is G -Lipschitz with respect to ∥·∥ 2 2 , and that q i ≤ q max = max q ∈Q ∥ q ∥ ∞ . This completes the proof.

The corresponding dual noise bound in Lem. 6 follows similarly, using cyclic updates in the loss table.

Lemma 6 (Dual Noise Bound) . For t ≥ 2 ,

Proof. Then, we may exploit the coordinate structure of the noise term to write

̸

̸

The red term is zero because j = k . The sum in the last line has n/b -1 terms because our order of updates forces one of the blocks of the ˆ ℓ t -1 vector to have values equal to w t -1 before defining δ D t .

Controlling the Recency of the Loss Table In this section, we bound the ∥ q t -1 -ˆ q t -2 ∥ 2 2 term appearing in the primal noise bound Lem. 5. Controlling this term is essential to achieving the correct rate, as we comment toward the end of this section.

Recall that the [ n ] indices are partitioned into blocks ( B 1 , . . . , B M ) for M := n/b , where b is assumed to divide n . For any t ≥ 1 , we first decompose

and analyze block-by-block. Our goal is to be able to count this quantity in terms of ∥ q t -q t -1 ∥ 2 2 terms. The main result is given in Lem. 10, which is built up in the following lemmas. Consider a block index K ∈ [ M ] . Define the number t K = M ⌊ ( t -K ) /M ⌋ + K when t -1 ≥ K and and t K = 0 otherwise.

Lemma 7. It holds that

Proof. We define t K to be the earliest time index τ on or before t -1 when block K of q τ was used to update ˆ q τ . When t -1 < K , then t K = 0 . When t -1 ≥ K , we can compute this number

where the last line follows by Young's inequality.

While we will not be able to cancel these terms on every iterate, we will be able to when aggregating over time and then redistributing them. Recall ( a t ) t ≥ 1 as described in Appx. C.1. Indeed, by summing across iterations, we see that

We can start by swapping the first two sums and only considering the values of t that are greater than or equal to K .

where we use in the last line that t K = 0 when t < K . We handle the terms S 0 and S 1 separately. In either case, we have to match the sums over K and over i in order to create complete vectors, as opposed to differences between coordinates. We also maintain the update rules of the sequence ( a t ) t ≥ 1 that will be used in the proof.

Lemma 8. Assume that α ≤ 1 M and a t ≤ (1 + α ) a t -1 . It holds that S 0 as defined in (31) satisfies

Proof. Write

the result as desired.

Thus, S 0 contributes about M 3 of such terms over the entire trajectory, which can be viewed as an initialization cost. Next, we control S 1 .

Lemma 9. Assume that α ≤ 1 M and a t ≤ (1 + α ) a t -1 . It holds that S 1 as defined in (31) satisfies

Proof. We can reparametrize t in terms of r = t -K , which will help in reasoning with t K . Define r K = M ⌊ r/M ⌋ + K , and let

The last line follows because the inner sum ∑ min { r + M,T } s = M ⌊ r/M ⌋ a s ∥ q s -q s -1 ∥ 2 2 can be thought of as a sliding window of size at most 2 M centered at r , which essentially repeats each element in the sum 2 M times at most.

In either case, the contribution over the entire trajectory is order M 2 = ( n/b ) 2 terms. Combining the bounds for S 0 and S 1 yields the following.

Lemma 10. Letting M = n/b be the number of blocks. Assume that α ≤ 1 M and a t ≤ (1+ α ) a t -1 . We have that

Proof. Use Lem. 7 to achieve (31) and apply Lem. 8 and Lem. 9 on each term to get

completing the proof.

We close this subsection with comments on how Lem. 10 can be used. The term E t ∥ ∥ δ D t ∥ ∥ 2 2 is multiplied in (36) by a factor na 2 t -1 G 2 µC t -1 . Thus, if we apply Lem. 10 when redistributing over time a term ∥ q t -1 -q t -2 ∥ 2 2 which will be multiplied by a factor (ignoring absolute constants) of

In order to cancel such a term, we require the use of -A t -1 ν 2 ∥ q t -q t -1 ∥ 2 2 . We can reserve half to be used up by (49), and be left with the condition

as we have applied that a t -1 ≤ 2 αC t -1 and a t -1 ≤ αA t -1 . Thus, this introduces a dependence of order n b √ nG 2 µν on α , which propagates to the learning rate α . We now proceed to the main logic of the convergence analysis.

C.3 Proof of Main Result

C.3.1 Lower Bound on Dual Objective

We first quantify the gap between L ( w ⋆ , q t ) and L ( w ⋆ , q ⋆ ) by providing a lower bound in expectation on L ( w ⋆ , q t ) , given in Lem. 11. As in Lem. 5, recall the notation π ( t, i ) to satisfy ˆ q t,i = q π ( t,i ) ,i and g t,i = ∇ ℓ i ( w π ( t,i ) ) , that is, the time index of the last update of table element i on or before time t , with π ( t, i ) = 0 for t ≤ 0 .

Lemma 11. Assume that α ≤ µ/ (24 eLκ Q ) . Then, for t ≥ 2 , we have that:

Proof. We use convexity and smoothness (1), then add and subtract (2) elements from the primal update, and finally use the definition of the proximal operator (3) with the optimality of w t for the problem that defines it.

Next, we are able to use Lem. 4 with v t = v P t , y t = ∇ ℓ ( w t ) ⊤ q t , ˆ y t = ˆ g ⊤ t ˆ q t , x t = w t , x ⋆ = w ⋆ , and γ = µC t -1 / 2 which yields that

Then, apply Lem. 5 to achieve

In the following, the blue terms indicate what changes from line to line. Combine the previous two steps to collect all terms for the lower bound. That is, apply (34) to write

then apply (35) to the blue term above to write

Finally, apply (36) to the term (37) to achieve

Next, we apply a t -1 ≤ 2 αC t -1 and α ≤ µ 24 eLκ Q to achieve:

For terms that contain π ( t -2 , i ) , we recall that π ( t -2 , i ) can be at most n/b timesteps behind t -2 , so we have that

We use here that a t ≤ (1 + α 4 ) a t -1 and impose the condition α 4 ≤ b n in the rate. Substituting this back into the lower bound achieves the desired claim.

C.3.2 Upper Bound on Gap Criterion

Next, we quantify the gap between L ( w t , q ⋆ ) and L ( w ⋆ , q ⋆ ) by upper bounding a t L ( w t , q ⋆ ) , as given in Lem. 12. When combined with the lower bound Lem. 11, we may then control the gap.

Lemma 12. For t ≥ 2 , we have that:

Proof. Weadd and subtract (1) terms in the dual update step and apply the definition of the proximal operator (2) with Bregman divergence, and the optimality of q t for the maximization problem that defines it.

We can combine the upper bound from Lem. 12 and lower bound from Lem. 11 in Lem. 13. We identify telescoping terms in blue and non-positive terms in red. The green term is canceled after aggregation across time t . This bound, like before applies for t ≥ 2 .

Lemma 13. Assume that α ≤ µ/ (24 eLκ Q ) . For t > 2 , we have that: E t [ γ t ] = E t [ a t ( L ( w t , q ⋆ ) -L ( w ⋆ , q t ) -a t µ 2 ∥ w t -w ⋆ ∥ 2 2 -a t ν 2 ∥ q t -q ⋆ ∥ 2 2 ] ≤ a t E t [ ( q ⋆ -q t ) ⊤ ( ℓ ( w t ) -ˆ ℓ t ) ] -a t -1 ( q ⋆ -q t -1 ) ⊤ ( ℓ ( w t -1 ) -ˆ ℓ t -1 ) (38) + A t -1 ν ∆ D ( q ⋆ , q t -1 ) -A t ν E t [∆ D ( q ⋆ , q t )] (39) -a t 2 L E t [ n ∑ i =1 q t,i ∥∇ ℓ i ( w t ) -∇ ℓ i ( w ⋆ ) ∥ 2 2 ] (40) + a t -1 4 L n ∑ i =1 q t -1 ,i ∥∇ ℓ i ( w t -1 ) -∇ ℓ i ( w ⋆ ) ∥ 2 2 + n ∑ i =1 a π ( t -2 ,i ) 4 L q π ( t -2 ,i ) ,i ∥ ∥ ∇ ℓ i ( w π ( t -2 ,i ) -∇ ℓ i ( w ⋆ )) ∥ ∥ 2 2 (41) -a t E t [ 〈 ∇ ℓ ( w t ) ⊤ q t -ˆ g ⊤ t -1 ˆ q t -1 , w ⋆ -w t 〉 ] + a t -1 〈 ∇ ℓ ( w t -1 ) ⊤ q t -1 -ˆ g ⊤ t -2 ˆ q t -2 , w ⋆ -w t -1 〉 (42) + 6 nG 2 αa t -1 µ ∥ q t -1 -ˆ q t -2 ∥ 2 2 (43) -µA t 2 E t ∥ w ⋆ -w t ∥ 2 2 + µC t -1 2 ∥ w ⋆ -w t -1 ∥ 2 2 + c t -1 µ 2 t -2 ∑ τ = t -n/b ∥ w ⋆ -w τ ∨ 0 ∥ 2 2 (44) + na t -1 αG 2 ν t -3 ∑ τ = t -n/b ∥ w t -1 -w τ ∨ 0 ∥ 2 2 -c t -1 µ 2 t -2 ∑ τ = t -n/b E t ∥ w t -w τ ∨ 0 ∥ 2 2 (45) + na t -1 αG 2 ν ∥ w t -1 -w t -2 ∥ 2 2 -µC t -1 4 E t ∥ w t -w t -1 ∥ 2 2 (46) -a t µ 2 E t [ ∥ w t -w ⋆ ∥ 2 2 ] -a t ν 2 E t ∥ q t -q ⋆ ∥ 2 2 -A t -1 ν 2 E t [∆ D ( q t , q t -1 )] . (47) 2

For t = 2 , the above holds with the addition of the term ν 4 ∥ q 1 -q 0 ∥ 2 2 + nG ν ∥ w 0 -w 1 ∥ 2 2 .

Proof. First, combine Lem. 11 and Lem. 12 to write:

Bound the cross term identified above. In the case that t = 2 , use Lem. 4 with v t = v D t , y t = ℓ ( w t ) , ˆ y t +1 = ˆ ℓ t , x t = q t , x ⋆ = q ⋆ , and γ = νA t -1 = ν which yields that

where the fourth term above can be bounded as

Using the definition of the update, we have that ∥ w 1 -w 0 ∥ 2 2 = (1 /µ 2 ) ∥ ∥ ∇ ℓ ( w 0 ) ⊤ q 0 ∥ ∥ 2 2 . In the case that t > 2 , use Lem. 4 as above but instead with γ = νA t -2 which yields that

We may then apply Lem. 6 and a t -1 ≤ αA t -2 to get that

We may use strong convexity to get the A t -2 ν 4 E t [ ∥ q t -q t -1 ∥ 2 2 ] term to cancel with -A t -1 ν 2 ∆ D ( q t , q t -1 ) and that A t -2 ≤ A t -1 to complete the proof.

C.3.3 Determining Constants

In this section, we provide derivations that determine the values of the constant c t that allow for cancellation of errors. We slightly adjust the notation in this subsection, in that we assume that for some η > 0

and determine η such that (28) is satisfied. We will see that η is simply a constant factor away from α , so the resulting condition we actually be on α . The latter is given formally in Lem. 14. We assume here that n/b ≥ 2 , which is taken as an assumption of Thm. 2.

In the statement of Lem. 13, the lines above (43) will telescope without additional conditions. For (44), we set c t = a t /m for some parameter m . Note that this condition does not need to be checked when n/b < 1 , as the additional sum term over τ will not be included in the update. Counting all the terms that will appear when matched on the index t -1 , we have the condition that

The first term result from the 'good term' -a t -1 µ 4 ∥ w t -1 -w ⋆ ∥ 2 2 from the bottom. The rightmost term above results because a s ∥ w ⋆ -w τ ∨ 0 ∥ 2 2 will have τ = t -1 when s ∈ { t +1 , . . . , t + n/b -1 } . We will begin by requiring that require that a t ≤ (1 + β ) a t -1 for all t and some β > 0 , and then determine β below. The condition reads as

which can be summed and represented as

Rearranging and taking a logarithm on both sides, this is the same as

Next, using the inequality 2 x 2+ x ≤ ln(1 + x ) with x = β (4 n/b -4+ m ) 4(1+ β ) 2 which holds for all x ≥ 0 , we have

We can also apply the upper bound ln( x +1) ≤ x with x = β (which also holds for any x ≥ 0 ) to write

which means that (51) will be satisfied (using (52)) if

In order to satisfy the inequality, substitute m = 4 cβ ( n/b -1) 2 for some c > 0 to be determined and assume that β ≤ 1 n/b -1 , so that β ( n/b -1) ≤ 1 . The LHS reads as

The RHS can be upper-bounded as

which makes the inequality satisfied when

so we can set c = 16 . We now have the flexibility to control β , and the telescoping of (44) is achieved. For (45), set β = α 4 and pass the condition of β ≤ 1 / ( n/b -1) onto α , which maintains the rate (and is already satisfied when α ≤ b n and n ≥ 2 ). Then, we can achieve

by requiring that α ≤ b √ µν 4 n 3 / 2 G , which achieves the telescoping of (45). Finally, to address (46), we may satisfy it if

which we can achieve by incorporating the condition a t ≤ 4 α ( n/b -1) 2 C t -1 into the update of a t , because na t αG 2 ν ≤ µa t m by the previous condition on α . Having chosen c t , we are prepared to produce a learning rate parameter η to capture all conditions on α in one, as given in Lem. 14.

Lemma 14. For all t ≥ 1 , we have the following.

Setting c t = a t / [16 α ( n/b -1) 2 ] implies that a t ≤ 2 αC t .

Using the update scheme

when

we have that (28) holds.

Proof. Noting that m = 16 α ( n/b -1) 2 we confirm that

This condition is satisfied when α ≤ 1 32( n/b -1) , so we incorporate α ≤ b 32 n into the rate, implying that a t ≤ 2 αC t . This concludes the proof of the first bullet point.

Next, we show that a t ≤ (1+ α/ 4) a t -1 will imply that a t ≤ αA t -1 , which is the second part of (28). First, we define a 2 = αa 1 as an initialization (which also satisfies a 2 ≤ (1+ α/ 4) a 1 when α ≤ 4 / 3 ), and show inductively that if a t -1 ≤ αA t -2 , then a t ≤ (1 + α/ 4) a t -1 = ⇒ a t ≤ αA t -1 . In the base case, A 1 = a 1 , so the condition a 2 = αa 1 satisfies a 2 ≤ αA 1 . Next, fixing t and assuming that 1) a t -1 ≤ αA t -2 and 2) that a t ≤ (1 + α/ 4) a t -1 , we have that

the desired result. Finally, we consider the condition a t ≤ 4 α ( n/b -1) 2 C t -1 , the third part of (28). We wish to show that the following inequality holds:

which is implied by the inequality

When n/b ≥ 2 , we have that (1 + α/ 4) ≤ 4( n/b -1) 2 ( α +1 -1 4( n/b -1) ) , so we require that b ≤ n/ 2 . Thus, our final updates are given by

Because each condition was satisfied when using a t = (1+ α/ 4) a t -1 , we define η = α/ 4 to achieve the claimed result.

C.3.4 Bound on Sum of Successive Gaps

Lem. 15 is an upper estimate for the expected sum of the gap function over T iterates. Recall that E 1 [ · ] the full expectation over { ( I t , J t ) } T t =1 . The green term is a quantity that remain as an initialization term, whereas the blue terms have to be bounded from above. The terms directly below the blue terms account for all of the 'negative π ( t -1 , i ) ' terms are not yet used up by the telescoping in lines (39), (40), and (41), and there are in fact between 1 and n copies of those terms in each iteration, even though we will use only 1 .

Lemma 15 (Progress Bound) . Assume that

For any T ≥ 1 , we have that

Proof. We proceed by first deriving an upper bound on γ 1 , the gap function for t = 1 . Note that w 1 is non-random, as a 0 = 0 implies that v 0 = ∇ ℓ ( w 0 ) . Using that w 1 is the optimum for the proximal operator that defines it, the upper bound can be written as

where we use that ˜ ℓ 1 = ˆ ℓ 1 . For the lower bound, use a similar argument to Lem. 11 to achieve

where we use that v P 0 = ∇ ℓ ( w 0 ) ⊤ q 0 . We combine them to get

where the last two terms are the result of the additional quadratic slack terms in γ 1 . All of the terms from the display above will be telescoped. Thus, we apply Lem. 13 and collect the unmatched terms from the t ≥ 2 one-step bound (using that A 0 = 0 ). The term (57) can be viewed as counting the remainder of (40) after it has telescoped some but not all terms

C.3.5 Completing the Proof

We use similar techniques as before to bound the remaining terms from the T -step progress bound given in Lem. 15. We may now prove the main result.

Theorem 2. For a constant α > 0 , define the sequence

along with its partial sum A t = ∑ t τ =1 a τ . Under Asm. 1, there is an absolute constant C such that using the parameter

the iterates of Algorithm 1 satisfy:

We can compute a point ( w T , q T ) achieving an expected gap no more than ε with bigO complexity

Proof. We first apply Lem. 15, and proceed to bound the inner product terms (55) and (56). Apply Young's inequality with parameter νA T -1 / 2 to get

The left-hand term will be canceled by (61) by applying strong concavity (leaving behind -νA T -1 4 E 0 ∥ q ⋆ -q T ∥ 2 2 ) and the right-hand term (because of the condition α ≤ √ µν 4 n 3 / 2 G ) will be canceled by (60). Next, consider (56). By Young's inequality with parameter µA T -1 / 2 , we have

where the second term will be canceled by (59). For the remaining term,

We may add the middle term above to (58), so that the remaining term to bound is

To show that this quantity is non-negative, we use that a 0 = 0 and Lem. 10 (recalling that M = n/b to see that

which will cancel with the rightmost term in (58) provided that α ≤ b 36 e 2 n √ µν nG 2 . Thus, plugging the previous displays into Lem. 15, we have that

The black lines will cancel given our conditions on α . Substituting the definition of γ t and moving the final non-positive terms on the last line, that is, ( A T -1 + a T ) µ 4 E 0 ∥ w T -w ⋆ ∥ 2 2 and ( A T -1 + a T ) ν 4 E 0 ∥ q T -q ⋆ ∥ 2 2 to the left-hand side achieves the claim.

C.4 Modification for Unregularized Objectives

For completeness, we describe a modification of DRAGO for unregularized objectives, or (2) when µ ≥ 0 and ν ≥ 0 . The analysis follows similarly to the previous subsections (regarding the µ, ν > 0 case), and we highlight the steps that differ in this subsection by presenting a slightly different upper bound on the gap criterion based on the modified primal and dual updates. This will result a different update for the sequence ( a t ) t ≥ 1 , subsequently affecting the rate.

C.4.1 Overview

The modified algorithm is nearly identical to Algorithm 1, except that the dual and primal updates can be written as

and

respectively, and µ 1 , µ 2 , ν 1 ≥ 0 are to-be-set hyperparameters. When ν > 0 , we may set ν 1 = 0 , and when µ > 0 , we may set µ 1 = µ 2 = 0 , which recover the Algorithm 1 updates exactly. While

we may set ν 1 = 1 when it is positive (and similarly for µ 1 and µ 2 ), they may be set to different values in order to balance the terms appearing in the rate below. As in Appx. C.1, we wish to upper bound the expectation of the quantity

which is still non-negative in the case of µ = 0 or ν = 0 . By using an appropriate averaging sequence ( a t ) t ≥ 1 and defining A T = ∑ T t =1 a t , we upper bound ∑ T t =1 a t E 0 [ γ t ] (see Thm. 2) by a constant value independent of T . Recall that the batch size is denoted by b . As we derive in Appx. C.4.3, our final update on the ( a t ) sequence is

Observe that when µ = 0 , we set a t = µ 1 12 enq max L to achieve a O (1 /t ) rate. We omit proofs in this subsection as they follow with the exact same steps as the corresponding lemmas in the strongly convex-strongly concave setting (which we point to for each result).

C.4.2 Upper Bound on Gap Criterion

Following the steps of Appx. C.3.1 and Appx. C.3.2, we will first derive lower and upper bounds on E t [ a t L ( w ⋆ , q t )] and a t L ( w t , q ⋆ ) and combine them to upper bound a t E t [ γ t ] . Recalling that E t [ · ] denotes the condition expectation given ( w t -1 , q t -1 ) , and we can then take the marginal expectation to upper bound a t E 0 [ γ t ] . The following lower bound is analogous to Lem. 11 and follows the exact same proof technique.

Lemma 16. For t ≥ 2 , assuming that a t ≤ C t -1 µ + µ 1 12 enq max L and a t ≤ (1 + b/n ) a t -1 , we have that:

Similarly, following the same steps as Lem. 12, one can derive the following upper bound.

Lemma 17. For t ≥ 2 , we have that:

By combining Lem. 17 and Lem. 16, we can upper bound the quantity E t [ a t ( L ( w t , q ⋆ ) -L ( w ⋆ , q t )] . Consequently, the following result follows the same steps as Lem. 13. As before, we identify telescoping terms in blue, non-positive terms in red, where green term is bounded after aggregation across time t .

Lemma 18. Assume that a t -1 ≤ C t -1 µ + µ 1 12 eLnq max and a t ≤ (1 + b/n ) a t -1 . For t > 2 , we have that:

-

+(

A

t

1

ν

+

ν

1

)∆

D

(

q

⋆

, q

t

-

1

)

(

A

t

ν

+

ν

1

)

E

t

[∆

D

(

q

⋆

, q

t

)]

(67)

-

(69)

For t = 2 , the above holds with the addition of the term nG 2 ν + ν 1 ∥ w 1 -w 0 ∥ 2 2 and without any term including A t -2 in the denominator.

We then select constants to achieve the desired telescoping in each line of Lem. 18.

C.4.3 Determining Constants

We now select the constants ( µ 1 , µ 2 , ν 1 ) , and the sequences ( a t ) and ( c t ) to complete the main part of the analysis.

In the statement of Lem. 18, the lines above (71) will telescope without additional conditions. For (72), we set c t = a t /m for some parameter m (just as in Appx. C.3.3). Note that this condition does not need to be checked when n/b < 1 , as the additional sum term over τ will not be included in the update. Counting all the terms that will appear when matched on the index t -1 , we have the condition that

where the first underbrace is non-positive based on the choice of m selected in Appx. C, and the equality is satisfied for all values of µ 1 and µ 2 . Lines (73) and (74) yield the conditions

for the telescoping to be achieved, which can equivalently be rewritten as

These can be accomplished by setting

We must also handle the term 3 nG 2 a 2 t -1 C t -1 µ + µ 1 ∥ q t -1 -ˆ q t -2 ∥ 2 2 . By the argument of Lem. 10 using a 2 t instead of a t , we have that when summing over t , we have that

where M = n/b is the number of blocks, and the e 4 term appears from the squared constants (as opposed to e 2 ). Using that we will sum the non-positive terms ∑ T t =1 A t -1 ν E t [∆ D ( q t , q t -1 , ]) , so that the final condition needed to cancel this term is

which can also be rewritten as

Thus, combining the previous conditions, our final update on the ( a t ) sequence is

as alluded to in Appx. C.4.1.

D Implementation Details

In this section, we provide additional background on implementing DRAGO in practice. This involves a description of the algorithm amenable for direct translation into code and procedures for computing the dual proximal mapping for common uncertainty sets and penalties. We assume in this section that W = R d and provide multiple options for the uncertainty set Q .

D.1 Algorithm Description

The full algorithm is given in Algorithm 2. We first describe the notation. Recall that M = n/b , or the number of blocks. We partition [ n ] into ( B 1 , . . . , B M ) , where each B K denotes a b -length list of contiguous indices. For any matrix u ∈ R n × m (including m = 1 ), we denote by u [ B K ] ∈ R b × m the rows of u corresponding to the indices in B K . Finally, for a vector s ∈ R b , we denote by se B K the vector that contains s k in indices k ∈ B K , and has zeros elsewhere. Next, we comment on particular aspects regarding the implementation version as compared to Algorithm 1 (Sec. 2).

We store two versions of each table, specifically ˆ ℓ, ˆ ℓ 1 ∈ R n , ˆ g 1 , ˆ g 2 ∈ R n × d , and ˆ q 1 , ˆ q 2 ∈ R n . For any iterate t , these variables are meant to store ˆ ℓ t , ˆ ℓ t -1 ∈ R n , ˆ g t -1 , ˆ g t -2 ∈ R n × d , and ˆ q t -1 , ˆ q t -2 ∈ R n .

The quantities ˆ g agg ∈ R d and ˆ w agg ∈ R d are introduced as to not recompute the sums in the primal update on each iteration (which would cost O ( nd ) operations). Instead, these aggregates are updated using O ( bd ) operations.

The loss and gradient tables are not updated immediately after the primal update. However, the values that fill the tables are computed, and the update occurs at the end of the loop. This is because ˆ ℓ [ B K t ] is used to fill ˆ ℓ 1 [ B K t ] at the end of the loop, we we must maintain knowledge of ˆ ℓ [ B K t ] temporarily.

While the proximal operator is specified for the primal in the case of W = R d , the proximal operator for the dual os computed by a subroutine DualProx , which we describe in the next subsection.

D.2 Solving the Maximization Problem

As discussed in Appx. B, the primary examples of DRO uncertainty sets Q are balls in f -divergence (specifically, KL and χ 2 ) and spectral risk measure sets. For the penalty D , it is also common to use f -divergences. We review these concepts in this section and provide recipes for computing the maximization problem.

f -Divergences We first recall the definition of f -divergences used throughout this section.

Definition 19. Let f : [0 , ∞ ) ↦→ R ∪{ + ∞} be a convex function such that f (1) = 0 , f ( x ) is finite for x > 0 , and lim x → 0 + f ( x ) = 0 . Let q and ¯ q be two probability mass functions defined on n atoms. The f -divergence from q to ¯ q generated by this function f is given by

where we define 0 f (0 / 0) := 0 . For any i such that ¯ q i = 0 but q i > 0 , we define D f ( q ∥ ¯ q ) =: + ∞ .

The two running examples we use are the χ 2 -divergence generated by f χ 2 ( x ) = x 2 -1 and the KL divergence generated by f KL ( x ) = x ln x on (0 , ∞ ) and define 0 ln 0 = 0 . For any convex set X ⊆ R k , we also introduce the convex indicator function

In either of the two cases below, we select the penalty D ( q ∥ 1 /n ) = D f ( q ∥ 1 /n ) to be an f -divergence. Denote in addition f ∗ as the Fenchel conjugate of f .

Algorithm 2 DRAGO: Implementation Version

Input: Learning rate parameter α > 0 , batch size b ∈ { 1 , . . . , n } , number of iterations T . Initialization:

for t = 1 to T do

Sample blocks I t and J t uniformly on [ n/b ] and compute K t = t mod ( n/b ) + 1

Primal Update:

Compute Loss and Gradient Table Updates:

Dual Update:

Update All Tables:

end for

return ( w,q ) .

D.2.1 Spectral Risk Measure Uncertainty Sets

As in Appx. B, the spectral risk measure uncertainty set is defined by a set of non-decreasing, nonnegative weights σ = ( σ 1 , . . . , σ n ) that sum to one. Our uncertainty set is given by

and we use D f has the penalty for either f χ 2 or f KL. The set Q ( σ ) is referred to the permutahedron on σ . In this case, the maximization problem can be dualized and solved via the following result.

Proposition 20. [Mehta et al., 2024, Proposition 3] Let l ∈ R n be a vector and π be a permutation that sorts its entries in non-decreasing order, i.e., l π (1) ≤ . . . ≤ l π ( n ) . Consider a function f strictly convex with strictly convex conjugate defining a divergence D f . Then, the maximization over the permutahedron subject to the shift penalty can be expressed as

where we define g i ( z ; l ) := σ i z + ν n f ∗ ( ( l π ( i ) -z ) /ν ) . The optima of both problems, denoted

are related as q opt ( l ) = ∇ ( νD f ( ·∥ 1 n /n )) ∗ ( l -z opt π -1 ( l )) , that is,

As described in Mehta et al. [2024, Appendix C], the minimization problem (78) is an exact instance of isotonic regression and can be solved efficiently with the pool adjacent violators (PA V) algorithm.

D.2.2 Divergence-Ball Uncertainty Sets

Another common uncertainty set format is a ball in f -divergence, or

We describe the case of the rescaled χ 2 -divergence in particular, in which the feasible set is an ℓ 2 -ball intersected with the probability simplex. Given a vector l ∈ R n , we aim to compute the mapping

where P n := { q ∈ R n : q ≥ 0 , 1 ⊤ q = 1 } denotes the n -dimensional probability simplex. We apply a similar approach to Namkoong and Duchi [2017], in which we take a partial dual of the problem above. Indeed, note first that for any q ∈ P n , we have that 1 2 ∥ q -1 n /n ∥ 2 2 = 1 2 ∥ q ∥ 2 2 -1 2 n . Thus, the optimal solution to (80) can be computed by solving

or equivalently, by strong duality via Slater's condition, solving

Notice that evaluation of the outer objective itself requires Euclidean projection onto the probability simplex as a subroutine, after which the maximization problem can be computed via the bisection method, as it is a univariate concave maximization problem over a convex set. In order to determine which half to remove in the bisection search, we also compute the derivative of λ ↦→ f ( λ ) , which is given by

where q opt ( λ ) achieves the minimum in min q ∈P n 1 2 ∥ q -l/ ( ν + λ ) ∥ 2 2 for a fixed λ ≥ 0 . For projection onto the probability simplex, we apply Algorithm 1 from Condat [2016], which is a solution relying on sorting the projected vector. The overall method consists of three steps.

Sorting: Projection onto the simplex relies on sorting the vector l/ ( ν + λ ) on each evaluation. However, because l/ ( ν + λ ) varies from evaluation to evaluation simply by multiplying by a positive scalar, we may pre-sort l and use the same sorted indices on each evaluation of ( f ( λ ) , f ′ ( λ )) listed below.

Two-Pointer Search: We find the upper and lower limits for λ by initializing λ min = 0 and λ max = 1 , and repeatedly making the replacement ( λ min , λ max ) ← ( λ max , 2 λ max ) until f ′ ( λ max ) < -ε for some tolerance ε > 0 . This, along with f ′ ( λ min ) > ε indicates that the optimal value of λ lies within ( λ min , λ max ) . For any λ with | f ′ ( λ ) | < ε , we return the associated q opt ( λ ) as the solution.

Binary Search: Finally, we repeatedly evaluate f ′ ( λ ) for λ = ( λ min + λ max ) / 2 . If f ′ ( λ ) > ε , we set λ min ← λ , whereas if f ′ ( λ ) < -ε , then we set λ max ← λ . We terminate when λ max -λ min < ε or | f ′ ( λ ) | < ε .

The parameter ε is set to 10 -10 in our experiments. Note that the same procedure can be used to compute the dual proximal operator in Algorithm 1. In particular, when ∆ D (( , q ) , q t -1 ) = 1 2 ∥ q -q t ∥ 2 2 , which is true when D is the χ 2 -divergence, then

which is a particular case of (80), and hence can be solved using the exact same procedure. The runtime of this subroutine is O ( n log n + n log(1 /ε )) , accounting for both the initial sorting at O ( n log n ) cost, and the O (log(1 /ε )) iterations of the exponential and binary searches. Each iteration requires a linear scan of n elements at cost O ( n ) .

Hardware Acceleration Finally, note that the computations in Appx. D.2 and Appx. D.2.2 involve primitives such as sorting, linear scanning through vectors, and binary search. Due to their serial nature (as opposed to algorithms that rely on highly parallelizable operations such as matrix multiplication), we also utilize just-in-time compilation on the CPU via the Numba package for increased efficiency.

E Experimental Details

We describe details of the experimental setup, including datasets, compute environment, and hyperparamater tuning. We largely maintain the benchmarks of Mehta et al. [2023].

E.1 Datasets

The sample sizes, dimensions, and source of the datasets are summarized in Tab. 5. The tasks associated with each dataset are listed below.

(a) yacht : predicting the residuary resistance of a sailing yacht based on its physical attributes Tsanas and Xifara [2012].

(b) energy : predicting the cooling load of a building based on its physical attributes Baressi Segota et al. [2020].

(c) concrete : predicting the compressive strength of a concrete type based on its physical and chemical attributes Yeh [2006].

(d) kin8nm : predicting the distance of an 8 link all-revolute robot arm to a spatial endpoint [Akujuobi and Zhang, 2017].

(e) power : predicting net hourly electrical energy output of a power plant given environmental factors [T¨ ufekci, 2014].

(f) acsincome : predicting income of US adults given features compiled from the American Community Survey (ACS) Public Use Microdata Sample (PUMS) [Ding et al., 2021].

(g) emotion : predicting the sentiment of sentence in the form of six emotions. Each input is a segment of text and we use a BERT neural network Devlin et al. [2019] as an initial feature map. This representation is fine-tuned using 2 epochs on a random half (8,000 examples) of the original emotion dataset, and then applied to the remaining half. We then apply principle components analysis (PCA) to reduce the dimension of each vector to 45 .

E.2 Hyperparameter Selection

We fix a minibatch size of 64 SGD and an epoch length of N = n for LSVRG. In practice, the regularization parameter µ and shift cost ν are tuned by a statistical metric, i.e. generalization error as measured on a validation set.

For the tuned hyperparameters, we use the following method. Let k ∈ { 1 , . . . , K } be a seed that determines algorithmic randomness. This corresponds to sampling a minibatch without replacement for SGD and SRDA and a single sampled index for LSVRG. Letting L k ( η ) denote the average value of the training loss of the last ten passes using learning rate η and seed k , the quantity L ( η ) = 1 K ∑ K k =1 L k ( η ) was minimized to select η . The learning rate η is chosen in the set { 1 × 10 -4 , 3 × 10 -4 , 1 × 10 -3 , 3 × 10 -3 , 1 × 10 -2 , 3 × 10 -2 , 1 × 10 -1 , 3 × 10 -1 , 1 × 10 0 , 3 × 10 0 } , with two orders of magnitude lower numbers used in acsincome due to its sparsity. We discard any learning rates that cause the optimizer to diverge for any seed.

0

50000

100000

10

-

8

10

-

5

10

-

2

Suboptimality

yacht

0

50000

100000

10

-

6

10

-

3

10

0

energy

0

50000

100000

10

-

5

10

-

2

concrete

0

1

2

10

-

6

10

-

3

10

0

yacht

0

1

2

10

-

8

10

-

5

10

-

2

energy

0

1

10

-

5

10

-

2

concrete

0

50000

100000

First-Order Oracle Evaluations

10

-

2

10

-

1

10

0

Suboptimality

kin8nm

0

50000

100000

First-Order Oracle Evaluations

10

-

2

10

-

1

10

0

power

0

200000

400000

First-Order Oracle Evaluations

10

-

1

10

0

emotion

0

5

Wall-Clock Time (Seconds)

10

-

5

10

-

2

kin8nm

0

5

Wall-Clock Time (Seconds)

10

-

6

10

-

3

10

0

power

0

20

Wall-Clock Time (Seconds)

10

-

5

10

-

2

emotion

SGD

Drago (

b

= 16

)

Drago (

b

=

n/d

)

Drago (

b

= 1

)

E.3 Compute Environment

Experiments were run on a CPU workstation with an Intel i9 processor, a clock speed of 2.80GHz, 32 virtual cores, and 126G of memory. The code used in this project was written in Python 3 using the Numba packages for just-in-time compilation. Run-time experiments were conducted without CPU parallelism. The algorithms are primarily written in PyTorch and support automatic differentiation.

E.4 Additional Experiments

We explore the sensitivity of the results to alterations of the objective and algorithm hyperparameters.

Sensitivity to Uncertainty Set Choice In Sec. 4, we mainly show performance on spectral riskbased uncertainty sets, in particular the conditional value-at-risk (CVaR). In this section, we also consider f -divergence ball-based uncertainty sets, with the procedure described in Appx. D.2.2. As in Namkoong and Duchi [2017], we use a radius that is inversely proportional to the sample size, namely ρ = 1 n , and the strong convexity-strong concavity parameter µ = ν = 1 . In Fig. 4, we demonstrate the performance of DRAGO with b = 1 , b = 16 (as chosen heuristically), and b = n/d . We compare against the biased stochastic gradient descent, which can be defined using oracle to compute the optimal dual variables given a vector of losses; however, note that LSVRG is designed only for spectral risk measures, so the method does not apply in the divergence ball setting. We observe that the optimization performance across both regression and multi-class classification tasks are qualitatively similar to that seen in Fig. 2 nad Fig. 3. The b = 1 variant performs well on smaller datasets ( n ≤ 1 , 000 ), whereas the b = 16 heuristic generally does not dominate in terms of gradient evaluations or wall time. While the number of gradient evaluations is significantly larger for the b = n/d variant, implementation techniques such as just-in-time complication (see Appx. D) allow for efficient computation, resulting in better overall optimization performance as a function of wall time.

Sensitivity to Batch Size In Fig. 5, we consider the datasets with the largest ratio of n to d (hence the largest theoretically prescribed batch size) and assess the performance of DRAGO with smaller batch sizes. For both datasets, we have that in magnitude, n/d ≈ 1000 . Intuitively, the smaller batch size methods would perform better in terms of oracle queries but the large batch methods would be more performant in terms of wall time. With only a batch size of b = 64 , this variant of DRAGO generally matches the best-performing setting when viewed from either oracle calls or direct wall time. This is approximately 16 × smaller than the n/d benchmark, indicating that tuning the batch size can significantly reduce the memory overhead of the algorithm while increasing speed.

0

50

100

150

200

10

-

8

10

-

6

10

-

4

10

-

2

10

0

kin8nm

µ

=

1.0,

ν

=

1.0

0

50

100

150

200

µ

=

10.0,

ν

=

1.0

0

50

100

150

200

µ

=

1.0,

ν

=

10.0

0

50

100

150

200

µ

=

10.0,

ν

=

10.0

0

50

100

150

200

First-Order Oracle Evaluations (K)

10

-

8

10

-

6

10

-

4

10

-

2

10

0

power

0

50

100

150

200

First-Order Oracle Evaluations (K)

0

50

100

150

200

First-Order Oracle Evaluations (K)

0

50

100

150

200

First-Order Oracle Evaluations (K)

0

20

40

60

10

-

8

10

-

6

10

-

4

10

-

2

10

0

kin8nm

µ

=

1.0,

ν

=

1.0

0

20

40

60

µ

=

10.0,

ν

=

1.0

0

20

40

60

µ

=

1.0,

ν

=

10.0

0

20

40

60

µ

=

10.0,

ν

=

10.0

0

20

40

60

Wall-Clock Time (Seconds)

10

-

8

10

-

6

10

-

4

10

-

2

10

0

power

0

20

40

60

Wall-Clock Time (Seconds)

0

20

40

60

Wall-Clock Time (Seconds)

0

20

40

60

Wall-Clock Time (Seconds)

Drago (

b

= 1024

)

Drago (

b

= 256

)

Drago (

b

= 64

)

Drago (

b

= 16

)

Drago (

b

= 4

)

Drago (

b

= 1

)

F NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: Sec. 3 for direct theoretical claims and Appx. B for detailed comparisons to other work.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: See Sec. 1 and Sec. 5 for general points, whereas Sec. 4 mentions specific cases of algorithm performance. The assumptions made are otherwise standard.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate 'Limitations' section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: This is done for every theoretical statement.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Appx. D contains detailed descriptions of experimental details and code with an associated environment and quickstart guide is provided.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The code is made public at https://github.com/ronakdm/drago.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: See Appx. D and Appx. E for implementation and experiment details.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: Training curves and hyperparameter selection experiments are averaged over multiple seeds, but error bars are not shown to make the plots visible.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer 'Yes' if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Yes, as written in Appx. E.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: NA

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: It is written after the main text of the paper.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [No]

Justification: Generative and related large-scale models are not used in this work.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: NA

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: Documented code is provided.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: Human participants are not used.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: No IRM approvals were needed for this work.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: SUMMARY:
To tackle the optimization problem where the objective is a weighted sum of losses from multiple domains, an improved optimization algorithm is designed by incorporating past history and appropriate regularization. Merits of the proposed algorithm are demonstrated by a rigorous theoretical analysis of convergence rate  and numerical analysis by synthetic and text classification data.

STRENGTHS:
Theoretical analysis seems solid.

WEAKNESSES:
The proposed method seems to be only relevant for problems with a low dimensionality $d$. Otherwise, the fact that the algorithm requires to keep track of historical weights would cause a high space complexity.

QUESTIONS:
1. Regarding the remark on global complexity on page 7, I do not follow where the $bd$ term comes from in the per iteration complexity. Is it possible to provide some details? My understanding may be inaccurate, but if each elementwise update of $v_t^D$ is of low computational cost, and we update all elements in a distributed fashion, can we improve the per iteration cost to be $O(n)$? 
2. The dimensionality considered in the empirical analysis is pretty low. Can additional empirical analysis be shown where $d$ is, e.g. of a similar magnitude as $n$, if not greater than $n$? I understand that there may be a challenge in updating $q$.

LIMITATIONS:
No direct societal impact as the paper is focused on theoretical aspects of the out of distribution inference problem.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-14

comments: SUMMARY:
The authors provide a state-of-the-art optimization algorithm with convergence guarantees for f-divergence DRO. The numerical study is also extensive

STRENGTHS:
1. Techincal results are strong and to the best of my knowledge, I haven't find mistakes
2. Numerical study is extensive

WEAKNESSES:
As the author has commented, they only solve for convex optimization formulation. The extension to non-convex setup also seems interesting and important.

QUESTIONS:
N/A

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-15

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes DRAGO, a variance-reduced optimization method for\ndistributionally robust optimization (DRO). In particular, the authors consider\nthe penalized DRO framework and view DRO as min-max optimization with a simple\nfinite-sum structure. Their algorithm maintains explicit histories for the\nindividual loss functions $l_i$, loss function gradients $\\nabla l_i$, and dual\nparameters $q_i$ (i.e. like SAG/SAGA); these estimates are used to control the\nvariance in both the primal and dual updates, which the authors prove leads to\na fast linear rate of convergence.  Unlike existing methods, DRAGO supports any\npenalty parameter $\\nu \\geq 0$ and its convergence rate on depends explicitly\non the different problem constants.\n\nSTRENGTHS:\n- DRAGO obtains a fast linear rate and the analysis shows its dependence on\n    batch size, dataset size, and the other problem constants.\n\n- DRAGO performs well in practice, particularly when the batch-size is tuned. \n\n- The analysis permits any non-negative penalization strength $\\nu$ for the\n    dual (i.e. uncertainty) parameters.\n\nWEAKNESSES:\n- The paper organization makes it quite difficult to understand the different\n    algorithmic components used in DRAGO.\n    \n- The assumptions necessary for convergence are somewhat unrealistic, although\n    these appear to be standard in the literature. \n\n- The experimental evaluation is small and limited to relatively low-dimensional\n    problems.\n\n### Detailed Comments\n\nThis is an interesting paper which tries to speed up DRO by combining variance reduction with proximal gradient updates for a structured min-max problem. While I'm not an expert in DRO, the paper seems to make a solid contribution, particularly by relaxing previous requirements on the dual regularization strength $\\nu$. I also appreciate that the convergence rate shows the role of each problem constant and tuning parameter, including batch-size. Please note that I did not check the appendix for correctness. \n\n**Paper Organization**:\nI found this paper quite confusing to read. Instead of introducing DRAGO and\nthen explaining it's various algorithmic components, the authors first try to\nexplain the intuition behind the updates. This is a good idea in theory, but in\npractice there are too many moving parts to DRAGO and I quickly got lost in the\n(mostly unimportant) details. The actually key parts of DRAGO --- proximal-gradient\nupdates with SAG/SAGA style variance reduction --- are barely mentioned and this\nmakes the algorithm much harder to understand.\n\nMy suggestion is to introduce Algorithm 1 first and then explain why \nadditional proximal terms are needed in the proximal-gradient update in order\nto control the variance. At least this way the reader has a reference point\nfor understanding the discussion.\n\n**Assumptions**:\nI think that assuming each $l_i$ is both Lipschitz continuous and Lipschitz\nsmooth is quite unrealistic. Typically only one or the other of these\nconditions holds. For instance, Lipschitz continuity implies that $l_i$ is\nbounded above by a linear function, while Lipschitz smoothness implies that\n$l_i$ is bounded above by a quadratic. While both can be true, the quadratic\nbound will generally be vacuous. However, I understand from Table 1 that these\nassumptions are typical for the literature, so I don't think this is too much\nof an issue.\n\n**Experimental Evaluation**: \nI would have liked to see experiments on a higher dimensional dataset.  Since\nthe entire Jacobian $\\nabla l(w)$ must be maintained in the history, SAG-type\nmethods usually have high memory requirements as $d$ increases. The authors\npartially handle this by using a large batch-size, but this is only feasible if\nthe mini-batches can still be fit on the GPU. When $d$ is sufficiently large,\nmini-batch gradients will be inefficient to compute and I am curious if the\nrelative performance of DRAGO and LSVRG ill switch. This is particularly \nworth considering because it is only in the large-batch setting that DRAGO\nseems to outperform LSVRG\n\n### Minor Comments\n\n- Line 52: This is a super-linear complexity, but a sub-linear convergence rate.\n    Contrast with Line 62, which shows a linear rate, but a logarithmic complexity.\n\n- Line 65: Should this read \"but do not\"?\n\n- Line 24/71: It seems like you are only considering $\\nu$ in the probability simplex.\n    It would be nice if you stated this explicitly somewhere.\n\n- Line 80: It's generally unrealistic for both $l_i$ and $\\nabla l_i$ to be \n    Lipschitz continuous over an unbounded set.\n\n- Line 86: You should mention that strong convexity comes only from the regularizer,\n    as otherwise it seems to conflict with Lipschitzness of $l_i$. \n\n- Line 102: Does $\\nabla l(w_{t-1})$ denote the Jacobian of the vector-valued map\n    $l$ such that $[l(w)]_i = l_i(w)$? It would be nice if you clearly defined\n    this notation.\n\n- Table 1: The font-size is much too small to read when printed. \n\n- Algorithm Description: The purposes of $\\bar{\\beta}$ and $\\beta_t$ and their\n    settings are not well described.\n\n- Display after Line 164: You seem to be using $D(q \\| 1/n)$ is\n    1-strongly-convex for this result, but this assumption is only introduced\n    in Assumption 1. Maybe move Assumption 1 to the start of the discussion\n    of convergence to fix this?\n\n- Figure 3: The font-size is much too small to be readable without zooming.\n\nQUESTIONS:\n- Line 491: Isn't $n q_{\\text{max}} \\leq n$ trivially satisfied since $q$\n    is constrained to the probability simplex?\n\n- Figure 2: Why is $b = 1$ only show for three datasets and only in \n    terms of oracle complexity?\n\n- Display after Line 124: It looks you are using the equality,\n\n$$\\\\mathbb{E}[(l_{j_t}(w_{t+1}) - \\hat l_{t, j_t}) (l_{j_t}(w_t) - \\hat l_{t-1, j_t})] = \\\\mathbb{E}[(l(w_{t+1}) - \\hat l_{t}) (l_{j_t}(w_t) - \\hat l_{t-1, j_t})]$$\n\nbut I don't think this is true because the terms in the product are correlated. \nCan you please explain this?\n\nLIMITATIONS:\nLimitations appear to be appropriate addressed in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-16"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper considers the DRO problem with a dual penalization term and primal regularization term. Their core contribution is a new primal-dual algorithm for this problem:\n- It achieves a new best first-order oracle complexity in certain parameter regimes. In particular, it beats previous algorithms when the dual penalization parameter $\\nu$ is small.\n- It achieves fine-grained dependence on the primal and dual condition numbers (i.e., not having a max over min term).\n- The uncertainty set $\\mathcal{Q}$ can be any closed and convex subset of the simplex. This is in contrast to previous works which specialize to more structured uncertainty sets.\n- The algorithm adapts ideas in a novel way - in particular, they show that coordinate-style updates can be applied even though the dual feasible set is non-separable. This might lead to applications in other minimax problems. It also only has a single hyperparameter $\\eta$ to tune, and comes with a batch-size parameter $b$ which can be used to trade-off between iteration count and per-iteration complexity.\n- Strong experimental results are shown (albeit, I should say that while they look very encouraging, I don't have the empirical background necessary to evaluate the methods used/experimental setup carefully).\n\nSTRENGTHS:\nI think this is a strong contribution to the DRO literature. The algorithm combines several ideas (variance reduction, minibatching, and cyclic coordinate-style updates) in a careful way. Of these, the application of cyclic coordinate-style updates to this setting (with a non-separable dual feasible set) seems the most novel and interesting to me. It is also very nice that it handles general uncertainty sets unlike previous work. The presentation is generally solid, although I think it could be improved in a few ways (see below).\n\nWEAKNESSES:\nRegarding presentation:\n1. I like Table 1 a lot, but I wish the comparison between your first-order complexity bound and previous complexity bounds was a bit more \"in your face\" and expansive. Currently, the comparisons are kind of spread out throughout the paper, requiring the reader to do more work than I think they should have to to understand where your method improves. For example, I like the discussion of Lines 471-479, but I think this should be in the main body. Also, I would compare your method directly to the sub-gradient method in the same place and carefully state what regimes you are winning in. \n2. I would formally put your first-order oracle complexity (and probably runtime too) into a theorem which also collects all of your assumptions together (or mentions Assumption 1). Currently, this is just in the body of the paper after Theorem 2. I think putting it into a theorem makes it more clear and easier to cite.\n3. The cases with $\\mu = 0$ or $\\nu = 0$ are handled in Appendix C.4, but not in a very formal way. E.g., a $O(1 / t)$ rate is mentioned, but I think the performance of your algorithm in these settings should be both stated carefully in a theorem and compared to prior work in this setting (e.g., the sub-gradient method as well as Levy et al. 2020, Carmon and Hausler 2022). Indeed, the $\\mu = 0$ and $\\nu = 0$ settings are important enough that entire papers have been written about them (and in fact, in some sense it is the original setting), so I think it is important that these are not an afterthought, and that the reader doesn't have to do any work to understand how the bounds you get in these settings compares to prior work. This is actually the single biggest weakness in my opinion.\n\nMinor/typos:\n1. Line 71 is missing a subscript. $\\ell(w)$ should be $\\ell_i(w)$.\n2. In the two inline lines above Line 431, there shouldn't be a square on the norm. I.e., it should just be $|| w_1 - w_2||_2$.\n3. Although it is clear what it is from context, I don't the the map $\\ell(w)$ is ever formally defined before you start using it in Line 97.\n4. \"witha\" in Line 7.\n5. It might be good to just mention that when you say \"smooth\" in the paper, you always mean it in the sense that the gradient is Lipschitz (and not in the sense of a $C^\\infty$ function).\n6. I would include the fact that $\\mathcal{Q}$ is any closed, convex, and nonempty set in Assumption 1. As far as I can tell, the only place you formally state this is in the text under Table 1.\n7. Maybe just use the term \"runtime\" instead of, e.g., \"global complexity.\" (Unless there is a reason you aren't using the term \"runtime\" that I'm missing?)\n8. It might be nice to mention some prior work which doesn't just have a single hyperparameter (or really, state which algorithms do and don't), just to make it more clear how significant of a contribution that is.\n\nQUESTIONS:\n1. Just to make sure I understand correctly, it is always best to choose $b = 1$ to get the best first-order oracle complexity, but your runtime may be better with a different choice of $b$?\n2. I'm confused by the $n q_{max} \\le n$ in Line 491. Isn't this always true? Similarly, Line 81 says that $n q_{max}$ is upper bounded by \"a constant in $n$,\" but isn't it just bounded by $n$ because $q_{max} \\le 1$? I thought $\\mathcal{Q}$ still needed to be a subset of the simplex per the text under Table 1. If not, I would fix that text and also include exactly what $\\mathcal{Q}$ is in Assumption 1.\n3. Regarding the case where $\\mu = 0$ or $\\nu = 0$, in Appendix C.4 you modify the entire algorithm/analysis to handle this case. It is good to know this works, but a classic trick to recover a non-strongly convex rate from the strongly-convex case is to just set something like, e.g., $\\mu = \\epsilon / D^2$ where $D$ is a bound on the diameter of $\\mathcal{W}$. Would this recover the same rate, and/or is there a reason you don't do this? Just curious.\n4. Just to make sure I understand correctly, is this the first work you are aware of which applies these cyclic coordinate-style updates to a non-separable dual feasible set, or is there any precedent for this?\n5. I'm a bit confused by Lines 199 to 204 and how they connect to Appendix B.3 and in particular the bound after Line 556. It is also claimed in Line 202 that $n q_{max}$ is constant - why is this? (Even after clarifying this to me, I would recommend connecting it more closely to Appendix B, i.e., cite the exact equation in Appendix B you are referencing as opposed to all of Appendix B.)\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies stochastic variance-reduced algorithms for regularized distribution robust optimization problems, formulated as a minimax problem. The paper proposed a novel primal-dual algorithm, and analyzed its iteration complexity. Under Lipschitz and smoothness assumptions, the algorithm achieves a linear convergence rate, with $O(n + d)$ per-iteration complexity, and total oracle calls of order $n + n q_{\\max} L/\\mu + n^{3/2} \\sqrt{G^2 / (\\mu \\nu)}$.\n\nSTRENGTHS:\nDistribution-robust estimation is an important question, and the minimax ERM formulation leads to a challenging optimization problem. This paper made solid progress in this direction. The algorithm idea is natural, and the result improves over several exisitng works along this line of research. Importantly, the result imposes no additional structural assumptions on the uncertainty set $\\mathcal{Q}$, except for convexity and an upper bound on the $\\ell^\\infty$ norm bounds for its elements.\n\nWEAKNESSES:\nWhen seeing $G$, $\\mu$ and $\\nu$ as constants, the oracle complexity scales as $n^{3/2}$ in terms of $n$. Although we get logarithmic dependence on $\\varepsilon$ and the result is already state-of-the-art, if we only track the dependence on $n$, it's worse than the vanilla subgradient methods, as listed in Table 1.\n\nAdditionally, the analysis relies on the Lipschitz smooth convex + strongly convex regularization structure. Is it possible to extend the analysis to non-strongly convex settings, or the case with convex individual functions but strongly convex summation?\n\nAnother weakness is the space complexity. Storing the gradient estimates $\\widehat{g}$ requires $O(n d)$ space. In many cases, this costs the same complexity as storing all the training data. So the advantage of stochastic algorithm is not very clear.\n\nQUESTIONS:\n- In order to minimize the total computational cost (assuming that the individual gradients can be computed in $O (d)$ time), what is the optimal choice of the minibatch size?\n- On a related note, if we choose the parameters optimally, how does the total computational cost compare to other algorithms. In particular, as the per-iteration cost depends on $n$ as well, it is important to compare with the full-batch gradient-based algorithms.\n\nLIMITATIONS:\nThe paper addressed limitations adequetly.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-30"}