[논문 리뷰] Few-Shot Electronic Health Record Coding through Graph Contrastive Learning

Shanshan Wang, Pengjie Ren외 6명|arXiv (Cornell University)|2021. 06. 29.

Machine Learning in Healthcare참고 문헌 53인용 수 6

한 줄 요약

이 논문은 소수의 샘플만 존재하는 희귀 ICD 코드와 일반적인 ICD 코드 모두에서 성능을 향상시키는 소수 샘플 전자 건강 기록(EHR) 코드화를 위한 그래프 대비 학습 프레임워크인 CoGraph를 제안한다. 이는 이질적 단어-엔티티(HEWE) 그래프를 구성하고, 그래프 샘플링 대비 학습(GSCL) 및 그래프 진화 대비 학습(GECL)이라는 두 가지 대비 학습 기법을 적용하여 구조적 패턴에서 이식 가능한 표현을 학습함으로써, 희귀 ICD 코드에서 최대 2.95%의 F1 향상과 함께 일반 코드에서의 뛰어난 성능를 유지한다.

ABSTRACT

Electronic health record (EHR) coding is the task of assigning ICD codes to each EHR. Most previous studies either only focus on the frequent ICD codes or treat rare and frequent ICD codes in the same way. These methods perform well on frequent ICD codes but due to the extremely unbalanced distribution of ICD codes, the performance on rare ones is far from satisfactory. We seek to improve the performance for both frequent and rare ICD codes by using a contrastive graph-based EHR coding framework, CoGraph, which re-casts EHR coding as a few-shot learning task. First, we construct a heterogeneous EHR word-entity (HEWE) graph for each EHR, where the words and entities extracted from an EHR serve as nodes and the relations between them serve as edges. Then, CoGraph learns similarities and dissimilarities between HEWE graphs from different ICD codes so that information can be transferred among them. In a few-shot learning scenario, the model only has access to frequent ICD codes during training, which might force it to encode features that are useful for frequent ICD codes only. To mitigate this risk, CoGraph devises two graph contrastive learning schemes, GSCL and GECL, that exploit the HEWE graph structures so as to encode transferable features. GSCL utilizes the intra-correlation of different sub-graphs sampled from HEWE graphs while GECL exploits the inter-correlation among HEWE graphs at different clinical stages. Experiments on the MIMIC-III benchmark dataset show that CoGraph significantly outperforms state-of-the-art methods on EHR coding, not only on frequent ICD codes, but also on rare codes, in terms of several evaluation indicators. On frequent ICD codes, GSCL and GECL improve the classification accuracy and F1 by 1.31% and 0.61%, respectively, and on rare ICD codes CoGraph has more obvious improvements by 2.12% and 2.95%.

연구 동기 및 목표

• 희귀 ICD 코드의 예측 성능이 부족한 이유인 훈련 샘플 부족으로 인한 심각한 클래스 불균형 문제를 해결한다.
• 기존 방법에서 자주 忽略되는 희귀 코드에 대해서도 일반 코드와 마찬가지로 높은 ICD 코드 예측 성능을 향상시키는 것을 목표로 한다.
• 자기지도 그래프 표현 학습을 통해 일반적인 ICD 코드에서 희귀 ICD 코드로 지식을 전이할 수 있는 소수 샘플 학습 프레임워크를 개발한다.
• 모델이 일반 ICD 코드에 유용한 특징만 학습하는 것을 방지하기 위해, 대비 학습을 위한 내재된 그래프 구조를 활용하여 전이 가능한 표현을 학습한다.
• 외부 의료 지식에 의존하지 않고도 EHR 내부의 내재된 상관관계를 활용하여 그래프 내부 및 간섭 상관관계를 탐색하는 대비 학습 기법을 설계한다.

제안 방법

• 각 EHR에 대해 단어와 의료 엔티티를 노드로, 공시현출 또는 언급 관계 등 관계를 간선으로 가지는 이질적 단어-엔티티(HEWE) 그래프를 구성한다.
• 각 ICD 코드 예측을 소수 샘플 학습 설정에서의 메타태스크로 간주하여 일반적인 코드에서 희귀 코드로 지식을 전이할 수 있도록 한다.
• 동일한 HEWE 그래프에서 샘플링한 부분 그래프를 대비하여 불변성을 학습하는 그래프 샘플링 대비 학습(GSCL)을 도입한다.
• 동일 환자의 임상 단계에 따라 이전 및 이후 HEWE 그래프를 대비하여 시간적 동적 변화를 캡처하는 그래프 진화 대비 학습(GECL)을 제안한다.
• ICD 분류에 대한 미세조정 이전에 이 대비 목적함수를 사용해 그래프 인코더를 사전학습하여 이식 가능한 특징을 학습한다.
• 훈련 중 소수 샘플 시나리오를 시뮬레이션하기 위해 지지세트와 쿼리세트를 활용한 메타학습 전략을 사용하여 희귀 ICD 코드에 대한 일반화 능력을 확보한다.

실험 결과

연구 질문

• RQ1 제한된 레이블이 있는 소수 샘플 환경에서 소수 샘플 학습 프레임워크가 희귀 코드 예측 성능을 효과적으로 향상시킬 수 있는가?
• RQ2 외부 의료 지식에 의존하지 않고도 EHR 내부의 내재된 구조적 패턴을 활용하여 이식 가능한 표현을 학습할 수 있는가?
• RQ3 GSCL 및 GECL과 같은 그래프 대비 학습 기법이 표준 지도 미세조정 대비 EHR 코드화의 표현 학습에 얼마나 효과적인가?
• RQ4 희귀 ICD 코드의 성능은 소수 샘플 지지세트 크기(K)가 증가함에 따라 향상되는가? 이는 일반 코드의 성능에 어떤 영향을 미치는가?
• RQ5 훈련 클래스에 대한 다양한 샘플링 전략(예: 무작위 대비 'on-top')이 희귀 ICD 코드로의 일반화 능력에 어떤 영향을 미치는가?

주요 결과

• 최신 기술(SOTA) 기법 대비 CoGraph는 일반 ICD 코드에서 분류 정확도가 1.31%p 향상되고 F1 점수는 0.61%p 향상되었다.
• 희귀 ICD 코드에서는 더 큰 성과를 보이며, 정확도는 2.12%p 향상되고 F1 점수는 2.95%p 향상되어 자원이 부족한 ICD 코드에 대한 효과성을 입증했다.
• 소수 샘플 지지세트 크기(K)가 1을 초과할수록 일반 ICD 코드의 성능은 약간 감소함을 확인하여 일반 코드와 희귀 코드 간의 학습 간 균형 문제가 존재함을 시사했다.
• 훈련용 ICD 클래스 수가 증가함에 따라 희귀 ICD 코드의 성능은 안정적이고 높은 수준을 유지하여 일반 코드에서 희귀 코드로의 효과적인 지식 전이가 이루어졌음을 시사했다.
• 희귀 ICD 코드의 정밀도는 재현율 및 정확도보다 뚜렷하게 높아, 모델이 희귀 코드를 보수적으로 예측하는 경향이 있음을 나타내며, 특히 매우 희귀한 코드의 재현율 향상 여지가 있음을 시사한다.
• 무작위 샘플링 대비 'on-top' 샘플링 전략은 훈련 분포의 클래스 불균형으로 인해 더 불안정한 훈련 동역학을 보였지만, 최종 성능에 있어 희귀 코드 예측에 큰 영향을 주지 않았다.