Causal Effect Estimation with Learned Instrument Representations

TL;DR

This work tackles unobserved confounding in observational causal inference by learning instrumental representations directly from data. It introduces ZNet, an encoder-based model that decomposes covariates into a confounding component and a learned instrument, enforcing IV moment conditions to produce valid representations for downstream two-stage IV estimators. Through both semi-synthetic and unstructured data experiments, ZNet recovers existing instruments when present and demonstrates strong performance with latent instruments, outperforming TARNet and other IV-generation baselines in estimating ATE and CATE. The approach broadens the applicability of IV methods to high-dimensional and non-tabular data, offering a practical plug-in for causal inference where explicit instruments are unavailable or uncertain.

Abstract

Instrumental variable (IV) methods mitigate bias from unobserved confounding in observational causal inference but rely on the availability of a valid instrument, which can often be difficult or infeasible to identify in practice. In this paper, we propose a representation learning approach that constructs instrumental representations from observed covariates, which enable IV-based estimation even in the absence of an explicit instrument. Our model (ZNet) achieves this through an architecture that mirrors the structural causal model of IVs; it decomposes the ambient feature space into confounding and instrumental components, and is trained by enforcing empirical moment conditions corresponding to the defining properties of valid instruments (i.e., relevance, exclusion restriction, and instrumental unconfoundedness). Importantly, ZNet is compatible with a wide range of downstream two-stage IV estimators of causal effects. Our experiments demonstrate that ZNet can (i) recover ground-truth instruments when they already exist in the ambient feature space and (ii) construct latent instruments in the embedding space when no explicit IVs are available. This suggests that ZNet can be used as a ``plug-and-play'' module for causal inference in general observational settings, regardless of whether the (untestable) assumption of unconfoundedness is satisfied.

Figures (7)

• Figure 1: Causal inference with learned instruments. We learn a feature representation of the observed data $X$ that decomposes them into a learned instrument $\widetilde{Z} = g(X)$ and a residual confounder $\widetilde{X} = f(X)$. This decomposition is induced by enforcing the moment conditions needed for the learned instrument $\widetilde{Z}$ to be valid.
• Figure 2: Illustration of the IV setting. (a) Causal graph for the IV setting; (b) When an IV exists, causal effects can be estimated under unobserved confounding by first fitting a model $\psi$ to predict $T$ from $(X,Z)$, and then training a model $\varphi$ that predicts outcomes based on predicted treatments. (c) Illustration for the latent instrument example discussed in Section \ref{['IC_decompose']}. Here, $P$ is the provider identity.
• Figure 3: Overview of the components of the ZNet Architecture.
• Figure 4: Regression plots for learned and true instruments. Learned instruments are correlated to true ones in the setup with true instrument candidate. Here, $X_{13}, X_{14}, X_{15}$ are true IVs ($x$-axis) plotted against the corresponding dimensions of the learned $\widetilde{Z}$ ($y$-axis).
• Figure 5: Learned instrument representations are correlated to existing instruments in latent categorical instrument dataset. Left most plot are regression predictions of the true categorical IV using learned $\widetilde{Z}$. Middle plot is a t-SNE visualization of $\widetilde{Z}$ colored by the true IV. Right most is a normalized confusion matrix of K-means clusters of $\widetilde{Z}$ compared to the true instruments.