World Model Augmentation for Imbalanced Multi-Label ECG Classification

Automated multi-label ECG classification struggles with severe class imbalance: rare co-occurring cardiac conditions are systematically underrepresented in clinical datasets. We address this with a world-model approach, pretraining a Joint-Embedding Predictive Architecture (LeJEPA) on over 700K longitudinal ECG pairs from MIMIC-IV-ECG, training it to predict how a patient's latent ECG representation changes between visits given the shift in their ICD label set. After pretraining, we repurpose the frozen dynamics model as a data augmentor -- given a normal embedding and a target condition combination, it synthesises the corresponding abnormal embedding entirely in representation space, without generating a single additional waveform. Training a lightweight MLP probe on the resulting 2.7M-embedding dataset (721K real + 2M synthetic) achieves a macro-averaged AUROC of 0.743 across 76 ICD-coded conditions, recovering 55\% of the gap between a real-data-only linear probe (0.687) and a fully fine-tuned encoder (0.789), with no encoder updates.