Stochastic Approximation Approaches to Group Distributionally Robust Optimization
Authors: Lijun Zhang
Peng Zhao
Zhenhua Zhuang
Tianbao Yang
Zhi-Hua Zhou
Venue: NeurIPS 2023 poster (2023)
Keywords: Group distributionally robust optimization
Stochastic mirror descent
Non-oblivious online learning
Sample complexity
Stochastic mirror-prox algorithm
Mini-batch
TL;DR: We developed four stochastic approximation approaches for Group DRO.

Abstract: This paper investigates group distributionally robust optimization (GDRO), with the purpose to learn a model that performs well over $m$ different distributions. First, we formulate GDRO as a stochastic convex-concave saddle-point problem, and demonstrate that stochastic mirror descent (SMD), using $m$ samples in each iteration, achieves an $O(m (\log m)/\epsilon^2)$ sample complexity for finding an $\epsilon$-optimal solution, which matches the $\Omega(m/\epsilon^2)$ lower bound up to a logarithmic factor. Then, we make use of techniques from online learning to reduce the number of samples required in each round from $m$ to $1$, keeping the same sample complexity. Specifically, we cast GDRO as a two-players game where one player simply performs SMD and the other executes an online algorithm for non-oblivious multi-armed bandits. Next, we consider a more practical scenario where the number of samples that can be drawn from each distribution is different, and propose a novel formulation of weighted GDRO, which allows us to derive distribution-dependent convergence rates.  Denote by $n_i$ the sample budget for the $i$-th distribution, and assume $n_1 \geq n_2 \geq \cdots \geq n_m$. In the first approach, we incorporate non-uniform sampling into SMD such that the sample budget is satisfied in expectation, and prove that the excess risk of the $i$-th distribution decreases at an $O(\sqrt{n_1 \log m}/n_i)$ rate. In the second approach, we use mini-batches to meet the budget exactly and also reduce the variance in stochastic gradients, and then leverage stochastic mirror-prox algorithm, which can exploit small variances, to optimize a carefully designed weighted GDRO problem. Under appropriate conditions, it attains an $O((\log m)/\sqrt{n_i})$ convergence rate, which almost matches the optimal $O(\sqrt{1/n_i})$ rate of only learning from the $i$-th distribution with $n_i$ samples.

