Original: MAEEG: Masked Auto-encoder for EEG Representation Learning
Revised: MAEEG: Masked Auto-encoder for EEG Representation Learning

Original: Anonymous Author(s)

Original: AffiliationAddress email

Original: Abstract
Revised: Abstract

Original: Decoding information from bio-signals such as EEG, using machine learninghas been a challenge due to the small data-sets and difficulty to obtain labels.
Revised: Decoding information from bio-signals such as EEG, using machine learning has been a challenge due to the small data-sets and difficulty to obtain labels.

Original: We propose a reconstruction-based self-supervised learning model, the maskedauto-encoder for EEG (MAEEG), for learning EEG representations by learningto reconstruct the masked EEG features using a transformer architecture.
Revised: We propose a reconstruction-based self-supervised learning model, the masked auto-encoder for EEG (MAEEG), for learning EEG representations by learning to reconstruct the masked EEG features using a transformer architecture.

Original: Wefound that MAEEG can learn representations that significantly improve sleep stageclassification ( ∼ 5% accuracy increase) when only a small number of labels aregiven.
Revised: We found that MAEEG can learn representations that significantly improve sleep stage classification ( ∼ 5% accuracy increase) when only a small number of labels are given.

Original: We also found that input sample lengths and different ways of maskingduring reconstruction-based SSL pretraining have a huge effect on downstreammodel performance.
Revised: We also found that input sample lengths and different ways of masking during reconstruction-based SSL pretraining have a huge effect on downstream model performance.

Original: Specifically, learning to reconstruct a larger proportion andmore concentrated masked signal results in better performance on sleep classification.
Revised: Specifically, learning to reconstruct a larger proportion and more concentrated masked signal results in better performance on sleep classification.

Original: Our findings provide insight into how reconstruction-based SSL could helprepresentation learning for EEG.
Revised: Our findings provide insight into how reconstruction-based SSL could help representation learning for EEG.

Original:  Introduction
Revised:  Introduction

Original: Decoding information from Electroencephalography (EEG) and other bio-signal modalities hasenabled health-related clinical applications, such as sleep pattern detection [Ghassemi et al., 2018].
Revised: Decoding information from Electroencephalography (EEG) and other bio-signal modalities has enabled health-related clinical applications, such as sleep pattern detection [Ghassemi et al., 2018].

Original: Recently, machine learning models have been effectively applied on classifying EEG signals, mostlyvia supervised learning [Lawhern et al., 2018].
Revised: Recently, machine learning models have been effectively applied on classifying EEG signals, mostly via supervised learning [Lawhern et al., 2018].

Original: However, such methods rely on a large numberof labels for bio-signal data, which are usually difficult and expensive to obtain.
Revised: However, such methods rely on a large number of labels for bio-signal data, which are usually difficult and expensive to obtain.

Original: One way toimprove accuracy when decoding EEG signals using only a small number of labels is to learn usefulrepresentations from unlabeled data via self-supervised learning (SSL).
Revised: One way to improve accuracy when decoding EEG signals using only a small number of labels is to learn useful representations from unlabeled data via self-supervised learning (SSL).

Original: SSL has been widely explored in the field of computer vision and speech processing.
Revised: SSL has been widely explored in the field of computer vision and speech processing.

Original: The idea of SSLis to boost the performance of a downstream task by learning meaningful representations throughan SSL task.
Revised: The idea of SSL is to boost the performance of a downstream task by learning meaningful representations through an SSL task.

Original: One way to learn representations is through masking and reconstructing the inputsignals (i e., reconstruction-based SSL).
Revised: One way to learn representations is through masking and reconstructing the input signals (i.e., reconstruction-based SSL).

Original: For example, the representations of speech learned from the
Revised: For example, the representations of speech learned from the

Original: Wav2Vec series of SSL models significantly improve the performance of various downstream tasks,such as automatic speech recognition and speech emotion recognition [Baevski et al., 2020, Mitraet al., 2022].
Revised: Wav2Vec series of SSL models significantly improve the performance of various downstream tasks, such as automatic speech recognition and speech emotion recognition [Baevski et al., 2020, Mitra et al., 2022].

Original: Given the success of reconstruction-based SSL methods in other domains, it is desirableto learn whether and how such methods could be adapted to physiological time-series data, such as
Revised: Given the success of reconstruction-based SSL methods in other domains, it is desirable to learn whether and how such methods could be adapted to physiological time-series data, such as

Original: EEG.
Revised: EEG.

Original: In this study, we explore representation learning using reconstruction-based SSL on EEG data.
Revised: In this study, we explore representation learning using reconstruction-based SSL on EEG data.

