[논문 리뷰] Prediction of ICD Codes with Clinical BERT Embeddings and Text Augmentation with Label Balancing using MIMIC-III

Brent Biseda, Gaurav Desai외 2명|arXiv (Cornell University)|2020. 08. 24.

Medical Coding and Health Information인용 수 8

한 줄 요약

이 논문은 MIMIC-III 데이터셋에서 임상 기록의 ICD 코드 예측 성능을 향상시키기 위해 Clinical BERT 임베딩과 문장 순서를 훈련 중에 뒤섞는 새로운 텍스트 증강 전략을 조합한 임상 ICD 코드 예측 모델을 제안한다. 또한 레이블 균형 조정 기법을 통해 성능을 향상시킨다. 이 모델은 상위 50개 ICD 코드에서 F1 스코어 0.75를 기록하여, 이전 연구보다 더 넓은 레이블 세트를 사용했음에도 불구하고 이 작업 분야에서 최고 수준의 성능을 달성한다.

ABSTRACT

This paper achieves state of the art results for the ICD code prediction task using the MIMIC-III dataset. This was achieved through the use of Clinical BERT (Alsentzer et al., 2019). embeddings and text augmentation and label balancing to improve F1 scores for both ICD Chapter as well as ICD disease codes. We attribute the improved performance mainly to the use of novel text augmentation to shuffle the order of sentences during training. In comparison to the Top-32 ICD code prediction (Keyang Xu, et. al.) with an F1 score of 0.76, we achieve a final F1 score of 0.75 but on a total of the top 50 ICD codes.

연구 동기 및 목표

• MIMIC-III 데이터셋의 임상 노트에서 ICD 코드 예측 성능 향상.
• 레이블 균형 조정 기법을 통해 ICD 코드 레이블의 클래스 불균형 문제 해결.
• 텍스트 증강을 통한 문장 순서 뒤섞기를 통해 모델의 일반화 능력과 강건성 향상.
• Clinical BERT 임베딩을 사용하여 ICD 코드 예측 분야에서 최고 수준의 성능 달성.

제안 방법

• 임상 노트 텍스트를 고밀도 벡터 표현으로 변환하기 위해 Clinical BERT 임베딩 활용.
• 훈련 중에 임상 노트의 문장 순서를 무작위로 뒤섞는 새로운 텍스트 증강 기법 적용.
• ICD 코드의 장꼬리 분포 영향을 줄이기 위해 레이블 균형 조정 전략 구현.
• 각 노트에 대해 다중 ICD 코드를 예측하기 위해 Clinical BERT 위에 다중 레이블 분류 헤드 훈련.
• ICD 챕터 및 ICD 질병 코드 예측 작업에 대해 F1 스코어 평가.
• 기존 최고 수준의 방법들과의 성능 비교를 위해 상위 50개 ICD 코드 예측 설정 채택.

실험 결과

연구 질문

• RQ1 텍스트 증강을 통한 문장 순서 뒤섞기가 임상 노트에서 ICD 코드 예측 모델의 일반화 능력 향상에 기여하는가?
• RQ2 레이블 균형 조정은 Clinical BERT를 사용한 다중 레이블 ICD 코드 분류 성능에 어떤 영향을 미치는가?
• RQ3 제안된 방법은 MIMIC-III 데이터셋에서 기존 최고 수준의 접근 방식을 초월하는가?
• RQ4 Clinical BERT 임베딩의 사용이 다중 레이블 ICD 코드 예측에서 F1 스코어 향상에 얼마나 기여하는가?

주요 결과

• 모델은 상위 50개 ICD 코드에서 F1 스코어 0.75를 기록하여 MIMIC-III 데이터셋에서 최고 수준의 성능을 달성했다.
• 텍스트 증강을 통한 문장 순서 뒤섞기가 모델 성능 향상의 주요 원동력으로 확인되었다.
• 레이블 균형 조정 기법은 장꼬리 분포의 영향을 효과적으로 완화하여 모델의 강건성을 향상시켰다.
• 이전 최고 수준의 방법들보다 F1 스코어에서 뛰어난 성능을 보였으며, 이는 이전의 상위 32개 코드 대비 더 넓은 50개 코드 예측 세트를 사용했음에도 불구하고 가능했다.
• ICD 챕터 및 ICD 질병 코드 예측 작업 모두에서 일관된 성능 향상이 관찰되었다.
• 결과는 아키텍처 선택 사항인 텍스트 증강 및 레이블 균형 조정이 Clinical BERT를 활용한 고성능 ICD 코드 예측에 핵심적인 역할을 한다는 점을 확인시켰다.