Stochastic Approximation Approaches to
Group Distributionally Robust Optimization and Beyond

In the first extention, 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\cdots\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. To achieve this goal, we propose a novel formulation of weighted GDRO, which weights each risk $R_{i}(\cdot)$ in (4) by a scale factor $p_{i}$. 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. Consequently, the budget will be satisfied in expectation after $n_{1}$ rounds, and our algorithm can be regarded as SMD for an instance of weighted GDRO. With the help of scale factors, we demonstrate 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 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{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 $O(n_{m})$ times.

In the second extension, we delve into another scenario where distributions exhibit heterogeneity, indicating significant variations in their risks (Li et al., 2019). The widely acknowledged sensitivity of the max operation to outliers implies that GDRO could be dominated by a single outlier distribution, while neglecting others (Shalev-Shwartz and Wexler, 2016). Inspired by the average top-$k$ loss for supervised learning (Fan et al., 2017), we modify our objective from the maximum risk $\mathcal{L}_{\max}(\mathbf{w})$ in GDRO to the average top-$k$ risk:

where $\mathcal{B}_{m,k}$ is the set of subsets of $[m]$ with size $k$, i.e., $\mathcal{B}_{m,k}=\{\mathcal{I}\subseteq[m]||\mathcal{I}|=k\}$. This modification aims to reduce the impact of outliers in heterogeneous distributions while still including GDRO as a special case.

We refer to the minimization of $\mathcal{L}_{k}(\mathbf{w})$ as average top-$k$ risk optimization (AT_kRO), and develop two stochastic algorithms. Similar to GDRO, AT_kRO can be formulated as a stochastic convex-concave saddle-point problem, akin to (4), with the only difference being that the domain of $\mathbf{q}$ is the capped simplex instead of the standard simplex. Therefore, we can employ SMD to update $\mathbf{w}$ and $\mathbf{q}$, which uses $m$ samples in each round. Theoretical analysis demonstrates that this approach achieves an $O(\sqrt{(\log(m/k))/T})$ convergence rate, implying an $O((m\log(m/k))/\epsilon^{2})$ sample complexity. Furthermore, to circumvent the limitation of predefining the total number of iterations $T$, we introduce an anytime version that attains an $\widetilde{O}(\sqrt{(\log(m/k))/t})$ convergence rate.

Following the second approach for GDRO, we reduce the number of samples required in each round from $m$ to $k$ by casting AT_kRO as a two-player game. In each round, we use the Dependent Rounding (DepRound) algorithm (Gandhi et al., 2006) to select $k$ distributions based on the current value of $\mathbf{q}$, and then draw $1$ sample from each selected distribution. Then, we construct unbiased stochastic gradients for $\mathbf{w}$, and apply SMD for updates. Since the domain of $\mathbf{q}$ is the capped simplex, we model the online problem for $\mathbf{q}$ as an instance of non-oblivious combinatorial semi-bandits, and extend Exp3-IX to develop its update rule. We prove that our algorithm achieves an $O(\sqrt{m(\log m)/(kT)})$ convergence rate, yielding an $O(m(\log m)/\epsilon^{2})$ sample complexity. Similarly, we have also designed an anytime approach, which uses $1$ sample per round and achieves an $\widetilde{O}(\sqrt{m(\log m)/t})$ rate.

This paper extends our previous conference version (Zhang et al., 2023) by developing anytime algorithms, investigating a new scenario, and conducting more experiments, as detailed below.

• •

  First, we adapt the two SA algorithms for GDRO to operate in an anytime manner. In the conference paper, our algorithms for GDRO required predefining the total number of iterations $T$ to set step sizes. By adopting time-varying step sizes, we design anytime algorithms and provide the corresponding theoretical analysis.

• •

  Second, we explore the scenario of heterogeneous distributions, which involves outlier distributions with significantly high risks. To mitigate the impact of these outliers, we propose to solve the AT_kRO problem and develop two algorithms: one employs SMD with $m$ samples per round, achieving a sample complexity of $O((m\log(m/k))/\epsilon^{2})$; the other combines SMD with an algorithm for non-oblivious combinatorial semi-bandits, achieving a sample complexity of $O(m(\log m)/\epsilon^{2})$ and using $k$ samples in each iteration. Furthermore, we have also extended these two algorithms into anytime versions.

• •

  Last, we construct a heterogeneous data set and perform experiments to verify the advantages of AT_kRO. Additionally, we compare the performance of the anytime algorithms with their non-anytime counterparts, demonstrating the benefits of the anytime capability.

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

For the simplex $\Delta_{m}$, we choose the negative entropy (neg-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.

For $\Delta_{m}$, it is easy to verify that its diameter measured by the neg-entropy function is bounded by $\sqrt{\ln m}$.

To simplify presentations, we assume the loss belongs to $[0,1]$, and its gradient is also bounded.

Note that it is possible to relax the bounded assumptions in (8) and (9) to light tail conditions such as the sub-Gaussian Property (Vershynin, 2018).

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 (4) by the error

which directly controls the optimality of $\bar{\mathbf{w}}$ to the original problem (3), since

To apply SMD, the key is to construct stochastic gradients of the function $\phi(\mathbf{w},\mathbf{q})$ in (4). We first present its true gradients with respect to $\mathbf{w}$ and $\mathbf{q}$:

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

Obviously, they are unbiased estimators of the true gradients:

where $\mathrm{E}_{t-1}[\cdot]$ represents the expectation conditioned on the randomness until round $t-1$. It is worth mentioning that the construction of $\mathbf{g}_{w}(\mathbf{w}_{t},\mathbf{q}_{t})$ can be further simplified to

where $i_{t}\in[m]$ is drawn randomly according to the probability $\mathbf{q}_{t}$.

Then, we use SMD to update $\mathbf{w}_{t}$ and $\mathbf{q}_{t}$:

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)$. For example, if $\nu_{w}(\mathbf{w})=\frac{1}{2}\|\mathbf{w}\|_{2}^{2}$, (14) becomes stochastic gradient descent (SGD), i.e.,

where $\Pi_{\mathcal{W}}[\cdot]$ denotes the Euclidean projection onto the nearest point in $\mathcal{W}$. Since $B_{q}(\mathbf{q},\mathbf{q}_{t})$ is defined in terms of the neg-entropy, (15) is equivalent to

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 complete procedure is summarized in Algorithm 1.

Based on the theoretical guarantee of SMD for stochastic convex-concave optimization (Nemirovski et al., 2009, § 3.1), we have the following theorem for Algorithm 1.

In this section, we explore methods to reduce the number of samples used in each iteration from $m$ to $1$. As shown in (13), 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 (17) is upper bounded by $m$. The reason is that we insist on the unbiasedness of the stochastic gradient, which leads to a large variance. To control the variance, Carmon and Hausler (2022) have applied gradient clipping to $\hat{\mathbf{g}}_{q}(\mathbf{w}_{t},\mathbf{q}_{t})$, and established an $O(m(\log m)/\epsilon^{2})$ sample complexity that holds in expectation. Different from their approach, we borrow techniques from online learning to balance the bias and the variance.

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; Roux et al., 2021). This transformation allows us to exploit no-regret algorithms developed in online learning to bound the optimization error. To solve problem (4), we ask the 1st player to minimize a sequence of convex functions

