[논문 리뷰] RETAIN: An Interpretable Predictive Model for Healthcare using Reverse Time Attention Mechanism

Edward Choi, Mohammad Taha Bahadori외 4명|arXiv (Cornell University)|2016. 08. 19.

Machine Learning in Healthcare참고 문헌 29인용 수 217

한 줄 요약

RETAIN은 EHR 데이터에서 역시간 주의와 두 개의 RNN으로 의료 결과를 예측하며, 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)에서 정확하면서도 해석 가능한 예측 모델의 필요성을 동기 부여합니다.
• 시간성 및 특징 중요도를 모델링하여 높은 정확도와 해석 가능성을 결합한 예측 프레임워크를 개발합니다.
• 최근 방문과 중요한 임상 변수에 역시간 순서로 주의를 기울여 의사 행동을 모방합니다.
• 예측에 기여하는 방문과 변수들을 해석할 수 있는 투명한 메커니즘을 제공합니다.

제안 방법

• 각 방문의 임상 코드를 학습 가능한 임베딩을 통해 고정 차원의 벡터로 임베딩합니다.
• 역시간 순서로 방문 수준 주의(alpha)와 변수 수준 주의(beta)를 생성하기 위해 두 개의 RNN을 사용합니다.
• 알파와 베타로 조절된 방문 임베딩의 가중합으로 컨텍스트 벡터를 계산하고, 소프트맥스/로지스틱 층으로 결과를 예측합니다.
• 예측 정확도와 안정성을 최적화하기 위해 교차 엔트로피 손실(또는 실수 값 출력에 적합한 손실)을 사용해 학습합니다.
• Appendix A에 설명된 대로 성능을 다소 향상시키기 위해 타임스탬프를 선택적으로 포함할 수 있습니다.

실험 결과

연구 질문

• RQ1 RETAIN이 EHR 데이터에서 RNN 기반 모델과 비교할 만큼의 예측 성능을 달성할 수 있습니까?
• RQ2 역시간 주의 메커니즘이 정확도를 희생하지 않으면서 전통적 모델 대비 해석 가능성을 향상시키나요?
• RQ3 방문 수준 주의와 변수 수준 주의가 모델 예측에 어떻게 기여하며 임상적으로 어떻게 해석될 수 있나요?
• RQ4 대규모 EHR 데이터셋에서 RETAIN의 확장성은 어느 정도인가요?

주요 결과

모델	테스트 음의 로그 우도	AUC	학습 시간 / 에포크	테스트 시간
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은 심부전 예측 작업에서 RNN 변형과 비교할 만한 예측 성능을 달성합니다.
• RETAIN은 예측에 영향을 주는 중요한 방문 및 변수 수준 기여자를 식별하여 해석 가능한 출력을 제공합니다.
• 실험에서 RETAIN의 음의 로그우도와 AUC는 RNN 기반선과 대등하여 정확도와 효율성을 모두 입증합니다.
• RETAIN의 학습 및 추론 시간은 RNN 기반 모델과 비슷하며 대규모 EHR 데이터에도 확장 가능성이 있습니다.
• 역시간 주의는 의미 있는 해석을 가능하게 합니다. 임상의는 어떤 코드를 어떤 방문이 특정 예측을 이끄는지 추적할 수 있습니다.