[논문 리뷰] Hierarchical BERT for Medical Document Understanding

Ning Zhang, Maciej Jankowski|arXiv (Cornell University)|2022. 03. 11.

Topic Modeling인용 수 10

한 줄 요약

이 논문은 하향식 방식으로 토큰, 문장, 문서 수준에서 텍스트를 처리함으로써 장문의 의료 문서에 대한 효과적이고 효율적이며 일반적인 표현 학습을 가능하게 하는 계층적 BERT 아키텍처인 Medical Document BERT (MDBERT)를 제안한다. MDBERT는 ICD 코드 할당 작업에서 최신 기술 수준의 성능을 달성하며(20%의 상대적 향상, MIMIC-III 기준), 기존 BERT의 자기주의 복잡도를 1/100 미만으로 줄이고, 다양한 후행 NLP 작업을 위한 이식 가능한 문장 및 문서 수준의 임베딩을 제공한다.

ABSTRACT

Medical document understanding has gained much attention recently. One representative task is the International Classification of Disease (ICD) diagnosis code assignment. Existing work adopts either RNN or CNN as the backbone network because the vanilla BERT cannot handle well long documents (>2000 to kens). One issue shared across all these approaches is that they are over specific to the ICD code assignment task, losing generality to give the whole document-level and sentence-level embedding. As a result, it is not straight-forward to direct them to other downstream NLU tasks. Motivated by these observations, we propose Medical Document BERT (MDBERT) for long medical document understanding tasks. MDBERT is not only effective in learning representations at different levels of semantics but efficient in encoding long documents by leveraging a bottom-up hierarchical architecture. Compared to vanilla BERT solutions: 1, MDBERT boosts the performance up to relatively 20% on the MIMIC-III dataset, making it comparable to current SOTA solutions; 2, it cuts the computational complexity on self-attention modules to less than 1/100. Other than the ICD code assignment, we conduct a variety of other NLU tasks on a large commercial dataset named as TrialTrove, to showcase MDBERT's strength in delivering different levels of semantics.

연구 동기 및 목표

• 표준 BERT의 512토큰 시퀀스 길이 제약로 인해 장문의 임상 노트(2000토큰 초과)를 처리하는 데에 한계가 존재하는 문제를 해결하기 위해.
• 기존 ICD 코드 할당 모델이 전용 작업에 특화되어 있어 일반적인 문서 및 문장 수준의 임베딩을 제공하지 못하는 문제를 해결하기 위해.
• 다양한 NLP 작업 간의 이식 가능성을 향상시키기 위해 토큰, 문장, 문서 수준의 다중 수준 의미 표현을 동시에 학습하는 통합형 엔드 투 엔드 모델을 개발하기 위해.
• 자기주의 복잡도를 줄여 장문의 의료 문서를 효율적으로 인코딩하기 위해.
• MDBERT가 학습한 임베딩의 일반화 능력을 다양한 임상 NLP 작업(예: ICD 코딩, 환자 세그먼트 분류, 표준화된 적응증 정의)에 걸쳐 입증하기 위해.

제안 방법

• MDBERT는 서로 다른 문서의 문장을 재배치하여 문서 수준에서의 공동 학습을 가능하게 하는 하향식 계층적 아키텍처를 사용한다.
• Transformer 블록은 다중 헤드 자기주의 복합체를 사용하여 상대적 위치 인코딩을 적용하여 각 의미 수준(토큰, 문장, 문서)의 표현을 처리한다.
• 의미 표현은 계층적으로 평균 풀링을 통해 집계되며, 임의의 수준에 작업 전용 헤드를 부착함으로써 엔드 투 엔드 학습이 가능하다.
• 모델은 코드 설명 증강을 학습 과정에 직접 통합하여 별도의 코드 임베딩 네트워크가 필요 없도록 한다.
• 문장 수준의 임베딩은 첫 번째 평균 풀링 레이어에서 유도되며, 문서 수준의 임베딩은 최종 풀링 표현에서 확보된다.
• 여러 후행 작업(예: ICD 코딩, 성별 분류, 환자 세그먼트 그룹화) 간의 다중 작업 학습을 적용하여 일반화 능력과 강건성을 향상시킨다.

실험 결과

연구 질문

• RQ1 표준 BERT의 512토큰 제한을 초월하는 장문의 의료 문서를 처리할 수 있는 BERT 기반 모델이 성능 저하 없이 효과적으로 작동할 수 있는가?
• RQ2 통합형 모델이 다양한 임상 NLP 작업 간 이식 가능한 일반 목적의 문서 및 문장 수준의 임베딩을 학습할 수 있는가?
• RQ3 계층적 BERT 아키텍처가 장문의 문서 이해 작업에서 성능을 유지하거나 향상시키면서도 계산 복잡도를 줄일 수 있는가?
• RQ4 MDBERT의 통합 임베딩 공간이 ICD 코드 할당, 환자 세그먼트 분류, 표준화된 적응증 정의와 같은 후행 작업을 지원하는 데 얼마나 효과적인가?
• RQ5 MDBERT의 다중 작업 학습 전략이 전용 작업에 맞춰 미세조정하는 것보다 자원이 적거나 복잡한 임상 NLP 작업에서 성능을 향상시킬 수 있는가?

주요 결과

• MDBERT는 MIMIC-III 데이터셋에서 ICD 코드 할당 작업에서 최대 20%의 상대적 성능 향상을 달성하며, 최신 기술 수준의 결과를 재현한다.
• 모델은 자기주의 모듈의 계산 복잡도를 기존 BERT의 1/100 미만으로 줄여 장문의 문서 처리를 효율적으로 가능하게 한다.
• MDBERT의 문서 수준 임베딩은 여러 후행 작업으로 잘 일반화되며, TrialTrove 데이터셋에서 연령대 분류에 대해 98.4% AUC-ROC, 환자 세그먼트 분류에 대해 99.8%를 기록한다.
• MDBERT의 문장 수준 임베딩은 표준화된 적응증 정의 작업에서 기존 베이스라인을 초월하여 SNOMED 매칭 작업에서 96.8% AUC 스코어를 기록하며, SBERT(95.3%)와 BioBERT(94.5%)를 모두 앞선다.
• 다중 작업 학습은 ICD10 코드 예측 성능을 향상시켜 관련 임상 NLP 작업 간 공유 표현 학습의 이점을 입증한다.
• MDBERT는 과적합의 징후를 보이며, 특히 작은 데이터셋에서는 약 10에포크 이후 성능이 정체되는 경향이 있어, 향후 연구에서는 더 강력한 정규화나 문장 수준의 제약 조건이 필요할 것으로 보인다.