Deep representation learning for individualized treatment effect estimation using electronic health records

Abstract

Utilizing clinical observational data to estimate individualized treatment effects (ITE) is a challenging task, as confounding inevitably exists in clinical data. Most of the existing models for ITE estimation tackle this problem by creating unbiased estimators of the treatment effects. Although valuable, learning a balanced representation is sometimes directly opposed to the objective of learning an effective and discriminative model for ITE estimation. We propose a novel hybrid model bridging multi-task deep learning and K-nearest neighbors (KNN) for ITE estimation. In detail, the proposed model firstly adopts multi-task deep learning to extract both outcome-predictive and treatment-specific latent representations from Electronic Health Records (EHR), by jointly performing the outcome prediction and treatment category classification. Thereafter, we estimate counterfactual outcomes by KNN based on the learned hidden representations. We validate the proposed model on a widely used semi-simulated dataset, i.e. IHDP, and a real-world clinical dataset consisting of 736 heart failure (HF) patients. The performance of our model remains robust and reaches 1.7 and 0.23 in terms of Precision in the estimation of heterogeneous effect (PEHE) and average treatment effect (ATE), respectively, on IHDP dataset, and 0.703 and 0.796 in terms of accuracy and F1 score respectively, on HF dataset. The results demonstrate that the proposed model achieves competitive performance over state-of-the-art models. In addition, the results reveal several findings which are consistent with existing medical domain knowledge, and discover certain suggestive hypotheses that could be validated through further investigations in the clinical domain.