[Paper Review] Invariant Risk Minimization Games

Kartik Ahuja, Karthikeyan Shanmugam+2 more|arXiv (Cornell University)|Feb 11, 2020

Gaussian Processes and Bayesian Inference35 references31 citations

TL;DR

The paper reframes invariant risk minimization as a Nash equilibrium problem among environments, introducing an ensemble game (EIRM) whose equilibria correspond to invariant predictors, and shows a BRD-based training algorithm with competitive performance and lower variance.

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.

Motivation & Objective

• Motivate robust prediction under distribution shift by exploiting invariances across environments.
• Formulate IRM as a multi-environment ensemble game to identify invariant predictors.
• Propose a practical best-response dynamics algorithm to compute Nash equilibria in the EIRM game.
• Prove theoretical equivalence between Nash equilibria and invariant predictors under mild assumptions.
• Demonstrate empirical gains in accuracy and variance on synthetic and real datasets compared to IRM and ERM baselines.

Proposed method

• Define an ensemble classifier w^av aggregating environment-specific classifiers w^e.
• Formulate a game where each environment minimizes its risk via BRD with respect to w^e and a shared representation Phi.
• Show equivalence: NE of the EIRM game corresponds to invariant predictors across environments (Theorem 1).
• Propose two training regimes: fixed Phi (F-IRM) and learned Phi (V-IRM) with BRD updates for w^e and Phi.
• Provide algorithm (Algorithm 1) using BRD with periodic Phi updates and environment-wise SGD steps.
• Discuss existence of NE under affine-closure (Assumption 1) and linear/hilbertian function classes (Theorem 3).

Experimental results

Research questions

• RQ1 Can invariant predictors across multiple environments be characterized as Nash equilibria in a multi-enviroment game?
• RQ2 Does a BRD-based training regime converge to invariant predictors, and under what conditions does NE exist?
• RQ3 What is the relationship between NE of EIRM and the invariant predictor set from IRM?
• RQ4 Does fixing or learning the representation Phi affect the ability to recover invariant predictors?
• RQ5 Do proposed methods generalize to nonlinear classifiers and provide robust performance under distribution shifts?

Key findings

• The ensemble game yields the same set of predictors as invariant predictors across environments (Theorem 1 and Corollary 1).
• BRD-based training with either fixed or learned representation achieves comparable or better accuracy with lower variance than prior IRM (e.g., Colored MNIST results).
• The approach accommodates nonlinear classifiers, not restricted to linear Phi or w, while preserving invariance guarantees.
• Fixing Phi (identity) and solving the EIRM game suffices to recover invariant predictors with bounded L^p norms (Theorem 2).
• Pure Nash equilibria exist for linear models under convex, continuous losses; in broader models, mixed equilibria exist (Theorem 3).
• Empirical results show F-IRM and V-IRM variants often outperform standard IRM and ERM baselines in robustness and stability across datasets.