SlowGuess/NeurIPS_2023_OCR huggingface dataset:
subject to the constraint  $\mathbf{q} \in \Delta_{m}$ . We highlight that there exists an important difference between our stochastic convex-concave problem and its deterministic counterpart. Here, the two players cannot
Algorithm 2 Non-oblivious Online Learning for GDRO
Input: Two step sizes:  $\eta_w$  and  $\eta_q$ ; IX coefficient  $\gamma$
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
2: Generate  $i_t \in [m]$  according to  $\mathbf{q}_t$ , and then draw a sample  $\mathbf{z}_t^{(i_t)}$  from  $\mathbf{q}_{i_t}$ .
3: Construct the stochastic gradient in (10) and the IX loss estimator in (17)
4: Update  $\mathbf{w}_t$  and  $\mathbf{q}_t$  according to (16) and (18), respectively
5: end for
6: 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$
directly observe the loss function, and can only approximate  $R_{i}(\mathbf{w}) = \mathrm{E}_{\mathbf{z}\sim \mathcal{P}_{i}}[\ell (\mathbf{w};\mathbf{z})]$  by drawing samples from  $\mathcal{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  $\mathcal{P}_i$ , the 2nd player can estimate the value of  $R_{i}(\mathbf{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 "non-oblivious 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 \in [m]$  from the probability distribution  $\mathbf{q}_t$ , and then draw a sample  $\mathbf{z}_t^{(i_t)}$  from the distribution  $\mathbf{q}_{i_t}$ . With the stochastic gradient in (10), we use SMD to update  $\mathbf{w}_t$ :
$$
\mathbf {w} _ {t + 1} = \underset {\mathbf {w} \in \mathcal {W}} {\operatorname {a r g m i n}} \left\{\eta_ {w} \left\langle \tilde {\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 {16}
$$
Then, we reuse the sample  $\mathbf{z}_t^{(it)}$  to update  $\mathbf{q}_t$  according to Exp3-IX, which first constructs the Implicit-eXploration (IX) loss estimator [Kocak et al., 2014]:
$$
\tilde {s} _ {t, i} = \frac {1 - \ell \left(\mathbf {w} _ {t} , \mathbf {z} _ {t} ^ {\left(i _ {t}\right)}\right)}{q _ {t , i} + \gamma} \cdot \mathbb {I} [ i _ {t} = i ], \forall i \in [ m ], \tag {17}
$$
where  $\gamma > 0$  is the IX coefficient and  $\mathbb{I}[A]$  equals to 1 when the event  $A$  is true and 0 otherwise, and then performs a mirror descent update:
$$
\mathbf {q} _ {t + 1} = \underset {\mathbf {q} \in \Delta_ {m}} {\operatorname {a r g m i n}} \left\{\eta_ {q} \left\langle \tilde {\mathbf {s}} _ {t}, \mathbf {q} - \mathbf {q} _ {t} \right\rangle + B _ {q} (\mathbf {q}, \mathbf {q} _ {t}) \right\}. \tag {18}
$$
We present the theoretical guarantee of Algorithm 2.
Theorem 2 Under Assumptions 1, 2, 3 and 4, and setting  $\eta_w = \frac{2D}{G\sqrt{5T}}$ ,  $\eta_q = \sqrt{\frac{\ln m}{mT}}$  and  $\gamma = \frac{\eta_q}{2}$  in Algorithm 2, with probability at least  $1 - \delta$ , we have
$$
\begin{array}{l} \epsilon_ {\phi} (\bar {\mathbf {w}}, \bar {\mathbf {q}}) \\ \leq D G \sqrt {\frac {1}{T}} \left(2 \sqrt {5} + 8 \sqrt {\ln \frac {2}{\delta}}\right) + 3 \sqrt {\frac {m \ln m}{T}} + \sqrt {\frac {1}{2 T}} + \left(\sqrt {\frac {m}{T \ln m}} + \sqrt {\frac {1}{2 T}} + \frac {1}{T}\right) \ln \frac {6}{\delta}. \tag {19} \\ \end{array}
$$
Remark 3 The above theorem shows that with 1 sample per iteration, Algorithm 2 is able to achieve an  $O(\sqrt{m(\log m) / T})$  convergence rate, thus maintaining the  $O(m(\log m) / \epsilon^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  $\mathbf{w}$  and  $\mathbf{q}$  that are independent from  $\bar{\mathbf{w}}$  and  $\bar{\mathbf{q}}$  [Soma et al., 2022, Theorem 3],
$$
\operatorname {E} \left[ \phi (\bar {\mathbf {w}}, \mathbf {q}) - \phi (\mathbf {w}, \bar {\mathbf {q}}) \right] = O \left(\sqrt {\frac {m}{T}}\right). \tag {20}
$$
However, (20) cannot be used to bound  $\epsilon_{\phi}(\bar{\mathbf{w}},\bar{\mathbf{q}})$  in (7), because of the dependency issue. To be more clear, we have  $\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}}) = \phi (\bar{\mathbf{w}},\widehat{\mathbf{q}}) - \phi (\widehat{\mathbf{w}},\bar{\mathbf{q}})$ , where  $\widehat{\mathbf{w}} = \mathrm{argmin}_{\mathbf{w}\in \mathcal{W}}\phi (\mathbf{w},\bar{\mathbf{q}})$  and  $\widehat{\mathbf{q}} = \mathrm{argmax}_{\mathbf{q}\in \Delta_m}\phi (\bar{\mathbf{w}},\mathbf{q})$  depend on  $\bar{\mathbf{w}}$  and  $\bar{\mathbf{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  $\mathcal{P}_i$ . Without loss of generality, we assume that  $n_1 \geq n_2 \geq \dots \geq n_m$ . Note that we have a straightforward Baseline which just runs Algorithm 1 for  $n_m$  iterations, and the optimization error  $\epsilon_{\phi}(\bar{\mathbf{w}}, \bar{\mathbf{q}}) = O(\sqrt{(\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_t^{(1)},\ldots ,b_t^{(m)}\}$  with  $\operatorname*{Pr}[b_t^{(i)} = 1] = p_i$  to determine whether to sample from each distribution. If  $b_{t}^{(i)} = 1$ , we draw a sample  $\mathbf{z}_t^{(i)}$  from  $\mathcal{P}_i$ . Now, the question is how to construct stochastic gradients from those samples. Let  $\mathcal{C}_t = \{i|b_t^{(i)} = 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
$$
\mathbf {g} _ {w} \left(\mathbf {w} _ {t}, \mathbf {q} _ {t}\right) = \sum_ {i \in C _ {t}} \frac {q _ {t , i}}{p _ {i}} \nabla \ell \left(\mathbf {w} _ {t}; \mathbf {z} _ {t} ^ {(i)}\right), \text {a n d} \left[ \mathbf {g} _ {q} \left(\mathbf {w} _ {t}, \mathbf {q} _ {t}\right) \right] _ {i} = \left\{ \begin{array}{l l} \ell \left(\mathbf {w} _ {t}; \mathbf {z} _ {t} ^ {(i)}\right) / p _ {i}, & i \in \mathcal {C} _ {t} \\ 0, & \text {o t h e r w i s e} \end{array} \right. \tag {21}
$$
to ensure unbiasedness. Then, they can be used by SMD to update  $\mathbf{w}_t$  and  $\mathbf{q}_t$ . After  $n_1$  iterations, the expected number of samples drawn from  $\mathcal{P}_i$  will be  $n_1p_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  $\| \mathbf{g}_w(\mathbf{w}_t,\mathbf{q}_t)\|_{w,*}\leq Gn_1 / n_m$  and  $\| \mathbf{g}_q(\mathbf{w}_t,\mathbf{q}_t)_i\|_\infty \leq n_1 / n_m$ . Following the arguments of Theorem 1, we can prove that the error  $\epsilon_{\phi}(\bar{\mathbf{w}},\bar{\mathbf{q}}) = O(\sqrt{(\log m) / n_1}\cdot n_1 / n_m) = O(\sqrt{n_1\log m} /n_m)$ , which is even larger than the  $O(\sqrt{(\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$ . A natural idea is to ignore  $1 / p_i$ , and define the following stochastic gradients:
$$
\mathbf {g} _ {w} \left(\mathbf {w} _ {t}, \mathbf {q} _ {t}\right) = \sum_ {i \in C _ {t}} q _ {t, i} \nabla \ell \left(\mathbf {w} _ {t}; \mathbf {z} _ {t} ^ {(i)}\right), \text {a n d} \left[ \mathbf {g} _ {q} \left(\mathbf {w} _ {t}, \mathbf {q} _ {t}\right) \right] _ {i} = \left\{ \begin{array}{l l} \ell \left(\mathbf {w} _ {t}; \mathbf {z} _ {t} ^ {(i)}\right), & i \in \mathcal {C} _ {t} \\ 0, & \text {o t h e r w i s e .} \end{array} \right. \tag {22}
$$
In this way, they are no longer stochastic gradients of (3), but can be treated as stochastic gradients of a weighted GDRO problem:
$$
\min  _ {\mathbf {w} \in \mathcal {W}} \max  _ {\mathbf {q} \in \Delta_ {m}} \left\{\varphi (\mathbf {w}, \mathbf {q}) = \sum_ {i = 1} ^ {m} q _ {i} p _ {i} R _ {i} (\mathbf {w}) \right\} \tag {23}
$$
Algorithm 3 Stochastic Mirror Descent for Weighted 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  $n_1$  do  
3: For each  $i \in [m]$ , generate a Bernoulli random variable  $b_t^{(i)}$  with  $\operatorname*{Pr}[b_t^{(i)} = 1] = p_i$ , and if  $b_t^{(i)} = 1$ , draw a sample  $\mathbf{z}_t^{(i)}$  from distribution  $\mathcal{P}_i$   
4: Construct the stochastic gradients defined in (22)  
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}{n_1} \sum_{t=1}^{n_1} \mathbf{w}_t$  and  $\bar{\mathbf{q}} = \frac{1}{n_1} \sum_{t=1}^{n_1} \mathbf{q}_t$
where each risk  $R_{i}(\cdot)$  is scaled by a factor  $p_i$ . Based on the gradients in (22), we still use (11) and (12) to update  $\mathbf{w}_t$  and  $\mathbf{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  $\eta_w = D^2\sqrt{\frac{8}{5n_1(D^2G^2 + \ln m)}}$  and  $\eta_q = (\ln m)\sqrt{\frac{8}{5n_1(D^2G^2 + \ln m)}}$  in Algorithm 3, with probability at least  $1 - \delta$ , we have
$$
R _ {i} (\bar {\mathbf {w}}) - \frac {n _ {1}}{n _ {i}} p _ {\varphi} ^ {*} \leq \frac {1}{p _ {i}} \mu (\delta) \sqrt {\frac {1 0 (D ^ {2} G ^ {2} + \ln m)}{n _ {1}}} = \mu (\delta) \frac {\sqrt {1 0 (D ^ {2} G ^ {2} + \ln m) n _ {1}}}{n _ {i}}, \forall i \in [ m ]
$$
where  $p_{\varphi}^{*}$  is the optimal value of (23) and  $\mu (\delta) = 8 + 2\ln \frac{2}{\delta}$ .
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_1p_\varphi^* /n_i$ , and the faster the convergence rate  $O(\sqrt{n_1\log m} /n_i)$ . Note that its rate is always better than the  $O(\sqrt{n_1\log m} /n_m)$  rate of SMD with (21) as gradients. Furthermore, it converges faster than the Baseline when  $n_i \geq \sqrt{n_1n_m}$ . In particular, for distribution  $\mathcal{P}_1$ , Algorithm 3 attains an  $O(\sqrt{(\log m) / n_1})$  rate, which almost matches the optimal  $O(\sqrt{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-Proxy 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  $\bar{n}$  that is no larger than  $n_m$ . Then, in each iteration, we draw a mini-batch of  $n_i / \bar{n}$  samples from every distribution  $\mathcal{P}_i$ . For  $\mathcal{P}_i$  with more samples, we can estimate the associated risk  $R_i(\cdot)$  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:  $(\mathbf{w}_t,\mathbf{q}_t)$  and  $(\mathbf{w}_t',\mathbf{q}_t')$ . In each round  $t$ , we first draw  $n_i / n_m$  samples from every distribution  $\mathcal{P}_i$ , denoted by  $\mathbf{z}_t^{(i,j)}$ ,  $j = 1,\dots ,n_i / n_m$ . Then, we use them to construct stochastic gradients at  $(\mathbf{w}_t',\mathbf{q}_t')$  of a weighted GDRO problem (23), where the
Algorithm 4 Stochastic Mirror-Prox Algorithm for Weighted GDRO
1: Initialize  $\mathbf{w}_1^{\prime} = \mathrm{argmin}_{\mathbf{w}\in \mathcal{W}}\nu_w(\mathbf{w})$ , and  $\mathbf{q}_1^\prime = [1 / m,\dots ,1 / m]^\top \in \mathbb{R}^m$
2: for  $t = 1$  to  $n_m / 2$  do
3: For each  $i\in [m]$ , draw  $n_i / n_m$  samples  $\{\mathbf{z}_t^{(i,j)}:j = 1,\dots ,n_i / n_m\}$  from distribution  $\mathcal{P}_i$
4: Construct the stochastic gradients defined in (24)
5: Calculate  $\mathbf{w}_{t + 1}$  and  $\mathbf{q}_{t + 1}$  according to (25) and (26), respectively
6: For each  $i\in [m]$ , draw  $n_i / n_m$  samples  $\{\hat{\mathbf{z}}_t^{(i,j)}:j = 1,\dots ,n_i / n_m\}$  from distribution  $\mathcal{P}_i$
7: Construct the stochastic gradients defined in (27)
8: Calculate  $\mathbf{w}_{t + 1}^{\prime}$  and  $\mathbf{q}_{t + 1}^{\prime}$  according to (28) and (29), respectively
10: return  $\bar{\mathbf{w}} = \frac{2}{n_m}\sum_{t = 2}^{1 + n_m / 2}\mathbf{w}_t$  and  $\bar{\mathbf{q}} = \frac{2}{n_m}\sum_{t = 2}^{1 + n_m / 2}\mathbf{q}_t$
value of  $p_i$  will be determined later. Specifically, we define
$$
\mathbf {g} _ {w} \left(\mathbf {w} _ {t} ^ {\prime}, \mathbf {q} _ {t} ^ {\prime}\right) = \sum_ {i = 1} ^ {m} q _ {t, i} ^ {\prime} p _ {i} \left(\frac {n _ {m}}{n _ {i}} \sum_ {j = 1} ^ {n _ {i} / n _ {m}} \nabla \ell \left(\mathbf {w} _ {t} ^ {\prime}; \mathbf {z} _ {t} ^ {(i, j)}\right)\right), \tag {24}
$$
$$
\mathbf {g} _ {q} \left(\mathbf {w} _ {t} ^ {\prime}, \mathbf {q} _ {t} ^ {\prime}\right) = \left[ p _ {1} \frac {n _ {m}}{n _ {1}} \sum_ {j = 1} ^ {n _ {1} / n _ {m}} \ell \left(\mathbf {w} _ {t} ^ {\prime}; \mathbf {z} _ {t} ^ {(1, j)}\right), \dots , p _ {m} \ell \left(\mathbf {w} _ {t} ^ {\prime}; \mathbf {z} _ {t} ^ {(m)}\right) \right] ^ {\top}.
$$
Let's take the stochastic gradient  $\mathbf{g}_q(\mathbf{w}_t',\mathbf{q}_t')$ , whose variance will be measured in terms of  $\| \cdot \|_{\infty}$  as an example to explain the intuition of inserting  $p_i$ . Define  $u_{i} = \frac{n_{m}}{n_{i}}\sum_{j = 1}^{n_{i} / n_{m}}\ell (\mathbf{w}_{t}^{\prime};\mathbf{z}_{t}^{(i,j)})$ . With a larger mini-batch size  $n_i / n_m$ ,  $u_{i}$  will approximate  $R_{i}(\mathbf{w}_{t}^{\prime})$  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  $\| \mathbf{g}_q(\mathbf{w}_t',\mathbf{q}_t')\|_{\infty}$ , since  $\| \cdot \|_{\infty}$  is insensitive to perturbations of small elements. Similar to the case in Theorem 3, the convergence rate of  $R_{i}(\cdot)$  depends on  $1 / p_i$ , and becomes faster if  $p_i$  is larger.
Based on (24), we use SMD to update  $(\mathbf{w}_t', \mathbf{q}_t')$ , and denote the solution by  $(\mathbf{w}_{t+1}, \mathbf{q}_{t+1})$ :
$$
\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} ^ {\prime}, \mathbf {q} _ {t} ^ {\prime}\right), \mathbf {w} - \mathbf {w} _ {t} ^ {\prime} \right\rangle + B _ {w} (\mathbf {w}, \mathbf {w} _ {t} ^ {\prime}) \right\}, \tag {25}
$$
$$
\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} ^ {\prime}, \mathbf {q} _ {t} ^ {\prime}\right), \mathbf {q} - \mathbf {q} _ {t} ^ {\prime} \right\rangle + B _ {q} \left(\mathbf {q}, \mathbf {q} _ {t} ^ {\prime}\right) \right\}. \tag {26}
$$
Next, we draw another  $n_i / n_m$  samples from each distribution  $\mathcal{P}_i$ , denoted by  $\hat{\mathbf{z}}_t^{(i,j)}$ ,  $j = 1,\dots ,n_i / n_m$ , to construct stochastic gradients at  $(\mathbf{w}_{t + 1},\mathbf{q}_{t + 1})$ :
$$
\mathbf {g} _ {w} \left(\mathbf {w} _ {t + 1}, \mathbf {q} _ {t + 1}\right) = \sum_ {i = 1} ^ {m} q _ {t + 1, i} p _ {i} \left(\frac {n _ {m}}{n _ {i}} \sum_ {j = 1} ^ {n _ {i} / n _ {m}} \nabla \ell \left(\mathbf {w} _ {t + 1}; \hat {\mathbf {z}} _ {t} ^ {(i, j)}\right)\right), \tag {27}
$$
$$
\mathbf {g} _ {q} (\mathbf {w} _ {t + 1}, \mathbf {q} _ {t + 1}) = \left[ p _ {1} \frac {n _ {m}}{n _ {1}} \sum_ {j = 1} ^ {n _ {1} / n _ {m}} \ell (\mathbf {w} _ {t + 1}; \hat {\mathbf {z}} _ {t} ^ {(1, j)}), \dots , p _ {m} \ell (\mathbf {w} _ {t + 1}; \hat {\mathbf {z}} _ {t} ^ {(m)}) \right] ^ {\top}.
$$
Then, we use them to update  $(\mathbf{w}_t', \mathbf{q}_t')$  again, and denote the result by  $(\mathbf{w}_{t+1}', \mathbf{q}_{t+1}')$ :
$$
\mathbf {w} _ {t + 1} ^ {\prime} = \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 + 1}, \mathbf {q} _ {t + 1}\right), \mathbf {w} - \mathbf {w} _ {t} ^ {\prime} \right\rangle + B _ {w} \left(\mathbf {w}, \mathbf {w} _ {t} ^ {\prime}\right) \right\}, \tag {28}
$$
$$
\mathbf {q} _ {t + 1} ^ {\prime} = \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 + 1}, \mathbf {q} _ {t + 1}\right), \mathbf {q} - \mathbf {q} _ {t} ^ {\prime} \right\rangle + B _ {q} \left(\mathbf {q}, \mathbf {q} _ {t} ^ {\prime}\right) \right\}. \tag {29}
$$
To meet the budget constraints, we repeat the above process for  $n_m / 2$  iterations. Finally, we return  $\bar{\mathbf{w}} = \frac{2}{n_m} \sum_{t=2}^{1+n_m/2} \mathbf{w}_t$  and  $\bar{\mathbf{q}} = \frac{2}{n_m} \sum_{t=2}^{1+n_m/2} \mathbf{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}(\cdot)$  is smooth, and the dual norm  $\| \cdot \|_{w,*}$  satisfies a regularity condition.
Assumption 5 All the risk functions are  $L$ -smooth, i.e.,
$$
\left\| \nabla R _ {i} (\mathbf {w}) - \nabla R _ {i} \left(\mathbf {w} ^ {\prime}\right) \right\| _ {w, *} \leq L \| \mathbf {w} - \mathbf {w} ^ {\prime} \| _ {w}, \forall \mathbf {w}, \mathbf {w} ^ {\prime} \in \mathcal {W}, i \in [ m ]. \tag {30}
$$
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  $\| \cdot \|_{w,*}$  is  $\kappa$ -regular for some small constant  $\kappa \geq 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  $\ell_p$ -norm and the infinity norm.
Then, we have the following theorem for Algorithm 4.
Theorem 4 Define
$$
p _ {\max } = \max  _ {i \in [ m ]} p _ {i}, \quad \omega_ {\max } = \max  _ {i \in [ m ]} \frac {p _ {i} ^ {2} n _ {m}}{n _ {i}}, \tag {31}
$$
$$
\widetilde {L} = 2 \sqrt {2} p _ {\max } \left(D ^ {2} L + D ^ {2} G \sqrt {\ln m}\right), a n d \sigma^ {2} = 2 c \omega_ {\max } \left(\kappa D ^ {2} G ^ {2} + \ln^ {2} m\right)
$$
where  $c > 0$  is an absolute constant. Under Assumptions 1, 2, 3, 4, 5 and 6, and setting
$$
\eta_ {w} = 2 D ^ {2} \min \left(\frac {1}{\sqrt {3} \widetilde {L}}, \frac {2}{\sqrt {7 \sigma^ {2} n _ {m}}}\right), a n d \eta_ {q} = 2 \min \left(\frac {1}{\sqrt {3} \widetilde {L}}, \frac {2}{\sqrt {7 \sigma^ {2} n _ {m}}}\right) \ln m
$$
in Algorithm 4, with probability at least  $1 - \delta$ , we have
$$
R _ {i} (\bar {\mathbf {w}}) - \frac {1}{p _ {i}} p _ {\varphi} ^ {*} = \frac {1}{p _ {i}} \left(\frac {7 \widetilde {L}}{n _ {m}} + \sqrt {\frac {\sigma^ {2}}{n _ {m}}} \left(1 4 \sqrt {\frac {2}{3}} + 7 \sqrt {3 \log {\frac {2}{\delta}}} + \frac {1 4}{n _ {m}} \log {\frac {2}{\delta}}\right)\right)
$$
where  $p_{\varphi}^{*}$  is the optimal value of (23).
Furthermore, by setting  $p_i$  as
$$
p _ {i} = \frac {1 / \sqrt {n _ {m}} + 1}{1 / \sqrt {n _ {m}} + \sqrt {n _ {m} / n _ {i}}}, \tag {32}
$$
with high probability, we have
$$
R _ {i} (\bar {\mathbf {w}}) - \frac {1}{p _ {i}} p _ {\varphi} ^ {*} = O \left(\left(\frac {1}{n _ {m}} + \frac {1}{\sqrt {n _ {i}}}\right) \sqrt {\kappa + \ln^ {2} m}\right).
$$
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) / \sqrt{n_i})$  rate for all distributions  $\mathcal{P}_i$  such that  $n_i \leq n_m^2$ , which is much better than the  $O(\sqrt{n_1 \ln m} / n_i)$  rate of Algorithm 3, and the  $O(\sqrt{(\log m) / n_m})$  rate of the Baseline. For distributions with a larger number of budget, i.e.,  $n_i > n_m^2$ , 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].