[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

Machine Learning in Healthcare29 references217 citations

TL;DR

RETAIN introduces reverse time attention with two RNNs to predict healthcare outcomes from EHR data, achieving accuracy comparable to RNNs while providing interpretable visit- and variable-level insights.

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 electronic health records (EHR).
• Develop a predictive framework that combines high accuracy with interpretability by modeling temporality and feature importance.
• Mimic physician behavior by attending to recent visits and important clinical variables in reverse time order.
• Provide a transparent mechanism to interpret which visits and variables contribute to predictions.

Proposed method

• Embed each visit’s clinical codes into a fixed-dimensional vector via a learnable embedding.
• Use two RNNs to generate visit-level attention (alpha) and variable-level attention (beta) in reverse time order.
• Compute a context vector as a weighted sum of visit embeddings modulated by alpha and beta, then predict outcomes with a softmax/ logistic layer.
• Train with cross-entropy (or appropriate loss for real-valued outputs) to optimize prediction accuracy and stability.
• Optionally incorporate timestamps to slightly improve performance as described in Appendix A.

Experimental results

Research questions

• RQ1 Can RETAIN achieve predictive performance comparable to RNN-based models on EHR data?
• RQ2 Does the reverse-time attention mechanism provide improved interpretability over traditional models without sacrificing accuracy?
• RQ3 How do visit-level and variable-level attentions contribute to model predictions and how can they be interpreted clinically?
• RQ4 What is the scalability of RETAIN in large-scale EHR datasets?

Key findings

Model	Test Neg Log Likelihood	AUC	Train Time / epoch	Test Time
LR	0.3269±0.0105	0.7900±0.0111	0.15s	0.11s
MLP	0.2959±0.0083	0.8256±0.0096	0.25s	0.11s
RNN	0.2577±0.0082	0.8706±0.0080	10.3s	0.57s
RNN+ αM	0.2691±0.0082	0.8624±0.0079	6.7s	0.48s
RNN+ αR	0.2605±0.0088	0.8717±0.0080	10.4s	0.62s
RETAIN	0.2562±0.0083	0.8705±0.0081	10.8s	0.63s

• RETAIN achieves predictive performance comparable to RNN variants on the heart failure prediction task.
• RETAIN provides interpretable outputs by identifying influential visits and variable-level contributors to predictions.
• In experiments, RETAIN’s negative log-likelihood and AUC are competitive with RNN baselines, demonstrating both accuracy and efficiency.
• Training and inference times for RETAIN are comparable to RNN-based models, with scalability suitable for large EHR datasets.
• Reverse-time attention enables meaningful interpretation; clinicians can trace which codes and visits drive a given prediction.