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

Authors: ronak  mehta
jelena  diakonikolas
zaid  harchaoui
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 04:58:27

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.

Main paper: 
 introduction: 
 Contemporary machine learning research is increasingly exploring the phenomenon of distribution shift, in which predictive models encounter differing 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 min w∈W max q∈Q L 0 (w, q) := n i=1 q i ℓ i (w) ,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 asEQUATIONwhere µ, ν ≥ 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] ,Figure 1 : Visualization of Uncertainty Sets and Penalties. Each plot is a probability simplex in n = 3 dimensions with the uncertainty set as the colored portion. The black dots are optimal dual variables q ⋆ ν := arg max q∈Q n i=1 q i ℓ i (w) -νD(q∥1/n) for a fixed w ∈ W. As ν decreases, q ⋆ ν may shift toward the boundary of the uncertainty set. The combination of ν and D determines an "effective" uncertainty set, whose shape is given by the level sets of D. Our methods apply to both.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 on each example is accessed via a collection of n primal first-order oracles {(ℓ i , ∇ℓ i )} n i=1 . Quantitatively, we measure the performance of algorithms by the runtime or global complexity of elementary operations to reach within ε of the minimum of L(w) = max q∈Q L(w, q), whereas 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 q ⋆ (w) = arg max q∈Q L(w, q), followed by ∇ L(w) = n i=1 q ⋆ i (w)∇ℓ i (w) + µw.(3)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., 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 in 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] . While 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 We propose DRAGO, a minibatch primal-dual algorithm for the penalized DRO problem (2) that achieves ε-suboptimality inO n b + κ Q L µ + n b nG 2 µν ln 1 ε (4)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 + nq max 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 components, 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. 
