MELP-Encoder-repro

MELP ECG encoder (reproduction)

The ECG tower of a MELP model trained from this fork: HKU-MedAI/MELP#4. Same architecture and packaging as fuyingw/MELP_Encoder, so it is a drop-in replacement. Original paper: From Token to Rhythm: A Multi-Scale Approach for ECG-Language Pretraining (ICML 2025), Wang, Xu & Yu — all credit for the method belongs to the authors.

import torch
from transformers import AutoModel

model = AutoModel.from_pretrained("xjc1022/MELP-Encoder-repro", trust_remote_code=True).eval()
ecg = torch.rand(1, 12, 5000)          # 12 leads, 10 s at 500 Hz, scaled to [0, 1]
with torch.no_grad():
    out = model(ecg)
# out["proj_ecg_emb"]    (B, 256)     rhythm level, the zero-shot embedding
# out["ecg_beat_emb"]    (B, 12, 256) beat level
# out["ecg_token_emb"]   (B, 128, 768) token level

Lead order is I, II, III, aVR, aVF, aVL, V1-V6, and each record is min-max scaled to [0, 1] over the whole 12x5000 array — the same preprocessing the training data used.

Zero-shot results

Six-dataset mean AUROC from scripts/zeroshot/test_zeroshot.py (the paper's protocol), test splits:

Those numbers describe the full MELP model this encoder came from: zero-shot needs the text tower as well, which is not part of this repo. What is published here is the ECG tower alone, for use as a feature extractor.

How it was trained

Text tower fuyingw/heart_bert, ECG tower initialised from an ECG-FM wav2vec2-CMSC checkpoint, trained on MIMIC-IV-ECG report pairs. 4x RTX 3090, batch 64/device, lr 1e-4, n_queries_contrast=12, loss weights 1.0 / 2.0 / 0.2. Best checkpoint by validation zero-shot AUROC, epoch 4.

Three things mattered more than any hyperparameter here, and all three are fixes in the linked PR:

1. The shipped LR scheduler has no warmup, and without it both towers collapse within ~25 steps — every zero-shot AUROC comes out at exactly 0.5.
2. Training on denoised waveforms while evaluating on raw ones opens a domain gap worth 5.6 AUROC points (70.75 -> 78.40 once the raw wfdb records are used).
3. The peak is at epoch 4. Runs stopped earlier, or validated only at epoch boundaries, miss it.

Data

Pretrained on MIMIC-IV-ECG, which is PhysioNet credentialed-access data. These are model weights rather than data, but check the PhysioNet data use agreement before redistributing anything derived from them.