# Weighted Risk Invariance: Domain Generalization under Invariant Feature Shift

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 47 of its 73 references that are in the map.
- First posted on arXiv 25 Jul 2024.

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

- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [Domain-Adversarial Training of Neural Networks](https://mlsanity.com/paper/1505.07818/) (2015)
- [Benchmarking Neural Network Robustness to Common Corruptions and Perturbations](https://mlsanity.com/paper/1903.12261/) (2019)
- [Shortcut Learning in Deep Neural Networks](https://mlsanity.com/paper/2004.07780/) (2020)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Moment Matching for Multi-Source Domain Adaptation](https://mlsanity.com/paper/1812.01754/) (2018)
- [Learning Transferable Features with Deep Adaptation Networks](https://mlsanity.com/paper/1502.02791/) (2015)
- [Adversarial Discriminative Domain Adaptation](https://mlsanity.com/paper/1702.05464/) (2017)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Deep CORAL: Correlation Alignment for Deep Domain Adaptation](https://mlsanity.com/paper/1607.01719/) (2016)
- [Deep Hashing Network for Unsupervised Domain Adaptation](https://mlsanity.com/paper/1706.07522/) (2017)
- [In Search of Lost Domain Generalization](https://mlsanity.com/paper/2007.01434/) (2020)
- [Deeper, Broader and Artier Domain Generalization](https://mlsanity.com/paper/1710.03077/) (2017)
- [Aleatoric and Epistemic Uncertainty in Machine Learning: An Introduction to Concepts and Methods](https://mlsanity.com/paper/1910.09457/) (2019)
- [Learning to Generalize: Meta-Learning for Domain Generalization](https://mlsanity.com/paper/1710.03463/) (2017)
- [Out-of-Distribution Generalization via Risk Extrapolation (REx)](https://mlsanity.com/paper/2003.00688/) (2020)
- [Recognition in Terra Incognita](https://mlsanity.com/paper/1807.04975/) (2018)
- [Domain Generalization via Invariant Feature Representation](https://mlsanity.com/paper/1301.2115/) (2013)
- [Causal inference using invariant prediction: identification and confidence intervals](https://mlsanity.com/paper/1501.01332/) (2015)
- [The Risks of Invariant Risk Minimization](https://mlsanity.com/paper/2010.05761/) (2020)
- [Invariance Principle Meets Information Bottleneck for Out-of-Distribution Generalization](https://mlsanity.com/paper/2106.06607/) (2021)
- [Invariant Risk Minimization Games](https://mlsanity.com/paper/2002.04692/) (2020)
- [A Fine-Grained Analysis on Distribution Shift](https://mlsanity.com/paper/2110.11328/) (2021)
- [Posterior Network: Uncertainty Estimation without OOD Samples via Density-Based Pseudo-Counts](https://mlsanity.com/paper/2006.09239/) (2020)
- [On Calibration and Out-of-domain Generalization](https://mlsanity.com/paper/2102.10395/) (2021)
- [Domain Generalization by Mutual-Information Regularization with Pre-trained Models](https://mlsanity.com/paper/2203.10789/) (2022)
- [Does Invariant Risk Minimization Capture Invariance?](https://mlsanity.com/paper/2101.01134/) (2021)
- [Heterogeneous Risk Minimization](https://mlsanity.com/paper/2105.03818/) (2021)
- [OoD-Bench: Quantifying and Understanding Two Dimensions of Out-of-Distribution Generalization](https://mlsanity.com/paper/2106.03721/) (2021)
- [Learning Invariant Representations and Risks for Semi-supervised Domain Adaptation](https://mlsanity.com/paper/2010.04647/) (2020)

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

- [Environment-Conditioned Tail Reweighting for Invariant Learning under Mixed Shifts](https://mlsanity.com/paper/2601.22944/) (2026)
- [S2Aligner: Pair-Efficient and Transferable Pre-Training for Sparse Text-Attributed Graphs](https://mlsanity.com/paper/2605.18579/) (2026)

## Abstract (arXiv)

Learning models whose predictions are invariant under multiple environments is a promising approach for out-of-distribution generalization. Such models are trained to extract features Xinv where the conditional distribution Y mid Xinv of the label given the extracted features does not change across environments. Invariant models are also supposed to generalize to shifts in the marginal distribution p(Xinv) of the extracted features Xinv, a type of shift we call an textitinvariant covariate shift. However, we show that proposed methods for learning invariant models underperform under invariant covariate shift, either failing to learn invariant modelsunicodex2014even for data generated from simple and well-studied linear-Gaussian modelsunicodex2014or having poor finite-sample performance. To alleviate these problems, we propose textitweighted risk invariance (WRI). Our framework is based on imposing invariance of the loss across environments subject to appropriate reweightings of the training examples. We show that WRI provably learns invariant models, i.e. discards spurious correlations, in linear-Gaussian settings. We propose a practical algorithm to implement WRI by learning the density p(Xinv) and the model parameters simultaneously, and we demonstrate empirically that WRI outperforms previous invariant learning methods under invariant covariate shift.