the drago algorithm: 
 We present here the Distributionally Robust Annular Gradient Optimizer (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 define the Bregman divergence in this way for purely technical reasons, to gracefully account for cases such as computing ∆ ψ (q, q) when ψ is the negative entropy function and q lies on the boundary of the n-dimensional probability simplex. Finding a minimizer of ( 2) is equivalent to finding a saddle-point (w ⋆ , q ⋆ ) ∈ W × Q which satisfies max q∈Q L(w ⋆ , q) = L(w ⋆ , q ⋆ ) = min w∈W L(w, q ⋆ ).Uncertainty Set Runtime / Global Complexity (Big-Õ) Carmon and Hausler, 2022] 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/ε) LCVaR-SGD † Lχ2-SGD † 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 * [ℓ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/2LSVRG 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κ) dlog(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) max nδ + κnqmax , 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 := Batch Size Support Constrained n d + κQL/µ + d nG 2 /(µν) log(1/ε)Table 1 : Complexity Bounds of DRO Methods. Runtime or global complexity (i.e., the total number of elementary operations required to compute w satisfying max q∈Q L(w, q) -L(w ⋆ , q ⋆ ) ≤ ε.Throughout, we assume that each ℓ i is convex and µ, ν ≥ 0. The "Support Constrained" uncertainty set refers to all closed, convex sets of probability mass vectors and any 1-strongly convex penalty. This includes f -divergences and spectral risk measures, but not general Wasserstein balls. * Bounds hold in high probability. † Complexity is measured in our framework; see Appx. B for details.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.We specify 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 updatesw t := arg min w∈W a t v P t , w + a t µ 2 ∥w∥ 2 2 + C t-1 µ 2 ∥w -w t-1 ∥ 2 2(5)+ c t-1 µ 2 t-2 s=t-n/b ∥w -w s∨0 ∥ 2 2(6)EQUATIONFor 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 lt ∈ R n , qt ∈ R n , and ĝ ∈ 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 = (Kb + 1, . . . , Kb). On each iteration t, we randomly sample independent block indices I t , J t ∼ Unif[n/b] and defineEQUATIONAs 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 ( l0 , q0 , ĝ0 ) = (ℓ(w 0 ), q 0 , ∇ℓ(w 0 )), and for iteration t ≥ 1 update blockK t := (n/b) mod t + 1 via ( lt,k , qt,k , ĝt,k ) = (ℓ k (w t ), q t,k , ∇ℓ k (w t )) if k ∈ B Kt ( lt-1,k , qt-1,k , ĝt-1,k ) if k / ∈ B KtWhile q t in particular is always known by the algorithm, we use the approximation qt 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.We also 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 ĝ⊤ t-1 qt-1 in (8) when t = 1, this quantity can be maintained with O(bd) operations in every subsequent iteration as only b rows of qt and ĝt are editted in each iteration. For space complexity, we may consider storing the entire gradient table ĝ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 ĝ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] ). 
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.We measure suboptimality using the primal-dual gap:γ t := L(w t , q ⋆ ) -L(w ⋆ , q t ) - µ 2 ∥w t -w ⋆ ∥ 2 2 - ν 2 ∥q t -q ⋆ ∥ 2 2 ,where the saddle point (w ⋆ , q ⋆ ) exists under Asm. 1. Note that γ t ≥ 0, as it is the sum of nonnegative quantitiesL(w t , q ⋆ ) -L(w ⋆ , q ⋆ ) - µ 2 ∥w t -w ⋆ ∥ 2 2 ≥ 0 and L(w ⋆ , q ⋆ ) -L(w ⋆ , q t ) - ν 2 ∥q t -q ⋆ ∥ 2 2 ≥ 0.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 Distributionally Robust Annular Gradient Optimizer (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, l0 = ℓ(w 0 ), ĝ0 = ∇ℓ(w 0 ), and q0 = q 0 . for t = 1 to T do Sample blocks I t and J t uniformly on [n/b], and define K t = t mod (n/b) + 1. Primal Update:Set v P t-1 = ĝ⊤ t-1 qt-1 + at-1n atb i∈B I t (q t-1,i ∇ℓ i (w t-1 ) -qt-2,i ĝt-2,i ), update w t using (5). Update table ( lt,k , ĝt,k ) to (ℓ k (w t ), ∇ℓ k (w t )) if k ∈ B Kt or ( lt-1,k , ĝt-1,k ) if k / ∈ B Kt . Dual Update: Set v D t = lt + at-1n atb j∈B J t (ℓ j (w t ) -lt-1,j )e j , update q t using (7). Update table qt,k to q t,k if k ∈ B Kt or qt-1,k if k / ∈ B Kt . end for return (w T , q T ).Assumption 1. Let ℓ 1 , . . . , ℓ n be G-Lipschitz continuous and L-smooth, in that for all i ∈ [n],∥ℓ i (w) -ℓ i (w ′ )∥ 2 ≤ G ∥w -w ′ ∥ 2 and ∥∇ℓ i (w) -∇ℓ i (w ′ )∥ 2 ≤ L ∥w -w ′ ∥ 2 ,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 a 1 = 1, a 2 = 4α, and a t = (1 + α) a t-1 for t > 2, along with its partial sum A t = t τ =1 a τ . Under Asm. 1, there is an absolute constant C such that using the parameterα = C min b n , µ Lκ Q , b n µν nG 2 ,the iterates of Algorithm 1 satisfy:T t=1 a t E 1 [γ t ] + A T µ 4 E 1 ∥w T -w ⋆ ∥ 2 2 + A T ν 4 E 1 ∥q T -q ⋆ ∥ 2 2 ≤ nG 2 ν ∥w 0 -w 1 ∥ 2 2 .We can compute a point (w T , q T ) achieving an expected gap no more than ε with big-O complexityEQUATION)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:EQUATION)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 a t L(w ⋆ , q t ) ≥ a t q ⊤ t ℓ(w t ) + telescoping and non-positive terms + a t ∇ℓ(w t ) ⊤ q tv P t , w ⋆w t (12)+ c t µ 2 t-2 τ =t-n/b ∥w t -w τ ∨0 ∥ 2 2(13)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 writea t ∇ℓ(w t ) ⊤ q t -v P t , w ⋆ -w t = a t ∇ℓ(w t ) ⊤ q t -ĝ⊤ t-1 qt-1 , w ⋆ -w t - na t-1 b(1 + η) i∈It ⟨∇ℓ i (w t-1 )q t-1 -ĝt-2,i qt-2,i , w ⋆ -w t-1 ⟩ - na t-1 b(1 + η) i∈It ⟨∇ℓ i (w t-1 )q t-1 -ĝt-2,i qt-2,i , w t-1 -w t ⟩ .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 controlling1 b i∈It ∥∇ℓ i (w t-1 )q t-1 -ĝt-2,i qt-2,i ∥ 2 2in 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∈Jt (ℓ j (w t ) -lt-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 lt-1,i = ℓ i (w t-1-K∨0 ) for i ∈ B K .(EQUATIONBy computing the conditional expectationE t [•] := E[•|w t-1 ],we have thatE t   1 b j∈Jt (ℓ j (w t ) -lt-1,j ) 2 2   = 1 n n i=1 (ℓ i (w t ) -lt-1,i ) 2 2 = 1 n n/b K=1 i∈B K (ℓ i (w t ) -ℓ i (w t-1-K∨0 )) 2 2 ≤ G 2 b n t-2 τ =t-n/b ∥w t-1 -w τ ∨0 ∥ 2 2 ,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. 
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 squared error loss or multinomial cross-entropy loss, given byℓ i (w) := 1 2 y i -x ⊤ i w 2 and ℓ i (w) := -x ⊤ i w yi + log y∈Y exp x ⊤ i w y ,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 gapEQUATIONwhere 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).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üfekci, 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 In all plots, the y-axis measures the normalized primal (i.e., DRO risk) suboptimality gap, defined in (15). Columns represent a varying dual regularization parameter ν. On the first three columns the x-axis measures the number of individual first-order oracle queries to {(ℓ i , ∇ℓ i )} n i=1 and the remaining three the x-axis displays wall-clock time. The objective becomes ill-conditioned as ν decreases.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.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. 
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. 
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. 
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 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). 
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. 
 
 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. Q ∈ ∆ n Uncertainty set, or dual feasible set, which is closed and convex.Here,∆ n = {p ∈ [0, 1] n : n i=1 p i = 1}. q 0The uniform vector q 0 = 1/n = (1/n, . . . , 1/n). Used as a dual initialization.A statistical divergence between q ∈ Q and q 0 , which is1-strongly convex in its first argument with respect to ∥•∥ 2 . ℓ Loss function ℓ : W → R n , which is differentiable in each component on an open set O ⊆ R d such that W ⊆ O. µ ≥ 0 Primal regularization constant. ν ≥ 0 Dual regularization constant. L(w, q) Objective function L(w, q) := q ⊤ ℓ(w) -νD(q∥q 0 ) + µ 2 ∥w∥ 2 2 . (w ⋆ , q ⋆ ) Saddle point of L, which satisfies L(w ⋆ , q) ≤ L(w ⋆ , q ⋆ ) ≤ L(w, q ⋆ ) for all (w, q) ∈ W × Q. [n] Index set [n] = {1, . . . , n}. q max The value max q∈Q ∥q∥ ∞ . κ Q The condition number κ Q = nq max ≥ 1. G Lipschitz continuity constant of each ℓ i for i ∈ [n],for which|ℓ i (w) -ℓ i (w ′ )| ≤ G ∥w -w ′ ∥ 2 . L Lipschitz continuity constant of each ∇ℓ i for i ∈ [n],for which ∥∇ℓ i (w)-∇ℓ i (w ′ )∥ 2 ≤ L ∥w -w ′ ∥ 2 . ∇ℓ(w)Jacobian matrix of ℓ : R d → R n at w (shape = n × d).(a t ) Sequence of positive constants that weigh the average gap criterion.(A t )Sequence of partial sums of (a t ), or A t = t s=1 a s . The convergence rate will be given by A -1 t , so we have a t increase geometrically. b Batch or block size.M Number of blocks n/b. (w t , q t ) t≥0Sequence of primal and dual iterates.e j The j-th standard basis vector e j ∈ {0,1} n . lt Loss table, which approximates ℓ(w t ) ∈ R n . ĝt Gradient table, which approximates ∇ℓ(w t ) ∈ R n×d . qt Weight table, which approximates q t ∈ Q. E t [•] Shorthand for E [ • | H t-1 ], i.e., expectation conditioned on history H t-1 = σ (I s , J s )} t-1 s=1 .Table 2 : Notation used throughout the paper.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.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∇ℓ(w 1 ) ⊤ ∇h(ℓ(w 1 )) -∇ℓ(w 2 ) ⊤ ∇h(ℓ(w 2 )) 2 ≤ ∇ℓ(w 1 ) ⊤ (∇h(ℓ(w 1 )) -∇h(ℓ(w 2 ))) 2 + (∇ℓ(w 1 ) -∇ℓ(w 2 )) ⊤ ∇h(ℓ(w 2 )) 2 ] ≤ nG ν ∥ℓ(w 1 ) -ℓ(w 2 )∥ 2 + L ∥w 1 -w 2 ∥ 2 ≤ nG 2 ν + L ∥w 1 -w 2 ∥ 2 .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 ν .Examples of DRO Problems Our problemmin w∈W max q∈Q L(w, q) := q ⊤ ℓ(w) -νD(q∥1/n) + µ 2 ∥w∥ 2 2 (16)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 asQ = Q(ρ) := q : D f (q∥1/n) ≤ ρ n and D(q∥1/n) = D f (q∥1/n),whereD 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 thatQ = Q(σ) := conv ({permutations of σ)}) and D(q∥1/n) = D f (q∥1/n),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 ∥xx ′ ∥ 2 + β |yy ′ | for β > 0, and 3) the function Ψ is Lipschitz continuous with known constant.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 thatq i ∈ [0, 1/(nθ)] for all i ∈ [n]. Thus, nq i ≤ 1 θ for all i = 1, . . . , n.For the χ 2 -DRO uncertainty set, we have by direct computation thatn ∥q -1/n∥ 2 2 ≤ ρ n ⇐⇒ ∥q∥ 2 2 ≤ 1 + ρ n 2 .This implies that for any i ∈ [n], we have that q 2 i ≤ (1 + ρ)/n 2 which implies thatnq i ≤ 1 + ρ,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.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] × . . . × 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 directly apply 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 min-max 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:1. The objective has a finite-sum structure, in that it can be written asn i=1 f i (w, q), where f i (w, q) := q i ℓ i (w) -νD(q∥1/n) + µ 2 ∥w∥ 2 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. 3. The dual regularizer q → D(q∥1/n) is not necessarily smooth. This encompasses common statistical divergences such as the Kullback Leibler (KL). 4. 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 asEQUATION)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,n Lκ Q µ + nd nG 2 µνThis is an improvement over setting b = 1, in which case the same number is(n + d) Lκ Q µ + n(n + d) nG 2 µν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.Method Assumptions Uncertainty Set Runtime / Global Complexity (Big-Õ) 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 -SGD † 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 ε -2BROO * 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/2LSVRG 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 κ) dlog(1/ε)Prospect Prospect [Mehta et al., 2024 ] Lower Bound [Zhang et al., 2022] May not be decomposable May not be finite sum 4 : Complexity of Primal-Dual Saddle Point Methods. Half-life (defined as τ such that a linearly convergent method requires O(τ log(1/ε)) iterations to achieve ε-suboptimality) and periteration cost of linearly convergent optimization algorithms. The global number of arithmetic operations (under the assumption that (ℓ i , ∇ w ℓ i ) costs O(d) operations and ∇ q L costs O(n) operations) required to achieve a point (w, q) satisfying L(w, q ⋆ ) -L(w ⋆ , q) ≤ ε can be computed by multiplying the last two columns. The "Assumptions" column contains changes to the assumptions that each ℓ i is L-smooth and that D(•∥1/n) is ν-smooth.ℓ 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) max nδ + κnq max , 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 := Batch Size Support Constrained n d + κ Q L/µ + d nG 2 /(µν) log(1/ε)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.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 objectiveEQUATION)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 definew ⋆ = arg min w∈W max q∈Q L(w, q) and q ⋆ = arg max q∈Q min w∈W L(w, q).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. We assume 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 l0 = ℓ(w 0 ) ∈ R n , ĝ0 = ∇ℓ(w 0 ) ∈ R n×d , and q0 = 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:a 1 = 1, a 2 = 4α, and a t = (1 + α) a t-1 for t > 2.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.EQUATIONEQUATION)w t := arg min w∈W a t v P t , w + a t µ 2 ∥w∥ 2 2 + C t-1 µ 2 ∥w -w t-1 ∥ 2 2 + c t-1 µ 2 t-2 s=t-n/b ∥w -w s∨0 ∥ 2 2 (23)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 ct-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( lt,k , ĝt,k ) = (ℓ k (w t ), ∇ℓ k (w t )) if k ∈ B Kt ( lt-1,k , ĝt-1,k ) otherwise .Define e j to be the j-th standard basis vector. Then, sample a random block J t uniformly from [M ] computeEQUATION)EQUATION)EQUATION)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. Setqt,k = q t,k if k ∈ B Kt qt-1,k otherwise .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γ t = a t L(w t , q ⋆ ) -L(w ⋆ , q t ) - µ 2 ∥w t -w ⋆ ∥ 2 2 - ν 2 ∥q t -q ⋆ ∥ 2 2 (27)and aim to bound E 1 t s=1 γ s by a constant. Throughout the proof, we will use a free parameter α, with update ruleEQUATIONBecause 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.1. Lower bound the dual suboptimality a t L(w ⋆ , q t ).2. 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.3. Derive all conditions on the learning rate constant α and batch size b.4. Bound γ 1 and sum γ t over t for a T -step progress bound.5. 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.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 qt table to bound the term ∥q t-1q 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 , whereH t = {(I s , J s )} t s=1is 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.Lemma 4 (Cross Term Bound). Let (x t ) t≥1 , (y t ) t≥1 , and (ŷ t ) t≥1 denote random sequences of R mvalued vectors, and let x ⋆ ∈ R m be a vector. Denote by I t be the index of a uniform random blockEQUATIONFinally, let (x t , y t , ŷt ) t≥1 adapted to H t (as defined above). Then, for any positive constant γ > 0,a t E t [⟨y t -v t , x ⋆ -x t ⟩] ≤ a t E t [⟨y t -ŷt-1 , x ⋆ -x t ⟩] -a t-1 ⟨y t-1 -ŷt-2 , x ⋆ -x t-1 ⟩ + n 2 a 2 t-1 2γ E t 1 b i∈B I t (y t-1,i -ŷt-2,i ) 2 2 + γ 2 E t ∥x t -x t-1 ∥ 2 2 .Proof. By plugging in the value of v t and using x ⋆x t = x ⋆x t-1 + x t-1x t , we have thata t E t [⟨y t -v t , x ⋆ -x t ⟩] = a t E t [⟨y t -ŷt-1 , x ⋆ -x t ⟩] -na t-1 E t   1 b i∈B I t (y t-1 -ŷt-2 ), x ⋆ -x t   = a t E t [⟨y t -ŷt-1 , x ⋆ -x t ⟩] -a t-1 ⟨y t-1 -ŷt-2 , x ⋆ -x t-1 ⟩ + na t-1 E t   1 b i∈B I t (y t-1 -ŷt-2 ), x t -x t-1   ≤ a t E t [⟨y t -ŷt-1 , x ⋆ -x t ⟩] -a t-1 ⟨y t-1 -ŷt-2 , x ⋆ -x t-1 ⟩ + n 2 a 2 t-1 2γ E t 1 b i∈B I t (y t-1,i -ŷt-2,i ) 2 2 + γ 2 E t ∥x t -x t-1 ∥ 2 2 ,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 , ŷt = ĝ⊤ t qt , and x t = w t . In the dual case, we have that v t = v D t , y t = ℓ(w t ), ŷt = lt+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 qt,i = q π(t,i),i and ĝ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 qt in terms of past values of the iterates (e.g., q t ).Lemma 5 (Primal Noise Bound). When t ≥ 2, we have thatE t δ P t 2 2 ≤ 3q max n n i=1 q t-1,i ∥∇ℓ i (w t-1 ) -∇ℓ i (w ⋆ )∥ 2 2 + 3q max n n i=1 q π(t-2,i),i ∇ℓ i (w π(t-2,i) ) -∇ℓ i (w ⋆ ) 2 2 + 3G 2 n ∥q t-1 -qt-2 ∥ 2 2 (30)Proof. By definition, we have thatE t δ P t 2 2 = E t 1 b i∈B I t ∇ℓ i (w t-1 )q t-1,i -ĝt-2,i qt-2,i 2 2 = 1 b 2 E t i∈B I t ∇ℓ i (w t-1 )q t-1,i -ĝt-2,i qt-2,i 2 2 ≤ 1 b E t i∈B I t ∥∇ℓ i (w t-1 )q t-1,i -ĝt-2,i qt-2,i ∥ 2 2 = 1 n n i=1 ∥∇ℓ i (w t-1 )q t-1,i -ĝt-2,i qt-2,i ∥ 2 2 ,where we use that I t is drawn uniformly over n/b. Continuing again with the term above, we have1 n n i=1 ∥∇ℓ i (w t-1 )q t-1,i -ĝt-2,i qt-2,i ∥ 2 2 = 1 n n i=1 ∥(∇ℓ i (w t-1 ) -∇ℓ i (w ⋆ ))q t-1,i -(ĝ t-2,i -∇ℓ i (w ⋆ ))q t-2,i + (q t-1,i -qt-2,i )∇ℓ i (w ⋆ )∥ 2 2 ≤ 3 n n i=1 q t-1,i 2 ∥∇ℓ i (w t-1 ) -∇ℓ i (w ⋆ )∥ 2 2 + q2 t-2,i ∥ĝ t-2,i -∇ℓ i (w ⋆ )∥ 2 2 + (q t-1,i -qt-2,i ) 2 ∥∇ℓ i (w ⋆ )∥ 2 2 ≤ 3 n n i=1 q max q t-1,i ∥∇ℓ i (w t-1 ) -∇ℓ i (w ⋆ )∥ 2 2 + q max qt-2,i ∥ĝ t-2,i -∇ℓ i (w ⋆ )∥ 2 2 + (q t-1,i -qt-2,i ) 2 G 2 .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,E t δ D t 2 2 ≤ G 2 n t-2 τ =t-n/b ∥w t-1 -w τ ∨0 ∥ 2 2 .Proof. Then, we may exploit the coordinate structure of the noise term to writeE t δ D t 2 2 = E t 1 b j∈B J t (ℓ j (w t-1 ) -lt-1,j )e j 2 2 = 1 b 2 E t j∈B J t (ℓ j (w t-1 ) -lt-1,j )e j 2 2 + E t   j̸ =k (ℓ j (w t-1 ) -lt-1,j )(ℓ k (w t-1 ) -lt-1,k )⟨e j , e k ⟩   = 1 b 2 E t j∈B J t (ℓ j (w t-1 ) -lt-1,j )e j 2 2 = 1 b 2 E t j∈B J t |ℓ j (w t-1 ) -lt-1,j | 2 ≤ 1 bn n i=1 |ℓ i (w t-1 ) -lt-1,i | 2 ≤ G 2 n t-2 τ =t-n/b ∥w t-1 -w τ ∨0 ∥ 2 2 .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 lt-1 vector to have values equal to w t-1 before defining δ D t .In this section, we bound the ∥q t-1qt-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 decomposen i=1 (q t,i -qt-1,i ) 2 = M K=1 i∈B K (q t,i -qt-1,i ) 2 ,and analyze block-by-block. Our goal is to be able to count this quantity in terms of ∥q tq 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 i∈B K (q t,i -qt-1,i ) 2 ≤ i∈B K (t -t K ) t s=t K +1 (q s,i -q s-1,i ) 2 .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 numbert K = M ⌊[(t -1) -(K -1)]/M ⌋ + K = M ⌊(t -K)/M ⌋ + K. Then, write i∈B K (q t,i -qt-1,i ) 2 = i∈B K (q t,i -q t K ,i ) 2 = i∈B K t s=t K +1 q s,i -q s-1,i 2 ≤ i∈B K (t -t K ) t s=t K +1 (q s,i -q s-1,i ) 2 ,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 thatT t=1 a t n i=1 (q t,i -qt-1,i ) 2 ≤ T t=1 a t M K=1 i∈B K (t -t K ) t s=t K +1 (q s,i -q s-1,i ) 2 .We can start by swapping the first two sums and only considering the values of t that are greater than or equal to K.EQUATIONwhere 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) satisfiesS 0 ≤ eM (M -1) 2 M -1 t=1 a t ∥q t -q t-1 ∥ 2 2 . Proof. Write S 0 = M K=1 K-1 t=1 a t i∈B K t t s=1 (q s,i -q s-1,i ) 2 by definition = M -1 t=1 ta t M K=t+1 i∈B K t s=1(q s,iq s-1,i ) 2 swap sums over K and t= M -1 t=1 ta t t s=1 M K=t+1 i∈B K (q s,i -q s-1,i ) 2 move sum over s ≤ M -1 t=1 ta t t s=1 M K=1 i∈B K (q s,i -q s-1,i ) 2 M K=t+1 (•) ≤ M K=1 (•) = M -1 t=1 ta t t s=1 ∥q s -q s-1 ∥ 2 2 M K=1 i∈B K (•) = n i=1 (•) = M -1 s=1 M -1 t=s ta t ∥q s -q s-1 ∥ 2 2swap sums over s and t≤ M -1 s=1 M -1 t=s t(1 + α) t-s a s ∥q s -q s-1 ∥ 2 2 a t ≤ (1 + α)a t-1 ≤ M -1 s=1 a s (1 + α) M M -1 t=s t ∥q s -q s-1 ∥ 2 2 t -s ≤ M ≤ M -1 s=1 ea s M -1 t=s t ∥q s -q s-1 ∥ 2 2 α ≤ 1 M =⇒ (1 + α) M ≤ e ≤ M -1 s=1 ea s M -1 t=1 t ∥q s -q s-1 ∥ 2 2 M -1 t=s (•) ≤ M -1 t=1 (•) = eM (M -1) 2 M -1 s=1 a s ∥q s -q s-1 ∥ 2 2 ,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) satisfiesS 1 ≤ 2e 2 M 2 T t=1 a t ∥q t -q t-1 ∥ 2 2 .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 letS 1 = M K=1 T -K r=0 a r+K i∈B K (r + K -r K ) r+K s=r K +1 (q s,i -q s-1,i ) 2 by definition and r = t -K ≤ M M K=1 T -K r=0 a r+K i∈B K r+K s=r K +1 (q s,i -q s-1,i ) 2 (r -M ⌊r/M ⌋) ≤ M ≤ M M K=1 T -K r=0 (1 + α) K a r i∈B K r+K s=r K +1 (q s,i -q s-1,i ) 2 a t ≤ (1 + α)a t-1 ≤ eM M K=1 T -K r=0 a r i∈B K r+K s=r K +1 (q s,i -q s-1,i ) 2 K ≤ M and α ≤ 1 M = eM T -2 r=0 a r min{T -r,M } K=1 i∈B K r+K s=r K +1(q s,iq s-1,i ) 2 swap sums over r and K≤ eM T -2 r=0 a r min{T -r,M } K=1 i∈B K min{r+M,T } s=M ⌊r/M ⌋+1 (q s,i -q s-1,i ) 2 r+K s=r K +1 (•) ≤ min{r+M,T } s=M ⌊r/M ⌋+1 (•) = eM T -2 r=0 a r min{r+M,T } s=M ⌊r/M ⌋+1 min{T -r,M } K=1 i∈B K (q s,i -q s-1,i ) 2 move sum over s ≤ eM T -2 r=0 a r min{r+M,T } s=M ⌊r/M ⌋+1 M K=1 i∈B K (q s,i -q s-1,i ) 2 min{T -r,M } K=1 (•) ≤ M K=1 (•) = eM T -2 r=0 a r min{r+M,T } s=M ⌊r/M ⌋+1 ∥q s -q s-1 ∥ 2 2 M K=1 i∈B K (•) = n i=1 (•) ≤ eM T -2 r=0 min{r+M,T } s=M ⌊r/M ⌋ (1 + α) r-s a s ∥q s -q s-1 ∥ 2 2 a t ≤ (1 + α)a t-1 ≤ e 2 M T -2 r=0 min{r+M,T } s=M ⌊r/M ⌋ a s ∥q s -q s-1 ∥ 2 2 r -s ≤ M and α ≤ 1 M ≤ 2e 2 M 2 T -1 s=1 a s ∥q s -q s-1 ∥ 2 2 .The last line follows because the inner sum min{r+M,T } s=M ⌊r/M ⌋ a s ∥q sq s-1 ∥ 2 2 can be thought of as a sliding window of size at most 2M centered at r, which essentially repeats each element in the sum 2M 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 thatT t=1 a t n i=1 (q t,i -qt-1,i ) 2 ≤ 3e 2 M 2 T t=1 a t ∥q t -q t-1 ∥ 2 2 .Proof. Use Lem. 7 to achieve (31) and apply Lem. 8 and Lem. 9 on each term to getT t=1 a t n i=1 (q t,i -qt-1,i ) 2 ≤ eM (M -1) 2 M -1 t=1 a t ∥q t -q t-1 ∥ 2 2 + 2e 2 M 2 T t=1 a t ∥q t -q t-1 ∥ 2 2 ≤ 3e 2 M 2 T t=1 a t ∥q t -q t-1 ∥ 2 2 ,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 factorna 2 t-1 G 2µCt-1 . Thus, if we apply Lem. 10 when redistributing over time a term ∥q t-1q t-2 ∥ 2 2 which will be multiplied by a factor (ignoring absolute constants) ofna 2 t-1 G 2 µC t-1 • M 2 = n 3 a 2 t-1 G 2 b 2 µC t-1 .In order to cancel such a term, we require the use of -At-1ν2 ∥q t -q t-1 ∥ 2 2. can reserve half to be used up by (49), and be left with the conditionn 3 a 2 t-1 G 2 b 2 µC t-1 ≤ A t-1 ν ⇐⇒ α 2 ≤ b 2 n 2 µν 2nG 2 ,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 nG2 µν on α, which propagates to the learning rate α. We now proceed to the main logic of the convergence analysis.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 qt,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 α ≤ µ/(24eLκ Q ). Then, for t ≥ 2, we have that:-E t [a t L(w ⋆ , q t )] ≤ -a t E t [q ⊤ t ℓ(w t )] - a t 2L n i=1 q t,i E t ∥∇ℓ i (w t ) -∇ℓ i (w ⋆ )∥ 2 2 + a t-1 4L n i=1 q t-1,i ∥∇ℓ i (w t-1 ) -∇ℓ i (w ⋆ )∥ 2 2 + n i=1 a π(t-2,i) 4L q π(t-2,i),i ∇ℓ i (w π(t-2,i) -∇ℓ i (w ⋆ )) + 6nαa t-1 G 2 µ ∥q t-1 -qt-2 ∥ 2 2 -a t E t ∇ℓ(w t ) ⊤ q t -ĝ⊤ t-1 qt-1 , w ⋆ -w t + a t-1 ∇ℓ(w t-1 ) ⊤ q t-1 -ĝ⊤ t-2 qt-2 , w ⋆ -w t-1 - a t µ 2 E t ∥w t ∥ 2 2 + a t νE t [D(q t ||q 0 )] - µA t 2 E t ∥w ⋆ -w t ∥ 2 2 - µC t-1 4 ∥w t -w t-1 ∥ 2 2 + µC t-1 2 ∥w ⋆ -w t-1 ∥ 2 2 - c t-1 µ 2 t-2 τ =t-n/b E t ∥w t -w τ ∨0 ∥ 2 2 + c t-1 µ 2 t-2 τ =t-n/b ∥w ⋆ -w τ ∨0 ∥ 2 2 .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.a t L(w ⋆ , q t ) := a t q ⊤ t ℓ(w ⋆ ) + a t µ 2 ∥w ⋆ ∥ 2 2 -a t νD(q t ||q 0 ) (1) ≥ a t q ⊤ t ℓ(w t ) + a t ∇ℓ(w t ) ⊤ q t , w ⋆ -w t + a t 2L n i=1 q ∥∇ℓ i (w t ) -∇ℓ i (w ⋆ )∥ 2 2 + a t µ 2 ∥w ⋆ ∥ 2 2 -νD(q t ||q 0 ) (2) = a t q ⊤ t ℓ(w t ) + a t ∇ℓ(w t ) ⊤ q t -v P t , w ⋆ -w t + a t v P t , w ⋆ -w t + a t µ 2 ∥w ⋆ ∥ 2 2 -a t νD(q t ||q 0 ) + a t 2L n i=1 q t,i ∥∇ℓ i (w t ) -∇ℓ i (w ⋆ )∥ 2 2 + µC t-1 2 ∥w ⋆ -w t-1 ∥ 2 2 - µC t-1 2 ∥w ⋆ -w t-1 ∥ 2 2(32)+ c t-1 µ 2 t-2 τ =t-n/b ∥w ⋆ -w τ ∨0 ∥ 2 2 - c t-1 µ 2 t-2 τ =t-n/b ∥w ⋆ -w τ ∨0 ∥ 2 2 (3) ≥ a t q ⊤ t ℓ(w t ) + a t ∇ℓ(w t ) ⊤ q t -v P t , w ⋆ -w t cross term +a t v P t , w t -w t =0 + a t µ 2 ∥w t ∥ 2 2 -a t νD(q t ||q 0 ) + a t 2L n i=1 q t,i ∥∇ℓ i (w t ) -∇ℓ i (w ⋆ )∥ 2 2 + µC t-1 2 ∥w t -w t-1 ∥ 2 2 - µC t-1 2 ∥w ⋆ -w t-1 ∥ 2 2(33)EQUATION)Next, we are able to use Lem. 4 with v t = v P t , y t = ∇ℓ(w t ) ⊤ q t , ŷt = ĝ⊤ t qt , x t = w t , x ⋆ = w ⋆ , and γ = µC t-1 /2 which yields thatEQUATION)Then, apply Lem. 5 to achieveE t δ P t 2 2 ≤ 3q max n n i=1 q t-1,i ∥∇ℓ i (w t-1 ) -∇ℓ i (w ⋆ )∥ 2 2 + 3q max n n i=1 q π(t-2,i),i ∇ℓ i (w π(t-2,i) ) -∇ℓ i (w ⋆ ) 2 2 + 3G 2 n ∥q t-1 -qt-2 ∥ 2 2(36)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 writeE t [a t L(w ⋆ , q t )] := a t q ⊤ t ℓ(w ⋆ ) + a t µ 2 ∥w ⋆ ∥ 2 2 -a t νD(q t ||q 0 ) ≥ a t q ⊤ t ℓ(w t ) + E t [a t ∇ℓ(w t ) ⊤ q t -v P t , w ⋆ -w t ] + a t 2L n i=1 q t,i ∥∇ℓ i (w t ) -∇ℓ (w ⋆ )∥ 2 2 + a t µ 2 ∥w t ∥ 2 2 -a t νD(q t ||q 0 ) + µA t 2 ∥w ⋆ -w t ∥ 2 2 + µC t-1 2 ∥w t -w t-1 ∥ 2 2 - µC t-1 2 ∥w ⋆ -w t-1 ∥ 2 2 + c t-1 µ 2 t-2 τ =t-n/b ∥w t -w τ ∨0 ∥ 2 2 - c t-1 µ 2 t-2 τ =t-n/b ∥w ⋆ -w τ ∨0 ∥ 2 2 ,then apply (35) to the blue term above to writeE t [a t L(w ⋆ , q t )] ≥ a t q ⊤ t ℓ(w t ) + a t 2L n i=1 q t,i ∥∇ℓ i (w t ) -∇ℓ i (w ⋆ )∥ 2 2 - n 2 a t-1 2 µC t-1 E t ∇ℓ it-1 (w t-1 )q t-1,it-1 -g t-2,it-1 qt-2,it-1 2 2(37)+a t E t ∇ℓ(w t ) ⊤ q t -ĝ⊤ t-1 qt-1 , w ⋆ -w t -a t-1 ∇ℓ(w t-1 ) ⊤ q t-1 -ĝ⊤ t-2 qt-2 , w ⋆ -w t-1 + a t µ 2 ∥w t ∥ 2 2 -a t νD(q t ||q 0 ) + µA t 2 ∥w ⋆ -w t ∥ 2 2 + µC t-1 4 ∥w t -w t-1 ∥ 2 2 - µC t-1 2 ∥w ⋆ -w t-1 ∥ 2 2 + c t-1 µ 2 t-2 τ =t-n/b ∥w t -w τ ∨0 ∥ 2 2 - c t-1 µ 2 t-2 τ =t-n/b ∥w ⋆ -w τ ∨0 ∥ 2 2 .Finally, apply (36) to the term (37) to achieve≥ a t q ⊤ t ℓ(w t ) + a t 2L n i=1 q t,i ∥∇ℓ i (w t ) -∇ℓ i (w ⋆ )∥ 2 2 - 3κ Q a 2 t-1 µC t-1 n i=1 q t-1,i ∥∇ℓ i (w t-1 ) -∇ℓ i (w ⋆ )∥ 2 2 - 3κ Q a 2 t-1 µC t-1 n i=1 q π(t-2,i),i ∇ℓ i (w π(t-2,i) -∇ℓ i (w ⋆ )) 2 2 - 3nG 2 a 2 t-1 µC t-1 ∥q t-1 -qt-2 ∥ 2 2 + a t E t ∇ℓ(w t ) ⊤ q t -ĝ⊤ t-1 qt-1 , w ⋆ -w t -a t-1 ∇ℓ(w t-1 ) ⊤ q t-1 -ĝ⊤ t-2 qt-2 , w ⋆ -w t-1 + a t µ 2 ∥w t ∥ 2 2 -a t νD(q t ||q 0 ) + µA t 2 ∥w ⋆ -w t ∥ 2 2 + µC t-1 4 ∥w t -w t-1 ∥ 2 2 - µC t-1 2 ∥w ⋆ -w t-1 ∥ 2 2 + c t-1 µ 2 t-2 τ =t-n/b ∥w t -w τ ∨0 ∥ 2 2 - c t-1 µ 2 t-2 τ =t-n/b ∥w ⋆ -w τ ∨0 ∥ 2 2 .Next, we apply a t-1 ≤ 2αC t-1 and α ≤ µ 24eLκ Q to achieve:3nG 2 a 2 t-1 µC t-1 ≤ 6nG 2 αa t-1 µ , 3κ Q a 2 t-1 µC t-1 ≤ a t-1 4L , 3κ Q a 2 t-1 µC t-1 ≤ a π(t-2,i) 4L , ∀i ∈ [n]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 thata t-1 ≤ 1 + 1 n/b n/b a t-n/b ≤ ea t-n/b ≤ ea π(t-2,i) .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.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:a t L(w t , q ⋆ ) ≤ a t q ⋆ ⊤ (ℓ(w t ) -v D t ) + a t µ 2 ∥w t ∥ 2 2 + A t-1 ν∆ D (q ⋆ , q t-1 ) + a t q t ⊤ v Dta t νD(q t ||q 0 ) -A t-1 ν∆ D (q t , q t-1 ) -A t ν∆ D (q ⋆ , q t ).Proof. We add 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.a t L(w t , q ⋆ ) := a t q ⋆ ⊤ ℓ(w t ) -a t νD(q ⋆ ||q 0 ) + a t µ 2 ∥w t ∥ 2 2 (1) = a t q ⋆ ⊤ (ℓ(w t )-v D t ) + a t µ 2 ∥w t ∥ 2 2 +A t-1 ν∆ D (q ⋆ , q t-1 )) +a t q ⋆ ⊤ v D t -a t νD(q ⋆ ||q 0 )-A t-1 ν∆ D (q ⋆ , q t-1 )) (2) ≤ a t q ⋆ ⊤ (ℓ(w t ) -v D t ) + a t µ 2 ∥w t ∥ 2 2 + A t-1 ν∆ D (q ⋆ , q t-1 )) + a t q t ⊤ v D t -a t νD(q t ||q 0 ) -A t-1 ν∆ D (q t , q t-1 ) -A t ν∆ D (q ⋆ , q t ).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 α ≤ µ/(24eLκ 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 ) -lt ) -a t-1 (q ⋆ -q t-1 ) ⊤ (ℓ(w t-1 ) -lt-1 ) (38) +A t-1 ν∆ D (q ⋆ , q t-1 ) -A t νE t [∆ D (q ⋆ , q t )] (39) - a t 2L E t n i=1 q t,i ∥∇ℓ i (w t ) -∇ℓ i (w ⋆ )∥ 2 2(40)+ a t-1 4L n i=1 q t-1,i ∥∇ℓ i (w t-1 ) -∇ℓ i (w ⋆ )∥ 2 2 + n i=1 a π(t-2,i) 4L q π(t-2,i),i ∇ℓ i (w π(t-2,i) -∇ℓ i (w ⋆ )) (41) -a t E t ∇ℓ(w t ) ⊤ q t -ĝ⊤ t-1 qt-1 , w ⋆ -w t + a t-1 ∇ℓ(w t-1 ) ⊤ q t-1 -ĝ⊤ t-2 qt-2 , w ⋆ -w t-1 (42) + 6nG 2 αa t-1 µ ∥q t-1 -qt-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)EQUATION)For t = 2, the above holds with the addition of the term ν 4 ∥q 1 -q 0 ∥ 2 2 + nG 2 ν ∥w 0 -w 1 ∥ 2 2 .Proof. First, combine Lem. 11 and Lem. 12 to write: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 ≤ E t a t (q ⋆ -q t ) ⊤ (ℓ(w t ) -v D t ) cross term +E t [A t-1 ν∆ D (q ⋆ , q t-1 ) -A t ν∆ D (q ⋆ , q t )] -E t a t 2L n i=1 q t,i ∥∇ℓ i (w t ) -∇ℓ i (w ⋆ )∥ 2 2 + a t-1 4L n i=1 q t-1,i ∥∇ℓ i (w t-1 ) -∇ℓ i (w ⋆ )∥ 2 2 + n i=1 a π(t-2,i) 4L q π(t-2,i),i ∇ℓ i (w π(t-2,i) -∇ℓ i (w ⋆ )) 2 2 + 6nG 2 αa t-1 µ ∥q t-1 -qt-2 ∥ 2 2 -a t E t ∇ℓ(w t ) ⊤ q t -ĝ⊤ t-1 , w ⋆ -w t + a t-1 ∇ℓ(w t-1 ) ⊤ q t-1 -ĝ⊤ t-2 qt-2 , w ⋆ -w t-1 - µ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 - µC t-1 4 E t ∥w t -w t-1 ∥ 2 2 -A t-1 νE t [∆ D (q t , q t-1 )] - a t µ 2 E t ∥w t -w ⋆ ∥ 2 2 - c t-1 µ 2 t-2 τ =t-n/b E t ∥w t -w τ ∨0 ∥ 2 2 - a t ν 2 E it ∥q t -q ⋆ ∥ 2 2 .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 ), ŷt+1 = lt , x t = q t , x ⋆ = q ⋆ , and γ = νA t-1 = ν which yields thatEQUATION)where the fourth term above can be bounded asn 2 ν E 2 1 b j∈Jt (ℓ j (w 1 ) -l1,j )e j 2 2 ≤ nb ν 1 b b j=1 (ℓ j (w 1 ) -l1,j )e j 2 2 ≤ nG 2 ν ∥w 1 -w 0 ∥ 2 2 .Using the definition of the update, we have that ∥w 1w 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 thatEQUATION)We may then apply Lem. 6 and a t-1 ≤ αA t-2 to get thatEQUATION)We may use strong convexity to get the At-2ν4 E t ∥q t -q t-1 ∥ 2 2term to cancel with -At-1ν 2 ∆ D (q t , q t-1 ) and that A t-2 ≤ A t-1 to complete the proof.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 η > 0a t ≤ (1 + η)a t-1and 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- a t-1 4 -A t-1 + C t-1 + t+n/b-1 s=t+1 a s /m ≤ 0.The first term result from the "good term" -at-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 as4n/b -4 + m 4m a t-1 ≥ (1 + β) 2 m n/b-2 s=0 (1 + β) s a t-1 ,which can be summed and represented as(4n/b -4 + m) 4 ≥ (1 + β) 2 (1 + β) n/b-1 -1 β .Rearranging and taking a logarithm on both sides, this is the same asEQUATIONNext, using the inequality 2x 2+x ≤ ln(1 + x) with x = β(4n/b-4+m)4(1+β) 2which holds for all x ≥ 0, we haveEQUATIONWe can also apply the upper bound ln(x + 1) ≤ x with x = β (which also holds for any x ≥ 0) to write(n/b -1) log(1 + β) ≤ (n/b -1)β,which means that (51) will be satisfied (using (52)) ifEQUATIONIn order to satisfy the inequality, substitute m = 4cβ(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 asn/b -1 + m/2 n/b -1 = 1 + cβ(n/b -1).The RHS can be upper-bounded as(1 + β) 2 + (β/2)(n/b -1 + m/4) = (1 + β) 2 + (β/2)(n/b -1 + cβ(n/b -1) 2 ) ≤ 1 + β(2 + β) + β(n/b -1) 1 + c 2 ,which makes the inequality satisfied whenc ≥ 2 4 + 2β n/b -1 + 1 ,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 achievena t αG 2 ν ≤ µa t-1 m = µa t-1 4α(n/b -1) 2 by requiring that α ≤ b √ µν 4n 3/2 G ,which achieves the telescoping of (45). Finally, to address (46), we may satisfy it ifEQUATIONwhich we can achieve by incorporating the condition a t ≤ 4α(n/b -1) 2 C t-1 into the update of a t , because natαG 2 ν ≤ µat 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 a 2 = 4ηa 1 , anda t = (1 + η) a t-1 for t > 2, when η = 1 4 min b 32n , µ 24eLκ Q , b n nG 2 µν .we have that (28) holds.Proof. Noting that m = 16α(n/b -1) 2 we confirm thatC t ≥ A t-1 /2 ⇐⇒ A t - 1 16α(n/b -1) a t ≥ A t-1 /2 ⇐⇒ 1 - 1 16α(n/b -1) a t ≥ -A t-1 /2 ⇐⇒ 2 1 16α(n/b -1) -1 a t ≤ A t-1 .This condition is satisfied when α ≤ 1 32(n/b-1) , so we incorporate α ≤ b 32n 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αA t-1 = α(a t-1 + A t-2 ) ≥ (α + 1)a t-1 ≥ a t ,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:a t ≤ 4α(n/b -1) 2 C t-1 = 4α(n/b -1) 2 A t-1 - 1 4α(n/b -1) a t-1 = 4α(n/b -1) 2 a t-1 + A t-2 - 1 4α(n/b -1) a t-1which is implied by the inequalitya t ≤ 4α(n/b -1) 2 a t-1 + (1/α)a t-1 - 1 4α(n/b -1) a t-1 = 4(n/b -1) 2 α + 1 - 1 4(n/b -1) a t-1 .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 bya 2 = 4ηa 1 ≤ (1 + η) a 1 and a t = (1 + η) a t-1 for t > 2.Because each condition was satisfied when using a t = (1+α/4)a t-1 , we define η = α/4 to achieve the claimed result.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α ≤ min b 32n , µ 24eLκ Q , b 36e 2 n µν nG 2 .For any T ≥ 1, we have thatEQUATION(57)EQUATIONEQUATION-EQUATIONProof. 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 asa 1 L(w 1 , q ⋆ ) ≤ a 1 q ⋆ ⊤ (ℓ(w 1 ) -l1 ) + a 1 µ 2 ∥w 1 ∥ 2 2 + a 1 q 1 ⊤ l1 -a 1 νD(q 1 ||q 0 ) -A 1 ν∆ D (q ⋆ , q 1 ),where we use that l1 = l1 . For the lower bound, use a similar argument to Lem. 11 to achievea 1 L(w ⋆ , q 1 ) ≥ a 1 q ⊤ 1 ℓ(w 1 ) + a 1 ∇ℓ(w 1 ) ⊤ q 1 -∇ℓ(w 0 ) ⊤ q 0 , w ⋆ -w 1 + a 1 µ 2 ∥w 1 ∥ 2 2 -a 1 νD(q 1 ||q 0 ) + a 1 2L n i=1 q 1,i ∥∇ℓ i (w 1 ) -∇ℓ i (w ⋆ )∥ 2 2 + µA 1 2 ∥w ⋆ -w 1 ∥ 2 2 ,where we use that v P 0 = ∇ℓ(w 0 ) ⊤ q 0 . We combine them to getγ 1 ≤ a 1 (q ⋆ -q 1 ) ⊤ (ℓ(w 1 ) -l1 ) -a 1 ∇ℓ(w 1 ) ⊤ q 1 -∇ℓ(w 0 ) ⊤ q 0 , w ⋆ -w 1 - a 1 2L n i=1 q 1,i ∥∇ℓ i (w 1 ) -∇ℓ i (w ⋆ )∥ 2 2 - µA 1 2 ∥w ⋆ -w 1 ∥ 2 2 -A 1 ν∆ D (q ⋆ , q 1 ) - a 1 µ 2 ∥w 1 -w ⋆ ∥ 2 2 - a 1 ν 2 ∥q 1 -q ⋆ ∥ 2 2 ,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 termsa π(T -1,i) 4L E 0 q π(T -1,i) ∇ℓ i (w π(T -1,i) ) -∇ℓ i (w ⋆ )2 2 across iterations.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 sequencea 1 = 1, a 2 = 4α, and a t = (1 + α) a t-1 for t > 2,along with its partial sum A t = t τ =1 a τ . Under Asm. 1, there is an absolute constant C such that using the parameterα = C min b n , µ Lκ Q , b n µν nG 2 ,the iterates of Algorithm 1 satisfy:T t=1 a t E 1 [γ t ] + A T µ 4 E 1 ∥w T -w ⋆ ∥ 2 2 + A T ν 4 E 1 ∥q T -q ⋆ ∥ 2 2 ≤ nG 2 ν ∥w 0 -w 1 ∥ 2 2 .We can compute a point (w T , q T ) achieving an expected gap no more than ε with big-O complexityEQUATIONProof. 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 geta T E 0 (q ⋆ -q T ) ⊤ (ℓ(w T ) -lT ) ≤ νA T -1 4 E 0 ∥q ⋆ -q T ∥ 2 2 + a 2 T νA T -1 E 0 ∥ℓ(w T ) -lT ∥ 2 2 ≤ νA T -1 4 E 0 ∥q ⋆ -q T ∥ 2 2 + a 2 T G 2 νA T -1 T -2 τ =T -n/b E 0 ∥w T -w τ ∨0 ∥ 2 2 ≤ νA T -1 4 E 0 ∥q ⋆ -q T ∥ 2 2 + a T αG 2 ν T -2 τ =T -n/b E 0 ∥w T -w τ ∨0 ∥ 2 2 .The left-hand term will be canceled by ( 61) by applying strong concavity (leaving behindEQUATION2 ) and the right-hand term (because of the condition α ≤ √ µν 4n 3/2 G ) will be canceled by (60). Next, consider (56) . By Young's inequality with parameter µA T -1 /2, we have-a T E 0 ∇ℓ(w T ) ⊤ q T -ĝ⊤ T -1 qT -1 , w ⋆ -w T ≤ a 2 T µA T -1 E 0 ∇ℓ(w T ) ⊤ q T -ĝ⊤ T -1 qT -1 2 2 + µA T -1 4 E 0 ∥w ⋆ -w T ∥ 2 2 ≤ a T α µ E 0 ∇ℓ(w T ) ⊤ q T -ĝ⊤ T -1 qT -1 2 2 + µA T -1 4 E 0 ∥w ⋆ -w T ∥ 2 2 ,where the second term will be canceled by (59). For the remaining term,E 0 ∇ℓ(w T ) ⊤ q T -ĝ⊤ T -1 qT -1 2 2 ≤ E 0 (∇ℓ(w T ) -ℓ(w ⋆ )) ⊤ q T + ∇ℓ(w ⋆ ) ⊤ (q T -qT -1 ) + (∇ℓ(w ⋆ ) -ĝT -1 ) ⊤ qT -1 2 2 ≤ 3E 0 (∇ℓ(w T ) -∇ℓ(w ⋆ )) ⊤ q T 2 2 + 3E 0 ∇ℓ(w ⋆ ) ⊤ (q T -qT -1 ) 2 2 + 3E 0 (∇ℓ(w ⋆ ) -ĝT -1 ) ⊤ qT -1 2 2 ≤ 3σ n n i=1 E 0 q T,i ∥∇ℓ i (w T ) -∇ℓ i (w ⋆ )∥ 2 2 + 3nG 2 E 0 ∥q T -qT -1 ∥ 2 2 + 3σ n n i=1 E 0 qT -1,i ∥∇ℓ i (w ⋆ ) -ĝT -1,i ∥We may add the middle term above to (58), so that the remaining term to bound is6nG 2 α µ T +1 t=1 a t-1 E 0 ∥q t-1 -qt-2 ∥ 2 2 - T t=1 A t-1 ν 2 E 0 [∆ D (q t , q t-1 )] .To show that this quantity is non-negative, we use that a 0 = 0 and Lem. 10 (recalling that M = n/b to see that6nG 2 α µ E 0 T +1 t=1 a t ∥q t -qt-1 ∥ 2 2 ≤ 18e 2 n 3 G 2 α b 2 µ E 0 T +1 t=1 a t-1 ∥q t -q t-1 ∥ 2 2 ≤ 18e 2 n 3 G 2 α 2 b 2 µ E 0 T t=1 A t-1 ∥q t -q t-1 ∥ 2 2 ,which will cancel with the rightmost term in (58) provided that α ≤ b 36e 2 n µν nG 2 . Thus, plugging the previous displays into Lem. 15, we have thatE 0 T t=1 γ t ≤ n 2 G 2 νµ 2 ∇ℓ(w 0 ) ⊤ q 0 2 2 + 3a T σ n α 2µ n i=1 E 0 q T,i ∥∇ℓ i (w T ) -∇ℓ i (w ⋆ )∥ 2 2 - a T 2L n i=1 E 0 q T,i ∥∇ℓ i (w T ) -∇ℓ i (w ⋆ )∥ 2 2 + 3σ n a T α 2µ n i=1 E 0 q π(T -1,i),i ∇ℓ i (w ⋆ ) -∇ℓ i (w π(T -1,i) ) 2 2 - n i=1 (n -T + π(T -1, i)) a π(T -1,i) 4L E 0 q π(T -1,i) ∇ℓ i (w π(T -1,i) ) -∇ℓ i (w ⋆ ) 2 2 - T t=1 a t µ 4 E 0 ∥w t -w ⋆ ∥ 2 2 - T t=1 a t ν 4 E 0 ∥q t -q ⋆ ∥ 2 2 - A T -1 µ 4 E 0 ∥w T -w ⋆ ∥ 2 2 - A T -1 ν 4 E 0 ∥q T -q ⋆ ∥ 2 2 .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 2to the left-hand side achieves the claim.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.The modified algorithm is nearly identical to Algorithm 1, except that the dual and primal updates can be written asEQUATIONandEQUATIONrespectively, 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γ t = a t L(w t , q ⋆ ) -L(w ⋆ , q t ) - µ 2 ∥w t -w ⋆ ∥ 2 2 - ν 2 ∥q t -q ⋆ ∥ 2 2 (65)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 isa t = min C t-1 µ + µ 1 12enq max L , 1 + b n a t-1 , b 32n (A t-1 ν + ν 1 ) min {C t-1 µ + µ 1 , c t-1 µ + µ 2 } √ nG .Observe that when µ = 0, we set a t = µ1 12enqmax 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).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 bounda 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 ≤ Ct-1µ+µ1 12enqmax L and a t ≤ (1 + b/n)a t-1 , we have that:-E t [a t L(w ⋆ , q t )] ≤ -a t E t [q ⊤ t ℓ(w t )] - a t 2L n i=1 q t,i E t ∥∇ℓ i (w t ) -∇ℓ i (w ⋆ )∥ 2 2 + a t µ 2 ∥w t ∥ 2 2 + a t-1 4L n i=1 q t-1,i ∥∇ℓ i (w t-1 ) -∇ℓ i (w ⋆ )∥ 2 2 + n i=1 a π(t-2,i) 4L q π(t-2,i),i ∇ℓ i (w π(t-2,i) -∇ℓ i (w ⋆ )) 2 2 + 3nG 2 a 2 t-1 C t-1 µ + µ 1 ∥q t-1 -qt-2 ∥ 2 2 -a t E t ∇ℓ(w t ) ⊤ q t -ĝ⊤ t-1 qt-1 , w ⋆ -w t + a t-1 ∇ℓ(w t-1 ) ⊤ q t-1 -ĝ⊤ t-2 qt-2 , w ⋆ -w t-1 + a t νE t [D(q t ||q 0 )] - A t µ + µ 1 + µ 2 2 E t ∥w ⋆ -w t ∥ 2 2 - C t-1 µ + µ 1 4 ∥w t -w t-1 ∥ 2 2 + C t-1 µ + µ 1 2 ∥w ⋆ -w t-1 ∥ 2 2 - c t-1 µ + µ 2 2 t-2 τ =t-n/b E t ∥w t -w τ ∨0 ∥ 2 2 + c t-1 µ + µ 2 2 t-2 τ =t-n/b ∥w ⋆ -w τ ∨0 ∥ 2 2 .Similarly, following the same steps as Lem. 12, one can derive the following upper bound.Lemma 17. For t ≥ 2, we have that:a t L(w t , q ⋆ ) ≤ a t q ⋆ ⊤ (ℓ(w t ) -v D t ) + (A t-1 ν + ν 1 )∆ D (q ⋆ , q t-1 ) + a t q t ⊤ v D t -a t νD(q t ||q 0 ) -(A t-1 ν + ν 1 )∆ D (q t , q t-1 ) -(A t ν + ν 1 )∆ D (q ⋆ , q t ).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 ≤ Ct-1µ+µ1 12eLnqmax and a t ≤ (1 + b/n)a t-1 . 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 ) -lt )a t-1 (q ⋆q t-1 ) ⊤ (ℓ(w t-1 ) -lt-1 ) (66)+(A t-1 ν + ν 1 )∆ D (q ⋆ , q t-1 ) -(A t ν + ν 1 )E t [∆ D (q ⋆ , q t )] (67) - a t 2L E t n i=1 q t,i ∥∇ℓ i (w t ) -∇ℓ i (w ⋆ )∥ 2 2(68)+ a t-1 4L n i=1 q t-1,i ∥∇ℓ i (w t-1 ) -∇ℓ i (w ⋆ )∥ 2 2 + n i=1 a π(t-2,i) 4L q π(t-2,i),i ∇ℓ i (w π(t-2,i) -∇ℓ i (w ⋆ ))2 2(69)-a t E t ∇ℓ(w t ) ⊤ q t -ĝ⊤ t-1 qt-1 , w ⋆ -w t + a t-1 ∇ℓ(w t-1 ) ⊤ q t-1 -ĝ⊤ t-2 qt-2 , w ⋆ -w t-1 (70) + 3nG 2 a 2 t-1 C t-1 µ + µ 1 ∥q t-1 -qt-2 ∥ 2 2(71)EQUATION)For t = 2, the above holds with the addition of the term nG 2 ν+ν1 ∥w 1w 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.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µ   - a t-1 4 -A t-1 + C t-1 + t+n/b-1 s=t+1 a s /m   ≤0 + -(µ 1 + µ 2 (n/b -1)) + µ 1 + (n/b -1)µ 2 =0 ≤ 0,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 conditionsEQUATION)for the telescoping to be achieved, which can equivalently be rewritten asa t ≤ (A t-1 ν + ν 1 )(C t-1 µ + µ 1 ) 4nG 2 and a t ≤ (A t-1 ν + ν 1 )(c t-1 µ + µ 2 ) 2nG 2 (77)These can be accomplished by settinga t ≤ (A t-1 ν + ν 1 ) min {C t-1 µ + µ 1 , c t-1 µ + µ 2 } 2 √ nG .We must also handle the term3nG 2 a 2 t-1 Ct-1µ+µ1 ∥q t-1 -qt-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 thatT t=1 3nG 2 a 2 t µ t ∥q t -qt-1 ∥ 2 2 ≤ 3nG 2 C t-1 µ + µ 1 • 3e 4 M 2 T t=1 a 2 t ∥q t -q t-1 ∥ 2 2 ≤ 512nG 2 M 2 C t-1 µ + µ 1 T t=1 a 2 t ∥q t -q t-1 ∥ 2 2where 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 is512nG 2 M 2 C t-1 µ + µ 1 a 2 t ≤ A t-1 ν + ν 1 2which can also be rewritten asa t ≤ 1 32M (A t-1 ν + ν 1 )(C t-1 µ + µ 1 ) nG 2Thus, combining the previous conditions, our final update on the (a t ) sequence isa t = min C t-1 µ + µ 1 12enq max L , 1 + b n a t-1 , b 32n (A t-1 ν + ν 1 ) min {C t-1 µ + µ 1 , c t-1 µ + µ 2 } √ nG ,as alluded to in Appx. C.4.1.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.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 l, l1 ∈ R n , ĝ1 , ĝ2 ∈ R n×d , and q1 , q2 ∈ R n . For any iterate t, these variables are meant to store lt , lt-1 ∈ R n , ĝt-1 , ĝt-2 ∈ R n×d , and qt-1 , qt-2 ∈ R n .• The quantities ĝagg ∈ R d and ŵ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 l[B Kt ] is used to fill l1 [B Kt ] at the end of the loop, we we must maintain knowledge of l[B Kt ] 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.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 byD f (q∥q) := n i=1 f q i qi qi ,where we define 0f (0/0) := 0. For any i such that qi = 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ι X (x) := 0 if x ∈ X 1 otherwise .In either of the two cases below, we select the penalty D(q∥1/n) = D f (q∥1/n) to be an fdivergence. Denote in addition f * as the Fenchel conjugate of f . are related as q opt (l) = ∇(νD f (•∥1 n /n)) * (lz opt π -1 (l)), that is,EQUATIONAs 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 (PAV) algorithm.Another common uncertainty set format is a ball in f -divergence, orEQUATION1 ⊤ q = 1 .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 mappingEQUATIONwhere 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 2n .Thus, the optimal solution to (80) can be computed by solvingmax q∈Pn min λ≥0 ⟨l, q⟩ - ν 2 ∥q∥ 2 2 -λ 1 2 ∥q∥ 2 2 -ρ - 1 2n ,or equivalently, by strong duality via Slater's condition, solvingmax λ≥0 f (λ) := (ν + λ) min q∈Pn 1 2 ∥q -l/(ν + λ)∥ 2 2 -λ ρ + 1 2n - 1 2(ν + λ) ∥l∥ 2 2 .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 byf ′ (λ) := 1 2 q opt (λ) 2 2 -ρ - 1 2n ,where q opt (λ) achieves the minimum in min q∈Pn 1 2 ∥ql/(ν + λ)∥ 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.1. 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.2. 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.3. 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 ∥qq t ∥ 2 2 , which is true when D is the χ 2 -divergence, thenq t = arg max q∈Pn 1 2 ∥q-1/n∥ 2 2 ≤ ρ n ⟨l + νβ t q t , q⟩ - ν(1 + β t ) 2 ∥q -1/n∥ 2 2 ,which is a particular case of (80), and hence can be solved using the exact same procedure. We describe details of the experimental setup, including datasets, compute environment, and hyperparamater tuning. We largely maintain the benchmarks of Mehta et al. [2023] .The sample sizes, dimensions, and source of the datasets are summarized in Tab. 5. The tasks associated with each dataset are listed below. factors [Tüfekci, 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.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. Figure 4 : Benchmarks on the χ 2 Uncertainty Set. In both panels, the y-axis measure the primal suboptimality gap, defined in (15). Individual plots correspond to particular datasets. Left: The x-axis displays the number of individual first-order oracle queries to {(ℓ i , ∇ℓ i )} n i=1 . Right: The x-axis displays wall-clock time.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.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.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. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.38th Conference on Neural Information Processing Systems (NeurIPS 2024).2 


Decision: The paper received positive reviews, with a consensus that this is a solid contribution to the literature on optimization algorithms for distributionally robust optimization. Consequently, I recommend acceptance of this paper. When preparing the camera-ready revision, please pay close attention to the reviewers’ feedback and include the additional explanations provided during the rebuttal period. In addition, in Table 1, the complexities listed for Levy et al. do not agree with the best results obtained in that paper (and shown in Table 1 of Levy et al.) - please correct this in the revision.

Reviews: review: paper_summary: paper_summary: The 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)}$.
strength_weakness: strength_weakness: When 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.

