PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Edward Choi, Mohammad Taha Bahadori, Joshua A. Kulas, Andy Schuetz, Walter F. Stewart, Jimeng Sun

arXiv:1608.05745v4cs.LGcs.AIcs.NE

TL;DR

Clinical prediction seeks both accuracy and interpretability, but traditional models often lose temporal information and RNNs are difficult to interpret. RETAIN uses two-level reverse-time attention to expose influential visits and variables while retaining RNN-comparable predictive performance. On a large EHR evaluation, it matched RNN variants in accuracy and speed, outperformed traditional models, and supported intuitive case-level interpretation.

  • Problem

    Traditional interpretable models aggregate EHR features and ignore temporal relations, whereas accurate RNN-based models are difficult to interpret in clinical care.

  • Method

    RETAIN is a two-level neural attention model that uses interpretable visit embeddings and reverse-time RNNs to generate visit- and variable-level attention.

  • Results

    RETAIN achieved accuracy and speed comparable to RNN variants, significantly outperformed traditional models, and provided intuitive interpretation in a heart-failure case study.

  • Takeaways & Limitations

    RETAIN demonstrates that sequential EHR models can combine RNN-level predictive performance with detailed interpretation of influential visits and clinical variables.

Abstract

from arXiv · show

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.

1 Introduction

EHR prediction must balance accuracy with interpretability, but traditional models often discard temporal structure while RNNs remain difficult to interpret. RETAIN addresses this tension with reverse-time attention and achieves RNN-comparable performance with clinically useful interpretations on a large EHR dataset.

  • Motivation: Traditional interpretable models aggregate EHR variables and ignore temporal relations, producing sub-optimal accuracy.Latent-variable time-series models capture temporality but often rely on abstract state variables that limit interpretation.
  • Motivation: RNNs model sequential EHR data effectively, but their outputs are notoriously difficult to interpret for clinical care.Existing attempts to interpret RNNs were described as insufficiently developed for clinical application.
  • RETAIN: RETAIN uses two-level neural attention to identify influential visits and quantify visit-specific features contributing to predictions.Its attention-generation mechanism leverages sequence information while learning an interpretable representation.
  • RETAIN: RETAIN examines patient visits in reverse time order, facilitating more stable attention generation and emphasizing meaningful recent information.The model is designed to emulate how physicians explore patient records from the present toward the past.
  • Evaluation: 14 million visits from 263K patients over 8 years formed the large health-system EHR evaluation dataset.The study compared RETAIN with traditional machine-learning methods and RNN variants for future heart-failure diagnosis.
  • Evaluation: RETAIN matched RNN variants in predictive accuracy and speed, outperformed traditional models, and provided intuitive interpretation through a case study and visualization.The reported comparison covers both predictive performance and interpretability.

2 Methodology

RETAIN represents sequential EHR visits with interpretable embeddings and two reverse-time RNNs that generate visit- and variable-level attention. The resulting context vector supports prediction while exposing which visits and embedded variables influence the output.

  • EHR Structure and Notation: EHR data are modeled as time-labeled sequences of multivariate observations, with labels predicted at each step or at sequence end.Each patient has a sequence of visits, and the number of target labels may exceed one.
  • EHR Structure and Notation: Encounter sequence modeling uses binary visit vectors of medical codes, whereas learning-to-diagnose uses continuous clinical measures to predict one or more diseases.Learning-to-diagnose is treated as a special case of encounter sequence modeling with prediction at the end of the visit sequence.
  • Reverse Time Attention Model RETAIN: RETAIN uses a linear embedding of each input visit to preserve visit-level and variable-level influence for interpretation.The embedding maps the input vector x_i into an interpretable representation v_i.
  • Reverse Time Attention Model RETAIN: Separate RNNs generate visit-level α attention and variable-level β attention from visit embeddings in reverse time order.The α weights govern visit influence, while β vectors focus on individual coordinates of visit embeddings.
  • Reverse Time Attention Model RETAIN: Attention-weighted representation vectors are combined into a context vector, which is transformed into the predicted label.The prediction uses a Softmax layer for categorical outputs and cross-entropy loss; mean squared error can be used for real-valued outputs.
  • Reverse Time Attention Model RETAIN: Compared with standard NLP attention, RETAIN reverses the architecture by using an interpretable embedding and RNN-generated attention weights.The design recovers sequential information while mimicking physician behavior; timestamps were omitted from the formulation, though their use gave a small performance improvement.