Original: Wepropose an SSL model, masked auto-encoder for learning EEG representations (MAEEG), whichcan learn EEG representations by reconstructing the raw signal from masked features.
Revised: We propose an SSL model, masked auto-encoder for learning EEG representations (MAEEG), which can learn EEG representations by reconstructing the raw signal from masked features.

Original: We foundthat MAEEG pretraining learns meaningful EEG representations, which yield better performanceon sleep stage classification.
Revised: We found that MAEEG pretraining learns meaningful EEG representations, which yield better performance on sleep stage classification.

Original: We further explore how masking may affect SSL and downstream taskperformance, and found that in general, a higher probability and more concentrated masking yieldsbetter task performance
Revised: We further explore how masking may affect SSL and downstream task performance, and found that in general, a higher probability and more concentrated masking yields better task performance.

Original:  Background and Related Work
Revised:  Background and Related Work

Original: SSL methods for time-series data have been widely applied in domains such as natural languageand speech processing.
Revised: SSL methods for time-series data have been widely applied in domains such as natural language and speech processing.

Original: Recently, studies have also explored various SSL methods for bio-signaldata, such as EEG.
Revised: Recently, studies have also explored various SSL methods for bio-signal data, such as EEG.

Original: For example, combining various ways of EEG data augmentation and contrastivelearning, Mohsenvand et al.
Revised: For example, combining various ways of EEG data augmentation and contrastive learning, Mohsenvand et al.

Original: [2020] showed improved performance on several EEG classification tasks.
Revised: [2020] showed improved performance on several EEG classification tasks.

Original: On the other hand, Banville et al.
Revised: On the other hand, Banville et al.

Original: [2021] proposed three SSL methods aiming to learn representationsby discriminating the temporal relationships between the EEG samples.
Revised: [2021] proposed three SSL methods aiming to learn representations by discriminating the temporal relationships between the EEG samples.

Original: These studies showed that a
Revised: These studies showed that a

Original: SSL-pretrained encoder could produce features that outperformed supervised learning when there isonly a small number of labeled data.
Revised: SSL-pretrained encoder could produce features that outperformed supervised learning when there is only a small number of labeled data.

Original: Another type of SSL methods for time-series data learn representations through masking andreconstructing features, and we refer to them as reconstruction-based SSL in the current study.
Revised: Another type of SSL methods for time-series data learn representations through masking and reconstructing features, and we refer to them as reconstruction-based SSL in the current study.

Original: Reconstruction-based SSL usually contains several stages: (1) The raw input signal is first encodedas features using a convolutional encoder; (2) part of the features are “masked" by setting the featuresto a certain value; (3) the masked features are sent to a second encoder, usually a transformer encoder,which aims to reconstruct the masked part of the features using the unmasked features as context;(4) finally, contrastive loss, or reconstruction loss is calculated to optimize the SSL task.
Revised: Reconstruction-based SSL usually contains several stages: (1) The raw input signal is first encoded as features using a convolutional encoder; (2) part of the features are “masked" by setting the features to a certain value; (3) the masked features are sent to a second encoder, usually a transformer encoder, which aims to reconstruct the masked part of the features using the unmasked features as context; (4) finally, contrastive loss, or reconstruction loss is calculated to optimize the SSL task.

Original: Inspiredby the Wav2Vec2 model, Kostas et al.
Revised: Inspired by the Wav2Vec2 model, Kostas et al.

Original: [2021] proposed BENDR, a reconstruction-based SSL withcontrastive learning for processing EEG signals.
Revised: [2021] proposed BENDR, a reconstruction-based SSL with contrastive learning for processing EEG signals.

Original: While the idea of using massive data (i.e., the TUHdataset, [Obeid and Picone, 2016]) to learn EEG representations using BENDR seems to improvesome downstream tasks, the results are not consistent across different data-sets/tasks.
Revised: While the idea of using massive data (i.e., the TUH dataset, [Obeid and Picone, 2016]) to learn EEG representations using BENDR seems to improve some downstream tasks, the results are not consistent across different data-sets/tasks.

Original: This couldbe due to the inconsistency of data-sets and model settings between pretraining and downstreamstages.
Revised: This could be due to the inconsistency of data-sets and model settings between pretraining and downstream stages.

Original: Overall, it remains unclear if their reconstruction-based SSL methods could learn usefulrepresentations for EEG signals.
Revised: Overall, it remains unclear if their reconstruction-based SSL methods could learn useful representations for EEG signals.

