Title: Stochastic Approximation Approaches to Group Distributionally Robust Optimization
Venue: NeurIPS 2023
Accepted: yes
Keywords: Group distributionally robust optimization

Stochastic mirror descent

Non-oblivious online learning

Sample complexity

Stochastic mirror-prox algorithm

Mini-batch
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

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 ) /ϵ 2 ) sample complexity for finding an ϵ -optimal solution, which matches the Ω( m/ϵ 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 multiarmed 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 ≥ n 2 ≥ · · · ≥ 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 ( √ 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 ) / √ n i ) convergence rate, which almost matches the optimal O ( √ 1 /n i ) rate of only learning from the i -th distribution with n i samples.

1 Introduction

In the classical statistical machine learning, our goal is to minimize the risk with respect to a fixed distribution P 0 [Vapnik, 2000], i.e.,

where z ∈ Z is a sample drawn from P 0 , W denotes a hypothesis class, and ℓ ( w ; z ) is a loss measuring the prediction error of model w on z . During the past decades, various algorithms have been developed to optimize (1), and can 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 a set of samples drawn from P 0 , and in SA, we directly solve the original problem by using stochastic observations of the objective R 0 ( · ) .

However, a model that trained over a single distribution may lack robustness in the sense that (i) it 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 neighborhood of P 0 [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 a finite number of distributions [Oren et al., 2019, Sagawa et al., 2020]. Mathematically, GDRO can be formulated as a minimax stochastic problem:

where P 1 , . . . , P m denote m distributions. A 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 a stochastic convex-concave saddle-point problem:

where ∆ m = { q ∈ R m : q ≥ 0 , ∑ m i =1 q i = 1 } is the ( m -1 )-dimensional simplex, and then solve (3) by their mirror descent stochastic approximation method, namely stochastic mirror descent (SMD). In 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 et al., 2022, Soma et al., 2022], or hold only in expectation [Carmon and Hausler, 2022].

As a starting point, we first provide a routine application of SMD to (3), and discuss the theoretical guarantees. In each iteration, we draw 1 sample from every distribution to construct unbiased estimators of R i ( · ) and its gradient, and then update both w and q by SMD. The proposed method achieves an O ( √ (log m ) /T ) convergence rate in expectation and with high probability, where T is the number of iterations. As a result, we obtain an O ( m (log m ) /ϵ 2 ) sample complexity for finding an ϵ -optimal solution of (3), which matches the Ω( m/ϵ 2 ) lower bound [Soma et al., 2022, Theorem 5] up to a logarithmic factor, and tighter than the O ( m 2 (log m ) /ϵ 2 ) bound of Sagawa 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 m 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-oblivious 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-armed bandits (MAB) [Neu, 2015], with stochastic rewards to update q . In this way, our algorithm only needs 1 sample in each round and attains an O ( √ m (log m ) /T ) convergence rate, implying the same O ( m (log m ) /ϵ 2 ) 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 n i be the sample budget of the i -th distribution, and without loss of generality, we assume that n 1 ≥ n 2 ≥ · · · ≥ n m . Now, the goal is not to attain the 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.

In each iteration of the first approach, we draw 1 sample from every P i with probability n i /n 1 , and then construct stochastic gradients to perform mirror descent. In this way, the budget will be satisfied in expectation after n 1 rounds. To analyze its performance, we propose a novel formulation of weighted GDRO, which weights each risk R i ( · ) in (3) by a scale factor p i . Then, our algorithm can be regarded as SMD for an instance of weighted GDRO. With the help of scale factors, we

demonstrate that the proposed algorithm enjoys distribution-dependent convergence in the sense that it converges faster for distributions with more samples. In particular, the excess risk on distribution P i reduces at an O ( √ n 1 log m/n i ) rate, and for P 1 , it becomes O ( √ (log m ) /n 1 ) , which almost matches the optimal O ( √ 1 /n 1 ) rate of learning from a single distribution with n 1 samples.

On the other hand, for distribution P i with budget n i < n 1 , the above O ( √ n 1 log m/n i ) rate is worse than the O ( √ 1 /n i ) rate obtained by learning from P i 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 ≤ n m rounds, and in each iteration, draws a mini-batch of n i / ¯ n samples from every distribution P i . 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, to update our solutions, and again make use of the weighted GDRO formulation to obtain distribution-wise convergence rates. Theoretical analysis shows that the excess risk converges at an O (( 1 n m + 1 √ n i ) log m ) rate for each P i . Thus, we obtain a nearly optimal O ((log m ) / √ n i ) rate for distributions P i with n i ≤ n 2 m , and an O ((log m ) /n m ) rate otherwise. Note that the latter rate is as expected since the algorithm only updates O ( n m ) times.

Related work We briefly discuss the related works on the GDRO problem in (2)/(3) and will review the traditional DRO problem in Appendix A. Sagawa et al. [2020] have applied SMD [Nemirovski et al., 2009] to (3), but only obtain a 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 O ( m (log m ) /ϵ 2 ) 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 [2022] 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 distance-generating function ν w ( · ) , which is 1 -strongly convex with respect to certain norm ∥ · ∥ w . We define the Bregman distance associated with ν w ( · ) as B w ( u , v ) = ν w ( u ) -[ ν w ( v ) + ⟨∇ ν w ( v ) , u -v ⟩ ] . For the simplex ∆ m , we choose the entropy function ν q ( q ) = ∑ m i =1 q i ln q i , which is 1 -strongly convex with respect to the vector ℓ 1 -norm ∥ · ∥ 1 , as the distance-generating function. Similarly, B q ( · , · ) is the Bregman distance associated with ν q ( · ) .

Then, we introduce the standard assumptions about the domain, and the loss function.

Assumption 1 The domain W is convex and its diameter measured by ν w ( · ) is bounded by D , i.e.,

For ∆ m , it is easy to verify that its diameter measured by the entropy function is bounded by √ ln m .

Assumption 2 For all i ∈ [ m ] , the risk function R i ( w ) = E z ∼P i [ ℓ ( w ; z )] is convex.

Assumption 3 For all i ∈ [ m ] , we have

Assumption 4 For all i ∈ [ m ] , we have

where ∥ · ∥ w, ∗ is the dual norm of ∥ · ∥ w .

Algorithm 1 Stochastic Mirror Descent for GDRO

Input : Two step sizes: η w and η q

1: Initialize w 1 = argmin w ∈W ν w ( w ) , and q 1 = [1 /m,... , 1 /m ] ⊤ ∈ R m

3: For each i ∈ [ m ] , draw a sample z ( i ) t from distribution P i

2: for t = 1 to T do

4: Construct the stochastic gradients defined in (9)

5: Update w t and q t according to (11) and (12), respectively

6: end for

7: return ¯ w = 1 T ∑ T t =1 w t and ¯ q = 1 T ∑ T t =1 q t

Last, we discuss the performance measure. To analyze the convergence property, we measure the quality of an approximate solution ( ¯ w , ¯ q ) to (3) by the error

which directly controls the optimality of ¯ w to the original problem (2), since

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 round t , denote by w t and q t the current solutions. We draw one sample z ( i ) t from every distribution P i , and define stochastic gradients as

It is worth mentioning that the construction of g w ( w t , q t ) can be further simplified to

where i t ∈ [ m ] is drawn randomly according to the probability q t .

Then, we utilize SMD to update w t and q t :

where η w > 0 and η q > 0 are two step sizes that will be determined later. The updating rule of w t depends on the choice of the distance-generating function ν w ( · ) . Since B q ( q , q t ) is defined in terms of the negative entropy, (12) is equivalent to

which is the Hedge algorithm [Freund and Schapire, 1997] applied to a maximization problem. In the beginning, we set w 1 = argmin w ∈W ν w ( w ) , and q 1 = 1 m 1 m , where 1 m is the m -dimensional vector consisting of 1 's. In the last step, we return the averaged iterates ¯ w = 1 T ∑ T t =1 w t and ¯ q = 1 T ∑ T t =1 q t as final solutions. The completed procedure is summarized in Algorithm 1.

