# Out-of-distribution Prediction with Invariant Risk Minimization: The Limitation and An Effective Fix

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

## Key facts

- Cited by 27 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 35 references that are in the map.
- First posted on arXiv 16 Jan 2021.

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

- [Unsupervised Domain Adaptation by Backpropagation](https://mlsanity.com/paper/1409.7495/) (2014)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Learning Transferable Features with Deep Adaptation Networks](https://mlsanity.com/paper/1502.02791/) (2015)
- [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)
- [Causal inference using invariant prediction: identification and confidence intervals](https://mlsanity.com/paper/1501.01332/) (2015)
- [Wasserstein Auto-Encoders](https://mlsanity.com/paper/1711.01558/) (2017)
- [Environment Inference for Invariant Learning](https://mlsanity.com/paper/2010.07249/) (2020)
- [Learning to Learn Single Domain Generalization](https://mlsanity.com/paper/2003.13216/) (2020)
- [Invariant Risk Minimization Games](https://mlsanity.com/paper/2002.04692/) (2020)
- [Invariant Rationalization](https://mlsanity.com/paper/2003.09772/) (2020)
- [Invariant Causal Prediction for Block MDPs](https://mlsanity.com/paper/2003.06016/) (2020)
- [External Validity: From Do-Calculus to Transportability Across Populations](https://mlsanity.com/paper/1503.01603/) (2015)
- [Stable Prediction across Unknown Environments](https://mlsanity.com/paper/1806.06270/) (2018)
- [A Causal View on Robustness of Neural Networks](https://mlsanity.com/paper/2005.01095/) (2020)
- [Learning Invariances using the Marginal Likelihood](https://mlsanity.com/paper/1808.05563/) (2018)
- [Enforcing Predictive Invariance across Structured Biomedical Domains](https://mlsanity.com/paper/2006.03908/) (2020)
- [Selecting Data Augmentation for Simulating Interventions](https://mlsanity.com/paper/2005.01856/) (2020)
- [Novel Human-Object Interaction Detection via Adversarial Domain Generalization](https://mlsanity.com/paper/2005.11406/) (2020)

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

- [Generalizing to Unseen Domains: A Survey on Domain Generalization](https://mlsanity.com/paper/2103.03097/) (2021)
- [Improving Out-of-Distribution Robustness via Selective Augmentation](https://mlsanity.com/paper/2201.00299/) (2022)
- [Fishr: Invariant Gradient Variances for Out-of-Distribution Generalization](https://mlsanity.com/paper/2109.02934/) (2021)
- [On Calibration and Out-of-domain Generalization](https://mlsanity.com/paper/2102.10395/) (2021)
- [C-Mixup: Improving Generalization in Regression](https://mlsanity.com/paper/2210.05775/) (2022)
- [Causally motivated Shortcut Removal Using Auxiliary Labels](https://mlsanity.com/paper/2105.06422/) (2021)
- [Out-of-distribution Generalization in the Presence of Nuisance-Induced Spurious Correlations](https://mlsanity.com/paper/2107.00520/) (2021)
- [Freeze then Train: Towards Provable Representation Learning under Spurious Correlations and Feature Noise](https://mlsanity.com/paper/2210.11075/) (2022)
- [Regularizing towards Causal Invariance: Linear Models with Proxies](https://mlsanity.com/paper/2103.02477/) (2021)
- [Causal Balancing for Domain Generalization](https://mlsanity.com/paper/2206.05263/) (2022)
- [Optimization-based Causal Estimation from Heterogenous Environments](https://mlsanity.com/paper/2109.11990/) (2021)
- [Invariant Language Modeling](https://mlsanity.com/paper/2110.08413/) (2021)
- [Domain Generalization -- A Causal Perspective](https://mlsanity.com/paper/2209.15177/) (2022)
- [Improving Domain Generalization with Domain Relations](https://mlsanity.com/paper/2302.02609/) (2023)
- [Data Augmentations for Improved (Large) Language Model Generalization](https://mlsanity.com/paper/2310.12803/) (2023)
- [Learning for Counterfactual Fairness from Observational Data](https://mlsanity.com/paper/2307.08232/) (2023)
- [Finding Valid Adjustments under Non-ignorability with Minimal DAG Knowledge](https://mlsanity.com/paper/2106.11560/) (2021)
- [Fairness and robustness in anti-causal prediction](https://mlsanity.com/paper/2209.09423/) (2022)
- [Robustness to Spurious Correlations Improves Semantic Out-of-Distribution Detection](https://mlsanity.com/paper/2302.04132/) (2023)
- [Multi-Domain Long-Tailed Learning by Augmenting Disentangled Representations](https://mlsanity.com/paper/2210.14358/) (2022)
- [Knowledge Guided Representation Learning and Causal Structure Learning in Soil Science](https://mlsanity.com/paper/2306.09302/) (2023)
- [Conditionally Invariant Representation Learning for Disentangling Cellular Heterogeneity](https://mlsanity.com/paper/2307.00558/) (2023)
- [Joint covariate-alignment and concept-alignment: a framework for domain generalization](https://mlsanity.com/paper/2208.00898/) (2022)
- [Weighted Risk Invariance: Domain Generalization under Invariant Feature Shift](https://mlsanity.com/paper/2407.18428/) (2024)
- [Bridging Distribution Shift and AI Safety: Conceptual and Methodological Synergies](https://mlsanity.com/paper/2505.22829/) (2025)
- [HCDG: A Hierarchical Consistency Framework for Domain Generalization on Medical Image Segmentation](https://mlsanity.com/paper/2109.05742/) (2021)
- [Trade-off between reconstruction loss and feature alignment for domain generalization](https://mlsanity.com/paper/2210.15000/) (2022)

## Abstract (arXiv)

This work considers the out-of-distribution (OOD) prediction problem where (1)~the training data are from multiple domains and (2)~the test domain is unseen in the training. DNNs fail in OOD prediction because they are prone to pick up spurious correlations. Recently, Invariant Risk Minimization (IRM) is proposed to address this issue. Its effectiveness has been demonstrated in the colored MNIST experiment. Nevertheless, we find that the performance of IRM can be dramatically degraded under strong Lambda spuriousness -- when the spurious correlation between the spurious features and the class label is strong due to the strong causal influence of their common cause, the domain label, on both of them (see Fig. 1). In this work, we try to answer the questions: why does IRM fail in the aforementioned setting? Why does IRM work for the original colored MNIST dataset? How can we fix this problem of IRM? Then, we propose a simple and effective approach to fix the problem of IRM. We combine IRM with conditional distribution matching to avoid a specific type of spurious correlation under strong Lambda spuriousness. Empirically, we design a series of semi synthetic datasets -- the colored MNIST plus, which exposes the problems of IRM and demonstrates the efficacy of the proposed method.