Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Invariant Risk Minimization Games

Kartik Ahuja, Karthikeyan Shanmugam, Kush R. Varshney, Amit Dhurandhar

TL;DR

Invariant Risk Minimization Games recasts the pursuit of invariant predictors across multiple environments as a multi-agent Nash equilibrium problem (EIRM). Each environment selects a component of an ensemble predictor, and the ensemble's average forms the shared decision rule; under affine-closure, the NE solutions coincide with the invariant predictors, enabling nonlinear representations and simple best-response training. The approach yields comparable or improved accuracy with substantially reduced variance relative to the original IRM formulation and maintains generalization guarantees across unseen environments. Empirical results on colored-MNIST-like datasets and structured-noise variants demonstrate the method's robustness to spurious correlations and illuminate oscillatory dynamics inherent to best-response training, which can be stabilized through termination criteria and schedule design.

Abstract

The standard risk minimization paradigm of machine learning is brittle when operating in environments whose test distributions are different from the training distribution due to spurious correlations. Training on data from many environments and finding invariant predictors reduces the effect of spurious features by concentrating models on features that have a causal relationship with the outcome. In this work, we pose such invariant risk minimization as finding the Nash equilibrium of an ensemble game among several environments. By doing so, we develop a simple training algorithm that uses best response dynamics and, in our experiments, yields similar or better empirical accuracy with much lower variance than the challenging bi-level optimization problem of Arjovsky et al. (2019). One key theoretical contribution is showing that the set of Nash equilibria for the proposed game are equivalent to the set of invariant predictors for any finite number of environments, even with nonlinear classifiers and transformations. As a result, our method also retains the generalization guarantees to a large set of environments shown in Arjovsky et al. (2019). The proposed algorithm adds to the collection of successful game-theoretic machine learning algorithms such as generative adversarial networks.

Invariant Risk Minimization Games

TL;DR

Invariant Risk Minimization Games recasts the pursuit of invariant predictors across multiple environments as a multi-agent Nash equilibrium problem (EIRM). Each environment selects a component of an ensemble predictor, and the ensemble's average forms the shared decision rule; under affine-closure, the NE solutions coincide with the invariant predictors, enabling nonlinear representations and simple best-response training. The approach yields comparable or improved accuracy with substantially reduced variance relative to the original IRM formulation and maintains generalization guarantees across unseen environments. Empirical results on colored-MNIST-like datasets and structured-noise variants demonstrate the method's robustness to spurious correlations and illuminate oscillatory dynamics inherent to best-response training, which can be stabilized through termination criteria and schedule design.

Abstract

The standard risk minimization paradigm of machine learning is brittle when operating in environments whose test distributions are different from the training distribution due to spurious correlations. Training on data from many environments and finding invariant predictors reduces the effect of spurious features by concentrating models on features that have a causal relationship with the outcome. In this work, we pose such invariant risk minimization as finding the Nash equilibrium of an ensemble game among several environments. By doing so, we develop a simple training algorithm that uses best response dynamics and, in our experiments, yields similar or better empirical accuracy with much lower variance than the challenging bi-level optimization problem of Arjovsky et al. (2019). One key theoretical contribution is showing that the set of Nash equilibria for the proposed game are equivalent to the set of invariant predictors for any finite number of environments, even with nonlinear classifiers and transformations. As a result, our method also retains the generalization guarantees to a large set of environments shown in Arjovsky et al. (2019). The proposed algorithm adds to the collection of successful game-theoretic machine learning algorithms such as generative adversarial networks.

Paper Structure

This paper contains 32 sections, 13 theorems, 33 equations, 40 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Invariance Principles in Causality
  4. Sample Reweighting, Domain Adaptation, and Robust Optimization
  5. Preliminaries
  6. Game Theory Concepts
  7. Invariant Risk Minimization
  8. Ensemble Invariant Risk Minimization Games
  9. Game-Theoretic Reformulation
  10. Equivalence Between NE and Invariant Predictors
  11. Existence of NE of $\Gamma^{\mathsf{EIRM}}$ and Invariant Predictors
  12. Algorithms for Finding NE of $\Gamma^{\mathsf{EIRM}}$
  13. Experiments
  14. Benchmarks
  15. Datasets
  16. ...and 17 more sections

Key Result

Theorem 1

If Assumption ass:affine closure holds, then $\tilde{\mathcal{S}}^{\mathsf{IV}} = \tilde{\mathcal{S}}^{\mathsf{EIRM}}$

Figures (40)

  • Figure 1: Illustration of best response training with 2 environments and representation learner. Dotted lines for backpropagation and solid lines for forward pass.
  • Figure 2: F-IRM, Colored Fashion MNIST: Comparing accuracy of ensemble
  • Figure 3: F-IRM, Colored Fashion MNIST: Correlation of the ensemble model with color
  • Figure 4: F-IRM, Colored Fashion MNIST: Correlations of the individual models with color
  • Figure 5: F-IRM, Colored Fashion MNIST: Comparing accuracy of ensemble
  • ...and 35 more figures

Theorems & Definitions (19)

  • Theorem 1
  • Corollary 1
  • Theorem 2
  • Theorem 3
  • Theorem 1
  • proof
  • Corollary 1
  • proof
  • Lemma 1
  • proof
  • ...and 9 more