Masked Autoencoder-Based Self-Supervised Learning for Electrocardiograms to Detect Left Ventricular Systolic Dysfunction.
Abstract
The generalization of deep neural network algorithms to a broader population is an important challenge in the medical field.
In this study, we aimed to apply self-supervised learning using masked autoencoders (MAEs) to improve the performance of deep learning models that detect left ventricular systolic dysfunction (LVSD) from 12-lead electrocardiography data.
In our MAE approach, we first mask the vast majority, that is, 75% of the ECG time series.
Second, we pretrain a
Vision Transformer encoder by inferring the masked part.
Our proposed approach enables rich features that generalize well from unlabeled ECG data to be learned.
In fact, the reconstructed ECG maintains the relationships among the major ECG components.
Transfer performance in the detection of LVSD outperforms the baseline CNN model on external validation datasets and shows promising results for generalization that enables us to use the model for a broader population by solely using ECG data collected in a single medical institution.
Introduction
The generalization of deep neural network (DNN) algorithms to a broader population is an important challenge, and this problem can be a major barrier to the social implementation of DNN algorithms in the medical field (1-3).
Low generalization performance can occur when the dataset used to train the algorithm is not diverse enough.
Improving generalizability requires large amounts of accurately labeled data, however, obtaining such data is often difficult due to labeling costs requiring specialized knowledge and ethical considerations in the medical field.
To address this problem, self-supervised learning, which can make more effective use of limited data, is attracting attention in computer-vision tasks.
Self-supervised learning using masked autoencoders (MAEs; licensed under an Attribution
Non-Commercial 4.0 International license) learns very high-capacity models that generalize well (4).
Pretraining with an MAE enables data-hungry models such as Vision Transformers (ViT)-
Large/Huge (5) to be trained on ImageNet-1K to improve generalization performance.
Contrastlearning approaches such as MoCo (6, 7) and SimCLR (8) have also achieved high performance in computer-vision tasks but require image transformation to create a contrast to the original image.
In contrast, MAE learns a reusable representation of the input image by masking random patches and reconstructing missing pixels.
Therefore, the transformations used in contrast learning are unnecessary.
In this study, we aimed to apply self-supervised learning using MAE to the analysis of 12-lead electrocardiograms (ECGs), which is one of the most popular medical time-series data.
The deep learning approach to analyzing ECGs has been widely studied in recent years, enabling highly accurate detection of cardiac diseases that were previously difficult to diagnose from ECGs (9-11).
However, these approaches require large amounts of labeled electrocardiography data, which are difficult to collect at a single institution.
To solve this problem, several methods have been proposed to apply self-supervised pretraining to ECGs (12, 13).
However, no previous studies have applied
MAE-based pretraining to train a high-capability ViT-Large exclusively for ECG analysis.
We focus on detecting left ventricular systolic dysfunction (LVSD) from the ECG.
LVSD is a common disease that may increase the risk of sudden death (14, 15).
Therefore, early detection of LVSD is desirable.
In clinical practice, LVSD is currently diagnosed by echocardiography (16), and diagnosis from
ECGs has been difficult.
However, it has been reported that LVSD can be diagnosed from an ECG by applying deep learning (9), and the implementation of these deep-learning algorithms in clinical practice is expected.
Method 2.1 Datasets
To build and evaluate the models, we collected electrocardiographic data from eight academic medical institutions.
We created two datasets for training.
The first is a small dataset (Dataset1) comprising 37,456 ECGs collected at institution 1, and the second is a large dataset (Dataset2) comprising 126,203 ECGs collected from institutions 1, 2, and 3 (Tables 1 and 2).
LVSD was assessed by echocardiography and defined as an ejection fraction of less than 40%.
Patients who had LVSD were labeled as positive, and the rest were labeled as negative.
Tables 1 and 2 show the details of the datasets.
Dataset1 was divided into Train1 (containing train1 and valid1) and Test1, and Dataset was divided into Train2 (containing train2 and valid2) and Test2.
Note that to avoid data leakage, all
ECGs from one patient were assigned to the same split.
All ECG data from Dataset1 and Dataset were used separately for the self-supervised pretraining, in which the labels were not used.
The Train or Train2 dataset was then used for the downstream task, in which a ViT-Large model was trained to detect LVSD by supervised learning.
In addition, ECGs from the five external institutions that were not used for training were prepared as an external validation dataset, as detailed in Appendix A, to test the performance of the model.
The prevalence of LVSD was also included in Appendix A.
This study was conducted with the approval of the University of Tokyo’s Institutional Review Board and in a manner such that the data could not be used to identify individuals.
2.2 Self-Supervised Pretraining
The first step of our approach was to pretrain an efficient ViT-Large encoder by self-supervised learning using Dataset1 and Dataset2.
In this study, we used the MAE algorithm.
In an MAE for image data (e.g., 224×224), the input is divided into 16 × 16 patches.
These patches are randomly masked and the missing pixels are then reconstructed.
This process creates an encoder that has learned a useful representation (4).
Raw data from each ECG examination record were represented as a 12×5000 matrix of ECG voltage, in which the first dimension was a spatial dimension (each column represented one lead) and the second dimension was the temporal dimension (each row represented a specific time point).
To take advantage of the interrelationship of the 12 ECG leads, the patch size was changed to 1×250.
As a result, the ECG information per patch was 0.5 s. For self-supervised learning, we trained the ViT-Large encoder for 1600 epochs with a batch size of 1024.
Implementation details followed those of the previous study (4).
2.3 Model Architecture
We use a ViT-Large architecture to take advantage of its strong generalization capabilities.
Although
ViT-Huge could have been used, ViT-Large was selected for usability considerations.
The encoder
[Table] [Table]
Over
[Table]
takes a 12×5000 ECG matrix as input, and finally outputs a 1024-dimensional feature vector.
In the self-supervised pretraining, ECGs after normalization were input to the encoder using MAE.
In the downstream task, a normalization layer and a fully connected layer were added to the encoder to detect LVSD.
As a baseline, a CNN architecture was used because a previous study showed that this architecture can achieve high performance in detecting LVSD (9).
The encoder consists of six temporal convolution blocks, one spatial convolution block, and one fully connected layer.
The encoder takes a 12×5000 ECG matrix as input, and finally outputs a 128-dimensional feature vector.
In the downstream task, two fully connected layers and a sigmoid layer were added to the encoder to detect LVSD.
The binary cross-entropy loss was minimized.
To obtain a qualitative sense of our reconstruction task, see Figure 1 and Appendix C.
[Figure]
2.4 Evaluation experiments
We conducted two experiments: one with our proposed approach and one with the baseline (without pretraining) using four sets of Nvidia Tesla A100 80 GB graphics processing units (Nvidia
Corporation, Santa Clara, USA).
First, we conducted self-supervised pretraining on the ViT-Large encoders for each pretraining dataset: Dataset1 and Dataset2.
Then, we fine-tuned the model on each downstream dataset (Train1 or Train2) using the pretrained ViT-Large encoders.
Second, the baseline was trained, omitting the self-supervised pretraining step.
We directly trained the models on the downstream datasets using supervised learning.
The CNN model was trained with randomly initialized model weights, whereas the ViT-Large model was trained with the model weights pretrained by MAE using ImageNet-1K. Table 4 in Appendix A shows the correspondence of the datasets used in each experiment.
In all experiments, the Adam optimizer was used.
Initial learning rates of 1e-4, 3e-5, and 1e-5 were tested in each experiment; the model weights with the lowest validation loss were saved and later used for evaluation on the test set.
Models fine-tuned on Train1 (consisting of data from institution 1) were evaluated on Test1 and the data from the remaining seven institutions, whereas the models fine-tuned on Train2 (consisting of data from institutions 1, 2, and 3) were evaluated on Test and all data from the remaining five institutions.
Appendix B presents an overview of our experiment.
Results and Discussion
We analyze the fine-tuning results of the self-supervised ViT-Large.
For the models fine-tuned on
Train1, the model pretrained using Dataset1 performed better, with area under the receiver operating characteristic curve (AUROC) results of 0.96, 0.94, 0.92, 0.92, 0.95, 0.96, and 0.94 compared to the baseline CNN model, which obtained AUROC results of 0.95, 0.88, 0.89, 0.87, 0.92, 0.94, and 0.
for the Test1 and institution 2–8 cohorts.
Moreover, the model pretrained using Dataset2 obtained the highest AUROC results of 0.97, 0.95, 0.93, 0.92, 0.95, 0.97, and 0.94 (Table 3).
The proposed method performs better than the baseline CNN, indicating that the ViT encoder can learn feature representations that are effective in detecting LVSD through self-supervised pretraining.
The results also show that higher model performance can be obtained by performing self-supervised learning with MAE using a large dataset.
Regarding the models fine-tuned with Train2, the model pretrained using Dataset1 obtained AUROC results of 0.95, 0.92, 0.95, 0.97, and 0.94, which are higher than those of the baseline CNN, which obtained AUROC results of 0.93, 0.90, 0.93, 0.95, and 0.92 for the Test2 and institution
4–8 cohorts.
The model pretrained using Dataset2 performed the best, with
AUROC result of 0.96, 0.94, 0.96, 0.98, and 0.95.
Furthermore, the baseline ViT-Large pretrained by MAE using ImageNet-1K had the lowest AUROC results of 0.91, 0.88, 0.91, 0.94, 0.90 for the
Test 2, institutions 4–8 cohorts.
The baseline ViT-Large performed worse than the baseline CNN model (table 3).
These data show that the performance of the pretrained ViT-Large can be further improved using a larger training dataset.
By contrast, the low accuracy of the ViT-Large pretrained on ImageNet-1K indicates that, although fine-tuning was performed using similarly large training data, the training dataset used in this study was not sufficient to train the ViT-Large.
Therefore, we believe that our proposed method helped train ViT-Large without using additional labeled data and instead making effective use of limited data.
Furthermore, as shown in the results in Table 3,
ViT-Large using the model pretrained with self-supervised learning yields higher AUROC scores with external validation data than do the baseline models, indicating that generalizability was improved.
We highlight a comparison of the performance of ViT-Large pretrained using Dataset1 and fine-tuned on Train1 (bolded text in Table 3) and the baseline CNN model fine-tuned on Train2 (underlined text in Table 3).
ViT-Large, which was pretrained and fine-tuned using a small training dataset from a single institution, obtained higher AUROC results on the external validation dataset and better generalization performance than the baseline CNN model fine-tuned on a large training dataset collected from three institutions.
The total number of labels used for our proposed self-supervised approach is only 31,816 labels in Train1.
In contrast, the baseline CNN model used 107,571 labels in
Train2.
In summary, the proposed method achieves higher AUROC results than the baseline CNN with only 30% labels and improves generalization performance.
[Table]
Train
Test 0.
Test — — — 0.
Institution 0. Institution 0.
Institution 0. Institution 0.
Institution 0. Institution 0.
Institution 0.
Broader Impact
Although our approach allowed us to develop a deep learning model with high generalization performance that efficiently uses few labels and detects LVSD, more detailed studies are needed to determine whether our method can be widely used for other ECG-based tasks.
Based on this study, we aim to publish a powerful ECG pretraining model that is capable of multitasking.
Our goal is to contribute to research for the social implementation of deep learning models to analyze ECGs soon.
To the best of our knowledge, our work has no potential negative impacts.
Appendix A Dataset Details
[Table]
C Qualitatively Understanding of the ECG Features Learned in the Self-Supervised Pretraining
We qualitatively assessed the validity of the ECGs reconstructed by the physician.
As shown in Figure 2, the reconstructions maintained the relationships among the major ECG components such as the P wave, QRS wave, and T wave.
In addition, there was no significant discrepancy in the information between the II, III, and aVF leads (with information about the lower wall of the heart) and the V1, V2, V3, V4, and V5 leads (with information about the anterior wall), which indicate the position information of the ECG.
[Figure]