Original: In this study, we examine whether reconstruction-based SSL methods are effective for learning EEGrepresentations, and how different ways of masking would affect representation learning.
Revised: In this study, we examine whether reconstruction-based SSL methods are effective for learning EEG representations, and how different ways of masking would affect representation learning.

Original: Specifically,we propose an SSL model, MAEEG, inspired by both BENDR and the MAE model proposed incomputer vision [He et al., 2022], to compare with BENDR and to broaden our understanding ofreconstruction-based SSL for EEG.
Revised: Specifically, we propose an SSL model, MAEEG, inspired by both BENDR and the MAE model proposed in computer vision [He et al., 2022], to compare with BENDR and to broaden our understanding of reconstruction-based SSL for EEG.

Original:  Method3.1 Sleep EEG data-set
Revised:  Method 3.1 Sleep EEG Data-set

Original: To examine how reconstruction-based SSL may help learn useful representations for classifying sleepstages from EEG signals, we use the data-set provided from 2018 PhysioNet Challenge [Ghassemiet al., 2018].The data-set contains overnight polysomnography of 994 subjects (i.e., around 7,000hours of data, 65% male, mean age: 55y ± 14.4), which was split into a training set (596 subjects),a validation set (198 subjects) and a testing set (200 subjects) as in Ghassemi et al.
Revised: To examine how reconstruction-based SSL may help learn useful representations for classifying sleep stages from EEG signals, we use the data-set provided from 2018 PhysioNet Challenge [Ghassemi et al., 2018].The data-set contains overnight polysomnography of 994 subjects (i.e., around 7, hours of data, 65% male, mean age: 55y ± 14.4), which was split into a training set (596 subjects), a validation set (198 subjects) and a testing set (200 subjects) as in Ghassemi et al.

Original: Eachpolysomnogram contains simultaneous recordings of 6-channel EEG, ECG and respiratory signalsduring sleep and two sets of labels: 5 sleep stages and respiratory-effort related arousals.
Revised: Each polysomnogram contains simultaneous recordings of 6-channel EEG, ECG and respiratory signals during sleep and two sets of labels: 5 sleep stages and respiratory-effort related arousals.

Original: The labelswere assigned by a trained sleep technician for each 30-sec non-overlapping window.
Revised: The labels were assigned by a trained sleep technician for each 30-sec non-overlapping window.

Original: Here, wefocused on using the 6-channel EEG signals to classify 5 stages of sleeping (i.e., wake, N1, N2, N3,
Revised: Here, we focused on using the 6-channel EEG signals to classify 5 stages of sleeping (i.e., wake, N1, N2, N3,

Original: REM).
Revised: REM).

Original: All the models and experiments presented in this section are built and run on Pytorch 1.8.1.
Revised: All the models and experiments presented in this section are built and run on PyTorch 1.8.1.

Original: 3.2 Reconstruction-based SSL models
Revised: 3.2 Reconstruction-based SSL models

Original: BENDR.
Revised: BENDR.

Original: The model architecture of BENDR is shown in Figure 1A. First, the 6-channel 100Hzraw EEG signals are the input to the 6 convolutional layers, which encode the raw EEG into64-dimensional convolved features ( t i , ∼ 1.05Hz).
Revised: The model architecture of BENDR is shown in Figure 1A. First, the 6-channel 100Hz raw EEG signals are the input to the 6 convolutional layers, which encode the raw EEG into 64-dimensional convolved features ( t i , ∼ 1.05Hz).

Original: Dropout and GroupNorm are added to eachconvolutional layer and GELU is the output activation function.
Revised: Dropout and GroupNorm are added to each convolutional layer and GELU is the output activation function.

Original: A mask I m is generated to “mask out"part of the features t i where i ∈ I m (i.e., set the activation to Gaussian noise M ∼ N (0 , (cid:112) 1 / 64) ).
Revised: A mask I m is generated to “mask out" part of the features t i where i ∈ I m (i.e., set the activation to Gaussian noise M ∼ N (0 , (cid:112) 1 / 64) ).

Original: The masked features q i is then sent to an 8-layer transformer encoder and encoded as 192-dimensionaloutput features.
Revised: The masked features q i is then sent to an 8-layer transformer encoder and encoded as 192-dimensional output features.

Original: The positions are encoded using a convolutional layer, which acts as relative positionalembedding.
Revised: The positions are encoded using a convolutional layer, which acts as relative positional embedding.

Original: Finally, a final convolutional layer maps the features back to 64-dimensional contextualfeatures c i .
Revised: Finally, a final convolutional layer maps the features back to 64-dimensional contextual features c i .

