xjc1022/MELP-repro

MELP (reproduction): ECG and text towers for zero-shot classification

A full MELP model — ECG tower, text tower and both projections into the shared 256-d space — so ECG records can be scored against free-text class descriptions with no fine-tuning. It is the same training run as xjc1022/MELP-Encoder-repro, which publishes only the ECG tower; this repo adds the text side that zero-shot needs.

These weights come from a reproduction run of the public code (HKU-MedAI/MELP, with the fixes in #4), not from the checkpoint behind the paper's tables. Paper: From Token to Rhythm: A Multi-Scale Approach for ECG-Language Pretraining (ICML 2025), Fuying Wang, Jiacheng Xu and Lequan Yu.

Zero-shot results

Mean per-class AUROC on the test splits, using the protocol of scripts/zeroshot/test_zeroshot.py: one CKEPE description per class, cosine similarity as the score. The first row was produced by loading this repo with AutoModel and nothing else from the training code (on CPU; the original checkpoint on GPU gives the same numbers to within 0.01).

	PTB-XL Rhythm	PTB-XL Form	PTB-XL Sub	PTB-XL Super	CPSC2018	CSN	Average
This model	86.43	70.32	77.65	77.21	84.19	74.59	78.40
Paper (Table 3)	85.4	69.1	81.2	76.2	84.2	77.6	79.0

The average lands 0.6 below the paper. The gap is concentrated in PTB-XL sub-class (−3.6) and Chapman-Shaoxing-Ningbo (−3.0); rhythm, form and super-class come out slightly above.

Usage

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; outputs agree across the three to within 1.2e-6.

import json
import torch
from huggingface_hub import hf_hub_download
from transformers import AutoModel, AutoTokenizer

repo = "xjc1022/MELP-repro"
model = AutoModel.from_pretrained(repo, trust_remote_code=True).eval()
tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)

# Class descriptions used in the paper (131 labels); any free text works too.
prompts = json.load(open(hf_hub_download(repo, "prompts/CKEPE_prompt.json")))
classes = ["SR", "AFIB", "STACH", "SBRAD"]
weights = model.zeroshot_weights(tokenizer, [prompts[c] for c in classes])  # (256, 4)

ecg = torch.rand(2, 12, 5000)  # replace with real records, preprocessed as below
with torch.no_grad():
    scores = model.encode_ecg(ecg) @ weights  # (2, 4) cosine similarities

zeroshot_weights also accepts a list of templates per class and averages them. prompts/dataset_class_names.json lists which labels each benchmark uses.

The scores are cosine similarities, not probabilities. The benchmark numbers rank records within each class (AUROC); they do not threshold or argmax the scores.

Input format

• 12 leads, 10 s at 500 Hz, shape (batch, 12, 5000).
• Lead order I, II, III, aVR, aVF, aVL, V1–V6 — the MIMIC-IV-ECG order. PTB-XL, CPSC2018 and Chapman store aVL before aVF; swap leads 4 and 5 for those.
• Min-max scale each record to [0, 1] with a single min and max over the whole 12×5000 array, not per lead.
• For CPSC2018 the evaluation keeps the first 2500 samples and zero-pads to 5000 before scaling.

Text

Text has to go through model.tokenize, which reproduces the pretraining tokenisation: [BOS] prepended, padded to 128, the [CLS] moved to after the padding, and the embedding read from that last position. Plain tokenizer(...) output does not give the same embedding. zeroshot_weights calls it for you and lower-cases the prompts, as the evaluation did.

ids, mask = model.tokenize(tokenizer, ["atrial fibrillation", "sinus rhythm"])
with torch.no_grad():
    out = model(ecg=ecg, input_ids=ids, attention_mask=mask)

Outputs

model(ecg=..., input_ids=..., attention_mask=...) returns:

key	shape
ecg_emb	(B, 256)	rhythm-level ECG embedding, L2-normalised (the zero-shot embedding)
ecg_beat_emb	(B, 12, 256)	beat-level embeddings
ecg_token_emb	(B, 128, 768)	token-level embeddings
text_emb	(T, 256)	text embedding, L2-normalised
logits_per_ecg	(B, T)	logit_scale.exp() * ecg_emb @ text_emb.T

What is in this repo

part	parameters	source in the training checkpoint
ECG tower (wav2vec2-CMSC, 8 layers, attentional pooling)	65.6M	ecg_encoder.*
Text tower (BERT-base, causal mask)	108.8M	lm_model.bert.*
Text projection	0.3M	proj_t.*

Not included: the captioning decoder and the sentence projection, which only serve the pretraining losses. Report generation is not possible with this repo.

The export was checked against the Lightning checkpoint on CPU in float32: ECG embeddings are bit-identical, and the 131 class embeddings differ by at most 1.2e-6. Weights are stored in float32 (the encoder-only repo stores bfloat16).

Training

Pretrained on MIMIC-IV-ECG ECG-report pairs (745,447 train / 15,171 validation), reading the raw wfdb records. Text tower initialised from fuyingw/heart_bert with the first 6 layers frozen; ECG tower initialised from an ECG-FM wav2vec2-CMSC checkpoint. 4× RTX 3090, batch 64 per GPU, lr 1e-4 with a 2-epoch warmup, weight decay 0.2, 12 contrastive queries, loss weights 1.0 (contrastive) / 2.0 (caption) / 0.2 (beat alignment). Validation every quarter epoch; the released checkpoint is the best one by mean validation zero-shot AUROC, 0.776 at epoch 4, step 13104.

Without the warmup, both towers collapse within about 25 steps and every zero-shot AUROC is 0.5. Training on the denoised .npy waveforms instead of the raw records cost 5.6 points on the average above (70.75 vs 78.40), because the downstream datasets are read raw.

Data and license

MIMIC-IV-ECG is open access on PhysioNet under the Open Data Commons Open Database License v1.0. Evaluation used PTB-XL, CPSC2018 (ICBEB) and Chapman-Shaoxing-Ningbo. The MELP code is MIT-licensed.

Citation

@article{wang2025token,
  title={From Token to Rhythm: A Multi-Scale Approach for ECG-Language Pretraining},
  author={Wang, Fuying and Xu, Jiacheng and Yu, Lequan},
  journal={arXiv preprint arXiv:2506.21803},
  year={2025}
}