Additionally, 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?

Another 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.
questions: questions: - 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?
- 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.
limitations: limitations: The paper addressed limitations adequetly.
paper_summary: paper_summary: The 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:
- 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.
- It achieves fine-grained dependence on the primal and dual condition numbers (i.e., not having a max over min term).
- 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.
- 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.
- 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).
strength_weakness: strength_weakness: Regarding presentation:
1. 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. 
2. 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.
3. 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.

Minor/typos:
1. Line 71 is missing a subscript. $\ell(w)$ should be $\ell_i(w)$.
2. 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$.
3. 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.
4. "witha" in Line 7.
5. 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).
6. 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.
7. 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?)
8. 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.
questions: questions: 1. 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$?
2. 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.
3. 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.
4. 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?
5. 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.)
limitations: limitations: Yes.
paper_summary: paper_summary: This paper proposes DRAGO, a variance-reduced optimization method for
distributionally robust optimization (DRO). In particular, the authors consider
the penalized DRO framework and view DRO as min-max optimization with a simple
finite-sum structure. Their algorithm maintains explicit histories for the
individual loss functions $l_i$, loss function gradients $\nabla l_i$, and dual
parameters $q_i$ (i.e. like SAG/SAGA); these estimates are used to control the
variance in both the primal and dual updates, which the authors prove leads to
a fast linear rate of convergence.  Unlike existing methods, DRAGO supports any
penalty parameter $\nu \geq 0$ and its convergence rate on depends explicitly
on the different problem constants.
strength_weakness: strength_weakness: - The paper organization makes it quite difficult to understand the different
    algorithmic components used in DRAGO.
    
