Data distribution impacts the performance and generalisability of contrastive learning-based foundation models of electrocardiograms

Abstract

Contrastive learning is a widely adopted self-supervised pretraining strategy, yet its dependence on cohort composition remains underexplored. We present Contrasting by Augmented Patient Electrocardiograms (CAPE) foundation model and pretrain on four cohorts (n = 5,203,269), from diverse populations across three continents (North America, South America, Asia). We systematically assess how cohort demographics, health status, and population diversity influence the downstream performance for prediction tasks also including two additional cohorts from another continent (Europe). We find that downstream performance depends on the distributional properties of the pretraining cohort, including demographics and health status. Moreover, while pretraining with a multi-centre, demographically diverse cohort improves in-distribution accuracy, it reduces out-of-distribution (OOD) generalisation of our contrastive approach by encoding cohort-specific artifacts. To address this, we propose the In-Distribution Batch (IDB) strategy, which preserves intra-cohort consistency during pretraining, discourages learning of spurious cohort-specific features, and instead promotes clinically meaningful variability within cohorts. This leads to improved out-of-distribution robustness, with gains of 9–40% in downstream label prediction performance. This work provides insights into pretraining strategies for more clinically deployable and generalisable foundation models.