[Paper Review] Meta-Learned Invariant Risk Minimization

Jun-Hyun Bae, Inchul Choi+1 more|arXiv (Cornell University)|Mar 24, 2021

Domain Adaptation and Few-Shot Learning27 references4 citations

TL;DR

This paper proposes Meta-IRM, a meta-learning approach that directly optimizes the ideal bi-level Invariant Risk Minimization (IRM) objective without assuming a linear classifier, using Model-Agnostic Meta-Learning (MAML). By treating each environment as a task and enforcing invariance through disjoint inner and outer loop optimization, Meta-IRM achieves superior out-of-distribution (OOD) generalization—especially when data are scarce or spurious correlations are abundant—outperforming IRMv1 and all variants in both linear and non-linear regimes.

ABSTRACT

Empirical Risk Minimization (ERM) based machine learning algorithms have suffered from weak generalization performance on data obtained from out-of-distribution (OOD). To address this problem, Invariant Risk Minimization (IRM) objective was suggested to find invariant optimal predictor which is less affected by the changes in data distribution. However, even with such progress, IRMv1, the practical formulation of IRM, still shows performance degradation when there are not enough training data, and even fails to generalize to OOD, if the number of spurious correlations is larger than the number of environments. In this paper, to address such problems, we propose a novel meta-learning based approach for IRM. In this method, we do not assume the linearity of classifier for the ease of optimization, and solve ideal bi-level IRM objective with Model-Agnostic Meta-Learning (MAML) framework. Our method is more robust to the data with spurious correlations and can provide an invariant optimal classifier even when data from each distribution are scarce. In experiments, we demonstrate that our algorithm not only has better OOD generalization performance than IRMv1 and all IRM variants, but also addresses the weakness of IRMv1 with improved stability.

Motivation & Objective

• Address the poor OOD generalization of Empirical Risk Minimization (ERM) when training data contain spurious correlations.
• Overcome the limitations of IRMv1, which fails when the number of environments is smaller than the number of spurious correlations or when training data per environment are scarce.
• Develop a more robust and stable method to learn invariant predictors by directly optimizing the ideal bi-level IRM objective without simplifying assumptions.
• Enable effective OOD generalization in non-linear settings by removing the linearity constraint on the classifier.
• Improve optimization stability and performance through auxiliary regularization on meta-loss variance and second-order gradient updates.

Proposed method

• Formulate the ideal IRM objective as a bi-level optimization problem, treating each training environment as a separate task in meta-learning.
• Apply Model-Agnostic Meta-Learning (MAML) to solve the bi-level problem, using disjoint environments for inner loop (task-specific adaptation) and outer loop (meta-update).
• Relax the linear classifier assumption of IRMv1 by allowing non-linear models, enabling better representation learning and generalization.
• Introduce an auxiliary loss that minimizes the standard deviation of meta-losses across all environments to enforce invariance and improve training stability.
• Use second-order gradients in the meta-update to better approximate the true bi-level optimization, enhancing convergence and performance.
• Ensure invariance by optimizing the classifier to perform consistently across diverse environments, even when test data shift significantly.

Experimental results

Research questions

• RQ1 Can a meta-learning framework effectively optimize the ideal bi-level IRM objective without assuming a linear classifier?
• RQ2 Does the proposed meta-IRM method achieve better OOD generalization than IRMv1 and its variants, especially under data scarcity or high spurious correlation?
• RQ3 How critical is the use of second-order gradients and disjoint environment sampling in achieving stable and invariant predictors?
• RQ4 To what extent does the auxiliary loss based on meta-loss variance improve the robustness and performance of the model?
• RQ5 Can meta-IRM generalize effectively in non-linear settings where IRMv1 fails due to its restrictive linear assumption?

Key findings

• Meta-IRM achieves significantly better OOD generalization than IRMv1 and all IRM variants, with test accuracy of 54.5% on a test environment with p_e = 0.9 and p_e2 = 0.9, compared to IRMv1’s 13.4%.
• In the Colored MNIST task with two spurious features (color and patch), meta-IRM maintains consistent performance (54.5% ± 4.0) across all test environments, while IRMv1 collapses to 13.4% on the most out-of-distribution setting.
• Ablation studies show that removing the standard deviation loss reduces OOD performance (64.2% vs. 70.4% for meta-IRM), confirming its role in stabilizing invariance learning.
• Using first-order approximation for meta-updates degrades performance (59.1% vs. 70.4%), demonstrating the importance of second-order gradients for accurate bi-level optimization.
• Using the same environment for inner and outer loops leads to poor OOD generalization (13.6% accuracy), proving that disjoint environment sampling is essential for invariance.
• Meta-IRM achieves stable and invariant predictions even when the true causal features are noisy, showing robustness beyond the limitations of the original IRM objective.