# Out-of-distribution Generalization for Total Variation based Invariant Risk Minimization

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 16 of its 27 references that are in the map.
- First posted on arXiv 27 Feb 2025.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Do ImageNet Classifiers Generalize to ImageNet?](https://mlsanity.com/paper/1902.10811/) (2019)
- [Out-of-Distribution Generalization via Risk Extrapolation (REx)](https://mlsanity.com/paper/2003.00688/) (2020)
- [Environment Inference for Invariant Learning](https://mlsanity.com/paper/2010.07249/) (2020)
- [Invariant Risk Minimization Games](https://mlsanity.com/paper/2002.04692/) (2020)
- [Heterogeneous Risk Minimization](https://mlsanity.com/paper/2105.03818/) (2021)
- [ZIN: When and How to Learn Invariance Without Environment Partition?](https://mlsanity.com/paper/2203.05818/) (2022)
- [In-N-Out: Pre-Training and Self-Training using Auxiliary Information for Out-of-Distribution Robustness](https://mlsanity.com/paper/2012.04550/) (2020)
- [Harnessing Out-Of-Distribution Examples via Augmenting Content and Style](https://mlsanity.com/paper/2207.03162/) (2022)
- [Meta-Learning for Few-Shot Land Cover Classification](https://mlsanity.com/paper/2004.13390/) (2020)
- [Spurious Feature Diversification Improves Out-of-distribution Generalization](https://mlsanity.com/paper/2309.17230/) (2023)
- [Invariant Learning via Probability of Sufficient and Necessary Causes](https://mlsanity.com/paper/2309.12559/) (2023)
- [On the Connection between Invariant Learning and Adversarial Training for Out-of-Distribution Generalization](https://mlsanity.com/paper/2212.09082/) (2022)
- [Invariant Risk Minimization Is A Total Variation Model](https://mlsanity.com/paper/2405.01389/) (2024)

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

- [Out-of-Distribution Generalization in Time Series: A Survey](https://mlsanity.com/paper/2503.13868/) (2025)
- [Generalizable Multimodal Large Language Model Editing via Invariant Trajectory Learning](https://mlsanity.com/paper/2601.19700/) (2026)
- [Environment-Conditioned Tail Reweighting for Invariant Learning under Mixed Shifts](https://mlsanity.com/paper/2601.22944/) (2026)

## Abstract (arXiv)

Invariant risk minimization is an important general machine learning framework that has recently been interpreted as a total variation model (IRM-TV). However, how to improve out-of-distribution (OOD) generalization in the IRM-TV setting remains unsolved. In this paper, we extend IRM-TV to a Lagrangian multiplier model named OOD-TV-IRM. We find that the autonomous TV penalty hyperparameter is exactly the Lagrangian multiplier. Thus OOD-TV-IRM is essentially a primal-dual optimization model, where the primal optimization minimizes the entire invariant risk and the dual optimization strengthens the TV penalty. The objective is to reach a semi-Nash equilibrium where the balance between the training loss and OOD generalization is maintained. We also develop a convergent primal-dual algorithm that facilitates an adversarial learning scheme. Experimental results show that OOD-TV-IRM outperforms IRM-TV in most situations.