[논문 리뷰] Predicting Patient Outcomes with Graph Representation Learning

Emma Rocheteau, Catherine Tong외 3명|arXiv (Cornell University)|2021. 01. 11.

Machine Learning in Healthcare인용 수 5

한 줄 요약

이 논문은 시간적 생리적 특징을 위한 장기 단기 기억망(LSTM)과 진단을 통해 환자 유사성을 모델링하는 그래프 신경망(GNN)을 조합한 하이브리드 모델인 LSTM-GNN을 제안한다. 이 모델은 eICU 데이터셋에서 퇠퇴 기간 예측 성능을 크게 향상시킨다. 그래프 기반 접근 방식은 이웃 정보를 통해 희귀 질환 패턴을 포착함으로써 표준 LSTM 기반 모델을 능가하며, 진단 인코딩과의 보완적 이점을 보여준다.

ABSTRACT

Recent work on predicting patient outcomes in the Intensive Care Unit (ICU) has focused heavily on the physiological time series data, largely ignoring sparse data such as diagnoses and medications. When they are included, they are usually concatenated in the late stages of a model, which may struggle to learn from rarer disease patterns. Instead, we propose a strategy to exploit diagnoses as relational information by connecting similar patients in a graph. To this end, we propose LSTM-GNN for patient outcome prediction tasks: a hybrid model combining Long Short-Term Memory networks (LSTMs) for extracting temporal features and Graph Neural Networks (GNNs) for extracting the patient neighbourhood information. We demonstrate that LSTM-GNNs outperform the LSTM-only baseline on length of stay prediction tasks on the eICU database. More generally, our results indicate that exploiting information from neighbouring patient cases using graph neural networks is a promising research direction, yielding tangible returns in supervised learning performance on Electronic Health Records.

연구 동기 및 목표

• 진단 및 약물과 같이 희소하고 고차원적인 전자 의무기록(EHR) 데이터를 효과적으로 활용하는 데에 기존 모델의 한계를 해결하기 위해.
• 환자 간의 관계적 정보를 통합하여 입원 중 사망률(IHM)과 중환자실 체류 기간(LOS) 예측 성능을 향상시키기 위해.
• 기존 인코더가 간과하는 희귀 공존 질환 패턴에 대해 그래프 기반의 환자 유사성 모델링이 성능 향상에 기여하는지 탐색하기 위해.
• 시간적 모델링과 그래프 기반 이웃 정보를 융합한 종단 간, 해석 가능한 모델을 개발하기 위해.

제안 방법

• 환자 간 진단 유사성을 기반으로 하여 노드가 환자이고 간선이 진단 벡터 간 유사성으로 구성된 정적 환자 그래프를 구축한다. 이때 진단 벡터에 대해 k-최근접 이웃 기반 접근 방식을 사용한다.
• 첫 24시간 동안의 시간적 생리적 시계열 데이터를 단방향 LSTM을 사용하여 처리하여 순차적 패턴을 추출한다.
• 환자 그래프에 대해 네 가지 GNN 아키텍처(GCN, GAT, GraphSAGE, MPNN)를 적용하여 이웃 정보를 집계하고 환자 표현을 개선한다.
• LSTM로 인코딩된 시간적 특징과 GNN로 인코딩된 환자 표현을 연결하여 최종 예측을 위해 완전 연결층을 통과시킨다.
• 진단 특징을 별도로 인코딩하는 진단 인코더를 사용하고, 이 구성 요소가 있는지 여부에 따라 성능을 비교하여 기여도를 평가한다.
• GAT의 주목적 가중치를 활용하여 모델 결정 원리를 해석하고, 예측에 가장 영향을 미치는 유사한 환자들을 시각화한다.

실험 결과

연구 질문

• RQ1 공통 진단을 공유하는 환자 간의 유사성을 그래프로 모델링하는 것이 표준 LSTM 기반 모델 대비 결과 예측 성능 향상에 기여하는가?
• RQ2 희귀 공존 질환 패턴을 다룰 때 그래프 기반 접근 방식이 전통적인 진단 인코딩 방식보다 더 우수한가?
• RQ3 시간적 모델링(LSTM)과 그래프 기반 이웃 정보를 융합하면 개별 구성 요소만 사용할 때보다 성능이 향상되는가?
• RQ4 환자 그래프 표현이 퇠퇴 기간 예측에 대해 사망률 예측보다 얼마나 더 향상시키는가?
• RQ5 GNN 내 주목적 메커니즘을 통해 모델 예측에 영향을 미치는 유사 환자들을 해석 가능한 방식으로 식별할 수 있는가?

주요 결과

• LSTM-GNN 모델은 LSTM 기반 모델 대비 퇠퇴 기간 예측에서 성능을 뛰어나게 하며, GAT를 사용할 경우 AUC가 0.854 ± 0.001, SAGE를 사용할 경우 0.851 ± 0.003을 기록한다.
• 그래프 전용 구성 요소(LSTM-GNN*)는 LSTM 기반 모델(LSTM*)보다 퇠퇴 기간 예측에서 성능 향상을 크게 이룩하여, 진단을 통한 환자 유사성 정보가 유의미한 신호임을 입증한다.
• 통합된 LSTM-GNN 모델이 가장 우수한 성능을 보이며, LOS의 평균 절대 오차는 1.86, IHM의 AUC는 0.371을 기록하여 시간적 모델링과 관계 기반 모델링 간 상호보완적 효과를 입증한다.
• 진단 인코더는 입원 중 사망률 예측에 더 큰 기여를 하며, 그래프 구성 요소는 퇠퇴 기간 예측에서 더 큰 성능 향상을 이룬다. 이는 과제에 따라 다른 이점이 있음을 시사한다.
• LSTM-GAT의 주목적 가중치는 임상적으로 관련성이 있는 이웃 환자들—예를 들어 동일한 희귀 질환(예: 심부전)을 가진 환자들—을 정확히 식별한다. 이는 모델의 해석 가능성에 대한 타당성을 뒷받침한다.
• 절단 실험 결과, LSTM 구성 요소를 제거할 경우 성능이 심각하게 저하됨을 확인하여, 시간적 시계열 데이터에 대해 순차적 모델링이 필수적임을 입증한다.