Group Distributionally Robust Optimization with Flexible Sample Queries

Haomin Bai; Dingzhi Yu; Shuai Li; Haipeng Luo; Lijun Zhang

cs.LGMay 21, 2025

Group Distributionally Robust Optimization with Flexible Sample Queries

Authors: Haomin Bai, Dingzhi Yu, Shuai Li, Haipeng Luo, Lijun Zhang

Organizations: National Key Laboratory for Novel Software Technology, Nanjing University, China · School of Artificial Intelligence, Nanjing University, China · John Hopcroft Center, Shanghai Jiao Tong University, China · Thomas Lord Department of Computer Science, University of Southern California, USA

Abstract

Group distributionally robust optimization (GDRO) aims to learn models that perform well across $m$ distributions simultaneously. Existing methods either (i) query $m$ samples per round, which imposes strict implementation requirements, or (ii) query only a single sample per round, which requires many rounds to converge and therefore leads to long runtime. To avoid these limitations, we introduce a novel setting called flexible sample queries (FSQ), which allows the number of samples that can be queried to vary across rounds. Under the FSQ setting, we cast GDRO as a two-player game: one performs non-oblivious online convex optimization, while the other tackles a non-oblivious prediction with limited advice (PLA) problem. In this game, the first player updates its decision using follow-the-regularized-leader with stochastic gradients. For the second player, we develop a new PLA method equipped with an adaptive exploration strategy that can leverage an arbitrary per-round sample size. We then establish the first high-probability regret bound for non-oblivious PLA. By integrating these components, we develop an anytime GDRO algorithm that supports FSQ, and derive a high-probability optimization error bound of $O\left(\frac{1}{t}\sqrt{\sum_{j=1}^t \frac{m}{r_j}\log m}\right)$ , where $r_j$ is the sample size at round $j$ . This result demonstrates that the optimization error decreases as the per-round sample sizes increase, and implies the same near-optimal sample complexity of $O(m\log (m)/ε^2)$ for any fixed sample size $r\in[m]$ .

Figures & tables

Risk (Synthetic)	Time ( Online ( $1$ ))	Time ( Flex-GDRO )	Speedup	Risk (Diabetes)	Time ( Online ( $1$ ))	Time ( Flex-GDRO )	Speedup
0.525	98.192	9.107	10.78	0.99	4.842	0.569	8.50
0.52	129.708	12.129	10.69	0.985	6.479	0.876	7.39
0.515	175.201	15.152	11.56	0.98	9.406	1.236	7.61
0.51	256.080	21.191	12.08	0.975	13.495	1.902	7.09
0.505	432.405	33.274	13.00	0.97	22.839	3.488	6.55
0.5	1123.537	108.804	10.33	0.965	51.470	8.710	5.91

Table 1: Comparison of Online (

1

) and Flex-GDRO on the synthetic and diabetes datasets. The experiments use five independent runs. Time is measured in seconds. Speedup is the runtime ratio of Online (

1

) to Flex-GDRO for achieving the same risk. Gray shading highlights Flex-GDRO .

Risk	Samples ( Online ( $1$ ))	Samples ( Flex-GDRO )	Time ( Online ( $1$ ))	Time ( Flex-GDRO )	Speedup
0.68	61.35k	69.33k	40.26	4.79	8.40
0.67	67.55k	75.34k	44.34	5.20	8.52
0.66	77.75k	81.36k	51.03	5.62	9.08
0.65	97.81k	87.39k	64.19	6.03	10.64
0.64	120.00k	101.84k	78.80	7.03	11.20

