[논문 리뷰] CEHR-BERT: Incorporating temporal information from structured EHR data to improve prediction tasks

Chao Pang, Xinzhuo Jiang외 5명|arXiv (Cornell University)|2021. 11. 10.

Machine Learning in Healthcare참고 문헌 13인용 수 31

한 줄 요약

CEHR-BERT는 구조화된 EHR 데이터의 시간 정보를 인공 시간 토큰 및 시간 개념 임베딩과 함께 BERT 프레임워크에 통합하고, 방문 유형 예측 목표를 추가로 도입하여 여러 질환 예측 작업의 성능을 개선한다.

ABSTRACT

Embedding algorithms are increasingly used to represent clinical concepts in healthcare for improving machine learning tasks such as clinical phenotyping and disease prediction. Recent studies have adapted state-of-the-art bidirectional encoder representations from transformers (BERT) architecture to structured electronic health records (EHR) data for the generation of contextualized concept embeddings, yet do not fully incorporate temporal data across multiple clinical domains. Therefore we developed a new BERT adaptation, CEHR-BERT, to incorporate temporal information using a hybrid approach by augmenting the input to BERT using artificial time tokens, incorporating time, age, and concept embeddings, and introducing a new second learning objective for visit type. CEHR-BERT was trained on a subset of Columbia University Irving Medical Center-York Presbyterian Hospital's clinical data, which includes 2.4M patients, spanning over three decades, and tested using 4-fold cross-validation on the following prediction tasks: hospitalization, death, new heart failure (HF) diagnosis, and HF readmission. Our experiments show that CEHR-BERT outperformed existing state-of-the-art clinical BERT adaptations and baseline models across all 4 prediction tasks in both ROC-AUC and PR-AUC. CEHR-BERT also demonstrated strong transfer learning capability, as our model trained on only 5% of data outperformed comparison models trained on the entire data set. Ablation studies to better understand the contribution of each time component showed incremental gains with every element, suggesting that CEHR-BERT's incorporation of artificial time tokens, time and age embeddings with concept embeddings, and the addition of the second learning objective represents a promising approach for future BERT-based clinical embeddings.

연구 동기 및 목표

• 구조화된 EHR 데이터의 시간적 구조를 활용하여 다운스트림 예측을 개선하려는 동기를 제시한다.
• 인공 시간 토큰과 시간 임베딩으로 시간을 인코딩하는 BERT 기반 모델 CEHR-BERT를 개발한다.
• 예측 성능을 높이기 위해 두 번째 사전 학습 목표(Visit Type Prediction)를 도입한다.
• 대규모 CUIMC-NYP OMOP 데이터 세트를 사용하여 여러 임상 예측 작업에서 CEHR-BERT를 평가한다.

제안 방법

• 환자 이력을 방문으로 표현하고 방문 간에 VS/VE 토큰 및 인공 시간 토큰(ATT)을 사용한다.
• 개념 임베딩을 시간 및 연령 임베딩과 연결하여 FC 레이어를 통해 시간적 개념 임베딩을 형성한다.
• Masked Language Modeling (MLM) 및 보조 방문 유형 예측(VTP) 목표로 사전 학습한다.
• 4-fold 평가를 통해 CEHR-BERT를 BEHRT, MedBERT 및 baselines와 비교한다. 4개의 예측 작업에서.
• 시간 토큰, 시간/연령 임베딩, 및 VTP의 기여를 평가하기 위한 제거 연구를 수행한다.

실험 결과

연구 질문

• RQ1 인공 시간 토큰과 시간 임베딩을 통해 시간 정보를 도입하는 것이 구조화된 EHR 데이터에 대해 BERT 기반 표현을 개선할 수 있는가?
• RQ2 방문 유형 예측 목표가 다운스트림 질병 예측 성능을 추가로 향상시키는가?
• RQ3 CEHR-BERT가 기존의 EHR-BERT 적응(BEHRT, MedBERT) 및 전통적인 베이스라인들과 여러 예측 작업에서 어떻게 비교되는가?
• RQ4 적은 수의 라벨 데이터에서 CEHR-BERT가 소수 샷 학습 상황에서 효과적인가?

주요 결과

• CEHR-BERT는 네 가지 예측 작업(t2dm HF, HF 입원, 가정으로의 퇴원, 사망)에서 BEHRT, MedBERT 및 베이스라인을 모두 능가했다.
• t2dm HF에서 CEHR-BERT는 AUC 80.7% 및 PR-AUC 32.3%를 달성했다.
• HF 재입원에서 CEHR-BERT는 AUC 66.3% 및 PR-AUC 38.6%를 달성했다.
• 가정으로의 퇴원 및 사망 결과에서 CEHR-BERT는 AUC 94.6% 및 PR-AUC 52.7%를 달성했다.
• 입원 예측에서 CEHR-BERT는 AUC 75.9% 및 PR-AUC 31.1%를 달성했다.
• 소수 샷 학습에서 학습 데이터의 5%로 훈련된 CEHR-BERT가 전체 데이터로 학습된 경쟁자들을 능가했다(예: t2dm HF의 AUC ~0.78, PR-AUC ~0.29).
• 제거 연구는 시간 토큰, 시간/연령 임베딩, 및 VTP 목표로부터 점진적 이점을 보였으며, 이들의 가법적 가치를 시사했다.