Learning the Robustness Mechanism with Bilevel Optimization

Yiyang Shen; Qihang Lin; Weiran Wang

cs.LGSep 28, 2026

Learning the Robustness Mechanism with Bilevel Optimization

Authors: Yiyang Shen, Qihang Lin, Weiran Wang

Organizations: Department of Informatics University of Iowa · Tippie College of Business University of Iowa · Department of Computer Science University of Iowa

Abstract

We propose a distributionally robust learning framework where parameters defining the robustness mechanism are learned from held-out data instead of extensively tuned. Using bilevel optimization with both upper and lower level minimax problems, we create two instances of our framework to tackle setups with and without group labels in the training set. Theoretically, we provide sample complexity analysis for our robustness mechanism learning paradigm, showing that it achieves generalization guarantees comparable to exhaustive grid search while being more computationally efficient. Empirically, we evaluate our framework under a challenging setup when both intra-group and inter-group test distribution shifts occur at the same time, thereby demonstrating the efficacy and scalability of our method.

Figures & tables

	CMNIST		CelebA		CivilComments
Method	Worst	Avg	Worst	Avg	Worst	Avg
GDRO	71.3 $\pm$ 1.3	72.7 $\pm$ 1.6	59.2 $\pm$ 0.1	92.7 $\pm$ 0.1	34.8 $\pm$ 6.0	87.8 $\pm$ 1.4
DFR Tr	62.9 $\pm$ 8.3	68.9 $\pm$ 4.3	65.5 $\pm$ 4.8	89.4 $\pm$ 0.3	40.8 $\pm$ 2.7	87.9 $\pm$ 0.9
PDE	62.8 $\pm$ 7.8	69.1 $\pm$ 3.8	35.9 $\pm$ 3.4	92.0 $\pm$ 0.6	39.0 $\pm$ 3.9	81.8 $\pm$ 0.8
HDRO	72.7 $\pm$ 0.2	76.3 $\pm$ 3.1	72.4 $\pm$ 3.0	91.4 $\pm$ 0.2	40.8 $\pm$ 3.1	88.0 $\pm$ 1.8
Fixed $\epsilon$

Table 1: Accuracy over 3 runs under shifted distributions for Bi-HDRO and its baselines

	CMNIST		CelebA		CivilComments
Method	Worst	Avg	Worst	Avg	Worst	Avg
ERM	1.6 $\pm$ 1.7	15.7 $\pm$ 6.1	25.0 $\pm$ 2.6	95.4 $\pm$ 0.1	36.8 $\pm$ 4.6	91.3 $\pm$ 0.8
GIC ${}^{C_{y}}$	25.7 $\pm$ 8.8	48.0 $\pm$ 13.4	47.1 $\pm$ 9.2	90.8 $\pm$ 0.8	54.2 $\pm$ 4.9	86.8 $\pm$ 1.4
XRM	68.8 $\pm$ 3.8	72.2 $\pm$ 3.2	51.4 $\pm$ 2.8	89.7 $\pm$ 0.2	25.9 $\pm$ 6.4	89.6 $\pm$ 1.5
AGRO	22.8 $\pm$ 4.9	32.4 $\pm$ 2.2	22.1 $\pm$ 1.8	95.4 $\pm$ 0.2	24.5 $\pm$ 3.5	91.4 $\pm$ 0.4
SSA	70.0 $\pm$ 2.9	72.3 $\pm$ 1.7	58.6 $\pm$ 2.3	89.7 $\pm$ 0.1	27.4 $\pm$ 8.6	87.1 $\pm$ 1.9

Table 2: Accuracy over 3 runs under shifted subpopulations for Bi-PG-DRO and its baselines

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

\bar{\mathcal{P}}_{k}=P_{\rho}(\alpha,p,W,q,\widetilde{W},\widetilde{q})+\frac{\rho_{1}}{2}\|(\alpha,W,q)-(\alpha^{k},W^{k},q^{k})\|^{2}-\frac{\rho_{2}}{2}\|(p,\widetilde{W},\widetilde{q})-(p^{k},\widetilde{W}^{k},\widetilde{q}^{k})\|^{2}

Appendix

Algorithm 1 A First-Order Method for ( 13 )

Group	Train	Val	Test
$y=0$ , green	2,998	2,591	8,966
$y=1$ , green	11,781	2,513	1,013
$y=0$ , red	12,130	2,465	1,068
$y=1$ , red	3,091	2,431	8,953
Total	30,000	10,000	20,000

