Recovering Causal Features for Instrumental Variable Regression with Contrastive Learning

Instrumental Variable (IV) regression is a standard technique for estimating causal effects in the presence of unobserved confounders. The classic IV setting assumes access to an external variable - called the instrument - which directly influences the treatment variable. In this work, we consider a more challenging yet realistic setting where the treatment is latent, and we can only observe a nonlinear (potentially high-dimensional) transformation of it. Particularly, using insights from the independent component analysis (ICA) literature, we propose a general contrastive learning framework to recover the latent treatment up to an affine transformation when it is linearly related to the instrument. We prove that the recovered representation is compatible with classical IV techniques. Empirically, we demonstrate the effectiveness of our method using control function and two-stage least squares (2SLS) estimators and evaluate the robustness of the learned estimators in distribution shift setting.