Paper: Guiding Masked Representation Learning to Capture Spatio-Temporal Relationship of Electrocardiogram (ICLR-2024)

Reviewer comment (Novelty):
The problem setting is specialized to the ECG, and the scope might be too narrow for the general audience of ICLR.

Author response:
While our focus is introducing the general ECG representation, we believe that our approach holds the potential for conducting diverse time-series data. Conceptually, the lead can be regarded as a general feature in the time-series domain. To explore this applicability, we demonstrated experiments on human activity recognition [1] (_HAR_) time-series data (as presented in Table 21), comparing the performance of ST-MEM against the supervised model and the masked time autoencoder (MTAE) model. Additionally, we provided an extra section discussing future prospects in Section 8.

ㅤ

__Change__: We included the future work in Section 8. Moreover, we provided the additional experiment on other time-series data, human activity recognition (_HAR_) data, and Appendix J.

Reviewer comment (Experiments):
To claim "general ECG representation", it is important to clarify what actual task was tested. Could you provide more details of classification of myocardial infarction as well as classification of cardiac arrhythmia, e.g., what are the diagnostic criteria of myocardial infarction? What type of arrhythmia are classified?

Author response:
__Reply__: In Appendix A, we included explanations for all labels in each of the downstream datasets. Table 6, 7, and 8 provide details on the target labels assigned to our three downstream datasets. Multiple arrhythmia such as atrial fibrillation (AFIB), first degree atrioventricular block (1AVB), premature ventricular contraction (PVC) and ST-segment depression (STD) are included as labels. Taking the myocardial infarction (MI) label in _PTB-XL_ as an example, as explained in Table 6, MI is a merged label consisting of 14 individual labels interpreted by cardiologists using the corresponding ECG and the patient's medical record.

ㅤ

__Change__: We updated Appendix A.

Reviewer comment (Experiments):
The authors tested Lead I for the single lead task. In a clinical setting, Lead II is the first choice in many situations. Did the author evaluate for Lead II as well?

Author response:
__Reply__: Thank you for the valuable comment. We only evaluated the lead I in that the lead I is easily accessible through the smartwatch. However, as you suggested, we provided additional single-lead experiments on the other five limb leads in Appendix H, Table 19.

ㅤ

__Change__: We provided Appendix H and Table 19.

Reviewer comment (Presentation):
In Figure 2, the authors assigned the same color code to two different segmentations. Could you change one of them to a different color code such as grayscale?

Author response:
__Reply__: There are some similar comments regarding Figure 2 from the other reviewer as well. Therefore, we decided to modify Figure 2. Hopefully, this illustration can provide you better understanding.
 
ㅤ

__Change__: Modified Figure 2.

Reviewer comment (Novelty):
The idea of using MAE to learn representation for ECG signal has been explore previous work (Sawano et al., 2022 [1] in the reference). The author should highlight the difference.

Author response:
Although we discussed the methodology of this work in Section 6.1, we did not tagged the reference and adequately emphasize the differences from our paper. It solely applies the spatio-temporal patchifying mentioned in Figure 2(c), i.e., excluding the lead indicating modules and lead-wise decoder. 

Our comparison with this work is elaborated in Section 5.5. In first row of Table 5, the average AUROC of [1] in _PTB-XL_ and _CPSC2018_ is 0.805 and 0.861, respectively. In contrast, ST-MEM outperforms [1] with average AUROC values of 0.838 and 0.938 in _PTB-XL_ and _CPSC2018_, respectively. We updated Section 6.1 to reflect changes. 

Moreover, in Appendix B, by drawing t-SNE plots as Figure 6, we quantitatively analyzed how well the representation was learned. This analysis was conducted on both regular and irregular rhythms. [1], which utilizes only spatio-temporal patchifying, forms clusters effectively within the same lead for (b) and (d), but struggles to form clusters among spatially similar leads. On the contrary, ST-MEM for (a) and (c) not only forms clusters effectively within the same lead but also, limb leads (I, II, III, aVR, aVL, and aVF) cluster together, and precordial leads (V1, V2, V3, V4, V5, and V6) cluster together.

