[논문 리뷰] Three-level Hierarchical Transformer Networks for Long-sequence and Multiple Clinical Documents Classification

Yuqi Si, Kirk Roberts|arXiv (Cornell University)|2021. 04. 17.

Machine Learning in Healthcare인용 수 5

한 줄 요약

이 논문은 문장에서 환자 수준의 결과까지의 정보를 단어 → 문장, 문장 → 임상 노트, 노트 → 환자 수준의 결과로 처리하는 삼단계 계층적 트랜스포머 네트워크를 제안한다. BERT의 컨텍스트 길이를 연장하고 시간 인식형 적응형 세그먼테이션을 통합함으로써, 모델은 최신 기술 수준의 성능을 달성하여 MIMIC-III에서 입원 중 사망률 및 표현형 예측 작업에서 AUC는 최대 21% 향상되고, PRC는 51%, F1은 46% 향상된다.

ABSTRACT

We present a Three-level Hierarchical Transformer Network (3-level-HTN) for modeling long-term dependencies across clinical notes for the purpose of patient-level prediction. The network is equipped with three levels of Transformer-based encoders to learn progressively from words to sentences, sentences to notes, and finally notes to patients. The first level from word to sentence directly applies a pre-trained BERT model as a fully trainable component. While the second and third levels both implement a stack of transformer-based encoders, before the final patient representation is fed into a classification layer for clinical predictions. Compared to conventional BERT models, our model increases the maximum input length from 512 tokens to much longer sequences that are appropriate for modeling large numbers of clinical notes. We empirically examine different hyper-parameters to identify an optimal trade-off given computational resource limits. Our experiment results on the MIMIC-III dataset for different prediction tasks demonstrate that the proposed Hierarchical Transformer Network outperforms previous state-of-the-art models, including but not limited to BigBird.

연구 동기 및 목표

• 장기적인 임상 노트를 모델링할 때 BERT의 512토큰 입력 길이 제한을 해결한다.
• 환자당 다수의 임상 문서 간의 계층적이고 시간적 종속성을 포착한다.
• 입원 중 사망률 및 표현형 분류와 같은 임상 결과에 대한 환자 수준의 예측 정확도를 향상시킨다.
• 하이퍼파라미터 튜닝과 분산 학습을 통해 계산 자원 제약 조건 하에서 모델 성능을 최적화한다.
• 고립된 노트(예: 퇴원 요약서)에 의존하는 것보다 전체 장기적 노트 시퀀스를 통합하는 것이 필수적임을 입증한다.

제안 방법

• 단어에서 문장으로의 전이에 사전 학습된 BERT 모델을 사용하여 문맥적 표현을 추출한다.
• 문장에서 노트로의 집계를 위한 두 번째 트랜스포머 기반 인코더 스택과, 노트에서 환자 분류로의 전이를 위한 또 다른 스택을 구현한다.
• 시간적 근접도를 기반으로 임상 노트를 그룹화하는 시간 인식형 적응형 세그먼테이션을 도입한다. 이는 장기적 구조를 유지한다.
• 문장 및 노트 수준에서 표현을 집계하기 위해 mean_max 풀링을 사용하여 계층적 특징 학습을 향상시킨다.
• 모든 세 수준에서 위치 인코딩을 적용하여 순서 인식 능력을 유지한다.
• 분산 학습을 활용하여 GPU 메모리 제약을 초월하고, BERT의 표준 제한을 넘는 더 긴 입력 시퀀스 처리를 가능하게 한다.

실험 결과

연구 질문

• RQ1 계층적 트랜스포머 아키텍처는 다수의 환자 노트를 포함하는 장기적인 임상 시퀀스를 효과적으로 모델링할 수 있는가?
• RQ2 적응형 세그먼테이션을 통한 시간적 구조 통합은 노트를 독립적으로 취급하는 것과 비교해 임상 예측 성능을 어떻게 향상시키는가?
• RQ3 임상 시퀀스 모델링에서 모델 크기, 입력 길이, 계산 비용 간 최적의 트레이드오프는 무엇인가?
• RQ4 퇴원 요약서 외에 모든 가용 임상 노트를 사용할 경우 성능 향상이 뚜렷한가?
• RQ5 풀링 전략, 인코더 레이어 수, 위치 인코딩 등의 아키텍처 구성 요소 중 어느 것이 모델 성능에 가장 큰 영향을 미치는가?

주요 결과

• 제안된 계층적 트랜스포머 네트워크는 입원 중 사망률 예측에서 이전 최신 기술 수준 모델인 BigBird 대비 최대 21% 향상된 AUC를 달성한다.
• 표현형 예측 작업에서 PRC는 51% 향상되고 F1 스코어는 46% 향상되어 다양한 지표에서 뚜렷한 성과 향상을 보인다.
• 모든 가용 임상 노트를 사용할 경우 성능 향상이 뚜렷하며, 단 한 개의 노트(예: 퇴원 요약서)만 사용할 경우 AUC는 약 0.85로 감소한다.
• 성능이 가장 뛰어난 모델은 문장 및 문서 수준에서 각각 2개의 인코더 레이어를 사용하며, 이는 더 큰 모델이 항상 더 좋은 결과를 내는 것은 아니며 최적화 장애물의 영향을 받기 때문이다.
• 적응형 세그먼테이션은 필수적이다. 이를 제거할 경우 AUC는 0.8558로 감소하고 PRC는 0.7887로 떨어지며, 시간적 관계 유지 기능이 뚜렷하게 드러난다.
• mean_max 풀링이 다른 풀링 전략보다 우수하며, 위치 인코딩은 계층 수준 간의 순서 인식 능력을 유지하는 데 필수적이다.