# Robust Invariant Representation Learning by Distribution Extrapolation

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 33 references that are in the map.
- First posted on arXiv 22 May 2025.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Domain-Adversarial Training of Neural Networks](https://mlsanity.com/paper/1505.07818/) (2015)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Can You Trust Your Model's Uncertainty? Evaluating Predictive Uncertainty Under Dataset Shift](https://mlsanity.com/paper/1906.02530/) (2019)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [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)
- [Out-of-Distribution Generalization via Risk Extrapolation (REx)](https://mlsanity.com/paper/2003.00688/) (2020)
- [Measuring Calibration in Deep Learning](https://mlsanity.com/paper/1904.01685/) (2019)
- [Revisiting the Calibration of Modern Neural Networks](https://mlsanity.com/paper/2106.07998/) (2021)
- [The Effectiveness of Data Augmentation in Image Classification using Deep Learning](https://mlsanity.com/paper/1712.04621/) (2017)
- [Ensemble deep learning: A review](https://mlsanity.com/paper/2104.02395/) (2021)
- [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)
- [On Calibration and Out-of-domain Generalization](https://mlsanity.com/paper/2102.10395/) (2021)
- [Scalable Marginal Likelihood Estimation for Model Selection in Deep Learning](https://mlsanity.com/paper/2104.04975/) (2021)
- [OoD-Bench: Quantifying and Understanding Two Dimensions of Out-of-Distribution Generalization](https://mlsanity.com/paper/2106.03721/) (2021)
- [Out of Distribution Generalization in Machine Learning](https://mlsanity.com/paper/2103.02667/) (2021)
- [Pareto Invariant Risk Minimization: Towards Mitigating the Optimization Dilemma in Out-of-Distribution Generalization](https://mlsanity.com/paper/2206.07766/) (2022)
- [What Is Missing in IRM Training and Evaluation? Challenges and Solutions](https://mlsanity.com/paper/2303.02343/) (2023)
- [An Empirical Study of Pre-trained Model Selection for Out-of-Distribution Generalization and Calibration](https://mlsanity.com/paper/2307.08187/) (2023)
- [Towards Understanding Variants of Invariant Risk Minimization through the Lens of Calibration](https://mlsanity.com/paper/2401.17541/) (2024)

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

- [Invariant Gradient Alignment for Robust Reasoning Distillation](https://mlsanity.com/paper/2606.05025/) (2026)

## Abstract (arXiv)

Invariant risk minimization (IRM) aims to enable out-of-distribution (OOD) generalization in deep learning by learning invariant representations. As IRM poses an inherently challenging bi-level optimization problem, most existing approaches -- including IRMv1 -- adopt penalty-based single-level approximations. However, empirical studies consistently show that these methods often fail to outperform well-tuned empirical risk minimization (ERM), highlighting the need for more robust IRM implementations. This work theoretically identifies a key limitation common to many IRM variants: their penalty terms are highly sensitive to limited environment diversity and over-parameterization, resulting in performance degradation. To address this issue, a novel extrapolation-based framework is proposed that enhances environmental diversity by augmenting the IRM penalty through synthetic distributional shifts. Extensive experiments -- ranging from synthetic setups to realistic, over-parameterized scenarios -- demonstrate that the proposed method consistently outperforms state-of-the-art IRM variants, validating its effectiveness and robustness.