ㅤ

__Change__: We modified Section 6.1.

ㅤ

ㅤ
---
### __Reference__

[1] Shinnosuke Sawano, Satoshi Kodera, Hirotoshi Takeuchi, Issei Sukeda, Susumu Katsushika, and Is- sei Komuro. Masked autoencoder-based self-supervised learning for electrocardiograms to detect left ventricular systolic dysfunction. In NeurIPS 2022 Workshop on Learning from Time Series for Health, 2022.

Reviewer comment (Experiments):
The experiment session seems missing some previous works. For example, Temesgen et al., 2022 [2] also report metrics on PTB-XL and seems to have higher AUROC score.

Author response:
The paper [2] proposes Contrastive Predictive Coding (CPC) to ECG data. However, the experimental setting for [2] is different from ours which causes the discrepancy in the AUROC score in the _PTB-XL_ dataset. For instance, their problem setting is a multi-label classification for all 71 labels; however, we define the problem as multi-class classification using super-class labels provided from the original _PTB-XL_ paper. Our detailed experimental settings can be found in Appendix A.1. Sorry for the unclearness of the experimental setting.

Since the experimental setting is different from [2], we utilized the pre-trained weights provided in the original paper. Then, we fine-tuned the CPC model on our settings. We could validate that ST-MEM still shows comparable performance and surpasses all the low-resource settings even though the CPC model utilized _PTB-XL_ and _CPSC2018_ to pre-train the model. We provided the additional experiments in Appendix K. 

ㅤ

__Change__: We added the new section as Appendix K.

ㅤ

ㅤ

---
### __Reference__

[2] Mehari, Temesgen, and Nils Strodthoff. "Self-supervised representation learning from 12-lead ECG data." Computers in biology and medicine 141 (2022): 105114.

Reviewer comment (Evaluation):
In the same section, all methods seem to have larger 6-lead vs 1-lead gap on PTB-XL than CPSC2018. Does the author have insight on why this is the case?

Author response:
We think the performance gaps depend on the difficulty of the task. As one can see in Table 3, the overall performance metrics on _CPSC2018_ are higher than _PTB-XL_, which may indicate that the _CPSC2018_ tasks are easier than _PTB-XL_. Specifically, single lead alone may be sufficient to classify particular classes of _CPSC2018_, such as atrial fibrillation (AFIB), premature ventricular complex (PVC), but this may not be the case for certain classes of _PTB-XL_, such as myocardial infarction (MI).

ㅤ

__Change__: None.

Reviewer comment (Theory):
Despite favorable empirical results, the intuition of the spatiotemporal patching based self-supervised learning is not very clear. It is not clear why such patching will help the learning of spatiotemporal representations. Furthermore, the relation with two baselines — MLAE that reconstructs spatial patches and MTAE that reconstructs temporal patches — need to be better clarified.

Author response:
### __[Better explanation of the spatiotemporal patching]__
 
ㅤ

__Reply__: The spatiotemporal patching allows the encoder to consider both temporal and spatial relationships concurrently. On the other hand, temporal patching leads the encoder to consider only the temporal relationship. The spatial patching (i.e., lead patching) discourages the encoder from capturing temporal relationships. For a better understanding of the patching approach, we modified Figure 2.

ㅤ
 
__Change__: We modified Figure 2.

ㅤ

ㅤ


### __[Clarifying the baselines, Masked Time Auto Encoder (MTAE) and Masked Lead Auto Encoder (MLAE)]__

ㅤ
 
__Reply__: Since the MTAE and MLAE take different patching strategies (i.e., temporal or spatial patching), the encoder can only consider either temporal or spatial relationships. In other words, the encoder of MTAE utilizes temporal patching to capture the temporal relationship only, while the encoder of MLAE takes spatial patches to consider the spatial relationship only. Furthermore, MTAE and MLAE cannot explicitly distinguish the lead type information, whereas ST-MEM can encapsulate the lead type information through lead embedding and [SEP] patch embedding. We revised the manuscript to clarify the relation of baselines (MTAE and MLAE) with ST-MEM.
 