Original: Contrastive loss is calculated by comparing c
Revised: Contrastive loss is calculated by comparing c

Original: i with the unmasked features t i .
Revised: i with the unmasked features t i .

Original: MAEEG.
Revised: MAEEG.

Original: The MAEEG model has a similar architecture to BENDR (Figure 1B), but has twoadditional layers to map the transformer output back to the raw EEG dimensions.
Revised: The MAEEG model has a similar architecture to BENDR (Figure 1B), but has two additional layers to map the transformer output back to the raw EEG dimensions.

Original: A reconstructionloss is calculated by comparing the reconstructed EEG ( ˆx ) and input EEG ( x ) signals as 1 − ˆx · x ∥ ˆx ∥∥ x ∥(Figure 1 B).
Revised: A reconstruction loss is calculated by comparing the reconstructed EEG ( ˆx ) and input EEG ( x ) signals as 1 − ˆx · x ∥ ˆx ∥∥ x ∥ (Figure 1 B).

Original: The key difference between BENDR and MAEEG is that instead of using contrastivelearning, MAEEG learns representations by minimizing the reconstruction loss.
Revised: The key difference between BENDR and MAEEG is that instead of using contrastive learning, MAEEG learns representations by minimizing the reconstruction loss (Appendix A.1).

Original: 3.3 Features and Analyses
Revised: 3.3 Features for Downstream classification

Original: After SSL pretraining, we examined downstream task performances by adding a linear classificationlayer on top of either of the two features: the convolved features t i vs. the contextual features
Revised: After SSL pretraining, we examined downstream task performances by adding a linear classification layer on top of either of the two features: the convolved features, t i , which are the output of the convolutional encoder, and the contextual features, c i , which are the output of the transformer encoder.

Revised: To build downstream classification model for classifying sleep stages, we preserve most of the SSL architecture with the pretrained encoders, and simply add a fully-connected classification layer which takes the SSL features as input and generates probabilities for predicting the 5 sleep stages (Figure A.1).

Revised: The choice of adding a linear classifier instead of a non-linear one is to prevent the classifier from being too powerful and making it difficult to evaluate the representations learned from

Revised: SSL pretraining.

Original: c i ,under either of the two settings: freeze vs. fine-tune the encoder (See Appendix A.1, Figure A.1).
Revised: To better understand representations learned from different stages of training (i.e., SSL pretraining with no labels and downstream supervised training with labels), we conducted two downstream classification analyses: (1) To examine low-level features t i learned solely from the SSL pre-training, we use t i while “freezing" the encoder (meaning that the weights of the encoder are fixed) during training for the downstream classification. Only the classification layer gets trained in this condition (Figure A.1 A); (2) To examine the pre-trained model as an initialization for supervised training, we use c i while fine-tuning the encoders and the classification layer during downstream classification (Figure A.1 B). Specifically, we extracted c

Revised: i from layer 2 as the SSL features to conduct all the analyses given that it showed the best results among all the transformer layers.

Revised: We conducted these analyses on MAEEG, BENDR and a baseline supervised model which has the same architecture as other models but the encoders are randomly initialized, i.e., no SSL pretraining prior to the downstream classification task.

Revised: For each downstream task, we trained 100 epochs in total and validated the model every 2 epochs.

Revised: We selected the models with highest validation accuracy to run the test set and get the final accuracy shown in Figure 2.

Original: Wealso investigated how masking in reconstruction-based SSL could affect EEG representation learningby modifying the masking probability and lengths (See Appendix A.2).
Revised: We also investigated how masking in reconstruction-based SSL could affect EEG representation learning by modifying the masking probability and lengths (See Appendix A.2).

Revised: We applied analyses with various number of subjects or percent of sessions to examine model performance given different amount of labels.

Original:  Results4.1 MAEEG learned useful representations for classifying sleep stages
Revised:  Results 4.1 MAEEG Learned Useful Representations for Classifying Sleep Stages

Original: We first examine convolved features  learned solely by SSL pretraining by freezing the convolutionalencoders when performing the sleep stage classification task.
Revised: We first examine convolved features t i learned solely by SSL pretraining by freezing the convolutional encoders during training for the sleep stage classification task.

Original: We found that representations pretrainedby MAEEG (both 30s and 100s samples input) significantly outperformed the representations trainedby BENDR and the supervised baseline model across different amount of subjects for training (Figure2A).
Revised: We found that representations pretrained by MAEEG (both 30s and 100s samples input) significantly outperformed the representations trained by BENDR and the supervised baseline model across different amount of subjects for training (Figure 2A).