Based on the theoretical guarantee of SMD [Nemirovski et al., 2009, §3.1], we have the following theorem.

Theorem 1 Under Assumptions 1, 2, 3 and 4, and setting η w = D 2 √ 8 5 T ( D 2 G 2 +ln m ) and η q =

(ln m ) √ 8 5 T ( D 2 G 2 +ln m ) in Algorithm 1, with probability at least 1 -δ , we have .

Remark 1 Theorem 1 shows that Algorithm 1 achieves an O ( √ (log m ) /T ) convergence rate. Since it consumes m samples per iteration, the sample complexity is O ( m (log m ) /ϵ 2 ) , which nearly matches the Ω( m/ϵ 2 ) lower bound [Soma et al., 2022, Theorem 5]. Due to space limitations, we defer all the proofs to Appendix B, and omit expectation bounds.

Comparisons with Sagawa et al. [2020] In their stochastic algorithm, Sagawa et al. [2020] generate a random index i t ∈ [ m ] uniformly in each round t , and draw 1 sample z ( i t ) t from P i t . The stochastic gradients are constructed as follows:

where ˆ g q ( w t , q t ) is a vector with mℓ ( w t ; z ( i t ) t ) in position i t and 0 elsewhere. Then, the two stochastic gradients are used to update w t and q t , in the same way as (11) and (12). However, it only attains a slow convergence rate: O ( m √ (log m ) /T ) , leading to an O ( m 2 (log m ) /ϵ 2 ) sample complexity, which is higher than that of Algorithm 1 by a factor of m . The slow convergence is due to the fact that the optimization error depends on the dual norm of the stochastic gradients in (14), which blows up by a factor of m , compared with the gradients in (9).

Comparisons with Haghtalab et al. [2022] To reduce the number of samples required in each round, Haghtalab et al. [2022] propose to reuse samples for multiple iterations. To approximate ∇ w ϕ ( w t , q t ) , they construct the stochastic gradient ˜ g w ( w t , q t ) in the same way as (10), which needs 1 sample. To approximate ∇ q ϕ ( w t , q t ) , they draw m samples z (1) τ , . . . , z ( m ) τ , one from each distribution, at round τ = mk +1 , k ∈ Z , and reuse them for m iterations to construct

Then, they treat ˜ g w ( w t , q t ) and g ′ q ( w t , q t ) as stochastic gradients, and update w t and q t by SMD. In this way, their algorithm uses 2 samples on average in each iteration. However, the gradient in (15) is no longer an unbiased estimator of the true gradient g q ( w t , q t ) at rounds t = τ +2 , . . . , τ + m -1 , making their analysis ungrounded. To see this, from the updating rule of SMD, we know that w τ +2 depends on q τ +1 , which in turn depends on the m samples drawn at round τ , and thus

̸

Remark 2 As shown in (10), we can use 1 sample to construct a stochastic gradient for w t with small norm, since ∥ ˜ g w ( w t , q t ) ∥ w, ∗ ≤ G under Assumption 4. Thus, it is relatively easy to control the error related to w t . However, we do not have such guarantees for the stochastic gradient of q t . Recall that the infinity norm of ˆ g q ( w t , q t ) in (14) is upper bounded by m . The reason is that we insist on the unbiasedness of the stochastic gradient, which leads to a large variance. In the next section, we borrow techniques from online learning to better balance the bias and the variance.

2.3 Non-oblivious Online Learning for GDRO

In the studies of convex-concave saddle-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 to bound the optimization error. To solve problem (3), we ask the 1st player to minimize a sequence of convex functions

subject to the constraint w ∈ W , and the 2nd player to maximize a sequence of linear functions

subject to the constraint q ∈ ∆ m . We highlight that there exists an important difference between our stochastic convex-concave problem and its deterministic counterpart. Here, the two players cannot

Input : Two step sizes: η w and η q ; IX coefficient γ

1: Initialize w 1 = argmin w ∈W ν w ( w ) , and q 1 = [1 /m,... , 1 /m ] ⊤ ∈ R m

2: Generate i t ∈ [ m ] according to q t , and then draw a sample z ( i t ) t from q i t

2: for t = 1 to T do

3: Construct the stochastic gradient in (10) and the IX loss estimator in (17)

4: Update w t and q t according to (16) and (18), respectively

5: end for

6: return ¯ w = 1 T ∑ T t =1 w t and ¯ q = 1 T ∑ T t =1 q t

directly observe the loss function, and can only approximate R i ( w ) = E z ∼P i [ ℓ ( w ; z ) ] by drawing samples from P i . The stochastic setting makes the problem more challenging, and in particular, we need to take care of the non-oblivious 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 can easily obtain a stochastic gradient with small norm by using 1 sample. So, we model the problem faced by the 1st player as 'non-oblivious online convex optimization (OCO) with stochastic gradients', and still use SMD to update its solution. In each round t , with 1 sample drawn from P i , the 2nd player can estimate the value of R i ( w t ) which is the coefficient of q i . Since the 2nd player is maximizing a linear function over the simplex, the problem can be modeled as 'nonoblivious multi-armed bandits (MAB) with stochastic rewards'. And fortunately, we have powerful online algorithms for non-oblivious MAB [Auer et 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 its theoretical guarantee to stochastic rewards.

The complete procedure is presented in Algorithm 2, and we explain key steps below. In each round t , we generate an index i t ∈ [ m ] from the probability distribution q t , and then draw a sample z ( i t ) t from the distribution q i t . With the stochastic gradient in (10), we use SMD to update w t :

Then, we reuse the sample z ( i t ) t to update q t according to Exp3-IX, which first constructs the Implicit-eXploration (IX) loss estimator [Kocák et al., 2014]:

where γ > 0 is the IX coefficient and I [ A ] equals to 1 when the event A is true and 0 otherwise, and then performs a mirror descent update:

We present the theoretical guarantee of Algorithm 2.

Theorem 2 Under Assumptions 1, 2, 3 and 4, and setting η w = 2 D G √ 5 T , η q = √ ln m mT and γ = η q 2 in Algorithm 2, with probability at least 1 -δ , we have

Remark 3 The above theorem shows that with 1 sample per iteration, Algorithm 2 is able to achieve an O ( √ m (log m ) /T ) convergence rate, thus maintaining the O ( m (log m ) /ϵ 2 ) sample complexity.

Comparisons with Soma et al. [2022] In a recent work, Soma et al. [2022] have deployed online algorithms to optimize w 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 q that are independent from ¯ w and ¯ q [Soma et al., 2022, Theorem 3],

However, (20) cannot be used to bound ϵ ϕ ( ¯ w , ¯ q ) in (7), because of the dependency issue. To be more clear, we have ϵ ϕ ( ¯ w , ¯ q ) = max q ∈ ∆ m ϕ ( ¯ w , q ) -min w ∈W ϕ ( w , ¯ q ) = ϕ ( ¯ w , ̂ q ) -ϕ ( ̂ w , ¯ q ) , where ̂ w = argmin w ∈W ϕ ( w , ¯ q ) and ̂ q = argmax q ∈ ∆ m ϕ ( ¯ w , q ) depend on ¯ w and ¯ q .

Remark 4 After we pointed out the issue of reusing samples, Haghtalab et 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 et al. [2022], and subsequently refined by us.

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., 2020], 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 i the number of samples that can be drawn from P i . Without loss of generality, we assume that n 1 ≥ n 2 ≥ · · · ≥ n m . Note that we have a straightforward Baseline which just runs Algorithm 1 for n m iterations, and the optimization error ϵ ϕ ( ¯ w , ¯ q ) = O ( √ (log m ) /n m ) .

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 t , we first generate a set of Bernoulli random variables { b (1) t , . . . , b ( m ) t } with Pr[ b ( i ) t = 1] = p i to determine whether to sample from each distribution. If b ( i ) t = 1 , we draw a sample z ( i ) t from P i . Now, the question is how to construct stochastic gradients from those samples. Let C t = { i | b ( i ) t = 1 } be the indices of selected distributions. If we stick to the original problem in (3), then the stochastic gradients should be constructed in the following way

