[Paper Review] Learning Overlapping Representations for the Estimation of Individualized Treatment Effects

Yao Zhang, Alexis Bellot+1 more|arXiv (Cornell University)|Jan 14, 2020

Advanced Causal Inference Techniques42 citations

TL;DR

The paper argues that domain-invariant representations are often unsuitable for causal effect estimation from observational data and introduces a deep kernel learning framework (DKLITE) that regularizes for counterfactual variance and preserves information via invertible representations to improve ITE estimation.

ABSTRACT

The choice of making an intervention depends on its potential benefit or harm in comparison to alternatives. Estimating the likely outcome of alternatives from observational data is a challenging problem as all outcomes are never observed, and selection bias precludes the direct comparison of differently intervened groups. Despite their empirical success, we show that algorithms that learn domain-invariant representations of inputs (on which to make predictions) are often inappropriate, and develop generalization bounds that demonstrate the dependence on domain overlap and highlight the need for invertible latent maps. Based on these results, we develop a deep kernel regression algorithm and posterior regularization framework that substantially outperforms the state-of-the-art on a variety of benchmarks data sets.

Motivation & Objective

• Motivate the limitations of domain-invariant representations for counterfactual inference under covariate shift.
• Propose regularization via posterior counterfactual variance and invertible representations to improve ITE generalization.
• Develop a deep kernel learning framework (DKLITE) with posterior regularization for flexible loss terms.
• Provide empirical evidence that counterfactual-variance regularization outperforms state-of-the-art methods on benchmark datasets.

Proposed method

• Formulate ITE estimation under the potential outcomes framework with consistency/ignorability/overlap assumptions.
• Introduce a deep kernel learning model where f_t(x)=w_t^T φ(x) and φ is learned by a neural network; place priors on w_t and model posterior distributions.
• Regularize the posterior to encourage counterfactual overlap via the counterfactual variance term Var_hatρ_t(X_{1−t}) and preserve information with an invertibility constraint (decoder ψ) to reconstruct X from φ(X).
• Derive a final loss: L_fin = L_lik + α1 L_var + α2 L_rec, where L_lik is negative log-likelihood of factual data, L_var penalizes counterfactual variance, and L_rec enforces representation invertibility.
• Optimize within a regularized Bayes framework, yielding predictive distributions f_t(x) ~ N(μ(x|D_t,Θ_t), σ^2(x|D_t,Θ_t)).
• Demonstrate that including counterfactual-variance regularization improves generalization and uncertainty estimation.

Experimental results

Research questions

• RQ1 Can counterfactual variance regularization improve identifiability/estimation of individualized treatment effects from observational data?
• RQ2 Is enforcing full distributional equality between treated and control groups (domain invariance) unnecessarily strict for ITE generalization, and can overlap suffice?
• RQ3 Does an invertible representation (via a decoder) preserve information content and improve ITE predictions?
• RQ4 How does the proposed DKLITE framework perform against state-of-the-art baselines on benchmark causal-inference datasets?

Key findings

• DKLITE outperforms state-of-the-art methods on IHDP, Twins, and Jobs datasets in both in-sample and out-of-sample settings.
• Inclusion of counterfactual-variance regularization and invertibility leads to substantial performance gains, especially in small data regimes.
• Jointly optimizing likelihood, counterfactual variance, and reconstruction loss yields synergistic improvements over optimizing any single term.
• Predicted uncertainty can be leveraged (DKLITE-U) to further boost performance by focusing attention on uncertain cases.