Table 2: Comparison of Online (

1

) and Flex-GDRO on CivilComments-WILDS. Time is measured in seconds. Speedup is the runtime ratio of Online (

1

) to Flex-GDRO for achieving the same risk. Gray shading highlights Flex-GDRO .

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Risk	Samples ( $\times 10^{3}$ ) ( Online ( $1$ ))	Samples ( $\times 10^{3}$ ) ( Flex-GDRO )	Time ( Online ( $1$ ))	Time ( Flex-GDRO )	Speedup
0.98	$52.06\pm 0.65$	$63.43\pm 3.33$	$78.45\pm 0.90$	$23.61\pm 0.94$	3.32
0.95	$71.22\pm 1.19$	$73.78\pm 3.81$	$106.50\pm 1.69$	$27.46\pm 1.18$	3.88
0.94	$79.53\pm 0.30$	$79.51\pm 5.01$	$118.76\pm 0.54$	$29.48\pm 1.65$	4.03
0.93	$89.73\pm 1.68$	$83.68\pm 6.89$	$133.66\pm 2.70$	$31.00\pm 2.37$	4.31

Appendix

Table 3: Comparison of Online (

1

) and Flex-GDRO on FEMNIST with

100

writer groups under dynamic resources. Entries are mean

\pm

sample standard deviation over five seeds at the first saved checkpoint reaching each threshold. Time is measured in seconds for head optimization. Speedup is the ratio of mean optimization times of Online (

1

) to Flex-GDRO for achieving the same risk. Gray shading highlights Flex-GDRO .

$m/r$	$m$	$r$	$T=1{,}000$	$T=5{,}000$	$T=10{,}000$	$T=20{,}000$
5	50	10	$2.476\pm 0.133$	$2.586\pm 0.081$	$2.680\pm 0.092$	$2.688\pm 0.130$
5	100	20	$2.502\pm 0.009$	$2.669\pm 0.038$	$2.694\pm 0.048$	$2.703\pm 0.073$
5	200	40	$2.525\pm 0.081$	$2.646\pm 0.038$	$2.691\pm 0.082$	$2.730\pm 0.107$
10	50	5	$5.464\pm 0.111$	$5.687\pm 0.160$	$5.932\pm 0.194$	$6.060\pm 0.225$
10	100	10	$5.538\pm 0.159$	$5.927\pm 0.115$	$5.925\pm 0.064$	$6.066\pm 0.090$
10	200	20	$5.751\pm 0.053$	$5.862\pm 0.147$	$6.043\pm 0.138$	$6.044\pm 0.214$

Appendix

Table 4: Empirical

\mathbf{q}_{t}

-weighted variance

\widehat{V}_{q,t}

of the loss estimator on synthetic data. Entries are mean

\pm

standard deviation over five seeds.

$m/r$	$m$	$r$	$T=1{,}000$	$T=5{,}000$	$T=10{,}000$	$T=20{,}000$
5	50	10	$2.537\pm 0.127$	$2.635\pm 0.068$	$2.725\pm 0.069$	$2.739\pm 0.136$
5	100	20	$2.549\pm 0.010$	$2.682\pm 0.054$	$2.719\pm 0.049$	$2.747\pm 0.076$
5	200	40	$2.534\pm 0.049$	$2.667\pm 0.041$	$2.693\pm 0.057$	$2.753\pm 0.109$
10	50	5	$5.597\pm 0.092$	$5.961\pm 0.239$	$6.146\pm 0.174$	$6.228\pm 0.250$
10	100	10	$5.690\pm 0.205$	$6.099\pm 0.140$	$6.088\pm 0.084$	$6.181\pm 0.153$
10	200	20	$5.801\pm 0.054$	$5.966\pm 0.091$	$6.134\pm 0.141$	$6.217\pm 0.232$

Appendix

Table 5: Unweighted mean per-coordinate variance

\widehat{V}_{t}

of the loss estimator on synthetic data. Entries are mean

\pm

standard deviation over five seeds.