Original: This suggests that representations learned from MAEEG pretraining do capture features that areuseful for classifying sleeping stages.
Revised: This suggests that representations learned from MAEEG pretraining do capture features that are useful for classifying sleeping stages.

Original: We next examined the model performance using contextual features when fine-tuning the encodersduring the sleep stage classification task.
Revised: Next, to examine how SSL pretraining may improve the supervised learning performance, we used contextual features c i and fine-tuned the model during the sleep stage classification task.

Original: We found that when only given one subject’s labels(containing all sleep stages), MAEEG pretrained on 100s (MAEEG-100s) outperformed other models,while BENDR pretrained on 30s (BENDR-30s) performed the best on 10 subjects’ and 50 subjects’labels (Figure2B).
Revised: We found that when only given one subject’s labels (containing all sleep stages), MAEEG pretrained on 100s (MAEEG-100s) outperformed other models, while BENDR pretrained on 30s (BENDR-30s) performed the best on 10 subjects’ and 50 subjects’ labels (Figure2B).

Original: The best classification performance (90% accuracy) was from MAEEG-100s.
Revised: The best classification performance (90% accuracy) was from MAEEG-100s.

Original: Interestingly, MAEEG-100s performed significantly better compared to MAEEG-30s, and an oppositeeffect on BENDR was observed - BENDR-30s performed significantly better compared to BENDR100s, suggesting that representations learned from the two reconstruction-based SSL models are distinct and the sample length for pretraining has a big effect on downstream performance.
Revised: Interestingly, MAEEG-100s performed significantly better compared to MAEEG-30s, and an opposite effect on BENDR was observed - BENDR-30s performed significantly better compared to BENDR-100s, suggesting that representations learned from the two reconstruction-based SSL models are distinct and the sample length for pretraining has a big effect on downstream performance.

Original: Tounderstand how the pretraining differs under different experimental settings, we further examined theattention maps learned during pretraining.
Revised: To understand how the pretraining differs under different experimental settings, we further examined the attention maps learned during pretraining.

Original: We found that models  learned to attend  distant featuresperformed better than models  learned to attend neighboring features (See Appendix A.3, FigureA.3).
Revised: We found that models that learned to attend to distant features (e.g. 20 seconds away) performed better than models that learned to attend neighboring features (See Appendix A.3, FigureA.3).

Original: Also, the advantage of MAEEG-pretrained convolved features did not reflect in model performancewhen the encoders were fine-tuned, suggesting that while MAEEG may learn some useful featuresto discriminate sleep stages through SSL pretraining, such pretraining may also hinder the models’ability to learn at the same time.
Revised: Also, the advantage of MAEEG-pretrained convolved features did not reflect in performance when the encoders were fine-tuned, suggesting that while MAEEG may learn some useful features to discriminate sleep stages through SSL pretraining, but may also hinder the models’ ability to learn at the same time.

Original: 4.2 Masking effects in MAEEG
Revised: 4.2 Masking Effects in MAEEG

Original: [Figure] Next, we examined how different ways of masking during reconstruction-based SSL pretraining mayaffect representation learning and downstream classification performance using contextual feature c iwith finetuning, by varying the mask rate and the number of mask chunks (See Appendix A.2).
Revised: Next, we examined how different ways of masking during reconstruction-based SSL pretraining may affect representation learning and downstream classification performance using contextual feature c i with finetuning, by varying the mask rate and the number of mask chunks (See Appendix A.2).

Original: Wefirst examined the SSL loss during pretraining, and found that higher percentage of masking with lesschunks (which yields longer masks) yielded higher SSL loss (i.e. more difficult to train).
Revised: We first examined the SSL loss during pretraining, and found that higher percentage of masking with less chunks (which yields longer masks) yielded higher SSL loss (i.e. more difficult to train).

Original: We thenexamined the downstream model performance and found that the model performances differ mostwhen given the least number of labels - such as 0.1% (i.e. ∼ 4 hours of labels) of downstream trainingdata (Figure A.2B).
Revised: We then examined the downstream model performance and found that the model performances differ most when given the least number of labels - such as 0.1% (i.e. ∼ 4 hours of labels) of downstream training data (Figure A.2B).

Original: With such small number of labels, we found that masking 75% of the tokensperformed better overall compared to masking 50% and 25% of the tokens, and models with a singlechunk performed better than models with multiple chunks.
Revised: With such small number of labels, we found that masking 75% of the tokens during SSL pretraining performed better overall compared to masking 50% and 25% of the tokens, and models with a single chunk performed better than models with multiple chunks.

