[Paper Review] Bayesian Inference of Individualized Treatment Effects using Multi-task Gaussian Processes

Ahmed M. Alaa, Mihaela van der Schaar|arXiv (Cornell University)|Apr 10, 2017

Gaussian Processes and Bayesian Inference32 references90 citations

TL;DR

Proposes a Bayesian multi-task Gaussian process framework to infer individualized treatment effects from observational data, with uncertainty quantification and risk-based empirical Bayes to mitigate selection bias; demonstrated on IHDP and UNOS datasets showing improved PEHE performance over baselines.

ABSTRACT

Predicated on the increasing abundance of electronic health records, we investi- gate the problem of inferring individualized treatment effects using observational data. Stemming from the potential outcomes model, we propose a novel multi- task learning framework in which factual and counterfactual outcomes are mod- eled as the outputs of a function in a vector-valued reproducing kernel Hilbert space (vvRKHS). We develop a nonparametric Bayesian method for learning the treatment effects using a multi-task Gaussian process (GP) with a linear coregion- alization kernel as a prior over the vvRKHS. The Bayesian approach allows us to compute individualized measures of confidence in our estimates via pointwise credible intervals, which are crucial for realizing the full potential of precision medicine. The impact of selection bias is alleviated via a risk-based empirical Bayes method for adapting the multi-task GP prior, which jointly minimizes the empirical error in factual outcomes and the uncertainty in (unobserved) counter- factual outcomes. We conduct experiments on observational datasets for an inter- ventional social program applied to premature infants, and a left ventricular assist device applied to cardiac patients wait-listed for a heart transplant. In both experi- ments, we show that our method significantly outperforms the state-of-the-art.

Motivation & Objective

• Motivate estimation of individualized treatment effects (ITE) from observational data and address selection bias in EHRs.
• Model potential outcomes as a vector-valued function and cast ITE estimation as a multi-task learning problem in vvRKHS.
• Develop CMGPs with a linear coregionalization kernel to capture shared and task-specific structure between treated and control outcomes.
• Introduce risk-based empirical Bayes to adapt kernel hyperparameters by balancing factual error and counterfactual uncertainty.
• Provide Bayesian credible intervals for ITE and evaluate against multiple baselines on semi-synthetic IHDP and UNOS datasets.

Proposed method

• Formulate Y(i)(w) = f_w(X_i) + ε_i,w with Gaussian noise to connect ITE to a vector-valued function f = [f_0, f_1].
• Model f as a Gaussian process with a vector-valued kernel K_theta and use a linear model of coregionalization to allow f_0 and f_1 to have different covariances (K_theta = A_0 k_0 + A_1 k_1).
• Represent ITE T(x) = e^T f(x) with e = [-1, 1]^T and derive a representer theorem showing T is in the span of kernel evaluations; connect to Bayesian GP posterior mean.
• Introduce Bayesian PEHE risk and risk-based empirical Bayes to select kernel hyperparameters by minimizing empirical factual error plus posterior counterfactual uncertainty.
• Describe algorithm for CMGPs with ADAM-based optimization of hyperparameters, computation of posterior ITE, and credible intervals for T(x).
• Optional: note use of leave-one-out cross-validation with Bayesian regularization to stabilize hyperparameter selection.

Experimental results

Research questions

• RQ1 Can we accurately estimate individualized treatment effects from observational data using a multi-task GP prior over potential outcomes?
• RQ2 How can we quantify and propagate uncertainty in counterfactuals to produce credible ITE intervals without explicit propensity modeling?
• RQ3 Does risk-based empirical Bayes for kernel hyperparameters improve robustness to selection bias compared to standard GP or single-task approaches?
• RQ4 How does the proposed CMGP approach compare to state-of-the-art baselines on semi-synthetic IHDP and real-world UNOS datasets in terms of PEHE?
• RQ5 What practical insights arise for precision medicine when applying CMGPs to heterogeneous treatment effects in clinical wait-list scenarios?

Key findings

• CMGPs outperform a broad range of baselines (GP, BART, CF, CFRF, Balancing regressors, kNN, MS, TML) on both IHDP and UNOS in-sample and out-of-sample PEHE.
• The risk-based empirical Bayes hyperparameter selection reduces selection bias by regularizing counterfactual uncertainty while fitting factual outcomes.
• The linear coregionalization kernel enables the model to capture differing heterogeneity and relevant features across f0 and f1, outperforming methods that ignore W–X interactions.
• Credible intervals for ITE are produced, enabling individualized confidence assessment for therapeutic decisions.
• The method demonstrates practical clinical utility in the UNOS LVAD wait-list context, illustrating potential improvements in allocation decisions by accounting for individual risk.