Distribution Shift Is Key to Learning Invariant Prediction

TL;DR

This paper investigates why ERM can outperform specialized out-of-distribution methods when training data exhibit distribution shift across domains. By combining analytical bounds with a Gaussian-deductive example, it shows that the degree of shift, quantified via KL divergence, can drive learned models toward invariant predictors, especially under causality-related data assumptions. Theoretical results (Theorems 1–2) bound cross-environment misalignment under KL-bounded and Massart-noisy conditions, and corollaries indicate that sufficiently large shift enables ERM to approximate invariant prediction. Empirically, synthetic and CMNIST experiments demonstrate that more domains and greater distribution shift yield predictions closer to ground-truth or oracle invariant models, with counterfactual analyses revealing reduced reliance on non-causal features; these findings highlight distribution shift as a key determinant of robust cross-domain generalization.

Abstract

An interesting phenomenon arises: Empirical Risk Minimization (ERM) sometimes outperforms methods specifically designed for out-of-distribution tasks. This motivates an investigation into the reasons behind such behavior beyond algorithmic design. In this study, we find that one such reason lies in the distribution shift across training domains. A large degree of distribution shift can lead to better performance even under ERM. Specifically, we derive several theoretical and empirical findings demonstrating that distribution shift plays a crucial role in model learning and benefits learning invariant prediction. Firstly, the proposed upper bounds indicate that the degree of distribution shift directly affects the prediction ability of the learned models. If it is large, the models' ability can increase, approximating invariant prediction models that make stable predictions under arbitrary known or unseen domains; and vice versa. We also prove that, under certain data conditions, ERM solutions can achieve performance comparable to that of invariant prediction models. Secondly, the empirical validation results demonstrated that the predictions of learned models approximate those of Oracle or Optimal models, provided that the degree of distribution shift in the training data increases.

Figures (3)

• Figure 1: Models trained using ERM on the datasets (a) $D1$ and (b) $D2$. As observed, since the shift degree of distributions in $D2$ is larger than that in $D1$, the $\omega$ in (b) approaches $\omega^*$ more closely than the model in (a). Here, $\omega^*$ denotes the preset GT model.
• Figure 2: Accuracies ($y_{un}$ and $y_{e1}$) as the shits $\Delta V = \| v_{e_2}-v_{e_1}\|$ increase. All results are obtained from ERM using CMNIST data, where $v_{e_2}$ and $v_{e_1}$ denote the distribution shift settings for spurious correlation colors in domains $e_2$ and $e_1$, respectively. In $10$ sampling trials, we fixed $v_{e_1}$ and varied $v_{e_2}$; thus, the shifts $\Delta V$ increased as sampling progressed. Correspondingly, we obtained $10$ prediction results $y_{un}$ for the unseen domain and $10$ results $y_{e1}$ for domain $e_1$. One observed phenomenon is that as the shifts $\Delta V$ increase, all prediction accuracies tend to converge to values between the oracle and optimal levels.
• Figure 3: Models trained using different learning algorithms on the dataset $D1$.

Theorems & Definitions (15)

• Example 1
• Proposition 1
• Theorem 1
• Theorem 2
• Corollary 1
• Lemma 1: The invariance for Causality b1
• Lemma 2: Fano’s lemma b41
• Lemma 3: b40
• proof : Proof of Proposition 1
• proof : Proof of Theorem 1