[Paper Review] Guiding Masked Representation Learning to Capture Spatio-Temporal Relationship of Electrocardiogram

Yeongyeon Na, Minje Park+2 more|arXiv (Cornell University)|Feb 2, 2024

ECG Monitoring and AnalysisMedicine3 citations

TL;DR

This paper proposes ST-MEM, a self-supervised masked autoencoder framework that learns general electrocardiogram (ECG) representations by reconstructing spatio-temporal patches of 12-lead ECG data. By jointly modeling temporal and spatial relationships through patchifying and masked reconstruction, ST-MEM achieves state-of-the-art performance in arrhythmia classification, outperforming contrastive and generative SSL baselines, and generalizes effectively to reduced-lead and low-resource settings.

ABSTRACT

Electrocardiograms (ECG) are widely employed as a diagnostic tool for monitoring electrical signals originating from a heart. Recent machine learning research efforts have focused on the application of screening various diseases using ECG signals. However, adapting to the application of screening disease is challenging in that labeled ECG data are limited. Achieving general representation through self-supervised learning (SSL) is a well-known approach to overcome the scarcity of labeled data; however, a naive application of SSL to ECG data, without considering the spatial-temporal relationships inherent in ECG signals, may yield suboptimal results. In this paper, we introduce ST-MEM (Spatio-Temporal Masked Electrocardiogram Modeling), designed to learn spatio-temporal features by reconstructing masked 12-lead ECG data. ST-MEM outperforms other SSL baseline methods in various experimental settings for arrhythmia classification tasks. Moreover, we demonstrate that ST-MEM is adaptable to various lead combinations. Through quantitative and qualitative analysis, we show a spatio-temporal relationship within ECG data. Our code is available at https://github.com/bakqui/ST-MEM.

Motivation & Objective

• To address the challenge of limited labeled ECG data in training models for diverse cardiac disease detection.
• To improve self-supervised representation learning in ECG by explicitly modeling both temporal and spatial relationships inherent in 12-lead ECG signals.
• To develop a flexible framework that generalizes across different lead configurations, including reduced-lead and single-lead ECGs.
• To validate that the learned representations capture meaningful spatio-temporal patterns through quantitative and qualitative analysis.

Proposed method

• Spatio-temporal patchifying is applied to 12-lead ECG signals, dividing them into temporal and spatial patches for model input.
• A masked autoencoder (MAE) architecture is used, where 60% of patches are randomly masked during pre-training and reconstructed using a shared decoder with learnable lead embeddings.
• The model employs a dual-branch encoding strategy with separate temporal and spatial attention mechanisms to capture dynamic and anatomical dependencies.
• A lead-wise shared decoder and separation embeddings are introduced to preserve spatial relationships across leads during reconstruction.
• Pre-training is conducted for 200 epochs with a fixed learning rate of 0.001, minimizing mean squared error between original and reconstructed signals.
• Fine-tuning is performed using cross-entropy loss for 30 epochs with a fixed learning rate of 0.001 on downstream classification tasks.

Experimental results

Research questions

• RQ1 Can a masked autoencoder framework effectively learn general ECG representations by capturing both temporal and spatial dependencies in 12-lead ECG data?
• RQ2 How does ST-MEM perform compared to contrastive and generative SSL baselines in arrhythmia classification under standard and low-resource settings?
• RQ3 To what extent can ST-MEM generalize to reduced-lead ECG configurations, such as limb leads or single-lead recordings?
• RQ4 What evidence supports that the learned representations encode meaningful spatio-temporal relationships in ECG signals?

Key findings

• ST-MEM achieves an AUROC of 0.897 ± 0.025 on CPSC2018 and 0.815 ± 0.012 on PTB-XL under 1% training data, outperforming all reproduced baselines including CPC and MLAE.
• On the human activity recognition (HAR) benchmark, ST-MEM achieves 98.4% F1 score, demonstrating strong transferability to general multivariate time series tasks.
• In low-resource settings with only 1% of training data, ST-MEM maintains superior performance, achieving 0.815 AUROC on PTB-XL and 0.897 AUROC on CPSC2018.
• Qualitative analysis confirms that the model learns to reconstruct meaningful ECG morphology, indicating effective capture of spatio-temporal patterns.
• ST-MEM generalizes effectively to reduced-lead configurations, maintaining high performance even when fine-tuned on limb leads or single leads.
• The model’s performance is comparable to supervised learning in some settings, indicating strong generalization capacity from self-supervised pre-training.