[논문 리뷰] Bootstrapping Your Own Positive Sample: Contrastive Learning With Electronic Health Record Data

Tingyi Wanyan, Jing Zhang외 4명|arXiv (Cornell University)|2021. 04. 07.

Machine Learning in Healthcare참고 문헌 37인용 수 8

한 줄 요약

이 논문은 전자 건강 기록(EHR) 데이터를 위한 새로운 양성 샘플링 전략—특성 기반 및 특성 기반—을 갖춘 대비 학습 프레임워크를 제안하여 코로나19 환자의 사망 예측 성능을 향상시킨다. 특성 공간 내 환자 유사도와 사전 계산된 유사도 지표를 활용함으로써, 이는 클래스 간 분리도를 향상시키고 클래스 내 이질성을 유지하며, AUROC 0.959를 달성하여 교차 엔트로피(0.873)와 포칼 손실(0.931)을 능가한다.

ABSTRACT

Electronic Health Record (EHR) data has been of tremendous utility in Artificial Intelligence (AI) for healthcare such as predicting future clinical events. These tasks, however, often come with many challenges when using classical machine learning models due to a myriad of factors including class imbalance and data heterogeneity (i.e., the complex intra-class variances). To address some of these research gaps, this paper leverages the exciting contrastive learning framework and proposes a novel contrastive regularized clinical classification model. The contrastive loss is found to substantially augment EHR-based prediction: it effectively characterizes the similar/dissimilar patterns (by its "push-and-pull" form), meanwhile mitigating the highly skewed class distribution by learning more balanced feature spaces (as also echoed by recent findings). In particular, when naively exporting the contrastive learning to the EHR data, one hurdle is in generating positive samples, since EHR data is not as amendable to data augmentation as image data. To this end, we have introduced two unique positive sampling strategies specifically tailored for EHR data: a feature-based positive sampling that exploits the feature space neighborhood structure to reinforce the feature learning; and an attribute-based positive sampling that incorporates pre-generated patient similarity metrics to define the sample proximity. Both sampling approaches are designed with an awareness of unique high intra-class variance in EHR data. Our overall framework yields highly competitive experimental results in predicting the mortality risk on real-world COVID-19 EHR data with a total of 5,712 patients admitted to a large, urban health system. Specifically, our method reaches a high AUROC prediction score of 0.959, which outperforms other baselines and alternatives: cross-entropy(0.873) and focal loss(0.931).

연구 동기 및 목표

• 임상 예측 작업을 위한 EHR 데이터에서 클래스 불균형과 높은 클래스 내 분산 문제를 해결한다.
• 시각 작업과는 달리 데이터 증강이 덜 가능하기 때문에, 전통적인 대비 학습이 EHR에서 가지는 한계를 극복한다.
• 의미적 유사성을 반영하면서도 동일한 결과 클래스 내에서 환자 다양성을 유지하는 EHR 데이터에 특화된 양성 샘플링 전략을 개발한다.
• 최소한의 감독 하에 대비 정규화를 사용하여 불균형한 EHR 데이터에서 임상 모델의 예측 성능을 향상시킨다.
• 대비 학습과 정보 있는 양성 샘플링이 실제 EHR 응용에서 더 나은 일반화 및 강인성을 이끌어낸다는 것을 입증한다.

제안 방법

• 클래스 불균형을 다루고 특성 학습을 향상시키기 위해 포칼 손실과 대비 손실을 결합한 대비 정규화된 분류 모델을 도입한다.
• 학습된 특성 공간 내에서 이웃을 선택하여 양성 쌍을 형성하는 특성 기반 양성 샘플링 전략을 제안하며, 이는 국소적 이웃 학습을 향상시킨다.
• 의료 특성 기반 등으로 계산된 사전 유사도 지표(예: 임상 특성 기반)를 사용하여 의미적으로 유의미한 양성 쌍을 정의하는 특성 기반 양성 샘플링 전략을 개발한다.
• ‘끌어당기기 및 밀어내기’ 메커니즘을 갖춘 대비 손실 함수를 적용한다: 유사한 환자(양성 샘플)는 임베딩 공간에서 가까이 있도록 하고, 비유사한 환자(음성 샘플)는 멀어지도록 한다.
• UMAP 시각화와 임베딩 분리 점수(ESS)를 사용하여 학습된 표현 공간에서의 클래스 간 및 클래스 내 군집 품질을 평가한다.
• 선형 분류기를 대비 임베딩에 적용하여 후행 예측 성능를 평가하고, 학습된 표현의 품질을 검증한다.

실험 결과

연구 질문

• RQ1 표준 감독 손실에 비해 맞춤형 양성 샘플링 전략을 갖춘 대비 학습이 불균형한 EHR 데이터에서 예측 성능을 향상시킬 수 있는가?
• RQ2 특성 기반 및 특성 기반 샘플링 전략은 랜덤 샘플링에 비해 EHR 데이터의 클래스 내 이질성을 얼마나 잘 포착하는가?
• RQ3 제안된 샘플링 전략이 학습된 임베딩 공간에서 클래스 간 분리도를 어느 정도 향상시키는가?
• RQ4 대비 정규화는 EHR에서 높은 데이터 이질성과 클래스 불균형이 존재하는 환경에서 더 강인하고 일반화 가능한 표현을 만들어내는가?
• RQ5 제안된 방법은 코로나19 환자 입원 시 사망 예측과 같은 실제 임상 예측 작업에서 뛰어난 성능을 달성할 수 있는가?

주요 결과

• 제안된 대비 정규화 모델은 실제 코로나19 EHR 데이터에서 AUROC 0.959를 달성하여 교차 엔트로피(0.873)와 포칼 손실(0.931)을 크게 능가했다.
• 특성 기반 양성 샘플링 전략이 가장 높은 성능을 보였으며, 랜덤 샘플링 대비 선형 분류기 정확도에서 4.5% 향상되었다.
• 특성 기반 및 특성 기반 샘플링 전략은 랜덤 샘플링보다 높은 임베딩 분리 점수(ESS)를 달성하여 더 나은 클래스 간 분리도를 나타냈다.
• 두 새로운 샘플링 전략은 랜덤 샘플링보다 더 높은 클래스 내 분산(클러스터 중심으로부터의 거리 표준편차)을 유지하여 환자 이질성을 더 잘 모델링했다.
• UMAP 시각화 결과, 제안된 방법은 랜덤 샘플링에 비해 사망자 및 비사망자 환자에 대해 더 명확하고 잘 분리된 군집을 생성했다.
• 대비 정규화는 랜덤 양성 샘플링에서 흔히 발생하는 임베딩 공간의 붕괴를 완화시켰으며, 각 클래스 내 다양한 서브군 표현을 유지함으로써 효과를 보였다.