ㅤ

__Change__: We clarified in Section 4.2 that MTAE and MLAE do not have lead indicating modules and differ in their approach to patchifying input ECGs with ST-MEM.

Reviewer comment (Theory):
The assumption of the method, in particular to what type of ECG signals and tasks this will apply, needs to be better discussed. For instance, the use of temporal patching seems to rely on the assumption that the rhythm within the input signal length (e.g., 10s) — whether it is normal or abnormal — is regular and periodic. Is that true? How would this work for rhythms that are irregular (e.g., atrial fibrillation), or rhythms where the abnormal rhythm only shows up in transient beats within a longer segment (which is typical for some PVCs and tachycardias). If the method is not designed for these scenarios, it should be clearly clarified.

Author response:
### __[Applicability of the method to regular and irregular ECGs]__

ㅤ
 
__Reply__: We used the ECGs of both regular and irregular rhythms during pre-training. We assume that the masked parts of the input signal can be reconstructed by utilizing the information extracted from the unmasked portions, regardless of the regularity of the rhythms. We did not mention the detailed characteristics of pre-training data, so we clarified in Section 4.1 that we used all ECGs of pre-training data without focusing on specific labels (e.g., normal rhythm).
 
Additionally, for more clarification, we demonstrated the self-attention maps of pre-trained ST-MEM for irregular rhythm ECGs in Appendix D. Notably, ST-MEM consistently assigned high attention weights to temporal patches that have similar shape to the query patch and neighboring spatial patches, regardless of the rhythm type of ECG. From this observation, we conjecture that pre-trained ST-MEM can effectively recognize the morphologies of ECG signals.
 
ㅤ

__Change__: We clarified that we used ECGs of all rhythm types in Section 4.1, and demonstrated the self-attention maps of pre-trained ST-MEM for regular and irregular rhythm ECGs in Appendix D.

Reviewer comment (Experiments):
Related to the above comments, more details about the data and classification tasks need to be given. What are the labels of rhythms being classified? Does each 10-s segment only have one label? The paper did a good job in highlighting that it’s important to develop methodology for the application problem at hand — this should also be reflected when describing data and experimental settings, as different ECG tasks can mean that the features one is looking for are very different (e.g. is it a regular rhythm, is it an irregular rhythm, is it a transient rhythm, etc).

Author response:
### __[Explanation of dataset and classification task]__

ㅤ

__Reply__: In appendix A, we added description for all labels of all 3 downstream datasets. In table 6, 7, and 8, we have outlined the target labels for our three downstream datasets. As apparent from the tables, we employed a variety of ECGs in both training and evaluation, encompassing regular rhythms such as Normal ECG and irregular rhythms like atrial fibrillation (AFIB), first degree atrioventricular block (1AVB) and premature ventricular contraction (PVC). Notably, during pretraining, we used all ECGs without focusing on specific labels. Typically, an ECG may have more than one label, but for our multiclass classification setting, we used only ECGs with a single label, discarding the rest. This processing step is also depicted in Section 4.1 and Appendix A. 

ㅤ

__Change__: We modified Section 4.1 and Appendix A.

Reviewer comment (Evaluation):
It would be interesting to see Fig 4, the embedding of the ECG signals, across different rhythm types as well, in order to appreciate the clustering across rhythm types versus spatial locations of the lead.

Author response:
### __[t-SNE plot for ECGs with different rhythm types]__
 ㅤ

__Reply__: We demonstrated the t-SNE plots only for ECGs with a regular rhythm. We added the t-SNE plots for ECGs with irregular rhythm as well on Appendix B. The appearance of t-SNE plots was similar across different rhythm types, showing that the representations of neighboring leads were gathered closely.
 ㅤ

__Change__: We added t-SNE plots of ECGs with irregular rhythm in Appendix B.

Reviewer comment (Presentation):
Figure 2 (c) is not clear.

Author response:
### __[Modifying Figure 2]__

ㅤ

__Reply__: Sorry for the unclearness. We modified Figure 2 to illustrate the intuition of different patch strategies.

