[Paper Review] Learning Representations for Counterfactual Inference

Fredrik Johansson, Uri Shalit+1 more|arXiv (Cornell University)|May 12, 2016

Machine Learning and Algorithms30 references191 citations

TL;DR

The paper links counterfactual inference with domain adaptation and representation learning, proposing balanced representations to improve counterfactual predictions in observational studies, and shows deep learning variants outperform prior methods.

ABSTRACT

Observational studies are rising in importance due to the widespread accumulation of data in fields such as healthcare, education, employment and ecology. We consider the task of answering counterfactual questions such as, "Would this patient have lower blood sugar had she received a different medication?". We propose a new algorithmic framework for counterfactual inference which brings together ideas from domain adaptation and representation learning. In addition to a theoretical justification, we perform an empirical comparison with previous approaches to causal inference from observational data. Our deep learning algorithm significantly outperforms the previous state-of-the-art.

Motivation & Objective

• Motivate counterfactual inference in observational studies and its distinction from standard supervised learning.
• Formulate counterfactual prediction as a covariate shift/domain adaptation problem.
• Introduce representation-learning methods that balance treated and control populations to reduce generalization error.
• Present linear and deep learning approaches to learn balanced representations for counterfactual tasks.
• Provide theoretical justification and empirical evaluation against existing causal inference methods.

Proposed method

• Learn a representation Phi: X -> R^d and a predictor h: R^d x T -> R that jointly minimize prediction error, counterfactual regularization, and distributional imbalance.
• Use discrepancy distance between the factual and counterfactual distributions to encourage balance between treated and control groups.
• In the linear setting, derive closed-form discrepancy and propose balancing variable selection via a sparse weighting W so Phi(x)=Wx.
• Extend to deep neural networks by embedding the discrepancy term in the network architecture and training objective.
• Provide a two-stage optimization: first optimize Phi and h with the imbalance penalties, then fit a final ridge regression on factual data.
• Discuss linear discrepancy as a matching of means in feature space and relate to covariate balance.

Experimental results

Research questions

• RQ1 How can counterfactual inference be framed as a domain adaptation problem under covariate shift?
• RQ2 Do balanced representations that minimize discrepancy between treated and control groups improve counterfactual predictions?
• RQ3 How can linear and deep learning representations be learned to balance populations while preserving predictive accuracy?
• RQ4 What theoretical guarantees connect representation balancing to counterfactual generalization error?

Key findings

• Balancing representations reduces counterfactual generalization error by decreasing distributional discrepancy between treated and control groups.
• Learning representations that trade off predictive accuracy against imbalance leads to better counterfactual predictions than methods that rely solely on re-weighting samples.
• Both linear (Balancing Linear Regression) and neural network (Balancing Neural Network) variants outperform prior causal inference methods in experiments.
• Theoretical results bound relative counterfactual generalization error using a discrepancy term and a nearest-neighbor regression-style bound.
• Empirical evaluation on IHDP and a News dataset demonstrates improved ITE and ATE estimation without overfitting to the factual distribution.