Original: All models with MAEEG pretrainingperformed better than the models without pretraining (Figure A.2C).
Revised: All models with MAEEG pretraining performed better than the models without pretraining (Figure A.2C).

Original: Finally, we examined whethersuch a trend (single chunk with longer mask performs better) preserved when we vary the maskspan of a single mask chunk, and whether it generalizes to BENDR.
Revised: Finally, we examined whether such a trend (single chunk with longer mask performs better) preserved when we vary the mask span of a single mask chunk, and whether it generalizes to BENDR.

Original: We found a consistent patternacross both MAEEG and BENDR, that models pretrained with longer mask span performs better thanmodels pretrained with shorter mask span when there only a small number of labels (Figure A.2D).
Revised: We found a consistent pattern across both MAEEG and BENDR, that models pretrained with longer mask span performs better than models pretrained with shorter mask span when there only a small number of labels (Figure A.2D).

Original: This finding is likely due to the fact that a single and long mask during pretraining forced the modelto learn a more difficult task and yield more useful representations for sleep stage classification.  
Revised: This finding is likely due to the fact that a single and long mask during pretraining forced the model to learn a more difficult task and yield more useful representations for sleep stage classification.  Discussion and Limitations

Revised: In the current study, we present MAEEG and compare it with another SSL model, BENDR, to examine how reconstruction-based SSL could be applied to learn meaningful EEG representations.

Revised: Specifically, we examined how the representations learned from these two SSL approaches could be applied to different downstream conditions by varying the amount of downstream training data.

Original: Conclusion
Revised: To compare model performance while controlling the subject variance, we showed the results in which we sampled data from different number of subjects (Figure 2). While we showed that representations learned from MAEEG yielded better performance, we could not exclude the possibility that some of the effect might still come from the specific subjects that we sampled from. We also conducted analyses by sampling different sessions from every subject as shown in the masking analyses (Figure A.2).

Revised: However, we are still sensitive to potential label imbalance and to the large subject variance in a small amount of data which could make the training more difficult.

Revised: Subject variance and imbalanced classification are common challenges for learning from bio-signal data.

Revised: Future studies should take these issues into account while examining representation learning for EEG or bio-signals.

Revised: One interesting observation in the current study is that masking out more signal (e.g., 75%) in the latent space yielded better downstream performance compared to masking out less signal (e.g., 25%).

Revised: While this result is not intuitive, the same trend was observed in computer vision where masking more of the signal during SSL pretraining could result in better downstream performance

Revised: [He et al., 2022].

Revised: The reason behind this could be that asking the model to reconstruct a larger amount of missing samples forces the model to learn the relationship between the remaining context and the missing part of the signal, instead of just interpolating it.

Revised: This trend can be dependent on the downstream task.

Revised: We suggest future studies to always evaluate the SSL model based on its downstream performance instead of pretraining performance.

Revised: Finally, the computational cost for MAEEG and BENDR are similar in the current study, given that both models have the signal masked in the latent space.

Revised: Such choice of masking might be the reason why MAEEG only learns to reconstruct slow signal variations (Section A.1).

Revised: Masking in the raw signal space may help MAEEG to learn to reconstruct more fine-grained signals but could also make the computation more expensive.

Revised: Alternatively, masking in the time-frequency spectrum domain could be beneficial for EEG representation learning [Xu et al., 2022].

Revised:  Conclusion

Original: Representation learning for EEG can be useful given the difficulty in collecting large amount of dataand labels.
Revised: Representation learning for EEG can be useful given the difficulty in collecting large amount of data and labels.

Original: In this study, we proposed a reconstruction-based SSL model, MAEEG, for learning
Revised: In this study, we proposed a reconstruction-based SSL model, MAEEG, for learning

Original: EEG representations by reconstructing EEG signals using masked auto-encoder architecture.
Revised: EEG representations by reconstructing EEG signals using masked auto-encoder architecture.

Original: Wetested our model on sleep stage classification and found that MAEEG can learn useful representationssolely from the pre-text task which boost the downstream classification performance.
Revised: We tested our model on sleep stage classification and found that MAEEG can learn useful representations solely from the pre-text task which boost the downstream classification performance.

Original: We also foundthat different ways of masking and sample lengths selected during SSL pre-training can significantlyaffect the downstream classification performance.
Revised: We also found that different ways of masking and sample lengths selected during SSL pre-training can significantly affect the downstream classification performance.

Original: We encourage future studies to examine howreconstruction-based SSL may help representation learning on other time-series datafor health.
Revised: We encourage future studies to examine how reconstruction-based SSL may help representation learning on other time-series data

