xjc1022/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.

For zero-shot classification use xjc1022/MELP-repro, which packages the same run with its text tower.

pip install torch transformers timm
pip install git+https://github.com/fuying-wang/fairseq-signals

Tested with transformers 4.47.0, 4.53.2 and 5.15.1.

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.

Weights are stored in bfloat16 and load as float32 by default; pass torch_dtype=torch.bfloat16 to keep them in bfloat16.

Zero-shot results

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

	Rhythm	Form	Sub	Super	CPSC2018	CSN	Average
This run	86.43	70.32	77.65	77.22	84.18	74.58	78.40
Paper (Table 3)	85.4	69.1	81.2	76.2	84.2	77.6	79.0

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; the full model is xjc1022/MELP-repro, and loading it reproduces the table above.

How it was trained

Text tower fuyingw/heart_bert, ECG tower initialised from an ECG-FM wav2vec2-CMSC checkpoint (xjc1022/MELP-ECG-FM-init), 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 open access on PhysioNet under the Open Data Commons Open Database License v1.0.