# UMIX: Improving Importance Weighting for Subpopulation Shift via Uncertainty-Aware Mixup

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

## Key facts

- Cited by 34 papers in the map (data through 14 Oct 2026).
- Builds on 50 of its 85 references that are in the map.
- First posted on arXiv 19 Sep 2022.

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

- [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)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features](https://mlsanity.com/paper/1905.04899/) (2019)
- [What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?](https://mlsanity.com/paper/1703.04977/) (2017)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics](https://mlsanity.com/paper/1705.07115/) (2017)
- [Class-Balanced Loss Based on Effective Number of Samples](https://mlsanity.com/paper/1901.05555/) (2019)
- [Deep CORAL: Correlation Alignment for Deep Domain Adaptation](https://mlsanity.com/paper/1607.01719/) (2016)
- [Learning Imbalanced Datasets with Label-Distribution-Aware Margin Loss](https://mlsanity.com/paper/1906.07413/) (2019)
- [Focal Loss for Dense Object Detection](https://mlsanity.com/paper/1708.02002/) (2017)
- [Out-of-Distribution Generalization via Risk Extrapolation (REx)](https://mlsanity.com/paper/2003.00688/) (2020)
- [An Empirical Study of Example Forgetting during Deep Neural Network Learning](https://mlsanity.com/paper/1812.05159/) (2018)
- [A Simple Baseline for Bayesian Uncertainty in Deep Learning](https://mlsanity.com/paper/1902.02476/) (2019)
- [Just Train Twice: Improving Group Robustness without Training Group Information](https://mlsanity.com/paper/2107.09044/) (2021)
- [Nuanced Metrics for Measuring Unintended Bias with Real Data for Text Classification](https://mlsanity.com/paper/1903.04561/) (2019)
- [Fairness Without Demographics in Repeated Loss Minimization](https://mlsanity.com/paper/1806.08010/) (2018)
- [Snapshot Ensembles: Train 1, get M for free](https://mlsanity.com/paper/1704.00109/) (2017)
- [An Investigation of Why Overparameterization Exacerbates Spurious Correlations](https://mlsanity.com/paper/2005.04345/) (2020)
- [Gradient Matching for Domain Generalization](https://mlsanity.com/paper/2104.09937/) (2021)
- [What is the Effect of Importance Weighting in Deep Learning?](https://mlsanity.com/paper/1812.03372/) (2018)
- [Fairness without Demographics through Adversarially Reweighted Learning](https://mlsanity.com/paper/2006.13114/) (2020)
- [Adversarial Domain Adaptation with Domain Mixup](https://mlsanity.com/paper/1912.01805/) (2019)
- [Invariance Principle Meets Information Bottleneck for Out-of-Distribution Generalization](https://mlsanity.com/paper/2106.06607/) (2021)

## Built on by (most cited first, 30 of 34)

- [Change is Hard: A Closer Look at Subpopulation Shift](https://mlsanity.com/paper/2302.12254/) (2023)
- [Provable Dynamic Fusion for Low-Quality Multimodal Data](https://mlsanity.com/paper/2306.02050/) (2023)
- [Discover and Cure: Concept-aware Mitigation of Spurious Correlation](https://mlsanity.com/paper/2305.00650/) (2023)
- [Avoiding spurious correlations via logit correction](https://mlsanity.com/paper/2212.01433/) (2022)
- [It's Never Too Late: Fusing Acoustic Information into Large Language Models for Automatic Speech Recognition](https://mlsanity.com/paper/2402.05457/) (2024)
- [The Pitfalls of Memorization: When Memorization Hurts Generalization](https://mlsanity.com/paper/2412.07684/) (2024)
- [Nuisances via Negativa: Adjusting for Spurious Correlations via Data Augmentation](https://mlsanity.com/paper/2210.01302/) (2022)
- [Improving Medical Multi-modal Contrastive Learning with Expert Annotations](https://mlsanity.com/paper/2403.10153/) (2024)
- [A Short Survey on Importance Weighting for Machine Learning](https://mlsanity.com/paper/2403.10175/) (2024)
- [Learning Robust Classifiers with Self-Guided Spurious Correlation Mitigation](https://mlsanity.com/paper/2405.03649/) (2024)
- [Skip n: A Simple Method to Reduce Hallucination in Large Vision-Language Models](https://mlsanity.com/paper/2402.01345/) (2024)
- [Data Augmentation via Subgroup Mixup for Improving Fairness](https://mlsanity.com/paper/2309.07110/) (2023)
- [Spuriousness-Aware Meta-Learning for Learning Robust Classifiers](https://mlsanity.com/paper/2406.10742/) (2024)
- [Black Sheep in the Herd: Playing with Spuriously Correlated Attributes for Vision-Language Recognition](https://mlsanity.com/paper/2502.15809/) (2025)
- [Sufficient Invariant Learning for Distribution Shift](https://mlsanity.com/paper/2210.13533/) (2022)
- [Data Optimization in Deep Learning: A Survey](https://mlsanity.com/paper/2310.16499/) (2023)
- [Out-Of-Distribution Detection with Diversification (Provably)](https://mlsanity.com/paper/2411.14049/) (2024)
- [Selective Learning: Towards Robust Calibration with Dynamic Regularization](https://mlsanity.com/paper/2402.08384/) (2024)
- [Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift](https://mlsanity.com/paper/2505.23027/) (2025)
- [Error Slice Discovery via Manifold Compactness](https://mlsanity.com/paper/2501.19032/) (2025)
- [Unsupervised Variational Translator for Bridging Image Restoration and High-Level Vision Tasks](https://mlsanity.com/paper/2408.08149/) (2024)
- [Reweighted Mixup for Subpopulation Shift](https://mlsanity.com/paper/2304.04148/) (2023)
- [FairlyUncertain: A Comprehensive Benchmark of Uncertainty in Algorithmic Fairness](https://mlsanity.com/paper/2410.02005/) (2024)
- [Mitigating Spurious Correlations via Disagreement Probability](https://mlsanity.com/paper/2411.01757/) (2024)
- [MaskMedPaint: Masked Medical Image Inpainting with Diffusion Models for Mitigation of Spurious Correlations](https://mlsanity.com/paper/2411.10686/) (2024)
- [NeuronTune: Towards Self-Guided Spurious Bias Mitigation](https://mlsanity.com/paper/2505.24048/) (2025)
- [Boosting Test Performance with Importance Sampling--a Subpopulation Perspective](https://mlsanity.com/paper/2412.13003/) (2024)
- [Enhancing Multi-view Open-set Learning via Ambiguity Uncertainty Calibration and View-wise Debiasing](https://mlsanity.com/paper/2508.01227/) (2025)
- [Semantic Equivariant Mixup](https://mlsanity.com/paper/2308.06451/) (2023)
- [Effective Guidance for Model Attention with Simple Yes-no Annotations](https://mlsanity.com/paper/2410.22312/) (2024)

## Abstract (arXiv)

Subpopulation shift widely exists in many real-world machine learning applications, referring to the training and test distributions containing the same subpopulation groups but varying in subpopulation frequencies. Importance reweighting is a normal way to handle the subpopulation shift issue by imposing constant or adaptive sampling weights on each sample in the training dataset. However, some recent studies have recognized that most of these approaches fail to improve the performance over empirical risk minimization especially when applied to over-parameterized neural networks. In this work, we propose a simple yet practical framework, called uncertainty-aware mixup (UMIX), to mitigate the overfitting issue in over-parameterized models by reweighting the ''mixed'' samples according to the sample uncertainty. The training-trajectories-based uncertainty estimation is equipped in the proposed UMIX for each sample to flexibly characterize the subpopulation distribution. We also provide insightful theoretical analysis to verify that UMIX achieves better generalization bounds over prior works. Further, we conduct extensive empirical studies across a wide range of tasks to validate the effectiveness of our method both qualitatively and quantitatively. Code is available at https://github.com/TencentAILabHealthcare/UMIX.