to ensure unbiasedness. Then, they can be used by SMD to update w t and q t . After n 1 iterations, the expected number of samples drawn from P i will be n 1 p i = n i , and thus the budget is satisfied in expectation. To analyze the optimization error, we need to bound the norm of stochastic gradients in (21). To this end, we have ∥ g w ( w t , q t ) ∥ w, ∗ ≤ Gn 1 /n m and ∥ g q ( w t , q t ) i ∥ ∞ ≤ n 1 /n m . Following the arguments of Theorem 1, we can prove that the error ϵ ϕ ( ¯ w , ¯ q ) = O ( √ (log m ) /n 1 · n 1 /n m ) = O ( √ n 1 log m/n m ) , which is even larger than the O ( √ (log m ) /n m ) 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 (21) is caused by the inverse probability 1 /p i . Anatural idea is to ignore 1 /p i , and define the following stochastic gradients:

In this way, they are no longer stochastic gradients of (3), but can be treated as stochastic gradients of a weighted GDRO problem:

Algorithm 3 Stochastic Mirror Descent for Weighted GDRO

Input

: Two step sizes: η w and η q

1: Initialize w 1 = argmin w ∈W ν w ( w ) , and q 1 = [1 /m,... , 1 /m ] ⊤ ∈ R m

2: for t = 1 to n 1 do

3: For each i ∈ [ m ] , generate a Bernoulli random variable b ( i ) t with Pr[ b ( i ) t = 1] = p i , and if b ( i ) t = 1 , draw a sample z ( i ) t from distribution P i

4: Construct the stochastic gradients defined in (22)

5: Update w t and q t according to (11) and (12), respectively

6: end for

7: return ¯ w = 1 n 1 ∑ n 1 t =1 w t and ¯ q = 1 n 1 ∑ n 1 t =1 q t

where each risk R i ( · ) is scaled by a factor p i . Based on the gradients in (22), we still use (11) and (12) to update w t and q t . We summarize the complete procedure in Algorithm 3.

We omit the optimization error of Algorithm 3 for (23), since it has exactly the same form as Theorem 1. 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 1, 2, 3 and 4, and setting η w = D 2 √ 8 5 n 1 ( D 2 G 2 +ln m ) and η q =

(ln m ) √ 8 5 n 1 ( D 2 G 2 +ln m ) in Algorithm 3, with probability at least 1 -δ , we have

where p ∗ φ is the optimal value of (23) and µ ( δ ) = 8 + 2 ln 2 δ .

Remark 5 We see that Algorithm 3 exhibits a distribution-dependent convergence behavior: The larger the number of samples n i , the smaller the target risk n 1 p ∗ φ /n i , and the faster the convergence rate O ( √ n 1 log m/n i ) . Note that its rate is always better than the O ( √ n 1 log m/n m ) rate of SMD with (21) as gradients. Furthermore, it converges faster than the Baseline when n i ≥ √ n 1 n m . In particular, for distribution P 1 , Algorithm 3 attains an O ( √ (log m ) /n 1 ) rate, which almost matches the optimal O ( √ 1 /n 1 ) 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-batches

In Algorithm 3, distributions with more samples take their advantage by appearing more frequently in the stochastic gradients. In this section, we propose a 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 a small number of iterations ¯ n that is no larger than n m . Then, in each iteration, we draw a mini-batch of n i / ¯ n samples from every distribution P i . For P i with more samples, we can estimate the associated risk R i ( · ) and its gradient more accurately, i.e., 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, we resort to a more advanced SA approach-stochastic mirror-prox algorithm (SMPA), whose convergence rate depends on the variance [Juditsky et al., 2011]. To overcome the second challenge, we again introduce scale factors into the optimization problem and the stochastic gradients.

In SMPA, we need to maintain two sets of solutions: ( w t , q t ) and ( w ′ t , q ′ t ) . In each round t , we first draw n i /n m samples from every distribution P i , denoted by z ( i,j ) t , j = 1 , . . . , n i /n m . Then, we use them to construct stochastic gradients at ( w ′ t , q ′ t ) of a weighted GDRO problem (23), where the

Algorithm 4 Stochastic Mirror-Prox Algorithm for Weighted GDRO

Input

: Two step sizes: η w and η q

1: Initialize w ′ 1 = argmin w ∈W ν w ( w ) , and q ′ 1 = [1 /m,... , 1 /m ] ⊤ ∈ R m

2: for t = 1 to n m / 2 do

3: For each i ∈ [ m ] , draw n i /n m samples { z ( i,j ) t : j = 1 , . . . , n i /n m } from distribution P i

4: Construct the stochastic gradients defined in (24)

5: Calculate w t +1 and q t +1 according to (25) and (26), respectively

6: For each i ∈ [ m ] , draw n i /n m samples { ˆ z ( i,j ) t : j = 1 , . . . , n i /n m } from distribution P i

7: Construct the stochastic gradients defined in (27)

8: Calculate w ′ t +1 and q ′ t +1 according to (28) and (29), respectively

9: end for

10: return ¯ w = 2 n m ∑ 1+ n m / 2 t =2 w t and ¯ q = 2 n m ∑ 1+ n m / 2 t =2 q t

value of p i will be determined later. Specifically, we define

Let's take the stochastic gradient g q ( w ′ t , q ′ t ) , whose variance will be measured in terms of ∥ · ∥ ∞ , as an example to explain the intuition of inserting p i . Define u i = n m n i ∑ n i /n m j =1 ℓ ( w ′ t ; z ( i,j ) t ) . With a larger mini-batch size n i /n m , u i will approximate R i ( w ′ t ) more accurately, and thus have a smaller variance. Then, it allows us to insert a larger value of p i , without affecting the variance of ∥ g q ( w ′ t , q ′ t ) ∥ ∞ , since ∥ · ∥ ∞ is insensitive to perturbations of small elements. Similar to the case in Theorem 3, the convergence rate of R i ( · ) depends on 1 /p i , and becomes faster if p i is larger.

Based on (24), we use SMD to update ( w ′ t , q ′ t ) , and denote the solution by ( w t +1 , q t +1 ) :

Next, we draw another n i /n m samples from each distribution P i , denoted by ˆ z ( i,j ) t , j = 1 , . . . , n i /n m , to construct stochastic gradients at ( w t +1 , q t +1 ) :

Then, we use them to update ( w ′ t , q ′ t ) again, and denote the result by ( w ′ t +1 , q ′ t +1 ) :

To meet the budget constraints, we repeat the above process for n m / 2 iterations. Finally, we return ¯ w = 2 n m ∑ 1+ n m / 2 t =2 w t and ¯ q = 2 n m ∑ 1+ n m / 2 t =2 q t as solutions. The completed procedure is summarized in Algorithm 4.

To analysis the performance of Algorithm 4, we further assume the risk function R i ( · ) is smooth, and the dual norm ∥ · ∥ w, ∗ satisfies a regularity condition.

Assumption 5 All the risk functions are L -smooth, i.e.,

Note that even in the studies of stochastic convex optimization (SCO), smoothness is necessary to obtain a variance-based convergence rate [Lan, 2012].

Assumption 6 The dual norm ∥ · ∥ w, ∗ is κ -regular for some small constant κ ≥ 1 .

The regularity condition is used when analyzing the effect of mini-batches on stochastic gradients. For a formal definition, please refer to Juditsky and Nemirovski [2008]. Assumption 6 is satisfied by most of papular norms considered in the literature, such as the vector ℓ p -norm and the infinity norm.

