Estimating individual treatment effect: generalization bounds and algorithms

TL;DR

The paper tackles estimating individualized treatment effects from observational data under strong ignorability by introducing CFR, a representation-learning framework that minimizes distributional differences between treated and control groups. It provides a theoretical IPM-based bound linking ITE error to standard factual loss plus a balance term, and presents a practical end-to-end algorithm with two outcome heads that regularizes for distributional imbalance. Empirical results on semi-synthetic IHDP and real-world Jobs data show CFR often matches or exceeds state-of-the-art methods, with notable gains under imbalance and in policy-risk scenarios. This work contributes a stability-oriented criterion for causal inference and a scalable non-linear approach for ITE estimation in observational settings.

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.

Figures (5)

• Figure 1: Neural network architecture for ITE estimation. $L$ is a loss function, $\text{IPM}_\mathrm{G}$ is an integral probability metric. Note that only one of $h_0$ and $h_1$ is updated for each sample during training.
• Figure 2: Out-of-sample ITE error versus IPM regularization for CFR Wass, relative to the error at $\alpha=0$, on 500 realizations of IHDP, with high ($q=1$), medium and low (artificial) imbalance between control and treated.
• Figure 3: Policy risk on Jobs as a function of treatment inclusion rate. Lower is better. Subjects are included in treatment in order of their estimated treatment effect given by the various methods. CFR Wass is similar to CFR and is omitted to avoid clutter.
• Figure 4: t-SNE visualizations of the balanced representations of IHDP learned by our algorithms CFR, CFR MMD and CFR Wass. We note that the nearest-neighbor like quality of the Wasserstein distance results in a strip-like representation, whereas the linear MMD results in a ball-like shape in regions where overlap is small.
• Figure 5: Out-of-sample error in estimated ITE, as a function of IPM regularization parameter for CFR Wass, on 500 realizations of IHDP, with high ($q=1$), medium and low (artificial) imbalance between control and treated.

Theorems & Definitions (49)

• Definition 1
• Definition 2
• Definition 3
• Definition 4
• Definition 5
• Definition 6
• Definition 7
• Lemma 1
• Theorem 1
• Definition A1