- The assumptions necessary for convergence are somewhat unrealistic, although
    these appear to be standard in the literature. 

- The experimental evaluation is small and limited to relatively low-dimensional
    problems.

### Detailed Comments

This 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. 

**Paper Organization**:
I found this paper quite confusing to read. Instead of introducing DRAGO and
then explaining it's various algorithmic components, the authors first try to
explain the intuition behind the updates. This is a good idea in theory, but in
practice there are too many moving parts to DRAGO and I quickly got lost in the
(mostly unimportant) details. The actually key parts of DRAGO --- proximal-gradient
updates with SAG/SAGA style variance reduction --- are barely mentioned and this
makes the algorithm much harder to understand.

My suggestion is to introduce Algorithm 1 first and then explain why 
additional proximal terms are needed in the proximal-gradient update in order
to control the variance. At least this way the reader has a reference point
for understanding the discussion.

**Assumptions**:
I think that assuming each $l_i$ is both Lipschitz continuous and Lipschitz
smooth is quite unrealistic. Typically only one or the other of these
conditions holds. For instance, Lipschitz continuity implies that $l_i$ is
bounded above by a linear function, while Lipschitz smoothness implies that
$l_i$ is bounded above by a quadratic. While both can be true, the quadratic
bound will generally be vacuous. However, I understand from Table 1 that these
assumptions are typical for the literature, so I don't think this is too much
of an issue.

