Causal Effect Estimation with Learned Instrument Representations

1. Problem Statement

The Challenge of Unobserved Confounding:
In observational causal inference, estimating the causal effect of a treatment ($T$) on an outcome ($Y$) is often biased due to unobserved confounders ($U$) that influence both $T$ and $Y$. Standard regression methods fail here because $E[Y|T] \neq E[Y|do(T)]$.

The Limitation of Traditional Instrumental Variables (IV):
Instrumental Variable (IV) regression is the standard solution for unobserved confounding. It requires a valid instrument $Z$ that satisfies three conditions:

1. Relevance: $Z$ is correlated with $T$.
2. Exclusion Restriction: $Z$ affects $Y$ only through $T$ (no direct effect).
3. Unconfoundedness: $Z$ is independent of the unobserved confounder $U$.

The Gap:
In practice, finding a valid, explicit instrument $Z$ is often difficult or impossible, especially in high-dimensional, unstructured data domains (e.g., medical images, text, ECGs). Existing methods typically assume an instrument is already observed or require specific structural assumptions (like multiple candidate IVs) that may not hold. This paper addresses the problem of constructing valid instrumental representations directly from observed covariates ($X$) when no explicit instrument exists.

---

2. Methodology: ZNet

The authors propose ZNet, a representation learning framework that decomposes observed features $X$ into two latent components: a confounding component ($\tilde{X}$) and an instrumental component ($\tilde{Z}$).

A. Architecture

ZNet mirrors the Structural Causal Model (SCM) of IVs using an encoder-decoder structure:

• Encoders: Two neural networks, $f(\cdot)$ and $g(\cdot)$, map the input features $X$ to:
  • $\tilde{X} = f(X)$: The learned confounding representation.
  • $\tilde{Z} = g(X)$: The learned instrumental representation.
• Downstream Models:
  • $\pi(\cdot)$: Predicts treatment $T$ given $\tilde{Z}$ (First Stage).
  • $\phi(\cdot)$: Predicts outcome $Y$ given $\tilde{X}$ and $T$ (Second Stage).
  • $\Phi(\cdot)$: A pre-trained network to estimate the residual error $\tilde{\epsilon}_Y = Y - E[Y|X, T]$.

B. Learning Objective (Loss Function)

The core innovation is enforcing the empirical moment conditions of a valid IV through the loss function. The model is trained to minimize supervised prediction losses while satisfying specific covariance constraints:

1. Supervised Losses:
   • Minimize error in predicting $T$ using $\tilde{Z}$ (ensures Relevance).
   • Minimize error in predicting $Y$ using $\tilde{X}$ and $T$.
2. Moment Conditions (Constraints):
   • Exclusion Restriction: Enforce $\text{Cov}(\tilde{Z}, \tilde{\epsilon}_Y) = 0$. The instrument must be uncorrelated with the outcome's error term.
   • Orthogonality: Enforce $\text{Cov}(\tilde{Z}, \tilde{X}) = 0$. The instrument and confounder representations must be independent.
   • Distributional Regularization: Use KL-divergence penalties to ensure $\tilde{Z}$ and $\tilde{X}$ follow standard normal distributions, preventing trivial solutions.

The total loss function ($L_{ZNet}$) combines Mean Squared Error (MSE), Binary Cross Entropy (BCE), correlation penalties, and KL-divergence terms, optimized via Lagrangian relaxation.

C. Training Strategy

Training occurs in three stages to ensure stability:

1. Residual Estimation: Train $\Phi$ to estimate $\tilde{\epsilon}_Y$.
2. Pre-training: Train the full network in a supervised setting (ignoring moment constraints) to ensure initial representations are predictive of $T$ and $Y$.
3. Fine-tuning: End-to-end training with all loss terms active, including the moment constraints, to refine the decomposition.

---

3. Key Contributions

1. Automated Instrument Generation: ZNet is the first method to learn instrumental representations directly from high-dimensional observed data without requiring pre-existing, explicit instruments.
2. Decomposition Framework: It introduces a principled way to decompose feature spaces into "confounding" and "instrumental" subspaces using deep learning, effectively automating the selection of IVs.
3. Compatibility: The learned representations ($\tilde{Z}$) are designed to be plug-and-play inputs for standard two-stage IV estimators (e.g., TSLS, DeepIV, DFIV).
4. Handling Unstructured Data: The method is specifically validated on unstructured data (ECG signals), demonstrating that latent patterns (e.g., provider-specific variations) can be extracted as valid instruments.

---

4. Experimental Results

The authors evaluated ZNet on the IHDP (semi-synthetic) dataset and a real-world ECG dataset.

A. Instrument Recovery (IHDP)

• Recovery of Explicit IVs: In settings where true instruments existed in the data, ZNet successfully learned representations highly correlated with the ground-truth instruments.
• Latent Instrument Discovery: In settings with no explicit IVs, ZNet constructed latent representations that satisfied IV conditions (high relevance to $T$, low correlation with $U$ and $\epsilon_Y$).
• Metrics: ZNet achieved high F-statistics (indicating strong relevance) and near-zero correlations with prohibited variables (confounders and errors) compared to baselines like AutoIV, VIV, and GIV.

B. Causal Effect Estimation

• Performance: ZNet significantly outperformed baselines (TARNet, AutoIV, VIV, GIV) in estimating Average Treatment Effects (ATE) and Conditional Average Treatment Effects (CATE).
• Robustness: ZNet was the only method that consistently outperformed TARNet (which ignores unobserved confounding) across all settings, including those where unobserved confounders influenced the observed data ($U \to X$).
• No-IV Scenarios: In the most challenging setting (no candidate instruments), ZNet reduced ATE error by a significant margin compared to other IV generation methods.

C. Real-World Application (ECG Data)

• Using 35,000+ ECGs, ZNet extracted a representation from the raw signals that acted as a valid instrument.
• Result: ZNet + DeepIV achieved an ATE error of 0.126, vastly outperforming Ordinary Least Squares (OLS) which had an error of 0.665. This proved the ability to leverage latent information in unstructured medical data for causal inference.

---

5. Significance and Implications

• Democratizing Causal Inference: ZNet removes the "bottleneck" of manual instrument selection, making IV regression applicable to a much wider range of observational datasets, particularly those with high-dimensional or unstructured data (images, text, signals).
• Bridging Deep Learning and Causality: It provides a concrete architecture that integrates the structural constraints of causal models (SCM) into deep learning training objectives, moving beyond purely predictive models.
• Practical Utility: The method offers a pragmatic solution for domains like healthcare and economics where valid instruments are theoretically known (e.g., provider variation) but not explicitly recorded in the data.
• Limitations: The authors note that while ZNet satisfies empirical moment conditions, it does not provide theoretical guarantees of identifiability without assumptions about the data generation process (specifically regarding the independence of $X$ and $U$). However, its empirical success suggests it is a powerful tool for bias mitigation in general settings.

In summary, ZNet represents a significant step forward in automated causal inference, enabling researchers to estimate causal effects in complex, real-world scenarios where traditional IV methods would previously have been infeasible.