Original: A AppendixA.1 Features for Downstream classification
Revised: A Appendix A.1 EEG Signal Reconstruction

Original: As mentioned, two features are learned by the SSL models - the convolved features, t i , which arethe output of the convolutional encoder, and the contextual features, c i , which are the output ofthe transformer encoder. To build downstream classification model for classifying sleep stages, wepreserve most of the SSL architecturewith the pretrained encoders, and simply add a fully-connectedclassification layer which takes the SSL features as input and generates probabilities for predictingthe 5 sleep stages (Figure A.1).
Revised: Note that the current MAEEG architecture (Figure 1B) is inspired by the MAE methods proposed in

Original: To better understand representations learned from different stagesof training (i.e., SSL pretraining with no labels and downstream supervised training with labels),we conducted two types of downstream classification analyses: (1) To examine low-level featurest i learned solely from the SSL pre-training, we use t i while “freezing" the convolutional encoder(meaning that the weights of the encoder are fixed) during downstream classification.
Revised: He et al.

Original: Only theclassification layer gets trained in this condition (Figure A.1 A); (2) To examine the pre-trainedmodel as an initialization for supervised training, we use c i while fine-tuning both convolutional andtransformer encoders and the classification layer during downstream classification (Figure A.1 B).
Revised: [2022] rather than adapted the model architecture. Instead of using transformer layers as decoder as in He et al.

Original: For (2), to determine which transformer layer contains the most useful SSL features for classification,we first probed the features from each layer and found that layer 2 showed the best classificationperformance out of all layers. We therefore extracted c
Revised: [2022], we simply add two convolution layers in MAEEG to allow the network to reconstruct the raw EEG signal temporally and spatially. By visualizing the reconstructed signal, we observed that MAEEG mainly learned to reconstruct the mean of the short-timescale signals, and capture the long-timescale fluctuations.

Original: i from layer 2 as the SSL features to conductall following analyses. We conducted these analyses onMAEEG, BENDR and a baseline supervisedmodel which has the same architecture as other models but the encoders are randomly initialized,i.e., no SSL pretraining prior to the  downstream classification task. We applied analyses with various number of subjects or percent of sessions to examine model performance given different amount oflabels.
Revised: Based on the results, capturing long-timescale fluctuations seems already beneficial for learning representations for sleep stage classification, possibly due to the fact that such specific task relies more on long-timescale information. It is likely that other downstream tasks would require the model to learn representations by reconstructing more finegrained information.

Original: A.2 Masking effect analyses
Revised: A.2 Masking Effect Analyses

Original: The masking during SSL pretraining determines the difficulty of the SSL task, and the representationslearned from such task.
Revised: The masking during SSL pretraining determines the difficulty of the SSL task, and the representations learned from such task.

Original: Therefore, it is critical to understand what could be an ideal way of maskingfor learning EEG signals.
Revised: Therefore, it is critical to understand what could be an ideal way of masking for learning EEG signals.

Original: In the original BENDR, they set a probability (p=0.065) for each tokento be the beginning of a mask with a length of 10 tokens.
Revised: In the original BENDR, they set a probability (p=0.065) for each token to be the beginning of a mask with a length of 10 tokens.

Original: The problem with this is that for eachsequence, the mask can vary a lot - sometimes most of the tokens can be masked with overlap, andsometimes there is no token being masked - which may cause the instability of representation learning.
Revised: The problem with this is that for each sequence, the mask can vary a lot - sometimes most of the tokens can be masked with overlap, and sometimes there is no token being masked - which may cause the instability of representation learning.

Original: Here, we tried to use a more systematic way to generate the mask, by determining the mask rate andthe number of mask chunks.
Revised: Here, we tried to use a more systematic way to generate the mask, by determining the mask rate and the number of mask chunks.

Original: These two factors would determine how long each mask looks like, asillustrated in Figure A.

Revised: These two factors would determine how long each mask looks like, as illustrated in Figure A.

Original: A.
Revised: A.

Original: We did two analyses:
Revised: We did two analyses:

Original: Pretrained the models with 3 mask rates: 75%, 50% and 25%, and 3 numbers of maskchunks: 1, 5, and 10. 
Revised: Pretrained the models with 3 mask rates: 75%, 50% and 25%, and 3 numbers of mask chunks: 1, 5, and 10. 2.

Original: Pretrained models with single mask chunk but vary the mask span to further confirm the 
Revised: Pretrained models with single mask chunk but vary the mask span to further confirm the effects.

