Adjustment for Confounding using Pre-Trained Representations
Rickmer Schulte, David Rügamer, Thomas Nagler
TL;DR
The paper extends average treatment effect (ATE) estimation to settings where non-tabular data, such as images and text, may confound treatment and outcome. It introduces pre-trained representations Z=φ(W) as adjustment tools within the doubly robust double machine learning (DML) framework, and develops conditions under which these latent features yield valid inference despite non-identifiability under invertible linear transforms (ILTs). Central contributions include formalizing P-valid and P-ODS notions for sufficiency, establishing that convergence rates for nuisance estimation depend on intrinsic dimension rather than sparsity/additivity, and proving that neural networks can adapt to low-dimensional manifolds to achieve fast rates. The work combines theory with extensive experiments on text (IMDb) and imaging (X-ray) to show that pre-trained representations enable unbiased ATE estimates and valid confidence intervals, while highlighting the importance of pre-training and intrinsic dimension. Overall, it provides both theoretical foundations and practical guidance for integrating pre-trained representations into causal inference with non-tabular confounding, with implications for healthcare and other domains where rich data modalities are prevalent.
Abstract
There is growing interest in extending average treatment effect (ATE) estimation to incorporate non-tabular data, such as images and text, which may act as sources of confounding. Neglecting these effects risks biased results and flawed scientific conclusions. However, incorporating non-tabular data necessitates sophisticated feature extractors, often in combination with ideas of transfer learning. In this work, we investigate how latent features from pre-trained neural networks can be leveraged to adjust for sources of confounding. We formalize conditions under which these latent features enable valid adjustment and statistical inference in ATE estimation, demonstrating results along the example of double machine learning. We discuss critical challenges inherent to latent feature learning and downstream parameter estimation arising from the high dimensionality and non-identifiability of representations. Common structural assumptions for obtaining fast convergence rates with additive or sparse linear models are shown to be unrealistic for latent features. We argue, however, that neural networks are largely insensitive to these issues. In particular, we show that neural networks can achieve fast convergence rates by adapting to intrinsic notions of sparsity and dimension of the learning problem.