MedRep

The official source code for MedRep (Published in JAMIA).
Paper doi: https://doi.org/10.1093/jamia/ocag032
Preprint doi: https://doi.org/10.48550/arXiv.2504.08329

Release Notes

• April 24, 2025: Full version of OMOP concept representations released.

• August 15, 2025: Graph-free version of MedRep added.

• August 26, 2025: DeBERTa weights pretrained with OMOP concept descriptions released.

  
  

Downloads

• MedRep (full version, 22.1GB): Download

• MedRep (graph-free version, 22.1GB): Download

• Concept index (65.0MB): Download

• DeBERTa weights (1.6GB): Download

  
  

Abstract

Electronic health record (EHR) foundation models have been an area ripe for exploration with their improved performance in various medical tasks. Despite the rapid advances, there exists a fundamental limitation: Processing unseen medical codes out of vocabulary. This problem limits the generalizability of EHR foundation models and the integration of models trained with different vocabularies. To alleviate this problem, we propose a set of novel medical concept representations (MedRep) for EHR foundation models based on the observational medical outcome partnership (OMOP) common data model (CDM). For concept representation learning, we enrich the information of each concept with a minimal definition through large language model (LLM) prompts and complement the text-based representations through the graph ontology of OMOP vocabulary. Our approach outperforms the vanilla EHR foundation model and the model with a previously introduced medical code tokenizer in diverse prediction tasks. We also demonstrate the generalizability of MedRep through external validation.

Illustration of Concept Representation Learning

Data Preparation

File Tree

Place the following files in the usedata/representation folder:

• concept_idx.csv
• concept_representation_description.npy
• concept_representation_medrep.npy

.
├── codes
├── data
│   ├── concepts
│   │   ├── CONCEPT.csv
│   │   └── CONCEPT_RELATIONSHIP.csv
│   ├── ehrshot
│   │   ├── condition_occurrence.csv
│   │   ├── death.csv
│   │   ├── drug_exposure.csv
│   │   ├── measurement.csv
│   │   ├── person.csv
│   │   ├── procedure_occurrence.csv
│   │   └── visit_occurrence.csv
│   └── mimic
│       ├── condition_occurrence.csv
│       ├── death.csv
│       ├── drug_exposure.csv
│       ├── measurement.csv
│       ├── patients.csv
│       ├── person.csv
│       ├── procedure_occurrence.csv
│       └── visit_occurrence.csv
├── usedata
│       ├── mimic
│       ├── ehrshot
│       └── representation
│           ├── concept_idx.csv
│           ├── concept_representation_description.npy
│           └── concept_representation_medrep.npy
├── results
..

Requirements

- Python 3.9.19
- torch 2.6.0
- transformers 4.49.0
- torch-geometric 2.6.1
- torch_scatter 2.1.2
- torch_sparse 0.6.18

How to Run

Data Preprocessing
The datasets for experiments will be generated in the usedata folder.

cd ./codes
bash preprocessing/run.sh

Learning Representations
The representations for experiments will be generated in usedata/representation folder.

cd ./codes
bash representation/run.sh

Model Pretraining
50 epochs on MIMIC-IV using 8 NVIDIA RTX A6000 GPUs (~15 hours)

cd ./codes
python experiments/pretrain.py --model behrt --rep-type medrep --gpu-devices 0 1 2 3 4 5 6 7

Model Finetuning
up to 50 epochs on MIMIC-IV using 1 NVIDIA RTX A6000 GPU (~3–6 hours)

cd ./codes
python experiments/finetune.py -d 0 --model behrt --rep-type medrep --outcome MT

Extract important features
Important features can be extracted after finetuning as follows:

cd ./codes
python experiments/finetune.py -d 0 --model behrt --rep-type medrep --outcome MT --extract-attention-score

Run the MedGemma-27B implementation as follows:

cd ./codes
CUDA_VISIBLE_DEVICES=0 python experiments/medgemma.py --rep medrep --outcome MT --seed 100 --ex mimic