Then, we have the following theorem for Algorithm 4.

Theorem 4 Define

where c > 0 is an absolute constant. Under Assumptions 1, 2, 3, 4, 5 and 6, and setting

in Algorithm 4, with probability at least 1 -δ , we have

where p ∗ φ is the optimal value of (23).

Furthermore, by setting p i as with high probability, we have

Remark 6 Compared with Algorithm 3, Algorithm 4 has two advantages: (i) the budget constraint is satisfied exactly; (ii) we obtain a faster O ((log m ) / √ n i ) rate for all distributions P i such that n i ≤ n 2 m , which is much better than the O ( √ n 1 ln m/n i ) rate of Algorithm 3, and the O ( √ (log m ) /n m ) rate of the Baseline. For distributions with a larger number of budget, i.e., n i > n 2 m , it maintains a fast O ((log m ) /n m ) rate. Since it only updates n m times, and the best we can expect is the O (1 /n m ) rate of deterministic settings [Nemirovski, 2004]. So, there is a 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 et al., 2011, Zhang et 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 O ( m (log m ) /ϵ 2 ) sample complexity, but with different number of samples used in each round, which are m and 1 respectively. Then, we formulate a weighted GDRO problem to handle the scenario in which different distributions have different sample budgets. We first incorporate non-uniform sampling into SMD to satisfy the sample budget in expectation, and deliver distribution-dependent convergence rates. Then, we propose to use mini-batches to meet the 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 are presented in Appendix C, and align closely with our theories.

Acknowledgments and Disclosure of Funding

This work was partially supported by the National Key R&D Program of China (2022ZD0114801), NSFC (62122037, 61921006), National Postdoctoral Program for Innovative Talent, China Postdoctoral Science Foundation (2023M731597), and the major key project of PCL (PCL2021A12).

References

Agarwal and T. Zhang. Minimax regret optimization for robust machine learning under distribution shift. In Proceedings of 35th Conference on Learning Theory , pages 2704-2729, 2022.

Amodei, S. 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 on Machine Learning , pages 173-182, 2016.

Auer, N. Cesa-Bianchi, Y. Freund, and R. E. Schapire. The nonstochastic multiarmed bandit problem. SIAM Journal on Computing , 32(1):48-77, 2002.

Becker and R. Kohavi. Adult. UCI Machine Learning Repository, 1996. DOI: https://doi.org/10.24432/C5XW20.

Ben-Tal, L. E. Ghaoui, and A. Nemirovski. Robust Optimization . Princeton University Press, 2009.

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.

Ben-Tal, E. Hazan, T. Koren, and S. Mannor. Oracle-based robust optimization via online learning. Operations Research , 63(3):628-638, 2015.

Bertsimas, V. Gupta, and N. Kallus. Robust sample average approximation. Mathematical Programming , 171:217-282, 2018.

Blum, N. Haghtalab, A. D. Procaccia, and M. Qiao. Collaborative PAC learning. In Advances in Neural Information Processing Systems 30 , pages 2389-2398, 2017.

Bubeck and N. Cesa-Bianchi. Regret analysis of stochastic and nonstochastic multi-armed bandit problems. Foundations and Trends in Machine Learning , 5(1):1-122, 2012.

Carmon and D. Hausler. Distributionally robust optimization via ball oracle acceleration. In Advances in Neural Information Processing Systems 35 , pages 35866-35879, 2022.

Cesa-Bianchi and G. Lugosi. Prediction, Learning, and Games . Cambridge University Press, 2006.

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.

Delage and Y. Ye. Distributionally robust optimization under moment uncertainty with application to data-driven problems. Operations Research , 58(3):595-612, 2010.

C. Duchi and H. Namkoong. Learning models with uniform performance via distributionally robust optimization. The Annals of Statistics , 49(3):1378 - 1406, 2021.

C. Duchi, P. W. Glynn, and H. Namkoong. Statistics of robust optimization: A generalized empirical likelihood approach. Mathematics of Operations Research , 46(3):946-969, 2021.

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.

D. Flaxman, A. T. Kalai, and H. B. McMahan. Online convex optimization in the bandit setting: Gradient descent without a gradient. In Proceedings of the 16th Annual ACM-SIAM Symposium on Discrete Algorithms , pages 385-394, 2005.

Freund and R. E. Schapire. A decision-theoretic generalization of on-line learning and an application to boosting. Journal of Computer and System Sciences , 55(1):119-139, 1997.

Freund and R. E. Schapire. Adaptive game playing using multiplicative weights. Games and Economic Behavior , 29(1):79-103, 1999.

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.

Haghtalab, M. I. Jordan, and E. Zhao. On-demand sampling: Learning optimally from multiple distributions. ArXiv e-prints , arXiv:2210.12529v2, 2023.

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.

Hu, G. Niu, I. Sato, and M. Sugiyama. Does distributionally robust supervised learning give robust classifiers? In Proceedings of the 35th International Conference on Machine Learning , volume 80, pages 2029-2037, 2018.

Jin, B. Zhang, H. Wang, and L. Wang. Non-convex distributionally robust optimization: Nonasymptotic analysis. In Advances in Neural Information Processing Systems 34 , pages 2771-2782, 2021.

Juditsky, A. Nemirovski, and C. Tauvel. Solving variational inequalities with stochastic mirror-prox algorithm. Stochastic Systems , 1(1):17-58, 2011.

B. Juditsky and A. S. Nemirovski. Large deviations of vector-valued martingales in 2-smooth normed spaces. ArXiv e-prints , arXiv:0809.0813, 2008.

Kocák, G. Neu, M. Valko, and R. Munos. Efficient learning by implicit exploration in bandit problems with side observations. In Advances in Neural Information Processing Systems 27 , pages 613-621, 2014.

Kuhn, P. M. Esfahani, V. A. Nguyen, and S. Shafieezadeh-Abadeh. Wasserstein distributionally robust optimization: Theory and applications in machine learning. Operations Research & Management Science in the Age of Analytics , pages 130-166, 2019.

J. Kushner and G. G. Yin. Stochastic Approximation and Recursive Algorithms and Applications . Springer, second edition, 2003.

Lan. An optimal method for stochastic composite optimization. Mathematical Programming , 133: 365-397, 2012.

Lattimore and C. Szepesvári. Bandit Algorithms . Cambridge University Press, 2020.

Levy, Y. Carmon, J. C. Duchi, and A. Sidford. Large-scale methods for distributionally robust optimization. In Advances in Neural Information Processing Systems 33 , pages 8847-8860, 2020.

Mohri, G. Sivek, and A. T. Suresh. Agnostic federated learning. In Proceedings of the 36th International Conference on Machine Learning , pages 4615-4625, 2019.

Namkoong and J. C. Duchi. Stochastic gradient methods for distributionally robust optimization with f -divergences. In Advances in Neural Information Processing Systems 29 , pages 2216-2224, 2016.

Namkoong and J. C. Duchi. Variance-based regularization with convex objectives. In Advances in Neural Information Processing Systems 30 , pages 2971-2980, 2017.

Nemirovski. Prox-method with rate of convergence O (1 /t ) for variational inequalities with lipschitz continuous monotone operators and smooth convex-concave saddle point problems. SIAM Journal on Optimization , 15(1):229-251, 2004.

Nemirovski, A. Juditsky, G. Lan, and A. Shapiro. Robust stochastic approximation approach to stochastic programming. SIAM Journal on Optimization , 19(4):1574-1609, 2009.

Neu. Explore no more: Improved high-probability regret bounds for non-stochastic bandits. In Advances in Neural Information Processing Systems 28 , pages 3168-3176, 2015.

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.13213(v5), 2019.

Oren, S. Sagawa, T. B. Hashimoto, and P. Liang. Distributionally robust language modeling. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing , pages 4227-4237, 2019.

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 in Neural Information Processing Systems 34 , pages 10067-10080, 2021.