ㅤ
 
__Change__: Modified Figure 2.  

ㅤ

ㅤ

Reviewer comment (Writing):
What is “seasonality”?

Author response:
### __[Clarifying the word “seasonality” in Section 5.6]__

ㅤ
 
__Reply__: Sorry for the confusion in Section 5.6. We misused the word "seasonality". We attempted to define the repeatability of QRS segments as the word seasonality since the query patch contains QRS segments. Note that we tried to avoid the use of clinical terms. We modified the sentence and removed the word "seasonality".

ㅤ
 
__Change__: Revising the sentence in section 5.6.
 
“Moreover, from a temporal perspective, the ECG signal contains seasonality such that the patch with similar seasonality shows high attention scores.” 

->
 
“Moreover, from a temporal perspective, the patches show high attention scores if the signal shape is similar to the query patch in that the ECG signal contains a periodic rhythm.”

ㅤ

ㅤ

Reviewer comment (Experiments):
Some experiments could be stronger. For example, PTB XL dataset have 4 different conditions but it seems that only a couple were used in the tests.

Author response:
### __[Different problem settings of downstream dataset _PTB-XL_]__

ㅤ

__Reply__: We focused solely on the task of distinguishing the 5 super-classes of diagnostic labels for _PTB-XL_. We conducted additional experiments for the other three settings: 23 diagnostic sub-class labels, 19 form labels, and 12 rhythm labels. In this multi-label classification experiment, our proposed method demonstrated superiority over other baselines (supervised learning and MTAE). Further details have been included in the Appendix F.


|     Methods       |     Categories       |                      |                      |
|-------------------|----------------------|----------------------|----------------------|
|                   |     __Sub-class__         |     __Form__             |     __Rhythm__           |
|     __Supervised__    |     0.914 ± 0.002    |     0.829 ± 0.021    |     0.934 ± 0.003    |
|     __MTAE__          |     0.911 ± 0.001    |     0.793 ± 0.014    |     0.920 ± 0.005    |
|     __ST-MEM (Ours)__      |     __0.929 ± 0.001__    |     __0.895 ± 0.008__    |     __0.966 ± 0.004__    |

ㅤ

__Change__: We added new section as Appendix F.

ㅤ

ㅤ

Reviewer comment (Experiments):
Having both the SEP token and the leads embedding seems redundant. Have the authors considered an ablation where only the leads embedding are used but not the SEP? If SEP is needed why is it needed twice and not just at the beginning or just at the end?

Author response:
### __[Ablation on SEP embedding]__

ㅤ

|    Ablation    |               |_PTB-XL_|
|:--------------:|:-------------:|:------:|
| __SEP embedding__ | __Lead embedding__ |        |
|        X       |       X       |  0.822 |
|        X       |       O       |  0.830 |
|        O       |       X       |  0.820 |
|        O       |       O       |  __0.838__ |

ㅤ

__Reply__: We assessed the effectiveness of the lead indicating modules (e.g., SEP embedding, and lead embedding) by removing them one by one. We added Table 13, the _PTB-XL_ linear evaluation performances of the model trained with or without each lead indicating module in Appendix C. The highest AUROC score is observed when both modules are employed simultaneously, indicating that they mutually contribute to enhancing the model's ability to learn spatio-temporal relationships.

ㅤ

#### __Table caption__: Our proposed case
|     Patch embedding    |     [SEP]       |     Patch (I)    |     [SEP]       |     [SEP]        |     Patch(II)    |     [SEP]        |
|------------------------|---------------|------------------|---------------|----------------|------------------|----------------|
|     __Lead embedding__     |     Lead I    |     Lead I       |     Lead I    |     Lead II    |     Lead II      |     Lead II    |

ㅤ

#### __Table caption__: alternative case
|     Patch embedding    |     [SEP]       |     Patch (I)    |     [SEP]        |     Patch(II)    |     [SEP]        |
|------------------------|---------------|------------------|----------------|------------------|----------------|
|     __Lead embedding__     |     Lead I    |     Lead I       |     Lead II    |     Lead II      |     Lead II    |

