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

TL;DR

This paper introduces a Bayesian multi-task Gaussian process framework for estimating individualized treatment effects from observational health data, providing confidence measures and addressing selection bias.

Contribution

It proposes a novel nonparametric Bayesian method using multi-task GPs with a risk-based empirical Bayes approach for treatment effect inference.

Findings

Outperforms state-of-the-art methods on health datasets

Provides credible intervals for treatment effect estimates

Effectively mitigates selection bias in observational data

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.