[Paper Review] Quantification of BERT Diagnosis Generalizability Across Medical Specialties Using Semantic Dataset Distance

Mihir Khambete, William Su+2 more|arXiv (Cornell University)|Aug 14, 2020

Artificial Intelligence in Healthcare and Education7 citations

TL;DR

This study evaluates BERT-based NLP model generalizability across medical specialties using semantic dataset distance (MCD) in MIMIC-III EHR data. It finds that models trained on multiple specialties generalize better, with performance strongly correlated to semantic similarity between training and test sets, and shows that model AUC drops from 0.92 (internal) to 0.83 (external) when tested across specialties.

ABSTRACT

Deep learning models in healthcare may fail to generalize on data from unseen corpora. Additionally, no quantitative metric exists to tell how existing models will perform on new data. Previous studies demonstrated that NLP models of medical notes generalize variably between institutions, but ignored other levels of healthcare organization. We measured SciBERT diagnosis sentiment classifier generalizability between medical specialties using EHR sentences from MIMIC-III. Models trained on one specialty performed better on internal test sets than mixed or external test sets (mean AUCs 0.92, 0.87, and 0.83, respectively; p = 0.016). When models are trained on more specialties, they have better test performances (p < 1e-4). Model performance on new corpora is directly correlated to the similarity between train and test sentence content (p < 1e-4). Future studies should assess additional axes of generalization to ensure deep learning models fulfil their intended purpose across institutions, specialties, and practices.

Motivation & Objective

• To assess how well BERT-based diagnosis sentiment classifiers generalize across different medical specialties in EHR notes.
• To investigate whether increasing training data diversity across specialties improves model generalization.
• To quantify the relationship between semantic dataset distance (MCD) and model performance on unseen test sets.
• To evaluate the limitations of current NLP models in handling variation in note styles across specialties and institutions.
• To provide a quantitative metric (MCD) for predicting model performance on new, unseen medical corpora.

Proposed method

• Used MIMIC-III EHR dataset containing over 2 million notes from Beth Israel Deaconess Medical Center’s ICU units.
• Extracted 2,955 sentences containing diagnosis terms from three specialties: oncology, cardiology, and pulmonology using ICD-10 and TCGA disease lists.
• Preprocessed sentences via SpaCy for tokenization and sentence segmentation, then embedded using SciBERT for semantic representation.
• Calculated Mean Cosine Distance (MCD) between training and test set sentence embeddings as a proxy for dataset similarity.
• Fine-tuned a BERT-based sentiment classifier on training sets with one, two, or all three specialties and evaluated AUC on internal, partial, and external test sets.
• Used t-SNE visualization to confirm that note type and semantic similarity were primary signals in embedding space.

Experimental results

Research questions

• RQ1 How does model performance vary when tested on data from the same specialty (internal) versus different specialties (external) in MIMIC-III?
• RQ2 To what extent does increasing training data diversity across specialties improve model generalization?
• RQ3 How well does semantic dataset distance (MCD) predict model performance on unseen test sets?
• RQ4 What role does note type play in embedding similarity and model generalization across specialties?
• RQ5 Can MCD serve as a reliable proxy for dataset similarity in predicting NLP model generalization in clinical NLP?

Key findings

• Models trained on a single specialty achieved a mean AUC of 0.92 on internal test sets, but dropped to 0.87 on mixed specialties and 0.83 on external specialties (p = 0.016).
• Increasing training set diversity across specialties significantly improved test performance (p < 1e-4), indicating better generalization with broader data.
• Model performance was strongly correlated with semantic dataset distance (MCD) between training and test sets (p < 1e-4), showing MCD is a reliable predictor of generalization.
• The oncology specialty had a disproportionate number of discharge summary notes (68%), which may have inflated performance on oncology test sets due to note-type bias.
• MCD effectively captured differences in dataset similarity, with internal pairs showing lower MCD than external pairs, though MCD may oversimplify complex similarity distributions.
• The study highlights that note type is a dominant signal in embedding space, potentially confounding semantic similarity and affecting generalization performance.