RETAIN: An Interpretable Predictive Model for Healthcare using Reverse Time Attention Mechanism (1608.05745v4)

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.

RETAIN: An Interpretable Predictive Model for Healthcare using Reverse Time Attention Mechanism

The paper introduces RETAIN, a predictive model aimed towards healthcare applications, which leverages a Reverse Time Attention mechanism to balance the tradeoff between accuracy and interpretability—a common impediment in the adoption of predictive models in clinical settings. The primary goal of this model is to provide high predictive accuracy while maintaining clarity in its decision-making process, crucial for medical applications.

Key Contributions

1. Reverse Time Attention Mechanism: RETAIN utilizes a two-level neural attention mechanism, analyzing patient Electronic Health Records (EHR) in reverse chronological order. This reverse time ordering is intended to reflect the clinical reality where recent events typically wield greater influence on patient outcomes.
2. Dual-RNN Structure: Two RNNs are employed in RETAIN—one for generating attention weights for visits (RNN_{\alpha}) and another for medical features within those visits (RNN_{\beta}). This bifurcation helps in disentangling the contributions of individual visits and their respective medical features.
3. Scalability and Efficiency: RETAIN is evaluated on a dataset encompassing 14 million visits from 263K patients over eight years. The model demonstrates robustness both in scalability and computational efficiency, comparable to traditional RNN models, while offering enhanced interpretability.

Methodology

Data Representation:

EHR data for each patient is represented as a time-labeled sequence of multivariate observations, encoded for different clinical variables like diagnoses, medications, and procedures. The model sequences visits and focuses on extracting meaningful temporal patterns to make predictions.

Attention Generation:

RETAIN uses linear embeddings for input features, followed by attention weights calculation through RNNs. This separation facilitates improved interpretability, as the visit- and variable-level attentions (\alpha_i and \beta_i) denote the influence of visit timing and specific medical features, respectively.

Prediction:

The context vector, created using weighted embeddings of past visits, is fed into a Softmax layer for final prediction. The goal can range from predicting future diagnoses to specific outcomes such as heart failure.

Experimental Setup

Dataset:

The primary case paper focuses on predicting heart failure among patients aged 50-80 using EHR data aggregated into diagnosis, medication, and procedure codes. The dataset consists of 14 million visits and employs clinical groupers to reduce dimensionality without losing significant clinical information.

Models for Comparison:

The paper evaluates RETAIN against several baselines including Logistic Regression (LR), Multilayer Perceptron (MLP), standard RNN, and RNNs with various attention mechanisms (RNN+$\alpha_M$ and RNN+$\alpha_R$). Each baseline employs different strategies to manage interpretability and sequence modeling.

Results

Performance Metrics:

Two primary metrics are used: Negative Log-Likelihood (NLL) for model loss and Area Under the ROC Curve (AUC) for classification performance. RETAIN shows competitive performance close to RNNs but with added interpretability.

Efficiency:

Training times per epoch and test times indicate that RETAIN is efficient, with computational demands similar to other advanced RNN models, handling large-scale EHR data adeptly.

Interpretability:

RETAIN's ability to break down its predictions into visit-level and variable-level contributions adds a layer of transparency critical for clinical deployment. Visualizations provided in the paper illuminate how RETAIN can contextualize recent medical events and their impacts on patient outcomes.

Discussion and Future Work

RETAIN's innovative approach bridges a significant gap in medical AI—offering a model that meshes the interpretability of simpler models with the accuracy of deep learning techniques.

Implications:

• Practical Significance: RETAIN's interpretability facilitates clinician trust and enables better integration into clinical workflows, potentially improving patient care through more informed decision-making.
• Theoretical Significance: The model pushes forward the boundaries of how attention mechanisms can be effectively applied in time-sequenced predictive tasks in healthcare.

Future Developments:

The paper hints at potential areas for refinement, including enhancing interpretability interfaces and extending RETAIN's application beyond heart failure to other complex, temporal prediction tasks in healthcare.

Conclusion:

While RETAIN doesn't claim to revolutionize the field, its balanced approach provides a pragmatic solution for deploying neural networks in sensitive and high-stakes environments like healthcare, ensuring both technical rigor and user accessibility.