[Paper Review] Med-BERT: pre-trained contextualized embeddings on large-scale structured electronic health records for disease prediction

Laila Rasmy, Yang Xiang+3 more|arXiv (Cornell University)|May 22, 2020

Machine Learning in Healthcare43 references61 citations

TL;DR

Med-BERT adapts the BERT framework to large-scale structured EHR data to produce contextualized embeddings, boosting disease-prediction performance, especially in small fine-tuning sets.

ABSTRACT

Deep learning (DL) based predictive models from electronic health records (EHR) deliver impressive performance in many clinical tasks. Large training cohorts, however, are often required to achieve high accuracy, hindering the adoption of DL-based models in scenarios with limited training data size. Recently, bidirectional encoder representations from transformers (BERT) and related models have achieved tremendous successes in the natural language processing domain. The pre-training of BERT on a very large training corpus generates contextualized embeddings that can boost the performance of models trained on smaller datasets. We propose Med-BERT, which adapts the BERT framework for pre-training contextualized embedding models on structured diagnosis data from 28,490,650 patients EHR dataset. Fine-tuning experiments are conducted on two disease-prediction tasks: (1) prediction of heart failure in patients with diabetes and (2) prediction of pancreatic cancer from two clinical databases. Med-BERT substantially improves prediction accuracy, boosting the area under receiver operating characteristics curve (AUC) by 2.02-7.12%. In particular, pre-trained Med-BERT substantially improves the performance of tasks with very small fine-tuning training sets (300-500 samples) boosting the AUC by more than 20% or equivalent to the AUC of 10 times larger training set. We believe that Med-BERT will benefit disease-prediction studies with small local training datasets, reduce data collection expenses, and accelerate the pace of artificial intelligence aided healthcare.

Motivation & Objective

• Motivate the use of pre-trained contextualized embeddings for structured EHR data to improve disease prediction when labeled data is limited.
• Leverage a large-scale EHR corpus to pre-train representations that can transfer to downstream clinical prediction tasks.
• Demonstrate gains in predictive accuracy for specific diseases across data-sparse settings.

Proposed method

• Adapt the BERT framework to structured diagnosis data from a very large EHR cohort (28,490,650 patients).
• Pre-train contextualized embeddings on the structured EHR data and fine-tune on downstream tasks.
• Evaluate on two disease-prediction tasks: heart failure in patients with diabetes and pancreatic cancer from two clinical databases.
• Compare performance to non-pre-trained baselines to quantify gains from pre-training.
• Report improvements in AUC across varying fine-tuning data sizes, highlighting small-data benefits.

Experimental results

Research questions

• RQ1 Can Med-BERT pre-trained on structured EHR data improve disease-prediction performance compared with non-pretrained models?
• RQ2 How does pre-training impact performance when fine-tuning data are scarce (e.g., 300–500 samples)?
• RQ3 Do gains generalize across different diseases and clinical databases?

Key findings

• Med-BERT substantially improves prediction accuracy, boosting AUC by 2.02–7.12% across tasks.
• Pre-trained Med-BERT substantially improves performance on very small fine-tuning sets (300–500 samples), boosting AUC by more than 20%.
• The small-data gains are comparable to what would be achieved with 10x larger training sets.
• The approach supports disease-prediction studies with small local datasets and can reduce data collection costs.