[論文レビュー] 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 を得る。
• 逆時系列順で訪問レベルの α と変数レベルの β の注意重みを生成するために、2つのRNNを用いる。
• 重み α_j と β_j を用いて v_j の加重和としてコンテキストベクトル c_i を計算する。
• コンテキストベクトル上でソフトマックス/シグモイドによりラベルを予測し、クロスエントロピーで学習する。

実験結果

リサーチクエスチョン

• RQ1 RETAINは心不全予測などのEHRベースのタスクにおいて、RNNと同等の予測精度を達成できるか？
• RQ2 逆時刻注意機構は、訪問および臨床変数ごとの解釈可能な寄与を提供し、臨床的に意味のあるものになるか？
• RQ3 従来のベースラインやRNNモデルと比較して、トレーニング時間と推論時間の点でRETAINはどの程度スケールするか？
• RQ4 異なるデータ表現やタスク設定（例：ESM対 L2D）で、RETAINは性能を維持するか？

主な発見

• RETAINは心不全予測でRNN系と同等の予測性能を達成し、AUCはほぼ同等、負の対数尤度はより低い。
• RETAINの注意機構は、訪問レベルと訪問内の変数レベルの両方で解釈可能な寄与度を可能にする。
• 心不全タスクでは、RETAINのテストAUCは0.8705、テスト負の対数尤度は0.2562で、RNNにほぼ一致し、LR/MLPベースラインを上回る。
• RETAINの訓練時間とテスト時間はRNNベースラインと同程度のオーダーであり、巨大なEHRデータセットに対するスケーラビリティを示す。
• 専用のアブレーション（RNN+α_R）では、スカラーな訪問レベルの注意は、RETAINの二重でよりリッチな注意より情報量が劣ることを示している。
• 本モデルは、時間とコードにわたる寄与の可視化をサポートし、解釈可能な臨床洞察を示す。