SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ke Yi*
cs_kerry@mail.scut.edu.cn
Yansen Wang†
yansenwang@microsoft.com
Kan Ren
kanren@microsoft.com
Dongsheng Li
dongsli@microsoft.com
Large-scale pre-training has shown great potential to enhance models on downstream tasks in vision and language. Developing similar techniques for scalp electroencephalogram (EEG) is suitable since unlabelled data is plentiful. Meanwhile, various sampling channel selections and inherent structural and spatial information bring challenges and avenues to improve existing pre-training strategies further. In order to break boundaries between different EEG resources and facilitate cross-dataset EEG pre-training, we propose to map all kinds of channel selections to a unified topology. We further introduce MMM, a pre-training framework with Multi-dimensional position encoding, Multi-level channel hierarchy, Multi-stage pre-training strategy built on the unified topology to obtain topology-agnostic representations. Experiments demonstrate that our approach yields impressive improvements over previous state-of-the-art techniques on emotional recognition benchmark datasets.
Scalp electroencephalography (EEG), a physiological signal that shows the human brain's status and can be collected in a non-invasive and low-cost manner, is widely used in disease diagnosis, medical monitoring, rehabilitation, and some recent brain-computer interface systems [1, 2, 3]. Although collecting EEG signals is relatively easy, interpreting the signals and providing meaningful labels often require huge expert efforts. Therefore, a few recent works have emerged to learn representations of EEG segments with self-supervised learning methods and apply the pre-trained model to several downstream tasks, including emotion recognition, sleep staging, epileptic seizure detection, and so on [4, 5, 6, 7, 8]. While there are many publicly available EEG datasets for pre-training, transferring across them is very challenging because these EEG recordings can vary significantly in terms of montage (the number and the places of electrodes placed on the scalp) and sample rate. These factors will completely change the format and the EEG contents' meaning. For example, a typical setting of the emotion recognition task involves 62 channels based on the 10-10 electrode system to track the fine-grained changes with very high spatial resolution, and a more practical setting to monitor seizure events may only have 19 channels based on the 10-20 system. Nevertheless, previous self-supervised learning methods avoid handling this problem and thus can only work within a single dataset, which hinders the model from being trained on more available resources.
*The work was conducted during Ke Yi's internship at Microsoft Research Asia.
$\dagger$  Correspondence to Yansen Wang and Kan Ren.
<sup>1</sup>Project link: https://seqml.github.io/MMM/.
Another challenge of EEG pre-training is to encode the structure and model the relationships among EEG signals from different locations. These signals are collected by the electrodes placed on the scalp, which can be viewed as a 2-D manifold with spatial information. This spatial information is often crucial for the expert to provide precise labels. For example, lesions can be located by comparing the amplitudes and phases of abnormal discharges between the left and right brain, and changes in emotional states may be reflected by asymmetrical EEG activity [9, 10, 11].
In summary, to get a capable and powerful EEG pre-trained model, we need 1) well encoding spatial information into the representation and 2) pre-training models with EEG corpus having various sensor configurations. In this paper, we propose MMM, a self-supervised learning framework based on the masked auto-encoder [12], which maps EEG data with different sensor configurations into unified representations. We introduce region-wise tokens to extract local information from the EEG channels, where region-wise tokens and origin channels form a Multi-level hierarchy learning. The well-learned region-wise tokens form the unified representation, which can be applied in downstream tasks with rich spatial information. We demonstrate that introducing Multi-dimensional position encoding to encode geometric information of sensors benefits pre-training, and Multi-stage mask strategy with random mask and region-wise mask can enhance the hidden representation's robustness. To further evaluate the effectiveness of our proposed framework, we validate our approach on emotion recognition tasks, which is a typical EEG downstream task and requires the model to capture the spatial information well.
- We propose to use a unified topology to model the scalp and map different montages onto the same topology.
- We propose a novel pre-training framework to learn unified geometry-aware EEG representations on top of the unified topology that can be generalized to different EEG channel configurations.
- Our pre-trained model gains state-of-the-art performance on the emotion recognition benchmarks. More experiments demonstrated its strong ability to transfer between different datasets even with different montages.
Self-supervised learning for EEG Electroencephalograms (EEGs) are commonly used to diagnose neurological, psychiatric, and sleep disorders and in applications involving brain-machine interfaces. In the EEG domain, self-supervised learning has emerged as a promising approach. In [8], the author designs two pretext tasks, temporal context prediction, and contrastive predictive coding, to perform representation learning of EEG signals and applies them to two downstream tasks: sleep staging and pathology detection. [6] introduce a set of data augmentations for EEG and extend the SimCLR framework to extract channel-wise features on time-series EEG data. The above methods focus on temporal information modeling while ignoring spatial information. [13, 14] attempt to solve the cross-montage problem by selecting a fixed channel set, while it might cause spatial information loss when faced with different configurations. [7] use GNN and RNN to model the spatiotemporal dependencies while using the generative pre-training method to improve model performance further. However, their GNN-based model assumes a fixed channel set in their experiments and might naturally be worse at transferring than the attention-based model.
To the best of our knowledge, there is no method trying to use cross-dataset EEG as their pre-training corpus.
EEG-based emotion recognition Before the deep learning methods were extensively adopted in this area, spectral EEG features were commonly investigated, such as power spectral density (PSD), differential entropy (DE), and differential asymmetry (DASM), from which DE feature is proved to be the most precise and stable one in EEG-based emotion recognition tasks [11]. Several EEG-based emotion recognition methods have been proposed based on the DE feature. [15, 16, 17] apply recurrent neural networks to traverse signals and capture the spatial dependency between the EEG channel. However, RNN-based modeling takes the EEG channel as a fixed sequence, which loses sensors' geometric information and makes transferring hard. [18, 19] apply graph neural network to model spatial relationship among EEG channel and GCN-based model achieve state-of-the-art in
EEG-based emotion recognition. However, these GCN-based models are only claimed to be used in supervised training. [4] has shown that Masked EEG modeling is an effective self-supervised pre-training paradigm for EEG data. However, this work takes EEG data as a sequence while ignoring the sensors' physical location information. Again, the above methods are all not trying to pre-training/training on cross-dataset EEG.
3.1 Task Formulation
Learning a good EEG representation involves modeling from both temporal and spatial aspects. While temporal modeling is better explored in the previous works, our main focus is on spatial structures. We use DE feature [20] extracted from the raw EEG to aggregate the temporal information. This simplification can bring several benefits: 1) DE feature is a well-studied feature that is demonstrated to be effective for temporal modeling of EEG, especially in EEG-based emotion recognition tasks [11, 21]. 2) The process of DE feature extraction doesn't involve learnable parameters and can exclude the eliminating impact of sample rate variations. This allows the feature to be seamlessly transferred across datasets with different time resolutions. 3) The extracted features are widely used across a wide range of baseline models [22, 23, 24], thereby facilitating equitable comparisons among various frameworks. We describe how DE features are extracted in the Supplementary file. It is worth noting that our proposed method is also applicable to raw EEG signals and other kinds of features as long as the temporal information is handled in a consistent manner across datasets.
In this way, the EEG pre-training task can be formulated as learning a unified hidden representation  $h = \operatorname{Encoder}(x) \in \mathbb{R}^H$  from a dataset  $\mathcal{X}$  containing the extracted DE features  $x \in \mathcal{X} \subset \mathbb{R}^{C_x \times F}$ , which can be afterward used in downstream tasks.  $H$  is the dimension of the hidden representations.  $C_x$  is the number of channels corresponding to a set of sensors  $\mathcal{E}_x = \{E_1, E_2, \dots, E_{C_x}\}$  used to collect the signal.  $F$  is the dimension of the DE feature, which specifically equals 5 in our case, representing the DE feature extracted on five frequency bands, i.e.,  $\delta$  band (1-4Hz),  $\theta$  band (4-8Hz),  $\alpha$  band (8-14Hz),  $\beta$  band (14-31Hz),  $\gamma$  band (31-50Hz).
Unlike vision pre-training, EEG data within a single task is limited, while cross-task EEG data varies significantly in the sensor configuration. Our target is to 1) pre-train a general feature extractor for EEG signals on an irregular dataset  $\mathcal{X}$ , which breaks the constraint that data have to share the same sensor configuration, and 2) encode the geometric information of the sensors properly which reflect the locations of the EEG signals collected.
3.2 Unified Topology
The core idea of addressing the problems is to map different sensors' configurations onto a unified topology and extract a generalized representation that relies on the geometric information of sensors. Recent research in neuroscience reveals that the biological nervous system forms several functional groups in different brain regions [25], and these regions are interconnected with physical, functional, and effective connectivity [26]. These functional regions provide us with a plausible basis for modeling EEG signals which provide a macroscopic view of brain activities.
Drawing inspiration from the functional organization of the brain, we abstract the scalp as a unified topology that can be segmented into several regions  $\mathcal{R} = \{R_1,R_2,\dots,R_S\}$ . Each region can be represented by a vector  $z\in \mathbb{R}^{D}$ . Consequently, the whole scalp can be represented by the collection of vectors from all regions as  $Z\in \mathbb{R}^{S\times D}$ , where  $S$  denotes the number of regions and  $D$  denotes the dimension of a regional representation. In this paper, we follow the previous findings [19] and divide the scalp into  $S = 17$  regions, as shown in Figure 1 (b). We leave the analysis of the influence of different region settings in the supplementary material.
Hence, the challenge of processing data with various montages can be viewed as mapping data  $x \in \mathcal{X}$  with the varying sensor configuration  $\mathcal{E}_x$  into the representations of the unified topology  $Z = \operatorname{Encoder}(x)$ . These representations in the same shape can thus be used to transfer among different datasets without any burdens. Moreover, because we design a plausible mapping between channels and regions, the geometric information can be naturally incorporated.
Figure 1: (a) Overview of our methods. Circles and diamonds represent original nodes and region-wise nodes. Nodes of the same color belong to the same region of the unified topology. White patterns get masked/re-masked before being fed into the encoder/decoder. (b) Mapping the original EEG data into a unified topology based on channel locations. (c) Illustration of the multi-dimensional positional encoding. (d) Attention masks. Green means the two tokens interact with attention, and white means no direct interaction.
3.3 MMM: an EEG Pre-training Framework on Top of Unified Topology
We introduce MMM, an EEG pre-training framework based on the unified topology that follows the Masked Auto-Encoder (MAE) [12] training schema, i.e., during pre-training, the model first encodes partially-masked irregular input tokens (DE features) to a unified representation with an encoder and then aims to reconstruct the masked tokens from the unified representation with a decoder.
Our designs of MMM to better support EEG pre-training are 3-fold: 1) Multi-dimensional positional encoding, which injects the geometric information into tokens; 2) Multi-level channel hierarchy, where extra tokens representing aggregated regions are added to the original set of channels. Tokens and channels interact hierarchically; 3) Multi-stage pre-training, where two different masking strategies, global random masking and regional masking, are used alternatively to learn the hierarchical representations effectively.
3.3.1 Multi-dimensional Positional Encoding
Traditional positional encoding only considers injecting sequence order information. Nevertheless, the injecting order of channels can be quite random and doesn't have actual meanings, and taking channels as a simple sequence will lose the spatial information for the sensors.
As aforementioned, sensors can be conceptualized as placed on a 2-D manifold, so the actual position of each sensor  $E$  can be represented as a 2-D coordinate  $(x,y)$ . We introduce a multi-dimensional positional encoding on this manifold to allow the model to render its learned representations spatially position-aware. The multi-dimensional positional encoding is defined as:
$$
x _ {p} = R (x), \quad y _ {p} = R (y),
$$
$$
\mathsf {E n c o d e} (x _ {p}, 2 i) = \sin \frac {x _ {p}}{\Omega^ {\frac {2 i}{d}}}, \quad \mathsf {E n c o d e} (x _ {p}, 2 i + 1) = \cos \frac {x _ {p}}{\Omega^ {\frac {2 i}{d}}},
$$
$$
\mathsf {E n c o d e} (y _ {p}, 2 i) = \sin \frac {y _ {p}}{\Omega^ {\frac {2 i}{d}}}, \quad \mathsf {E n c o d e} (y _ {p}, 2 i + 1) = \cos \frac {y _ {p}}{\Omega^ {\frac {2 i}{d}}},
$$
$$
\mathrm {P E} _ {x, y, i} = \operatorname {C O N C A T} [ \operatorname {E n c o d e} (x _ {p}, i), \operatorname {E n c o d e} (y _ {p}, i) ].
$$
Here,  $(x,y)$  is a 2-D coordinate from the scalp,  $d$  is the number of sensors, and  $d$  pairs of  $(x,y)$  corresponding to the grid of the 10-10 system form a space. Function  $R$  projects  $x,y$  to their rank
(a)Origin
(b)Random
(c)Region-wise
Figure 2: Examples of random mask and region-wise mask. For the region-wise mask strategy, channels in the same group will be all visible or all masked.
value in its original space. The  $i$  is the index of feature dimension in the encoding, and  $\Omega$  is a large constant (normally set to 10000). The final encoding is generated by concatenating the encoding of the  $x$  and  $y$  coordinates.
The multi-dimensional positional encoding is used in two places. Firstly, it is added after the DE features from different channels are processed through a linear layer, thereby injecting geometric information into the data. Secondly, the encoding is re-applied after obtaining the representations. Position encoding is performed on the re-masked representations to inform the decoder how to distinguish re-masked tokens.
3.3.2 Multi-level Channel Hierarchy
Overall, we employ a transformer-based encoder-decoder as the backbone model of MMM under the masked auto-encoder schema. The encoder-decoder processes representations of tokens with multi-head attention for several rounds. Besides the origin nodes from existing channels, a set of region-wise tokens are accompanied to model the regions of the unified topology.
Encoder Our encoder comprises a stack of  $N$  Transformer blocks. Under the pretext of reconstruction, we set a fraction of all tokens to zero and then attempt to reconstruct them with the remaining tokens. We introduce region-wise tokens  $s \in \mathbb{R}^{17 \times D}$  and append them to the end of the sequence. This process allows extracting local information from the original channels through the attention mechanism, as shown in Figure 1 (a). In order to tackle potential redundancy that could arise during the interaction of region-wise tokens with nodes from other regions, we assign a fixed attention mask to the encoders/decoders. This ensures precise messages pass between the region-wise nodes and their corresponding origin nodes, as well as amongst region-wise nodes and origin nodes. An example of an attention mask can be seen in Figure 1 (d). The encoder maps a sequence of tokens  $x \in \mathbb{R}^{(C + 17) \times F}$  to  $Z \in \mathbb{R}^{(C + 17) \times D}$ , where the last 17 tokens form the hidden representation.
Decoder Our decoder is symmetrical to encoder with  $N$  Transformer blocks. Tokens from the origin EEG channels are re-masked, and the decoder can only access information contained in the region-wise nodes. The bottleneck design forces the encoder to model EEG well with the unified topology.
3.3.3 Multi-stage Pre-training
MMM reconstructs the inputs and predicts the DE value of EEG channels, which consists of the previous  $C$  decoder's output tokens. We use the mean square error (MSE) loss to measure the difference between the predicted DE value and the ground truth.
As a new hierarchy is introduced and only the information from region-wise nodes is used to reconstruct the origin nodes, the unified representation must be capable of both intra-region and inter-region reconstruction. On the one hand, every vector in the representation of the unified topology contains aggregated information about a specific region and serves to reconstruct fine-grained features within the region. On the other hand, the interactions between different regions should also be learned in case information from one region is completely missing. Therefore, to further boost the training of
the representations, we introduce the following two masking strategies accordingly and use them in turns during pre-training, as seen in Figure 2, and we call it multi-stage pre-training.
Global Random Masking Channels are masked randomly and independently, and the number of masked channels is fixed. This provides the chance for the region-wise nodes to see partial information within the region and use it to reconstruct the masked intra-region channels.
Regional Masking Channels within the same region are either fully observable or completely masked, and the number of masked regions is fixed. In this way, we force the masked regions to learn representations from region-wise nodes of other regions since they are no longer directly linked to any observable origin nodes with attention.
3.4 Apply MMM to Downstream Tasks
To apply MMM to downstream tasks, such as the emotion recognition task, we extract the unified representation of the DE feature and feed it into a 2-layer MLP. As a classification task, the cross-entropy loss is calculated to fine-tune the model. It's noteworthy that the unified representation, which is a hidden vector, can serve in several types of downstream applications.
In this section, we first introduce the datasets we used and the overall settings of our experiments. Then, we compared our model with a series of competitive baselines under different settings to demonstrate the strong modeling and transferring ability between different datasets of our proposed MMM framework. We also include ablation studies to investigate the contributions of the designs separately and showcase the effectiveness of the unified representation for EEG modeling.
4.1 Datasets and Experiment Setup
SEED SEED [11] is a discrete EEG emotion dataset elicited by videos. The SEED dataset contains EEG signals of fifteen subjects (eight females and seven males) recorded by 62 EEG sensors while watching fifteen Chinese film clips for three types of emotions. For each subject, the EEG signal recording process is repeated in three different periods corresponding to three sessions. Each session includes the EEG signals while watching the fifteen film clips corresponding to fifteen trials.
SEED-IV SEED-IV dataset [27] is another version of the SEED dataset. It comprises EEG data of fifteen subjects (eight females and seven males, not overlapping with SEED) recorded in 62 channels. Each subject has three sessions, and each session contains 24 trials.
SEED-Union We combine the SEED and SEED-IV datasets to form the SEED-Union dataset, which is only used for pre-training. Its train set consists of the training samples in SEED and SEED-IV and so as the test set. There is no label in SEED-Union since it's only for the reconstruction pre-training task.
Lite series We evenly select 32 channels from the 62 channels in SEED and SEED-IV datasets to form the SEED-Lite and SEED-IV-Lite datasets. The selected 32 channels are evenly distributed on the scalp. Then, we combine the SEED-Lite and SEED-IV-Lite datasets to form the SEED-Union-Lite dataset, which is also only used for pre-training.
TUEG TUEG [28] dataset consists of clinical recordings using a mostly conventional recording configuration (monopolar electrodes in a 10-20 configuration, sample rate varying from  $250\mathrm{Hz}$  to  $1024\mathrm{Hz}$ ) of 14,987 people for 27,063 hours in total. The subjects were  $51\%$  female, and ages ranged from under 1 year old to over 90.
Experiment setup Previous baselines on SEED series datasets all use public DE features, while SEED and SEED-IV provide DE features with different STFT window sizes. It's more reasonable to pre-train the model on the DE features with the same STFT window size. Unless specifically mentioned, the experiments of our model are conducted on the 1-second DE features on both SEED and SEED-IV. Refer to the Appendix B.2 for baseline descriptions and training details.
4.2 Subject-dependent Classification
We compare our models in different settings on SEED and SEED-IV with established baselines, and the results are shown in Table 1.
Table 1: Comparison of our method with baseline methods on SEED and SEED-IV dataset on the subject-dependent classification task. Window size is the length of features to construct the input. For example,  $4\mathrm{s} \times 5$  means every DE feature aggregates 4-second raw EEG signals, and five consecutive DE features are used as the model input. The dagger symbol  $(\dagger)$  represents the exclusion of the worst-performing session out of three for each subject when calculating the final average accuracy.
Method	Pre-trained	Window Size	Accuracy (%) / Std (%)
SEED	SEED-IV	SEED	SEED-IV
STRNN [17]	X	1s×9	-	89.50 / 7.63 †	-
DGCNN [18]	X	1s×1	4s×1	90.40 / 8.49 †	69.88 / 16.29 †
BiDANN [23]	X	1s×1	4s×1	92.38 / 7.04 †	70.29 / 12.63 †
BiHDM [29]	X	1s×1	4s×1	93.12 / 6.06 †	74.35 / 14.09 †
R2G-STNN [15]	X	1s×9	-	93.34 / 5.96 †	-
RGNN [24]	X	1s×1	4s×1	94.24 / 5.95 †	79.37 / 10.54 †
MD-AGCN [22]	X	1s×5	4s×5	94.81 / 4.52	87.63 / 5.77
MV-SSTMA [4]	✓	1s×10	4s×10	95.32 / 3.05	92.82 / 5.03
V-IAG [30]	X	1s×1	-	95.64 / 5.08 †	-
HDGCN [19]	X	1s×1	-	96.40 / 4.54 †	-
MAE [12]	X	1s×1	4s×1	87.89 / 8.06	78.22 / 14.21
MAE [12]	✓	1s×1	4s×1	91.85 / 7.72	81.09 / 10.96
MMM (ours)	X	1s×1	1s×1	94.59 / 6.49	86.04 / 10.04
MMM (ours)	✓	1s×1	1s×1	95.76 / 5.53	87.92 / 9.24
MMM (ours)	✓	1s×1	1s×1	97.80 / 3.28 †	91.78 / 7.05 †
MMM (ours)	✓	-	4s×1	-	93.62 / 7.74