Learning When the Concept Shifts: Confounding, Invariance, and Dimension Reduction

TL;DR

This paper introduces a method for learning invariant representations under unobserved confounding to improve prediction across distribution shifts, combining causal modeling and dimension reduction.

Contribution

It proposes a data-driven representation learning approach that optimizes a non-convex objective on the Stiefel manifold to find invariant subspaces resilient to shifts.

Findings

Nearly all local optima align with invariant subspaces under regularization.

The method achieves a small gap between target and source risk.

Empirical validation on real-world data demonstrates effectiveness.

Abstract

Practitioners often face the challenge of deploying prediction models in new environments with shifted distributions of covariates and responses. With observational data, such shifts are often driven by unobserved confounding, and can in fact alter the concept of which model is best. This paper studies distribution shifts in the domain adaptation problem with unobserved confounding. We postulate a linear structural causal model to account for endogeneity and unobserved confounding, and we leverage exogenous invariant covariate representations to cure concept shifts and improve target prediction. We propose a data-driven representation learning method that optimizes for a lower-dimensional linear subspace and a prediction model confined to that subspace. This method operates on a non-convex objective -- that interpolates between predictability and stability -- constrained to the Stiefel manifold, using an analog of projected gradient descent. We analyze the optimization landscape and prove that, provided sufficient regularization, nearly all local optima align with an invariant linear subspace resilient to distribution shifts. This method achieves a nearly ideal gap between target and source risk. We validate the method and theory with real-world data sets to illustrate the tradeoffs between predictability and stability.