Learning Invariant Representations with Missing Data

TL;DR

This paper develops methods to learn invariant data representations even when some nuisance variables are missing, ensuring better generalization across different test distributions.

Contribution

It introduces MMD estimators for invariance objectives that work with missing nuisance data, enabling robust model training without complete nuisance observations.

Findings

Achieves test performance comparable to full-data methods in simulations.

Effectively enforces invariance despite missing nuisance variables.

Demonstrates applicability on clinical data.

Abstract

Spurious correlations allow flexible models to predict well during training but poorly on related test distributions. Recent work has shown that models that satisfy particular independencies involving correlation-inducing \textit{nuisance} variables have guarantees on their test performance. Enforcing such independencies requires nuisances to be observed during training. However, nuisances, such as demographics or image background labels, are often missing. Enforcing independence on just the observed data does not imply independence on the entire population. Here we derive \acrshort{mmd} estimators used for invariance objectives under missing nuisances. On simulations and clinical data, optimizing through these estimates achieves test performance similar to using estimators that make use of the full data.