under the constraint $\mathbf{w}\in\mathcal{W}$, and the 2nd player to maximize a sequence of linear functions

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 directly observe the loss function, and can only approximate $R_{i}(\mathbf{w})=\mathrm{E}_{\mathbf{z}\sim\mathcal{P}_{i}}\big{[}\ell(\mathbf{w};\mathbf{z})\big{]}$ 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 players.

Next, we discuss the online algorithms that will be used by the two players. As shown in Section 3.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. In contrast, if we apply SMD with $\hat{\mathbf{g}}_{q}(\mathbf{w}_{t},\mathbf{q}_{t})$ in (17), the regret scales at least linearly with $m$.

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 $\mathcal{P}_{i_{t}}$. With the stochastic gradient in (13), we use SMD to update $\mathbf{w}_{t}$:

Then, we reuse the sample $\mathbf{z}_{t}^{(i_{t})}$ to update $\mathbf{q}_{t}$ according to Exp3-IX, which first constructs the Implicit-eXploration (IX) loss estimator (Kocák et al., 2014):

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:

Compared with (15), the only difference is that the stochastic gradient $-\mathbf{g}_{q}(\mathbf{w}_{t},\mathbf{q}_{t})$ is now replaced with the IX loss estimator $\tilde{\mathbf{s}}_{t}$. However, it is not an instance of SMD, because $\tilde{\mathbf{s}}_{t}$ is no longer an unbiased stochastic gradient. The main advantage of $\tilde{\mathbf{s}}_{t}$ is that it reduces the variance of the gradient estimator by sacrificing a little bit of unbiasedness, which turns out to be crucial for a high probability guarantee, and thus can deal with non-oblivious adversaries. Since we still use the entropy regularizer in (26), it also enjoys an explicit form that is similar to (16).

We present the theoretical guarantee of Algorithm 2. To this end, 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).

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

Combining the above two theorems directly leads to the following optimization error bound.