[논문 리뷰] Interpretable Disease Prediction based on Reinforcement Path Reasoning over Knowledge Graphs

Zhoujian Sun, Wei Dong외 2명|arXiv (Cornell University)|2020. 10. 16.

Machine Learning in Healthcare참고 문헌 39인용 수 5

한 줄 요약

이 논문은 지식 그래프(KG)를 기반으로 한 강화 학습 기반의 해석 가능한 질병 예측 방법을 제안한다. 정책 기반 에이전트가 환자 데이터에서 시작하여 임상적으로 의미 있는 경로를 통해 질병을 예측하는 랜덤 워크를 수행한다. 모델은 두 개의 실생활 전자건강기록(EHR) 데이터셋에서 각각 0.743과 0.639의 매크로 AUC 스코어를 달 đạt했으며, 임상적으로 검증된 해석 가능한 진행 경로를 제공한다.

ABSTRACT

Objective: To combine medical knowledge and medical data to interpretably predict the risk of disease. Methods: We formulated the disease prediction task as a random walk along a knowledge graph (KG). Specifically, we build a KG to record relationships between diseases and risk factors according to validated medical knowledge. Then, a mathematical object walks along the KG. It starts walking at a patient entity, which connects the KG based on the patient current diseases or risk factors and stops at a disease entity, which represents the predicted disease. The trajectory generated by the object represents an interpretable disease progression path of the given patient. The dynamics of the object are controlled by a policy-based reinforcement learning (RL) module, which is trained by electronic health records (EHRs). Experiments: We utilized two real-world EHR datasets to evaluate the performance of our model. In the disease prediction task, our model achieves 0.743 and 0.639 in terms of macro area under the curve (AUC) in predicting 53 circulation system diseases in the two datasets, respectively. This performance is comparable to the commonly used machine learning (ML) models in medical research. In qualitative analysis, our clinical collaborator reviewed the disease progression paths generated by our model and advocated their interpretability and reliability. Conclusion: Experimental results validate the proposed model in interpretably evaluating and optimizing disease prediction. Significance: Our work contributes to leveraging the potential of medical knowledge and medical data jointly for interpretable prediction tasks.

연구 동기 및 목표

• 의료 지식과 전자건강기록(EHR)을 통합하여 해석 가능성 향상을 도모하는 질병 예측 모델을 개발하기 위해.
• 블랙박스 의료 AI 모델의 투명성 부족 문제를 해결하기 위해 추적 가능하고 임상적으로 의미 있는 질병 진행 경로를 생성하기 위해.
• 지식 그래프를 활용하여 검증된 의료 지식에서 유래한 알려진 질병-위험 요인 관계를 인코딩하기 위해.
• 환자 프로필에서 질병 결과로 향하는 경로 선택을 안내하기 위해 EHR 데이터를 사용해 강화 학습 정책을 훈련하기 위해.
• 임상 전문가의 검토를 통해 생성된 예측 경로의 임상적 관련성과 해석 가능성 검증하기 위해.

제안 방법

• 검증된 의료 지식을 사용하여 질병과 위험 요인 간의 관계를 인코딩한 지식 그래프(KG)를 구축하기 위해.
• 각 환자를 질병 또는 위험 요인을 통해 연결된 KG 내의 노드로 표현하기 위해.
• 질병 예측을 KG 상의 랜덤 워크로 공식화하며, 환자 엔티티에서 시작하여 질병 엔티티로 끝나는 경로로 설정하기 위해.
• 예측 정확도를 최대화하는 최적의 보행 경로를 학습하기 위해 정책 기반 강화 학습(RL) 에이전트를 사용하기 위해.
• 이력 환자 결과 기반으로 경로 선택을 최적화하기 위해 실생활 EHR 데이터를 사용해 RL 정책을 훈련하기 위해.
• 환자 특징에서 예측된 질병으로 향하는 단계별 추론 시퀀스로서 해석 가능한 예측 경로를 생성하기 위해.

실험 결과

연구 질문

• RQ1 지식 그래프 상에서 강화 학습을 통해 해석 가능하고 임상적으로 의미 있는 질병 예측 경로를 생성할 수 있는가?
• RQ2 기존의 표준 기계학습 기반 모델 대비 특정 질병 위험 예측 성능은 어느 정도인가?
• RQ3 임상 전문가들이 생성된 예측 경로의 신뢰성과 해석 가능성에 대해 어느 정도 긍정적인 평가를 내릴 수 있는가?
• RQ4 의료 지식과 EHR 데이터의 통합이 질병 예측의 성능 향상과 해석 가능성 향상에 기여하는가?
• RQ5 다양한 실생활 EHR 데이터셋 간에 모델의 경로 추론 능력은 얼마나 견고한가?

주요 결과

• 첫 번째 EHR 데이터셋에서 심혈관계 질환 53종을 예측하는 데 매크로 AUC 0.743을 달달하며 뛰어난 예측 성능를 입증했다.
• 두 번째 EHR 데이터셋에서는 매크로 AUC 0.639를 기록하여 다양한 임상 데이터에서 일관된 성능을 보였다.
• 임상 전문가들이 생성된 질병 진행 경로를 검토한 결과, 그 해석 가능성과 임상적 신뢰성이 확인되었다.
• 강화 학습 정책은 지식 그래프를 효과적으로 탐색하여 의미 있는 환자 중심 예측 경로를 생성하는 데 성공했다.
• 의료 지식과 EHR 데이터의 통합은 정확한 동시에 설명 가능한 질병 위험 예측을 가능하게 했다.
• 모델의 경로 추론은 투명하고 추적 가능한 설명을 제공하여 임상 의사결정 지원에 기여한다.