**Experimental Evaluation**: 
I would have liked to see experiments on a higher dimensional dataset.  Since
the entire Jacobian $\nabla l(w)$ must be maintained in the history, SAG-type
methods usually have high memory requirements as $d$ increases. The authors
partially handle this by using a large batch-size, but this is only feasible if
the mini-batches can still be fit on the GPU. When $d$ is sufficiently large,
mini-batch gradients will be inefficient to compute and I am curious if the
relative performance of DRAGO and LSVRG ill switch. This is particularly 
worth considering because it is only in the large-batch setting that DRAGO
seems to outperform LSVRG

### Minor Comments

- Line 52: This is a super-linear complexity, but a sub-linear convergence rate.
    Contrast with Line 62, which shows a linear rate, but a logarithmic complexity.

- Line 65: Should this read "but do not"?

- Line 24/71: It seems like you are only considering $\nu$ in the probability simplex.
    It would be nice if you stated this explicitly somewhere.

- Line 80: It's generally unrealistic for both $l_i$ and $\nabla l_i$ to be 
    Lipschitz continuous over an unbounded set.

- Line 86: You should mention that strong convexity comes only from the regularizer,
    as otherwise it seems to conflict with Lipschitzness of $l_i$. 

- Line 102: Does $\nabla l(w_{t-1})$ denote the Jacobian of the vector-valued map
    $l$ such that $[l(w)]_i = l_i(w)$? It would be nice if you clearly defined
    this notation.

