[논문 리뷰] KG-MTT-BERT: Knowledge Graph Enhanced BERT for Multi-Type Medical Text Classification

Yong He, Cheng Wang외 4명|arXiv (Cornell University)|2022. 10. 08.

Topic Modeling인용 수 13

한 줄 요약

KG-MTT-BERT는 진단관계그룹(DRG) 예측을 위한 지식 그래프 증강 BERT 모델을 제안한다. 다양한 임상 텍스트 유형(예: 진단, 절차, 병원 경과)에 공통된 BERT 인코더를 적용하고, 토큰 수준의 인코딩과 전역 최대 풀링을 사용하며, 의료 지식 그래프를 통합함으로써 DRG 분류에서 최신 기술 수준(SOTA) 성능을 달성한다. 정확도 91.79%와 AUC 93.10%를 기록하며 기준 모델들을 크게 앞서며 성능을 뛰어나게 한다.

ABSTRACT

Medical text learning has recently emerged as a promising area to improve healthcare due to the wide adoption of electronic health record (EHR) systems. The complexity of the medical text such as diverse length, mixed text types, and full of medical jargon, poses a great challenge for developing effective deep learning models. BERT has presented state-of-the-art results in many NLP tasks, such as text classification and question answering. However, the standalone BERT model cannot deal with the complexity of the medical text, especially the lengthy clinical notes. Herein, we develop a new model called KG-MTT-BERT (Knowledge Graph Enhanced Multi-Type Text BERT) by extending the BERT model for long and multi-type text with the integration of the medical knowledge graph. Our model can outperform all baselines and other state-of-the-art models in diagnosis-related group (DRG) classification, which requires comprehensive medical text for accurate classification. We also demonstrated that our model can effectively handle multi-type text and the integration of medical knowledge graph can significantly improve the performance.

연구 동기 및 목표

• 전자 건강 기록(EHR)에서 유형이 다양한 장문의 의료 텍스트를 진단관계그룹(DRG)으로 분류하는 문제를 해결하기 위해.
• BERT가 장문의 임상 노트, 다양한 텍스트 유형, 영역 특화 의료 용어를 처리하는 데 한계를 보이는 점을 극복하기 위해.
• BERT 기반 아키텍처에 의료 지식 그래프를 통합하여 DRG 분류 성능을 향상시키기 위해.
• 도메인 지식과 이질적인 임상 텍스트 입력을 효과적으로 인코딩하고 융합하는 통합 모델을 개발하기 위해.

제안 방법

• 각 임상 텍스트 유형(예: 진단, 절차, 병원 경과)을 별도로 처리하기 위해 공통된 BERT 인코더를 적용하여 입력 모odal 특이성을 유지한다.
• 텍스트 수준의 인코딩 대신 토큰 수준의 인코딩을 사용하여 세밀한 의미 표현을 캡처함으로써 F1 점수에서 1.6% 이상 향상된다.
• DRG 할당에 핵심적인 의료 용어를 강조하기 위해 토큰 표현에 대해 전역 최대 풀링을 적용한다.
• 의료 지식 그래프를 통합하여 모델에 도메인 특화 지식을 주입함으로써 의료 용어와 관계에 대한 이해를 향상시킨다.
• 다양한 텍스트 유형의 풀링된 표현을 통합하여 분류를 위한 통합 표현 행렬로 결합한다.
• 최종 분류 레이어를 사용하여 DRG 카테고리 예측을 수행하며, 교차 엔트로피 손실을 사용해 엔드 투 엔드로 훈련한다.

실험 결과

연구 질문

• RQ1 장문이거나 의미 유형이 다른 EHR의 다중 유형 의료 텍스트를 BERT 기반 모델이 DRG 카테고리로 효과적으로 분류할 수 있는가?
• RQ2 DRG 분류에 관련된 의료 정보를 캡처하는 데 있어 토큰 수준의 인코딩이 텍스트 수준의 인코딩보다 우월한가?
• RQ3 의료 지식 그래프를 통합할 경우 BERT의 임상 텍스트 분류 성능에 어떤 영향을 미치는가?
• RQ4 전역 최대 풀링, 평균 풀링, 또는 어텐션 기반 풀링 중 어떤 전략이 이질적인 임상 텍스트 표현을 DRG 예측에 가장 잘 요약하는가?
• RQ5 도메인 특화 지식을 통합할 경우 의료 텍스트 분류의 일반화 능력과 정확도에 얼마나 기여하는가?

주요 결과

• KG-MTT-BERT는 DRG-A 데이터셋에서 91.79% 정확도, 89.61% F1, 93.10% AUC를 기록하며, 모든 기준 모델과 최신 기술 수준 모델을 크게 앞서 성능을 뛰어나게 한다.
• 토큰 수준의 인코딩은 텍스트 수준의 인코딩 대비 F1 점수에서 1.6% 이상 향상되며, 세밀한 표현의 이점을 입증한다.
• 전역 최대 풀링은 모든 풀링 전략 중에서 가장 높은 성능을 보이며, AUC 0.8741을 기록하여 평균 풀링 및 어텐션 기반 풀링을 앞선다.
• 의료 지식 그래프 통합으로 정확도가 4.76%포인트(87.03%에서 91.79%로) 향상되고 AUC는 5.69점 향상되어 도메인 적응에서의 핵심적 역할을 입증한다.
• 제거 실험(ablation study) 결과 지식 그래프와 토큰 수준의 인코딩이 가장 영향력 있는 구성 요소이며, 지식 그래프가 가장 큰 성능 향상을 기여한다.
• 어텐션 메커니즘을 통해 모델은 'stoma', 'anorectal'과 같은 핵심 의료 용어를 성공적으로 식별하며, 임상적 관련성과 일치하여 해석 가능성도 지원한다.