Adversarial Masking for Pretraining ECG Data Improves Downstream Model Generalizability
Abstract
Medical datasets often face the problem of data scarcity, as ground truth labels must be generated by medical professionals.
One mitigation strategy is to pretrain deep learning models on large, unlabelled datasets with self-supervised learning (SSL), but this introduces the issue of domain shift if the pretraining and task dataset distributions differ.
Data augmentations are essential for improving the generalizability of SSL-pretrained models, but they tend to be either handcrafted or randomly applied.
We use an adversarial model to generate masks as augmentations for 12-lead electrocardiogram (ECG) data, where masks learn to occlude diagnostically-relevant regions.
Compared to random augmentations, models pretrained with adversarial masking reaches better accuracy under a domain shift condition and in data-scarce regimes on two diverse downstream tasks, arrhythmia classification and patient age estimation.
Adversarial masking is competitive with and even reaches further improvements when combined with state-of-art ECG augmentation methods, 3KG and random lead masking (RLM) , demonstrating the generalizability of our method 1 .
Introduction
Across medical applications, deep learning is increasingly used to automate disease diagnosis [Miotto et al., 2018].
In some cases, neural networks have even reached or exceeded the performance of expert physicians [Hannun et al., 2019].
One such application is with 12-lead electrocardiogram (ECG) data, which is commonly collected to screen for various cardiovascular disorders [Fesmire et al., 1998].
There has been a recent surge in ECG-based deep learning research, largely enabled by challenges like the annual PhysioNet/Computing in Cardiology Challenge [Perez Alday et al.,
Reyna et al., 2021].
While these research outcomes show substantial progress within the field, their performance only reflects training on large-scale, labelled dataset.
In contrast, real world medical datasets are likely much smaller due to the extensive resources required to collect medical labels.
Data scarcity is a well-documented issues that effect deep learning training.
Models trained on small datasets lack generalizability to unseen data and cannot be deployed reliably [Kelly et al., 2019].
To mitigate issues associated with small training datasets, large yet unlabelled dataset can be leveraged to pretrain deep learning models with robust representations, which is commonly done in the ECG domain [Sarkar and Etemad, 2020, Weimann and Conrad, 2021, Liu et al., 2021, Kiyasseh et al., 2021,
Diamant et al., 2022, Mehari and Strodthoff, 2022, Oh et al., 2022].
However, if the two datasets are collected under different environments, with different sensors, or across different populations, the transferrability of the model to the downstream task can be greatly impacted [Koh et al., 2021].
Contrastive self-supervised learning (SSL) is a pretraining technique that does not require a labelled dataset and can induce robust representations in the model
[Chen et al., 2020a].
The encoding model is trained to maximize the similarity between latent representations of augmented pairs of data while separating from the representations of other data samples.
SLL-pretrained models learn the underlying structure of the data invariant to the augmentations, which means that selecting good augmentations is crucial for learning useful representations.
State-of-art within the ECG domain is patient-level SSL (PL-SSL), where pairs of data are taken from the same patient at different points in time [Kiyasseh et al., 2021, Diamant et al., 2022].
The performance of PL-SSL techniques are powerful, but can be improved in combination with other augmentations.
For time-series data, augmentations do not reach uniform and predictable gains across different datasets and tasks, making the augmentation selection process difficult [Iwana and Uchida, 2021].
Augmentations are typically selected from a standard pool like noise injection, baseline shift, timedomain masking, and more recent advances using generative adversarial models (GANs) to create synthetic data [Wen et al., 2021].
In the ECG data domain, a recent work 3KG develops physiologically-consistent spatial augmentations for ECGs that reach good performance on small datasets [Gopal et al., 2021].
Another introduces random lead masking (RLM) , where leads of the
ECG are fully masked out at random [Oh et al., 2022].
However, such augmentations are either handcrafted with manually tuned parameters or randomly applied, while we investigate the use of an adversarial method to optimize augmentation parameters.
In this work, we elect to focus on time-domain masking, a type of augmentation with high potential but is underexploited for periodic time-series data like ECGs:
• We adapt the image-based adversarial masking method of Shi et al.
[2022] to generate masks as augmentations for each given ECG sample during PL-SSL pretraining.
To our knowledge, this is the first work to implement adversarial masking for time-series data.
• We show that incorporating adversarial masking improves performance in two downstream tasks compared to baseline augmentations.
We also find orthogonal benefits when combined with state-of-art ECG augmentations, 3KG and RLM .
Methods
The overview of the adversarial masking pretraining framework with the encoding model E and adversarial masking model A and the downstream transfer learning step is show in Figure 1.
2.1 Self-Supervised Pretraining
CMSC Objective: We use Contrastive Multi-segment Coding (CMSC) as the PL-SSL strategy from Kiyasseh et al.
Given a batch of 12-lead ECG data (cid:8) x ∈ R D (cid:9) B i =1 , the positive pair is created as x i = x 1 .. D 2 , x ′ i =
x D 2 ..D , where x i , x ′ i represent temporally non-overlapping ECG segments from the same patient.
We further transform a data sample T ( x ′ i ) , where T represents one or more augmentations.
The encoding model E is trained to align the feature representations of the positive pair of data, h i =
E ( x i ) and h ′ i
=
E ( T ( x ′ i )) , while separating them from all other negative samples within the batch where x i ̸
= x j .
We perform training with the CMSC objective (same as the SimCLR objective [Chen et al., 2020b]) described in Equation 1, where τ = 0 .
1 is a temperature scaling term and the sim function represents cosine similarity.
[Equation]
Adversarial Objective: We add an adversarial masking model A that generates a set of N masks for a given data sample, m
i ∈ R N × D =
A ( x i ) .
Multiple masks ensure that the masked regions of the ECG are alternated, as only one mask is sampled, binarized, and applied in each training step.
To overcome potential catastrophic forgetting due to overactive masking, we only apply the sampled mask with a probability of p=0.8 . We use a differentiable soft binarization function to drive the mask values towards either 0 or 1, shown in Equation 2 with γ = 25 . This enables back-propagation while simulating a shape-preserving binary mask without significant distortions to the ECG signals. [Equation]
The masking model A acts in opposition to the encoding model E by generating difficult augmentations, hence maximizing the CMSC objective.
We also adapt the sparse penalty from Shi et al.
[2022] to limit the amount of masking, described by Equation 3 with α = 0 .
1 as a weighting term.
This prevents A from fully occluding the data or not occluding anything at all.
The full min-max loss objective optimizing E and A is described by Equation 4:
[Equation]
Architectures: For the encoder backbone, we use a 1D [-ResNet-
 [He-]ResNet-18 with a hidden dimension of[He et al., 2016], which is the best performing architecture across multiple ECG tasks out of eight backbones [Nonaka and Seita, 2021].
We also use a two-layer projection head typical in contrastive
SSL to convert the encoder’s outputs to a 128-dimension space, but the projector is not transferred to the downstream task.
For the adversarial model, we adapt an open-source 1D U-Net with four downsampling and upsampling layers.
The number of out channels is equal to the number of masks N . The outputs are passed through either a softmax or sigmoid function depending on N .
Pretraining Details: All models are implemented with PyTorch Lightning.
Training and testing are performed with a single NVIDIA Volta V100 GPU on the MIT Supercloud [Reuther et al., 2018].
Pretraining is performed with 12-lead ECG datasets from the PhysioNet/Computing in Cardiology
Challenge
[Perez Alday et al.], which comprises over 66 thousand patient recordings from six institutions in four countries. [Table]
We use a learning rate of 0.0001 and an Adam optimizer for both the encoding model and adversarial model.
The models are updated one after another with their respective losses within the same batch.
The batch size is 32 with gradient accumulation over 4 batches, which is equivalent to an effective batch size of 128.
To reduce computation efforts, we also use mixed precision training.
An early stopping condition based on minimizing L CMSC is implemented to reduce the training time. 2.2 Downstream Transfer Learning We transfer the encoding models’s learned representations to two downstream classification tasks by training a linear layer on top of the frozen model weights. The tasks, arrhythmia classification and age regression, are used as benchmarks by Diamant et al.
The former correlates to disease characteristics and the latter to patient identity, representing two diverse objectives.
The downstream transfer dataset is the 12-lead Chapman-Shaoxing ECG dataset
[Table]
which has over 10 thousand patients and four major classes of cardiac rhythm labels [Zheng et al., 2020].
The arrhythymia class balance is 36% atrial fbrillation (AFIB), 22% supraventricular tachycardia (GSVT), 21% sinus bradycardia (SB), and 21% sinus rhythm (SR).
Patient ages range from 4 to 98 years.
Transfer learning is performed with a batch size of 256, learning rate of 0.01, and an Adam optimizer.
A single Linear layer projects the output dimension of the encoding model to the number of output classes.
There are no overlaps between the pretraining and downstream datasets, so our transfer learning setup constitutes a domain shift .
We present transfer learning results using a train-validation-test split of 80%, 10% and 10%.
We simulate real world data scarcity conditions by reducing the transfer training dataset to 100%, 10%, and 1% of the original size (8516 samples).
The size of the test set is held constant and the test accuracy for all experiments is reported as the average across 3 random seeds.
Table 1 shows the performance metrics and standard deviation the tasks with models pretrained on all baseline and adversarial augmentations. The baseline augmentations are adapted from literature and details are available in Appendix A. Classification accuracy is reported for arrhythmia classification (higher is better) and the absolute error is reported for age regression (lower is better). For each task at each dataset size reduction, the best performing augmentation method for pretraining are in bold.
Results from a finetuned Scratch baseline is also reported.
Adversarial masking ( Adv Mask ) results are given with N = 2 , which performs the best.
For both downstream tasks, Adv Mask yields superior results to most baseline augmentations across multiple training dataset sizes, including other masking techniques.
It is competitive to but does not consistently outperform 3KG [Gopal et al., 2021] and RLM [Oh et al., 2022], two methods which are specific to ECGs and encourages learning invariances across leads.
However, Adv Mask combined with 3KG and RLM shows significant improvements over the respective augmentations on their own, with either Adv Mask + 3KG or Adv Mask + RLM reaching the best performance in nearly all transfer trials.
This shows that benefits introduced by Adv Mask is orthogonal to other augmentations and consistent data-scarce and domain-shifted conditions.
3.2 Analysis of Augmentations
We visualize the generated masks in Figure 2, with more examples in Appendix B. Lead II of the ECG is displayed in red and the masked signal in green. Since we use a soft binarization method, the mask can take any value between 0 and 1.
The blue overlay represent regions that are heavily driven to dropout.
The shape of the original ECG is well-preserved outside of the masked regions, but may be augmented with a slight shift and scaling factor. In an analysis of the effect of data augmentation on model generalization, Huang et al. [2021] show richer augmentations lead to higher generalizability in downstream tasks, which supports why Adv Mask is effective. With N = 2 , one mask consistently covers the QRS complex and T-wave, while the other mask covers the remaining areas.
We hypothesize this is similar to capturing diagnostically-relevant “semantic" content of the ECG, in the same way that Shi et al.
[2022] demonstrates adversarial image masks cover semantically-coherent regions of the image.
According to an analysis of the salient regions of
ECG data, the QRS complex often carry high levels of disease-diagnostic information [Jones et al., 2020].
This suggests that the encoding model emphasizes learning the structual information of the highly salient regions of the ECGs during pretraining.
While the results may indicate that learnt masks converge to peaks, we show that Adv Mask outperforms the Peakmask baseline, which uses statistical methods to mask out peaks in as a naive replication of Adv Mask ’s behaviour.
Conclusions and Future Work
Adversarial masking for ECG data can be utilized positively in SSL pretraining of deep learning models.
It induces the learning of generalizable representations for downstream transfer learning in highly data-scarce and domain-shifted tasks.
Furthermore, masking as an augmentation scheme is agnostic to the choice of architecture and training scheme, so the method can be extended to other
SSL frameworks beyond SimCLR, and even to other time-series data modalities.
Future work includes evaluating across more datasets to understand where adversarial masking can bring the most benefit.
To strengthen our results, we would also like to extend our experiments to include generative models as an augmentation baseline and tailor the sparsity penalty to leverage specific properties of time-series data.
A Augmentations
The baseline augmentations are sourced from ECG deep learning research.
While the list is not exhaustive, we aim to be representative of typical random ECG augmentation techniques used in existing works.
Note that a limitation of our results is that we only test with isolated augmentations, whereas combining two or more augmentations may intuitively yield better results.
We also do not have generative models (GANs) as baseline, but we aim to include them in the future.
Gaussian Noise: A vector of noise v ( t ) ∼ N (0 , 0 . 05) is sampled and added to each lead of the ECG.
Gaussian noise injection is a very common augmentation used in SSL training for both time-series and image data.
Powerline Noise: Powerline noise n ( t ) =
αcos (2 πtkf p + ϕ ) , with α ∼ U (0 , 0 . 5) , ϕ ∼ N (0 , 2 π ) , and f p =
Hz is added to each lead of the ECG [Mehari and Strodthoff, 2022, Oh et al., 2022].
Short-time Fourier Transform: STFT involves computing the Fourier transform of short segments of the time-series signal to generate a spectrogram.
A random mask with values sampled from a beta distribution B ( α = 5 , β = 2) is applied to the spectrogram.
Finally, the STFT operation is inversed to recover the time domain signal.
Baseline Wander: The ECG signal is perturbed with a very low frequency signal to simulate drifting: n ( t ) =
C (cid:80)
Kk
=1 a cos (2 πtk ∆ f + ϕ ) , with C ∼ N (1 , 0 .
5 2 ) , α ∼ U (0 , 0 . 5) , ϕ ∼ N (0 , 2 π ) , and ∆ f ∼ U (0 . 01 , 0 . 2)
[Mehari and Strodthoff, 2022, Oh et al., 2022].
Baseline Shift: A fraction p = 0 .
2 of the baseline of each lead of the ECG signal is shifted by a factor of α ∼ N ( − 0 .
[Mehari and Strodthoff, 2022, Oh et al., 2022].
Mask:
Two types of random masking are implemented, where Mask refers to any timestep being masked out to 0 with the probability p = 0 .
Blockmask: Blockmask masks out a continuous portion p = 0 .
2 of each lead to 0.
The main difference is that Blockmask would occlude larger structural regions of the ECG, whereas Mask only occludes local details.
Peakmask: We introduce Peakmask as a statistical baseline that aims to replicate the peak-finding behaviour of the adversarially generated masks. Timesteps where the average value across leads exceeds the average value over the entire ECG are masked. We either sample the computed mask directly, or the opposite regions (non-peak areas). This emulates the Adv Mask technique with N=2.
3KG: Developed by Gopal et al.
[2021], 3KG augments the 3D spatial representation of the ECG in vectorcardiogram (VCG) space and reprojects it back into ECG space, mimicking natural variations in cardiac structure and orientation.
We take the best parameters reported in the paper, a random rotation − 45 ◦ ≤ θ ≤ 45 ◦ and a scaling factor 1 ≤ s ≤ 1 . 5 .
RLM: Introduced by Oh et al.
[2022], RLM fully occludes individual ECG leads with probability p=0.5 , which was the parameter used in their paper.
This reduces the dependency on requiring allleads to extract useful information.
B More Adversarial Masks
Figures 3-5 depict more examples of adversarially generated masks with N = 1 , 3 , 12 .
When N < 12 , one mask is randomly sampled and applied to all leads of the ECG.
When N = 12 , each mask is applied separately to each lead.
As mentioned in the paper, N ̸ = 2 trials saw reduced performance so their results are not reported.
Increasing N does not seem to be helpful as masks tend to cover either peaks or flat regions, which can be achieved with N =2.