Bayesian Inference of Individualized Treatment Effects using Multi-task Gaussian Processes

Ahmed M. Alaa; Mihaela van der Schaar

Source-linked AI summary

Bayesian Inference of Individualized Treatment Effects using Multi-task Gaussian Processes

Ahmed M. Alaa, Mihaela van der Schaar

arXiv:1704.02801v2Apr 10, 2017cs.LG

TL;DR

The paper addresses individualized treatment-effect estimation from observational data, where counterfactual outcomes are unobserved. It uses a Bayesian multi-task GP framework with risk-based empirical Bayes, outperforming benchmark methods in experiments while leaving credible-set coverage unresolved.

Problem
Individualized treatment effects must be estimated from observational data even though each subject’s counterfactual outcome is unobserved and treatment assignments are feature-dependent.
Method
The method models potential outcomes as a vector-valued function in a vvRKHS and places a multi-task Gaussian-process prior on them, adapting its hyperparameters with risk-based empirical Bayes.
Results
CMGPs outperform all other benchmarks on PEHE in both the IHDP and UNOS datasets.
Takeaways & Limitations
Individualized confidence intervals and patient-specific effect estimates support precision-medicine applications, including more personalized assessment of LVAD benefit.
Takeaways & Limitations
Bayesian credible sets do not guarantee frequentist coverage, and validating coverage guarantees in this setting remains future work.

Abstract

from arXiv · show

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.

1 Introduction

The paper targets individualized treatment-effect estimation from observational EHR data, where factual outcomes are observed but counterfactuals are missing and selection bias complicates inference. It formulates causal inference as Bayesian multi-task learning and addresses uncertainty and selection bias through a multi-task GP and risk-based empirical Bayes.

Multi-phase clinical trials cost $115.3 million in respiratory care, $105.4 million in anesthesia, and $78.6 million in oncology on average.
Observational EHR data offer cheaper, quicker alternatives for using machine learning to infer heterogeneous treatment effects.
Individualized treatment-effect inference observes each subject’s factual outcome but not the corresponding counterfactual outcome, while selection bias shifts treated and control feature distributions.
The paper models potential outcomes as a vector-valued function and learns treatment effects with a Bayesian multi-task Gaussian-process prior over a vvRKHS.
Risk-based empirical Bayes adapts the multi-task GP by combining factual-outcome error with uncertainty in unobserved counterfactual outcomes.
Pointwise credible intervals provide individualized confidence measures, but Bayesian credible sets do not guarantee frequentist coverage.

2 Problem Setup

The problem is to estimate individualized treatment effects from observational data when each subject contributes only one observed potential outcome. Treatment assignment depends on features, creating selection bias, and the analysis assumes unconfoundedness and overlap.

Each subject has features X_i and two random potential outcomes representing responses with and without treatment.
The individualized treatment effect is defined for subjects with features X_i = x using the conditional expected difference between treated and untreated potential outcomes.
The observational dataset contains independent samples of features, treatment assignments, and the outcome under the received treatment.
Because treatment assignments generally depend on features, the propensity score P(W_i = 1|X_i = x) reflects the treatment-assignment policy.
The analysis adopts unconfoundedness and overlap while acknowledging that selection bias can make estimates inaccurate.
Estimating individualized effects is difficult because only one potential outcome is observed for each subject, so standard supervised learning cannot directly estimate T(x).

3 Multi-task Learning for Causal Inference

The paper formulates individualized causal inference as multi-task learning over potential outcomes, using a Bayesian multi-task GP to estimate effects and counterfactual uncertainty. Risk-based empirical Bayes selects hyperparameters by balancing factual fit against posterior counterfactual uncertainty, addressing selection bias.

Multi-task Learning: The potential outcomes function is modeled as a two-output vvRKHS function, with treatment assignment indexing the two learning tasks.The outputs are the response surfaces f_0(.) and f_1(.), corresponding to control and treatment outcomes.
Multi-task Learning: The PEHE quantifies treatment-effect estimation quality, but its empirical value cannot be computed because each subject has only one observed potential outcome.Evaluating factual outcomes alone can fail under strong selection bias because factual performance may not generalize to counterfactual outcomes.
Multi-task Learning: The representer theorem reduces the PEHE-optimal treatment-effect estimator to a finite kernel expansion over the observed subjects.The coefficients are determined by the kernel matrix, treatment assignments, observed outcomes, and regularization.
A Bayesian Perspective: The multi-task GP prior makes the kernel estimator a posterior mean while providing pointwise credible intervals and Bayesian reasoning about unobserved counterfactuals.The Bayesian PEHE risk marginalizes the oracle empirical loss over the posterior distribution of counterfactual outcomes without explicit propensity modeling.
A Bayesian Perspective: Risk-based empirical Bayes selects a kernel that fits factual outcomes while regularizing posterior counterfactual variance, thereby targeting better counterfactual generalization under selection bias.Its feature-space interpretation jointly fits factual response surfaces and minimizes uncertainty in counterfactual outcomes.