- Table 1: The font-size is much too small to read when printed. 

- Algorithm Description: The purposes of $\bar{\beta}$ and $\beta_t$ and their
    settings are not well described.

- Display after Line 164: You seem to be using $D(q \| 1/n)$ is
    1-strongly-convex for this result, but this assumption is only introduced
    in Assumption 1. Maybe move Assumption 1 to the start of the discussion
    of convergence to fix this?

- Figure 3: The font-size is much too small to be readable without zooming.
questions: questions: - Line 491: Isn't $n q_{\text{max}} \leq n$ trivially satisfied since $q$
    is constrained to the probability simplex?

- Figure 2: Why is $b = 1$ only show for three datasets and only in 
    terms of oracle complexity?

- Display after Line 124: It looks you are using the equality,

$$\\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})]$$

but I don't think this is true because the terms in the product are correlated. 
Can you please explain this?
limitations: limitations: Limitations appear to be appropriate addressed in the paper.
paper_summary: paper_summary: The authors provide a state-of-the-art optimization algorithm with convergence guarantees for f-divergence DRO. The numerical study is also extensive
strength_weakness: strength_weakness: As the author has commented, they only solve for convex optimization formulation. The extension to non-convex setup also seems interesting and important.
questions: questions: N/A
limitations: limitations: N/A
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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: limitations: No direct societal impact as the paper is focused on theoretical aspects of the out of distribution inference problem.

