Meta-Learned Invariant Risk Minimization

Abstract

Invariant Risk Minimization (IRM) aims to find predictors that generalize across out-of-distribution (OOD) environments by learning invariant correlations rather than spurious ones. We propose meta-IRM, a meta-learning based approach that instantiates the IRM objective using MAML. The key idea is assigning different training environments to the inner and outer optimization loops — since spurious correlations vary across environments, this naturally drives the model toward invariant features. We empirically evaluate meta-IRM on Colored MNIST variants including multi-class, insufficient data, and multiple spurious correlation settings.

Overview

We directly solve the ideal IRM bi-level optimization using the MAML framework, assigning different training environments to inner and outer loops to learn invariant features.

1. Environment-specific adaptation — Adapt model parameters to each training environment \(e_i\) in the inner loop.
2. Cross-environment evaluation — Evaluate the adapted parameters \(\theta_i'\) on a different environment \(e_j\) to compute the meta-loss.
3. Invariant convergence — Combining gradients from different environments cancels spurious correlations, converging toward invariant features.

Diagram of meta-IRM. Gradients from environment-adapted parameters are combined to converge toward the invariant feature direction ($\theta_{optimal}$).

Method

IRM aims to find an invariant predictor across multiple training environments. Meta-IRM approaches this using the MAML framework, with a key insight: assigning different training environments to the inner loop and outer loop. Since spurious correlations vary across environments, combining gradients from different environments drives convergence toward invariant features.

Meta-Learning Framework

The inner/outer structure of MAML maps onto the IRM objective. The inner loop adapts to one environment, while the outer loop (meta-loss) evaluates on a different environment.

Inner optimization: Adapt model parameters to each training environment \(e_i\) :

Meta-optimization: Update \(\theta\) using meta-losses computed from different environments \(e_j\) with adapted parameters \(\theta_i'\) :

Schematic of meta-IRM learning process. Model parameters $\theta$ are adapted to training environments in the inner optimization; each adapted parameter computes meta-loss from different environments.

Auxiliary Standard Deviation Loss

An auxiliary loss based on the standard deviation \(\sigma\) of meta-losses encourages uniform performance across all environments. This regularization improves training stability and convergence speed.

Results

Colored MNIST

Test accuracy (%) on Colored MNIST. $p_e=0.9$ is the OOD environment (reversed color-label correlation).

Accuracy across test environments ($p_e$ from 0.1 to 0.9).

Multi-Class Problem

Multi-class Colored MNIST ($k$=5, 10) test accuracy (%).

Insufficient Data

Results when training data is halved (25,000 to 12,500 per environment).

Test accuracy (%) with reduced training data (12,500 per environment).

Effect of decreasing training data (25,000 → 12,500 → 6,250).

Two Spurious Features

Results with 2 spurious correlations (color + patch) and 2 training environments.

Test accuracy with two spurious features (color + patch) across environments.

Ablation Study

BibTeX

@article{bae2021meta,
  author    = {Bae, Jun-Hyun and Choi, Inchul and Lee, Minho},
  title     = {Meta-Learned Invariant Risk Minimization},
  journal   = {arXiv preprint arXiv:2103.12947},
  year      = {2021}
}