SuPreME: A Supervised Pre-training Framework for Multimodal ECG Representation Learning

This repository contains the official implementation of SuPreME, a supervised pre-training framework for multimodal ECG representation learning that leverages both ECG signals and clinical text reports to learn robust representations for downstream ECG analysis tasks.

📢 News: This paper has been accepted to EMNLP 2025! 🎉

🏗️ Overview

🚀 Installation

Prerequisites

• Python 3.9+
• CUDA 11.0+ (for GPU acceleration)
• PyTorch 2.3.1+

Environment Setup

1. Clone the repository:

git clone https://github.com/your-username/SuPreME.git
cd SuPreME

2. Create conda environment:

conda env create -f environment.yaml
conda activate supreme

3. Install additional dependencies:

pip install -r requirements.txt

Key Dependencies

• torch==2.3.1 - Deep learning framework
• transformers==4.44.0 - Pre-trained language models
• scikit-learn==1.5.1 - Machine learning utilities
• wandb==0.17.5 - Experiment tracking
• wfdb==4.1.2 - ECG data processing

🎯 Quick Start

1. Pre-training

# Pre-train with ViT backbone
python scripts/pretrain.py --config configs/vit.yaml

# Pre-train with ResNet backbone  
python scripts/pretrain.py --config configs/resnet.yaml

2. Fine-tuning

# Fine-tune on PTB-XL dataset
python scripts/finetune.py \
    --dataset ptbxl_super_class \
    --backbone vit_tiny \
    --pretrain_path checkpoints/best_model.pth \
    --batch_size 16 \
    --epochs 100

3. Zero-shot Evaluation

# Evaluate zero-shot performance
python scripts/evaluate.py

📖 Usage

Configuration

The framework uses YAML configuration files for easy customization:

network:
  ecg_model: vit_tiny          # ECG encoder: resnet18/34/50/101, vit_tiny/small/middle/base
  num_leads: 12                # Number of ECG leads
  patch_size: 125              # Patch size for ViT
  text_model: ncbi/MedCPT-Query-Encoder  # Text encoder
  freeze_layers: 12            # Number of frozen BERT layers
  dropout: 0.1                 # Dropout rate

training:
  dataset: "mimic"             # Dataset name
  batch_size: 256              # Batch size
  epochs: 100                  # Training epochs
  learning_rate: 0.001         # Learning rate

Data Preparation

The metadata files, annotation CSV files, and downstream dataset metadata are available at: https://huggingface.co/datasets/mingscai/SuPreME

1. MIMIC-IV Dataset: Pre-process ECG signals and clinical reports
2. Downstream Datasets: PTB-XL, ICBEB, Chapman datasets
3. Data Format:
   • ECG signals: .npy files with shape (N, leads, length)
   • Text reports: CSV files with clinical descriptions
   • Labels: Multi-label binary classification format

Download Dataset Metadata

# Download metadata and annotation files from Hugging Face
git lfs install
git clone https://huggingface.co/datasets/mingscai/SuPreME

Download Pre-trained Models

Pre-trained model checkpoints are available at: https://huggingface.co/mingscai/SuPreME

# Download pre-trained model checkpoints
git lfs install
git clone https://huggingface.co/mingscai/SuPreME

Training Pipeline

Pre-training

from models.supreme import SuPreME
from utils.trainer import Trainer

# Initialize model
model = SuPreME(device, network_config)

# Start training
trainer = Trainer(model, dataset, train_loader, val_loader, ...)
trainer.train(epochs=100)

Fine-tuning

# Load pre-trained weights from Hugging Face
checkpoint = torch.load('path/to/downloaded/checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])

# Fine-tune on downstream task
# (See scripts/finetune.py for complete example)

Using Pre-trained Models

# Fine-tune using pre-trained SuPreME model
python scripts/finetune.py \
    --dataset ptbxl_super_class \
    --backbone vit_tiny \
    --pretrain_path path/to/downloaded/checkpoint.pth \
    --batch_size 16 \
    --epochs 100

📈 Evaluation

Metrics

• AUROC: Area Under the Receiver Operating Characteristic curve
• F1-Score: Harmonic mean of precision and recall
• Accuracy: Classification accuracy

Performance

Framework	Evaluation Approach	Zero-shot (0%)	Linear Probing
			1%	10%	100%
From Scratch
Random Init (CNN)	L	-	55.09	67.37	77.21
Random Init (Transformer)	L	-	53.53	65.54	75.52
ECG Only
SimCLR	L	-	58.24	66.71	72.82
BYOL	L	-	55.78	70.61	74.92
BarlowTwins	L	-	58.92	70.85	75.39
MoCo-v3	L	-	57.92	72.04	75.59
SimSiam	L	-	59.46	69.32	75.33
TS-TCC	L	-	54.66	69.37	76.95
CLOCS	L	-	56.67	70.91	75.86
ASTCL	L	-	57.53	71.15	75.98
CRT	L	-	56.62	72.03	76.65
ST-MEM	L	-	56.42	63.39	69.60
Multimodal Learning
MERL	Z & L	73.54	63.57	78.35	83.68
SuPreME (Ours)	Z & L	77.20	63.24	72.34	84.48

Performance of SuPreME and eSSLs in AUROC, with 'Z' for zero-shot and 'L' for linear probing. Best results are bolded.

📁 Project Structure

SuPreME/
├── configs/                 # Configuration files
│   ├── resnet.yaml         # ResNet configuration
│   └── vit.yaml            # ViT configuration
├── models/                  # Model implementations
│   ├── supreme.py          # Main SuPreME model
│   ├── resnet1d.py         # 1D ResNet variants
│   ├── vit1d.py            # 1D Vision Transformer
│   └── cfn.py              # Transformer Query Network
├── scripts/                 # Training and evaluation scripts
│   ├── pretrain.py         # Pre-training script
│   ├── finetune.py         # Fine-tuning script
│   └── evaluate.py         # Evaluation script
├── utils/                   # Utility functions
│   ├── dataset.py          # Dataset classes
│   ├── trainer.py          # Training utilities
│   ├── logger.py           # Logging utilities
│   └── initializer.py      # Initialization utilities
├── zeroshot/               # Zero-shot evaluation
│   ├── zeroshot_dataset.py # Zero-shot datasets
│   └── zeroshot_evaluator.py # Zero-shot evaluation
├── requirements.txt        # Python dependencies
├── environment.yaml        # Conda environment
└── README.md              # This file

📄 Citation

If you use this code in your research, please cite our paper:

@article{cai2025supreme,
  title={SuPreME: A Supervised Pre-training Framework for Multimodal ECG Representation Learning},
  author={Cai, Mingsheng and Jiang, Jiuming and Huang, Wenhao and Liu, Che and Arcucci, Rossella},
  journal={arXiv preprint arXiv:2502.19668},
  year={2025}
}

📜 License

This project is licensed under the MIT License - see the LICENSE file for details.

---

Note: This is a research implementation. For production use, please ensure proper validation and testing on your specific datasets and use cases.