[논문 리뷰] BERTSurv: BERT-Based Survival Models for Predicting Outcomes of Trauma Patients

Yun Zhao, Qinghang Hong외 4명|arXiv (Cornell University)|2021. 03. 19.

Machine Learning in Healthcare참고 문헌 28인용 수 11

한 줄 요약

BERTSurv는 구조화되지 않은 임상 노트와 구조화된 임상 측정치를 통합하여 외상 환자의 사망 및 생존 예측을 위한 새로운 딥러닝 프레임워크이다. 비구조화된 임상 노트와 구조화된 데이터에 대한 어텐션 메커니즘을 활용하여 사망 예측의 AUC-ROC가 0.86에 도달하고 생존 분석의 C-index가 0.7을 기록하며 기준 모델을 능가하며 어텐션 시각화를 통해 해석 가능성도 향상시켰다.

ABSTRACT

Survival analysis is a technique to predict the times of specific outcomes, and is widely used in predicting the outcomes for intensive care unit (ICU) trauma patients. Recently, deep learning models have drawn increasing attention in healthcare. However, there is a lack of deep learning methods that can model the relationship between measurements, clinical notes and mortality outcomes. In this paper we introduce BERTSurv, a deep learning survival framework which applies Bidirectional Encoder Representations from Transformers (BERT) as a language representation model on unstructured clinical notes, for mortality prediction and survival analysis. We also incorporate clinical measurements in BERTSurv. With binary cross-entropy (BCE) loss, BERTSurv can predict mortality as a binary outcome (mortality prediction). With partial log-likelihood (PLL) loss, BERTSurv predicts the probability of mortality as a time-to-event outcome (survival analysis). We apply BERTSurv on Medical Information Mart for Intensive Care III (MIMIC III) trauma patient data. For mortality prediction, BERTSurv obtained an area under the curve of receiver operating characteristic curve (AUC-ROC) of 0.86, which is an improvement of 3.6% over baseline of multilayer perceptron (MLP) without notes. For survival analysis, BERTSurv achieved a concordance index (C-index) of 0.7. In addition, visualizations of BERT's attention heads help to extract patterns in clinical notes and improve model interpretability by showing how the model assigns weights to different inputs.

연구 동기 및 목표

• 구조화되지 않은 임상 노트와 구조화된 측정치를 통합하여 외상 환자 예후 예측을 위한 딥러닝 프레임워크를 개발한다.
• 기존 생존 모델이 텍스트 기반 임상 노트를 통합하지 못하는 격차를 해결한다.
• BERT 기반 표현 학습을 활용하여 사망 예측 및 생존 분석 정확도를 향상시킨다.
• 임상 노트의 어텐션 패턴을 시각화하여 모델 결정의 설명 가능성을 향상시킨다.
• 실제 MIMIC-III 외상 환자 데이터를 기반으로 프레임워크를 평가한다.

제안 방법

• BERTSurv는 미세조정된 BERT를 사용하여 구조화되지 않은 임상 노트를 문맥 기반 임bedding으로 인코딩한다.
• ICU 입원 후 첫 4시간 동안의 임상 측정치를 평균화하고 BERT 임베딩과 연결한다.
• 사망 예측(이元 결과)에 대해 이진 교차 엔트로피(BCE) 손실을 사용한다.
• 생존 분석(시간-이벤트 결과)에 대해 부분 로그우도(PLL) 손실을 사용한다.
• BERT의 자체 어텐션 메커니즘을 시각화하여 모델 결정을 해석하고 임상적으로 의미 있는 패턴을 추출한다.
• 21개의 임상 측정치와 4개의 인구통계학적 특성을 포함한 MIMIC-III 외상 환자 데이터를 기반으로 모델를 훈련 및 평가한다.

실험 결과

연구 질문

• RQ1 BERT 기반 모델은 구조화되지 않은 임상 노트와 구조화된 측정치를 효과적으로 통합하여 외상 환자의 생존 예측 성능을 향상시킬 수 있는가?
• RQ2 BERT를 통해 임상 노트를 통합함으로써 구조화된 데이터만 사용하는 모델에 비해 성능 향상이 뚜렷한가?
• RQ3 BERT의 어텐션 시각화를 통해 사망 예측과 관련된 임상적으로 의미 있는 패턴을 노트에서 식별할 수 있는가?
• RQ4 전통적인 생존 모델인 코크스 회귀 모델에 비해 BERTSurv의 예측 성능는 어떠한가?
• RQ5 임상 텍스트의 어텐션 헤드 분석을 통해 모델의 설명 가능성은 어느 정도 향상되는가?

주요 결과

• BERTSurv는 사망 예측에서 AUC-ROC가 0.86을 기록하여 임상 노트를 포함하지 않은 다층 퍼셉트론 기준 모델 대비 3.6% 향상된 성능을 보였다.
• 생존 분석에서 BERTSurv는 C-index가 0.7을 기록하여 C-index 0.68를 기록한 코크스 모델을 능가했다.
• 어텐션 시각화 결과, 사망이 빠르게 발생한 환자에서 "심한 흉부 외상"과 같은 임상적으로 관련성이 높은 표현에 높은 어텐션 가중치가 부여된 것으로 나타났으며, 유리한 결과를 보인 환자에서는 "좋은 자세"와 같은 표현에 집중된 것으로 확인되었다.
• 모델는 사망이 빠르게 발생한 환자의 노트에서 "심한 흉부"와 "외상"과 같은 패턴을 식별하여 심각한 질환을 나타내는 것으로 확인했다.
• 반대로, 후기 퇴원 환자는 "좋음" 및 "좋은 자세"와 같은 긍정적 표현에 어텐션 집중을 보여, 임상적으로 의미 있는 패턴을 반영했다.
• BERT와 임상 측정치의 통합은 외상 예후 모델링에서 예측 성능 및 설명 가능성을 크게 향상시켰다.