Radivojac, N. V. Chawla, A. K. Dunker, and Z. Obradovic. Classification and knowledge discovery in protein databases. Journal of Biomedical Informatics , 37(4):224-239, 2004.

Rafique, M. Liu, Q. Lin, and T. Yang. Weakly-convex-concave min-max optimization: Provable algorithms and applications in machine learning. Optimization Methods and Software , 37(3): 1087-1121, 2022.

Rakhlin and K. Sridharan. Optimization, learning, and games with predictable sequences. In Advances in Neural Information Processing Systems 26 , pages 3066-3074, 2013.

N. Rothblum and G. Yona. Multi-group agnostic PAC learnability. In Proceedings of the 38th International Conference on Machine Learning , pages 9107-9115, 2021.

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 P. Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. In International Conference on Learning Representations , 2020.

Samuel and G. Chechik. Distributional robustness loss for long-tail learning. In Proceedings of the 2021 IEEE/CVF International Conference on Computer Vision , pages 9475-9484, 2021.

Scarf. A min-max solution of an inventory problem. Studies in the Mathematical Theory of Inventory and Production , pages 201-209, 1958.

Shapiro. Distributionally robust stochastic programming. SIAM Journal on Optimization , 27(4): 2258-2275, 2017.

Sinha, H. Namkoong, and J. Duchi. Certifying some distributional robustness with principled adversarial training. In International Conference on Learning Representations , 2018.

Słowik and L. Bottou. On distributionally robust optimization and data rebalancing. In Proceedings of the 25th International Conference on Artificial Intelligence and Statistics , pages 1283-1297, 2022.

Soma, K. Gatmiry, and S. Jegelka. Optimal algorithms for group distributionally robust optimization and beyond. ArXiv e-prints , arXiv:2212.13669, 2022.

Syrgkanis, A. Agarwal, H. Luo, and R. E. Schapire. Fast convergence of regularized learning in games. In Advances in Neural Information Processing Systems 28 , pages 2989-2997, 2015.

N. Vapnik. The Nature of Statistical Learning Theory . Springer, second edition, 2000.

Vershynin. High-Dimensional 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.11926, 2021.

Xu, C. Dan, J. Khim, and P. Ravikumar. Class-weighted classification: Trade-offs and robust approaches. In Proceedings of the 37th International Conference on Machine Learning , pages 10544-10554, 2020.

Zhang, A. K. Menon, A. Veit, S. Bhojanapalli, S. Kumar, and S. Sra. Coping with label shift via distributionally robust optimisation. In International Conference on Learning Representations , 2021.

Zhang, T. Yang, R. Jin, and X. He. O (log T ) projections for stochastic optimization of smooth and strongly convex functions. In Proceedings of the 30th International Conference on Machine Learning , pages 1121-1129, 2013.

Z.-H. Zhou. A theoretical perspective of machine learning with computational resource concerns. ArXiv e-prints , arXiv:2305.02217, 2023.

A Related Work

Distributionally robust optimization (DRO) stems from the pioneering work of Scarf [1958], and has gained a lot of interest with the advancement of robust optimization [Ben-Tal 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, 2021], label shift [Zhang et al., 2021], 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.,

where S ( P 0 ) denotes a set of probability distributions around P 0 . In the literature, there mainly exist three ways to construct S ( P 0 ) : (i) enforcing moment constraints [Delage and Ye, 2010], (ii) defining a neighborhood around P 0 by a 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., 2021], and (iii) hypothesis testing of goodness-of-fit [Bertsimas et al., 2018].

By drawing a set of samples from P 0 , 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-Read 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 (1). Namkoong and Duchi [2017] show that empirical DRO with the χ 2 -divergence has the effect of variance regularization, leading to better generalization w.r.t. distribution P 0 . Later, Duchi et al. [2021] demonstrate similar behaviors for the f -divergence constrained neighborhood, and provide one- and two-sided confidence intervals for the minimal risk in (1). Based on the Wasserstein distance, Esfahani and Kuhn [2018] establish an upper confidence bound on the risk of the empirical solution.

Since (33) is more complex than (1), considerable research efforts were devoted to develop efficient algorithms for DRO and its empirical version. For P 0 with finite support, Ben-Tal et al. [2013, Corollary 3] have demonstrated that (33) with f -divergences is equivalent to a convex optimization problem, provided that the loss ℓ ( w ; z ) is convex in w . Actually, this conclusion is true even when P 0 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-in some cases 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 [Namkoong and Duchi, 2016], and DRO with convex loss [Levy et al., 2020] and non-convex losses [Jin et al., 2021, Qi et al., 2021, Rafique et al., 2022].

The main focus of this paper is the GDRO problem in (2)/(3), instead of the traditional DRO in (33). Sagawa et al. [2020] have applied SMD [Nemirovski et al., 2009] to (3), but only obtain a sub-optimal sample complexity of O ( m 2 (log m ) /ϵ 2 ) , 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 O ( m (log m ) /ϵ 2 ) 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 [2022] propose a variant of GDRO named as minimax regret optimization (MRO), which replaces the risk R i ( w ) with 'excess risk' R i ( w ) -min w ∈W R i ( w ) . More generally, we can introduce calibration terms in DRO to prevent any single distribution to dominate the maximum [Słowik and Bottou, 2022].

Finally, we note that GDRO has a similar spirit with collaborative PAC learning [Blum et al., 2017, Nguyen and Zakynthinou, 2018, Rothblum and Yona, 2021] in the sense that both aim to find a single model that performs well on multiple distributions.

B Analysis

In this section, we present proofs of main theorems.

B.1 Proof of Theorem 1

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 w and q as a single variable [ w ; q ] ∈ W × ∆ m , and redefine the norm and the distance-generating function [Nemirovski et al., 2009, §3.1]. Let E be the space that W lies in. We equip the Cartesian product E × R m with the following norm and dual norm:

We use the notation x = [ w ; q ] , and equip the set W× ∆ m with the distance-generating function

It is easy to verify that ν ( x ) is 1 -strongly convex w.r.t. the norm ∥ · ∥ . Let B ( · , · ) be the Bregman distance associated with ν ( · ) :

where x ′ = [ w ′ ; q ′ ] .

Then, we consider the following version of SMD for updating x t :

where η > 0 is the step size. In the beginning, we set x 1 = argmin x ∈W× ∆ m ν ( x ) = [ w 1 ; q 1 ] . From the decomposition of the Bregman distance in (36), we observe that (37) is equivalent to (11) and (12) by setting

Next, we show that the stochastic gradients are well-bounded. Under our assumptions, we have

As a result, the concatenated gradients used in (37) is also bounded in term of the dual norm ∥ · ∥ ∗ :

Now, we are ready to state our theoretical guarantees. By setting

Lemma 3.1 and (3.13) of Nemirovski et al. [2009] imply that

Furthermore, from Proposition 3.2 of Nemirovski et al. [2009], we have, for any Ω > 1

We complete the proof by setting δ = 2exp( -Ω) .

B.2 Proof of Theorem 2

We first bound the regret of the 1st player. In the analysis, we address the non-obliviousness by the 'ghost iterate' technique of Nemirovski et al. [2009].

Theorem 5 Under Assumptions 1, 2 and 4, by setting η w = 2 D G √ 5 T , we have

and with probability at least 1 -δ ,

By extending Exp3-IX to stochastic rewards, we have the following bound for the 2nd player.

Theorem 6 Under Assumption 3, by setting η q = √ ln m mT and the IX coefficient γ = η q 2 , we have

and with probability at least 1 -δ ,

From Jensen's inequality and the outputs ¯ w = 1 T ∑ T t =1 w t and ¯ q = 1 T ∑ T t =1 q t , we have

We obtain (19) by substituting the high probability bound in Theorems 5 and 6 into (38), and taking the union bound.

B.3 Proof of Theorem 5

