# Distributionally Robust Optimization with Probabilistic Group

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

## Key facts

- Cited by 14 papers in the map (data through 14 Oct 2026).
- Builds on 31 of its 45 references that are in the map.
- First posted on arXiv 10 Mar 2023.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization](https://mlsanity.com/paper/1610.02391/) (2016)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [A Broad-Coverage Challenge Corpus for Sentence Understanding through Inference](https://mlsanity.com/paper/1704.05426/) (2017)
- [ImageNet-trained CNNs are biased towards texture; increasing shape bias improves accuracy and robustness](https://mlsanity.com/paper/1811.12231/) (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)
- [Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks](https://mlsanity.com/paper/1703.10593/) (2017)
- [Learning to Reweight Examples for Robust Deep Learning](https://mlsanity.com/paper/1803.09050/) (2018)
- [Annotation Artifacts in Natural Language Inference Data](https://mlsanity.com/paper/1803.02324/) (2018)
- [Just Train Twice: Improving Group Robustness without Training Group Information](https://mlsanity.com/paper/2107.09044/) (2021)
- [Agnostic Federated Learning](https://mlsanity.com/paper/1902.00146/) (2019)
- [Nuanced Metrics for Measuring Unintended Bias with Real Data for Text Classification](https://mlsanity.com/paper/1903.04561/) (2019)
- [An Investigation of Why Overparameterization Exacerbates Spurious Correlations](https://mlsanity.com/paper/2005.04345/) (2020)
- [Environment Inference for Invariant Learning](https://mlsanity.com/paper/2010.07249/) (2020)
- [Data-driven Distributionally Robust Optimization Using the Wasserstein Metric: Performance Guarantees and Tractable Reformulations](https://mlsanity.com/paper/1505.05116/) (2015)
- [What is the Effect of Importance Weighting in Deep Learning?](https://mlsanity.com/paper/1812.03372/) (2018)
- [Hidden Stratification Causes Clinically Meaningful Failures in Machine Learning for Medical Imaging](https://mlsanity.com/paper/1909.12475/) (2019)
- [No Subclass Left Behind: Fine-Grained Robustness in Coarse-Grained Classification Problems](https://mlsanity.com/paper/2011.12945/) (2020)
- [Large-Scale Methods for Distributionally Robust Optimization](https://mlsanity.com/paper/2010.05893/) (2020)
- [Correct-N-Contrast: A Contrastive Approach for Improving Robustness to Spurious Correlations](https://mlsanity.com/paper/2203.01517/) (2022)
- [Variance-based regularization with convex objectives](https://mlsanity.com/paper/1610.02581/) (2016)
- [An Empirical Study on Robustness to Spurious Correlations using Pre-trained Language Models](https://mlsanity.com/paper/2007.06778/) (2020)
- [Spread Spurious Attribute: Improving Worst-group Accuracy with Spurious Attribute Estimation](https://mlsanity.com/paper/2204.02070/) (2022)
- [Data-Driven Robust Optimization](https://mlsanity.com/paper/1401.0212/) (2013)
- [Model Patching: Closing the Subgroup Performance Gap with Data Augmentation](https://mlsanity.com/paper/2008.06775/) (2020)

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

- [ARGS: Alignment as Reward-Guided Search](https://mlsanity.com/paper/2402.01694/) (2024)
- [Process Reward Model with Q-Value Rankings](https://mlsanity.com/paper/2410.11287/) (2024)
- [Feed Two Birds with One Scone: Exploiting Wild Data for Both Out-of-Distribution Generalization and Detection](https://mlsanity.com/paper/2306.09158/) (2023)
- [FOOGD: Federated Collaboration for Both Out-of-distribution Generalization and Detection](https://mlsanity.com/paper/2410.11397/) (2024)
- [On the Learnability of Out-of-distribution Detection](https://mlsanity.com/paper/2404.04865/) (2024)
- [Mitigating Spurious Correlation via Distributionally Robust Learning with Hierarchical Ambiguity Sets](https://mlsanity.com/paper/2510.02818/) (2025)
- [Mitigating Subgroup Disparities in Multi-Label Speech Emotion Recognition: A Pseudo-Labeling and Unsupervised Learning Approach](https://mlsanity.com/paper/2505.14449/) (2025)
- [Subgroups Matter for Robust Bias Mitigation](https://mlsanity.com/paper/2505.21363/) (2025)
- [Bridging Distribution Shift and AI Safety: Conceptual and Methodological Synergies](https://mlsanity.com/paper/2505.22829/) (2025)
- [Group Distributionally Robust Machine Learning under Group Level Distributional Uncertainty](https://mlsanity.com/paper/2509.08942/) (2025)
- [Generalizable Multimodal Large Language Model Editing via Invariant Trajectory Learning](https://mlsanity.com/paper/2601.19700/) (2026)
- [Closing the loop in learning with missing data](https://mlsanity.com/paper/2608.09030/) (2026)
- [Learning the Robustness Mechanism with Bilevel Optimization](https://mlsanity.com/paper/2609.35541/) (2026)
- [Optimal Transport Reweighting for Robust Learning under Spurious Correlations and Label Noise](https://mlsanity.com/paper/2610.01028/) (2026)

## Abstract (arXiv)

Modern machine learning models may be susceptible to learning spurious correlations that hold on average but not for the atypical group of samples. To address the problem, previous approaches minimize the empirical worst-group risk. Despite the promise, they often assume that each sample belongs to one and only one group, which does not allow expressing the uncertainty in group labeling. In this paper, we propose a novel framework PG-DRO, which explores the idea of probabilistic group membership for distributionally robust optimization. Key to our framework, we consider soft group membership instead of hard group annotations. The group probabilities can be flexibly generated using either supervised learning or zero-shot approaches. Our framework accommodates samples with group membership ambiguity, offering stronger flexibility and generality than the prior art. We comprehensively evaluate PG-DRO on both image classification and natural language processing benchmarks, establishing superior performance