Learning the Robustness Mechanism with Bilevel Optimization
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 1.3 | 72.7 1.6 | 59.2 0.1 | 92.7 0.1 | 34.8 6.0 | 87.8 1.4 |
| DFR Tr | 62.9 8.3 | 68.9 4.3 | 65.5 4.8 | 89.4 0.3 | 40.8 2.7 | 87.9 0.9 |
| PDE | 62.8 7.8 | 69.1 3.8 | 35.9 3.4 | 92.0 0.6 | 39.0 3.9 | 81.8 0.8 |
| HDRO | 72.7 0.2 | 76.3 3.1 | 72.4 3.0 | 91.4 0.2 | 40.8 3.1 | 88.0 1.8 |
| Fixed | ||||||
| CMNIST | CelebA | CivilComments | ||||
| Method | Worst | Avg | Worst | Avg | Worst | Avg |
| ERM | 1.6 1.7 | 15.7 6.1 | 25.0 2.6 | 95.4 0.1 | 36.8 4.6 | 91.3 0.8 |
| GIC | 25.7 8.8 | 48.0 13.4 | 47.1 9.2 | 90.8 0.8 | 54.2 4.9 | 86.8 1.4 |
| XRM | 68.8 3.8 | 72.2 3.2 | 51.4 2.8 | 89.7 0.2 | 25.9 6.4 | 89.6 1.5 |
| AGRO | 22.8 4.9 | 32.4 2.2 | 22.1 1.8 | 95.4 0.2 | 24.5 3.5 | 91.4 0.4 |
| SSA | 70.0 2.9 | 72.3 1.7 | 58.6 2.3 | 89.7 0.1 | 27.4 8.6 | 87.1 1.9 |
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
| Group | Train | Val | Test |
|---|---|---|---|
| , green | 2,998 | 2,591 | 8,966 |
| , green | 11,781 | 2,513 | 1,013 |
| , red | 12,130 | 2,465 | 1,068 |
| , red | 3,091 | 2,431 | 8,953 |
| Total | 30,000 | 10,000 | 20,000 |
| 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 | |
| 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 |
| Setting | CMNIST | CelebA | CivilComments |
| Model | ResNet-50 | ResNet-50 | DistilBERT |
| Weight decay | |||
| Lower Train/Upper Validation batch size | 256 | 128 | 64 |
| in Shen et al. (2026) | 10000 | ||
| Iterations of SAPD per outer iteration | 1 | ||
| 1.0 | |||
| CMNIST | CelebA | CivilComments | |
|---|---|---|---|
| HDRO | 405 | 3145 | 625 |
| Bi-HDRO ( ) | 6674 | 129 | 3912 |