[논문 리뷰] Contrastive Learning of Temporal Distinctiveness for Survival Analysis in Electronic Health Records

Mohsen Nayebi Kerdabadi, Arya Hadizadeh Moghaddam외 3명|arXiv (Cornell University)|2023. 08. 24.

Machine Learning in HealthcareComputer Science인용 수 3

한 줄 요약

이 논문은 전자 건강 기록에서 생존 분석을 위한 새로운 대비 학습 프레임워크인 OTCSurv를 제안한다. 이 프레임워크는 생존 지속 시간의 차이로 측정되는 시간적 독창성에 기반하여 디지털 음성 쌍을 생성함으로써 어려움을 고려하는 음성 쌍을 구성한다. 온톨로지 인식 시퀀셜 인코더와 감독된 가중 대비 손실을 통합함으로써 OTCSurv는 위험 순서 정렬과 사건 발생 시간 예측 모두에서 향상된 성능을 달성하며, 주목사의 분석을 통해 모델의 해석 가능성도 향상시켰다. 이는 AKI 예측에서 최신 기술 수준의 성능을 달성하였다.

ABSTRACT

Survival analysis plays a crucial role in many healthcare decisions, where the risk prediction for the events of interest can support an informative outlook for a patient's medical journey. Given the existence of data censoring, an effective way of survival analysis is to enforce the pairwise temporal concordance between censored and observed data, aiming to utilize the time interval before censoring as partially observed time-to-event labels for supervised learning. Although existing studies mostly employed ranking methods to pursue an ordering objective, contrastive methods which learn a discriminative embedding by having data contrast against each other, have not been explored thoroughly for survival analysis. Therefore, in this paper, we propose a novel Ontology-aware Temporality-based Contrastive Survival (OTCSurv) analysis framework that utilizes survival durations from both censored and observed data to define temporal distinctiveness and construct negative sample pairs with adjustable hardness for contrastive learning. Specifically, we first use an ontological encoder and a sequential self-attention encoder to represent the longitudinal EHR data with rich contexts. Second, we design a temporal contrastive loss to capture varying survival durations in a supervised setting through a hardness-aware negative sampling mechanism. Last, we incorporate the contrastive task into the time-to-event predictive task with multiple loss components. We conduct extensive experiments using a large EHR dataset to forecast the risk of hospitalized patients who are in danger of developing acute kidney injury (AKI), a critical and urgent medical condition. The effectiveness and explainability of the proposed model are validated through comprehensive quantitative and qualitative studies.

연구 동기 및 목표

• 기존 생존 모델이 코어닝된 데이터를 다루는 데서 약점이 있으며, 종단적 전자 건강 기록에서 분류 가능한 환자 표현을 학습하는 데에 한계가 있음을 해결하고자 한다.
• 특히 생존 지속 시간의 차이를 기반으로 하여 음성 샘플링의 어려움을 평가하는 대비 학습을 생존 분석에 적용해보고자 한다.
• 다중 작업 학습 접근법을 통해 생존 모델링에서 위험 순서 정렬과 정밀한 사건 발생 시간 예측을 동시에 향상시키고자 한다.
• 주목사 메커니즘을 활용하여 환자 기록 내에서 중요한 방문 기록과 의료 코드를 강조함으로써 모델의 해석 가능성 향상을 도모하고자 한다.
• 의료 온톨로지 지식을 표현 학습에 통합하여, 특히 희귀하거나 빈도가 낮은 의료 코드에 대해 더 나은 일반화 성능을 확보하고자 한다.

제안 방법