Our goal is to analyze SMD for non-oblivious OCO with stochastic gradients. In the 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 of SGD for non-oblivious OCO. But it is insufficient for our purpose, so we provide our proof by following the analysis of SMD for stochastic convex-concave

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

Summing the above inequality over t = 1 , . . . , T , we have

where the last step is due to [Nemirovski et al., 2009, (2.42)]

By Jensen's inequality, we have

Maximizing each side over w ∈ W , we arrive at

Next, we bound the last term in (41), i.e., max w ∈W F ( w ) . Because E t -1 [˜ g w ( w t , q t )] = ∇ w ϕ ( w t , q t ) , F ( w ) is the sum of a martingale difference sequence for any fixed w . However, it is not true for ˜ w = argmax w ∈W F ( w ) , because ˜ w depends on the randomness of the algorithm. Thus, we cannot directly apply techniques for martingales to bounding max w ∈W 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 dependency [Nemirovski et al., 2009, proof of Lemma 3.1]. Imagine there is an online algorithm which performs SMD by using ∇ w ϕ ( w t , q t ) -˜ g w ( w t , q t ) as the gradient:

where v 1 = w 1 . By repeating the derivation of (39), we can show that

where in the last inequality, we make use of (40) and

Then, we have

From the updating rule of v t in (42), we know that v t is independent from ∇ w ϕ ( w t , q t ) -˜ g w ( w t , q t ) , and thus V 1 , . . . , V T is a martingale difference sequence.

Substituting (45) into (41), we have

Taking expectation over both sides, we have

where we set η w = 2 D G √ 5 T .

To establish a high probability bound, we make use of the Hoeffding-Azuma inequality for martingales stated below [Cesa-Bianchi and Lugosi, 2006].

Lemma 1 Let V 1 , V 2 , . . . be a martingale difference sequence with respect to some sequence X 1 , X 2 , . . . such that V i ∈ [ A i , A i + c i ] for some random variable A i , measurable with respect to X 1 , . . . , X i -1 and a positive constant c i . If S n = ∑ n i =1 V i , then for any t > 0 ,

To apply the above lemma, we need to show that V t is bounded. Indeed, we have

From Lemma 1, with probability at least 1 -δ , we have

We complete the proof by substituting (47) into (46).

B.4 Proof of Theorem 6

Because we can only observe ℓ ( w t , z ( i t ) t ) instead of R i t ( w t ) , the theoretical guarantee of Exp3-IX [Neu, 2015] cannot be directly applied to Algorithm 2. To address this challenge, we generalize the regret analysis of Exp3-IX to stochastic rewards.

By the definition of ϕ ( w , q ) in (3) and the property of linear optimization over the simplex, we have

where j ∗ ∈ argmax j ∈ [ m ] ∑ T t =1 R j ( w t ) and the vector s t ∈ R m is defined as

To facilitate the analysis, we introduce a vector ˆ s t ∈ R m with

where z ( i ) t denotes a random sample drawn from the i -th distribution. Note that ˆ s t is only used for analysis , with the purpose of handling the stochastic rewards. In the algorithm, only ˆ s t,i t = 1 -ℓ ( w t ; z ( i t ) t ) is observed in the t -th iteration.

Following the proof of Theorem 1 of Neu [2015], we have

which makes use of the property of online mirror descent with local norms [Bubeck and Cesa-Bianchi, 2012, Theorem 5.5]. From (5) of Neu [2015], we have

Combining (50) and (51), we have

From (48), we have

We proceed to 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 [Neu, 2015, Lemma 1], which we further generalize to the setting with stochastic rewards.

Lemma 2 Let ξ t,i ∈ [0 , 1] for all t ∈ [ T ] and i ∈ [ m ] , and ˜ ξ t,i be its IX-estimator defined as ˜ ξ t,i = ˆ ξ t,i p t,i + γ t I [ i t = i ] , where ˆ ξ t,i is such that E[ ˆ ξ t,i ] = ξ t,i with ˆ ξ t,i ∈ [0 , 1] and the index i t is sampled from [ m ] according to the distribution p t ∈ ∆ m . Let { γ t } T t =1 be a fixed non-increasing sequence with γ t ≥ 0 and let α t,i be non-negative F t -1 -measurable random variables satisfying α t,i ≤ 2 γ t for all t ∈ [ T ] and i ∈ [ m ] . Then, for any δ > 0 , with probability at least 1 -δ ,

Furthermore, when γ t = γ ≥ 0 for all t ∈ [ T ] , the following holds with probability at least 1 -δ ,

simultaneously for all i ∈ [ m ] .

Notice that our construction of ˜ s t in (17) satisfies that ˜ s t,i = ˆ s t,i q t,i + γ I [ i t = i ] and i t is drawn from [ m ] according to q t ∈ ∆ m as well as E[ˆ s t,i ] = s t,i , which meets the conditions required by Lemma 2. As a result, according to (55), we have

for all j ∈ [ m ] (including j ∗ ) with probability at least 1 -δ .

To bound term B , we can directly use Lemma 1 of Neu [2015], because our setting η q 2 = γ satisfies its requirement. Thus, with probability at least 1 -δ , we have

We now consider term C in (53). Let V t = ⟨ q t , s t ⟩-ˆ s t,i t . Then, it is easy to verify that E t -1 [ V t ] = 0 . So, the process { V t } T t =1 forms a martingale difference sequence and it also satisfies | V t | ≤ 1 for all t . Hence, we can apply Lemma 1 and have

with probability at least 1 -δ .

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 -δ

where the third line holds because of our parameter settings γ = η q 2 and η q = √ ln m mT .

Based on this high probability guarantee, we can then obtain the expected regret upper bound using the formula that

holds for any real-valued random variable X [Bubeck and Cesa-Bianchi, 2012, § 3.2]. In particular, taking

yields E[ X ] ≤ 3 , which implies that

B.5 Proof of Theorem 3

For the stochastic gradients in (22), their norm can be upper bounded in the same way as (9). That is,

So, with exactly the same analysis as Theorem 1, we have

and with probability at least 1 -δ ,

Next, we discuss how to bound the risk of ¯ w on every distribution P i , i.e., R i ( ¯ w ) . Following the derivation in (8), we know

Thus, for every distribution P i , R i ( ¯ w ) can be bounded in the following way:

Taking the high probability bound in (57) as an example, we have with probability at 1 -δ

B.6 Proof of Theorem 4

We first provide some simple facts that will be used later. From Assumption 3, we immediately know that each risk function R i ( · ) also belongs to [0 , 1] . As a result, the difference between each risk function and its estimator is well-bounded, i.e., for all i ∈ [ m ] ,

From Assumption 4, we can prove that each risk function R i ( · ) is G -Lipschitz continuous. To see this, we have

As a result, we have

Furthermore, the difference between the gradient of R i ( · ) and its estimator is also well-bounded, i.e., for all i ∈ [ m ] ,

Recall the definition of the norm ∥ · ∥ and dual norm ∥ · ∥ ∗ for the space E × R m in (34), and the distance-generating function ν ( · ) in (35). Following the arguments in Section B.1, the two updating rules in (25) and (26) can be merged as

where η w = 2 ηD 2 and η q = 2 η ln m . Similarly, (28) and (29) are equivalent to

Let F ([ w ; q ]) be the monotone operator associated with the weighted GDRO problem in (23), i.e.,

From our constructions of stochastic gradients in (24) and (27), we clearly have

Thus, Algorithm 4 is indeed an instance of SMPA [Juditsky et al., 2011, Algorithm 1], and we can use their Theorem 1 and Corollary 1 to bound the optimization error.

Before applying their results, we show that all the preconditions are satisfied. The parameter Ω defined in (16) of Juditsky et al. [2011] can be upper bounded by

Next, we need to demonstrate that F ([ w ; q ]) is continuous.

Lemma 3 For the monotone operator F ([ w ; q ]) , we have

