EHR-Diffusion

Learn useful patient representations from longitudinal EHR data with just a few commands.

• Self supervised learning for EHR data
• Diffusion models for EHR representation

If you have electronic health records (EHRs) containing sequences of diagnoses, procedures, medications, or other clinical codes, this repository helps you learn high-quality patient embeddings that can be used for downstream machine learning tasks.

Instead of manually engineering features, EHR-DiffRepr learns a compact representation for each patient directly from their medical history using self-supervised diffusion-based learning.

The learned embeddings can be used for tasks such as:

• Disease prediction
• Readmission prediction
• Patient phenotyping
• Patient similarity search
• Clustering and cohort discovery
• Any downstream machine learning model that benefits from meaningful patient representations

Simply train the model on your EHR data, then extract one embedding vector per patient and use it with your favorite machine learning pipeline.

Features

• Learns patient representations without requiring labels
• Designed for longitudinal EHR data
• Supports ICD, CPT, NDC, and other categorical medical codes
• End-to-end pipeline from raw EHR data to patient embeddings
• Easy to integrate with downstream ML workflows
• Includes both the original method and practical improvements for real-world datasets

Expected Input

The model expects a long-format table with one row per (patient, visit, code):

patient_id	visit_idx	code
1001	0	E11.9
1001	0	I10
1001	1	99213
1002	0	J45.909

A preprocessing script is included to show the expected format and generate a small synthetic dataset so you can try everything before using your own data.

Quick Start

pip install -r requirements.txt

# Train the model
python -m ehr_diffrepr.training.train \
    --config configs/default.yaml \
    --data your_ehr.csv

# Extract patient embeddings
python -m ehr_diffrepr.training.extract_representations \
    --checkpoint runs/default/best.pt \
    --data your_ehr.csv \
    --out embeddings.parquet

Reference

This code is based on the method described in:

@article{tabejamaat2025cycle,
  title={Cycle Conditioning for Robust Representation Learning from Categorical Data},
  author={Tabejamaat, Mohsen and Etminani, Farzaneh and Ohlsson, Mattias},
  journal={Transactions on Machine Learning Research},
  year={2025}
}