Comments: title: Response
Response to Reviewer VCV6
Thank you for answering my questions
Thank you for your feedback / Addressing additional concerns
Response to Reviewer VCV6
Response to Reviewer em1q
Response to Reviewer VCV6 (continued)
comment: I thank the authors for answering my questions. I will keep my score.
Thanks for the response. I appreciate the detailed explanation about storage costs and comparison to other methods, and I request the authors to add those discussion in the revised version. I will keep my score.
Thank you once again for your review and timely replies! We will reflect this discussion in the final version and include the associated experiments and figures for varying batch size.
Thanks for providing these additional experimental details. It seems clear that the batch-size is a critical tuning parameter for Drago and I think including a larger-scale ablation study would definitely strengthen the submission. At the very least, an enhanced discussion of the batch-size (similar to what we've had here) should be included in experimental section of the paper. 

Overall, I think this paper is a solid submission that should be accepted. I will maintain my score for now and reevaluate during the reviewer discussion period if necessary.
Given that my questions have been satisfactorily addressed, I'd like to raise my score.
This addresses my question - thank you! (I definitely think adding this explanation or something similar to the paper would be good - especially people who come from a more purely optimization as opposed to statistical background may be confused otherwise.)
Dear reviewers,

As the discussion period draws to a close on August 13 (2 business days from now) we kindly request that you take a moment to review and acknowledge our responses to your questions and comments.

In response to the comments raised by Reviewer WLa8, P67z, VCV6 we have added additional experiments that demonstrate the performance of Drago for higher-dimensional problems with small batch sizes. We have better clarified our theoretical assumptions, such as those on the losses $\ell_i$ and the size of the uncertainty set $\mathcal{Q}$. Please see the responses for details.

If any further concerns or questions arise that we could address, please do not hesitate to reach out. We appreciate your time and feedback.

Sincerely,

The authors
Thank you for these comments and your thorough review. Please let us know if we have addressed them fully.

>**Using cyclic updates makes your method somewhat more similar to the incremental aggregated gradient method (IAG) [1].**

Thank you for this highly relevant reference. We have updated the discussion of related work to include it and give context to our method.

>**I don't like compactness assumptions for general machine learning problems, but this is worth adding as a comment to the paper.**

We have added discussion on cases in which Lipschitzness and smoothness of each $\ell_i$ are valid to be clear about the limitations of our assumptions. 

>**I was hoping to see an ablation study on the role of $b=1$ in the UCI experiments.**

We may have misunderstood your original comment on experimental evaluations when pursuing the "large $d$" experiments&mdash;thank you for the clarification. When running $b=1$ on the UCI datasets, we see that it is only competitive with baselines and other Drago variants when looking at suboptimality against gradient evaluations for the small $n$ datasets (yacht, energy, concrete), which is shown in Figure 2 of the manuscript. In terms of suboptimalty against wall time, the $b=1$ variant generally performs worse than other Drago variants and LSVRG. We collect some of these results in the table below for concreteness. The first table is from yacht, which is generally representative of the three smaller UCI datasets. The second table is from kin8nm, which is generally representative of the larger UCI datasets and ascincome (which is sourced from Fairlearn).

**Performance against Wall Time (yacht)**
|  Algorithm | Suboptimality after 2 seconds |
| ------- | ------- |
| LSVRG  | $10^{-8.30}$   |
| Drago ($b=1$) | $10^{-4.13}$   |
| Drago ($b=16$)    | $10^{-7.96}$   |
| Drago ($b=n/d$)    | $10^{-8.21}$   |

**Performance against Wall Time (kin8nm)**
|  Algorithm | Suboptimality after 5 seconds |
| ------- | ------- |
| LSVRG  | $10^{-8.23}$   |
| Drago ($b=1$) | $10^{-1.66}$    | 
| Drago ($b=16$)    | $10^{-4.97}$    | 
| Drago ($b=n/d$)    | $10^{-8.23}$    | 

We hypothesize that the key measure of problem size is $n/d$ (which happens to be the theoretically optimal batch size). When both $n$ and $d$ are large (as seen in the experiments on emotion in the supplemental PDF), a small batch size ($b \approx 2$) can perform quite well. When $d$ is small relative to $n$, which is generally true of the UCI datasets, large batch sizes perform drastically better. It is worth mentioning that because $d$ is small in these scenarios, it is usually not computationally prohibitive to use larger batch sizes.

Thank you again and please let us know if there are more questions we can answer.
Thank you for this recommendation and for improving the clarity of our work&mdash;we have added the assumption that the risk parameters (such as $\theta$) will be of constant order to make this claim. 

To give some background on this assumption, both spectral risk measures and $f$-divergence balls can be described using statistical functionals applied to the empirical distribution of the data. The dependence on $n$ typically comes only from the empirical distribution, and not the functional itself (although we will make this explicit in an assumption as you correctly pointed out). 

For any real-valued random variable $Z$ with cumulative distribution function (CDF) $F$ and finite first moment, the (population) $\theta$-CVaR of $Z$ is defined as $\mathbb{E}[Z | Z \geq F^{-1}(1-\theta)]$, where $F^{-1}$ is the quantile function or generalized inverse of $F$. In other words, it is the conditional mean of $Z$ given that $Z$ is above its $(1-\theta)$-quantile. For i.i.d. data $Z_1, \ldots, Z_n$ with empirical CDF $F_n$, we may define the (sample) $\theta$-CVaR as $\mathbb{E}[Z | Z \geq F_n^{-1}(1-\theta)]$ (which is roughly the mean of the top $n\theta$ order statistics). When applied to the empirical distribution of training losses, this recovers the description of CVaR given in Appendix B.2. When $\theta = \Theta(1)$, statistical arguments yield the convergence of the sample quantity to the population quantity, which is why it is typically assumed.

Thank you again for this discussion and please let us know if we have addressed your questions fully.
> SAG-type methods usually have high memory requirements as $d$ increases.

Yes, I'm very familiar with the least-squares-type arguments for reducing the memory requirements for SAG. Using cyclic updates makes your method somewhat more similar to the incremental aggregated gradient method (IAG) [1]. This algorithm is the conceptual predecessor of SAG, so it's probably worth referencing in your paper. 

> The experimental evaluation is small and limited to relatively low-dimensional problems.

Thanks for including these new experiments. It's nice to see that Drago works well when the batch-size is quite small for at least one problem, but I think I was talking more about the regression experiments in Figure 2. Here, Drago only seems to work well in the large-batch setting (which your comment about $b = 1$ seems to confirm). Thus, I was hoping to see an ablation study on the role of $b$ in the UCI experiments. 

> The paper organization makes it quite difficult to understand the different algorithmic components used in Drago.

Sounds good. 

> The assumptions necessary for convergence are somewhat unrealistic, although these appear to be standard in the literature. I think that assuming each $\ell_i$ is both Lipschitz continuous and Lipschitz smooth is quite unrealistic. Typically only one or the other of these conditions holds.

I don't like compactness assumptions for general machine learning problems, but this is worth adding as a comment to the paper. The Huber loss example is nice, but not particularly used in modern ML, although maybe it remains popular in robust statistics --- I'm not sure. 

Thanks for the clarifications on the remaining points.  

[1] Blatt, Doron, Alfred O. Hero, and Hillel Gauchman. "A convergent incremental gradient method with a constant step size." SIAM Journal on Optimization 18.1 (2007): 29-51.
Thank you for these clarifications. The one answer I still don't fully follow is the one related to $n q_{max}$. I now realize Line 81 is connected to Line 468 in Appendix B. Line 468 reads: "$n q_{max}$ is upper bounded by a constant independent of $n$ in common cases" and then goes on to cite $\theta$-CVaR and $\chi^2$-DRO with radius $\rho$ as examples. I don't agree with this though in general. For $\theta$-CVaR for example, the constraint is $q_{max} \le \frac{1}{\theta n}$ as you wrote in your response. Then it is true that $n q_{max}$ is a constant if you restrict to $\theta = \Omega(1)$. But to my knowledge, the $\theta$-CVaR DRO problem encompasses all $\theta \in [1 / n, 1]$, in which case it is not correct to claim that $n q_{max}$ is always a constant for $\theta$-CVaR. (Indeed, the entire range $\theta \in [1 / n, 1]$ is considered in Levy et al. 2020. And it is even pointed out later in your Appendix B - Line 489 - that previous methods in the literature don't beat gradient descent when $\theta = 1 / n$.)

It doesn't seem to me that the strength of your contribution relies on $n q_{max}$ being a constant as opposed to being as large as $n$ per the rates in Table 1. However, I would avoid claiming $n q_{max}$ is a constant in general for $\theta$-CVaR, $\chi^2$-DRO, etc. It seems to me like the point you want to make (?) is that if $\theta = \Omega(1)$, then your rate improves, but then you must make the $\theta = \Omega(1)$ restriction clear (and similarly for $\chi^2$-DRO, etc.).

Please let me know if I'm still misunderstanding something however.

(Also just to add to my original question - I would suggest avoiding the phrase "a constant in $n$" and just say "a constant" or "a constant independent of $n$." I understand it now, but when I first read "a constant in $n$," I thought it meant $cn$ for some constant $c$.)
>**Minor Comments: ...**

Thank you for the additional “Minor Comments”. We have incorporated them into the manuscript. To address the questions:
- Line 102: This is the full Jacobian - thank you for the suggestion.
- Algorithm Description: These constants are defined by simplifying the algorithm description used in the analysis. We have made the algorithm description consistent with the analysis so that these constants can be avoided.