where ˜ L is defined in (31).

We proceed to show the variance of the stochastic gradients satisfies the light tail condition. To this end, we introduce the stochastic oracle used in Algorithm 4:

where

and z ( i,j ) is the j -th sample drawn from distribution P i . The following lemma shows that the variance is indeed sub-Gaussian.

Lemma 4 For the stochastic oracle g ([ w ; q ]) , we have

where σ 2 is defined in (31).

Based on (66), Lemma 3, and Lemma 4, we can apply the theoretical guarantee of SMPA. Recall that the total number of iterations is n m / 2 in Algorithm 4. From Corollary 1 of Juditsky et al. [2011], by setting

we have

for all Λ > 0 . Choosing Λ such that exp( -Λ 2 / 3) ≤ δ/ 2 and exp( -Λ n m / 2) ≤ δ/ 2 , we have with probability at least 1 -δ

Following the derivation of (58), we have

Inspired by Juditsky et al. [2011, §4.3.1], we use the value of p i in (32) to simplify (68). It is easy to verify that

10

2

10

3

10

4

10

5

10

6

# of iterations

0.30

0.40

0.50

0.60

0.70

SMD(1)

SMD(

m

)

Online(1)

(a) The synthetic data

10

2

10

3

# of iterations

0

.

45

0

.

50

0

.

55

0

.

60

0

.

65

0

.

70

SMD(1)

SMD(

m

)

Online(1)

(b) The Adult dataset

10

3

10

4

10

5

10

6

# of samples

0.30

0.40

0.50

0.60

0.70

SMD(1)

SMD(

m

)

Online(1)

(a) The synthetic data

10

2

10

3

# of samples

0

.

45

0

.

50

0

.

55

0

.

60

0

.

65

0

.

70

SMD(1)

SMD(

m

)

Online(1)

(b) The Adult dataset

Substituting (69) and (70) into (68), we have

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 et al., 2022], we use both synthetic and real-world datasets. First, we construct a synthetic data with group number m = 20 . For each group i ∈ [ m ] , we draw a model w ∗ i ∈ R 1000 from the uniform distribution over the unit sphere. For distribution P i , the sample ( x , y ) is generated with x ∼ N (0 , I ) and y = sign( x ⊤ w ∗ i ) with probability 0 . 9 and y = -sign( x ⊤ w ∗ i ) with probability 0 . 1 .

We 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 a combination of {black, white, others} with {female, male}.

10 0

10 1

10 2

10 3

10 4

# of iterations

0.10

0.25

0.40

0.55

0.70

Risk

SMD(

m

)

SMD

r

SMPA

m

(a) Risk on

P

1

10 0

10 1

10 2

10 3

10 4

# of iterations

0.20

0.30

0.40

0.50

0.60

0.70

Risk

SMD(

m

)

SMD

r

SMPA

m

(b) Risk on

P

4

10 0

10 1

10 2

10 3

10 4

# of iterations

0.20

0.30

0.40

0.50

0.60

0.70

Risk

SMD(

m

)

SMD

r

SMPA

m

(c) Risk on

P

6

10 0

10 1

10 2

10 3

10 4

# of iterations

0.20

0.30

0.40

0.50

0.60

0.70

Risk

SMD(

m

)

SMD

r

SMPA

m

(d) Risk on

P

8

10 0

10 1

10 2

10 3

10 4

# of iterations

0.20

0.30

0.40

0.50

0.60

0.70

Risk

SMD(

m

)

SMD

r

SMPA

m

(e) Risk on

P

10

10 0

10 1

10 2

10 3

10 4

# of iterations

0.20

0.30

0.40

0.50

0.60

0.70

Risk

SMD(

m

)

SMD

r

SMPA

m

(f) Risk on

P

12

10 0

10 1

10 2

10 3

10 4

# of iterations

0.20

0.30

0.40

0.50

0.60

0.70

Risk

SMD(

m

)

SMD

r

SMPA

m

(g) Risk on

P

16

10 0

10 1

10 2

10 3

10 4

# of iterations

0.20

0.30

0.40

0.50

0.60

0.70

Risk

SMD(

m

)

SMD

r

SMPA

m

(h) Risk on

P

18

10 0

10 1

10 2

10 3

10 4

# of iterations

0.30

0.40

0.50

0.60

0.70

Risk

SMD(

m

)

SMD

r

SMPA

m

(i) Risk on

P

20

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 i is defined as the empirical distribution over the data in the i -th group.

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 ℓ ( · ; · ) to be the logistic loss and utilize different methods to train a linear model. In the balanced setting, we use the max risk, i.e., max i ∈ [ m ] R i ( w ) , as the performance measure. To estimate 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 ( √ (log m ) /T ) , O ( √ m (log m ) /T ) , and O ( m √ (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 ) /ϵ 2 ) . By contrast, SMD( 1 ) has a much higher sample complexity, i.e., O ( m 2 (log m ) /ϵ 2 ) .

10 0

10 1

10 2

10 3

10 4

# of iterations

0

.

50

0

.

55

0

.

60

0

.

65

0

.

70

Risk

SMD(

m

)

SMD

r

SMPA

m

(a) Risk on

P

1

10 0

10 1

10 2

10 3

10 4

# of iterations

0.30

0.40

0.50

0.60

0.70

Risk

SMD(

m

)

SMD

r

SMPA

m

(b) Risk on

P

2

10 0

10 1

10 2

10 3

10 4

# of iterations

0.40

0.50

0.60

0.70

Risk

SMD(

m

)

SMD

r

SMPA

m

(c) Risk on

P

3

10 0

10 1

10 2

10 3

10 4

# of iterations

0.20

0.30

0.40

0.50

0.60

0.70

Risk

SMD(

m

)

SMD

r

SMPA

m

(d) Risk on

P

4

10 0

10 1

10 2

10 3

10 4

# of iterations

0.50

0.55

0.60

0.65

0.70

Risk

SMD(

m

)

SMD

r

SMPA

m

(e) Risk on

P

5

10 0

10 1

10 2

10 3

10 4

# of iterations

0.30

0.40

0.50

0.60

0.70

Risk

SMD(

m

)

SMD

r

SMPA

m

(f) Risk on

P

6

C.3 GDRO with Imbalanced Data

For experiments on the synthetic dataset, we set the number of samples as n i = 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 3 as SMD( m ). We designate our Algorithms 3 and 4 as SMD r and SMPA 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 r and SMPA m are n m , n 1 and n m / 2 , respectively.

We present the risk on the individual distribution in Fig. 3 and Fig. 4. First, we observe that our SMPA m is faster than both SMD( m ) and SMD r across all distributions, and finally attains the lower risk is most cases. This behavior aligns with our Theorem 4, which reveals that SMPA m achieves a nearly optimal rate of O ((log m ) / √ n i ) for all distributions P i , after n m / 2 iterations. We also note that on distribution P 1 , although SMD r 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 r on P 1 reduces at a nearly optimal O ( √ (log m ) /n 1 ) rate, after n 1 iterations. From Fig. 3(i) and Fig. 4(f), we can see that the final risk of SMD( m ) on the last distribution P m matches that of SMPA m . This outcome is anticipated, as they exhibit similar convergence rates of O ( √ (log m ) /n m ) and O ((log m ) / √ n m ) , respectively.

D Supporting Lemmas

D.1 Proof of Lemma 2

The proof follows the argument of Neu [2015, Proof of Lemma 1], and we generalize it to the setting with stochastic rewards. First, observe that for any i ∈ [ m ] and t ∈ [ T ] ,

where the last step is due to the inequality z 1+ z/ 2 ≤ log(1 + z ) for z ≥ 0 and we introduce the notations β t = 2 γ t and ¯ ξ t,i = ( ˆ ξ t,i /p t,i ) · I [ i t = i ] to simplify the presentation.

Define the notation ˜ λ t = ∑ m i =1 α t,i ˜ ξ t,i and λ t = ∑ m i =1 α t,i ξ t,i . Then, we have