Stochastic Approximation Approaches to Group Distributionally Robust Optimization
Lijun Zhang 1.2
Peng Zhao', Zhen-Hua Zhuang', Tianbao Yang3, Zhi-Hua Zhou' National Key Laboratory for Novel Software Technology. Nanjing University, Nanjing, China
?Peng Cheng Laboratory, Shenzhen 51 81505, China
& Department of Computer Science and Engineering, Texas A&M University, College Station, USA {zhanglj, zhaop, zhuangzh, zhouzh) @lamda, nju. edu. en tianbao- yataaao-aannttmmu.e.. edu
Abstract
This paper investigates group distributionally robust optimization (GDRO), with the purpose to learn model that performs well over m different distributions. First, formulate GDRO stochastic convex-concave saddle-point problem, and demonstrate that stochastic mirror descent (SMD), using m samples in each iteration, achieves an O(m (log m)/²²) sample complexity for finding e-optimal solution, which matches the ((m/e²) lower bound up to a logarithmic factor. Then, make use of techniques from online learning reduce the number of samples required in each round from m to 1, keeping the same sample complexity. Specifically, we cast GDRO as a two-players game where one player simply performs SMD and the other executes an online algorithm for non-oblivious multi- armed bandits. Next, we consider a more practical scenario where the number of samples that can be drawn from each distribution is different, and propose novel formulation of weighted GDRO, which allows us to derive distribution-dppendde convergence rates. Denote by n, the sample budget for the i-th distribution, and assume 1 2 n2 2 2 Rom In the first approach, we incorporate non-uniform sampling into SMD such that the sample budget is satisfied in expectation, and prove that the excess risk of the i-th distribution decreases at an O log m/ m/n,) rate. In the second approach, and also reduce the variance in stochastic gradients, and then leverage stochastic mirror- prox algorithm. which can exploit small variances, to optimize carefully designed weighted GDRO problem. Under appropriate conditions, attains an ((log m)/ /n:)) convergence rate, which almost matches the optimal O ((/1/n,) rate of only learning from the i-th distribution with n, samples.
use mini-batches meet the budget exactly
Introduction
In the classical statistical machine learning, our goal is to minimize the risk with respect to a fixed distribution Po [Vapnik 2000). i.c.,
min WEW
(Ro(w) ExIPo [e(w;
where Z is a sample drawn from Po. W denotes hypothesis class, and E(w; z) loss measuring the prediction error of model w on z. During the past decades, various algorithms have been developed to optimize ( and be grouped in two categories: sample average approximation (SAA) and stochastic approximation (SA) [Kushner and Yin 2003)- In SAA, we minimize an empirical risk defined as the average loss over set of samples drawn from Po. and in SA, we directly solve the original problem by using stochastic observations of the objective Ro(.).
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

However, a model that trained over a single distribution may lack robustness in the sense that (i) could suffer high error on minority subpopulations, though the average loss is small; (ii) its performance could degenerate dramatically when tested on a different distribution. Distributionally robust optimization (DRO) provides a principled way to address those limitations by minimizing the worst-case risk in a neighhorhood of Po [Ben- Tal et al. 2013). Recently, it has attracted great interest in optimization [Shapiro 2017). statistics [Duchi and Namkoong 2021], operations research (Duchi et al. 2021], and machine learning (Hu et al. 2018 Jin et al. 2021 Agarwal and Zhang 2022 In this paper, we consider an emerging class of DRO problems, named as Group DRO (GDRO) which optimizes the maximal risk over finite number of distributions [Oren et al. 2019 Sagawa et al. 2020 Mathematically, GDRO can be formulated as a minimax stochastic problem:
min max WEWiE [m))
{ {Ra(w) = Ex [[[w;))]}
where Pi...., P... denote m distributions. motivating example is federated learning, where a centralized model is deployed at multiple clients, each of which faces a (possibly) different data distribution [Mohri et al. 2019
Supposing that samples can be drawn from all distributions freely, we develop efficient SA approaches for (2), in favor of their light computations over SAA methods. As elaborated by Nemirovski et al. 2009 $3.2], we can cast 2 as stochastic convex-coneat saddle-poid aodle--oint problem:
m
min max WEW q
o(w, q)
Σ 9: Re (w)
(3)
where A,, = Rm Σ 9: = E. by their mirror descent stochastic approximation method, namely stochastic mirror descent (SMD). In
is the (m- -)dddminsiioonnl simplex, and then solve
3
fact, several recent studies have adopted this (or similar) strategy to optimize (3). But, unfortunately,
we found that existing results are unsatisfactory because they either deliver a loose sample complexity
[Sagawa et al. 2020), suffer subtle dependency issues in their analysis (Haghtalab al. 2022 Soma
et al. 2022]. or hold only in expectation [Carmon and Hausler 2022].
As starting point, we first provide routine application of SMD to 33. and discuss the thcoretical guarantees. In each iteration, we draw 1 sample from every distribution to construct unbiased estimators of (-) and its gradient, and then update both w and by SMD. The proposed method achieves an O( (1v /log m) convergence rate in expectation and with high probability, where T is
the number of iterations. As result, we obtain (m m)/e²) sample complexity for finding an c-optimal solution of (3). which matches the 2(m/ lower bound (Soma et al. 2022 Theorem 5]
up to logarithmic factor, and tighter than the (m² (log m)/c3 bound of Saaawa et al. (2020) by an m factor. While being straightforward, this result seems new for GDRO.
Then, we proceed to reduce the number of samples used in each iteration from to 1. We remark that a naive uniform sampling over m distributions does not work well, and yields a worse sample complexity (Sagawa et al. 2020). As an alternative. we borrow techniques from online learning with stochastic observations, and explicitly tackle the non-obe con-obiivious nature of the online process, which distinguishes our method from that of Soma et al. [2022). Specifically, we use SMD to update w, and Exp3-IX, an algorithm for non-oblivious multi-armer bandits (MAB) [Neu] 2015), with stochastic rewards to update q. In this way, our algorithm only needs sample in each round and attains an O(/m(log m) convergence rate, implying the same (mm(lo m)/e²) sample complexity.
Next, we investigate a more practical and challenging scenario in which there are different budgets of samples that can be drawn from each distribution, a natural phenomenon encountered in learning with imbalanced data (Amodei et al., 2016). Let n4 be the sample budget of the i-th distribution, and without loss of generality, we assume that n4 2 R2 2 optimal sample complexity, but to reduce the risk on all distributions as much as possible, under the budget constraint. For GDRO with different budgets, we develop two SA approaches based on non-uniform sampling and mini-batches, respectively.
Rum Now, the goal is not to attain the
In each iteration of the first approach, we draw 1 sample from every Pa with probability n./n1. and then construct stochastic gradients to perform mirror descent. In this way, the budget will be satisfied in expectation after n1 rounds. To analyze its performance, we propose a novel formulation of weighted GDRO, which weights each risk (.) in (3) by a scale factor pt. Then, our algorithm can be regarded as SMD for an instance of weighted GDRO. With the help of scale factors,

---

demonstrate that the proposed algorithm enjoys distributiinn-eeenddee convergence in the sense that it converges faster for distributions with more samples. In particular, the excess risk on distribution P. reduces at an O(vn, log m/n.) rate, and for Pi, becomes O(v (Vo m) /n1)), which almost matches the optimal O( (T/mi) rate of learning from a single distribution with nn samples.
On the other hand, for distribution P, with budget n, < n1, the above O(ynj m/n:) rate is worse than the O( O(vI/nz) rate obtained by learning from P. alone. In shape contrast with this limitation, our second approach yields nearly optimal convergence rates for multiple distributions across a large range of budgets. To meet the budget constraint, it runs for n. S R'm rounds, and in each iteration, draws mini-batch of n samples from every distribution P.. As a result, (i) the budget constraint is satisfied exactly; (ii) for distributions with a larger budget, the associated risk function can be estimated more accurately, making the variance of the stochastic gradient smaller. To benefit from the small variance, we leverage stochastic mirror-prox algorithm (Juditsky et al. 2011]. instead of SMD, update our solutions, and again make use of the weighted GDRO formulation to obtain distribution wirrbbuttoowwiseeee convergence rates. Theoretical analysis shows that the excess risk converges an O(( 1 1 log m) rate for each P.. Thus, we obtain a nearly optimal O(log m)/ nn,)
Ppm /ms
rate for distributions P. with n, I n and an O(log m) /mm) rate otherwise. Note that the latter rate is as expected since the algorithm only updates (1)) times.
Related work We briefly discuss the related works on the GDRO problem in (2 and will review the traditional DRO problem in Appendix A Sagawa al. 2020) have applied SMD [Nemrrovski et al..1. 2009] to (3). but only obtain sub-optimal sample complexity. In the sequel, |Haghtalab et al. [2022] and Soma et al. [2022] have tried to improve the sample complexity by reusing samples and applying techniques from MAB respectively, but their analysis suffers dependency issues. Carmon and Hausler [2022 Proposition 2] successfully established an Omm(log m)/e²) sample complexity by combining SMD and gradient clipping, but their result holds only in expectation. To deal with heterogeneous noise in different distributions, Agarwal and Zhang [022) propose a variant of GDRO named as minimax regret optimization (MRO), which replaces the risk with "excess risk".
2 SA Approaches to GDRO
In this section, we provide two efficient SA approaches for GDRO, which are equipped with the same sample complexity, but with different number of samples used in each round (m versus 1).
2.1 Preliminaries
First, we state the general setup of mirror descent Nemirovski et al., 2009 We equip the domain W with a -enneratgggnerating function (.)(.), which is I-strongly convex with respect to certain norm Ilw We define the Bregman distance associated with (-) as (w, v) = 1/o (u) [M (v)
(VDn(v), v)]. For the simplex Am, we choose the entropy function (q) = Σm qi in q.
which is I-strongly convex with respect to the vector f--norm ||. I| as the distance-generatingg function. Similarly, B,((,) is the Bregman distance associated with (.))).
Then, we introduce the standard assumptions about the domain, and the loss function.
Assumption The domain W is convex and its diameter measured by v.)(.) is bounded by D, i.e.,
max (w) wEW
min (w) < D
WEW
For Am. it is easy to verify that its diameter measured by the entropy function is bounded by Vln..
Assumption For all i E [m], the risk function Rs(w) = Ez-P, [[( is convex.
Assumption 3 For all i E [m]) we have
0 < ((w;z) S 1, Vw E W, 2 2~ ~ P..
(5)
Assumption 4 For all i E [m], we have
VVE(;; 2) I,,, S G, Vw E W, z~
(6)
where || -ll,, is the dual norm of llw.
3

---

Algorithm Stochastic Mirror Descent for GDRO
Input: Two step sizes: The and Tl
1: Initialize W1 argmin Un (w), and qi [1/m,.... 1/m] T E I
2: for f. = 1 to T do
3: For each i E [m]] draw a sample Z; from distribution P.
4: Construct the stochastic gradients defined in (
5: Update W, and qr according to (11) and 12 respectively
6: end for
7:
return w - We and q = El= q
Last, we discuss the performance measure. To analyze the convergence property, we measure the quality an approximate solution (w, ② to 3 by the error
Es <w, q) max p(w.q) q)
min <(w, (w q)
(7)
qe
weW
which directly controls the optimality of to the original problem 2
since
max (w) min max (w) < max
min WEW
(8)
km(m]
WEW im(m]
2.2 Stochastic Mirror Descent for GDRO
To apply SMD, the key is to construct stochastic gradients of the function @(w, q) in 3 In each (6) round t, denote by w, and qe the current solutions. We draw one sample z; from every distribution P., and define stochastic gradients as
m
Σ tr.s Ve(w,, and g (w,, = (w;;;; z, ("),... e(w;; (T))))
ga We, q:)
(9)
It is worth mentioning that the construction of gu (We, q:) can be further simplified 5w. qr) = vl(w,
(10)
where i, [m] drawn randomly according to the probability q.-
Then, we utilize SMD to update w, and q::
Wx+1 argmin {The (g (gw (w,, q:), we) Bw w.)]
(11)
wEW
9++1 argmin {Na( ( (w,. qi). qt) + (g,
(12)
qEAm
where 7u > 0 and f7 > are two step sizes that will be determined later. The updating rule w depends on the choice of the distance- distance-generann generating function (-). Since By q:) is defined in terms
of the negative entropy, (12 equivalent to
exp (nrl(wizz,)))) (2))
4.+111
Vi [m]]
(13)
Σラ= 9t exp (nge 2())
which is the Hedge algorithm Freund and Schapire 1997 applied to maximization problem. In the beginning, we set W1 = argmin wew Vio (w), and qı I.m where 1, is the m-dimennional vector consisting of 1's. In the last step, return the averaged iterates
ΣI W, and
q= 7 Σ final solutions. The completed procedure is summarized in Algorithm
Based on the theoretical guarantee of SMD Nemirovski et al. 2009 $3.1], have the following theorem.
Theorem Under Assumptions I
3 and and setting The D ST and flg
(Inm) V GT(D" Hlmm)
Algorithm
with probability at least we have
E6 (w, q) < + 21 )
10(D2G2 Inm)
T

---

Remark Theorem shows that Algorithm achieves an Ol (Voo m) (T) convergence rate. Since it consumes m samples per iteration, the sample complexity is (mm m)/e which nearly matches the (mm/=2) lower bound [Soma et 2022 Theorem 5]. Due space limitations, defer all the proofs to Appendix B and omit expectation bounds.
Comparisons with Sagawa et al. 2020] In their stochastic algorithm, Sagawa et 2020 generate a random index & E [m uniformly in each round f, and draw sample z; from P... The stochastic gradients are constructed follows:
gw Wt. q1) = 9t,s, m V (()),, and gg (We, qt) = [0,.............. ml mel(wii
OJT
(14)
where (w., qt) ) is a vector with me(w,;; z, ) in position i, and 0 elsewhere. Then, the two stochastic gradients
used update w, and qr, in the same way as [1] and (12). However, only attains slow convergence rate: O(m / (log m//T") leading an O(m (log (loo)) (log sample
complexity, which is higher than that of Algorithm| by factor of m. The slow convergence is due the fact that the optimization error depends on the dual norm of the stochastic gradients in (I4). which blows up by a factor of m, compared with the gradients in
Comparisons with Haghtalab et al. [2022] To reduce the number samples required in each round, Haghtalab et al. [2022] propose to reuse samples for multiple iterations. To approximate V Wt qr), they construct the stochastic gradient gw gw (w,, qt) in the same way (IOp. which needs sample. To approximate V V wr, 9t), they draw samples z; ..... Zr one from each distribution, at round T mk + 1, h E Z. and reuse them for m iterations to construct
(m)
(,,,,,, q:) [e(we; 2 ) e( ((w;;("")) (1))
T,
(15)
Then, they treat gw q:) and g (We, qe) as stochastic gradients, and update w, and qe by SMD. In this way, their algorithm uses samples on average in each iteration. However, the gradient in (15
no longer an unbiased estimator of the true gradient g, (wr, q:) rounds = + 2....... T + m
making their analysis ungrounded. To see this, from the updating rule of SMD, we know that W--2 depends on 4r+1, which in turn depends on the m samples drawn at round т, and thus
E [( 22 A R,(w i I.....,,
Remark As shown (10). we can use sample to construct a stochastic gradient for w, with (wz,
small norm, since (w:, 9:) the error related to W,. However, we do not have such guarantees for the stochastic gradient of q Recall that the infinity norm of gy (We, q:) in (14) is upper bounded by m. The reason is that insist on the unbiasedness of the stochastic gradient, which leads to a large variance. the next section, we borrow techniques from online learning to better balance the bias and the variance.
Im,,, S G under Assumption 4 Thus, it is relatively easy to control
2.3 Non-oblivious Online Learning for GDRO
the studies convex-concave anddle-point problems, it is now well-known that they can be solved by playing two online learning algorithms against each other [Freund and Schapire 1999 Rakhlin and Sridharan| 2013] Syrgkanis et al.) 2015]. With this transformation, we can exploit no-regret algorithms developed in online learning bound the optimization error. To solve problem 3 ask the 1st player to minimize a sequence of convex functions
o(w, q1)
Σ (ww), o(w, 92)
Σ (w)),, ".. ((w,9r) Σ QT, (w)
subject to the constraint W, and the 2nd player to maximize sequence of linear functions
Φ( W q) = Σ (wi)). of 6(w2.9)
Σ (w2). o(wr. 9) = Σ ((())
subject the constraint q An. highlight that there exists an important difference between our stochastic convex-concave problem and its deterministic counterpart. Here, the two players cannot
5

---

Algorithm 2 Non-oblivious Online Learning for GDRO
Input: Two step sizes: The and fly IX coefficient
1: Initialize W1 = argmin Vn (w), and qi = (1/m,..... 1/m] T E IRm
2: for = 1 to T do
2: Generate i E [m]] according to qr. and then draw a sample z (fe) from q,
3: Construct the stochastic gradient (TO) and the IX loss estimator in (17)
4: Update W, and q, according to (66 and (18) respectively
5: end for
6: return W - E We and q = 1==1 q
directly observe the loss function, and can only approximate (,(w) = Ez- [<(w;)) 22] by drawing samples from P,. The stochastic setting makes the problem more challenging, and in particular, we need to take care of the non-oblivions nature of the learning process. Here, "non-oblivious" refers to the fact that the online functions depend on the past decisions of the learner.
Next, we discuss the online algorithms that will be used by the two players. As shown in Section 2.2 the 1st player easily obtain stochastic gradient with small norm by using 1 sample. So, model the problem faced by the 1st player as "non-oblivious online convex optimization (OCO) with stochastic gradients", and still SMD update its solution. In each round i, with 1 sample drawn from P., the 2nd player can estimate the value of (w,,)) which is the coefficient of qr. Since the 2nd player is maximizing linear function over the simplex, the problem can be modeled "non- oblivious multi-armed bandits (MAB) with stochastic rewards". And fortunately, we have powerful online algorithms for non-oblivious MAB [Auer al.) 2002 Lattimore and Szepesvári 2020]. whose regret has a sublinear dependence on m. In this paper, we choose the Exp3-IX algorithm Neu 2015). and generalize theoretical guarantee to stochastic rewards.
The complete procedure is presented Algorithm 2 and we explain key steps below. In each round
t, we generate index & E [m] from the probability distribution qr, and then draw a sample z; from the distribution q. With the stochastic gradient in (10), we use SMD to update wy:
W++1 argmin { yw (gw W., q:), w we) + Bw we)
(16)
wEW
Then, reuse the sample 2((x) to update q, according to Exp3-IX, which first constructs the Implicit-exploration (IX) loss estimator Kocák al. 2014]:
Ili
i), Vi E [m],
(17)
4e,i +
where y > 0 is the IX coefficient and I[AA] equals to 1 when the event A is true and 0 otherwise, and then performs a mirror descent update:
9++1 argmin (Tg (5., q - q.) + Bg q:)).
(18)
q
We present the theoretical guarantee of Algorithm
Theorem
Under Assumptions 2
3
and
4 and setting Yhr GvOT 11
and 7 = =7g in
mf
Algorithm 2
with probability at least
8. we have
E (W,W)))
I
2 In
m In
I 2T
6 (19)
<DGV 225 8
3
In
2T
The above theorem shows that with sample per iteration, Algorithm is able to achieve an Olv /mogg m) convergence rate, thus maintaining the O(m(lo sample complexity.
Remark 3
Comparisons with Soma et al. [002 In a recent work, Soma et al. 2022 have deployed online algorithms optimize and q. but did not consider the non-oblivious property. As a result, their

---

theoretical guarantees, which build upon the analysis for oblivious online learning [Orabona 2019]. cannot justify the optimality of their algorithm for (3). Specifically, their results imply that for any fixed w and that are independent from w and q [Soma et al.. 2022 Theorem 3],
m T
E [o( W. q) - Φ( q) =
(20)
cannot be used to bound eo (w q) in (7). because of the dependency issue. To be more clear, we have <0(w, q) mm o( W. q) millwew o(w, q) = (w, q) o(W, q), where w- argmin wEw Φ( w. q) and q argmax qEA.. (w, q) depend on w and q.
However, 20
After we pointed out the issue of reusing samples, HHaghtalb al. 2023 modified their method by incorporating bandits algorithms to optimize q. From our understanding, the idea of applying bandits to GDRO is firstly proposed by [Soma al. [2022]. and subsequently refined by us.
Remark 4
3 Weighted GDRO and SA Approaches
When designing SA approaches for GDRO, it is common to assume that the algorithms are free to draw samples from every distribution (Sagawa et al., 20200. as we do in Section 2 However, this assumption may not hold in practice. For example, data collection costs can vary widely among distributions [Radivojac et al. 2004), and data collected from various channels can have different throughputs [Zhou] 2023). In this section, we investigate the scenario where the number of samples can be drawn from each distribution could be different. Denote by n, the number of samples that can be drawn from P.. Without loss of generality, we assume that nı n2 2 2 7mm Note that we have a straightforward Baseline which just runs Algorithm for nmm iterations, and the optimization error Eo w. q) = O (log n)mmm).
3.1 Stochastic Mirror Descent with Non-uniform Sampling
To meet the budget, we propose to incorporate non-uniform sampling into SMD. Specifically, in
round we first generate set of Bernoulli random variables (b(2) :-.., ((m") with Pr(() (0) 1] p
to determine whether to sample from each distribution. If 63 = 1, we draw sample zi from P,. Now, the question is how construct stochastic gradients from those samples. Let C. {d)b()) 1}
be the indices of selected distributions. we stick to the original problem in (3), then the stochastic gradients should be constructed in the following way
4t,i VE(W:; (6)). and [g q:)]; (:), = { E( zi (3) )/p,. Ps
Su (W:, 9:)
Σ IEC.
i C. otherwise
(21)
0,
to ensure unbiasedness. Then, they can be used by SMD to update w, and q. After iterations, the expected number of samples drawn from P. will be nıp; n4, and thus the budget is satisfied in expectation. To analyze the optimization eтror, we need to bound the norm of stochastic gradients in
(1... To this end, we have [|g.. (Wr q:) mm.. S Gn, n'm and ||g, (w,, I.)xllso I n1/nm- Following the arguments of Theorem we can prove that the error co Co(w, (, 9) = ((V(log m)/n1 n1 n1/nm) nmm) =
O(ynn m/gm/ which is even larger than the O(v mmmmmm.) error of the Baseline.
In the following, we demonstrate that a simple twist of the above procedure can still yield meaningful results that are complementary to the Baseline. We observe that the large norm of the stochastic gradients in 221 is caused by the inverse probability 1/p,. A natural idea is to ignore 1/p, and define the following stochastic gradients:
Ew gw (Wr, qi) = Σ qr,s Vll ("), and [g, (We. qt)l
= =
e(wri (2)), i C. otherwise.
(22)
O,
sEC.
In this way, they are no longer stochastic gradients of (3), b but can be treated as stochastic gradients of a weighted GDRO problem:
m
min max ((w,q) = Σ gip Ri(w)
(23)
wEW qm
7

---

Algorithm 3 Stochastic Mirror Descent for Weighted GDRO
Input: Two step sizes: The and flu
1: Initialize W1 = argmin Un (w),, and qi = [1/m,..... 1/m] T E I
2: for = 1 to n1 do
3: For each i E [m], generate a Bernoulli random variable 6(0) with Pr((,) 1] = Ps. and
b(() 1, draw a sample (*) from distribution P;
4: Construct the stochastic gradients defined in (22)
5: Update w, and qr according to [11 and 12 respectively
6: end for
7: return w TI 21=1 and q 1
n1 1 q
where each risk (.(-) is scaled by factor p.. Based on the gradients in (22). we still use |1] and (12) to update We and qr- We summarize the complete procedure in Algorithm
We omit the optimization eπror of Algorithm 3 for (23 since it has exactly the same form as Theorem | What we are really interested in is the theoretical guarantee of its solution on multiple distributions. To this end, we have the following theorem.
Theorem 3 Under Assumptions / 3 and and setting The D2
5m D m) and fl
(Inm) V (D in Algorithm with probability at least have
T1 P R((w)-
u())
G + Inm) m1()
V 1I0(DD22 + Imm)n
Vi [m]
Pi
where Pig is the optimal value of 23 and v(0) 2
Remark 5 We see that Algorithm exhibits distribution dependde convergence behavior: The larger the number of samples n4, the smaller the target risk T1 /nn, and the faster the convergence rate Ol nı log m/n.). Note that its rate is always better than the O(ynz log m/ /mmmm) rate of SMD) with 21 as gradients. Furthermore, it converges faster than the Baseline when n; 2 /ntn 7mm In particular, for distribution P1, Algorithm attains an D(V(log m)/n1) rate, which almost matches the optimal 0(/1/n) rate of learning from a single distribution. Finally, we would like to emphasize that a similar idea of introducing "scale factors" has been used by Juditsky et al. [2011] $4.3.1] for stochastic semidefinite feasibility problems and Agarwal and Zhang 2022 for empirical MRO.
3.2 Stochastic Mirror-Prox Algorithm with Mini-batchess
In Algorithm 3, distributions with more samples take their advantage by appearing more frequently in the stochastic gradients. In this section, we propose different way, which lets them reduce the variance in the elements of stochastic gradients by mini-batches [Roux et al. 2008).
The basic idea is as follows. We run our algorithm for small number of iterations ñ that is no larger than Rmm Then, in each iteration, we draw a mini-batch of n: samples from every distribution P, For P, with more samples, we can estimate the associated risk (-(-)) and its gradient more accurately, ie., with a smaller variance. However, to make this idea work, we need to tackle two obstacles: (i) the performance of the SA algorithm should depend on the variance of gradients instead of the norm, and for this reason SMD is unsuitable; (ii) even some elements of the stochastic gradient have small variances, the entire gradient may still have a large variance. To address the first challenge, resort to a more advanced SA approach- stochastic mirror-prox algorithm (SMPA), whose convergence rate depends the variance [Juditsky et al.) 2011). To overcome the second challenge, again introduce scale factors into the optimization problem and the stochastic gradients.
In SMPA, need to maintain two sets of solutions: (w., q:) and (w', q'). In each round t, we first draw n;/ / nm samples from every distribution P;, denoted by x! (t.5) j 1,.... them to construct stochastic gradients at (w', qf)) of weighted GDRO problem 23 where the
......,.. n4/ m///mm Mm.. Then, we

---

Algorithm 4 Stochastic Mirror-Prox Algorithm for Weighted GDRO
Input: Two step sizes: The and f/u
1: Initialize w argmin Un (w), and qí =
1/m]] E IRm
2: for = 1 to m2m do
3: For each E [m.] draw ni/nmm samples {zi ((.5) jj 1,... n4 ni/nm from distribution P;
4: Construct the stochastic gradients defined in 24
5: Calculate Wi+1 and qu+1 according to (25 and 26 respectively
6: For each E [m]] draw ni/To, samples {2
(())
...., ni/nm) from distribution P,
7: Construct the stochastic gradients defined in 27
8: Calculate w/+1 and 44+1 according (88 and 29
respectively
9: end for
return nim =t=2 Wt and q nm 2t=2
1 12
q:
value p. will be determined later. Specifically, we define
gu W/y )))
m
Le
Tmm
fm
Velw'(z(( (101))
(24)
na///mm
8y (w',, 9.)
l(wliz( (1,1)) )...).......................)).. ..... Pm (w/s 2 (m))
p1
Let's take the stochastic gradient gy (w', q.). whose variance will be measured in terms of || llae. as example to explain the intuition of inserting pi. Define u,
Ens/"" L(w!; (f.f). With larger mini-batch size nz/nm- 02 will approximate (a(w)) more accurately, and thus have smaller variance. Then, it allows us to insert a larger value of pt. without affecting the variance of [gy (w: 1lo,. since ol is insensitive perturbations of small elements. Similar the case in Theorem the convergence rate of R.(.) depends on 1/p;, and becomes faster if ps is larger.
Based on (24 we use SMD to update (w', q.),, and denote the solution by (We+1s 9:41):
Wx+1 = argmin { {Th (g.. (w/s 9'), w w() + Bw w())
(25)
wEW
4++1 = argmin { 8y q.). q Φ) + (,(g,
(26)
qEam.
Next, we draw another nz/Um samples from each distribution Pi, denoted by 2((
1. n4/nm to construct stochastic gradients at (Wr+1, 9++1):
g (We11, g++1)
Σ
90+1.120
Tbm
(27)
/
gg (We11, qu+1)
=
p1 Bem
( (1:1)
Pmn Pml(We+11;2) 2, (m))
Then, them to update (w/, 9%) again, and denote the result by (w/ Wi Ib+1):
w W argmin TTw (gw (We+1, 94+1). - -w") + B (w, w()
(28)
wEW
4+1 argmin {na{ ga (w (0 1.+1), d'd B, q')
(29)
qea..
To meet the budget constraints,
we repeat the above process for nmm / iterations. Finally, we q, solutions. The completed procedure is
return Tm Lt=2 1+nm//2 W, and q 7m 2tm2
2. -1tm//2
summarized in Algorithm
To analysis the performance of Algorithm we further assume the risk function R;(-) is smooth, and the dual norm I| satisfics regularity condition.

---

Assumption All the risk functions are L-smooth, i.e.,
(VR((w) (R,(w') L || w'llw: Vw, w/ E W, e [m]]
(30)
Note that even the studies of stochastic convex optimization (SCO), smoothness is necessary to obtain a variance-based convergence rate [Lan 2012
Assumption The dual norm ||
w,, is k-regular for some small constant & 2 1.
The regularity condition is used when analyzing the effect of mini-batches on stochastic gradients. For formal definition, please refer to Duditsky and Nemirovski (2008). Assumption [ [ satisfied by most of papular norms considered in the literature, such as the vector f--norm and the infinity norm.
Then, have the following theorem for Algorithm
Theorem Define
P Pm
Pmax max Pix (E(m]
Wmax
max 16 [m(m] ni
(31)
L. = 2/ 2 DAGVIn and 0² = 20 (R G2 + +In m)
where e is an absolute constant. Under Assumptions 7 and and setting
(1
Tu 2D2 min
and fg = 2 min
Inm m
V3L
vi
v3L To
in Algorithm with probability at least
8, we have
)
1
1
7I.
0²
2 +
2
14 2 g
(.(w)
P
+
14
7v
3 log
log
Pi
Di
1mm
Tym
3
7mm
where Pip the optimal value of Furthermore, by setting pi
1/ 1//mm
1mm +
Pi
(32)
1/ Tmm
v Vnm/nt
with high probability, we have
(w( 1 P4 o
((1) 7um
Vk
In
TTu
Remark Compared with Algorithm 3 Algorithm has two advantages: (i) the budget constraint
is satisfied exactly; (ii) we obtain faster O (log m)/ /n.) rate for all distributions P; such that < n which is much better than the O(yni (mm/n:) rate of Algorithm 3 and the O VVlo m)/mmm)
rate of the Baseline. For distributions with larger number of budget, i.e., ni n it maintains fast (llog m))nm, rate. Since it only updates n., times, and the best we can expect is the O(1 (()/)m) rate of deterministic settings [Nemirovski] 2004]. So, there is performance limit for mini-batch based methods, and after that increasing the batch-size cannot reduce the rate, which consists with the usage of mini-batches in SCO [Cotter al. 2011 Zhang al. 2013).
4 Conclusion
For the GDRO problem, we develop two SA approaches based on SMD and non-oblivious MAB. Both of them attain the nearly optimal (mmiog m) sample complexity, but with different number of samples used in each round, which are m and 1 respectively. Then, we formulate weighted GDRO problem handle the scenario in which different distributions have different sample budgets. We first incorporate non-uni non-uniform sampling into SMD to satisfy the sample budget in expectation. and deliver distribution-dependeee
convergence rates. Then, we propose to use mini-batches to meet budget exactly, deploy SMPA to exploit the small variances, and establish nearly optimal rates for multiple distributions. We have conducted experiments to evaluate our proposed algorithms. The empirical results presented in Appendix C and align closely with our theories.
10

---

Acknowledgments and Disclosure of Funding
This work was partially supported by the National Key R&D Program of China (20227D01 14801), NSFC (62122037, 619210066 National Postdoctoral Program for Innovative Talent, China Postdoc- toral Science Foundation (2023M731597), and the major key project of PCL (PCL2021A12).
References
A. Agarwal and T. Zhang. Minimax regret optimization for robust machine learning under distribution shift. In Proceedings of 35th Conference Learning Theory, pages 2704 2729, 2022.
D. Amodei, Ananthanarayanan, R. Anubhai, and et al. Deep speech 2 End-to-end speech recognition in english and mandarin. In Proceedings of the 33rd International Conference Machine Learning, pages 173,182,20 2016.
P. Auer, N. Cesa-Bianchi, Y. Freund, and R. E. Schapire. The nonstochastic multiarmed bandit problem. SIAM Journal on Computing. 32(1):48-77, 2002.
B. Becker and R. Kohavi. Adult.
UCI Machine Learning Repository, 1996.
DOI:
A. Ben- Tal, L. E. Ghaoui, and A. Nemirovski. Robust Optimization. Princeton University Press, 2009.
A. Ben-Tal, D. den Hertog, A. D. Waegenaere, B. Melenberg, and G. Rennen. Robust solutions of optimization problems affected by uncertain probabilities. Management Science, 59(2):341-357, 2013.
A. Ben-Tal, E. Hazan, T. Koren, and S. Mannor. Oracle-based robust optimization via online learning. Operations Research, 63(3):628-638, 2015.
D. Bertsimas, V. Gupta, and N. Kallus. Robust sample average approximation. Mathematical Programming, 171:217-282, 2018.
A. Blum, N. Haghtalab. A. D. Procaccia, and M. Qiao. Collaborative PAC learning. In Advances in Neural Information Processing Systems 30, pages 2389-2398, 2017.
S. Bubeck and N. Cesa-Bianchi. Regret analysis of stochastic and nonstochastic multi multi-armed bandit problems. Foundations and Trends in Machine Learning, 5(1):1-122, 2012.
Y. Carmon and D. Hausler. Distributionally robust optimization via ball oracle acceleration. In Advances in Neural Information Processing Systems 35, pages 35866-35879, 2022.
N. Cesa-Bianchi and G. Lugosi. Prediction, Learning, and Games. Cambridge University Press, 2006.
A. Cotter, O. Shamir, N. Srebro, and K. Sridharan. Better mini-batch algorithms via accelerated gradient methods. In Advances in Neural Information Processing Systems 24, pages 1647-1655, 2011.
E. Delage and Y. Ye. Distributionally robust optimization under moment uncertainty with application to data-driven problems. Operations Research, 58(3):595-611 -612, 2010.
J. C. Duchi and H. Namkoong. Learning models with uniform performance via distributionally robust optimization. The Annals of Statistics, 49(3):1378 - 1406, 2021.
J. C. Duchi, P. W. Glynn, and H. Namkoong. Statistics of robust optimization: A generalized empirical likelihood approach. Mathematics of Operations Research, 46(3):946-96 969, 2021.
P. M. Esfahani and D. Kuhn. Data-driven distributionally robust optimization using the Wasserstein metric: performance guarantees and tractable reformulations. Mathematical Programming, 171: 115-166, 2018.
11

---

A. D. Flaxman, A. T. Kalai, and H. B. McMahan. Online convex optimization in the bandit setting: Gradient descent without gradient. In Proceedings of the 16th Annual ACM-SIAM Symposium on Discrete Algorithms, pages 385-394, 2005.
Y. Freund and E. Schapire. decision-theoretic generalization of on-line learning and an application to boosting. Journal of Computer and System Sciences, 55(1): (011):119-1333 19-139, 1997.
Y. Freund and R. Schapire. Adaptive game playing using multiplicative weights. Games and Economic Behavior, 29 1):79-103, 1999.
N. Haghtalab, M. I. Jordan, and E. Zhao. On-demand sampling: Learning optimally from multiple distributions. In Advances in Neural Information Processing Systems 35, pages 406-419, 2022.
N. Haghtalab, M. I. Jordan, and E. Zhao. On-demand sampling: Learning optimally from multiple distributions. ArXiv e-prints, HiXiv:2100.1259922 2023.
T. Hashimoto, M. Srivastava, H. Namkoong, and P. Liang. Fairness without demographics in repeated loss minimization. In Proceedings of the 35th International Conference on Machine Learning. pages 1929-1938, 2018.
W. Hu, G. Niu, Sato, and M. Sugiyama. Does distributionally robust supervised learning give robust classifier?? In Proceedings of the 35th International Conference on Machine Learning, volume 80, pages 2029-2037, 2018
J. Jin, B. Zhang, H. Wang, and L. Wang. Non-convex distributionally robust optimization: Non- asymptotic analysis. In Advances in Neural Information Processing Systems 34, pages 2771-2782, 2021.
A. Juditsky, A. Nemirovski, and C. Tauvel. Solving variational inequalities with stochastic mirror-pre prox algorithm. Stochastic Systems, 1(1):17-58, 2011.
A. B. Juditsky and A. S. Nemirovski. Large deviations of vector-valued martingales in 2-smooth normed spaces, ArXiv e-prints, arXiv:0809.0m 08113, 2008
T. Kocák, G. Neu, M. Valko, and R. Munos. Efficient learning by implicit exploration in bandit problems with side observations. Advances in Neural Information Processing Systems 27, pages 613-621, 2014.
D. Kuhn, P. M. Esfahani, V. A. Nguyen, and S. Shafieezadeh-Abude Wasserstein distributionally robust optimization: Theory and applications in machine learning. Operations Research & Management Science in the Age of Analytics, pages 130-166, 2019.
H. J. Kushner and G. G. Yin. Stochastic Appreximation and Recursive Algorithms and Applications. Springer, second edition, 2003.
G. Lan. An optimal method for stochastic composite optimization. Mathematical Programming. 133: 365-397, 2012.
T. Lattimore and C. Szepesvári. Bandit Algorithmss. Cambridge University Press, 2020.
D. Levy, Y. Carmon, J. C. Duchi, and A. Sidford. Large-scale methods for distributionally robust optimization. In Advances Neural Information Processing Systems 33, pages 8847-8860, 2020.
M. Mohri, G. Sivek, and A. T. Suresh. Agnostic federated learning. In Proceedings of the 36th International Conference on Machine Learning, pages 4615 4625, 2019.
H. Namkoong and J. C. Duchi. Stochastic gradient methods for distributionally robust optimization with fidivergences. In Advances Neural Information Processing Systems 29, pages 2216-2224 2016.
H. Namkoong and Duchi. Variance-based regularization with convex objectives. In Advances in Neural Information Processing Systems 30, pages 2971 2971-2980, 2017.
12

---

A. Nemirovski. Prox-method with rate of convergence O(1/t) for variational inequalities with lipschitz continuous monotone operators and smooth convex-coneave saddle point problems. SIAM Journal on Optimization, 15(1):229-255) 2004.
A. Nemirovski, A. Juditsky, G. Lan, and A. Shapiro. Robust stochastic approximation approach to stochastic programming. SIAM Journal on Optimization, 19(4):1574 (114):1574-100 1609, 2009.
G. Neu. Explore no more: Improved high-probability regret bounds for non-stochastic bandits. In Advances in Neural Information Processing Systems 28, pages 3168-311 3176, 2015.
H. L. Nguyen and L.. Zakynthinou. Improved algorithms for collaborative PAC learning. In Advances in Neural Information Processing Systems 31, pages 7642-7650, 2018.
Orabona. A modern introduction to online learning. ArXiv e-prints, arXiv:1912.113. arXiv:1912.152133v5) 2019.
Y. Oren, S. Sagawa, T. B. Hashimoto, and P. Liang. Distributionally robust language modeling. In Proceedings of the 2019 Conference Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing, pages 4227-4237, 2019.
Q. Qi. Z. Guo, Y. Xu, R. Jin, and T. Yang. An online method for a class of distributionally robust optimization with non-convex objectives. In Advances Neural Information Processing Systems 34, pages 10067-10080, 2021.
P. Radivojac, N. V. Chawla, A. K. Dunker, and Z. Obradovic. Classification and knowledge discovery in protein databases. Journal of Biomedical Informatics, 37(4))224 2004.
H. Rafique, M. Liu, Q. Lin, and T. Yang. Weakly-convee-concae min-max optimization: Provable algorithms and applications in machine learning. Optimization Methods and Software, 37(3): 1087-1121, 2022.
S. Rakhlin and K. Sridharan. Optimization, learning, and games with predictable sequences. In Advances in Neural Information Processing Systems 26, pages 3066-3074, 2013.
G. N. Rothblum and G. Yona. Multi-group agnostic PAC learnability. In Proceedings of the 38th International Conference Machine Learning, pages 9107-9115, 2021.
N. L. Roux, P.-A. Manzagol, and Y. Bengio. Topmoumoute online natural gradient algorithm. In Advances in Neural Information Processing Systems 20, pages 849-856 2008.
Sagawa, P. W. Koh, T. B. Hashimoto, and Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. In International Conference on Learning Representations, 2020.
D. Samuel and G. Chechik. Distributional robustness loss for long-tail learning. In Proceedings of the 2021 IEEE/CVF International Conference Computer Vision, pages 9475-9484, 2021.
H. Scarf. A min-max solution of inventory problem. Suddies in the Mathematical Theory of Inventory and Production, pages 201-209, 1958.
A. Shapiro. Distributionally robust stochastic programming. SIAM Journal on Optimization, 27(4): 2258-2275, 2017.
A. Sinha, H. Namkoong, and J. Duchi. Certifying some distributional robustness with principled adversarial training. In International Conference on Learning Representations, 2018.
A. Stowik and L. Bottou. On distributionally robust optimization and data rebalancing. In Proceedings of the 25th International Conference Artificial Intelligence and Statistics, pages 1283-1297, 2022.
T. Soma, K. Gatmiry, and S. Jegelka. Optimal algorithms for group distributionally robust optimization and beyond. ArXiv e-prints, arXiv:2212.1366) 2022.
V. Syrgkanis, A. Agarwal, H. Luo, and R. F. Schapire. Fast convergence of regularized learning in games. In Advances in Neural Information Processing Systems 28, pages 2989-2997, 2015.
13

---

V. N. Vapnik. The Nature of Statistical Learning Theory. Springer, second edition, 2000.
R. Vershynin. Ligh--Diensiooal Probability; An Introduction with Applications in Data Science. Cambridge University Press, 2018.
Wang, R. Gao, and Y. Xie. Sinkhorn distributionally robust optimization. ArXiv e-prints, arXiv:2109.1 1926, 2021.
Z. Xu, C. Dan, J. Khim, and P. Ravikumar. Class-weighted classification: Trade-off's and robust approaches. In Proceedings of the 37th International Conference on Machine Learning, pages 10544-10554, 2020.
Zhang, A. K. Menon, Veit, S. Bhojanapalli, S. Kumar, and S. Sra. Coping with label shift via distributionally robust optimisation. In International Conference on Learning Representations, 2021.
Zhang, Yang, Jin, and He. O(log projections for stochastic optimization smooth and strongly convex functions. In Proceedings of the 30th International Conference Machine Learning, pages 1121-1129, 2013.
Z.-H. Zhou. theoretical perspective of machine learning with computational resource concerns. ArXiv e-prints, arXiv:2305.02217,
2023. 2023.
14

---

A Related Work
Distributionally robust optimization (DRO) stems from the pioneering work of Scart
[1558] and has
gained lot of interest with the advancement of robust optimization (Ben-Tall et al. 2009
2015 It
has been successfully applied to a variety of machine learning tasks, including adversarial training
(Sinha et al. 2018 algorithmic fairness [Hashimoto et al., 2018). class imbalance [Xu et al. 2020 long-tail learning Samuel and Chechik [021], label shift (Zhang et al. [021]. etc.
In general, DRO is formulated to reflect our uncertainty about the target distribution. To ensure good performance under distribution perturbations, it minimizes the risk w.r.t. the worst distribution in an uncertainty set, i.e.,
(Ez-p [[(w;))] WEW PES(Ps)
min
sup
(33)
where S(Po) denotes a set of probability distributions around Po- In the literature, there mainly exist three ways to construct (())): (i) enforcing moment constraints [Delage and Ye 2010), (ii) defining neighborhood around Pn by distance function such as the f-divergence [Ben- Tal et al. 2013].
the Wasserstein distance [Kuhn et al. 2019), and the Sinkhorn distance [Wang et al. 202 and (ii)
hypothesis testing of goodness-of-fi [Bertsimas et al. 2018).
By drawing a set of samples from Po. we can also define an empirical DRO problem, which can be regarded as an SAA approach for solving (33). When the uncertainty set is defined in terms of the Cressie-Reat family of f-divergences, |Duchi and Namkoong [2021] have studied finite sample and asymptotic properties of the empirical solution. Besides, it has been proved that empirical DRO can
also benefit the risk minimization problem in ( Namkoong and Duchi [0077] show that empirical
DRO with the x" -ivveeeree has the effect of variance regularization, leading to better generalization w.r.t. distribution Po- Later, Ducch et al. [202T] demonstrate similar behaviors for the f-divergence constrained neighborhood, and provide one- and two-sided confidence intervals for the minimal risk in (T. Based on the Wasserstein distance, |Esfahani and Kuhn (2018) establish an upper confidence bound on the risk of the empirical solution.
Since (33 more complex than ( considerable research efforts were devoted to develop efficient algorithms for DRO and its empirical version. For Po with finite support, | Tal et al. 20013 Corollary 3] have demonstrated that (33) with f-divergences equivalent to convex optimization
problem, provided that the loss { 2) is convex in w. Actually, this conclusion is true even when Po is continuous [Shapiro] 2017 $3.2]. Under mild assumptions, [Esfahani and Kuhn 2018 show that DRO problems over Wasserstein balls can be reformulated as finite convex programs- even as linear programs. Besides the constrained formulation in (33), there also exists a penalized (or regularized) form of DRO [Sinha et al. 2018), which makes the optimization problem more tractable. In the past years, a series of SA methods have been proposed for empirical DRO with convex losses [Nakkoong and Duchi 2016), , and DRO with convex loss [Levy et al. 2020 and non-convex losses Qi et al. 2021 Rafique et al. 2022).
some cases
iJin et al. 2021
The main focus of this paper is the GDRO problem in 333 instead of the traditional DRO in (33) Saaawa et al.) [2020] have applied SMD (Nemirovski et al.) 2009] to (3). but only obtain a sub-optimal sample complexity of ( (mog (logm)/ m)/ because of the large variance in their gradients.
In the sequel. Haghtalab et al. (2022) and [Soma et al. (2022) have tried to improve the sample complexity by reusing samples and applying techniques from MAB respectively, but their analysis suffers dependency issues. Carmon and Hausler [2022, Proposition 2] successfully established an Omm(log m²)=²) sample complexity by combining SMD and gradient clipping, but their result holds only in expectation. To deal with heterogeneous noise in different distributions, Agaawal and Zhang [002 propose a variant of GDRO named minimax regret optimization (MRO), which replaces the risk (,(w) with "excess risk" R((w) min (w). More generally, we can introduce calibration terms in DRO prevent any single distribution to dominate the maximum (llowik and
Bottou 2022].
Finally, we note that GDRO has similar spirit with collaborative PAC learning Blum et al. 2017 Nguyen and Zakynthinou 2018 Rothblum and Yona 202T in the sense that both aim to find single model that performs well on multiple distributions.
15

---

B Analysis
In this section, we present proofs of main theorems.
B.1 Proof of Theorem
The proof is based on Lemma 3.1 and Proposition 3.2 of Nemirovski et al. 2009 To apply them, we show that their preconditions are satisfied under our assumptions.
Although two instances of SMD are invoked to update w and q separately, they can be merged as 1 instance by concatenating and and the distance-gent -enneraiing function [Nemirovski al. 2009 $3.1J. Let be the space that W lies in. We equip the Cartesian product & x Ie with the fol fowlowing norm and dual norm:
single variable [w; ql E W x Am. and redefine the norm
[w: elll
2D |wl
l1 li, and [u vell 2D | + v Inm. (34)
2nn m
We the notation
[w; q). and equip the W Am with the distance distance-generating function
u(x) = (((w;))) 2D "n (w) + 2 (g(g).
(35)
It is easy verify that vxx) is 1-strongly convex the norm || II- Let B(:,)) be the Bregman distance associated with v(-):
B(x, x = vxxx - [u(x') + (Vu(x), x']
2D² (Mw(w) - [w') + (w'), www')
(36)
2 2lmm (Vg(g) - -[v(9') (Vv((')) q')))
Bn, (w, w') + 2lmm B, B q')
2D
where x' = q')
Then, we consider the following version of SMD for updating x,:
X
{1<(5= (we, qt): gy (We, q.)], x1) + B(x, x1) xeWxam.
argmin
(37)
where 1 0 is the step size. In the beginning, we X argmin v(x) [wi;gi] From the decomposition of the Bregman distance in 36), we observe that (37 is equivalent to 1T and by setting
Dn and fl = 2ny Inm.
Next, we show that the stochastic gradients well-bounded. Under our assumptions, have
lg..( (w:: Im). Σ nr, Ve(w:;22
Ve(ws; (0))
Σ Gr.G = G,
[lgy q:) loc [|[[(w;; ( ... (Wr; (m)
As result, the concatenated gradients used in 37 is also bounded in term of the dual norm ||::
[g (We, q(): Be (we, q1]]ll = 2D g (Wr. 101. + (1lg,(w, 9:]1.. Inm
< /2D + 22m.
M
Now, we ready state our theoretical guarantees. By setting
2 5T(D2G2 + lmm)
11 MV51
16

---

Lemma 3,1 and (3.13) of Nemirovski et al. (2099]
imply that
5 ( q)] - 2M
10(D2 + Inm)
2
T
T
Furthermore, from Proposition 3.2 of Nemirovski et al.
20099) we have, for any Ω>
5
10(D2G2 + +mm)
Pr ( (8 M V (8 -22)
2eexp(-S))
T
We complete the proof by setting 8 = exp(
B.2
Proof of Theorem
We first bound the regret of the 1st player. In the analysis, address the non-obliviousness by the "ghost iterate" technique of Nemirovski et al. 2009
Theorem Under Assumptions
/ and by setting The GVST
2D
have
E
Φ Wr, qr)
min
o( qi)
< 2DGV5T
1=1
and with probability at least -5
T
T ((w, qd)
DGVT 2V5 + 225++8 88
1
We q1)
min wEW
By extending Exp3-IX to stochastic rewards, we have the following bound for the 2nd player.
by setting na = V Im and the IX coefficient mI
Theorem 6 Under Assumption
Yo 2"
we have
E max qeam
d W. q)
Σ
@(wes a:) < 33vmThmm + V
T +3
Inlm
mf
+
V
1
T
+1
1=1
and with probability at least
8. 8.
max g
Σ
T
o(wi: q)
Σ
T
T + o(w,, qi) s 33mmmmmmm + $
(
Vmmm
mT
+ V +1
T
In
3
2
From Jensen's 's inequality and the outputs w = E and q 1 Σ qr, we have
cØ(w, q) = max o q) min ó(w, q) qm
weW
T
T
max
o wr q)
min wEW
Σ
Øl (wW, q:)
qm
(38)
t=1
t=1
1
max qE m
((W, 9
Σ olw,, gi)
1 +
Σ (w,, 9c) - min
d w qr)
T
wew
We obtain (19 by substituting the high probability bound in Theorems the union bound.
and
into 38 and taking
B.3 Proof of Theorem 5
Our goal is to analyze SMD for non-oblivious OCO with stochastic gradients. In literature, we did not find a convenient reference for it. A very close one is the Lemma 3.2 of Flaxman et al. [2005), which bounds the expected regret SGD non-oblivious OCO. But is insufficient for our purpose, provide proof by following the analysis of SMD for stochastic convex-concave
17

---

optimization Nemirovski et al., 2009 $3]. Notice that we cannot use the theoretical guarantee of SMD for SCO [Nemirovski et al. 2009 $2.3]. because the objective function is fixed in SCO.
From the standard analysis of mirror descent, e.g., Lemma 2.1 of Nemirovski et al. 2009 we have
(w,, q:), w) Bu (w, wt) - - (w( W +
gg (We, (we, 9r)l
1/e
Summing the above inequality over I,...,T, we have
T
(w,, (w,,, qt), We - w) Ber wi)
Th
T
Ban (w,, q:)l
The
2
(39)
10
B W, wi)
T
D2
The TG
nw
2
TBe
2
where the last step is due to Nemirovski et al.
2009 (2.42))
B (w, w1) S max V (w) WEW
max
min Vw (w)
D)a
(40)
wEW
wEW
By Jensen's 's inequality, we have
T V Vwo(we:::) We w)
| Wt, qt) - -e(w.q:]]
T
(o Wrr 9r), Wr W () +
(We,
Σ
V woolwer q:) - -B (t), W. w)
D2
nw TG
V w
#wo(wr, Wr, qr) (w,, w. ww)
The
Maximizing each side over W, we arrive
max Σ WEW
o p(w, (w., qt) To(we, 4t)
min wEW
Σ
ø(w, qp)
1=1
(41)
D2
nu
TG
T
max WEW
V wwol(w Wes 9t) - Bw (We, gc), We w/
2
=F((w
Next, bound the last term in 41 i.e., maxwew F(w). Because F 10-1/@_ (We, qe) V Vwo(w,, qt), F(w) is the sum of martingale difference sequence for any fixed w. However, is not for w argmax F'w). because w depends the randomness of the algorithm. Thus, we cannot directly apply techniques for martingales to bounding maxwew F F(w). This is the place where the analysis differs from that of SCO.
To handle the above challenge, we introduce a virtual sequence of variables to decouple the depen- dency (Nemirovski al. 2009 proof of Lemma 3.1]. Imagine there is online algorithm which performs SMD by using Vw Wr, qt) gw (w, q() as the gradient:
Vt+1 argmin { (Neu (W wr, q.) g.. (We, qt), w vz) B (w, vo))
(42)
weW
where v, Wi. By repeating the derivation of 39 we can show that
Σ V Vwo(wa, qe) gw We. .... w)
(43)
Bud W. W1)
ne
w db W., q:) - B.. 9.) S
D2
TG2
7u
qu
18

---

where in the last inequality, we make use of 40) and
Wwo(Ww, qc) BB (Wo, ww, (1o 9:
9:)
1,, + (wo,
(44)
< (we, qilmm)) + (we:
2G.
Then, have
T
max WEW
V Wwo(w,, Wr, qt) - B q:), w. w)
max WEW
wwb(we, We, q:)
Bio (w,, V2 w)
Σ V
Vwo(We: Bwo W. v.)
43
D2
T
+ 2m
TG
Σ Wwol(ws q.) gm (w:: 9:); We Vv).
The
(45)
From the updating rule of v, in (42), we know that v, independent from V Ww, -:)--g... gm (We, q:). and thus Vi...., Vr is a martingale difference sequence.
Substituting
45
into 411
have
2D
Th2
d( q:)
min Σ
(w, q:)
+
V2
(46)
7/u
2
1=1
1=1
1=1
Taking expectation over both sides, have
2D
5n
TC
E
Σ o o(w,, q:)
min
Σ o(w, q:)
+
2DGV5T
wEW
2
where we The
2D GVST
To establish high probability bound, make of the Hoeffding-Azuma inequality for martingales stated below [Cesa-Bianchi and Lugosi 2006).
Lemma Let V1, V2,... be a martingale difference sequence with respect to some sequence X1, X2,... such that Vi [A,,A, + ca) some random variable A,, measurable with respect XI,..., Ni-1 and a positive constant C₁ Sf Σim1 V,, then for any t > 0,
212
Pr[S, > exp
exp
Σ
T2 =i=
To apply the above lemma, we need to show that V, is bounded. Indeed, we have
(V o(w,, q:) (we, q:), W. veYl < V Woolww,cc qi) gw (We, qe)llu |w, vellw
2G||w: - vellu S 2G (I Wi filllu. w + l|ve willw)
40
<2G 2V (We. w1) + 2V w1)
4V2DG.
From Lemma I with probability at least 1 - 8, we have
T
V1 s 8DG Σ
Th
(47)
We complete the proof by substituting
into 46
19

---

B.4 Proof of Theorem 6
Because we can only observe ((w,, zi ( instead of R., (We). the theoretical guarantee of Exp3-IX
[Neu 2013] cannot be directly applied to Algorithm 2 regret analysis of Exp3-IX to stochastic rewards.
To address this challenge, we generalize the
By the definition of of(w ww q) in 3 and the property of linear optimization over the simplex, we have
m
m
max qeam
WW, ②
Φ We. qt)
max qmam
R.(w:)
ge, Re (we)
T
T
T
T
(W,))
Σ
qt Ra(we)
= Σ
FumP, [((we:z)
ye Fan [((W;;)) (48)
1=1
1=1 i=1
1=1
1=1
T
11,5
81,
Σ
(q. si)
81,
=1i=
1=1
t=1
t=1
where j E aargmax
T -1=1 Ríí (we) and the vector S; E IR' is defined as
81 Fampy (!(We;))) [0,1]. 1]. Vi [m].
To facilitate the analysis, we introduce vector $t E IR with
8,, - Wr:
0, Vi E [m]]
(49)
where (t) denotes a random sample drawn from the i-th distribution. Note that & is only used for analysis, with the purpose of handling the stochastic rewards. In the algorithm, only $t.i.
1- 1( ((+) is observed in the t-th iteration.
Following the proof of Theorem Neu 2015]. we have
T
Σ (q:.
8t.
In m
fg
St,
(50)
7/
2
t=1
1=1
i=1
which makes use of the property of online mirror descent with local norms 2012 Theorem 5.5]. From (5) of [2015], we have
[Bubeck and BisacBianchi
(q) 81)
Σ 1=1
8t,5
St,
(51)
i=1
Combining (50p and (1).
we have
St.i
+
17 +7)
+
Imm
(52)
2
1g
t=1
t=
tmm
From 48
we have
T
T Σ
max
(WW,,og q)
o(Wi, q:)
qe
1=1
1=1
T
St,ic
81,j"
+
qu St))
St
(53)
1=1
t=1
t=1
T
st.d
Tl
+
7
m
T
(qq: st)
St,te)
In
2
71
:=B B
We proceed bound the above three terms A, B and C respectively.
To bound term A, we need the following concentration result concerning the IX loss estimates 2015 Lemma 1]. which we further generalize the setting with stochastic rewards.
Neu
20

---

Lemma 2 Let Et, E [0, for all [T] and i E [m]] and Eτ, be its IX-estimator defined as E., = E., Ihi = 8). where Et is such that ElSr.i] = Er,i with fr. E [0, and the index i, is Pe.c+he
sampled from [m]] according to the distribution Ps E Am Let {Te be a fixed non-increasing
sequence with 7t. 2 0 and let or, be non-negative F-me-measurable random variables satisfying
ot,i S 2n.for all & (T) and i E [m]. Then, for any > O, with probability at least 1 6,
m ΣΣ
(Er, Ee,s) In
(54)
Furthermore, when "/- x for all t E [2', the following holds with probability at least 1 - 8,
T
Σ
(Et. Ee,)) S 24
(55)
simultuneously for all i E [m]]
Notice that our construction of 5t satisfies that $t,4 R1. li = i] and drawn from [m]
according to qe E Am as well as E[ = se,s, which meets the conditions required by Lemma
As result, according (55). we have
Bi. Σ t=1
8t,5)
1 S In 2-
m
8
for all j [m]] (including j*) with probability least 1 8. To bound term B, we can directly use Lemma of Neu 2015). because our setting 2a y satisfies its requirement. Thus, with probability at least 1 - 0, we have
fl4 2
+
m
1lg 2
+
m
S + In
49
11q 2
+ + mT + 5
(56)
We now consider term C' (33)
Let V₁ = (qe s,) - S,,,, Then, it easy to verify that [V]]
So, the process {V:)== forms a martingale difference sequence and also satisfies VVl S 1 for t. Hence, can apply Lemma and have
T
T 2
[++nn] 1n
(i) - S.,..) S Σ
2T
2T In
with probability at least 1 8.
Combining the three upper bounds for the terms A, B and C, and further taking the union bound, we have, with probability at least 1
o(::,q) q)
((we, q:)
A
t=1
t=1
3m
[g
3
T
In m
<
ln
+
+
mT
+ In
1
+
In
+
2"
8
2
8
2
8
na
mT
3m
T
T
2 22m2l
mf Inm
In
+
+
1
+ 1
In
Inm
8
2
2
T
mf
T
333ml In
+
+
+
1
In
2
Ium
2
8
where the third line holds because of our parameter settings 7
Ug 2
and f)
In m mf
21

---

Based on this high probability guarantee, we can then obtain the expected regret upper bound using the formula that
! Pr x x
E[X]
-
PrPxxxx
holds for any real-valued random variable taking
Bubeeck and Cesa- Cesa-Bianchi
2012 $ 3.2]. particular,
X
(v
mf Inm
V
T
+
max gE
Σ
o Wi q)
Σ o wo; (() 3V mmmmlmm 1/2
T
yields E[X]
s 3, which implies that
T
T Σ
mT
- +1
E
max PEam
o w. q)
o W. q:)
S mm +
V +3
lum +
1/2 2 +1
B.5 Proof of Theorem 3
For the stochastic gradients in 22 their norm be upper bounded in the same way 9
That is,
(wz, 9e)lu,
Σ
ye, Ve(waz)
Σ iECt
qu,
|
G.. = G,
LECt
ic
w.+
[, 9e]l = max
[e(
1.
iEC,
So. with exactly the same analysis as Theorem
we have
G0 + lmm)
( q)D 2
n1
and with probability least 1 - 5,
E Sw q) (8 2 (8+2ln? 2)
10 Inm)
(57)
f1
Next, discuss how bound the risk of on every distribution P., i.e., (w),). Following the derivation in (), we know
max (w() min max (w) S (w, #E(m)
wEW ic(m
Thus, for every distribution P,, (w) can be bounded in the following way:
1 min max p.R,(w) _ (w,()))). q)
(,(w)= <
Pi wEW ie(m)
Taking the high probability bound in (77 as an example, we have with probability at 1 1
1
1
8 (8+21m 2 5))
10 + +nm)
(w( < min max p: R2 (w)
P3 NEWie(()
Pi
(58)
7g
min max R, (w) +
+ 2 2)
10(D2G2 + mmm)n
ny weW sc|m]
n4
B.6 Proof of Theorem 4
We first provide some simple facts that will used later. From Assumption 3 we immediately know that each risk function (-(-) also belongs to [0, 1]. As result, the difference between each risk function and its estimator is well boll-bounded, i.c., for all i [m].
< (,(w) ((w;z) < 1, W, z~ Pt.
(59)
22

---

From Assumption this, have
we can prove that each risk function R.(-) is G-Lipschitz continuous. To see
[W)(w)) ||Ezp, Ve(w; S Ez-P; E G, E W, E [m]. (60)
As result, we have
[R((w) (((w)) S Gl[ -wllw: Vw, w E W,ie [m].
(61)
Furthermore, the difference between the gradient of Rz(-) and its estimator is also well-bounded, i.e. for all [m]]
60
(V):()) V S [[VR, [|V()
[IV
VW E W. P.. (62) 2G,
Recall the definition of the norm I. and dual norm || for the space & I in 34 and the
distance-gente gennnacingnrrringgggg function vi.) in (35). Following the arguments in Section B.1 the two updating
rules in 25) and 26 can be merged
[We+1:q+1]] argmin {I<lBs. (W), By (w d)), x [w': ((l) + B(x, (w/: (63) *EWXA..
where the = 22 and 111 = 2n Similarly, 28 and 29 are equivalent to
( argmin
{u(lBw (We+1, 9:+1): go (1+1 (++1)) 9-11), xEWxAm
[w/: (w/ a
(64)
Let q(( be the monotone operator associated with the weighted GDRO problem 23 i.e.
q)(w:q]) V w q): Vq q)]
(65)
Σ g.P. VR,: - [p2 R1 Pm Rm (w))
From our constructions of stochastic gradients in and (7 we clearly have
E { [g (w/, q(): gy (/ 90) F( F(Ww: 91),
E [g. (W+11. 9--1): 8y (w+11 9c+1)]} = F([We+1 qu41))
Thus, Algorithm is indeed an instance of SMPA JJuitsky et al. 201T Algorithm 1], and we can use their Theorem and Corollary
bound the optimization error.
Before applying their results, we show that all the preconditions are satisfied. The parameter Ω defined in (16) of Juditsky al. 2011 be upper bounded by
Ω
max B(x, (())
56
max Bw w() + max wEW
Bg q()
xE XEWXAm
D2
qm Inm
(66)
40
max (w) min Vw (w)
+
max pg(g) qeam
min (q(9)
2
D
WEW
wEW
Imm
q
Next, we need to demonstrate that F F(w:; q]) is continuous.
Lemma For the monotone operator F([w; q)). we have
|| [F(w:q)) F(w':[[[[]]]] g/) s ī[[(w- w-: q
where L is defined in 31
We proceed to show the variance of the stochastic gradients satisfies the light tail condition. To this end, introduce the stochastic oracle used in Algorithm
g( g(ww: q])
Ig (w. q): By (, (w.g))
(67)
23

---

where
m
ni
Bio
(. q)
Σ
qePp
7mm
ve(w;z((0))
i=1
==1
n1/Rm
8y W: q)
=
P1
,....,,.........,............. f(w: z (...)...... Pmn Pene (ww; z (m))
and (+,+) is the j-th sample drawn from distribution P.. The following lemma shows that the variance indeed sub-Gaussian.
Lemma For the stochastic oracle g( g(ww; q)). have
[IF((w: ql) g((w: ellll a -
]
E
EXp
o²
where o² is defined in 331
apply the theoretical guarantee of SMPA. Recall that the total number of iterations is nm/ / in Algorithm From Corollary of of et al. 2011). by setting
Based (66), Lemma and Lemma 4
2
11 min
V3L
ViTo2 nm
we have
7İ
20²
02
A2
Anm
Pr
co w. q) 2
1
+ 7A
I
exp
+
exp
7mm
3nm
7mm
3
for all A 0. Choosing A such that (xxp(-A//3) s 8/2 and exp( probability at least 1 1-
8/2, we have with
7T.
202
(1(
333gg 177311o0
2 + log
2
2
02
(w 9 <
1
+7
8/6
nm
3m
nm
nm
Following the derivation of (58.. we have
(w(w)
min max p. R;(w))
p; wEW iE(m)
(68)
7L
o²
2
2 +
14 2
14
+
7
7v
3log
log
P;
3
g
7hy
Inspired by Juditsky et al. verify that
2011
$4.3.1]. we use the value of pt in to simplify
668 It easy to
1/ TTm + /im/ni
Pmax
Thm
1 1+
Pi
1/ Tm
nm
/mm/n1
1 L
Pmax /nm
O
v /lmm
(69)
Pi Tom
Tmm
1m
Ps
+1
/ Um
Wmax
max i6(m]
p.
1
+ )
I'm
n4
V Pm
1
1/ Tm + v //mm/r
Rm
/Wmax
Vumax
pi
1/ /mm +1
V
n
1
o²
1
Wmaxe 18 + In? m)
((11))
1
V + Vn+
O
0
(70)
p
Tym
pi
7m
7m
/na

---

0.70
0.70
SMD(1)
0.65
SMD(m)
0.60
Online(1)
Max risk
0.60
0.50
0.55
0.40
SMD(1)
0.50
(MD(m)
0.30
Ouline(1)
0.45
10
10
10
# of iterations
# of iterations
(a) The synthetic data
(b) The Adult dataset
Figure 1: Balanced settings: max risk of different methods versus the number of iterations
0.70
0.70
SMD(I)
0.65
SMD(m)
0.60
Online(1)
Max risk
Max rick
0.60
0.55
0.40
SMD(1)
0.50
SMD(m)
Online(1)
0.45
10
10a
10
10a
10a
10
# of samples
# of samples
(a) The synthetic
(b) The Adult dataset
Figure 2: Balanced settings: max risk of different methods versus the number of samples
Substituting 69
and
70
into 688
we have
R₁ W
min max (w))
=0
11) /ns/
Vx+ In m
Ps wEW
Pmm
C
Experiments
In this section, we conduct empirical studies to evaluate our proposed algorithms.
C.1 Datasets
Following the setup in previous work [Namkoong and Duchi 2016 Soma al. 2022]. we use both synthetic and real-worie datasets. First, we construct a synthetic data with group number m = 20. For each group [m]. we draw model w e IR 1000 from the uniform distribution over the unit sphere. For distribution P., the sample (x, y) is generated with N 1) and y sign(x w:) (x))) with probability 0.9 and y sign(x) w.) with probability 0.1.
also use the Adult dataset [Becker and Kohavi, 1996). which includes attributes such as age, gender, race, and educational background of 48842 individuals. The objective is to determine whether an individual's income exceeds 50000 USD. We set up m = 6 groups based on the race and gender attributes, where each group represents combination of bllck,, white, others] with [female, male). J.
25

---

0.70
0.70
0.70
(MD(m)
(MD(m)
SMD(m))
0.60
0.60
0.55
SMD,
SAID,
SMD,
0.50
0.50
Risk
SMPA
Risk
SMPP
Risk
SMIA
0.40
0.40
0.40
0.30
0.30
0.25
0.20
0.20
0.10
10a
10
10a
10
10
# of iterations
# of iterations
# of iterations
(a) Risk on Pı
(b) Risk on Pa
(c) Risk on Pe
0.70
0.70
0.70
(MD(m)
(MD(m)
SMD(M))
0.60
0.60
0.60
SMD,
SMD,
SMD
0.50
0.50
0.50
Risk
SMRAA
SMPA
SMPA
0.40
0.40
0.40
0.30
0.30
0.30
0.20
0.20
0.20
10"
10
10a
10a
10
10
10
10a
10
10
10
10a
10"
10
# of iterations
# of iterations
# of iterations
(d) Risk on Ps.
(e) Risk Pio
(f) Risk P12
0.70
0.70
0.70
(M)())
(MD(())
SAID(m)
0.60
SMD, SMPA m
0.60
SMD
0.60
SMD, SMPA
Risk
0.50
0.50
SMPA
Rick
0.50
0.40
0.40
0.40
0.30
0.30
0:20
0.20
0.30
10
10
10a
10a
10
10
10
10
10
10a
10P
10
10
10
10
# of iterations
# of iterations
# of iterations
(g) Risk Pie
(h) Risk Pis
Risk P20
Figure 3: Imbalanced settings with the synthetic dataset: individual risk versus the number of
iterations
C.2 GDRO with Balanced Data
For experiments on the synthetic dataset, we will generate the random sample on the fly, according to the protocol above. For those on the Adult dataset, we will randomly select samples from ech group. In other words, P: is defined as the empirical distribution over the data in the i-th group.
We refer to the method of Sagawa al. 2020 and our Algorithm by SMD(1) and SMD(m) to underscore that they are instances of SMD with sample and m samples in each iteration, respectively. We denote Algorithm 2 as Online(1) learning and uses 1 sample per iteration. We set E(::-) to be the logistic loss and utilize different methods train linear model. In the balanced setting, use the max risk, i.e., max R, (w(,,, as the performance measure. To estimate the risk value, we will draw substantial number of samples, and use the empirical average approximate the expectation.
emphasize that it is based
techniques from online
We first report the max risk with respect to the number of iterations in Fig. I We observe that SMD(m) is faster than Online(1), which in turn outperforms SMD(1). This observation is con- sistent with our theories, since their convergence rates are O ((Volo m)/T). ((1m0((mm) /mooo m)/T), and (0m (log m)/T), respectively. Next, we plot the max risk against the number of samples consumed by each algorithm in Fig. As can be seen, the curves of SMD(mm) and Online(1) are very close, indicating that they share the same sample complexity, i.e., Omm(log m)) By contrast, SMD(1) has much higher sample complexity, i.c., O(m² (log mleg))
26

---

0.70
0.70
0.70
SMD(m))
SMD(m))
SMD(m)
0.65
SMD,
0.60.
SMD,
0.60
SMD,
Risk
0.60
SMPA
Risk
0.50:
SMPA..
Risk
SMPA
0.50
0.55
0.40
0.50
0.30
0.40
10
10
10
10a
10
10
10
10
10"
10-
10
10
10
10"
10*
# of iterations
# of iterations
# of iterations
(a) Risk Pi
(b) Risk P2
(c) Risk Ps.
0.70
0.70
0.70
(MD(m)
SMD(m)
SMD(M)
0.60
SMD,
0.65
SND,, SMIAA.
0.60
SMA,
Risk
0.50
SMPA
Risk
SMPA
0.60:
0.50
0.40
0.30
0.55
0.40
0.20
0.50
0.30
10
10
10
10
10
10
10
10
10
10
10"
# of iterations
# of iterations
# of iterations
(d) Risk on P4
(e) Risk on Ps.
(I) Risk on Pa
Figure 4: Imbalanced settings with the Adult dataset: individual risk versus the number of iterations
C,3 GDRO with Imbalanced Data
For experiments the synthetic dataset, we set the number of samples n, 800 (21 - i), and generate each sample as before. For those on the Adult dataset, we first select 364 samples randomly from each group, reserving them for later use in estimating the risk of each group. Then, we visit the remaining samples in each group once to simulate the imbalanced setting, where the numbers of samples in 6 groups are 26656, 11518, 1780, 1720, 998, and 364.
Similarly, we label the Baseline mentioned in the first paragraph of Section as SMD(m). We designate Algorithms 3 and as SMD, and SMPA to highlight that the former combines SMD with random sampling and the latter one integrates SMPA and mini-batches. In the imbalanced setting, we will examine how the risk each individual distribution decreases with respect the number iterations. Recall that the total number iterations of SMD(m), SMD, and SMPA,, are Mm n and / respectively.
We present the risk on the individual distribution Fig. and Fig. First, we observe that our SMPA.. is faster than both SMD(m) and SMD, across all distributions, and finally attains the lower risk is most cases. This behavior aligns with our Theorem 4 which reveals that SMPA achieves a nearly optimal rate of O m /Rz) for all distributions P,, after n., iterations. We also note that on distribution P.. although SMD, converges slowly, its final risk is the lowest, as illustrated in Fig. 3(a) and Fig. 4(a) This phenomenon is again in accordance with our Theorem 3 which shows that the risk of SMD, on Pi reduces at nearly optimal (1v D(V(log m)/na //)) rate, after nı iterations. From Fig. (3)) and Fig. 4(t we can see that the final risk of SMD(m) on the last distribution P... matches that of SMPAm. This outcome is anticipated, as they exhibit similar convergence rates of O(v ((V(log) (log m)/n..) and O(log m)/ /nm), respectively.
D Supporting Lemmas
D.1 Proof of Lemma 2
The proof follows the argument Neu 2015 Proof of Lemma 1], and we generalize to the setting with stochastic rewards. First, observe that for any i E [m]] and E [T].
27

---

Ee,t
Eo
-1[ir =
Pc, + 7t
Ea.
-T(is
(.. E [0, 1)
Pt, ToSii
27t²" Ss./pe.i
Ilis i)
2%/s 1 + -/+ &t,i Pr,i
Σ log (1 + BeEl.s)
(71)
3.
where the last step is due the inequality 173/2 log(1 + 2) for z 2 and we introduce the notations 2% and É.,s (Pts)) Ilfie i] simplify the presentation. Define the notation d₁ Σm de, and A E a Then, have
m
E4-1 [oxp(ár)
E
cxp
m
Aiti log (1+ BrEt.s)
E
exp
Be
m
E
cxp
log (1 (1+as ansbess)
ari D.
I 1 by assumption)
E
11 + (E,5),
E
1+ Cr S exp
Σ anises = exp((tt
(72)
where the second incquality is by the inequality x log(1 + y) < log(1 + xy) that holds for all
y 2 and E [0, 1], and the equality E,- [1_ (1 St,i) = E,- [1+ Em [r, follows from the fact that Et,, E., = 0 holds whenever * j. The last line due to 16-1 =
(.-1((G.sPP.s) Ife o]] ε..: and the inequality 1 + - e for all E R.
As a result, from 72 we conclude that the process Z₃ = cxp (E: (λ, - x.)) is supermartingale.
Indeed, E2-
E,-1 exp (E!=l Xa)) exp(Ae A1)] s Z,-1- Thus, have E[Zr] <
F( S < F [.]
= 1. By Markov's 's inequality.
T Ed Edh-vs)
Pr
A1)
E
exp
eep((ee S exp(-e)
holds for any > 0. By setting exp(-e) = 8 and solving the value, we complete the proof for (54).
And the inequality (55) for the scenario 7/ = y can be immediately obtained by setting 0,,i = 2 and taking the union bound over all E [m].
28

---

D.2
Proof of Lemma 3
From the definition of norms in 34 we have
[ q
F([w':
q:pp (p.VRR(W)
Σ
(w) .
Pm (w') Pm Rm (w)] T
=2
4iPs V (w) Σ de (R,
1,.0
(P) Ri (w') - p1 (i)R((w), p Rm Rm (w') - Pm Rem (on(w) || In
+
m Σ qipp (w)) Σ dD.VR((W) Σ (ppVRR(()) -2
m
m
dep VR((w')
w,t
Pm Rum (w') Pm Ren (en(w) (w) [] 7| Inm
+2
(p) RI (w') - p1 ()(R((),
m
m hiPi (w)
m
m Σ 1=1
<DD²
+4D²
Σ
dep R.:('))
==A
+ max mm(m)
ps [R)(w) (w)][ Inm.
To bound term A, we have
4D2
Σ
VR.(w)
Σ V.P.R (w)
1=1
2
<4D²
Σ
19: dill (. Ihe,
Σ 1g:
< PD pmnx lq q'1l7
where Pmaxx is defined in 31
To bound B, we have
4D²
Σ P.P.
V R w) Σ dp. VP.RR
<4D²
Σ
A.P. (V)(w)) - VR;(w') Il
4D2 Σ dipp L - w'll
1=1
m
24 p2 || w/ Σ d 4D2 12p2
|| 1www21 w/1
To bound C, we have
max p: [w() R:(w") Inm
x mm]
- w' - Inm S 2G || w'll2 Inm.
max :E(m]
Putting everything together, we have
[[[[w:]] - F(('': < 22 + 2G lnm) 2 ww - w'1 + 4 pi l|q q
Sp (8D* Inm
2D |w - w/1
+
|1q - d/1l
2lmm
(8D*L2+ Inm) [[w 1i-w::q- q']l"
29

---

which implies
[P((w:))--F((ww)) dl) S Pmax D h q-''ll s willw-w'sqq-
where L defined in (31
D.3 Proof of Lemma
The light tail condition, required by Juditsky et al. 2011]. is essentially the sub-Gaussian condition. To this end, we introduce the following sub -abbssassian properties |Vershynin 2018
Proposition 2.5.2].
Proposition 1 (Sub-gaussian properties) Let X a random variable. Then the following proper- ties equivalent; the parameters K; > 0 appearing in these properties differ from each other by at most an absolute constant factor.
(i) The tails of X satisfy
2exp( Vt 0.
(ii) The moments of X satisfy
|X|lr = (E|X|)) s K2vP, Vp 2
(iii) The moment generating function (MGF) of X² satisfies
E| E[exp(>²x3]] S exp(K3X²), vX such that [|] S 1/K3.
(iv) The MGF of X2 is bounded at some point, namely
E[cxp(X*/K3)] 22
Moreover, F[ = 0 then properties (i)(iv) are also equivalent to the following property. (v) The MGF of X satisfies
E exp(AX)) s exp(K3>2), V.E..
From the above proposition, we observe that the exact value of those constant Ki,... K₃ is not important, and very tedious to calculate them. So, the following, we only focus on the order of those constants. To simplify presentations, we use to denote an absolute constant that is independent of the essential parameters, and its value may change from line line.
Since
[F([;; g((w; 1l)
= V q) (w, 12,, + 2 q) -g, (w, 9)1 Inm,
we proceed to analyze the behavior of Vwwpw, -)-g))gwwll To this end, we have the following lemma.
gw (w, q) 122, and // q) -,,,,(w, (1.
Lemma 5
We have
CR | Wwo(w..q q) g.. (w, 1)
exp
(73)
E exp
V (w, q) - (w, (w,
2
Cwimax
Inmm
where Wmax is defined in
31 and 0 is absolute constant.
30

---

From Lemma
we have
E
exp
1
F([w: g((w. (/t..al)ll?
Dc C + 2cw In m
2D
=F
exp
V /ws(ww,g) q) gw q)
2cx D' + 20wmmax In2 m
Imm
Vqpp(w -Bo Bw, a)ll
2Cn DeR
G2 Wmax + 2cwmax In m
RD I/V q)
gw (w, Dl
E
exp
xD2 + +ln²m m
CxGx2mmax [V B (w 1)
In m
KD2 + In2
Owmax In m
G2
V q)
gw (w, 1)
E
exp
kD + In
cxG2
In² m
| VigipW W. q) 8y (w, 12.
kDaG² + In² 17
E
exp
Cidmax In m
73
G22
In
2
kD + In m
xD2
+
2 m
where the first inequality follows from Jensen's inequality.
D.4 Proof of Lemma
To analyze Vwip(w, q) - gw. g. q) 12, we first consider the approximation error caused by samples from P::
/
nty
Tbm
vl(w;z((5)) VR, (w)
um
Vl(w: -
14
n
Under the regularity condition of || lllc, Assumption 6 we have, for any 9 2
nie ni//n
Ply
Ve(w; -VR:()] 10.**
2G(V22 + v Tm < <exp(->212) (74)
n4
ny
which is directly consequence of the concentration inequality of vector norms (Jditsky and Nemirovski] 2008 Theorem 2. 2iii) and (62). Then, we introduce the following lemma to simplify (74).
Lemma 6 Suppose we have
Pr 1x exp(->2/2),
2epp V7 0
where is nonnegative. Then, we have
Pr [ 2 < < (->2) max(6a?, by >
From
74
and Lemma 6 we have
Dee Rem
Pr
n2
Tmm
Vl(w: - 27
2v2G
nm
j=1
<2
(-22/ max (6K, 8)) < 2 exp (-?2/(8k)), V>>0
31

---

which satisfies the Proposition (i). From the equivalence between Proposition 1 (i) and Proposi- tion (iv), we have
E
exp
Py
ni//tm
Vl(w:z(00) -V--R (w)
CRG nm
2.
f4
n4
10,
Inserting the scaling factor pi, we have
/ Pmm
ny
cRG2
E
exp
nlm ps
( Ve(W;z((0)) [w)
)
$2.
(75)
n2
10.0
To simplify the notation, we define
/
Ve(w;z(1o1) VR:(w)) -VR(()))))))((]))]]
p T4m
pb
and Wmax
n
iE [m]
By Jensen's 's inequality, we have
E
exp
V (wp((w
-w (w, Dill
CKG2
Wmax
E
exp
Σ
quu
[onCAWwnax]]
E
lexp
(Iluellee / [cn[FWimxx]]
9:2 = 2.
where we use the fact that || I/ (-)2 and exp(-) are convex, and the last two functions are increasing +.
To analyze
|| Va soopw.gg q) - (w, 1) we consider the approximation error related to P;:
ne/nm
ne/nu
nmm
E(w; ((; ((r) R((w)
nmm
((wiz(b)) Rs Rw(w) (w)
f4
Note that the absolute value derivation of 775).
|. is 1-regular (Juditsky and Nemirovski
2008
Following
59) and the
have
exp
Ps
nm
TR, /
((W:2 -R(
cpt nm
<
2
(76)
To prove that ||V q (w, q)
(w, 1) is also sub-Gaussian, we need to analyze the effect of the infinity norm. To this end, we develop the following lemma.
Lemma Suppose
(exp (IX,P//K3) 2 Vj [m].
(77)
Then,
|X, [cK2 Bax In ml] ) 5 2.
exp
max jE(m]
where c 0 is an absolute constant, and Kmax maX Ky.
From
76 and Lemma 7 we have
Vqpp(w, q) - 8, (w. ) 22
exp
CWmax In m
32

---

D.5
Proof of Lemma 6
When y [0, 2a], have
- 22exp(->2/6el
When 7 2 2o, have
Pr|x 2 7 - a)²/2) - exp(->2/8)
where we use the fact
7 - ? 2
Thus, we always have
Pr|x 2 2
2xxp (-12/ 5
max .8), by >
D.6 Proof of Lemma
From
and the equivalence between Proposition (i) and Proposition (iv), we have
Pr [r[ 2 1] < 2exp (-2/)))), As 2 0, E [m]-
As result,
Pr
max
X2 2 = Pr 3j, X 2 1 Pr X < Σ
exp (-P/cK3)
_1(mm
f=1
22m exp (-t² /-t//K/s cK max) exp (-t² ee max + In (m(2m)).
Choosing t VcK2 max (In[2m) 222 have
Pr
max jE[m]
X,] 2 VcK2 max (In(2m) + (2/2)
eexp (-g212).
Thus
Pr
max jE[m]
X3l 2 ck2
[Vin[mm] + +///2) S Sexp (-72/2)
max
+ Pr
max max JE/mJ
Xyl 2 /21n(2mm +7 S Sexp (-72/2).
ck
By Lemma
we have
(22m).)). max 1x112> < 2exp (->2/ max (12 In [n[2m], 8)), Vy > 0.
Pr
ck max
jE(m]
From the equivalence between Proposition (i) and Proposition (iv), we have
|X,²² / [cK2 max m] )] 22
exp
max [E(m]
33


Reviews:
[{"review_id": "tiveK2nhBT", "reviewer": "Reviewer_41n7", "summary": "This paper provides algorithms for group distributionally robust optimization.\nFor the case in which the sample budget for each distribution is identical,\nthe proposed algorithm achieves nearly optimal sample complexity.\nThis paper also considers the sample budgets are different for each distribution and provide an algorithm with non-trivial sample-complexity bounds.", "questions": "- Can we remove extra $\\log m$ factors?\n- It is claimed that the $O(1/\\sqrt{n_i})$ rate is optimal, which is not obvious to me. \nIt is understood that for normal (non-distributionally-robust) learning over a single distribution, the optimal rate is $O(1 / \\sqrt{n_i})$.\nHowever, the definition of risk indicators employed in this paper (LHS in Theorem 4) differs from the usual.\nTherefore, in order to discuss what the optimal rate is, it is necessary to again prove a lower bound on this risk.\n- In the paper, one contribution is explained as reducing the number of samples in each iteration.\nThough I agree that total sample complexity matters,\nit was not clear to me what practical benefit there is in reducing the number of samples per iteration (while keeping the total number of samples the same). \nI would appreciate a more detailed explanation of the benefit.", "limitations": "The limitations and potential negative societal impact are adequately addressed.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- This paper is the first to consider a practical setting in which the sample budget varies by distribution.\n- The paper is well structured and easy to follow.", "weaknesses": "- Given the results of [Soma et al., 2022], the contribution for the uniform-sample-budget setting appears somewhat limited.\n- The obtained bounds in this work includes extra $\\log m$ factors, unlike ones by [Soma et al., 2022].", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Many thanks for the constructive reviews!\n\n---\nQ1: The obtained bounds in this work includes extra $\\log m$ factors, unlike ones by [Soma et al., 2022].\n\nA1: First, it is important to recognize that the original analysis of Soma et al. [2022] is *flawed*. For a detailed explanation, please refer to **A1 to Reviewer EKF4**. Second, after fixing the error in their analysis, we are only able to derive the bound shown in (19) of our paper. Unfortunately, this bound is not useful, as it cannot be employed to control the optimality gap. To put it more precisely, the bound in (19) cannot be directly compared with our upper bounds, such as that in (18), due to the difference in the left-hand side of the inequality. Consequently, although (19) does not contain the $\\log m$ factor, it is not stronger than our results.\n\n---\nQ2: Given the results of [Soma et al., 2022], the contribution for the uniform-sample-budget setting appears somewhat limited. \n\nA2: Below, we summarize the novelty of Algorithm 2 compared with Soma et al. [2022].\n1. It's crucial to realize the non-oblivious nature of the online process when selecting the appropriate algorithm and performing the analysis. In particular, the second player must employ an online algorithm designed for non-oblivious MAB. In contrast, the online algorithms and the analysis used by Soma et al. [2022] do not support the non-oblivious property. Therefore, they fail to bound the optimality gap.\n\n2. Although our Algorithm 2 is composed of SMD and Exp3-IX, existing theoretical guarantees for them cannot be applied directly. That is because the non-oblivious nature, coupled with the stochastic setting, makes the analysis much more challenging. To bound the regret of SMD, we need to make use of the “ghost iterate” technique, as detailed in Section C.3 of the supplementary material. Similarly, when bounding the regret of Exp3-IX, we need to construct new martingale sequences to deal with stochastic rewards, as shown in Section C.4.\n---\n\nQ3: Can we remove extra $\\log m$ factors?\n\nA3: We think it is impossible to remove the $\\log m$ factor. It's worth noting that even in the context of non-oblivious MAB, we do have an additional $\\log m$ factor, and it is conjectured to be unimprovable [Audibert and Bubeck, 2010, Remark 14]. Another evidence is that in a closely related problem—Collaborative PAC Learning [Blum et al., 2017], the lower bound contains a $\\log m$ factor.\n\nReference:\nJean-Yves Audibert and Sebastien Bubeck. Regret bounds and minimax policies under partial monitoring. Journal of Machine Learning Research, 11:2785–2836, 2010.\n\n---\n\nQ4: It is claimed that the $O(1/\\sqrt{n_i})$ rate is optimal, which is not obvious to me. … However, the definition of risk indicators employed in this paper (LHS in Theorem 4) differs from the usual. \n\nA4: Notice that when $m=1$, the LHS in Theorem 4 becomes the *standard* excess risk. Thus, if we can establish a rate that is faster than $O(1/\\sqrt{n_i})$, then applying it to the case $m=1$ means that the excess risk is reduced at a rate faster than $O(1/\\sqrt{n_i})$, which is impossible.\n\n---\nQ5: It was not clear to me what practical benefit there is in reducing the number of samples per iteration (while keeping the total number of samples the same).\n\nA5: It reduces the running time for a single iteration, enhancing the algorithm's real-time capability. In other words, Algorithm 2 can produce an output for every single sample received, whereas Algorithm 1 requires receiving $m$ samples before it can produce a new output.\n"}, {"review_id": "EAyHayLCEY", "reviewer": "Reviewer_evdu", "summary": "In this paper, the authors present a systematic investigation of group distributionally robust optimization (GDRO). First, they consider the standard standing where the algorithm can draw samples from all distributions freely. In this case, they demonstrate that stochastic mirror descent (SMD) attains an optimal $O(m \\log m/\\epsilon^2)$ sample complexity by drawing $m$ samples per iteration. Then, they make use of non-oblivious MAB to reduce the number of samples from m to 1, and obtain the same complexity.\n\nThe authors also consider a practical setting where different distributions may have different sample budgets. To this end, they introduce a weighted formulation of GDRO, and develop two stochastic approaches based on non-uniform sampling and mini-batches. By setting the weights appropriately, they are able to establish distribution-dependent convergence rates.", "questions": "1. The authors have mentioned that the work of Haghtalab et al. [2022] suffers a dependency issue. Could you please elaborate on this point? In particular, which part of their analysis is problematic?\n\n2. Is there any other work on GDRO that also investigates the setting in Section 3?\n\n3. When the risk functions have additional structure, such as strongly convex, it is possible to get better sample complexity?", "limitations": "The authors have discussed the limitations of their approaches in their remarks. I do not find any potential negative societal impact of this work.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. DRO/GDRO is an important topic in machine learning, and has wide applications. The authors identify technical flaws in existing works [Haghtalab et al., 2022, Soma et al., 2022], and present simple and principled stochastic approaches. i.e., Algorithms 1 and 2. The sample complexities of their methods are nearly optimal, and the idea of applying non-oblivious MAB to GDRO in Algorithms 2 is smart and novel. \n\n2. The new setting in Section 3 is important in practice, since it is common to encounter imbalanced data sets. The formulation of weighted GDRO is interesting, and provides one possible solution. The theoretical guarantee of Algorithm 4 is very strong in the sense that it demonstrates nearly optimal rates for multiple distributions simultaneously.\n\n3. The paper is well-written and acknowledges previous work properly. The authors present their techniques very clearly, and the results are sound to me.", "weaknesses": "There are no empirical studies. While theoretical contributions stand on its own right, I would like to suggest the authors to conduct experiments to verify the main theorems.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Many thanks for the constructive reviews!\n\n---\n\nQ1: The authors have mentioned that the work of Haghtalab et al. [2022] suffers a dependency issue. Could you please elaborate on this point? In particular, which part of their analysis is problematic?\n\nA1: In the work by Haghtalab et al. [2022], specifically on page 24, the proof following (15) is problematic. This stems from the fact that they *reuse* samples in their algorithm, which causes their stochastic gradients to be *biased*. As a result, they cannot invoke Fact C.3 to derive the final equation on Page 24. Though Haghtalab et al. label their stochastic gradients as “globally unbiased” on Page 6, this characterization is incorrect. As long as samples are reused, the gradient inherently becomes biased. Even within the context of traditional stochastic optimization, we are not aware of any established theory that allows for the reuse of samples across multiple iterations.\n\n---\n\nQ2: Is there any other work on GDRO that also investigates the setting in Section 3?\n\nA2: We have not identified any studies within GDRO that consider the imbalanced scenario. In the literature, we only find a related work in agnostic federated learning, where Mohri et al. [2019] examine the imbalanced scenario in the *offline* setting. However, their generalization error bound is dominated by the distribution with the fewest samples [Mohri et al., 2019, Corollary 1], and they did not provide distribution-dependent convergence rates.\n\n---\n\nQ3: When the risk functions have additional structure, such as strongly convex, it is possible to get better sample complexity?\n\nA3: If we adhere to the convex-concave problem in (3), it seems impossible to get a better sample complexity. This is due to the fact that the convergence rate is contingent on both $\\mathbf{w}$ and $\\mathbf{q}$. For example, if the risk function is strongly convex in $\\mathbf{w}$, we know that $\\phi(\\mathbf{w}, \\mathbf{q})$ is also strongly convex in $\\mathbf{w}$; However, $\\phi(\\mathbf{w}, \\mathbf{q})$ remains linear in $\\mathbf{q}$, making it hard to improve the convergence rate.\n\nOn the other hand, it may be possible to improve the sample complexity if we focus on the original problem in (2). This possibility arises from the fact that if all the risk functions exhibit strong convexity, their pointwise maximum will also be strongly convex. We intend to delve into this intriguing aspect in our future research.\n\n---\n\nQ4: I would like to suggest the authors to conduct experiments to verify the main theorems.\n\nA4: We have carried out preliminary experiments, the results of which can be found in the global response to all reviewers. These empirical findings align closely with our theories.\n"}, {"review_id": "vqAKYTytO1", "reviewer": "Reviewer_Gpdz", "summary": "The paper presents a comprehensive study on group distributionally robust optimization (GDRO) by formulating it as a convex-concave saddle point problem. The authors propose utilizing mirror descent algorithms and online convex optimization techniques to effectively address this problem. Furthermore, the authors investigate the case of unequal sample sizes drawn from each distribution, leading to the development of a weighted GDRO algorithm. This work provides valuable insights into GDRO and offers practical solutions to address various scenarios.", "questions": "There are several technical questions:\n1) Assumption 2 would be much more natural if it was stated as an assumption on the loss function and not the risk.\n2) When maxima and minima are taken within the paper - I don’t understand why they exist. For example, in (4) why are maxima and minima attained? Similarly in (7).\n3) Where is \\bar q needed in Theorem 3? Or is it used at all? If not why does Algorithm 3 return \\bar q?", "limitations": "The novelty and contribution of this paper, in my understanding, is questionable. ", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The GDRO problem presented in the paper is interesting and well-defined. It has broad applications and holds significant interest. The authors demonstrate a deep understanding of the literature on mirror descent algorithms and their variants, showcasing their expertise in the field.", "weaknesses": "Unfortunately, the paper has several weaknesses, which are listed below:\n\n1) Contribution is unclear: Upon careful review of the paper, I have concerns regarding the lack of clarity regarding the contributions. Section 2 appears to primarily present known results, where the formulation of GDRO as a convex-concave saddle point problem and its solution using stochastic mirror descent are well-established techniques. As a result, Theorem 1 seems to be directly derived from the work by Nemirovski et al. (2009). Similarly, Algorithm 2 appears to be a widely-known algorithm in the online convex optimization (OCO) literature, with known convergence guarantees (Theorem 2).\nIn this context, Section 3, which focuses on the \"weighted GDRO and SA approaches,\" seems to be the section where the potential contributions lie. However, it appears to be a straightforward extension of Algorithm 1 and Theorem 1, which have already been outlined in Juditsky et al. (2011).\nTo provide a clearer understanding of the novel and non-trivial aspects of Theorem 3 and Algorithm 3, it would be beneficial for the authors to explicitly highlight the specific advancements or modifications they have introduced. Clarifying the unique features, surprising results, or novel insights that differentiate Theorem 3 and Algorithm 3 from the existing literature will help in addressing the concerns raised regarding the originality and significance of the contributions.\n\n2) Numerical experiments are missing: I am perfectly fine if theoretical NeurIPS papers do not contain numerical experiments. This paper, however, presents an algorithm (or actually  algorithms) for solving a practically relevant GDRO problem. In my opinion, it is absolutely necessary that when presenting algorithms to solve a known problem - these algorithms are numerically verified. Also if something is new about Algorithm 3 (which I am not fully sure - see point above) - I would be curious to see if this algorithm behaves better then existing approaches.\n\n3) Technical issues: See questions in the next paragraph.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Many thanks for the constructive reviews! We provide detailed responses below, and hope the reviewer could reassess the significance of our results. We are looking forward to addressing any further question in the reviewer-author discussion period.\n\n---\nQ1: Theorem 1 seems to be directly derived from … Nemirovski et al. (2009)\n\nA1: We totally agree. As we detailed in Lines 47, 55 and 147, Algorithm 1 is a routine application of SMD to (3), and Theorem 1 directly follows from Nemirovski et al. [2009]. In fact, we are very *surprised* that recent studies overlooked the fact Nemirovski et al. [2009] had explored the same problem in Section 3.2 of their paper. So, we feel it is necessary to present Algorithm 1 and give credit to Nemirovski et al. [2009].\n\n---\nQ2: Algorithm 2 is a widely-known … with known convergence guarantees (Theorem 2)\n\nA2: At the first glance, Algorithm 2 might appear to be a straightforward combination of SMD and Exp3-IX. However, the design and analysis of Algorithm 2 present several challenges:\n1. It's crucial to realize the non-oblivious nature of the online process when selecting the appropriate algorithm and performing the analysis. In particular, the second player must employ an online algorithm designed for non-oblivious MAB. It's worth noting that this is not a straightforward task; previous work [Soma et al. 2022] did not recognize this issue, making their analysis flawed. For details, please refer to **A1 to Reviewer EKF4**.\n\n2. The non-oblivious nature, coupled with the stochastic setting, means that existing theoretical guarantees for SMD and Exp3-IX cannot be applied directly. To bound the regret of SMD, we need to make use of the “ghost iterate” technique, as detailed in Section C.3. Similarly, when bounding the regret of Exp3-IX, we need to construct new martingale sequences to deal with stochastic rewards, as shown in Section C.4.\n\nIn summary, the integration of SMD and Exp3-IX in Algorithm 2 goes beyond mere combination, and involves innovative problem-solving to yield the desired rates.\n\n---\nQ3: Section 3 …be a straightforward extension of Algorithm 1 and Theorem 1, which … outlined in Juditsky et al. (2011). To provide a clearer understanding of … Theorem 3 and Algorithm 3, it would be …\n\nA3: We believe the reviewer is referring to Theorem 4 and Algorithm 4, rather than Theorem 3 and Algorithm 3. Below, we highlight the novelty of them.\n1. While Juditsky et al. [2011, §4.3.1] introduced the concept of \"scale factors\" in solving the Stochastic Semidefinite Feasibility problem, their approach and context are distinct from ours. In their research, the number of samples drawn from different distributions is equal, but the losses can differ and the distributions are heterogeneous. Conversely, in Section 3 of our paper, the number of samples from various distributions may vary, but the losses are identical and the distributions are homogeneous. As such, we cannot directly apply the techniques of Juditsky et al. [2011].\n\n2. To adapt the concept of \"scale factors\" to our specific problem, we innovatively utilize mini-batches to reduce the variance of the stochastic gradient. As the sizes of the mini-batches differ across distributions, the variances become disparate as well. The larger the number of samples, the smaller the variance is. Through this approach, we can simulate the effect of heterogeneous distributions, allowing us to incorporate \"scale factors\" into our methodology. This novel adaptation of existing concepts represents a key aspect of our contribution to the imbalanced scenario.\n\n3. When analyzing the convergence behavior of Algorithm 4, determining the value of $p_i$ proves to be a significant challenge. It required extensive effort to ascertain the setting in (57), a discovery that led to elegant simplifications in (68) and (69). The advancements of Theorem 4 are the *distribution-dependent* convergence rates, which, to the best of our knowledge, represent a novel contribution to the field of GDRO.\n---\nQ4: Contribution is unclear\n\nA4: We summarize our contributions below.\n1. We regard Algorithm 2 as one main contribution, in which we tackle the *non-oblivious* challenge in both the algorithm's design and its theoretical analysis.\n2. We consider Algorithm 4 as another major contribution, where we integrate mini-batches and “scale factors” in a novel manner, leading to *distribution-dependent* convergence rates.\n3. We treat Algorithm 3 as a more modest contribution. Nevertheless, it holds value, since its convergence rate complements the Baseline, and it offers a clear motivation for the weighted GDRO formulation.\n\nUndoubtedly, Algorithm 1 is credited to Nemirovski et al. [2009]. We just want to bring this fact to the community's attention.\n\n---\n\nQ5: Numerical experiments are missing\n\nA5: We have carried out preliminary experiments, the results of which can be found in the global response to all reviewers. These empirical findings align closely with our theories.\n\n---\nQ6: Assumption 2 … more natural if it was stated as an assumption on the loss and not the risk\n\nA6: While it may seem more intuitive to assume the convexity of the loss function, it's important to note that this assumption is more *stringent* than assuming the convexity of the risk function, as the former implies the latter. Specifically, Assumption 2 underscores that our algorithms and theories can also accommodate cases where the loss function is non-convex, so long as the risk function remains convex.\n\n---\nQ7: When maxima and minima are taken … why they exist\n\nA7: The reason is because we assume $\\mathcal{W}$ is convex and bounded in Line 120. Note that we employ the *same* assumptions as Nemirovski et al. [2009], which are standard in stochastic approximation.\n\n---\nQ8: Where is $\\bar{\\mathbf{q}}$ needed in Theorem 3? …\n\nA8: We apologize for the confusion. $\\bar{\\mathbf{q}}$ is utilized solely within our analysis and is not required to be returned by Algorithm 3.\n"}, {"review_id": "VmIiUPEFUA", "reviewer": "Reviewer_EKF4", "summary": "The paper considers group DRO and shows an SMD algorithm to optimize. This algorithm requires a sample from every distribution in each iteration and matches the sample complexity lower bound up to a log factor. The second algorithm is designed for the case when the online functions depend on past decisions. For this, the authors use SMD and Exp3-IX. Finally, for the case with sample budgets from each distribution, the paper proposes two algorithms, one satisfies the constraints on the number of samples in expectation and the other using mini-batches satisfies the exact constraints. ", "questions": "1. Can the authors discuss more the comparison with Soma, Gatmiry & Jegelka (2022) in my concern above? \n2. I think it would be better to add some experiments.", "limitations": "The authors discuss some limitations of the proposed algorithms.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Despite some weaknesses, I think all algorithms proposed in this paper are simple and natural, which at the same time achieve optimal sample complexity. \n- The presentation is clear to follow with proper explanations and intuitions.", "weaknesses": "1. I think the weakness of the first algorithm is that in many cases, it might be impossible to obtain samples from all distributions.\n2. I’m not sure I understand the novelty of the second algorithm compared with Soma, Gatmiry & Jegelka (2022). In the comparison after remark 3, the authors point out an issue with this prior work. The authors use an alternative objective ($\\epsilon_\\phi$) to control the optimality gap, while Soma, Gatmiry & Jegelka (2022) directly bound the optimality gap instead of $\\epsilon_\\phi$, I don’t see any issue with not being able to bound $\\epsilon_\\phi$. Note that in Soma, Gatmiry & Jegelka (2022), the authors also propose a few algorithms in the same spirit as Algorithm 2, which I think work for the non-oblivious case as well.\n3.  Mini-batch sizes in Algorithm 4 can be very large.\n4. No experiments", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Many thanks for the constructive reviews! We believe there are some misunderstandings about the work of Soma et al. [2022], and provide detailed responses below. We hope the reviewer could examine them, and reassess the significance of our results. We are looking forward to addressing any further question in the reviewer-author discussion period.\n\n---\n\nQ1: In the comparison after remark 3, the authors point out an issue with this prior work. The authors use an alternative objective ($\\epsilon_\\phi$) to control the optimality gap, while Soma, Gatmiry & Jegelka (2022) directly bound the optimality gap instead of $\\epsilon_\\phi$. I don’t see any issue with not being able to bound $\\epsilon_\\phi$. Note that in Soma, Gatmiry & Jegelka (2022), the authors also propose a few algorithms in the same spirit as Algorithm 2, which I think work for the non-oblivious case as well.\n\nA1: In fact, the analysis of Soma et al. [2022] is *flawed*, and they fail to bound the optimality gap. Besides, their analysis is based on the theoretical guarantee for oblivious online learning [Orabona, 2019], and thus does not work for the non-oblivious case.\n\nBelow, we explain why their analysis is problematic. To facilitate understandings, we use the notations of Soma et al. [2022]. From (12) of their paper, we know that the key is to bound the expectation of two regret:\n$$\n\\mathbf{E}[R_\\theta(T)]= \\mathbf{E}\\left[ \\sum_{t=1}^T L(\\theta_t,q_t)\\right] - \\mathbf{E}\\left[\\min_{\\theta \\in \\Theta} \\sum_{t=1}^T L(\\theta,q_t) \\right]\n$$\nand\n$$\n\\mathbf{E} [R_q(T)]= \\mathbf{E}\\left[ \\max_{q \\in \\Delta_m} \\sum_{t=1}^T L(\\theta,q_t) \\right]- \\mathbf{E}\\left[ \\sum_{t=1}^T L(\\theta_t,q_t) \\right].\n$$\nThen, in Section B.1, they make use of Lemmas 2 and 3 to bound $\\mathbf{E}[R_\\theta(T)] $ and $\\mathbf{E} [R_q(T)]$, which is unfortunately incorrect. Let's take $R_\\theta(T)$ as an example to elucidate the reason. From Lemma 3, they can bound the regret in terms of the random function $\\ell(\\theta,z_t)$, where $z_t$ is the random sample generated in the $t$-th round:\n$$\nr_\\theta(T)=\\sum_{t=1}^T \\ell(\\theta_t,z_t)- \\min_{\\theta \\in \\Theta} \\sum_{t=1}^T \\ell(\\theta,z_t).\n$$\nHowever, the expectation of $r_\\theta(T)$ is *not* equal to $\\mathbf{E}[R_\\theta(T)]$, and thus Lemma 3 cannot be applied. To see this, we have\n$$\n\\mathbf{E}\\left[ r_\\theta(T) \\right] = \\mathbf{E}\\left[\\sum_{t=1}^T \\ell(\\theta_t,z_t)- \\min_{\\theta \\in \\Theta} \\sum_{t=1}^T \\ell(\\theta,z_t) \\right] = \\mathbf{E}\\left[\\sum_{t=1}^T L(\\theta_t,q_t) \\right]- \\mathbf{E}\\left[\\min_{\\theta \\in \\Theta} \\sum_{t=1}^T \\ell(\\theta,z_t) \\right].\n$$\nNotice that we *cannot switch the order of expectation and minimization*, and thus\n$$\n\\mathbf{E}\\left[\\min_{\\theta \\in \\Theta} \\sum_{t=1}^T \\ell(\\theta,z_t) \\right] { \\color{red} \\neq} \\mathbf{E}\\left[\\min_{\\theta \\in \\Theta} \\sum_{t=1}^T L(\\theta,q_t) \\right] \\Rightarrow \\mathbf{E}\\left[ r_\\theta(T) \\right] { \\color{red} \\neq} \\mathbf{E}[R_\\theta(T).\n$$\n\nWe have attempted to rectify the aforementioned issue in their analysis.  In doing so, we have to eliminate the $\\min$ and $\\max$ operators, leading to the bound shown in (19) of our paper. But this bound is not useful, as it cannot be employed to control the optimality gap.\n\n---\n\nQ2: I’m not sure I understand the novelty of the second algorithm compared with Soma, Gatmiry & Jegelka (2022).\n\nA2: At the first glance, our Algorithm 2, which integrates SMD and Exp3-IX, may seem similar to the approach used by Soma et al. [2022]. However, both the design and analysis of Algorithm 2 are fundamentally different from theirs. We outline the novel aspects below.\n\n1. It's crucial to realize the non-oblivious nature of the online process when selecting the appropriate algorithm and performing the analysis. In particular, the second player must employ an online algorithm designed for non-oblivious MAB. In contrast, the online algorithms and the analysis used by Soma et al. [2022] do not support the non-oblivious property. Therefore, they fail to bound the optimality gap, as discussed in A1.\n\n2. The non-oblivious nature, coupled with the stochastic setting, means that existing theoretical guarantees for SMD and Exp3-IX cannot be applied directly. To bound the regret of SMD, we need to make use of the “ghost iterate” technique, as detailed in Section C.3 of the supplementary material. Similarly, when bounding the regret of Exp3-IX, we need to construct new martingale sequences to deal with stochastic rewards, as shown in Section C.4.\n---\n\nQ3: I think the weakness of the first algorithm is that in many cases, it might be impossible to obtain samples from all distributions.\n\nA3: In that case, we could use Algorithm 2, which only draws 1 sample per iteration, or Algorithm 3, which uses less samples via random sampling.\n\n---\n\nQ4: Mini-batch sizes in Algorithm 4 can be very large\n\nA4: It only happens when the data is highly imbalanced. In such cases, we can combine the idea of random sampling and mini-batches, which will increase the number of iterations and reduce the size of mini-batches.\n\n---\n\nQ5: I think it would be better to add some experiments.\n\nA5: We have carried out preliminary experiments, the results of which can be found in the global response to all reviewers. These empirical findings align closely with our theories.\n"}, {"review_id": "nhu3iP97vP", "reviewer": "Reviewer_JWj5", "summary": "This paper provides near-optimal optimization algorithm to group distributionally robust optimization problem (GDRO). Since it is a special case of the stochastic convex-concave saddle point algorithm, the standard SMD could indeed achieve near-optimal sample complexity. Besides, the authors also apply online learning techniques to reduce the number of samples required in each iteration. Finally, the authors consider a complicated scenario where the number of samples that can be drawn from each target distribution could be different. Then the authors consider a weighted GDRO formulation and the convergence analysis of optimization algorithm.", "questions": "N/A", "limitations": "The authors properly discussed relate work regarding optimization algorithms for DRO. I think one reference could be added: Wang et. al, Sinkhorn Distributionally Robust Optimization, arXiv preprint arXiv:2109.11926, June 2023. \n\nReason: The authors listed some progress on solving Wasserstein DRO. The reference (Wang et. al, 2023) provides a SMD algorithm with biased gradient oracles that solves entropic-regularized DRO problem up to $\\delta$-optimality gap with complexity $\\tilde{O}(\\delta^{-2})$. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "- Results are presented in a clear way. The proof is also fluent. It should be accepted.\n- SMD and the online learning fashion of SMD are used in a cleaver way such that the authors achieved near-optimal sample complexity, and the per-iteration sample complexity cost is reduced (the per-iteraiton complexity is important because it relates to the storage cost of an optimization algorithm). \n- The notion of weighted GDRO is also of research interest. ", "weaknesses": "- No numerical study is presented to demonstrate the convergence rate of proposed algorithm.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Many thanks for the constructive reviews!\n\n---\n\nQ1: No numerical study is presented to demonstrate the convergence rate of proposed algorithm\n\nA1: We have carried out preliminary experiments, the results of which can be found in the global response to all reviewers. These empirical findings align closely with our theories.\n\n---\nQ2: I think one reference could be added …\n\nA2: Thank you for identifying the missing reference. We will read and cite this paper.\n"}]

Author rebuttal:
We have conducted preliminary experiments to evaluate our algorithms. Below, we present the experimental setting and empirical results. All related figures can be found in the accompanying PDF.

## Datasets 
We first use a **synthetic** data from Soma et al. [2022]. We set the group number $m=20$, and for each group $i \in [m]$, we draw a model $\mathbf{w}_i^* \in \mathbb{R}^{1000}$ from the uniform distribution over the unit sphere. For distribution $\mathcal{P}_i$, the sample $(\mathbf{x},y)$ is generated with $\mathbf{x} \sim \mathcal{N}(0, I)$ and $y = \mathrm{sign}(\mathbf{x}^\top \mathbf{w}_i^*)$ with probability $0.9$ and $y = -\mathrm{sign}(\mathbf{x}^\top \mathbf{w}_i^*)$ with probability $0.1$.

We also use the **Adult** dataset, which includes attributes such as age, gender, race, and educational background of $48842$ individuals. Following Soma et al. [2022], we construct $m=6$ groups based on the race and gender attributes, where each group represents a combination of {black, white, others} with {female, male}.  We divide this dataset into training and testing subsets, using a 3:1 ratio. 

## GDRO with balanced data
For the synthetic data, we will generate the random sample on the fly, according to the protocol above. For the Adult dataset, we will randomly draw samples from each group in the training set. In other words, $\mathcal{P_i}$ is defined as the empirical distribution over the $i$-th group in the training set. 

We refer to the method of Sagawa et al. [2020] and our Algorithm 1 by SMD($1$) and SMD($m$) to underscore that they are instances of SMD with $1$ sample and $m$ samples in each iteration, respectively. We denote our Algorithm 2 as Online($1$) to emphasize that it is based on techniques from online learning and uses $1$ sample per iteration. We set $\ell(\cdot;\cdot)$ to be the logistic loss and utilize different methods to train a linear model. In the balanced setting, we use the max risk in (2) as the performance measure. To evaulate the risk value, we will draw a substantial number of samples, and use the empirical average to approximate the expectation.

We first report the max risk with respect to the number of iterations in Fig. 1. We observe that SMD($m$) is faster than Online($1$), which in turn outperforms SMD($1$). This observation is consistent with our theories, since their convergence rates are $O(\sqrt{(\log m)/T})$, $O(\sqrt{m (\log m)/T})$, and $O(m\sqrt{(\log m)/T})$, respectively. Next, we plot the max risk against the number of samples consumed by each algorithm in Fig. 2. As can be seen, the curves of SMD($m$) and Online($1$) are very close, indicating that they share the same sample complexity, i.e., $O(m (\log m)/\epsilon^2) $. By contrast, SMD($1$) has a much higher sample complexity, i.e., $O(m^2 (\log m)/\epsilon^2)$.

## GDRO with imbalanced Data
For the synthetic data, we set the number of samples as $n_i = 800*(21-i)$. Notice that the Adult dataset is inherently imbalanced. So, we can visit each sample in the training set *once* to simulate the imbalanced setting, and the numbers of samples in $6$ groups are $18012, 7920, 1428, 1388, 908, 486$.

Similarly, we label the Baseline mentioned in the first paragraph of Section 3 as SMD($m$). We designate our Algorithms 3 and 4 as SMD$\_\mathrm{r}$ and SMPA$\_\mathrm{m}$ to highlight that the former combines SMD with random sampling and the latter one integrates SMPA and mini-batches. In the imbalanced setting, we will examine how the risk on each individual distribution decreases with respect to the number of iterations. Recall that the total number of iterations of SMD($m$), SMD$\_\mathrm{r}$ and SMPA$\_\mathrm{m}$ are $n_m$, $n_1$ and $n_m/2$, respectively.

Due to the limitation of space, we only present the risk on the first, middle and last distributions (i.e., $\mathcal{P}\_1$, $\mathcal{P}\_{m/2}$ and $\mathcal{P}\_m$) in Fig. 3 and Fig. 4. First, we observe that our SMPA$\_\mathrm{m}$ is faster than both SMD($m$) and SMD$\_\mathrm{r}$ across all distributions, and finally attains the lower risk is most cases. This behavior aligns with our Theorem 4, which reveals that SMPA$\_\mathrm{m}$ achieves a nearly optimal rate of $O((\log m)/\sqrt{n_i})$ for all distributions $\mathcal{P}\_i$, after $n_m/2$ iterations. We also note that on distribution $\mathcal{P}\_1$, although SMD$\_\mathrm{r}$ converges slowly, its final risk is the lowest. This phenomenon is again in accordance with our Theorem 3, which shows that the risk of SMD$\_\mathrm{r}$ on $\mathcal{P}\_1$ reduces at a nearly optimal $O(\sqrt{(\log m)/n_1})$ rate, after $n_1$ iterations. From Fig.3(c) and Fig.4(c), we can see that the final risk of SMD($m$) on the last distribution $\mathcal{P}\_m$ matches that of SMPA$\_\mathrm{m}$. This outcome is anticipated, as they exhibit similar convergence rates of $O(\sqrt{(\log m)/n_m})$ and $O((\log m)/\sqrt{n_m})$, respectively.


Meta-review: This paper makes significant contributions to the field of optimization, particularly in group distributionally robust optimization (GDRO). It introduces novel algorithms with near-optimal sample complexity, matching lower bounds within logarithmic factors. These contributions incorporate online learning techniques, address non-oblivious scenarios, and introduce a novel weighted GDRO formulation, enhancing the versatility and efficiency of GDRO approaches. These results have the potential to extend the utility of GDRO algorithms, making them applicable in a wider range of scenarios.

The reviews highlight several key concerns. One recurring issue is the lack of empirical experiments to validate the proposed algorithms, which could provide essential insights into their practical performance; while the authors have already added preliminary experiments, a more comprehensive experimental section would strengthen the paper. Reviewers also questioned the novelty of certain algorithmic aspects in comparison to prior work, particularly in the context of non-oblivious online learning. Furthermore, concerns were raised about the practicality of obtaining samples from all distributions in certain scenarios and the potential challenges posed by large mini-batch sizes in one of the proposed algorithms. 

Overall, despite these concerns, the reviewers generally acknowledge the paper's technical soundness, clear presentation, and significant contributions to robust optimization.  I recommend to accept the paper, given the authors’ commitment to adequately address the remaining concerns and suggestions from the reviewers in their final version.
Decision: Accept (poster)