# Enhancing Distributional Stability among Sub-populations

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 14 of its 39 references that are in the map.
- First posted on arXiv 7 Jun 2022.

## Builds on (most cited first, 14 of 14)

- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [In Search of Lost Domain Generalization](https://mlsanity.com/paper/2007.01434/) (2020)
- [Causal inference using invariant prediction: identification and confidence intervals](https://mlsanity.com/paper/1501.01332/) (2015)
- [Retiring Adult: New Datasets for Fair Machine Learning](https://mlsanity.com/paper/2108.04884/) (2021)
- [Certifying Some Distributional Robustness with Principled Adversarial Training](https://mlsanity.com/paper/1710.10571/) (2017)
- [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)
- [Accuracy on the Line: On the Strong Correlation Between Out-of-Distribution and In-Distribution Generalization](https://mlsanity.com/paper/2107.04649/) (2021)
- [Invariance Principle Meets Information Bottleneck for Out-of-Distribution Generalization](https://mlsanity.com/paper/2106.06607/) (2021)
- [Truncated Back-propagation for Bilevel Optimization](https://mlsanity.com/paper/1810.10667/) (2018)
- [Invariant Risk Minimization Games](https://mlsanity.com/paper/2002.04692/) (2020)
- [Towards a Theoretical Framework of Out-of-Distribution Generalization](https://mlsanity.com/paper/2106.04496/) (2021)
- [Heterogeneous Risk Minimization](https://mlsanity.com/paper/2105.03818/) (2021)
- [Incorporating Unlabeled Data into Distributionally Robust Learning](https://mlsanity.com/paper/1912.07729/) (2019)

## Built on by (most cited first, 3 of 3)

- [Scalable Out-of-distribution Robustness in the Presence of Unobserved Confounders](https://mlsanity.com/paper/2411.19923/) (2024)
- [Supervision Adaptation Balancing In-distribution Generalization and Out-of-distribution Detection](https://mlsanity.com/paper/2206.09380/) (2022)
- [Data Heterogeneity Modeling for Trustworthy Machine Learning](https://mlsanity.com/paper/2506.00969/) (2025)

## Abstract (arXiv)

Enhancing the stability of machine learning algorithms under distributional shifts is at the heart of the Out-of-Distribution (OOD) Generalization problem. Derived from causal learning, recent works of invariant learning pursue strict invariance with multiple training environments. Although intuitively reasonable, strong assumptions on the availability and quality of environments are made to learn the strict invariance property. In this work, we come up with the ``distributional stability" notion to mitigate such limitations. It quantifies the stability of prediction mechanisms among sub-populations down to a prescribed scale. Based on this, we propose the learnability assumption and derive the generalization error bound under distribution shifts. Inspired by theoretical analyses, we propose our novel stable risk minimization (SRM) algorithm to enhance the model's stability w.r.t. shifts in prediction mechanisms (Y|X-shifts). Experimental results are consistent with our intuition and validate the effectiveness of our algorithm. The code can be found at https://github.com/LJSthu/SRM.