3 Interpreting RETAIN

RETAIN interprets predictions at both the visit and variable levels by combining attention weights with input embeddings. Its decomposition identifies which visits and clinical variables contribute to each predicted outcome.

  • The largest α values identify visits that contribute most strongly to the prediction.α alone indicates which visit is influential, but not which variable within that visit explains the influence.
  • RETAIN fixes α and β attention values while perturbing original input variables to measure their effects on predicted label probabilities.The variable producing the largest change in a label probability is assigned the highest contribution.
  • The model output is computed from a context vector formed by attention-weighted visit embeddings and passed through a Softmax layer.The context vector c_i summarizes visits up to time i using α and β attention weights.
  • RETAIN interprets individual variables through their contribution coefficient, which combines α, β, and the embedding matrix.For variable x_j,k, the coefficient is α_jW(β_j ⊙ W_emb[:, k]).
  • For binary inputs, the contribution coefficient itself gives the variable contribution; non-binary inputs require multiplying it by x_j,k.This distinction ensures that variable magnitude is included when input values are not binary.

4 Experiments

The experiments evaluate RETAIN for heart-failure prediction against traditional models and RNN variants using longitudinal EHR data. RETAIN matches RNN predictive performance and training scalability while providing detailed sequence-aware interpretations.

  • 4.2 Heart Failure Prediction: RETAIN was compared with logistic regression, MLP, RNN variants, and traditional machine-learning baselines for heart-failure prediction.Evaluation included negative log-likelihood and AUC, with bootstrap estimates based on 10,000 runs.
  • 4.1 Experimental setting: 3,884 cases and 28,903 controls were evaluated using diagnosis, medication, and procedure codes from the preceding 18 months.The cohort was drawn from Sutter Health EHRs for adults aged 50 to 80 years.
  • 4.2 Heart Failure Prediction: Logistic regression and MLP underperformed the four temporal learning algorithms, while RETAIN matched other RNN variants in prediction performance.The reported comparison emphasizes RETAIN’s interpretation benefit alongside competitive predictive accuracy.
  • 4.2 Heart Failure Prediction: RETAIN’s training time was comparable to RNN, although RETAIN typically required 30 epochs to converge versus approximately 10 for RNN.For encounter sequence modeling, attention-model training time increases linearly with input-sequence length.
  • 4.3 Model Interpretation for Heart Failure Prediction: In a test-patient case study, recent heart-failure-related codes received greater attention than earlier skin-related codes.The visualization examined contributions of diagnosis codes across the patient’s visit sequence.
  • 4.3 Model Interpretation for Heart Failure Prediction: Reversing the visit sequence reduced the model’s heart-failure risk score to 9.0%, demonstrating sensitivity to temporal order.The reversed sequence produced lower contributions from past heart-failure-related codes than the original ordering.
  • 4.3 Model Interpretation for Heart Failure Prediction: Adding antiarrhythmics and anticoagulants reduced the predicted heart-failure risk from 0.2474 to 0.2165 in the case study.The medications also reduced positive contributions from heart-valve-disease and cardiac-dysrhythmia codes at the last visit.

5 Conclusion

RETAIN combines predictive power with interpretability by using attention generation to recover sequence information while keeping representation learning simple. The authors propose reverse-time recurrent processing and identify interactive visualization and broader healthcare evaluation as future work.

  • RETAIN is designed to preserve RNN predictive power while providing a higher degree of interpretation.Its attention mechanism improves prediction while the representation-learning component remains simple for interpretation.
  • RETAIN trains two RNNs in reverse time order to generate attention variables efficiently.
  • Future work includes an interactive RETAIN visualization system and evaluation in additional healthcare applications.