4 Causal Multi-task Gaussian Processes (CMGPs)

CMGPs use a multi-task GP prior with a flexible linear coregionalization kernel and risk-based empirical Bayes adaptation. The procedure estimates individualized effects and credible intervals through automated hyperparameter optimization, with cubic matrix inversion as its main computational cost.

Causal Inference via CMGPs: CMGP models the potential outcomes function with the Bayesian prior f ∼ GP(0, K_θ) and returns an ITE function with credible intervals.The model is presented as a recipe for Bayesian causal inference using a multi-task Gaussian process.
Constructing the CMGP Kernel: The linear model of coregionalization permits treatment-specific response surfaces to have different covariance functions, smoothness levels, and relevant features.Automatic relevance determination supplies treatment-specific length scales for feature relevance.
Adapting the Prior via Risk-based Empirical Bayes: Risk-based empirical Bayes uses leave-one-out factual prediction error with Bayesian regularization to adapt the prior and avoid overfitting factual outcomes.Jeffrey’s priors on the regularization parameters allow those parameters to be integrated out into a revised objective.
Causal Inference via CMGPs: Algorithm 1 initializes hyperparameters automatically, optimizes them with a re-parameterized ADAM procedure, and outputs posterior-mean ITE estimates with Bayesian credible intervals.The interval is computed from the posterior variance and the normal distribution’s inverse error function.
Causal Inference via CMGPs: O(n^3) matrix inversion dominates the computational burden, although conventional sparse approximations can reduce this cost for large observational studies.The inversion occurs when computing the posterior covariance-related matrix after hyperparameter optimization.

5 Experiments

The experiments evaluate CMGPs on semi-synthetic IHDP data and a newly introduced UNOS setup, using PEHE across repeated train-validation-test splits. CMGPs outperform the evaluated benchmarks on both datasets, with larger gains under stronger selection bias in UNOS.

Datasets: Experiments use the semi-synthetic IHDP dataset and a newly introduced UNOS observational setup involving cardiac patients awaiting transplantation.IHDP outcomes are simulated from real covariates and treatment assignments, while UNOS models LVAD effects and survival outcomes.
Datasets: The UNOS cohort contains 1,006 patients initially, with 14 covariates; removing 200 patients produces an 806-patient cohort with increased selection bias.The reduced cohort contains 774 control and 232 treated patients before the specified removals.
Baselines: The benchmark suite includes tree-based, balancing, propensity-based, matching, doubly robust, and Gaussian process methods.Comparisons include BART, causal forests, counterfactual regression methods, matching methods, TML, and separate treated-control GP regression.
Evaluation: Performance is evaluated with in-sample and out-of-sample PEHE over 1,000 realizations of each dataset using 60/20/20 train-validation-test splits.Counterfactual outcomes are withheld from every benchmark, including during training.
Results: CMGPs outperform all other benchmarks in PEHE on both IHDP and UNOS datasets.Their gains over ordinary GP regression increase on UNOS, which exhibits stronger selection bias.

6 Discussion: Towards Precision Medicine

The UNOS analysis examines whether individualized LVAD benefit estimates could improve transplant wait-list prioritization. It identifies patients for whom population-average assumptions may produce delayed high-priority assignment.

Clinical utility: CMGP analysis of UNOS patients wait-listed from 2005–2010 uses real survival times to assess clinical utility.The analysis evaluates patients under the existing transplant allocation context.
Clinical utility: 10.3% of wait-list patients with an implanted LVAD exhibited delayed assignment to high priority.The paper presents this as evidence that the current allocation approach can misclassify individual patients.
Patient example: One patient’s high-priority assignment was delayed because the allocation scheme overestimated her LVAD survival benefit using the population average.The algorithm estimated a lower individual benefit for the diabetic patient than the population average.
Implication: The authors envision prioritizing patients according to individualized risks rather than assuming equal LVAD benefit across patients.This proposed direction follows the reported patient-level discrepancy between population-average and individualized estimates.