CGLearn: Consistent Gradient-Based Learning for Out-of-Distribution Generalization

TL;DR

This work introduces a simple yet powerful approach, CGLearn, which relies on the agreement of gradients across various environments, serving as a powerful indication of reliable features, while disagreement suggests less reliability due to potential differences in underlying causal mechanisms.

Abstract

Improving generalization and achieving highly predictive, robust machine learning models necessitates learning the underlying causal structure of the variables of interest. A prominent and effective method for this is learning invariant predictors across multiple environments. In this work, we introduce a simple yet powerful approach, CGLearn, which relies on the agreement of gradients across various environments. This agreement serves as a powerful indication of reliable features, while disagreement suggests less reliability due to potential differences in underlying causal mechanisms. Our proposed method demonstrates superior performance compared to state-of-the-art methods in both linear and nonlinear settings across various regression and classification tasks. CGLearn shows robust applicability even in the absence of separate environments by exploiting invariance across different subsamples of observational data. Comprehensive experiments on both synthetic and real-world datasets highlight its effectiveness in diverse scenarios. Our findings underscore the importance of leveraging gradient agreement for learning causal invariance, providing a significant step forward in the field of robust machine learning. The source code of the linear and nonlinear implementation of CGLearn is open-source and available at: https://github.com/hasanjawad001/CGLearn.

Figures (3)

• Figure 1: Illustration of three environments generated by intervening on the variable $e$, which takes distinct values $e = 0.2$, $e = 2$, and $e = 5$ in environments $e_1$, $e_2$, and $e_3$, respectively. In each environment, $X_1$ acts as a causal factor for the target variable $Y$, while $X_2$ is a spurious (non-causal) factor with respect to $Y$. This figure exemplifies how different interventions on $e$ create distinct environments.
• Figure 2: Nonlinear MLP implementation of CGLearn. $X_1$ (causal) and $X_2$ (spurious) feed into the first hidden layer $h_1$. Weight updates in $h_1$ are performed based on gradient consistency (using $L^2$-norm) for each feature across all training environments. The rest of the weights such as weights in $h_2$, are updated similarly to ERM (without imposing any consistency constraints).
• Figure 3: Performance comparison of CGLearn, IRM, ICP, and ERM across various linear multiple environment setups. Each subplot represents different configurations of the data, showing the mean squared error (MSE) for causal and noncausal variables over 50 trials.