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

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

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 데이터를 사용한 의료 분야에서 정확하면서도 해석 가능한 예측 모델의 필요성을 제시한다.
• 이전에 방문의 영향력을 식별하고 해당 방문 내 중요한 변수들을 찾아내는 이중 수준 주의(attention) 기법을 개발한다.
• 최근 정보를 강조하기 위해 과거 방문에 역시간으로 주의를 기울여 의사 행동을 모방한다.
• 해석 가능한 출력과 함께 RNN에 필적하는 예측 성능과 확장성을 입증한다.

제안 방법

• 각 환자의 EHR을 고차원 특징 벡터를 가진 방문의 순서로 표현한다.
• 방문 정보를 선형 임베딩 W_emb x_i를 통해 v_i로 얻는다.
• 역시간 순서로 방문 수준(α) 및 변수 수준(β) 주의 가중치를 생성하기 위해 두 개의 RNN을 사용한다.
• 가중 합으로 컨텍스트 벡터 c_i를 계산한다. 가중치는 α_j와 β_j로 주어진다.
• 컨텍스트 벡터에 대해 소프트맥스/시그모이드를 사용하여 레이블을 예측하고 교차 엔트로피로 학습한다.

실험 결과

연구 질문

• 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의 이중으로 더 풍부한 주의보다 덜 정보적임을 보여준다.
• 모델은 시간과 코드에 따른 기여도 시각화를 지원하여 해석 가능한 임상 인사이트를 보여준다.