• 모델은 이중 인코더 아키텍처를 사용한다: ICD-9 코드를 계층적 의료 개념으로 매핑하는 온톨로지 인코더와 환자 방문 기록의 시간적 의존성을 모델링하는 시퀀셜 자기주도 인코더.
• 생존 지속 시간의 차이에 기반하여 환자 임베딩을 대비하는 감독된 가중 대비 손실(SupWCon)을 설계하였으며, 쌍 간의 시간 간격에 따라 온도 매개변수를 조정하여 샘플링의 어려움을 반영한다.
• 관찰된 환자와 코어닝된 환자의 생존 지속 시간을 기반으로 음성 쌍을 구성함으로써, 사건 발생 시간의 차이가 대비 학습에 영향을 미치도록 보장한다.
• 세 가지 목적을 동시에 최적화한다: SupWCon을 통한 대비 표현 학습, 콘cordance를 위한 순서 정렬 손실, 시간 예측을 위한 회귀 손실.
• 개별 인스턴스 수준과 방문 수준의 주목사 풀링 메커니즘을 사용하여 방문 및 개별 의료 코드에 대한 설명 가능한 주목사 가중치를 생성한다.
• 온톨로지 인코더는 의료 온톨로지(예: ICD-9)의 계층적 관계를 활용하여 희귀하거나 빈도가 낮은 코드의 표현을 향상시키며, 더 넓은 부모 개념에 주목함으로써 이를 달성한다.

Figure 1. Architecture of the proposed OTCSurv model. There are three main components: 1) Sequential attention-based ontological encoder, which mainly consists of an ontological encoder, two attention-pooling blocks, and a transformer encoder to learn the ultimate instance-level representations of p

실험 결과

연구 질문

• RQ1 생존 지속 시간의 차이로 측정되는 시간적 독창성에 기반한 대비 학습이 전자 건강 기록의 생존 분석에서 표현 학습을 향상시킬 수 있는가?
• RQ2 생존 지속 시간의 차이를 음성 샘플링의 어려움에 통합함으로써, 균일하거나 무작위 음성 샘플링 대비 성능 향상이 이루어지는가?
• RQ3 의료 온톨로지 지식을 통합함으로써 표현 품질이 향상되며, 특히 전자 건강 기록에서 희귀하거나 빈도가 낮은 의료 코드에 대해 어떤 정도의 향상이 이루어지는가?
• RQ4 대비, 순서 정렬, 회귀 손실을 동시에 통합한 다중 작업 학습 프레임워크가 위험 순서 정렬과 시간 예측 정확도를 동시에 향상시킬 수 있는가?
• RQ5 모델의 주목사 가중치가 임상적으로 의미 있는 패턴, 예를 들어 중요한 방문 기록이나 영향력이 큰 진단을 어떻게 반영하는가?

주요 결과

• OTCSurv는 대규모 실생활 전자 건강 기록 데이터셋에서 AKI 예측 성능이 최신 기술 수준에 도달하였으며, c-index와 시간 예측 정확도 모두에서 기존 방법을 능가하였다.
• 제안된 SupWCon 손실은 생존 지속 시간의 차이에 따라 온도 매개변수를 동적으로 조정함으로써 하드 음성 쌍에서의 학습을 향상시켜 모델 성능을 크게 향상시켰다.
• 주목사 분석 결과, 높은 주목사 가중치를 가진 제5회 방문과 572(간질환), 410(심근경색), 287(출혈성 질환) 등의 코드가 AKI 발병에 가장 예측 가능함을 확인하였으며, 임상 문헌과 일치하였다.
• 온톨로지 인코더는 희귀 코드의 표현을 향상시키기 위해 더 넓은 부모 개념에 주목함으로써 효과적으로 작용하였으며, 예를 들어 ICD-9 코드 514는 조상 코드 460-519를 통해 주목사에 영향을 받았다.
• 절단 실험을 통해 대비 손실 또는 온톨로지 구성 요소를 제거할 경우 성능 저하가 심각하게 발생함을 확인하였으며, 이는 각 구성 요소의 기여도를 검증하였다.
• 모델는 강력한 해석 가능성을 보였으며, 주목사 가중치가 임상적으로 관련성이 있는 방문 기록과 진단을 명확히 식별함으로써 임상적 신뢰도와 투명성을 높였다.

Figure 2. Comparison of the mean survival curves of proposed model and baselines with Kaplan-Meier curve (for all patients).