Adjustment for Confounding using Pre-Trained Representations

TL;DR

This paper explores how pre-trained neural network features can be used to adjust for confounding in causal inference with non-tabular data, addressing challenges of high dimensionality and non-identifiability.

Contribution

It formalizes conditions under which latent features from neural networks enable valid ATE adjustment and inference, highlighting neural networks' ability to adapt to sparsity and intrinsic data structure.

Findings

Neural networks can achieve fast convergence rates for latent feature learning.

Latent features enable valid adjustment for confounding in non-tabular data.

Structural assumptions for linear models are unrealistic for neural network features.

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.