ㅤ

Moreover, there is the reason that we appended the SEP embedding twice, i.e., at the beginning and at the end of each lead signal. The above table depicts two cases. Imagine that we do the single lead experiments for lead I ECGs and lead II ECGs. The discrepancy of input shape for lead I and II would discourage the model from distinguishing the lead type. On the other hand, the proposed method can provide a consistent input shape on each single lead experiment. Therefore, we appended SEP embeddings twice.

ㅤ

__Change__: We added Table 13 in Appendix C.

Reviewer comment (Experiments):
The initial statement “detecting various heart diseases” seems to imply that the proposed technique could do so, however, the tests only show a couple of heart conditions. For example, in the PTB XL dataset there are 4 different conditions, why not show the results on all of them?

Author response:
### __[Different problem settings of downstream dataset _PTB-XL_]__

ㅤ

__Reply__: We focused solely on the task of distinguishing the 5 super-classes of diagnostic labels for _PTB-XL_. We conducted additional experiments for the other three settings: 23 diagnostic sub-class labels, 19 form labels, and 12 rhythm labels. In this multi-label classification experiment, our proposed method demonstrated superiority over other baselines (supervised learning and MTAE). Further details have been included in the Appendix F.


|     Methods       |     Categories       |                      |                      |
|-------------------|----------------------|----------------------|----------------------|
|                   |     __Sub-class__         |     __Form__             |     __Rhythm__           |
|     __Supervised__    |     0.914 ± 0.002    |     0.829 ± 0.021    |     0.934 ± 0.003    |
|     __MTAE__          |     0.911 ± 0.001    |     0.793 ± 0.014    |     0.920 ± 0.005    |
|     __ST-MEM (Ours)__      |     __0.929 ± 0.001__    |     __0.895 ± 0.008__    |     __0.966 ± 0.004__    |

ㅤ

__Change__: We added new section as Appendix F.

ㅤ

ㅤ

Reviewer comment (Theory):
“The patches undergo flattening” — why is flattening needed here? Isn’t the signal already flat?

Author response:
### __[Use of word “flattening” in patchfying explanation]__

__Reply__: We are sorry for confusing the audience by misusing the word “flattening.” In spatio-temporal patchifying, in contrast to the temporal patchifying scenario, there is no flatten operation as it is already flattened. We deleted the word “flatten”.
 
__Change__: Revised the sentence.

“The patches undergo flattening and linear projection to form patch embeddings of dimension D.”

-> 

“The patches undergo linear projection to form patch embeddings of dimension D.”

Reviewer comment (Experiments):
The pre-training dataset comes with different sampling rates (two at 500 Hz and 1 at 400 Hz). How was this taken care of? Was the 500 Hz subsampled? This should be explained.

Author response:
### __[Sampling frequency of ECG dataset]__
 
ㅤ

__Reply__: For consistency, we resampled all 6 ECG datasets, including pre-training and downstream datasets, to 250Hz. It is explained in Section 4 and Appendix A.1.

ㅤ

__Change__: We modified Section 4 and Appendix A.1.

ㅤ

ㅤ

Reviewer comment (Reproducibility):
For signals that are longer than 10 seconds, I assume the signal was split into 10 seconds, but how was the outcome computed? Average? Voting? This should be clarified.

Author response:
### __[Explanation on evaluation of cropped ECGs]__
 
ㅤ

__Reply__: To maintain uniformity, we standardized all ECGs to a duration of 10 seconds. ECGs exceeding 10 seconds were disjointedly cropped into 10-second segments. For instance, if there was a 27-second ECG, we created two individual samples: one from 0 to 10 seconds and another from 10 to 20 seconds, discarding 20 to 27 seconds. These two ECGs were trained and evaluated separately, each retaining the labels of the original 27-second ECG. Therefore, these two ECGs are considered as two individual data points and computed independently, neither averaging nor voting. It is explained in Appendix A.
 
ㅤ

__Change__: We modified Section 4.1 and added detail explanation on Appendix A.