[Paper Review] Estimating individual treatment effect: generalization bounds and algorithms

Uri Shalit, Fredrik Johansson+1 more|arXiv (Cornell University)|Jun 13, 2016

Advanced Causal Inference Techniques391 citations

TL;DR

The paper derives a bound on individual treatment effect (ITE) estimation error under strong ignorability and introduces CFR, a representation-learning framework that balances treated and control groups to improve ITE estimation, with experiments showing competitive performance against state-of-the-art methods.

ABSTRACT

There is intense interest in applying machine learning to problems of causal inference in fields such as healthcare, economics and education. In particular, individual-level causal inference has important applications such as precision medicine. We give a new theoretical analysis and family of algorithms for predicting individual treatment effect (ITE) from observational data, under the assumption known as strong ignorability. The algorithms learn a "balanced" representation such that the induced treated and control distributions look similar. We give a novel, simple and intuitive generalization-error bound showing that the expected ITE estimation error of a representation is bounded by a sum of the standard generalization-error of that representation and the distance between the treated and control distributions induced by the representation. We use Integral Probability Metrics to measure distances between distributions, deriving explicit bounds for the Wasserstein and Maximum Mean Discrepancy (MMD) distances. Experiments on real and simulated data show the new algorithms match or outperform the state-of-the-art.

Motivation & Objective

• Motivate accurate estimation of individual treatment effects (ITEs) from observational data under strong ignorability.
• Derive a generalization-error bound for ITE estimation that decomposes into factual error and distributional discrepancy between treated and control groups.
• Propose a representation-learning framework that enforces balance between treated and control distributions to improve ITE estimation.
• Develop and evaluate end-to-end neural-network based algorithms for ITE estimation that optimize the bound via IPM-based regularization.
• Demonstrate empirical performance against existing methods on semi-synthetic and real data.

Proposed method

• Define a representation Phi and a hypothesis h over Phi for predicting outcomes under each treatment.
• Derive an IPM-based bound linking ITE error to the factual loss and a distributional distance between p(x|t=0) and p(x|t=1) in the Phi space.
• Use Wasserstein distance or MMD as computable IPMs to quantify distributional discrepancy.
• Propose CFR (Counterfactual Regression): end-to-end neural network that jointly learns Phi and two heads h0, h1 to predict control and treated outcomes, with a balance-regularization term based on IPM.
• Provide a TARNet variant without the distributional balance term.
• Train via stochastic gradient descent with a weighted empirical loss and IPM-based regularization to minimize the upper bound on PEHE (precision in estimation of heterogeneous effect).

Experimental results

Research questions

• RQ1 How large is the generalization error when estimating ITE from observational data under strong ignorability?
• RQ2 Can a learned representation reduce distributional mismatch between treated and control groups to improve ITE estimation?
• RQ3 Do IPM-based regularizations (Wasserstein or MMD) improve ITE estimation compared to standard covariate-adjusted models?
• RQ4 Is the CFR approach advantageous over existing methods (e.g., Causal Forests, TMLE, BLR/BART) on semi-synthetic and real datasets?
• RQ5 How does the proposed method perform in within-sample and out-of-sample ITE estimation tasks?

Key findings

• A bound shows the ITE estimation error is upper-bounded by the sum of the factual loss plus a distribution-distance term between treated and control representations.
• The bound uses Integral Probability Metrics (IPMs) and yields practical regularization via Wasserstein distance or MMD on the learned representation.
• A neural-network framework CFR (Counterfactual Regression) with separate heads for treatment and control improves ITE estimation by maintaining treatment influence in the representation.
• Empirical results on semi-synthetic IHDP and real Jobs data show CFR and its balanced variant outperform several baselines and competitive state-of-the-art methods.
• The TARNet variant without balance regularization is included for comparison to assess the impact of distributional balancing.
• The approach generalizes beyond linear models to deep representations and non-linear hypotheses while leveraging IPM-based distances for regularization.