Weighted Risk Invariance: Domain Generalization under Invariant Feature Shift

TL;DR

This paper introduces Weighted Risk Invariance (WRI), a new method that improves invariant model learning under covariate shifts by reweighting training examples, outperforming previous methods in linear-Gaussian settings.

Contribution

The paper proposes WRI, a novel framework with a practical algorithm that ensures invariant models under covariate shifts, addressing limitations of existing methods.

Findings

WRI provably learns invariant models in linear-Gaussian settings.

WRI outperforms previous methods under invariant covariate shift.

The algorithm learns density and model parameters simultaneously.

Abstract

Learning models whose predictions are invariant under multiple environments is a promising approach for out-of-distribution generalization. Such models are trained to extract features $X_{\text{inv}}$ where the conditional distribution $Y \mid X_{\text{inv}}$ of the label given the extracted features does not change across environments. Invariant models are also supposed to generalize to shifts in the marginal distribution $p(X_{\text{inv}})$ of the extracted features $X_{\text{inv}}$, a type of shift we call an $\textit{invariant covariate shift}$. However, we show that proposed methods for learning invariant models underperform under invariant covariate shift, either failing to learn invariant models$\unicode{x2014}$even for data generated from simple and well-studied linear-Gaussian models$\unicode{x2014}$or having poor finite-sample performance. To alleviate these problems, we propose $\textit{weighted risk invariance}$ (WRI). Our framework is based on imposing invariance of the loss across environments subject to appropriate reweightings of the training examples. We show that WRI provably learns invariant models, i.e. discards spurious correlations, in linear-Gaussian settings. We propose a practical algorithm to implement WRI by learning the density $p(X_{\text{inv}})$ and the model parameters simultaneously, and we demonstrate empirically that WRI outperforms previous invariant learning methods under invariant covariate shift.