Original: effects. The resulting pretrained models were evaluated by the downstream performance when using thelearned features for classification.
Revised: The resulting pretrained models were evaluated by the downstream performance when using the learned features for classification.

Original: We conducted all the analyses on BENDR and MAEEG pretrainedon 100s samples because models pretrained on 30s samples did not have enough tokens for examiningthe effect of these mask variations.
Revised: We conducted all the analyses on BENDR and MAEEG pretrained on 100s samples because models pretrained on 30s samples did not have enough tokens for examining the effect of these mask variations.

Revised: Also, we examined varying percentage of training samples during the downstream task to examine the effect of masking under different sample size conditions.

Revised: Note that instead of varying the number of subjects as in the main task, we varied the percentage of data sampled from each subject, which eliminates the potential bias coming from the specific sampled subjects.

Original: A.3 Attention learned from SSL
Revised: A.3 Attention Learned from SSL

Original: To understand how the EEG representation learned from reconstruction-based SSL help sleep classification, we visualized the attention from the transformer layer where features were extracted (i.e.layer 2) after SSL pretraining but before extensive downstream supervised training (we slightlytuned the models with 1 epoch with 1 subject’s data just to get stable representations.)
Revised: To understand how the EEG representation learned from reconstruction-based SSL help sleep classification, we visualized the attention from the transformer layer where features were extracted (i.e. layer 2) after SSL pretraining but before extensive downstream supervised training (we slightly tuned the models with 1 epoch with 1 subject’s data just to get stable representations.)

Original: We extractthe token-by-token attention maps for the baseline supervised model and the models pretrained by
Revised: We extract the token-by-token attention maps for the baseline supervised model and the models pretrained by

Original: MAEEG and BENDR.
Revised: MAEEG and BENDR.

Original: We visualized the attention maps from a sample where the label is N2 (Figure
Revised: We visualized the attention maps from a sample where the label is N2 (Figure

Original: A.3).
Revised: A.3).

Original: Compared to the attention from baseline models which is equally distributed as shown in
Revised: Compared to the attention from baseline models which is equally distributed as shown in

Original: Figure A.3 A, MAEEG- and BENDR- pretrained models showed distinct attention patterns learnedfrom pretraining.
Revised: Figure A.3 A, MAEEG- and BENDR- pretrained models showed distinct attention patterns learned from pretraining.

Original: Interestingly, we observed that MAEEG-100s and MAEEG-30s actually learnedvery different attention patterns, with MAEEG-30s learning to attend locally to neighboring tokenswhile MAEEG-100s learning to attend more distant tokens (Figure A.3 B).
Revised: Interestingly, we observed that MAEEG-100s and MAEEG-30s actually learned very different attention patterns, with MAEEG-30s learning to attend locally to neighboring tokens while MAEEG-100s learning to attend more distant tokens (Figure A.3 B).

Original: On the other hand,
Revised: On the other hand,

Original: BENDR-100s and BENDR-30s also learned very different attention patterns, with BENDR-100slearning to attend locally and BENDR-30s learning a more global distributed attention.
Revised: BENDR-100s and BENDR-30s also learned very different attention patterns, with BENDR-100s learning to attend locally and BENDR-30s learning a more global distributed attention.

Original: This suggeststhat representation learning can be very different based not only on the model architecture, butalso on subtle SSL settings such as the length of pretraining samples.
Revised: This suggests that representation learning can be very different based not only on the model architecture, but also on subtle SSL settings such as the length of pretraining samples.

Original: Moreover, MAEEG-30s and
Revised: Moreover, MAEEG-30s and

Original: BENDR-100s performed worse on the classification task compared to their counterparts (Figure 2B),suggesting that local attention acquired from pretraining may not be desirable for the sleep stageclassification task, possibly due to the fact that this specific task relies on long-timescale information.
Revised: BENDR-100s performed worse on the classification task compared to their counterparts (Figure 2B), suggesting that local attention acquired from pretraining may not be desirable for the sleep stage classification task, possibly due to the fact that this specific task relies on long-timescale information.

Original: Overall, we found that masking 75% of the signals with 1 chunk results in the best performance
Revised: Overall, we found that masking 75% of the signals with 1 chunk results in the best performance

Original: Classification results using 0.1% of labels when varying the masking span.
Revised: Classification results using 0.1% of labels when varying the masking span.

Original: We found that longer mask span yields better results for both MAEEG-100s and BENDR-100s.
Revised: We found that longer mask span yields better results for both MAEEG-100s and BENDR-100s.

Original: np no pretrain
Revised: np no pretrain