SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lijun Zhang $^{1,2}$ , Peng Zhao $^{1}$ , Zhen-Hua Zhuang $^{1}$ , Tianbao Yang $^{3}$ , Zhi-Hua Zhou $^{1}$
$^{2}$ Peng Cheng Laboratory, Shenzhen 518055, China
$^{3}$ Department of Computer Science and Engineering, Texas A&M University, College Station, USA {zhanglj, zhaop, zhuangzh, zhouzh}@lamda.nju.edu.cn, tianbao-yang@amu.edu
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 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 \geq n_2 \geq \dots \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.
In the classical statistical machine learning, our goal is to minimize the risk with respect to a fixed distribution  $\mathcal{P}_0$  [Vapnik, 2000], i.e.,
$$
\min  _ {\mathbf {w} \in \mathcal {W}} \left\{R _ {0} (\mathbf {w}) = \operatorname {E} _ {\mathbf {z} \sim \mathcal {P} _ {0}} \left[ \ell (\mathbf {w}; \mathbf {z}) \right] \right\}, \tag {1}
$$
where  $\mathbf{z} \in \mathcal{Z}$  is a sample drawn from  $\mathcal{P}_0$ ,  $\mathcal{W}$  denotes a hypothesis class, and  $\ell(\mathbf{w}; \mathbf{z})$  is a loss measuring the prediction error of model  $\mathbf{w}$  on  $\mathbf{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  $\mathcal{P}_0$ , and in SA, we directly solve the original problem by using stochastic observations of the objective  $R_0(\cdot)$ .
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  $\mathcal{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:
$$
\min  _ {\mathbf {w} \in \mathcal {W}} \max  _ {i \in [ m ]} \left\{R _ {i} (\mathbf {w}) = \operatorname {E} _ {\mathbf {z} \sim \mathcal {P} _ {i}} \left[ \ell (\mathbf {w}; \mathbf {z}) \right] \right\} \tag {2}
$$
where  $\mathcal{P}_1, \ldots, \mathcal{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:
$$
\min  _ {\mathbf {w} \in \mathcal {W}} \max  _ {\mathbf {q} \in \Delta_ {m}} \left\{\phi (\mathbf {w}, \mathbf {q}) = \sum_ {i = 1} ^ {m} q _ {i} R _ {i} (\mathbf {w}) \right\} \tag {3}
$$
where  $\Delta_{m} = \{\mathbf{q}\in \mathbb{R}^{m}:\mathbf{q}\geq 0,\sum_{i = 1}^{m}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}(\cdot)$  and its gradient, and then update both  $\mathbf{w}$  and  $\mathbf{q}$  by SMD. The proposed method achieves an  $O(\sqrt{(\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) / \epsilon^2)$  sample complexity for finding an  $\epsilon$ -optimal solution of (3), which matches the  $\Omega(m / \epsilon^2)$  lower bound [Soma et al., 2022, Theorem 5] up to a logarithmic factor, and tighter than the  $O(m^2 (\log m) / \epsilon^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  $\mathbf{w}$ , and Exp3-IX, an algorithm for non-oblivious multi-armed bandits (MAB) [Neu, 2015], with stochastic rewards to update  $\mathbf{q}$ . In this way, our algorithm only needs 1 sample in each round and attains an  $O(\sqrt{m(\log m) / T})$  convergence rate, implying the same  $O(m(\log m) / \epsilon^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 \geq n_2 \geq \dots \geq 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  $\mathcal{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}(\cdot)$  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  $\mathcal{P}_i$  reduces at an  $O(\sqrt{n_1\log m} /n_i)$  rate, and for  $\mathcal{P}_1$ , it becomes  $O(\sqrt{(\log m) / n_1})$ , which almost matches the optimal  $O(\sqrt{1 / n_1})$  rate of learning from a single distribution with  $n_1$  samples.
On the other hand, for distribution  $\mathcal{P}_i$  with budget  $n_i < n_1$ , the above  $O(\sqrt{n_1 \log m} / n_i)$  rate is worse than the  $O(\sqrt{1 / n_i})$  rate obtained by learning from  $\mathcal{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  $\bar{n} \leq n_m$  rounds, and in each iteration, draws a mini-batch of  $n_i / \bar{n}$  samples from every distribution  $\mathcal{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((\frac{1}{n_m} + \frac{1}{\sqrt{n_i}}) \log m)$  rate for each  $\mathcal{P}_i$ . Thus, we obtain a nearly optimal  $O((\log m) / \sqrt{\bar{n}_i})$  rate for distributions  $\mathcal{P}_i$  with  $n_i \leq n_m^2$ , and an  $O((\log m) / n_m)$  rate otherwise. Note that the latter rate is as expected since the algorithm only updates  $\hat{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) / \epsilon^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).
First, we state the general setup of mirror descent [Nemirovski et al., 2009]. We equip the domain  $\mathcal{W}$  with a distance-generating function  $\nu_{w}(\cdot)$ , which is 1-strongly convex with respect to certain norm  $\| \cdot \|_w$ . We define the Bregman distance associated with  $\nu_{w}(\cdot)$  as  $B_{w}(\mathbf{u},\mathbf{v}) = \nu_{w}(\mathbf{u}) - [\nu_{w}(\mathbf{v}) + \langle \nabla \nu_{w}(\mathbf{v}),\mathbf{u} - \mathbf{v}\rangle]$ . For the simplex  $\Delta_m$ , we choose the entropy function  $\nu_{q}(\mathbf{q}) = \sum_{i=1}^{m} q_{i} \ln q_{i}$  which is 1-strongly convex with respect to the vector  $\ell_1$ -norm  $\| \cdot \|_1$ , as the distance-generating function. Similarly,  $B_{q}(\cdot, \cdot)$  is the Bregman distance associated with  $\nu_{q}(\cdot)$ .
Then, we introduce the standard assumptions about the domain, and the loss function.
Assumption 1 The domain  $\mathcal{W}$  is convex and its diameter measured by  $\nu_{w}(\cdot)$  is bounded by  $D$ , i.e.,
$$
\max  _ {\mathbf {w} \in \mathcal {W}} \nu_ {w} (\mathbf {w}) - \min  _ {\mathbf {w} \in \mathcal {W}} \nu_ {w} (\mathbf {w}) \leq D ^ {2}. \tag {4}
$$
For  $\Delta_{m}$ , it is easy to verify that its diameter measured by the entropy function is bounded by  $\sqrt{\ln m}$ .
Assumption 2 For all  $i\in [m]$ , the risk function  $R_{i}(\mathbf{w}) = \operatorname{E}_{\mathbf{z}\sim \mathcal{P}_{i}}[\ell (\mathbf{w};\mathbf{z})]$  is convex.
Assumption 3 For all  $i\in [m]$  , we have
$$
0 \leq \ell (\mathbf {w}; \mathbf {z}) \leq 1, \forall \mathbf {w} \in \mathcal {W}, \mathbf {z} \sim \mathcal {P} _ {i}. \tag {5}
$$
Assumption 4 For all  $i \in [m]$ , we have
$$
\left\| \nabla \ell (\mathbf {w}; \mathbf {z}) \right\| _ {w, *} \leq G, \forall \mathbf {w} \in \mathcal {W}, \mathbf {z} \sim \mathcal {P} _ {i} \tag {6}
$$
where  $\| \cdot \|_{w,*}$  is the dual norm of  $\| \cdot \|_w$ .
Algorithm 1 Stochastic Mirror Descent for GDRO
Input: Two step sizes:  $\eta_w$  and  $\eta_q$
1: Initialize  $\mathbf{w}_1 = \mathrm{argmin}_{\mathbf{w} \in \mathcal{W}} \nu_w(\mathbf{w})$ , and  $\mathbf{q}_1 = [1/m, \dots, 1/m]^\top \in \mathbb{R}^m$
2: for  $t = 1$  to  $T$  do
3: For each  $i \in [m]$ , draw a sample  $\mathbf{z}_t^{(i)}$  from distribution  $\mathcal{P}_i$
4: Construct the stochastic gradients defined in (9)
5: Update  $\mathbf{w}_t$  and  $\mathbf{q}_t$  according to (11) and (12), respectively
6: end for
7: return  $\bar{\mathbf{w}} = \frac{1}{T}\sum_{t=1}^{T}\mathbf{w}_t$  and  $\bar{\mathbf{q}} = \frac{1}{T}\sum_{t=1}^{T}\mathbf{q}_t$
Last, we discuss the performance measure. To analyze the convergence property, we measure the quality of an approximate solution  $(\bar{\mathbf{w}},\bar{\mathbf{q}})$  to (3) by the error
$$
\epsilon_ {\phi} (\bar {\mathbf {w}}, \bar {\mathbf {q}}) = \max  _ {\mathbf {q} \in \Delta_ {m}} \phi (\bar {\mathbf {w}}, \mathbf {q}) - \min  _ {\mathbf {w} \in \mathcal {W}} \phi (\mathbf {w}, \bar {\mathbf {q}}) \tag {7}
$$
which directly controls the optimality of  $\bar{\mathbf{w}}$  to the original problem (2), since
$$
\max  _ {i \in [ m ]} R _ {i} (\bar {\mathbf {w}}) - \min  _ {\mathbf {w} \in \mathcal {W}} \max  _ {i \in [ m ]} R _ {i} (\mathbf {w}) \leq \max  _ {\mathbf {q} \in \Delta_ {m}} \sum_ {i = 1} ^ {m} q _ {i} R _ {i} (\bar {\mathbf {w}}) - \min  _ {\mathbf {w} \in \mathcal {W}} \sum_ {i = 1} ^ {m} \bar {q} _ {i} R _ {i} (\mathbf {w}) = \epsilon_ {\phi} (\bar {\mathbf {w}}, \bar {\mathbf {q}}). \tag {8}
$$
2.2 Stochastic Mirror Descent for GDRO
To apply SMD, the key is to construct stochastic gradients of the function  $\phi (\mathbf{w},\mathbf{q})$  in (3). In each round  $t$ , denote by  $\mathbf{w}_t$  and  $\mathbf{q}_t$  the current solutions. We draw one sample  $\mathbf{z}_t^{(i)}$  from every distribution  $\mathcal{P}_i$ , and define stochastic gradients as
$$
\mathbf {g} _ {w} \left(\mathbf {w} _ {t}, \mathbf {q} _ {t}\right) = \sum_ {i = 1} ^ {m} q _ {t, i} \nabla \ell \left(\mathbf {w} _ {t}; \mathbf {z} _ {t} ^ {(i)}\right), \text {a n d} \mathbf {g} _ {q} \left(\mathbf {w} _ {t}, \mathbf {q} _ {t}\right) = \left[ \ell \left(\mathbf {w} _ {t}; \mathbf {z} _ {t} ^ {(1)}\right), \dots , \ell \left(\mathbf {w} _ {t}; \mathbf {z} _ {t} ^ {(m)}\right) \right] ^ {\top}. \tag {9}
$$
It is worth mentioning that the construction of  $\mathbf{g}_w(\mathbf{w}_t,\mathbf{q}_t)$  can be further simplified to
$$
\tilde {\mathbf {g}} _ {w} \left(\mathbf {w} _ {t}, \mathbf {q} _ {t}\right) = \nabla \ell \left(\mathbf {w} _ {t}; \mathbf {z} _ {t} ^ {(i _ {t})}\right) \tag {10}
$$
where  $i_t \in [m]$  is drawn randomly according to the probability  $\mathbf{q}_t$ .
Then, we utilize SMD to update  $\mathbf{w}_t$  and  $\mathbf{q}_t$ :
$$
\mathbf {w} _ {t + 1} = \underset {\mathbf {w} \in \mathcal {W}} {\operatorname {a r g m i n}} \left\{\eta_ {w} \left\langle \mathbf {g} _ {w} \left(\mathbf {w} _ {t}, \mathbf {q} _ {t}\right), \mathbf {w} - \mathbf {w} _ {t} \right\rangle + B _ {w} \left(\mathbf {w}, \mathbf {w} _ {t}\right) \right\}, \tag {11}
$$
$$
\mathbf {q} _ {t + 1} = \underset {\mathbf {q} \in \Delta_ {m}} {\operatorname {a r g m i n}} \left\{\eta_ {q} \left\langle - \mathbf {g} _ {q} \left(\mathbf {w} _ {t}, \mathbf {q} _ {t}\right), \mathbf {q} - \mathbf {q} _ {t} \right\rangle + B _ {q} (\mathbf {q}, \mathbf {q} _ {t}) \right\} \tag {12}
$$
where  $\eta_w > 0$  and  $\eta_q > 0$  are two step sizes that will be determined later. The updating rule of  $\mathbf{w}_t$  depends on the choice of the distance-generating function  $\nu_w(\cdot)$ . Since  $B_q(\mathbf{q},\mathbf{q}_t)$  is defined in terms of the negative entropy, (12) is equivalent to
$$
q _ {t + 1, i} = \frac {q _ {t , i} \exp \left(\eta_ {q} \ell \left(\mathbf {w} _ {t} ; \mathbf {z} _ {t} ^ {(i)}\right)\right)}{\sum_ {j = 1} ^ {m} q _ {t , j} \exp \left(\eta_ {q} \ell \left(\mathbf {w} _ {t} ; \mathbf {z} _ {t} ^ {(j)}\right)\right)}, \forall i \in [ m ] \tag {13}
$$
which is the Hedge algorithm [Freund and Schapire, 1997] applied to a maximization problem. In the beginning, we set  $\mathbf{w}_1 = \operatorname*{argmin}_{\mathbf{w} \in \mathcal{W}} \nu_w(\mathbf{w})$ , and  $\mathbf{q}_1 = \frac{1}{m} \mathbf{1}_m$ , where  $\mathbf{1}_m$  is the  $m$ -dimensional vector consisting of  $1$ 's. In the last step, we return the averaged iterates  $\bar{\mathbf{w}} = \frac{1}{T} \sum_{t=1}^{T} \mathbf{w}_t$  and  $\bar{\mathbf{q}} = \frac{1}{T} \sum_{t=1}^{T} \mathbf{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  $\eta_w = D^2\sqrt{\frac{8}{5T(D^2G^2 + \ln m)}}$  and  $\eta_q = (\ln m)\sqrt{\frac{8}{5T(D^2G^2 + \ln m)}}$  in Algorithm 1, with probability at least  $1 - \delta$ , we have
$$
\epsilon_ {\phi} (\bar {\mathbf {w}}, \bar {\mathbf {q}}) \leq \left(8 + 2 \ln \frac {2}{\delta}\right) \sqrt {\frac {1 0 (D ^ {2} G ^ {2} + \ln m)}{T}}.
$$
Remark 1 Theorem 1 shows that Algorithm 1 achieves an  $O(\sqrt{(\log m) / T})$  convergence rate. Since it consumes  $m$  samples per iteration, the sample complexity is  $O(m(\log m) / \epsilon^2)$ , which nearly matches the  $\Omega (m / \epsilon^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 \in [m]$  uniformly in each round  $t$ , and draw 1 sample  $\mathbf{z}_t^{(i_t)}$  from  $\mathcal{P}_{i_t}$ . The stochastic gradients are constructed as follows:
$$
\hat {\mathbf {g}} _ {w} \left(\mathbf {w} _ {t}, \mathbf {q} _ {t}\right) = q _ {t, i _ {t}} m \nabla \ell \left(\mathbf {w} _ {t}; \mathbf {z} _ {t} ^ {(i _ {t})}\right), \text {a n d} \hat {\mathbf {g}} _ {q} \left(\mathbf {w} _ {t}, \mathbf {q} _ {t}\right) = [ 0, \dots , m \ell \left(\mathbf {w} _ {t}; \mathbf {z} _ {t} ^ {(i _ {t})}\right), \dots , 0 ] ^ {\top} \tag {14}
$$
where  $\hat{\mathbf{g}}_q(\mathbf{w}_t,\mathbf{q}_t)$  is a vector with  $m\ell (\mathbf{w}_t;\mathbf{z}_t^{(i_t)})$  in position  $i_t$  and 0 elsewhere. Then, the two stochastic gradients are used to update  $\mathbf{w}_t$  and  $\mathbf{q}_t$ , in the same way as (11) and (12). However, it only attains a slow convergence rate:  $O(m\sqrt{(\log m) / T})$ , leading to an  $O(m^2 (\log m) / \epsilon^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  $\nabla_{\mathbf{w}}\phi (\mathbf{w}_t,\mathbf{q}_t)$ , they construct the stochastic gradient  $\tilde{\mathbf{g}}_w(\mathbf{w}_t,\mathbf{q}_t)$  in the same way as (10), which needs 1 sample. To approximate  $\nabla_{\mathbf{q}}\phi (\mathbf{w}_t,\mathbf{q}_t)$ , they draw  $m$  samples  $\mathbf{z}_{\tau}^{(1)},\ldots ,\mathbf{z}_{\tau}^{(m)}$ , one from each distribution, at round  $\tau = mk + 1$ ,  $k\in \mathbb{Z}$ , and reuse them for  $m$  iterations to construct
$$
\mathbf {g} _ {q} ^ {\prime} \left(\mathbf {w} _ {t}, \mathbf {q} _ {t}\right) = \left[ \ell \left(\mathbf {w} _ {t}; \mathbf {z} _ {\tau} ^ {(1)}\right), \dots , \ell \left(\mathbf {w} _ {t}; \mathbf {z} _ {\tau} ^ {(m)}\right) \right] ^ {\top}, t = \tau , \dots , \tau + m - 1. \tag {15}
$$
Then, they treat  $\tilde{\mathbf{g}}_w(\mathbf{w}_t,\mathbf{q}_t)$  and  $\mathbf{g}_q^{\prime}(\mathbf{w}_t,\mathbf{q}_t)$  as stochastic gradients, and update  $\mathbf{w}_t$  and  $\mathbf{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  $\mathbf{g}_q(\mathbf{w}_t,\mathbf{q}_t)$  at rounds  $t = \tau +2,\dots ,\tau +m - 1$  making their analysis ungrounded. To see this, from the updating rule of SMD, we know that  $\mathbf{w}_{\tau +2}$  depends on  $\mathbf{q}_{\tau +1}$ , which in turn depends on the  $m$  samples drawn at round  $\tau$ , and thus
$$
\mathrm {E} \left[ \ell \left(\mathbf {w} _ {\tau + 2}; \mathbf {z} _ {\tau} ^ {(i)}\right) \right] \neq R _ {i} \left(\mathbf {w} _ {\tau + 2}\right), i = 1, \dots , m.
$$
Remark 2 As shown in (10), we can use 1 sample to construct a stochastic gradient for  $\mathbf{w}_t$  with small norm, since  $\| \tilde{\mathbf{g}}_w(\mathbf{w}_t,\mathbf{q}_t)\|_{w,*}\leq G$  under Assumption 4. Thus, it is relatively easy to control the error related to  $\mathbf{w}_t$ . However, we do not have such guarantees for the stochastic gradient of  $\mathbf{q}_t$ . Recall that the infinity norm of  $\hat{\mathbf{g}}_q(\mathbf{w}_t,\mathbf{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
$$
\phi (\mathbf {w}, \mathbf {q} _ {1}) = \sum_ {i = 1} ^ {m} q _ {1, i} R _ {i} (\mathbf {w}), \phi (\mathbf {w}, \mathbf {q} _ {2}) = \sum_ {i = 1} ^ {m} q _ {2, i} R _ {i} (\mathbf {w}), \dots , \phi (\mathbf {w}, \mathbf {q} _ {T}) = \sum_ {i = 1} ^ {m} q _ {T, i} R _ {i} (\mathbf {w})
$$
subject to the constraint  $\mathbf{w} \in \mathcal{W}$ , and the 2nd player to maximize a sequence of linear functions
$$
\phi (\mathbf {w} _ {1}, \mathbf {q}) = \sum_ {i = 1} ^ {m} q _ {i} R _ {i} (\mathbf {w} _ {1}), \phi (\mathbf {w} _ {2}, \mathbf {q}) = \sum_ {i = 1} ^ {m} q _ {i} R _ {i} (\mathbf {w} _ {2}), \dots , \phi (\mathbf {w} _ {T}, \mathbf {q}) = \sum_ {i = 1} ^ {m} q _ {i} R _ {i} (\mathbf {w} _ {T})
$$