# Learning Optimal Features via Partial Invariance

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 50 references that are in the map.
- First posted on arXiv 28 Jan 2023.

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

- [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)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Don't Stop Pretraining: Adapt Language Models to Domains and Tasks](https://mlsanity.com/paper/2004.10964/) (2020)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [In Search of Lost Domain Generalization](https://mlsanity.com/paper/2007.01434/) (2020)
- [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)
- [Agnostic Federated Learning](https://mlsanity.com/paper/1902.00146/) (2019)
- [Multi-Task Identification of Entities, Relations, and Coreference for Scientific Knowledge Graph Construction](https://mlsanity.com/paper/1808.09602/) (2018)
- [The Risks of Invariant Risk Minimization](https://mlsanity.com/paper/2010.05761/) (2020)
- [Word Embeddings Quantify 100 Years of Gender and Ethnic Stereotypes](https://mlsanity.com/paper/1711.08412/) (2017)
- [Domain Generalization using Causal Matching](https://mlsanity.com/paper/2006.07500/) (2020)
- [Invariance Principle Meets Information Bottleneck for Out-of-Distribution Generalization](https://mlsanity.com/paper/2106.06607/) (2021)
- [Deep Learning: A Critical Appraisal](https://mlsanity.com/paper/1801.00631/) (2018)
- [Invariant Risk Minimization Games](https://mlsanity.com/paper/2002.04692/) (2020)
- [Statistics of Robust Optimization: A Generalized Empirical Likelihood Approach](https://mlsanity.com/paper/1610.03425/) (2016)
- [Domain Generalization via Conditional Invariant Representation](https://mlsanity.com/paper/1807.08479/) (2018)
- [Invariant Causal Prediction for Nonlinear Models](https://mlsanity.com/paper/1706.08576/) (2017)
- [Does Invariant Risk Minimization Capture Invariance?](https://mlsanity.com/paper/2101.01134/) (2021)
- [Time Waits for No One! Analysis and Challenges of Temporal Misalignment](https://mlsanity.com/paper/2111.07408/) (2021)
- [MetaShift: A Dataset of Datasets for Evaluating Contextual Distribution Shifts and Training Conflicts](https://mlsanity.com/paper/2202.06523/) (2022)
- [Empirical or Invariant Risk Minimization? A Sample Complexity Perspective](https://mlsanity.com/paper/2010.16412/) (2020)
- [Fairness and Robustness in Invariant Learning: A Case Study in Toxicity Classification](https://mlsanity.com/paper/2011.06485/) (2020)
- [Enforcing Predictive Invariance across Structured Biomedical Domains](https://mlsanity.com/paper/2006.03908/) (2020)
- [On the benefits of representation regularization in invariance based domain generalization](https://mlsanity.com/paper/2105.14529/) (2021)
- [Invariant Language Modeling](https://mlsanity.com/paper/2110.08413/) (2021)

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

- [Distribution Shift Is Key to Learning Invariant Prediction](https://mlsanity.com/paper/2601.12296/) (2026)
- [Efficient Model-Agnostic Multi-Group Equivariant Networks](https://mlsanity.com/paper/2310.09675/) (2023)
- [Multi-Disease-Aware Training Strategy for Cardiac MR Image Segmentation](https://mlsanity.com/paper/2503.17896/) (2025)
- [How Many Samples Are Enough for Learning Across Domains?](https://mlsanity.com/paper/2609.39336/) (2026)
- [Can Domain Generalization be Guaranteed in Small-Sample Learning?](https://mlsanity.com/paper/2609.39512/) (2026)

## Abstract (arXiv)

Learning models that are robust to distribution shifts is a key concern in the context of their real-life applicability. Invariant Risk Minimization (IRM) is a popular framework that aims to learn robust models from multiple environments. The success of IRM requires an important assumption: the underlying causal mechanisms/features remain invariant across environments. When not satisfied, we show that IRM can over-constrain the predictor and to remedy this, we propose a relaxation via textitpartial invariance. In this work, we theoretically highlight the sub-optimality of IRM and then demonstrate how learning from a partition of training domains can help improve invariant models. Several experiments, conducted both in linear settings as well as with deep neural networks on tasks over both language and image data, allow us to verify our conclusions.