̸

where the second inequality is by the inequality x log(1 + y ) ≤ log(1 + xy ) that holds for all y ≥ -1 and x ∈ [0 , 1] , and the equality E t -1 [ Π i m =1 ( 1 + α t,i ¯ ξ t,i )] = E t -1 [ 1 + ∑ m i =1 α t,i ¯ ξ t,i ] follows from the fact that ¯ ξ t,i · ¯ ξ t,j = 0 holds whenever i = j . The last line is due to E t -1 [ ¯ ξ t,i ] = E t -1 [( ˆ ξ t,i /p t,i ) · I [ i t = i ]] = ξ t,i and the inequality 1 + z ≤ e z for all z ∈ R .

As a result, from (72) we conclude that the process Z t = exp ( ∑ t s =1 ( ˜ λ s -λ s ) ) is a supermartingale. Indeed, E t -1 [ Z t ] = E t -1 [ exp ( ∑ t -1 s =1 ( ˜ λ s -λ s ) ) · exp( ˜ λ t -λ t ) ] ≤ Z t -1 . Thus, we have E[ Z T ] ≤ E[ Z T -1 ≤ . . . ≤ E[ Z 0 ] = 1 . By Markov's inequality,

holds for any ϵ > 0 . By setting exp( -ϵ ) = δ and solving the value, we complete the proof for (54). And the inequality (55) for the scenario γ t = γ can be immediately obtained by setting α t,i = 2 γ and taking the union bound over all i ∈ [ m ] .

D.2 Proof of Lemma 3

From the definition of norms in (34), we have

To bound term A , we have

4 D 2 ∥ ∥ ∥ ∥ ∥ m ∑ i =1 q i p i ∇ R i ( w ) -m ∑ i =1 q ′ i p i ∇ R i ( w ) ∥ ∥ ∥ ∥ ∥ 2 w, ∗ ≤ 4 D 2 ( m ∑ i =1 | q i -q ′ i |∥ p i ∇ R i ( w ) ∥ w, ∗ ) 2 (60) ≤ 4 D 2 ( m ∑ i =1 | q i -q ′ i | p i G ) 2 ≤ 4 D 2 G 2 p 2 max ∥ q -q ′ ∥ 2 1 . where p max is defined in (31). To bound B , we have 2

To bound C , we have

Putting everything together, we have

which implies

where ˜ L is defined in (31).

D.3 Proof of Lemma 4

The light tail condition, required by Juditsky et al. [2011], is essentially the sub-Gaussian condition. To this end, we introduce the following sub-gaussian properties [Vershynin, 2018, Proposition 2.5.2].

Proposition 1 (Sub-gaussian properties) Let X be a random variable. Then the following properties are equivalent; the parameters K i > 0 appearing in these properties differ from each other by at most an absolute constant factor.

(i) The tails of X satisfy

(ii) The moments of X satisfy

(iii) The moment generating function (MGF) of X 2 satisfies

(iv) The MGF of X 2 is bounded at some point, namely

Moreover, if E[ X ] = 0 then properties (i)-(iv) are also equivalent to the following property. (v) The MGF of X satisfies

From the above proposition, we observe that the exact value of those constant K 1 , . . . , K 5 is not important, and it is very tedious to calculate them. So, in the following, we only focus on the order of those constants. To simplify presentations, we use c to denote an absolute constant that is independent of all the essential parameters, and its value may change from line to line.

Since

we proceed to analyze the behavior of ∥∇ w φ ( w , q ) -g w ( w , q ) ∥ 2 w, ∗ and ∥∇ q φ ( w , q ) -g q ( w , q ) ∥ 2 ∞ . To this end, we have the following lemma.

where ω max is defined in (31) and c > 0 is an absolute constant.

From Lemma 5, we have

where the first inequality follows from Jensen's inequality.

D.4 Proof of Lemma 5

To analyze ∥∇ w φ ( w , q ) -g w ( w , q ) ∥ 2 w, ∗ , we first consider the approximation error caused by samples from P i :

Under the regularity condition of ∥ · ∥ w, ∗ in Assumption 6, we have, for any γ ≥ 0 ,

which is a directly consequence of the concentration inequality of vector norms [Juditsky and Nemirovski, 2008, Theorem 2.1.(iii)] and (62). Then, we introduce the following lemma to simplify (74).

where X is nonnegative. Then, we have

From (74) and Lemma 6, we have

which satisfies the Proposition 1.(i). From the equivalence between Proposition 1.(i) and Proposition 1.(iv), we have

Inserting the scaling factor p i , we have

To simplify the notation, we define

By Jensen's inequality, we have

where we use the fact that ∥·∥ w, ∗ , ( · ) 2 and exp( · ) are convex, and the last two functions are increasing in R + .

To analyze ∥∇ q φ ( w , q ) -g q ( w , q ) ∥ 2 ∞ , we consider the approximation error related to P i :

Note that the absolute value | · | is 1 -regular [Juditsky and Nemirovski, 2008]. Following (59) and the derivation of (75), we have

To prove that ∥∇ q φ ( w , q ) -g q ( w , q ) ∥ 2 ∞ is also sub-Gaussian, we need to analyze the effect of the infinity norm. To this end, we develop the following lemma.

Lemma 7 Suppose

Then,

where c > 0 is an absolute constant, and K max = max j ∈ [ m ] K j .

From (76) and Lemma 7, we have

D.5 Proof of Lemma 6

When γ ∈ [0 , 2 α ] , we have

When γ ≥ 2 α , we have

where we use the fact

Thus, we always have

D.6 Proof of Lemma 7

From (77), and the equivalence between Proposition 1.(i) and Proposition 1.(iv), we have

As a result,

Choosing t = √ cK 2 max (ln[2 m ] + γ 2 / 2) , we have

Thus

From the equivalence between Proposition 1.(i) and Proposition 1.(iv), we have

Peer reviews:
comments: 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.

STRENGTHS:
- Results are presented in a clear way. The proof is also fluent. It should be accepted.
- 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). 
- The notion of weighted GDRO is also of research interest.

WEAKNESSES:
- No numerical study is presented to demonstrate the convergence rate of proposed 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. 

Reason: 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})$.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-20

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n1. 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\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nThe authors discuss some limitations of the proposed algorithms.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nUnfortunately, 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.\n\nQUESTIONS:\nThere 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?\n\nLIMITATIONS:\nThe novelty and contribution of this paper, in my understanding, is questionable.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\nThere 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe authors have discussed the limitations of their approaches in their remarks. I do not find any potential negative societal impact of this work.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

comments: SUMMARY:
This paper provides algorithms for group distributionally robust optimization.
For the case in which the sample budget for each distribution is identical,
the proposed algorithm achieves nearly optimal sample complexity.
This paper also considers the sample budgets are different for each distribution and provide an algorithm with non-trivial sample-complexity bounds.

STRENGTHS:
- This paper is the first to consider a practical setting in which the sample budget varies by distribution.
- 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.
- The obtained bounds in this work includes extra $\log m$ factors, unlike ones by [Soma et al., 2022].

QUESTIONS:
- Can we remove extra $\log m$ factors?
- It is claimed that the $O(1/\sqrt{n_i})$ rate is optimal, which is not obvious to me. 
It is understood that for normal (non-distributionally-robust) learning over a single distribution, the optimal rate is $O(1 / \sqrt{n_i})$.
However, the definition of risk indicators employed in this paper (LHS in Theorem 4) differs from the usual.
Therefore, in order to discuss what the optimal rate is, it is necessary to again prove a lower bound on this risk.
- In the paper, one contribution is explained as reducing the number of samples in each iteration.
Though I agree that total sample complexity matters,
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). 
I would appreciate a more detailed explanation of the benefit.

LIMITATIONS:
The limitations and potential negative societal impact are adequately addressed.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-26