A A method to use the timestamps

RETAIN can incorporate visit timestamps when generating attention weights. Timestamps provide temporal information without being embedded directly into each visit representation.

  • Visit timestamps may encode days from the first visit, intervals between visits, or days until an event such as heart-failure diagnosis.
  • The timestamp-modified procedure changes the attention-generation steps while retaining the original visit embedding for the context vector.
  • Temporal information is used to calculate attentions for the visit sequence rather than to embed each visit.The approach embeds visits with an MLP and recovers sequential information through RNN-generated attentions.

B.1 Hyper-parameter Tuning

RETAIN’s hyper-parameters were tuned on the validation set using random search over embedding, hidden-layer, regularization, and dropout settings.

  • The validation set tuned visit embedding size, both RNN hidden-layer sizes, L2 regularization, and dropout rates.
  • Random search explored m, p, and q values from 32 to 256, L2 values from 0.1 to 0.0001, and dropout rates from 0.0 to 0.8.
  • The final heart-failure model used m, p, q = 128, dropoutvi = 0.6, dropoutci = 0.6, and L2 = 0.0001.

B.2 Code Grouper

The dataset’s diagnosis, medication, and procedure codes were standardized with established medical coding systems and grouped to reduce dimensionality.

  • Diagnosis, medication, and procedure codes used ICD-9, GPI, and CPT, respectively.
  • Diagnosis grouping reduced approximately 14,000 ICD-9 codes to 283, while medication grouping reduced approximately 151,000 codes to 96.
  • Table 3 lists the qualifying ICD-9 codes used to define heart failure cases.

B.3 Training Specifics of the Basline Models

The baseline models used specified regularization, dropout, and hidden-layer configurations, while the heart-failure cohort was defined and matched using explicit case-control criteria.

  • Baseline configurations: Logistic regression used L2 regularization of 0.01 on its weights.
  • Baseline configurations: MLP used 0.6 dropout and L2 regularization of 0.0001 for hidden-layer and logistic-regression weights.
  • Baseline configurations: RNN used 0.6 dropout on both hidden-layer outputs, L2 = 0.0001 on the logistic-regression weight, and hidden layers of size 256.
  • Attention baselines: RNN+αM and RNN+αR used dropout rates of 0.4 for hidden-layer outputs and 0.6 for context vectors.
  • Attention baselines: Both attention baselines used L2 = 0.0001 and hidden-layer dimensions of 256 for their respective attention-generating networks.
  • Heart-failure cohort: Controls were matched by sex, age, and location, with up to ten controls per case and index dates aligned to the corresponding case.

C Results on encounter sequence modeling

Encounter-sequence modeling predicted next-visit diagnosis codes from visit histories and evaluated RETAIN against stationary and recurrent baselines. RETAIN remained interpretable while performing near RNN levels, whereas stationary models differed by task.

  • Task and setup: The task predicts diagnosis codes at the next visit from each patient’s sequence of visits.
  • Task and setup: The dataset was split into training, validation, and test sets at a 0.75:0.10:0.15 ratio.
  • Baselines and evaluation: LR and MLP aggregated up to ten past input vectors, with Softmax applied directly for LR and after a hidden layer for MLP.
  • Baselines and evaluation: Performance was evaluated with test-set negative log likelihood and Recall@k, using k = 5 and 10.
  • Results: RNN achieved the best encounter-diagnosis accuracy, while RETAIN was only slightly inferior and provided full interpretation of its predictions.
  • Results: MLP matched RNN+αM on Recall@10, suggesting this task depended more on disease frequency than occurrence order; stationary models performed worse on heart-failure prediction.
  • Results: Using all past input vectors performed slightly worse than using only ten.
Loading 1608.05745v4…