[논문 리뷰] CORE-BEHRT: A Carefully Optimized and Rigorously Evaluated BEHRT

Mikkel Odgaard, Kiril Vadimovic Klein외 4명|arXiv (Cornell University)|2024. 04. 23.

Advanced Neural Network Applications인용 수 4

한 줄 요약

CORE-BEHRT는 전자건강기록(EHR)을 위한 체계적으로 최적화되고 철저히 평가된 BERT 기반 모델을 소개하며, 약물 정보와 타임스탬프 포함 등의 향상된 데이터 표현 및 개선된 훈련 프로토콜을 통해 성능을 향상시킨다. 25개의 임상 과제에서 평균 AUROC 0.801을 달성했으며, 17개 과제에서 유의미한 향상과 24개 과제에서 전반적인 성능 향상을 보이며 임상 적용에 대한 강력한 일반화 능력과 신뢰성을 입증한다.

ABSTRACT

The widespread adoption of Electronic Health Records (EHR) has significantly increased the amount of available healthcare data. This has allowed models inspired by Natural Language Processing (NLP) and Computer Vision, which scale exceptionally well, to be used in EHR research. Particularly, BERT-based models have surged in popularity following the release of BEHRT and Med-BERT. Subsequent models have largely built on these foundations despite the fundamental design choices of these pioneering models remaining underexplored. Through incremental optimization, we study BERT-based EHR modeling and isolate the sources of improvement for key design choices, giving us insights into the effect of data representation, individual technical components, and training procedure. Evaluating this across a set of generic tasks (death, pain treatment, and general infection), we showed that improving data representation can increase the average downstream performance from 0.785 to 0.797 AUROC ($p<10^{-7}$), primarily when including medication and timestamps. Improving the architecture and training protocol on top of this increased average downstream performance to 0.801 AUROC ($p<10^{-7}$). We then demonstrated the consistency of our optimization through a rigorous evaluation across 25 diverse clinical prediction tasks. We observed significant performance increases in 17 out of 25 tasks and improvements in 24 tasks, highlighting the generalizability of our results. Our findings provide a strong foundation for future work and aim to increase the trustworthiness of BERT-based EHR models.

연구 동기 및 목표

• BERT 기반 EHR 모델에서 철저한 평가와 최적화의 부족이 임상 적용을 저해하는 문제를 해결하기 위해.
• 주요 설계 선택 사항—데이터 표현, 아키텍처, 훈련 프로토콜—이 최종 성능에 미치는 영향을 고립하고 정량화하기 위해.
• 단계적 실험과 다중 과제 검증을 통해 BEHRT를 체계적으로 최적화함으로써 모델의 신뢰성을 향상시키기 위해.
• 다양한 임상 예측 과제에서 일관된 성능 향상을 보이며, 좁은 벤치마크를 넘어서 일반화 가능성을 입증하기 위해.
• 향후 EHR 모델링 연구를 위한 재현 가능하고 최적화된 기반과 실질적인 통찰을 제공하기 위해.

제안 방법

• 데이터 표현(예: 약물, 타임스탬프, ICD-10 코드 포함)과 훈련 설정을 변화시켜 BEHRT에 대한 점진적 최적화를 수행하였다.
• 성능 지표의 안정성과 수렴 불안정성 감소를 위해 다중 사전 훈련 실행을 포함한 5중 교차검증을 사용하였다.
• 최적화를 이끄는 데 사용된 3개의 초도 과제(사망, 통증 치료, 감염)에서의 성능 평가 후, 25개의 다양한 임상 예측 과제에서 검증을 수행하였다.
• 마스크 전략과 풀링 방법의 영향을 평가하기 위해, EHR 과제 목표와의 일치를 중심으로 후행 AUROC에 영향을 분석하였다.
• 통계적 유의성 검정(p < 10⁻⁷)을 포함한 철저한 평가 파이프라인을 적용하여 성능 향상의 타당성을 검증하였다.
• 실제 임상 데이터를 기반으로 임상적 관련성을 확보하였으며, 연령 창과 모odal 조건을 적용하여 실제 운영 환경을 시뮬레이션하였다.

Figure 1: This figure depicts the selection process of the best configuration. Non-selected configurations are depicted in a lighter shade, indicating they were considered but not chosen based on their impact on the average Area Under the Receiver Operating Characteristic (AUROC), across three predi

실험 결과

연구 질문

• RQ1 약물과 타임스탬프 포함 등 데이터 표현을 향상시키면 EHR 기반 BERT 모델의 후행 성능에 어떤 영향을 미치는가?
• RQ2 데이터 표현을 초월해 아키텍처 및 훈련 프로토콜 개선이 모델 성능에 얼마나 기여하는가?
• RQ3 최적화 성과는 광범위하고 다양한 임상 예측 과제에 대해 얼마나 일반화 가능한가?
• RQ4 모델 크기와 피팅 데이터 양이 성능 포화 및 확장성에 미치는 영향은 무엇인가?
• RQ5 다양한 마스크 비율과 풀링 전략이 과제 간 일관성 없는 향상을 초래하는 이유는 무엇이며, 이를 어떻게 완화할 수 있는가?

주요 결과

• 데이터 표현 향상으로 인해 3개의 核 心 과제에서 평균 AUROC가 0.785에서 0.797로 상승하였으며(p < 10⁻⁷), 주로 약물 및 타임스탬프 특징 포함로 인한 것이다.
• 아키텍처 및 훈련 프로토콜의 추가 최적화로 25개의 다양한 임상 예측 과제에서 평균 AUROC가 0.801로 상승하였으며, 17개 과제에서 유의미한 향상과 24개 과제에서 전반적인 성능 향상을 보였다.
• 스토브 예측에서 최신 기술 수준 성능을 달성하여 AUROC 70.7%를 기록했으며, 이는 더 제한된 데이터 조건에서도 이전 연구(66.9%)를 초월하였다.
• 성능 향상은 일관되고 통계적으로 유의미하여, 최적화 과정이 특정 과제에 편향되지 않았음을 시사한다.
• 약물 및 타임스탬프의 포함이 성능에 지배적인 영향을 미치며, 임상적 맥락을 포착하는 데 핵심적인 역할을 함을 시사한다.
• 데이터 크기가 증가함에도 일부 과제에서는 성능 포화가 조기에 발생하여, 데이터 양만으로는 충분하지 않으며, 모델 설계와 표현 방식이 더 큰 영향을 미친다는 것을 시사한다.

Figure 2: Overview of the iterative optimization procedure, highlighting adapted settings in bold. We augmented the input data with medication codes, patient sex, event timestamps, full International Classification of Diseases (ICD-10), and Anatomical Therapeutic Chemical (ATC) codes. Improvements t