[Paper Review] RETAIN: An Interpretable Predictive Model for Healthcare using Reverse Time Attention Mechanism

Edward Choi, Mohammad Taha Bahadori+4 more|arXiv (Cornell University)|Aug 19, 2016

Artificial Intelligence in Healthcare841 citations

TL;DR

RETAIN provides an interpretable, reverse-time attention-based predictive model for EHR data that matches RNN accuracy while offering clear attributions to visits and variables.

ABSTRACT

Accuracy and interpretability are two dominant features of successful predictive models. Typically, a choice must be made in favor of complex black box models such as recurrent neural networks (RNN) for accuracy versus less accurate but more interpretable traditional models such as logistic regression. This tradeoff poses challenges in medicine where both accuracy and interpretability are important. We addressed this challenge by developing the REverse Time AttentIoN model (RETAIN) for application to Electronic Health Records (EHR) data. RETAIN achieves high accuracy while remaining clinically interpretable and is based on a two-level neural attention model that detects influential past visits and significant clinical variables within those visits (e.g. key diagnoses). RETAIN mimics physician practice by attending the EHR data in a reverse time order so that recent clinical visits are likely to receive higher attention. RETAIN was tested on a large health system EHR dataset with 14 million visits completed by 263K patients over an 8 year period and demonstrated predictive accuracy and computational scalability comparable to state-of-the-art methods such as RNN, and ease of interpretability comparable to traditional models.

Motivation & Objective

• Motivate the need for accurate yet interpretable predictive models in healthcare using EHR data.
• Develop a two-level attention mechanism that identifies influential past visits and key variables within those visits.
• Mimic physician behavior by attending to past visits in reverse time to emphasize recent information.
• Demonstrate predictive performance and scalability comparable to RNNs while providing interpretable outputs.

Proposed method

• Represent each patient’s EHR as a sequence of visits with high-dimensional feature vectors.
• Embed visit information via a linear embedding W_emb x_i to obtain v_i.
• Use two RNNs to generate visit-level (α) and variable-level (β) attention weights in reverse time order.
• Compute a context vector c_i as a weighted sum of v_j with weights α_j and β_j.
• Predict labels with a softmax/sigmoid over the context vector and train with cross-entropy.

Experimental results

Research questions

• RQ1 Can RETAIN achieve predictive accuracy comparable to RNNs on EHR-based tasks such as heart failure prediction?
• RQ2 Does the reverse-time attention mechanism provide interpretable attributions over visits and individual clinical variables that are clinically meaningful?
• RQ3 How does RETAIN scale in terms of training and inference time compared with traditional baselines and RNN models?
• RQ4 Does RETAIN maintain performance across different data representations and task settings (e.g., ESM vs L2D)?

Key findings

• RETAIN achieves predictive performance comparable to RNN variants for heart failure prediction, with similar AUC and lower-negative log-likelihood.
• RETAIN’s attention mechanism enables interpretable attribution at both the visit level and the variable level within visits.
• On the heart failure task, RETAIN’s test AUC is 0.8705 and test negative log-likelihood is 0.2562, closely matching RNN and outperforming LR/MLP baselines.
• Training and test times for RETAIN are in the same order of magnitude as RNN baselines, demonstrating scalability for large EHR datasets.
• A dedicated ablation (RNN+α_R) shows that scalar visit-level attention is less informative than the dual, richer attention in RETAIN.
• The model supports visualization of contributions over time and codes, illustrating interpretable clinical insights.