# Learning the Robustness Mechanism with Bilevel Optimization

arXiv: https://arxiv.org/abs/2609.35541 · HTML: https://mlsanity.com/paper/2609.35541/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 57 of its 73 references that are in the map.
- First posted on arXiv 28 Sep 2026.

## Builds on (most cited first, 30 of 57)

- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [DARTS: Differentiable Architecture Search](https://mlsanity.com/paper/1806.09055/) (2018)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Decoupling Representation and Classifier for Long-Tailed Recognition](https://mlsanity.com/paper/1910.09217/) (2019)
- [Out-of-Distribution Generalization via Risk Extrapolation (REx)](https://mlsanity.com/paper/2003.00688/) (2020)
- [Learning to Reweight Examples for Robust Deep Learning](https://mlsanity.com/paper/1803.09050/) (2018)
- [Meta-Learning with Implicit Gradients](https://mlsanity.com/paper/1909.04630/) (2019)
- [Towards Out-Of-Distribution Generalization: A Survey](https://mlsanity.com/paper/2108.13624/) (2021)
- [Bilevel Programming for Hyperparameter Optimization and Meta-Learning](https://mlsanity.com/paper/1806.04910/) (2018)
- [Meta-learning with differentiable closed-form solvers](https://mlsanity.com/paper/1805.08136/) (2018)
- [Last Layer Re-Training is Sufficient for Robustness to Spurious Correlations](https://mlsanity.com/paper/2204.02937/) (2022)
- [Confounding variables can degrade generalization performance of radiological deep learning models](https://mlsanity.com/paper/1807.00431/) (2018)
- [Gradient-based Hyperparameter Optimization through Reversible Learning](https://mlsanity.com/paper/1502.03492/) (2015)
- [Learning Models with Uniform Performance via Distributionally Robust Optimization](https://mlsanity.com/paper/1810.08750/) (2018)
- [Environment Inference for Invariant Learning](https://mlsanity.com/paper/2010.07249/) (2020)
- [Optimizing Millions of Hyperparameters by Implicit Differentiation](https://mlsanity.com/paper/1911.02590/) (2019)
- [Forward and Reverse Gradient-Based Hyperparameter Optimization](https://mlsanity.com/paper/1703.01785/) (2017)
- [No Subclass Left Behind: Fine-Grained Robustness in Coarse-Grained Classification Problems](https://mlsanity.com/paper/2011.12945/) (2020)
- [Truncated Back-propagation for Bilevel Optimization](https://mlsanity.com/paper/1810.10667/) (2018)
- [Statistics of Robust Optimization: A Generalized Empirical Likelihood Approach](https://mlsanity.com/paper/1610.03425/) (2016)
- [Correct-N-Contrast: A Contrastive Approach for Improving Robustness to Spurious Correlations](https://mlsanity.com/paper/2203.01517/) (2022)
- [On Feature Learning in the Presence of Spurious Correlations](https://mlsanity.com/paper/2210.11369/) (2022)
- [A Two-Timescale Framework for Bilevel Optimization: Complexity Analysis and Application to Actor-Critic](https://mlsanity.com/paper/2007.05170/) (2020)
- [Change is Hard: A Closer Look at Subpopulation Shift](https://mlsanity.com/paper/2302.12254/) (2023)
- [Curriculum Domain Adaptation for Semantic Segmentation of Urban Scenes](https://mlsanity.com/paper/1707.09465/) (2017)
- [Spread Spurious Attribute: Improving Worst-group Accuracy with Spurious Attribute Estimation](https://mlsanity.com/paper/2204.02070/) (2022)
- [Towards LLM Unlearning Resilient to Relearning Attacks: A Sharpness-Aware Minimization Perspective and Beyond](https://mlsanity.com/paper/2502.05374/) (2025)
- [Discover and Cure: Concept-aware Mitigation of Spurious Correlation](https://mlsanity.com/paper/2305.00650/) (2023)

## Abstract (arXiv)

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.