[논문 리뷰] MiME: Multilevel Medical Embedding of Electronic Health Records for Predictive Healthcare

Edward Choi, Cao Xiao외 2명|arXiv (Cornell University)|2018. 10. 22.

Machine Learning in Healthcare참고 문헌 38인용 수 99

한 줄 요약

MiME는 진단-치료 관계를 보조 작업으로 모델링함으로써 EHR 데이터에서 다층 임베딩을 학습하며, 특히 작은 데이터셋에서 강력한 예측 성능을 달성합니다.

ABSTRACT

Deep learning models exhibit state-of-the-art performance for many predictive healthcare tasks using electronic health records (EHR) data, but these models typically require training data volume that exceeds the capacity of most healthcare systems. External resources such as medical ontologies are used to bridge the data volume constraint, but this approach is often not directly applicable or useful because of inconsistencies with terminology. To solve the data insufficiency challenge, we leverage the inherent multilevel structure of EHR data and, in particular, the encoded relationships among medical codes. We propose Multilevel Medical Embedding (MiME) which learns the multilevel embedding of EHR data while jointly performing auxiliary prediction tasks that rely on this inherent EHR structure without the need for external labels. We conducted two prediction tasks, heart failure prediction and sequential disease prediction, where MiME outperformed baseline methods in diverse evaluation settings. In particular, MiME consistently outperformed all baselines when predicting heart failure on datasets of different volumes, especially demonstrating the greatest performance improvement (15% relative gain in PR-AUC over the best baseline) on the smallest dataset, demonstrating its ability to effectively model the multilevel structure of EHR data.

연구 동기 및 목표

• 코드의 고유한 다층 구조를 활용하여 EHR 데이터에 대한 딥 러닝의 데이터 양 부족 문제를 해결합니다.
• 방문 내에서 진단–치료 상호 작용을 포착하는 다층 임베딩을 학습합니다.
• 외부 라벨 없이 보조 예측 작업을 활용하여 심부전 예측과 같은 작업에서 예측 성능을 향상시킵니다.

제안 방법

• 방문을 Dx 코드가 포함된 Dx 객체들의 집합으로 표현하고 각 Dx 객체에는 관련 치료가 연결됩니다.
• Dx 코드 임베딩과 치료와의 상호작용을 통해 g(d_i, m_i)로 진단 객체 임베딩 o_i를 계산합니다.
• 건너뛰기 연결된 상향식 형태를 사용하여 방문 임베딩 v를 집계하고 최종 예측을 위한 환자 표현 h를 도출합니다.
• o_i로부터 Dx 및 치료 코드를 예측하는 보조 예측 작업을 사용하여 추가 라벨링 없이 임베딩을 유도합니다.
• g(d_i, m_i)와 W_m를 통해 Dx–Rx 상호작용 모형을 이차적 방식으로 포착합니다.
• 목표 작업(예: 심부전 예측)과 Dx/m 예측 코드에 대한 보조 손실(L_aux)을 함께 최종 학습합니다.

실험 결과

연구 질문

• RQ1 데이터가 제한적일 때 EHR 데이터의 고유한 다층 구조를 활용해 예측 성능을 향상시킬 수 있을까요?
• RQ2 Dx–치료 관계를 기반으로 한 보조 작업이 방문 및 환자 임베딩의 질과 일반화 능력을 향상시킬까요?
• RQ3 심부전 및 순차 질환 예측에서 코드를 평면화하거나 온톨로지 기반 지식을 주입하는 기준선과 MiME은 어떻게 비교될까요?

주요 결과

• MiME는 다양한 크기의 데이터셋에서 심부전 예측 작업에서 기준선보다 우수하며, 보조 작업을 사용하는 MiME의 경우 가장 작은 데이터셋에서 PR-AUC가 상대적으로 15% 향상되었습니다.
• MiME는 데이터 양과 방문 복잡성에 관계없이 심부전 예측 및 순차적 질병 예측에서 일관되게 기준선을 상회합니다.
• 보조 작업은 일반화 성능을 향상시키며, 특히 더 작거나 더 복잡한 데이터셋에서 MiME aux가 여러 설정에서 MiME 단독보다 더 높은 PR-AUC를 달성합니다.
• MiME은 데이터 축소에 대한 강인성을 보이며 오직 코드 계층 구조나 평면화된 표현을 사용하는 모델보다 Dx–치료 상호작용을 더 효과적으로 포착합니다.