[论文解读] 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 Healthcare被引用 841

一句话总结

RETAIN 提供一个可解释、基于反向时间注意力的预测模型用于 EHR 数据，其准确性可与 RNN 相当，同时对就诊和变量具有清晰的归因。

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.

研究动机与目标

• 动机：在医疗保健中利用 EHR 数据实现既准确又可解释的预测模型的需求。
• 开发一种二级注意力机制，识别影响力大的过去就诊和这些就诊中的关键变量。
• 通过以反向时间关注过去的就诊以强调最近的信息，模拟医生的行为。
• 展示预测性能和可扩展性，达到可与 RNN 相当的水平，同时提供可解释的输出。

提出的方法

• 将每个患者的 EHR 表示为一系列就诊，具有高维特征向量。
• 通过线性嵌入 W_emb x_i 将就诊信息嵌入，得到 v_i。
• 使用两个 RNN 在逆序生成就诊级注意力权重 α 和变量级注意力权重 β。
• 将上下文向量 c_i 计算为 v_j 的加权和，权重为 α_j 和 β_j。
• 使用对上下文向量的 softmax/sigmoid 进行标签预测，并用交叉熵训练。

实验结果

研究问题

• RQ1 RETAIN 能否在基于 EHR 的任务（如心力衰竭预测）上达到与 RNN 相当的预测准确性？
• RQ2 反向时间注意力机制是否提供对就诊和单个临床变量的可解释归因，且具有临床意义？
• RQ3 相对于传统基线和 RNN 模型，RETAIN 在训练和推理时间方面的扩展性如何？
• RQ4 RETAIN 是否在不同数据表示和任务设置下维持性能（如 ESM 与 L2D）？

主要发现

• RETAIN 在心力衰竭预测中达到与 RNN 变体相当的预测性能，AUC 相近且负对数似然更低。
• RETAIN 的注意力机制使就诊级和就诊内变量级的归因具有可解释性。
• 在心力衰竭任务中，RETAIN 的测试 AUC 为 0.8705，测试负对数似然为 0.2562，接近 RNN 且优于 LR/MLP 基线。
• RETAIN 的训练和测试时间数量级与 RNN 基线相同，展示了对 large EHR 数据集的可扩展性。
• 专门的消融研究（RNN+α_R）表明标量就诊级注意力不如 RETAIN 中的双重、更新颖的注意力更具信息量。
• 模型支持对随时间及代码的贡献可视化，展现可解释的临床洞察。