[논문 리뷰] DKEC: Domain Knowledge Enhanced Multi-Label Classification for Diagnosis Prediction

Xueren Ge, Abhishek, Satpathy외 3명|arXiv (Cornell University)|2023. 10. 10.

Text and Document Classification Technologies인용 수 4

한 줄 요약

DKEC는 희귀(꼬리) 클래스에서 성능을 향상시키기 위해 이질적인 의료 엔티티 관계를 레이블별 주의 메커니즘과 그룹별 학습 전략을 통해 통합하는 도메인 지식을 통합한 다중 레이블 의료 진단 예측 프레임워크를 제안한다. 이는 희소한 샘플을 가진 레이블에 특히 유리하며, 소형 언어 모델이 대형 모델과 유사한 성능을 달성할 수 있도록 한다. 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Multi-label text classification (MLTC) tasks in the medical domain often face the long-tail label distribution problem. Prior works have explored hierarchical label structures to find relevant information for few-shot classes, but mostly neglected to incorporate external knowledge from medical guidelines. This paper presents DKEC, Domain Knowledge Enhanced Classification for diagnosis prediction with two innovations: (1) automated construction of heterogeneous knowledge graphs from external sources to capture semantic relations among diverse medical entities, (2) incorporating the heterogeneous knowledge graphs in few-shot classification using a label-wise attention mechanism. We construct DKEC using three online medical knowledge sources and evaluate it on a real-world Emergency Medical Services (EMS) dataset and a public electronic health record (EHR) dataset. Results show that DKEC outperforms the state-of-the-art label-wise attention networks and transformer models of different sizes, particularly for the few-shot classes. More importantly, it helps the smaller language models achieve comparable performance to large language models.

연구 동기 및 목표

• 희귀 진단에 대한 훈련 샘플이 부족한 의료 다중 레이블 텍스트 분류에서의 장꼬리 레이블 분포 문제를 해결하기 위해.
• 의료 지침에서 구조화된 도메인 지식을 신경망 모델에 통합하여 자원이 적은 레이블의 표현 학습을 향상시키기 위해.
• 소형 언어 모델이 진단 예측 과제에서 대형 언어 모델과 유사한 성능을 달성할 수 있도록 하는 방법을 개발하기 위해.
• 기존 온톨로지 외에도 다양한 의료 엔티티 관계(예: 증상, 약물, 절차)를 통합하는 일반화 가능한 프레임워크를 구축하기 위해.
• 다양한 언어 모델과 실제 EHR 및 ePCR 데이터셋을 대상으로 방법을 평가하기 위해.

제안 방법

• 의료 지침에서 이질적인 의료 엔티티(예: 증상, 약물, 절차)와 이질적인 관계(예: '치료함', '의미함')를 통합한 이질적 그래프를 구축한다.
• 레이블별 주의 메커니즘이 이질적 그래프를 사용하여 동적 주의 가중치를 계산하여, 각 레이블에 대해 관련 의료 엔티티에 모델의 주의를 집중시킨다.
• 이질적 그래프에서의 임베딩을 사용하여 레이블을 의미적 유사도 기반으로 그룹화하여 관련 진단의 클러스터를 형성한다.
• 그룹별 학습은 최적화 동안 전체 레이블 그룹에 대해 이진 교차 엔트로피 손실을 계산하여 희귀 클래스의 효과적 샘플 수를 증가시킨다.
• 후처리 단계에서 논리 규칙(예: 연령 기반 구분)을 적용하여 그룹 내에서 겹치거나 모호한 레이블의 정밀도를 향상시킨다.
• 이 프레임워크는 소형 모델(예: TinyClinicalBERT)을 포함한 다양한 언어 모델과 통합되어 일반화 능력과 효율성을 평가한다.

실험 결과

연구 질문

• RQ1 공식 지침에서 유래한 이질적 의료 엔티티 관계를 통합함으로써 다중 레이블 의료 텍스트 분류에서 표현 학습이 향상되는가?
• RQ2 레이블의 의미적 유사도 기반으로 그룹화된 학습 전략이 데이터 불균형 문제를 효과적으로 완화하고 희귀(꼬리) 클래스의 성능을 향상시키는가?
• RQ3 도메인 지식과 그룹별 학습 전략을 통한 향상으로 소형 언어 모델이 대형 언어 모델과 유사한 성능을 달성할 수 있는가?
• RQ4 DKEC 프레임워크는 다양한 언어 모델과 실제 의료 데이터셋에 대해 얼마나 일반화 가능한가?
• RQ5 지침에서 유래한 도메인 지식 통합이 불균형한 의료 데이터셋에서 빈도가 높은(머리) 레이블에 대한 편향을 얼마나 줄이는가?

주요 결과

• RAA 데이터셋에서 DKEC는 최신 기술 수준의 방법보다 매크로 F1 스코어를 8.2% 향상시켜 32.8%까지 상승시키며, 특히 희귀(꼬리) 클래스에서 가장 큰 성과를 기록한다.
• MIMIC-III 데이터셋에서 DKEC-Comb는 TinyClinicalBERT(15M 파라미터)를 사용해 매크로 F1 0.499를 달성했으며, 기준 모델인 BioMedLM(27억 파라미터)의 0.459를 초월한다.
• MIMIC-III에서 TinyClinicalBERT에 DKEC-Comb를 적용하면 마이크로 F1 0.516을 기록했고, 기준 모델인 GatorTron(3억 2500만 파라미터)의 0.523을 초월하여 강력한 일반화 능력을 보였다.
• 동일한 모델(TinyClinicalBERT)을 사용할 경우 DKEC-Comb는 매크로 F1 0.234를 기록했으며, 기준 모델인 BioMedLM(0.248)보다 약간 뛰어나 소형 모델에 대한 효과성을 입증한다.
• 그룹별 학습과 이질적 레이블별 주의 메커니즘은 특히 자원이 적은 진단에 대한 소수 샘플 학습에서 크게 향상시킨다.
• 이 방법은 소형 모델이 대형 언어 모델과 유사한 성능을 달성할 수 있도록 하여 자원 제약이 있는 장치에의 구현에 적합하다.