[Paper Review] Masked Transformer for Electrocardiogram Classification

Zhou Ya, Xiaolin Diao+4 more|arXiv (Cornell University)|Aug 31, 2023

ECG Monitoring and AnalysisMedicine3 citations

TL;DR

This paper proposes MTECG, a masked Transformer model for electrocardiogram (ECG) classification that leverages masked autoencoding pre-training on a large-scale, expert-annotated ECG dataset of 220,251 recordings. By splitting ECG signals into non-overlapping segments and applying fluctuated reconstruction targets with a lightweight decoder, MTECG-T achieves state-of-the-art performance on both private and public ECG datasets, outperforming recent SOTA methods while maintaining a lightweight 5.7M-parameter architecture.

ABSTRACT

Electrocardiogram (ECG) is one of the most important diagnostic tools in clinical applications. With the advent of advanced algorithms, various deep learning models have been adopted for ECG tasks. However, the potential of Transformer for ECG data has not been fully realized, despite their widespread success in computer vision and natural language processing. In this work, we present Masked Transformer for ECG classification (MTECG), a simple yet effective method which significantly outperforms recent state-of-the-art algorithms in ECG classification. Our approach adapts the image-based masked autoencoders to self-supervised representation learning from ECG time series. We utilize a lightweight Transformer for the encoder and a 1-layer Transformer for the decoder. The ECG signal is split into a sequence of non-overlapping segments along the time dimension, and learnable positional embeddings are added to preserve the sequential information. We construct the Fuwai dataset comprising 220,251 ECG recordings with a broad range of diagnoses, annotated by medical experts, to explore the potential of Transformer. A strong pre-training and fine-tuning recipe is proposed from the empirical study. The experiments demonstrate that the proposed method increases the macro F1 scores by 3.4%-27.5% on the Fuwai dataset, 9.9%-32.0% on the PTB-XL dataset, and 9.4%-39.1% on a multicenter dataset, compared to the alternative methods. We hope that this study could direct future research on the application of Transformer to more ECG tasks.

Motivation & Objective

• To address the limited performance of Transformers in ECG classification despite their success in vision and NLP.
• To overcome the challenge of small, labeled ECG datasets by constructing a large-scale, expert-annotated dataset of 220,251 ECG recordings.
• To adapt masked autoencoding (MAE) techniques—proven in vision and NLP—to time-series ECG data for self-supervised pre-training.
• To investigate effective training recipes, including reconstruction targets, learning rate scheduling, and regularization, for lightweight Transformer models in ECG tasks.
• To evaluate whether vanilla Vision Transformer architectures can outperform specialized CNNs when combined with proper pre-training and fine-tuning strategies.

Proposed method

• ECG signals are split into non-overlapping, fixed-length segments to reduce sequence length and preserve structural information.
• Learnable positional embeddings are used to maintain sequential order and distinguish segment positions.
• A masked autoencoding (MAE) pre-training strategy is applied on unlabeled ECG data, where random segments are masked and the model reconstructs them using a lightweight decoder.
• Fluctuated reconstruction targets are introduced to encourage the model to learn meaningful wave shape features rather than trivial patterns.
• During fine-tuning, layer-wise learning rate decay and DropPath regularization are applied to reduce overfitting and improve generalization.
• The final model, MTECG-T, is a lightweight Vision Transformer-based architecture trained with this complete pre-training and fine-tuning recipe.

Experimental results

Research questions

• RQ1 Can masked autoencoding be effectively adapted to ECG time-series data to improve self-supervised representation learning?
• RQ2 How do training components such as reconstruction target design, training schedule length, and regularization techniques impact model performance in ECG classification?
• RQ3 Can a lightweight, vanilla Vision Transformer architecture achieve state-of-the-art performance in ECG classification when combined with proper pre-training and fine-tuning strategies?
• RQ4 Does the proposed masked pre-training method enable effective use of unlabeled ECG data in clinical scenarios where labels from other modalities (e.g., echocardiography) are scarce?
• RQ5 Is the performance gain of the Transformer-based model over CNNs due to architectural superiority or the effectiveness of the pre-training recipe?

Key findings

• MTECG-T achieves state-of-the-art performance on both private and public ECG datasets, outperforming recent SOTA methods in classification accuracy.
• The model with 5.7 million parameters performs stably across a wide range of masking ratios (5% to 75%), demonstrating robustness to pre-training hyperparameters.
• Ablation studies confirm that fluctuated reconstruction targets, longer training schedules, layer-wise learning rate decay, and optimal DropPath rates are critical for performance.
• Scaling the model beyond a certain size leads to performance degradation, suggesting that the dataset size limits the benefits of larger architectures.
• The results challenge the assumption that Transformers are unsuitable for lightweight ECG classification, showing that vanilla ViT with proper training recipes can surpass specialized CNNs.
• The self-supervised nature of MTECG enables effective utilization of unlabeled ECG data, making it suitable for novel clinical scenarios with limited labeled data.