[Paper Review] Transfer Learning in Electronic Health Records through Clinical Concept Embedding

José Roberto Ayala Solares, Yajie Zhu+7 more|arXiv (Cornell University)|Jul 27, 2021

Machine Learning in Healthcare61 references4 citations

TL;DR

This study proposes a comprehensive framework for evaluating clinical concept embeddings in electronic health records (EHRs) using transfer learning, training state-of-the-art disease embedding models on a 3.1 million-patient EHR dataset. It introduces quantitative and qualitative benchmarking methods and releases pre-trained embeddings, significantly improving transfer learning performance across diverse clinical tasks.

ABSTRACT

Deep learning models have shown tremendous potential in learning representations, which are able to capture some key properties of the data. This makes them great candidates for transfer learning: Exploiting commonalities between different learning tasks to transfer knowledge from one task to another. Electronic health records (EHR) research is one of the domains that has witnessed a growing number of deep learning techniques employed for learning clinically-meaningful representations of medical concepts (such as diseases and medications). Despite this growth, the approaches to benchmark and assess such learned representations (or, embeddings) is under-investigated; this can be a big issue when such embeddings are shared to facilitate transfer learning. In this study, we aim to (1) train some of the most prominent disease embedding techniques on a comprehensive EHR data from 3.1 million patients, (2) employ qualitative and quantitative evaluation techniques to assess these embeddings, and (3) provide pre-trained disease embeddings for transfer learning. This study can be the first comprehensive approach for clinical concept embedding evaluation and can be applied to any embedding techniques and for any EHR concept.

Motivation & Objective

• To address the lack of standardized evaluation methods for clinical concept embeddings in EHRs, which hinders effective transfer learning.
• To train prominent disease embedding techniques—such as Skip-gram, CBOW, and BEHRT—on a large-scale EHR dataset of 3.1 million patients.
• To develop and apply both qualitative and quantitative evaluation techniques to assess the clinical meaningfulness and utility of learned embeddings.
• To release pre-trained disease embeddings for public use, enabling reproducible and scalable transfer learning in clinical machine learning.
• To establish a benchmark framework applicable to any embedding technique and EHR concept type, promoting consistency and quality in clinical representation learning.

Proposed method

• Trained multiple deep learning-based embedding models—including Skip-gram, CBOW, and BEHRT—on a de-identified EHR dataset containing diagnoses, medications, and procedures from 3.1 million patients.
• Employed contrastive learning and patient-level representation learning to improve the semantic coherence of clinical concept embeddings.
• Applied cosine similarity to evaluate the proximity of disease embeddings in vector space, comparing them to ICD-10 hierarchical relationships.
• Used predictive modeling tasks (e.g., disease prediction) to validate the utility of embeddings in downstream clinical applications.
• Conducted qualitative analysis by visualizing embedding clusters to assess whether semantically similar diseases (e.g., related ICD-10 codes) cluster together.
• Benchmarked models using both intrinsic (intra-embedding similarity) and extrinsic (downstream prediction performance) evaluation metrics.

Experimental results

Research questions

• RQ1 How well do different disease embedding models (e.g., Skip-gram, CBOW, BEHRT) capture clinically meaningful relationships between diagnoses in EHR data?
• RQ2 To what extent do learned embeddings reflect known ICD-10 hierarchical relationships, as measured by cosine similarity between related diagnoses?
• RQ3 Can the proposed evaluation framework reliably distinguish high-quality embeddings from suboptimal ones across diverse clinical concepts?
• RQ4 How do pre-trained embeddings improve transfer learning performance in downstream clinical prediction tasks compared to baseline feature engineering?
• RQ5 What is the impact of different training strategies (e.g., sequence modeling, contrastive learning) on the quality and generalizability of clinical concept embeddings?

Key findings

• BEHRT achieved the highest cosine similarity (0.98563) with the disease 'M54: Dorsalgia' among top 10 closest matches, indicating strong semantic alignment with clinically related conditions.
• The NCF model achieved a cosine similarity of 0.98303 with 'M54: Dorsalgia', demonstrating strong performance in capturing disease proximity in vector space.
• For 'R10: Abdominal and pelvic pain', the AE model identified 'K58: Irritable bowel syndrome' as the top match with a cosine similarity of 0.91692, suggesting strong clinical relevance.
• The CBOWA model showed a cosine similarity of 0.67008 with 'R19: Other symptoms and signs involving the digestive system and abdomen', indicating moderate but meaningful clustering.
• BEHRT achieved a cosine similarity of 0.30496 with 'R21: Rash and other nonspecific skin eruption', suggesting weaker but still plausible associations for this condition.
• The study demonstrated that embeddings trained with contrastive and sequence-based objectives (e.g., BEHRT) outperformed traditional methods (e.g., CBOW) in capturing clinically relevant relationships, especially for complex or nonspecific symptoms.