Appendix

Table 3: CMNIST Statistics.

Group	Train	Val	Test
non-blond female	71,629	8,535	9,767
non-blond male	66,874	8,276	7,535
blond female	22,880	2,874	2,480
blond male (before shift)	1,387	182	180
blond male (after shift)	1,461	172	116
Total (before shift) 162,770		19,867	19,962

Appendix

Table 4: CelebA Statistics

Group	Train	Val	Test
non-toxic, non-black	231,738	39,006	115,223
non-toxic, black	6,785	1,119	3,335
toxic, non-black	27,404	4,522	13,687
toxic, black (before shift)	3,111	533	1,537
toxic, black (after shift)	2,738	533	1,910
Total (before shift)	269,038	45,180	133,782

Appendix

Table 5: CivilComments Statistics

Setting	CMNIST	CelebA	CivilComments
Model	ResNet-50	ResNet-50	DistilBERT
Weight decay $\lambda$	$10^{-3}$	$10^{-1}$	$10^{-3}$
Lower Train/Upper Validation batch size	256	128	64
$L_{\nabla h}$ in Shen et al. (2026)	10000
Iterations of SAPD per outer iteration $T$	1
$\eta_{\mathrm{init}}$	1.0

Appendix

Table 6: Settings and Hyperparameters for our methods.

	CMNIST	CelebA	CivilComments
HDRO	405	3145	625
Bi-HDRO ( $\epsilon_{g}$ )	6674	129	3912

Appendix

Table 7: Wall-clock time for one run in seconds.

Explore similar work

Statistical Properties of Robust Learning under Distributional Shifts
Aug 13, 2026Zhiyi Li, Xiaojie Mao, Yunbei Xu +1Distributionally-Robust Optimization Distribution Shifts
National University of Singapore · Tsinghua University · University College London
Unification and Optimization of Robust Supervised Learning
May 27, 2026Jonas Hanselle, Valentin Margraf, Clemens Damke +1Empirical Risk Minimization Learnability
LMU Munich, MCML · LMU Munich, MCML, DFKI
Distributionally-Robust Learning to Optimize
May 7, 2026Vinit Ranjan, Jisun Park, Bartolomeo StellatoDistributionally-Robust Optimization Convex Optimization
Department of Operations Research and Financial Engineering, Princeton University. · Department of Operations Research and Financial Engineering, Princeton University, and Research Institute of Mathematics, Seoul National University. · Research Institute of Mathematics, Seoul National University.

Aug 13, 2026, stat.MLJ/K move · Enter open · S save

Statistical Properties of Robust Learning under Distributional Shifts

Zhiyi Li, Xiaojie Mao, Yunbei Xu, Ruohan Zhan

National University of Singapore · Tsinghua University · University College London

Distributionally-Robust Optimization Distribution Shifts Learnability

Open paper AskPDF arXiv

Distributional shifts arise when the target deployment environment differs from the source environment that generated the training data. Robust learning frameworks such as Distributionally Robust Optimization (DRO) and Robust Satisficing (RS) aim to address this challenge, yet their finite-sample guarantees under such shifts, and their systematic comparison, remain underexplored: existing analyses typically establish guarantees either in the source environment or for adversarial worst-case performance over an ambiguity set. This paper instead studies generalization error in the target environment---the excess loss under the shifted target distribution. Our contributions are threefold. First, we derive finite-sample generalization error bounds in the shifted target environment for both DRO and RS. These bounds explicitly characterize the trade-off between reduced sensitivity to shift and the regularization penalty induced by each method's robustness hyperparameter, and they avoid the curse of dimensionality associated with Wasserstein empirical concentration. Second, when partial shift information such as shift magnitude or direction is available, we propose information-directed hyperparameter calibrations and compare the two methods given the same information. Under these calibrations, and in the partial-information regimes we study, DRO and RS exhibit complementary theoretical and empirical behavior. Finally, we apply the framework to a network lot-sizing problem, using it to interpret how robust policies respond to positive shifts in the demand distribution. Together, these results fill a gap in understanding the statistical properties of robust learning methods under distributional shifts and provide a principled basis for comparing DRO and RS.

Abstract

Explore similar work

Statistical Properties of Robust Learning under Distributional Shifts

Unification and Optimization of Robust Supervised Learning

Distributionally-Robust Learning to Optimize