$m/r$	$m$	$r$	$T=1{,}000$	$T=5{,}000$	$T=10{,}000$	$T=20{,}000$
5	50	10	$0.220\pm 0.012$	$0.189\pm 0.009$	$0.180\pm 0.013$	$0.183\pm 0.011$
5	100	20	$0.212\pm 0.003$	$0.194\pm 0.008$	$0.179\pm 0.006$	$0.177\pm 0.012$
5	200	40	$0.218\pm 0.010$	$0.192\pm 0.007$	$0.190\pm 0.019$	$0.183\pm 0.007$
10	50	5	$0.213\pm 0.004$	$0.192\pm 0.012$	$0.180\pm 0.010$	$0.183\pm 0.015$
10	100	10	$0.218\pm 0.011$	$0.191\pm 0.012$	$0.180\pm 0.008$	$0.172\pm 0.012$
10	200	20	$0.213\pm 0.005$	$0.189\pm 0.005$	$0.180\pm 0.011$	$0.182\pm 0.009$

Appendix

Table 6: Mean per-coordinate variance of the

w

-player stochastic gradient on synthetic data. Entries are mean

\pm

standard deviation over five seeds.

(p_{i},p_{j})=\begin{cases}(p_{i}+\alpha,p_{j}-\alpha)\text{ with probability }\frac{\beta}{\alpha+\beta}\\ (p_{i}-\beta,p_{j}+\beta)\text{ with probability }\frac{\alpha}{\alpha+\beta}\end{cases}

Appendix

Algorithm 4 DepRound

Group Distributionally Robust Optimization

See all 12 papers

Group Distributionally Robust Optimization over Decentralized Networks
Oct 21, 2026Haomin Bai, Junyan Sun, Quan Cheng +2Group Distributionally Robust Optimization Distributionally Robust Optimization
State Key Laboratory of Novel Software Technology, Nanjing University · School of Artificial Intelligence, Nanjing University · School of Software Technology, Zhejiang University
Generative Distributionally Robust Optimization
Jul 27, 2026Ziwei Zhang, Jonathan Yu-Meng Li, Zhihao JinSinkhorn Divergence Distributionally Robust Optimization
1Telfer School of Management, University of Ottawa · Department of Electrical and Computer Engineering, Western University
From Dual Tracking to Clipping: Provably Faster Distributionally Robust Multi-Objective Optimization
May 7, 2026Yufeng Yang, Fangning Zhuo, Ziyi Chen +2Gradient Clipping Gradient Descent
Texas A&M University · A*STAR · University of Maryland

Oct 21, 2026J/K move · Enter open · S save

Group Distributionally Robust Optimization over Decentralized Networks

Haomin Bai, Junyan Sun, Quan Cheng, Yuanyu Wan+1

State Key Laboratory of Novel Software Technology, Nanjing University · School of Artificial Intelligence, Nanjing University · School of Software Technology, Zhejiang University

Group Distributionally Robust Optimization Distributionally Robust Optimization Decentralized Optimization

Open paper AskPDF arXiv

We study group distributionally robust optimization (GDRO) over decentralized networks, where workers collaborate to learn robust models by minimizing the maximum local risk. Existing methods optimize only dual-regularized surrogates and establish guarantees for a globally averaged model that is not available to workers. In this paper, we directly optimize the original GDRO objective and establish guarantees for local models in semi-decentralized (SemiDec) and fully decentralized (Dec) settings. A key difficulty is that applying the primal-dual decomposition from centralized GDRO yields worker-dependent losses whose evaluation requires unavailable cross-worker information. To address this challenge, we develop a novel error decomposition that recasts decentralized GDRO as (i) a decentralized online convex optimization (D-OCO) problem and a prediction-with-expert-advice problem in SemiDec, and as (ii) two coupled D-OCO problems in Dec, with discrepancy terms capturing inconsistencies among local decisions. By isolating consensus errors in analysis, we develop modular frameworks that support different consensus routines. We then establish the first per-worker convergence guarantees for optimizing the original GDRO objective over decentralized networks. Instantiations with accelerated gossip and Choco-gossip achieve an

O(\sqrt{(\log m)/T})

convergence rate over

T

update rounds, while experiments demonstrate improved worst-worker performance over baselines.

Abstract

Group Distributionally Robust Optimization

Explore similar work

Group Distributionally Robust Optimization over Decentralized Networks

Generative Distributionally Robust Optimization

From Dual Tracking to Clipping: Provably Faster Distributionally Robust Multi-Objective Optimization