Annotation-Assisted Learning of Treatment Policies From Multimodal Electronic Health Records

TL;DR

This paper introduces AACE, a novel method that uses expert annotations to improve causal treatment policy learning from multimodal EHR data, leading to better treatment decision support.

Contribution

AACE leverages expert annotations during training to enhance confounding adjustment and predicts treatment benefits solely from multimodal data at inference.

Findings

AACE outperforms existing risk-based and representation-based causal methods.

The method demonstrates strong empirical results on synthetic, semi-synthetic, and real-world datasets.

AACE provides practical insights for applying causal ML in clinical settings.

Abstract

We study how to learn treatment policies from multimodal electronic health records (EHRs) that consist of tabular data and clinical text. These policies can help physicians make better treatment decisions and allocate healthcare resources more efficiently. Causal policy learning methods prioritize patients with the largest expected treatment benefit. Yet, existing estimators are designed for tabular covariates under causal assumptions that may be hard to justify in the multimodal setting. A pragmatic alternative is to apply causal estimators directly to multimodal representations, but this can produce biased treatment effect estimates when the representations do not preserve the relevant confounding information. As a result, predictive models of baseline risk are commonly used in practice to guide treatment decisions, although they are not designed to identify which patients benefit most from treatment. We propose AACE (Annotation-Assisted Coarsened Effects), an annotation-assisted approach to causal policy learning for multimodal EHRs. The method uses expert-provided annotations during training to support confounding adjustment, and then predicts treatment benefit from only multimodal representations at inference. We show that the proposed method achieves strong empirical performance across synthetic, semi-synthetic, and real-world EHR datasets, outperforming risk-based and representation-based causal baselines, and offering practical insights for applying causal machine learning in clinical practice.