EEGMamba: Bidirectional State Space Model with Mixture of Experts for EEG Multi-task Classification

## Abstract

Abstract In recent years, with the development of deep learning, electroencephalogram (EEG) classification networks have achieved certain progress. Transformer-based models can perform well in capturing long-term dependencies in EEG signals. However, their quadratic computational complexity poses a substantial computational challenge. Moreover, most EEG classification models are only suitable for single tasks and struggle with generalization across different tasks, particularly when faced with variations in signal length and channel count. In this paper, we introduce EEGMamba, the first universal EEG classification network to truly implement multi-task learning for EEG applications. EEGMamba seamlessly integrates the Spatio-Temporal-Adaptive (ST-Adaptive) module, bidirectional Mamba, and Mixture of Experts (MoE) into a unified framework. The proposed ST-Adaptive module performs unified feature extraction on EEG signals of different lengths and channel counts through spatial-adaptive convolution and incorporates a class token to achieve temporal-adaptability. Moreover, we design a bidirectional Mamba particularly suitable for EEG signals for further feature extraction, balancing high accuracy, fast inference speed, and efficient memory-usage in processing long EEG signals. To enhance the processing of EEG data across multiple tasks, we introduce task-aware MoE with a universal expert, effectively capturing both differences and commonalities among EEG data from different tasks. We evaluate our model on eight publicly available EEG datasets, and the experimental results demonstrate its superior performance in four types of tasks: seizure detection, emotion recognition, sleep stage classification, and motor imagery. The code is set to be released soon.

## 1 Introduction

Electroencephalogram (EEG) is a technique of recording brain activity using electrophysiological indicators, which captures the electrical wave changes during brain activity. EEG can be utilized to detect various human physiological activities such as seizure detection, emotion recognition, motor imagery, sleep stage classification, and other physiological related task (Shoeibi et al., ; Jafari et al., ; Altaheri et al., ; Sri et al., ).

In recent years, with the development of deep learning, EEG classification models based on deep learning have been widely used (Chen et al., ). Among them, models based on Convolutional Neural Networks (CNNs) and Transformers are the most representative, each with their own strengths and weaknesses. CNN-based EEG classification networks have the advantage of faster training and inference speeds, and they perform well on short EEG signals. However, due to the lack of global sequence modeling ability, their performance on long EEG signals cannot be guaranteed (Sakhavi et al., ; Thuwajit et al., ; Schirrmeister et al., ). In contrast, Transformer-based EEG classification networks have good capability of global sequence modeling, achieving excellent performance on both short and long EEG signals. Nevertheless, as the length of the EEG signals increases, the computational complexity of the model increases quadratically, significantly raising the training and inference costs (Dai et al., ; Xie et al., ; Wang et al., ).

Figure: Figure 1: Our proposed EEGMamba can simultaneously process EEG signals from multiple tasks including epilepsy detection, sleep stage classification, emotion recognition, and motor imagery. It achieves state-of-the-art (SOTA) performance on the majority of datasets.
Refer to caption: x1.png

Recently, State Space Models (SSM) with selection mechanism and efficient hardware-aware design, such as Mamba (Gu and Dao, ), have shown great potential in long sequence modeling. By utilizing selective state space model, it effectively captures the relationships between tokens in a sequence, addressing the limitation of CNNs in modeling long sequences. Moreover, it exhibits linear computational complexity, which outperforms the quadratic complexity of Transformers and provides a strong backbone network for training EEG classification models on long EEG signals.

Single-task learning (STL) is the most commonly used paradigm in current EEG classification models (O’Shea et al., ; Phan et al., ; Algarni et al., ; Autthasan et al., ), where each task is learned independently given a set of learning tasks. For example, EEGNet (Lawhern et al., ) has been validated on four different tasks but can only address one type of task in a single training session. In contrast, multi-task learning (MTL) trains models by simultaneously learning all tasks and sharing representations across related ones, which enabling the model to learn more robust and universal representations for multiple tasks compared to single-task model (Choo et al., ). Therefore, designing a classification network capable of handling multi-task EEG data simultaneously might be a promising approach.

Few previous studies have employed multi-task classification in EEG, and they all have certain limitations (Prodhan et al., ; Li et al., ). For instance, (Li et al., ) achieved simultaneous classification tasks across four emotion evaluation metrics using the same dataset, but its multi-task classification ability is limited to handling multiple labels within a single dataset. The lack of models capable of performing EEG classification across multiple different datasets may be due to the highly challenging problems.

One of the significant obstacles for multi-task EEG classification is that different EEG data have varying numbers of channels and signal lengths, which makes it difficult for networks to adapt during a single training. For example, MaskSleepNet (Zhu et al., ) can classify EEG signals with different numbers of channels by manually setting the channel parameter, but it uses a fixed-parameter Multi-scale CNN that can only process EEG signals with limited input lengths. While EEG ConvNet (Schirrmeister et al., ) is designed with a structure capable of adapting to arbitrary signal lengths, it still requires manual setting in different trainings. Therefore, enabling the model to adapt to different signal lengths and channel counts represents a significant challenge.

On the other hand, EEG data from different tasks show both differences and commonalities, making it challenging for models without specialized multi-task processing module to capture these relationships, ultimately leading to interference between tasks. Mixture of Experts (MoE) is a deep learning model with sparse gate-controlled architecture, consisting of a group of expert models and a gating network (Jacobs et al., ; Shazeer et al., ; Xue et al., ). The gating network can dynamically select experts to specifically process input data, enabling the network to accurately distinguish and better process multi-task data, thus reducing interference between tasks. Therefore, using MoE to achieve EEG multi-task classification might be a feasible solution.

In general, existing EEG classification models mainly face two challenges. First, these models find it difficult to balance high accuracy, fast inference speed, and efficient memory-usage when dealing with long EEG signals. Second, they often struggle to handle different EEG classification tasks and demonstrate poor generality.

To address the aforementioned two issues, we propose EEGMamba, which utilizes bidirectional Mamba suitable for EEG signals, as well as a Spatio-Temporal-Adaptive (ST-Adaptive) module and task-aware MoE for targeted processing of multi-task EEG classification. Our model enhances Mamba by employing bidirectional modeling to capture the relationships between tokens in a one-dimensional temporal sequence, achieving high accuracy and fast inference speed. Additionally, we propose an ST-Adaptive module that uses spatial-adaptive convolution to process EEG signals of varying channel numbers and a class token to achieve temporal adaptability without any additional processing. To efficiently capture differences and commonalities between EEG data from different tasks, we design a task-aware gating network that accurately directs different EEG task tokens to specific experts for processing, while also employing a universal EEG expert to exploit commonalities among different EEG tasks. In summary, our contributions are as follows:

- •
Bidirectional Mamba Design for EEG Signals. We introduce bidirectional Mamba specifically for EEG signals, achieving the balance between fast inference speed, efficient memory-usage and excellent global perception ability.
- •
First Implementation of Multi-task Learning in EEG application. EEGMamba is the first model to truly implement multi-task learning for EEG classification, enabling a more integrated and effective analysis of complex brain signal data.
- •
ST-Adaptive Module for Flexible EEG Processing. We propose an ST-Adaptive module that can automatically adapt to EEG signals of different lengths and channels, allowing for simultaneous processing in single training session.
- •
Task-aware MoE for EEG Data. We design Task-aware MoE with a universal expert, achieving the capture of both differences and commonalities between EEG data from different tasks.

## 2 Method

EEGMamba primarily consists of the ST-Adaptive module, BiMamba, and task-aware MoE. The ST-Adaptive module processes EEG signals of arbitrary lengths and channel numbers through spatial-adaptive convolution, tokenize layer, and temporal-adaptation based on the class token. The features extracted by the ST-Adaptive module are then processed by multiple BiMamba blocks and task-aware MoE modules. The BiMamba block allows the model to effectively capture long-term dependencies in EEG signals, while the task-aware MoE enables targeted processing of EEG features for different tasks. Finally, a task-aware classifier provides the classification results. The overall model architecture is illustrated in Figure 2.

Figure: Figure 2: Overall structure of EEGMamba. The model consists of ST-Adaptive module, Bidirectional Mamba (BiMamba) blocks and Task-aware MoE modules.
Refer to caption: x2.png

### 2.1 Preliminary Work

Mamba is inspired by continuous state space equations. For continuous input $x(t)\in\mathbb{R}$ in the time domain, the corresponding output $y(t)\in\mathbb{R}$ is determined by the current hidden state $h(t)$ and input $x(t)$ at time $t$ , as shown in Equation (1). Here, $A\in\mathbb{R}^{N\times N}$ is the state matrix, $B\in\mathbb{R}^{N\times 1}$ is related to the system’s hidden state, and $C\in\mathbb{R}^{1\times N}$ is a parameter associated with the input and output.

$$ $h^{\prime}(t)=Ax(t)+Bh(t)$ (1) $$

$$ $y(t)=Ch(t)$ $$

Mamba discretizes the continuous time $t$ into discrete time, transforming the continuous state space equations into discrete state space equations. Specifically, by introducing a time-scale parameter $\Delta$, $A$ and $B$ are transformed into discrete time parameters $\bar{A}$ and $\bar{B}$ respectively. The zero-order hold (ZOH) technique is used as the transformation rule, as shown in Equation (2).

$$ $\bar{A}=exp(\Delta A)$ (2) $$

$$ $\bar{B}={(\Delta A)}^{-1}(exp(\Delta A)-I)\Delta B$ $$

In practice, following the approach of (Gu and Dao, ), we approximate $\bar{B}$ using a first-order Taylor expansion, as show in Equation (3):

$$ $\bar{B}={(\Delta A)}^{-1}(exp(\Delta A)-I)\Delta B\approx\Delta B$ (3) $$

Finally, the discretized form of the continuous state space equation is shown in Equation (4).

$$ $h_{t}=\bar{A}h_{t-1}+\bar{B}x_{t}$ (4) $$

$$ $y_{t}=Ch_{t}$ $$

Based on the mentioned discrete state-space equations, Mamba further introduces data dependency into the model parameters, enabling the model to selectively propagate or forget information based on the sequential input tokens. In addition, it utilizes a parallel scanning algorithm to accelerate the equation solving process.

### 2.2 ST-Adaptive Module

Figure: Figure 3: Overall structure of ST-Adaptive module.
Refer to caption: x3.png

EEG signals from different datasets often have different lengths and channel numbers. To address this issue, we propose a Spatio-Temporal-Adaptive module that transforms input signals of arbitrary lengths and channel numbers into uniform feature dimension, as shown in Figure 3.

To handle the inconsistency in the number of input channels, we introduce a spatial-adaptive convolutional module, which standardizes the data to a fixed number of channels. This module consists of a series of 1D-CNN sub-modules, each designed with a uniform output channel count but adaptable to varying input channels. Through this approach, EEG data with different channel numbers are processed uniformly. Let $x\in\mathbb{R}^{B\times C_{i}\times L_{i}}$ represent the EEG signals, where $C_{i}$ denotes the number of EEG channels for the $i$-th task, and $L_{i}$ is the EEG signal length for the $i$-th task.

$$ $y_{SA}={CNN}_{SA}(x)\in\mathbb{R}^{B\times D\times L_{i}}$ (5) $$

As shown in Equation (5), $y_{SA}$ is the result obtained through spatial-adaptive convolution, where the channel dimension is changed from $C_{i}$ determined by the task $i$ to a unified $D$. Then, $y_{SA}$ is converted into an EEG token sequence through the tokenize layer. In order to better extract features from EEG signals, we design a dual-path structure utilizing a small kernel convolution module ${CNN}_{S}$ and a wide convolutional module ${CNN}_{W}$. Obtain the small kernel feature token sequence $z_{s}$ and the wide kernel feature token sequence $z_{w}$, respectively. Finally, we concatenate them in the time dimension to form the EEG token sequence $T$, as shown in Equation (6).

$$ $z_{s}=\mathcal{T}({CNN}_{s}(y_{SA}))\in\mathbb{R}^{B\times N_{s}\times D}$ (6) $$

$$ $z_{w}=\mathcal{T}\left({CNN}_{w}\left(y_{SA}\right)\right)\in\mathbb{R}^{B \times N_{w}\times D}$ $$

$$ $T=Concat(z_{s},z_{w},dim=1)\in\mathbb{R}^{B\times N\times D}$ $$

Among them, $\mathcal{T}$ represents the transpose operation, $N_{s}$, $N_{w}$, $N$ are the number of EEG small kernel feature tokens, EEG wide kernel feature tokens, and overall EEG tokens, respectively.

Due to the varying lengths of EEG signals, the number of EEG tokens (i.e., the length of the token sequence $T$) obtained from the tokenize layer is inconsistent. To address this issue, we introduce a temporal-adaptive module that incorporates a special class token (Dosovitskiy et al., ) for final classification. Specifically, we concatenate this class token with the previously extracted feature token sequence $t_{s}^{1},t_{s}^{2},...{,t}_{s}^{N_{s}}$ and $t_{w}^{1},t_{w}^{2},...{,t}_{w}^{N_{w}}$ to obtain the token sequence $T$, as shown in Equation (7).

$$ $T=[t_{cls},t_{s}^{1},t_{s}^{2},...{,t}_{s}^{N_{s}},t_{w}^{1},t_{w}^{2},...{,t} _{w}^{N_{w}}]\in\mathbb{R}^{B\times(N+1)\times D}$ (7) $$

Then, the input token sequence $T$ is processed through a network (using bidirectional Mamba blocks in this study) to integrate EEG token sequence information into the class token. This approach prevents the network from developing biases towards certain tokens in the EEG feature token sequence $T$ due to variations in input length, thereby achieving temporal adaptability.

### 2.3 Bidirectional Mamba Block for EEG

Mamba is designed for Natural Language Processing (NLP), with its output at each moment depends only on the current input and hidden state, without consideration for future time steps. Since NLP is primarily a generative autoregressive task that relies on previous information for judgment, Mamba’s single-directional modeling approach is sufficient to complete such tasks. However, EEG classification tasks require simultaneous processing of both preceding and following information, which cannot be learned by single-directional modeling. Therefore, for EEG signals, the original Mamba’s single-directional modeling is insufficient.

To address this issue, we design a bidirectional Mamba for one-dimensional temporal signals, which can model the input bidirectionally and more effectively learn the dependencies between time series tokens. We use the features extracted by the ST-Adaptive module as the input for the first bidirectional Mamba block.

Figure: Algorithm 1 Bidirectional Mamba Block Process

We denote the input of the bidirectional Mamba block as a sequence $T_{k-1}$ and the output as a sequence $T_{k}$. First, $T_{k-1}$ is normalized to $T_{k-1}^{norm}$ by layer normalization. Next, it is mapped by ${Linear}_{X}$ and ${Linear}_{Z}$ to $X_{k-1}$ and $Z_{k-1}$, respectively. Then, $X_{k-1}$ enters parallel forward and backward sequence modeling modules. The forward module includes forward 1D causal convolution ${Conv}_{f}$ and forward SSM module ${SSM}_{f}$. Similarly, the backward module includes backward 1D causal convolution ${Conv}_{b}$ and backward SSM module ${SSM}_{b}$. Then, the results of forward sequence modeling $Y_{k-1}^{f}$ and backward sequence modeling $Y_{k-1}^{b}$ are summed with $Z_{k-1}$ through gating and then projected through a linear layer $Linear_{D}$ to obtain $T_{k-1}^{\prime}$. Finally, the output sequence $T_{k}$ is obtained through residual connection. The detailed process is shown in Algorithm 1.

### 2.4 Task-aware MoE with Universal Expert

#### 2.4.1 Sparsely-activated MoE

A typical Mixture of Experts (MoE) usually consists of several experts, and each expert is typically represented as a Multi-Layer Perceptron (MLP) whose activation is controlled by a gating network (Shazeer et al., ). We define $N_{e}$ as the number of experts, $E_{i}$ as the $i$-th expert, and $G$ as the gating network. For each input EEG token sequence $T$, the output $T^{\ast}$ of MoE can be expressed as Equation (8):

$$ $T^{\ast}=\sum_{i=1}^{N_{e}}{e_{i}(T)\ast E_{i}(T)}$ (8) $$

$$ $e_{i}(T)=SoftMax({{Top}_{k}(G(T),k)})_{i}$ $$

$$ $\operatorname{Top}_{k}(V,k)_{i}=\begin{cases}v_{i},&\text{if }v_{i}\text{ is top }k\text{ value of }V\\ -\infty,&\text{otherwise}\end{cases}$ $$

#### 2.4.2 Task-aware Gating Networks

A gating network calculates gating values based on the input tokens and selects top $k$ experts for activation, typically implemented using a fully connected layer ${Linear}_{Gate}$. However, this can lead to the problem that only a few experts are trained. To avoid this, we adopted the method from (Shazeer et al., ), adding noise to the gating value computation process using a fully connected layer ${Linear}_{Noise}$, which increases randomness and helps in balancing the load among the experts.

Furthermore, we propose a task-aware gating network which helps improve the accuracy of experts in processing different types of EEG tokens. Specifically, we encode the EEG task into task tokens $t_{task}\in\mathbb{R}^{B\times D}$, then concatenate $t_{task}$ with the EEG token sequence $T$ to obtain $T_{cat}$, which is then sent to the gating network. The gating values calculated in this manner incorporate task information, allowing for better assignment of different tasks to different experts. The working process of the task-aware gating network is shown in Equation (9), where $\epsilon$ represents standard Gaussian noise.

$$ $T_{cat}=Concat(T,BroadCast(t_{task}),dim=-1)$ (9) $$

$$ $G(T,t_{task})={Linear}_{Gate}(T_{cat})+\epsilon\ast SoftPlus({Linear}_{Noise}( T_{cat}))$ $$

#### 2.4.3 EEG universal expert

EEG signals from different tasks exhibit both differences and commonalities. Only using different experts to process EEG tokens might overlook the connections between tokens from different tasks. Therefore, we design an EEG universal expert that can process EEG tokens from all different tasks and capture their commonalities. To achieve this function, the universal expert is activated for any inputs and not controlled by the gating network’s output values.

Overall, our MoE module includes both task experts and a universal expert. Task experts can accurately process EEG tokens from different tasks according to gating values, while universal experts can process all EEG tokens. The output of MoE is the weighted sum of these two types of experts. We adopted a weight design scheme similar to (Gou et al., ), as shown in Equation (10). Here, the output weight $\omega$ of the universal expert is determined by the maximum gating value:

$$ $T^{\ast}=\sum_{i=1}^{N_{e}}{e_{i}(T)\ast E_{i}(T)}+\omega\ast E^{u}(T)$ (10) $$

$$ $\omega=1-Max(e(T))$ $$

## 3 Experimental Setup

### 3.1 Dataset

We evaluate the proposed EEGMamba by using eight datasets from four different tasks, including Siena Scalp EEG Database (Detti et al., ), CHB-MIT (Shoeb, ), SleepEDF-20 (Kemp et al., ), SHHS (Quan et al., ), DEAP (Koelstra et al., ), SEED (Duan et al., ), Shu (Ma et al., ), and BCI-IV-2a (Brunner et al., ). Table 1 provides an overview of each dataset. For different tasks, the number of classes, the number of channels and the optimal EEG segment length tend to vary depending on the specific task performed. In the experiment, we predefine the number of channels and classes for each EEG dataset.

**Table 1: Dataset introduction. ‘# Sample’ refers to the total number of samples used for training and testing after preprocessing steps. More details about the datasets can be found in the appendix D.**
| Datasets | Tasks | # Subjects | # Sample | # Classes | # Channels | Rate | Duration |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Siena | Epilepsy detection | 13 from 14 | 78,958 | 2 | 29 | 512 Hz | 4 seconds |
| CHB-MIT | Epilepsy detection | 23 | 111,678 | 2 | 23 | 256 Hz | 4 seconds |
| SleepEDF-20 | Sleep stage classification | 20 | 33,847 | 5 | 1 | 100 Hz | 30 seconds |
| SHHS | Sleep stage classification | 329 from 6441 | 259,799 | 5 | 1 | 125 Hz | 30 seconds |
| DEAP | Emotion recognition | 32 | 1,040 | 2 | 4 | 128 Hz | 60 seconds |
| SEED | Emotion recognition | 15 | 60,912 | 3 | 62 | 200 Hz | 20 seconds |
| Shu | Motor imagery | 25 | 9,579 | 2 | 32 | 250 Hz | 4 seconds |
| BCI-IV-2a | Motor imagery | 9 | 3,948 | 4 | 22 | 250 Hz | 3 seconds |

### 3.2 Implementation Details

Data Preprocessing. We only employ minimal necessary preprocessing. First, we apply a band-pass filter to the EEG signals, retaining components between 0.1 Hz and 50 Hz to remove low-frequency drift and high-frequency noise. Then, we standardize the sampling rate of all EEG signals to 200 Hz. In addition, the public versions of some datasets have undergone some preprocessing. We include a detailed introduction in the Appendix D.

Data Division. In all experiments, including the baseline comparison experiments and ablation experiments, we employ five-fold cross-validation grouped by subjects, so that EEG data from the same subject only appear in one fold. Details of the subject division scheme are provided in the Appendix E.3.

Environments. The experiments are implemented by Python 3.9.18, PyTorch 2.0.1 + CUDA 12.2 on a Linux server with 256 GB memory. All models are trained on Intel(R) Xeon(R) Gold 6342 CPU and a Nvidia A100 GPU 80G.

Our detailed training strategy, hyperparameter settings, metrics, and baselines are provided in Appendix E.4, E.5, E.6, and F.

## 4 Results and Discussion

### 4.1 Single-task EEGMamba Performance Comparison

The single-task EEGMamba experiment aims to demonstrate the effectiveness of the Mamba-based model. In this experiment, we modify the model by removing MoE modules and redundant spatial-adaptive convolution branches, so the single-task EEGMamba only consists of the essential CNN modules and BiMamba modules. We compare the performance of single-task EEGMamba with previous classification models on eight datasets, as shown in Figure 1. Obviously, single-task EEGMamba outperforms the other non Mamba-based models on the majority of datasets.

We also discuss the memory-usage and inference speed of single-task EEGMamba and Transformer-based models, particularly for long sequences. Figure 4(a) and Figure 4(b) show the results for single-channel and multi-channel (here 20 channels) data, respectively. The Transformer-based models in baselines include AttnSleep, EEG Conformer and HCANN. As signal length increases, the memory-usage of Transformer-based models grows quadratically, while single-task EEGMamba grows linearly. In terms of inference speed, Transformer-based models slow down sharply with longer sequences, while the speed of single-task EEGMamba decreases gently. HCANN performs well on single-channel data due to structural modifications on classical Transformer, but it experiences a significant increase in memory-usage and a notable decrease in inference speed when handling multi-channel data. Overall, single-task EEGMamba comprehensively outperforms Transformer-based models in memory-usage and inference speed.

Figure: (a) Single-channel data
Refer to caption: x4.png

To summarize, compared with the previous classification networks, single-task EEGMamba achieves better performance, lower memory-usage and faster inference speed when dealing with long EEG signals, which roundly demonstrates the feasibility of the Mamba-based model on EEG signals.

### 4.2 EEGMamba for Multi-task EEG Classification

Table 2, 3, 4 and 5 show the performance of EEGMamba on different datasets compared with several state-of-the-art (SOTA) baselines. EEGMamba ranks among the top three on seven datasets and achieves the best performance on four datasets.

It is worth noting that all classification networks, except EEGMamba, are trained on a single dataset. Single datasets typically have consistency in data distribution, features, and labels, which allows the model to better adapt and optimize for specific patterns of that dataset, thus improving accuracy. Nevertheless, EEGMamba outperforms existing SOTA models across multiple datasets and showed superior overall performance, demonstrating its strong generalization ability to integrate EEG signals from different tasks.

**Table 2: Performance of EEGMamba compared with baselines on seizure detection task.**
|  |  | Siena | CHB-MIT |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Methods | Multi-task | ACC | AUROC | F1 | ACC | AUROC | F1 |
| EEGNet (Lawhern et al., ) | \usym2715 | 0.9886 $\pm$ 0.0033 | 0.8828 $\pm$ 0.0360 | 0.6905 $\pm$ 0.0185 | 0.9814 $\pm$ 0.0024 | 0.9064 $\pm$ 0.0607 | \ul0.7690 $\pm$ 0.0488 |
| AttnSleep (Eldele et al., ) | \usym2715 | 0.9895 $\pm$ 0.0032 | 0.9066 $\pm$ 0.0196 | 0.6918 $\pm$ 0.0588 | 0.9723 $\pm$ 0.0190 | 0.9048 $\pm$ 0.0465 | 0.7549 $\pm$ 0.0657 |
| EEGConformer (Song et al., ) | \usym2715 | 0.9878 $\pm$ 0.0044 | 0.8744 $\pm$ 0.0377 | 0.6366 $\pm$ 0.0273 | \ul0.9810 $\pm$ 0.0040 | 0.8917 $\pm$ 0.0927 | 0.7507 $\pm$ 0.0648 |
| BIOT (Yang et al., ) | \usym2715 | 0.9897 $\pm$ 0.0043 | 0.8986 $\pm$ 0.0223 | 0.7301 $\pm$ 0.0550 | 0.9678 $\pm$ 0.0284 | 0.8996 $\pm$ 0.0831 | 0.7278 $\pm$ 0.0886 |
| HCANN (Ji et al., ) | \usym2715 | 0.9906 $\pm$ 0.0026 | 0.9283 $\pm$ 0.0208 | 0.6714 $\pm$ 0.1115 | 0.9664 $\pm$ 0.0227 | 0.9110 $\pm$ 0.0572 | 0.7680 $\pm$ 0.1203 |
| Single-task EEGMamba | \usym2715 | 0.9897 $\pm$ 0.0053 | 0.9137 $\pm$ 0.0105 | 0.7106 $\pm$ 0.0326 | 0.9817 $\pm$ 0.0036 | \ul0.9084 $\pm$ 0.0437 | 0.7712 $\pm$ 0.0600 |
| EEGMamba | \usym2713 | 0.9897 $\pm$ 0.0038 | \ul0.9082 $\pm$ 0.0179 | \ul0.7070 $\pm$ 0.0260 | 0.9789 $\pm$ 0.0132 | 0.9126 $\pm$ 0.0492 | 0.7964 $\pm$ 0.0444 |

**Table 3: Performance of EEGMamba compared with baselines on sleep stage classification task.**
|  |  | SleepEDF-20 | SHHS |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Methods | Multi-task | ACC | AUROC | F1 | ACC | AUROC | F1 |
| EEGNet (Lawhern et al., ) | \usym2715 | 0.8165 $\pm$ 0.0254 | 0.9464 $\pm$ 0.0109 | 0.7322 $\pm$ 0.0225 | 0.8174 $\pm$ 0.0173 | 0.9351 $\pm$ 0.0078 | 0.6663 $\pm$ 0.0064 |
| AttnSleep (Eldele et al., ) | \usym2715 | 0.8172 $\pm$ 0.0346 | 0.9383 $\pm$ 0.0123 | 0.7244 $\pm$ 0.0270 | \ul0.8366 $\pm$ 0.0169 | 0.9557 $\pm$ 0.0053 | 0.7270 $\pm$ 0.0153 |
| EEGConformer (Song et al., ) | \usym2715 | 0.7998 $\pm$ 0.0486 | 0.9385 $\pm$ 0.0220 | 0.7118 $\pm$ 0.0392 | 0.8000 $\pm$ 0.0154 | 0.9343 $\pm$ 0.0069 | 0.6543 $\pm$ 0.0085 |
| BIOT (Yang et al., ) | \usym2715 | 0.8226 $\pm$ 0.0387 | 0.9536 $\pm$ 0.0147 | 0.7455 $\pm$ 0.0315 | 0.8331 $\pm$ 0.0152 | 0.9501 $\pm$ 0.0103 | 0.7243 $\pm$ 0.0287 |
| HCANN (Ji et al., ) | \usym2715 | \ul0.8316 $\pm$ 0.0396 | \ul0.9589 $\pm$ 0.0129 | \ul0.7573 $\pm$ 0.0387 | 0.8355 $\pm$ 0.0167 | 0.9581 $\pm$ 0.0077 | 0.7425 $\pm$ 0.0117 |
| Single-task EEGMamba | \usym2715 | 0.8387 $\pm$ 0.0399 | 0.9608 $\pm$ 0.0116 | 0.7681 $\pm$ 0.0359 | 0.8441 $\pm$ 0.0163 | \ul0.9578 $\pm$ 0.0074 | \ul0.7387 $\pm$ 0.0155 |
| EEGMamba | \usym2713 | 0.8486 $\pm$ 0.0276 | 0.9636 $\pm$ 0.0107 | 0.7738 $\pm$ 0.0293 | 0.8478 $\pm$ 0.0177 | 0.9587 $\pm$ 0.0077 | 0.7433 $\pm$ 0.0160 |

**Table 4: Performance of EEGMamba compared with baselines on emotion recognition task.**
|  |  | DEAP | SEED |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Methods | Multi-task | ACC | AUROC | F1 | ACC | AUROC | F1 |
| EEGNet (Lawhern et al., ) | \usym2715 | \ul0.5979 $\pm$ 0.0341 | \ul0.5906 $\pm$ 0.0325 | 0.5624 $\pm$ 0.0214 | 0.5739 $\pm$ 0.0544 | \ul0.7448 $\pm$ 0.0565 | \ul0.5561 $\pm$ 0.0486 |
| AttnSleep (Eldele et al., ) | \usym2715 | 0.5930 $\pm$ 0.0173 | 0.5941 $\pm$ 0.0346 | \ul0.5590 $\pm$ 0.0112 | 0.4808 $\pm$ 0.0232 | 0.6717 $\pm$ 0.0318 | 0.4900 $\pm$ 0.0295 |
| EEGConformer (Song et al., ) | \usym2715 | 0.5905 $\pm$ 0.0351 | 0.5500 $\pm$ 0.0275 | 0.5545 $\pm$ 0.0222 | 0.4861 $\pm$ 0.0172 | 0.6642 $\pm$ 0.0302 | 0.4846 $\pm$ 0.0302 |
| BIOT (Yang et al., ) | \usym2715 | 0.5900 $\pm$ 0.0165 | 0.5703 $\pm$ 0.0283 | 0.5495 $\pm$ 0.0310 | 0.5507 $\pm$ 0.0591 | 0.7363 $\pm$ 0.0666 | 0.5453 $\pm$ 0.0700 |
| HCANN (Ji et al., ) | \usym2715 | 0.5881 $\pm$ 0.0226 | 0.5878 $\pm$ 0.0350 | 0.5083 $\pm$ 0.0484 | 0.5284 $\pm$ 0.0282 | 0.7061 $\pm$ 0.0589 | 0.5101 $\pm$ 0.0361 |
| Single-task EEGMamba | \usym2715 | 0.5985 $\pm$ 0.0247 | 0.5721 $\pm$ 0.0184 | 0.5505 $\pm$ 0.0157 | 0.5779 $\pm$ 0.0584 | 0.7636 $\pm$ 0.0514 | 0.5718 $\pm$ 0.0580 |
| EEGMamba | \usym2713 | 0.5994 $\pm$ 0.0134 | 0.5957 $\pm$ 0.0209 | 0.5628 $\pm$ 0.0262 | \ul0.5646 $\pm$ 0.0366 | 0.7538 $\pm$ 0.0413 | 0.5583 $\pm$ 0.0326 |

**Table 5: Performance of EEGMamba compared with baselines on motor imagery task.**
|  |  | Shu | BCI-IV-2a |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Methods | Multi-task | ACC | AUROC | F1 | ACC | AUROC | F1 |
| EEGNet (Lawhern et al., ) | \usym2715 | 0.5971 $\pm$ 0.0454 | \ul0.6529 $\pm$ 0.0708 | \ul0.6077 $\pm$ 0.0538 | 0.4721 $\pm$ 0.0570 | 0.7449 $\pm$ 0.0591 | 0.4888 $\pm$ 0.0683 |
| AttnSleep (Eldele et al., ) | \usym2715 | \ul0.6105 $\pm$ 0.0454 | 0.6464 $\pm$ 0.0698 | 0.6061 $\pm$ 0.0515 | 0.3807 $\pm$ 0.0384 | 0.6376 $\pm$ 0.0240 | 0.3747 $\pm$ 0.0229 |
| EEGConformer (Song et al., ) | \usym2715 | 0.6014 $\pm$ 0.0392 | 0.6418 $\pm$ 0.0643 | 0.6064 $\pm$ 0.0494 | 0.4228 $\pm$ 0.0421 | 0.6856 $\pm$ 0.0359 | 0.4136 $\pm$ 0.0471 |
| BIOT (Yang et al., ) | \usym2715 | 0.5186 $\pm$ 0.0051 | 0.5183 $\pm$ 0.0050 | 0.5116 $\pm$ 0.0090 | 0.3398 $\pm$ 0.0483 | 0.5970 $\pm$ 0.0561 | 0.2983 $\pm$ 0.0307 |
| HCANN (Ji et al., ) | \usym2715 | 0.5302 $\pm$ 0.0229 | 0.5136 $\pm$ 0.0051 | 0.4131 $\pm$ 0.0530 | 0.3635 $\pm$ 0.0353 | 0.6112 $\pm$ 0.0336 | 0.3258 $\pm$ 0.0422 |
| Single-task EEGMamba | \usym2715 | 0.6169 $\pm$ 0.0467 | 0.6597 $\pm$ 0.0653 | 0.6145 $\pm$ 0.0437 | 0.4596 $\pm$ 0.0547 | 0.7180 $\pm$ 0.0541 | 0.4556 $\pm$ 0.0543 |
| EEGMamba | \usym2713 | 0.6207 $\pm$ 0.0505 | 0.6645 $\pm$ 0.0681 | 0.6183 $\pm$ 0.0525 | \ul0.4231 $\pm$ 0.0522 | \ul0.6873 $\pm$ 0.0542 | \ul0.4156 $\pm$ 0.0545 |

Additionally, the multi-task training of EEGMamba provides significant advantages in terms of convenience. First, it is an end-to-end system that does not require separate pre-training and fine-tuning stages, yet offers stronger generalization ability than the pre-trained model. Furthermore, to obtain the corresponding results presented in Table 2 to 5, EEGMamba only needs to be trained once. In contrast, other classification networks require multiple training sessions, each time involving manual adjustments to data length, channel count, and class numbers, making the process much more cumbersome.

### 4.3 Visualization of Task-aware MoE in Multi-task Classification

We explore the role of designed task-aware MoE in practical applications. Since the EEGMamba model contains eight independent MoE modules, we focus our discussion on the last one MoE module as an example. We calculate the activation probability of each expert for different tasks in the task-aware MoE, as shown in Figure 5. The x-axis represents the index of experts, and the y-axis represents their activation probabilities.

Figure: Figure 5: Activation probabilities of MoE experts in the final layer.
Refer to caption: x7.png

When using task-aware MoE, the model exhibits a clear preference for specific experts based on the given task, with different tasks evidently favoring different experts. Specifically, different tasks tend to activate different experts, while data from the same task show similar expert selection probabilities. For instance, experts 5 and 6 are preferred for epilepsy detection, while experts 0 and 5 are favored for sleep stage classification, demonstrating how task-aware MoE enhances flexibility by dynamically adapting to different tasks. This targeted expert selection not only improves task-specific performance but also maintains efficient processing by bypassing irrelevant experts, thereby reducing unnecessary computational overhead.

### 4.4 Ablation Study

To evaluate the effectiveness of each component in EEGMamba, we conduct ablation experiments on four model variants, including: (i) *Single-directional Mamba*: EEGMamba with Single-directional Mamba; (ii) *EEGMamba w/o MoE*: EEGMamba without the whole MoE module; (iii) *Vanilla MoE*: EEGMamba with the vanilla MoE; (iv) *EEGMamba w/o Task-aware Gating*: EEGMamba without the Task-aware Gating in MoE; (v) *EEGMamba w/o Universal Expert*: EEGMamba without the Universal Expert in MoE.

Figure 6 presents a comparison of ablation experiments on eight datasets across four tasks. EEGMamba outperforms other variants on all metrics for all tasks, demonstrating the contribution of each component in our framework. In comparison to the full EEGMamba, the performance of *Single-directional Mamba* shows a significant decline, emphasizing the importance of employing bidirectional Mamba for EEG classification task modeling. Moreover, the performance decline of *EEGMamba w/o MoE* indicates that MoE plays a role in learning the distinctions between different tasks in multi-task classification. In most tasks, the performance of *EEGMamba w/o Task-aware Gating* and *EEGMamba w/o Universal Expert* is similar but slightly lower than the full EEGMamba.

Figure: Figure 6: Results of the ablation study on different datasets.
Refer to caption: x8.png

## 5 Conclusion

In this paper, we propose EEGMamba, the first model that truly implements multi-task learning for EEG applications. EEGMamba integrates a Spatio-Temporal-Adaptive module to adaptively extract features of EEG data with different lengths and channel counts. We introduce bidirectional Mamba to achieve high accuracy and fast inference speed when processing long-term EEG datasets. Moreover, we design a task-aware Mixture of Experts (MoE) and an EEG universal expert, allowing the model to process multiple tasks simultaneously and better learn the commonalities among EEG signals from different tasks. Our experiments across eight publicly available EEG datasets from four tasks demonstrate the superior performance of our proposed model in multi-task classification scenarios. Our work fills the gap in multi-task classification research within EEG applications, paving the way for future development in this field.

## Appendix A Related Works

### A.1 EEG Classification

The development of deep learning has greatly advanced EEG classification tasks. CNNs are a classic type of neural network with mature applications in EEG classification. [Schirrmeister et al., ] proposed a shallow convolutional network with both spatiotemporal convolutional layers to decode task-related information from raw EEG signals. Similarly, [Lawhern et al., ] introduced EEGNet, a classic EEG classification network based on depthwise separable convolution, which has demonstrated stable and robust performance in various EEG classification tasks. Recurrent Neural Networks (RNNs) are proposed to capture temporal dependencies in time-series EEG signals. [Supratak et al., ] used the RNN architecture for sleep stage classification. [Chen et al., ] used CNN and Long Short Term Memory (LSTM) networks for sleep stage classification.

EEG classification networks based on Transformers have also made significant progress. [Eldele et al., ] introduced attention mechanisms into EEG classification networks for classifying sleep stages. [Song et al., ] proposed EEG Conformer, a EEG classification network based on spatio-temporal convolution and Transformers. EEG Conformer effectively extracts local and global features from EEG signals, and it performs well in tasks such as motor imagery and emotion recognition. HCANN [Ji et al., ] combined the multi-head mechanism with CNN to extract complementary representation information from multiple subspaces, making it more suitable for EEG signals. It has achieved state-of-the-art performance on three datasets from different tasks.

In recent years, there has been notable progress in pre-trained EEG classification networks. [Yang et al., ] proposed BIOT, a generic biosignal learning model that employs a tokenization module and was evaluated on several EEG, ECG, and human sensory datasets. [Yi et al., ] proposed a pre-training framework named MMM, which follows the approach of Masked Auto-Encoder (MAE) for pre-training and employs a multi-stage pre-training strategy to enhance the robustness of the representations.

### A.2 State Space Model

A state space model is a mathematical model that represents a physical system as a set of input, output, and state variables related by a first-order differential equation. [Gu et al., ] proposed the Structured State-Space Sequence Model (S4) to model long-term dependencies. [Smith et al., ] introduced a new S5 layer by incorporating Multiple Input Multiple Output (MIMO) SSM and efficient parallel scanning within the S4 layer. [Fu et al., ] designed a new SSM layer, H3, which further narrowed the performance gap between SSM and Transformers. Recently, [Gu and Dao, ] proposed a data-dependent SSM structure and built a universal language model backbone network: Mamba. Its selective mechanism and hardware-aware design allow it to maintain computational efficiency and excellent performance while scaling to billions of parameters.

### A.3 Mixture of Experts

The Mixture of Experts model was first introduced by [Jacobs et al., ], which controls a system composed of different networks called experts through a supervisory program, with each expert responsible for handling a specific subset of training samples. [Shazeer et al., ] introduced the concept of sparsity into MoE and applied it to LSTM models for translation tasks. With the development of large language models, [Fedus et al., ] extensively investigated the stability issues of MoE models during training and fine-tuning processes, and built a MoE model with 16 trillion parameters and 2048 experts. Recently, [Xue et al., ] proposed OpenMOE, which further explores the details of MoE using the power of the open-source community, thereby promoting the development of MoE.

## Appendix B Overall Structure of Single-task EEGMamba

Figure 7 shows the structure of the single-task EEGMamba model. Compared to EEGMamba, the single-task version removes the MoE modules and the redundant spatial-adaptive convolution branches, retaining only one convolution to process the raw EEG signals. The tokenize layer and BiMamba blocks are kept, with support for stacking any number of BiMamba layers. Additionally, the task-aware classifier in the original EEGMamba is replaced with a standard classifier. Overall, single-task EEGMamba is a lightweight Mamba-based model for EEG classification.

Figure: Figure 7: Overall structure of Single-task EEGMamba.
Refer to caption: x9.png

## Appendix C Notaion Table

Table 6 shows the notations used in the main text.

**Table 6: Notations used in EEGMamba.**
| Symbols | Descriptions |
| --- | --- |
| $B\in\mathbb{N}^{+}$ | Batch size |
| $C_{i}\in\mathbb{N}^{+}$ | Numbers of channels in EEG signals |
| $D\in\mathbb{N}^{+}$ | Hidden dimension of the model |
| $L_{i}\in\mathbb{N}^{+}$ | Numbers of data points in EEG signals |
| $x\in\mathbb{R}^{B\times C_{i}\times L_{i}}$ | EEG signals |
| ${CNN}_{SA}$ | Spatial-adaptive convolution module |
| ${CNN}_{S}$ | Small kernel convolution module |
| ${CNN}_{W}$ | Wide kernel convolution module |
| $y_{SA}$ | Features extracted by the spatial-adaptive convolutional module |
| $z_{s}\in\mathbb{R}^{B\times N_{s}\times D}$ | Small kernel feature token sequence |
| $z_{w}\in\mathbb{R}^{B\times N_{w}\times D}$ | Wide kernel feature token sequence |
| $T\in\mathbb{R}^{B\times(N+1)\times D}$ | EEG token sequence |
| $t_{s}^{j}\in\mathbb{R}^{B\times D}$ | Small kernel feature token |
| $t_{w}^{j}\in\mathbb{R}^{B\times D}$ | Wide kernel feature token |
| $t_{cls}\in\mathbb{R}^{B\times D}$ | Class token for EEG classification |
| $N_{s}\in\mathbb{N}^{+}$ | Numbers of small kernel feature tokens |
| $N_{w}\in\mathbb{N}^{+}$ | Numbers of wide kernel feature tokens |
| $N\in\mathbb{N}^{+}$ | Numbers of overall EEG tokens |
| ${Conv}_{f}$ | Forward causal convolution in BiMamba block |
| ${Conv}_{b}$ | Backward causal convolution in BiMamba block |
| ${SSM}_{f}$ | Forward SSM module in BiMamba block |
| ${SSM}_{b}$ | Backward SSM module in Bimamba block |
| $N_{e}$ | Numbers of experts in MoE |
| $E_{i}$ | The $i$-th expert in MoE |
| $E^{u}$ | Universal expert in MoE |
| $G$ | Gating network in MoE |
| $e_{i}$ | Gating score of the $i$-th expert |
| $\omega$ | Output weight of the universal expert |
| $t_{task}\in\mathbb{R}^{B\times D}$ | Task token for task-aware gating network |
| $L_{b}$ | Balance loss for loading balance |
| $L_{z}$ | Router z-loss for training stability |
| $L_{aux}$ | Auxiliary loss for loading balance and training stability |

## Appendix D Dataset

### D.1 Siena Scalp EEG Database

The Siena Scalp EEG Database consists of EEG recordings of 14 patients acquired at the Unit of Neurology and Neurophysiology of the University of Siena. Subjects include 9 males (ages 25-71) and 5 females (ages 20-58). Subjects were monitored with a Video-EEG with a sampling rate of 512 Hz, with electrodes arranged on the basis of the international 10-20 System. Most of the recordings also contain 1 or 2 EKG signals. The data were acquired employing EB Neuro and Natus Quantum LTM amplifiers, and reusable silver/gold cup electrodes. Patients were asked to stay in the bed as much as possible, either asleep or awake. The diagnosis of epilepsy and the classification of seizures according to the criteria of the International League Against Epilepsy were performed by an expert clinician after a careful review of the clinical and electrophysiological data of each patient. In our experiment, we removed non-EEG signals from each EDF record, retaining 29 EEG channels and ensuring that the signals from different subjects maintained the same channel order: Fp1, F3, C3, P3, O1, F7, T3, T5, Fc1, Fc5, Cp1, Cp5, F9, Fz, Cz, Pz, Fp2, F4, C4, P4, O2, F8, T4, T6, Fc2, Fc6, Cp2, Cp6, F10. We discarded the data from Subject 10 due to the lack of some necessary EEG channels. The data records, after channel unification, were segmented into 4-second segments to facilitate classification.

### D.2 CHB-MIT

The CHB-MIT Scalp EEG Database is collected by the Children’s Hospital Boston, which contains 24 cases of 23 patients with intractable seizures. The first 23 cases are from 22 patients (17 females, aged 1.5-19 years; 5 males, aged 3-22 years). For the last case, there is no clear gender or age record. the Children’s Hospital Boston evaluated the potential conditions for surgical intervention in all epilepsy patients after discontinuing medication for a period of time, and monitored the patients for several days. The original EEG record was obtained using 256 Hz sampling rate with 16-bit resolution from electrodes placed according to the international 10-20 EEG electrode positions and nomenclature [Janjarasjitt, ]. Given that the number of available channels varies among different patients, we select 23 common channels and discarded data from less than 23 channels. Due to the varying duration of the original data ranging from tens of minutes to several hours, we have truncated it into 4-second segments for easy classification.

### D.3 SleepEDF-20

SleepEDF-20 includes Polysomnography (PSG) records from each subject for two consecutive days and nights. The recording of subject 13 on the second night was lost due to a failing cassette or laserdisc. Sleep experts use R&K rules [Wolpert, ] to visually determine signal characteristics and label each 30 second period in the dataset as one of eight stages W, N1, N2, N3, N4, REM, MOVEMENT, UNKNOWN. Similar to previous work [Huy et al., ], N3 and N4 were merged into N3. In addition, the stages of "MOVEMENT" and "UNKNOWN" have also been removed. [Eldele et al., ] have preprocessed the raw data, retaining the Fpz-Cz channel with a sampling rate of 100 Hz, and make it publicly available at [https://researchdata.ntu.edu.sg/dataset.xhtml?persistentId=doi:10.21979/N9/MA1AVG](https://researchdata.ntu.edu.sg/dataset.xhtml?persistentId=doi:10.21979/N9/MA1AVG). We use this version.

### D.4 SHHS

Sleep Heart Health Study (SHHS) is a multi-center cohort study on the cardiovascular and other consequences associated with sleep apnea. The research subjects suffer from various diseases, including lung disease, cardiovascular disease, and coronary heart disease. [Eldele et al., ] have preprocessed the raw data, including retaining the C4-A1 channel with a sampling rate of 125 Hz, and make it publicly available at [https://researchdata.ntu.edu.sg/dataset.xhtml?persistentId=doi:10.21979/N9/EAMYFO](https://researchdata.ntu.edu.sg/dataset.xhtml?persistentId=doi:10.21979/N9/EAMYFO). Additionally, in order to reduce the impact of these diseases, only subjects who are considered to have regular sleep patterns (such as subjects with apnea hypopnea index (AHI) less than 5) are retained, and the evaluation criteria here refer to the research method of [Fonseca et al., ]. Finally, data from 329 participants out of 6441 are retained.

### D.5 DEAP

In the DEAP dataset, movies are used as emotional inducers in experiments. This dataset contains data from over 32 participants aged between 19 and 37, half of whom are females. Participants sit one meter away from the screen. The device records EEG signals at a sampling rate of 512 Hz. 40 selected music video clips were used to trigger emotions. At the end of each video, participants were asked to evaluate their level of arousal, valence, preference, and dominance. The self-assessment scale ranges from 1 to 9. The scores of the subjects are divided into two categories (low or high) based on a stable threshold of 4.5. During the preprocessing process, the EEG signal is downsampled to 128 Hz and a bandpass filter with a cutoff frequency of 4-45 Hz is applied. In this paper, we use the same channel selection as [Khateeb et al., ], which includes four electrodes: Fp1, Fp2, F3, and C4.

### D.6 SEED

The SEED dataset collects EEG data from 15 participants while watching emotional movies. It contains a total of 45 experiments. The EEG data is collected by 62 channels based on the international 10-20 system and a sampling rate of 1000 Hz. During the preprocessing process, the data is downsampled to 200 Hz and subjected to a bandpass filter ranging from 0 to 75 Hz. The extraction of EEG sections was based on the duration of each movie, and we further cut these EEG into segments of 20 seconds in length. Within each subject’s data file, there are 16 arrays, with 15 of these arrays containing 15 preprocessed segments of EEG data from the experiment. The label array includes corresponding emotional labels, where 1 for positive, 2 for negative, and 3 for neutral emotions.

### D.7 Shu

The motor imagery dataset experiment consists of three phases. The first phase (0-2 seconds) is the resting preparation period, during which subjects can rest, perform minor physical activities, and blink. The second phase (2-4 seconds) is the cue phase, where an animation of left or right hand movement appears on the monitor, indicating the upcoming task. The third phase (4-8 seconds) is the MI (Motor Imagery) phase, during which subjects perform the hand movement MI task as prompted, and EEG signals are recorded. We only use 4 seconds of data from the third phase (i.e. MI stage) for classification. Each session consists of 100 trials, with five sessions conducted for each subject every 2 to 3 days, resulting in a total of 500 trials per subject.

### D.8 BCI-IV-2a

The BCI-IV-2a dataset includes EEG signals obtained from trials involving 9 subjects. This experiment includes four different motor imagery tasks: left hand, right hand, foot, and tongue. Each participant participated in two training sessions, with six sessions per session. In each run, there were 48 trials, a total of 288 trials (12 trials per MI task, a total of 72 trials per task). A set of 25 Ag/AgCl electrodes were used in the experiment, of which 22 were dedicated to recording EEG signals, while the remaining three electrodes recorded eye movement signals (not used in our experiment). All recorded signals are processed through a bandpass filter of 0.5 to 100 Hz and a 50 Hz notch filter. The sampling frequency is set to 250 Hz. Similar to Shu, the experiment consists of three phases, with the EEG from the third phase being used for classification. This EEG data, which is for motor imagery, has a duration of 3 seconds and a sampling frequency of 75 Hz.

## Appendix E Experimental Related Supplements

### E.1 Load Balance and Model Stability in MoE

Training an MoE typically encounters two issues: (1) Load imbalance: the gating network tends to select only a few experts. (2) Training instability: excessively large gating values for a few experts lead to an unstable training process. To address these issues, we incorporate balance loss $L_{b}$ [Shazeer et al., ] and router z-loss $L_{z}$ [Zoph et al., ] as auxiliary losses for the model to mitigate load imbalance and training instability, as shown in Equation (11), where $B$ represents the batch size.

$$ $L_{b}=\frac{Std(e(T))}{Mean(e(T))}$ (11) $$

$$ $L_{z}=\frac{1}{B}\sum_{i=1}^{B}{(log(exp(T)))}^{2}$ $$

$$ $L_{aux}=L_{b}+L_{z}$ $$

### E.2 Task-aware Classifier

Figure: Figure 8: Overall structure of Task-aware Classifier.
Refer to caption: x10.png

To address the inconsistency in the number of classes, we introduce a task-aware classifier, consisting of sub-modules, each with a single linear layer configured to have a different number of output dimension corresponding to the specific number of classes, as shown in Figure 8. This approach enables uniform processing of EEG data with varying class counts. The number of classes for each dataset is pre-defined, and for data belonging to the same task, the task identifier is passed through the forward pass, ensuring that data from the same task produce outputs with consistent shapes.

Let $t_{cls}\in\mathbb{R}^{B\times D}$ represents the class token output from the final task-aware MoE block. As shown in Equation 12, $logits_{i}$ is the result obtained through task-aware classifier, where the output dimension is changed from the number of classes $K_{i}$ determined by the task $i$.

$$ $logits_{i}=\mathit{Linear}_{i}(t_{cls})\in\mathbb{R}^{B\times K_{i}}$ (12) $$

### E.3 Subject Division in EEGMamba Experiment

Table 7 presents the grouping and combination of subjects in our five-fold cross-validation experiment. The numbers in the table represent subject IDs in the dataset. Generally, ‘1 $\sim$ 5’ indicates five subjects, including subject 1 through subject 5. For the SHHS dataset, only a subset of subjects is used (D.4), and ‘10 - 2021’ refers to all selected subjects within the range of IDs from 10 to 2021, rather than all subjects in that range consecutively.

**Table 7: Division and combination of subjects in different datasets.**
| Group | Epilepsy detection | Sleep stage classification | Emotion recognition | Motor imagery |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Siena | CHB-MIT | SleepEDF-20 | SHHS | DEAP | SEED | Shu | BCI-IV-2a |  |
| 1 | 0, 1, 3 | 1 $\sim$ 5 | 0 $\sim$ 3 | 10 - 2021 | 1 $\sim$ 6 | 1 $\sim$ 3 | 1 $\sim$ 5 | 1, 2 |
| 2 | 5, 6, 7 | 6 $\sim$ 10 | 4 $\sim$ 7 | 1023 - 2956 | 7 $\sim$ 12 | 4 $\sim$ 6 | 6 $\sim$ 10 | 3, 4 |
| 3 | 9, 11, 12 | 11 $\sim$ 15 | 8 $\sim$ 11 | 2983 - 4047 | 13 $\sim$ 18 | 7 $\sim$ 9 | 11 $\sim$ 15 | 5, 6 |
| 4 | 13, 14, 15 | 16 $\sim$ 19 | 12 $\sim$ 15 | 4051 - 4781 | 19 $\sim$ 25 | 10 $\sim$ 12 | 16 $\sim$ 20 | 7, 8 |
| 5 | 16, 17 | 20 $\sim$ 23 | 16 $\sim$ 19 | 4783 - 5789 | 26 $\sim$ 32 | 13 $\sim$ 15 | 21 $\sim$ 25 | 9 |
| Total | 13 | 23 | 20 | 329 | 32 | 15 | 25 | 9 |

### E.4 Training Strategy

Training the EEGMamba model across multiple EEG datasets with varying tasks presents two primary challenges. First, the inconsistency in the number of channels and lengths across different EEG datasets prevents direct mixed-batch training. Second, training different datasets sequentially may lead to the model forgetting knowledge from earlier datasets.

To address these issues, we propose a dynamic sampling training strategy. Specifically, in each training iteration, we randomly select a batch from the same dataset based on the proportion of samples that have not yet participated in the training. This ensures that data within the same batch have consistent channel counts and lengths. Furthermore, as the probability of sampling each dataset is dynamically adjusted based on the amount of untrained data, larger datasets receive more attention at the beginning of training, while smaller datasets are primarily sampled later, effectively avoiding the model’s forgetting of smaller datasets.

### E.5 Parameter Settings

Table 8 shows the important hyperparameters we used in the experiment.

**Table 8: Hyperparameters for EEGMamba.**
| Hyperparameters | EEGMamba | Single-task EEGMamba |
| --- | --- | --- |
| Hidden dimension | 128 | 128 |
| BiMamba layers | 8 | 2 |
| MoE blocks | 8 | None |
| Experts | 8 | None |
| Experts activated each time | 2 | None |
| Batch size | 128 | 128 |
| Learning rate | 2e-4 | 2e-4 |
| Optimizer | Adamw | Adamw |
| Weight decay | 1e-6 | 1e-6 |
| Training epochs | 100 | 100 |

### E.6 Metrics

Accuracy is a fundamental performance metric for classification models, defined as the ratio of correctly classified samples to the total number of samples. It applies to both binary and multi-class tasks.

AUROC is a key metric for evaluating the performance of classification models, summarizing the model’s ability to distinguish between positive and negative classes across various thresholds by calculating the area under the ROC curve. The AUROC value ranges from 0 to 1, with a value closer to 1 indicating better classification performance.

F1 Score is the harmonic mean of precision and recall, particularly useful in scenarios where a balance between these two metrics is desired. Weighted F1 is used for both binary and multi-class classification in this paper, representing a weighted average of the individual F1 scores for each class, where each score is weighted according to the number of samples in that specific class.

## Appendix F Baselines

We consider the following representative models:

(i) EEGNet [Lawhern et al., ] is a classic EEG classification network based on depthwise separable convolution, which has a concise structure and demonstrated stable and robust performance in various EEG classification tasks.

(ii) AttnSleep [Eldele et al., ] is a deep learning model based on the attention mechanism, designed to automatically classify sleep stages by processing polysomnography (PSG) data including EEG.

(iii) EEG Conformer [Song et al., ] utilizes convolution modules and self-attention modules to capture local features and global dependencies in EEG signals respectively, enabling precise analysis of EEG data.

(iv) BIOT [Yang et al., ] is a pre-trained model that can be applied to various biosignals include EEG.

(v) HCANN [Ji et al., ] is a recently proposed EEG classification network featuring a multi-head mechanism that is adaptively modified for EEG signals. It has achieved state-of-the-art (SOTA) performance across three BCI tasks.

We conduct all baseline tests using publicly available pretrained weights and the open-source code. Generally, we use the same training hyperparameters as in Table 7 in the baseline experiments.

## Appendix G Visualization of Features Extracted by Single-task EEGMamba

Figure 9 shows t-distributed stochastic neighbor embedding (t-SNE) plots of features extracted by single-task EEGMamba from different datasets. The plot exhibits distinct distances between features of different classes and small distances within the same class, indicating the successful extraction of features from different classes by single-task EEGMamba. This may indicate its comprehensive performance superiority across different datasets.

Figure: Figure 9: Visualization results of feature extracted by single-task EEGMamba on different datasets.
Refer to caption: x11.png

## Appendix H Visualization of MoE Weights

Figure: Figure 10: The specific structure of an expert.
Refer to caption: x12.png

In Figure 2, each expert is essentially a Multi-Layer Perceptron (MLP) consisting of two linear layers. The detailed structure is shown in Figure 10, where hidden dimension $D=128$. We visualize the expert weight heatmap for the final MoE module of EEGMamba, where Figure 11 shows the weights of the first linear layer and Figure 12 shows those of the second linear layer. Clearly, the weight distributions vary across different experts, demonstrating that they specialize in handling different tasks.

Figure: Figure 11: The first linear layer weight visualization of experts in final MoE module.
Refer to caption: x13.png

Figure: Figure 12: The second linear layer weight visualization of experts in final MoE module.
Refer to caption: x14.png

## Appendix I Detailed Results on Memory-Usage and Inference Speed

Table 9, 10 present the detailed results of memory-usage and inference speed, where OOM indicates out of memory.

**Table 9: Detailed results on Memory-Usage and Inference Speed with single-channel data.**
| Sequence Length | 2000 | 3000 | 5000 | 10000 | 20000 | 40000 |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
|  | Memory-Usage | 646 MiB | 724 MiB | 834 MiB | 1096 MiB | 1608 MiB | 2648 MiB |
| EEGNet | Inference Speed | 427.20 iter/s | 289.08 iter/s | 174.92 iter/s | 88.11 iter/s | 44.08 iter/s | 21.95 iter/s |
|  | Memory-Usage | 3518 MiB | 6670 MiB | 16534 MiB | 61012 MiB | OOM | OOM |
| AttnSleep | Inference Speed | 133.97 iter/s | 81.98 iter/s | 42.65 iter/s | 10.87 iter/s | OOM | OOM |
|  | Memory-Usage | 3748 MiB | 6958 MiB | 16384 MiB | 62702 MiB | OOM | OOM |
| EEG Conformer | Inference Speed | 104.92 iter/s | 56.14 iter/s | 22.28 iter/s | 5.89 iter/s | OOM | OOM |
|  | Memory-Usage | 2340 MiB | 3318 MiB | 3936 MiB | 9868 MiB | 21108 MiB | 49724 MiB |
| HCANN | Inference Speed | 92.18 iter/s | 62.90 iter/s | 37.84 iter/s | 17.83 iter/s | 8.21 iter/s | 3.87 iter/s |
|  | Memory-Usage | 2864 MiB | 3936 MiB | 6202 MiB | 11600 MiB | 22938 MiB | 45174 MiB |
| Single-task<br>EEGMamba | Inference Speed | 101.43 iter/s | 81.78 iter/s | 46.49 iter/s | 21.21 iter/s | 10.15 iter/s | 5.30 iter/s |

**Table 10: Detailed results on Memory-Usage and Inference Speed with multi-channel data.**
| Sequence Length | 2000 | 3000 | 5000 | 10000 | 20000 | 40000 |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
|  | Memory-Usage | 1630 MiB | 2014 MiB | 2804 MiB | 4810 MiB | 8938 MiB | 17026 MiB |
| EEGNet | Inference Speed | 285.19 iter/s | 191.88 iter/s | 115.93 iter/s | 58.21 iter/s | 29.22 iter/s | 14.65 iter/s |
|  | Memory-Usage | 3532 MiB | 6682 MiB | 16554 MiB | 61028 MiB | OOM | OOM |
| AttnSleep | Inference Speed | 140.96 iter/s | 79.58 iter/s | 41.66 iter/s | 10.81 iter/s | OOM | OOM |
|  | Memory-Usage | 6838 MiB | 11590 MiB | 24430 MiB | 63650 MiB | OOM | OOM |
| EEG Conformer | Inference Speed | 63.44 iter/s | 36.58 iter/s | 16.58 iter/s | 4.65 iter/s | OOM | OOM |
|  | Memory-Usage | 27378 MiB | 68034 MiB | OOM | OOM | OOM | OOM |
| HCANN | Inference Speed | 44.07 iter/s | 5.18 iter/s | OOM | OOM | OOM | OOM |
|  | Memory-Usage | 2954 MiB | 4140 MiB | 6410 MiB | 12208 MiB | 23758 MiB | 46800 MiB |
| Single-task<br>EEGMamba | Inference Speed | 97.31 iter/s | 75.74 iter/s | 43.16 iter/s | 19.72 iter/s | 9.49 iter/s | 5.05 iter/s |

## Appendix J Limitations

Although the current experimental results show that EEGMamba can be well applied to EEG multi-task classification, it still has some limitations. On the one hand, this paper only covers four kinds of EEG tasks to verify the performance of EEGMamba, which is only a small part of the tasks that EEG can accomplish. On the other hand, it should be extended to other one-dimensional time signals besides EEG to prove the universality of the model in one-dimensional time signals.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: Could you elaborate on the choice of baseline models? How does EEGMamba’s performance compare with multi-task EEG models specifically designed for generalization across diverse tasks?

How does the ST-Adaptive module capture interactions between channels and temporal features, which are critical in EEG data? Have you considered modeling these dependencies more explicitly?

How would EEGMamba perform on out-of-distribution tasks or tasks not seen during training? Have you tested its ability to generalize to entirely new EEG task types?

Could you clarify how this work builds on or differs from previous applications of SSMs or Mamba models in EEG classification? Adding this context could help readers better understand the specific contributions of EEGMamba.

Can you evaluate this on longer contexts to better get a sense for the necessity and usefulness for Mamba?
soundness: 3
strengths: EEGMamba achieves state-of-the-art performance across multiple EEG tasks, demonstrating robust multi-task capability. Its bidirectional Mamba blocks enable efficient handling of long sequences, avoiding the high memory demands of Transformer models. The flexible ST-Adaptive module supports EEG signals of varied lengths and channels, and the task-aware Mixture of Experts (MoE) enhances task-specific accuracy, reducing interference across tasks. This adaptability and strong generalization across datasets position EEGMamba as a versatile model for EEG classification.
summary: The paper presents EEGMamba, a model tailored for multi-task EEG classification. It aims to overcome the limitations of existing models in terms of computational complexity and generalization across tasks with varying signal lengths and channel counts. EEGMamba integrates three main innovations: the Spatio-Temporal-Adaptive (ST-Adaptive) module for unified feature extraction, Bidirectional Mamba to balance accuracy and computational efficiency, and a Task-aware Mixture of Experts (MoE) to handle the differences and similarities across EEG tasks. Evaluated across eight public datasets and covering seizure detection, emotion recognition, sleep stage classification, and motor imagery, EEGMamba demonstrates strong performance and adaptability.
weaknesses: EEGMamba’s evaluation lacks comparison the newest baseline models (notably LaBRaM from ICLR 2024), which limits the interpretability of its reported gains. The Spatio-Temporal-Adaptive (ST-Adaptive) module, while flexible for varying signal lengths and channel counts, may not adequately capture complex channel-time dependencies crucial in EEG data. Furthermore, training the ST module on a per-task basis could lead to representations that are too specialized, reducing generalizability, particularly for out-of-distribution (OOD) tasks, which were not included in the study. Additionally, the paper lacks a related works section detailing prior applications of SSMs or Mamba in EEG classification, making it difficult to contextualize EEGMamba’s specific advancements in this space.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: In Section 4.2's experimental comparison, was the model trained using all the datasets at once? Could there be interactions between the datasets? Would training the model on each dataset separately improve the performance metrics? Have you conducted any experiments on this? I'm quite interested.
soundness: 3
strengths: ### **Originality**
The originality of EEGMamba lies in its novel approach to EEG classification through the integration of bidirectional Mamba, Spatio-Temporal-Adaptive (ST-Adaptive) modules, and task-aware Mixture of Experts (MoE). The innovative combination of these elements addresses the computational complexity and variability in signal length and channels, which are critical challenges in EEG classification. This creativity, particularly in applying multitask learning to EEG signals, represents a significant advancement in the field.

### **Quality**
The quality of this work is demonstrated through rigorous evaluations on multiple publicly available EEG datasets. EEGMamba's superior performance in seizure detection, emotion recognition, sleep quality, and emotion recovery highlights its robustness and effectiveness. The model's ability to handle long sequences and adapt to different feature extraction tasks while maintaining high accuracy and fast inference speed.

### **Clarity**
The authors provide a detailed description of the EEGMamba architecture and its components. The step-by-step explanation of how the bidirectional Mamba, ST-Adaptive module, and task-aware MoE are integrated and function together contributes to a well-structured and coherent narrative.

### **Significance**
The model's design, which allows for the efficient capture of both task-specific and general features, has the potential to transform how EEG data is processed and analyzed. This can lead to more accurate and comprehensive analyses of complex brain signal data, benefiting various applications such as medical diagnostics and cognitive research.
summary: The paper introduces EEGMamba, a novel EEG classification network designed for multitask learning. It integrates spatiotemporal adaptive (ST-Adaptive) modules, bidirectional Mamba, and a mixture of experts (MoE) approach. This addresses challenges in EEG classification, such as computational complexity and variations in signal length and channels. The model efficiently handles long sequences and adapts to feature extraction while capturing both task-specific and general features. Evaluations on multiple public EEG datasets demonstrate that EEGMamba outperforms existing models in seizure detection, emotion recognition, sleep quality, and emotion recovery.
weaknesses: Overall, the experiments in this paper are comprehensive; however, in Section 4.1, the discussion on single-channel and multi-channel models only compares memory usage and inference speed, without evaluating the impact of multi-channel models on performance metrics. Additionally, the t-SNE visualization lacks layer-by-layer analysis of the model's influence on clustering results, which does not adequately demonstrate the feature extraction capability of each layer. It is recommended to visualize feature clustering by module.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: Plz go and check weakness
soundness: 2
strengths: EEGMamba achieves superior memory efficiency and faster inference speed than traditional transformer-based models, especially on longer EEG sequences, thus proving its practical value. 

The proposed model demonstrates an approach to handling EEG data of varying lengths and channels, incorporating a class token for temporal adaptability and a task-aware MoE to distinguish task-specific features.
summary: The paper "EEGMAMBA: Bidirectional State Space Model with Mixture of Experts for EEG Multi-Task Classification" presents EEGMamba, a multi-task learning framework for EEG classification tasks, addressing challenges related to signal length and channel variability. EEGMamba integrates a Spatio-Temporal-Adaptive (ST-Adaptive) module, bidirectional Mamba blocks, and a task-aware Mixture of Experts (MoE) to enhance adaptability and task-specific processing across diverse EEG datasets. The ST-Adaptive module standardizes data of various lengths and channel numbers, while the bidirectional Mamba captures temporal dependencies in EEG sequences, and the task-aware MoE module selects experts based on task, enhancing classification accuracy and generalization. Tested on eight datasets spanning four task types (seizure detection, emotion recognition, sleep stage classification, and motor imagery), EEGMamba achieves state-of-the-art results, demonstrating superior efficiency, memory usage, and accuracy across tasks.
weaknesses: 1. Limited Contribution: The author claims that EEGmamba is the first universal EEG classification network to effectively implement multi-task learning for EEG applications. However, several established methods are available that can be applied to multi-task EEG learning, as cited in [1][2][3].

2. Inappropriate Comparisons: The choice of baselines for comparison lacks relevance. For instance, using AttnSleep [3] as a baseline for seizure detection/emotion recognition is incongruous, as it is specifically designed for sleep staging. Additionally, the author does not include multi-task learning methods as baselines, instead comparing against smaller models tailored to individual tasks. For fair assessment, each specific task should be compared to the most relevant model designed for that purpose. To evaluate cross-task capabilities, comparisons should involve multi-task learning methods like those in [1][2][3][5][6].

3. Efficiency Evaluation: The author should provide quantitative evidence demonstrating the proposed method's efficiency advantages over previous approaches.



[1] Jiang, W., Zhao, L., & Lu, B. L. Large Brain Model for Learning Generic Representations with Tremendous EEG Data in BCI. In The Twelfth International Conference on Learning Representations.

[2] Chen, Y., Ren, K., Song, K., Wang, Y., Wang, Y., Li, D., & Qiu, L. (2024). EEGFormer: Towards transferable and interpretable large-scale EEG foundation model. arXiv preprint arXiv:2401.10278.

[3] Jiang, W. B., Wang, Y., Lu, B. L., & Li, D. (2024). NeuroLM: A Universal Multi-task Foundation Model for Bridging the Gap between Language and EEG Signals. arXiv preprint arXiv:2409.00101.

[4] Eldele, E., Chen, Z., Liu, C., Wu, M., Kwoh, C. K., Li, X., & Guan, C. (2021). An attention-based deep learning approach for sleep stage classification with single-channel EEG. IEEE Transactions on Neural Systems and Rehabilitation Engineering, 29, 809-818.

[5] Zhang, D., Yuan, Z., Yang, Y., Chen, J., Wang, J., & Li, Y. (2024). Brant: Foundation model for intracranial neural signal. Advances in Neural Information Processing Systems, 36.

[6] Wang, C., Subramaniam, V., Yaari, A. U., Kreiman, G., Katz, B., Cases, I., & Barbu, A. (2023). BrainBERT: Self-supervised representation learning for intracranial recordings. arXiv preprint arXiv:2302.14367.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. Introduction: (line 49-50):

The authors claim that CNNs are unable to handle long EEG signals, citing three papers: (Sakhavi etal.,2018), (Thuwajit etal.,2021) and (Schirrmeister etal.,2017). However, none of these studies provide evidence to support such a conclusion. In fact, Thuwajit et al. (2021) proposed EEGWavenet, which utilizes a multiscale CNN-based spatiotemporal feature extraction module. This module gradually increases its receptive field to match the length of the input EEG signal, indicating that CNNs can handle the global input sequences in certain contexts. While the authors' claim is not entirely substantiated, it raises interesting questions that I would like to see addressed:

a. For EEG classification, what defines a "long" signal in terms of sample size? At what point does a short receptive field cause CNN performance to degrade?

b. Is global sequence modelling truly necessary for long-term EEG signals? From a neuroscience perspective, how much does brain activity from, say, 10 seconds ago, affect the current state? Which types of tasks specifically require such long-term modelling?

2. section 2.2 ST-Adaptive Module:

The authors propose a spatio-temporal-adaptive module that transforms arbitrary EEG inputs into a uniform feature dimension, as depicted in Figure 3. However, I have several concerns:

a. Scalability to New Datasets: Is this approach scalable to new datasets once the model is trained? Given that different Conv1d layers are used to transform varying numbers of EEG channels into a common hidden state in the Spatial-Adaptive Convolution, how flexible is this method for new tasks where the number of channels may differ from the training datasets? Even if the number of channels is the same, the channels themselves vary. It appears that the model learns dataset-specific spatial filters rather than a truly global spatial representation, which may limit its generalizability to new datasets, which is a big issue for a backbone model, since people would like to use it on different tasks later on.

b. Tokenize Layer: In the tokenize layer, two branches of CNN are used: one with a short kernel (15) and another with a long kernel (49). Was this choice based on experimental results? Why are there only two branches, and why were kernel sizes 15 and 49 specifically chosen? Since there’s no discussion of this configuration in the ablation study, it’s unclear whether this is the most optimal setup.

3. section 3.2 Data Division:

In this study, the 5-fold cross-validation experiment was implemented as leave-subject-out, which is not a common approach in the BCI field due to the significant subject variability. This evaluation approach faces two challenges: (1) developing a population model trained on a group of subjects, and (2) addressing subject transfer by evaluating the model on unseen subjects. This raises several concerns:

a. Subject Variability Impact: In tasks with minimal subject variability, such as seizure detection and sleep stage detection, the classification results are high, as shown in Figure 1. However, tasks like motor imagery and other BCI tasks exhibit high subject variability, which severely impacts model performance. This is evident in the low accuracies achieved on the BCI-IV-2a task (44%) and the SEED task (57.2%), which are insufficient for practical BCI applications. A discussion on this performance discrepancy is necessary to demonstrate how the proposed model addresses these challenges and whether it shows superiority in domains with high subject variability.

b. Performance Discrepancy with Benchmark Models: Many benchmark models, such as EEG Conformer (Song et al., 2022), were not designed to handle population transfer. EEG Conformer, for instance, was trained in a subject-specific manner and achieved state-of-the-art performance on the same BCI-IV-2a (78.66%) and SEED (95.30%) datasets. However, in this study, the EEG Conformer’s performance dropped significantly to 35.2% and below 50%, respectively. Could the authors explain this stark performance difference? Is it primarily due to the leave-subject-out evaluation setting? If so, would the proposed EEGMamba model also retain its high performance if evaluated in a subject-specific manner like EEG Conformer?

c. Use of Separate Test Sets: Datasets like BCI-IV-2a include a separate test set specifically intended for evaluation. In this study, it is unclear whether the authors utilized this test set. Since many studies report performance on this designated test set, comparing the classification accuracy reported here with those from other studies may not be straightforward. Clarification on whether the official test set was used, or if an alternative test split was applied, is needed to ensure a fair comparison with prior work.

4. section 4.1 Single-Task EEGMamba Performance Comparison

Line 377: The authors mention the memory and inference time challenges of transformer models when handling long sequences. Similar to my previous concern, what qualifies as a "long" sequence in terms of signal length for transformers to encounter this bottleneck? Is there a real-world application where such long sequences need to be tackled? Many EEG tasks are only a few seconds in length, so it would be helpful to clarify the practical need for handling significantly longer sequences.

5. section 4.2 EEGMamba for EEG Multi-Task Classification

While the idea of training a single model to perform well across multiple datasets is interesting, its practical application is unclear. From Figure 1, the single-task models appear to achieve similar or even better performance on 5 out of the 8 benchmarked datasets (Siena, BCI-IV-2a, Shu, SEED, CHB-MIT). This raises a few questions:

a. Benefit of Multi-Task Training: Is there a demonstrable benefit to multi-task training? Specifically, is there any statistical difference between the performance of the single-task model and the multi-task model on the datasets? It would be helpful to clarify whether multi-task training consistently improves performance or if its benefits are marginal.

6. Lines 454-456:

The authors argue that the model only needs to be trained once. However, the analysis was performed offline on 8 selected public datasets solely. To assess whether this model can be applied in real-world scenarios, the following questions need to be addressed:

a. Cross-Session Issues: In practice, many EEG-based applications require short calibration sessions to adjust for cross-session variability in subject-dependent models. Does the proposed model also require such calibration? If calibration is needed, would this involve further training or fine-tuning of the pre-trained model? In the case of the multi-task model, would calibration require data from other tasks as well, or could it be done independently?

b. Generalization to New Datasets: How well does the model perform on a new dataset that belongs to one of the pre-trained tasks? If the model only needs to be trained once, does this mean it can be directly applied to similar tasks without additional training? For example, how would the model handle BCI-IV-2B, which has only 3 channels but is still a motor imagery task, or another sleep stage classification dataset? If yes, how would the model manage inconsistencies in the number of channels, and what performance can be expected? If not, wouldn’t this imply that researchers would still need to retrain the model, making it not different from other models?

b. Domain-Specific Advantages: If multi-task training is advantageous, is this benefit domain-specific? For example, should a researcher developing a motor imagery decoder expect better results from multi-task training? How can researchers determine whether a single-task or multi-task approach will yield better performance for a specific domain or task?

c. Practicality of Multi-Task Training: In practice, most researchers focus on specific tasks and experiments, collecting data for single tasks. Does this multi-task approach suggest that researchers in the future should also record additional tasks or rely on public datasets to improve performance? Or is there a scenario where it would make sense for a model to simultaneously classify motor imagery, emotion, sleep stages, and seizure events? More guidance on when and why to use multi-task training would be valuable.

7. section 4.4 Ablation Study:

In Figure 6, it appears that the performances of the different configurations, except for the single-directional Mamba, are quite similar. Is there a statistically significant difference between these configurations? It would be helpful to include a discussion on whether the variations in performance are meaningful or simply within the margin of error.

8. Conclusion lines 526-527:

The authors claim that EEGMamba is the first model to truly implement multi-task learning for EEG applications. However, I am curious about how EEGMamba differs from a simpler approach, such as using a shared backbone model as a feature extractor with separate classification heads for different tasks, as done in [1]. Could the authors clarify the key differences between the proposed model and such an approach using MoE?

[1] Xie Y, Wang K, Meng J, Yue J, Meng L, Yi W, Jung TP, Xu M, Ming D. Cross-dataset transfer learning for motor imagery signal classification via multi-task learning and pre-training. J Neural Eng. 2023 Oct 20;20(5). doi: 10.1088/1741-2552/acfe9c. PMID: 37774694.

9. Conclusion line 531:

The authors claim that the proposed model can better learn the commonalities among EEG signals from different tasks. However, what specific commonalities are being referred to? Is there any interpretation or evidence to support this claim? It would be helpful to understand how these commonalities are identified and whether the model offers any insight into them.

Addressing the above questions could help clarify specific weaknesses and improve the overall impact of the study, more specifically:

Question 5: Answering this question would clarify Weakness Point 1, providing more insight into the motivation for multi-task training.

Questions 1, 2(a), and 6: Answering these would address Weakness Points 2 and 3, which would significantly enhance the study's potential impact and score by clarifying practical applications and the model’s performance on motor imagery datasets.

Questions 4 and 7: These would address Weakness Point 4 by supporting the comparisons with a statistical validation.
soundness: 3
strengths: 1. The authors successfully introduce the new Mamba architecture to EEG decoding, achieving strong results across 8 public datasets.

2. Mamba demonstrates memory efficiency and fast inference, making it advantageous for real-world applications.

3. The model effectively addresses the multi-task classification problem, showcasing the feasibility of training a single model for multiple downstream tasks.
summary: The authors introduce EEGMamba, a model designed for multi-task EEG classification. EEGMamba consists of an ST-Adaptive module that learns spatial filters for each task, transforming EEG inputs with varying channel counts into a uniform feature space. The module then tokenizes the data using both small and large kernels to capture short-term and long-term features. These tokens are processed by a BiMamba backbone with task-aware Mixture of Experts (MoE) layers, enabling the model to capture both task-specific and shared features. Finally, each task has a dedicated classification head.

EEGMamba allows for multi-task EEG classification in a single training session. The authors evaluated the model’s performance against five other models across eight public datasets, covering tasks such as epilepsy detection, sleep stage classification, emotion recognition, and motor imagery. The experiments used a 5-fold cross-validation approach, where specific subjects were reserved for the test set in each fold. Results show that EEGMamba outperformed competing models under this evaluation, and it demonstrated efficient memory usage and inference speed, particularly with long-sequence data.
weaknesses: 1. The motivation for multi-task training is somewhat unclear. The authors should clarify why multi-task training is necessary and how it can be beneficial for specific applications.

2. Certain questions remain unanswered regarding the practical use of the proposed model, which limits its potential impact. For instance, under the current evaluation scheme, it is unclear how the model would generalize to new datasets or subjects, particularly when new datasets vary in channel count. Does introducing a new dataset require retraining or additional multi-

task training even if it is single-task? Additionally, what would be the training strategy for developing a subject-dependent model within the multi-task framework if only one task is available from the subject? To address this, the authors could consider testing the model on an additional dataset to evaluate whether the pre-trained model can transfer effectively. If it cannot, are there still advantages to the EEGMamba multi-task approach?

3. The authors should discuss the relatively low classification accuracy on the motor imagery datasets, which is currently too low for practical motor imagery classification. If this is due to the evaluation setting, additional experiments with per-subject models should be conducted to assess performance, and these results should be compared to other models.

4. Statistical tests are needed to confirm whether the observed differences between models or modules are significant. For instance, statistical analysis should be conducted for the results in Figure 1 comparing EEGMamba and Single-task EEGMamba, as well as for the different ablation models in Figure 6.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. Could you please specify how the EEG tasks are encoded into task tokens?
2. I noticed that the DEAP dataset was utilized in this study, but only data from four electrodes were selected. What is the rationale behind this choice? Additionally, regarding the binary classification on the DEAP dataset, does it pertain to valence, arousal, liking, or dominance? Furthermore, in Table 1, the authors provide the optimal segment lengths for all datasets. What references were used to determine these durations? I observed that, unlike most existing works that employ shorter segments of 1s, 2s, or 4s for the DEAP and SEED datasets, this paper utilizes segment lengths of 60s and 20s. What is the reasoning for selecting such unusually long data lengths?
3. This work employs five-fold cross-validation for data partitioning, which does not appear to be a commonly used EEG dataset partitioning method. What is the rationale or basis for this choice?
soundness: 2
strengths: This paper presents a novel multi-task EEG classification model that incorporates the Bidirectional Mamba and MoE modules, offering a structurally innovative approach. Experiments were conducted on eight datasets, covering a wide range of downstream EEG tasks. The paper is clearly written and easy to follow, with a particularly well-explained method section.
summary: To address the issues of quadratic computational complexity in handling long-term dependencies in EEG classification models, and the lack of cross-task generalization as most models are designed for single tasks, this paper proposes the EEGMamba model. The model introduces the ST-Adaptive module to address the problem of varying EEG signal lengths and channel numbers across different datasets. It also proposes the Bidirectional Mamba to improve computational efficiency and incorporates the MoE (Mixture of Experts) module to simultaneously capture both the commonalities and differences in EEG signals. Relevant experiments were conducted on eight datasets across four EEG tasks.
weaknesses: This paper proposes a multi-task EEG classification model aimed at performing classification across multiple datasets (downstream tasks) using a single model. However, the experimental results do not demonstrate a clear advantage over other single-task models, making it difficult to convincingly argue for the benefits and necessity of the multi-task approach. Additionally, the ablation study results for the various modules lack significant and consistent differences, making it challenging to prove the effectiveness of each module. Moreover, the motivation for using the Bidirectional Mamba is insufficiently justified. The ST-Adaptive method, proposed to address the issue of varying EEG channel numbers across datasets, is essentially a module integration, lacking in innovation.
1. In this work, the ST-Adaptive module applies a different one-dimensional convolution for each task, transforming the varying original channel numbers into a fixed number of $D$ channels. However, this approach does not seem to fully achieve the concept of being "adaptive." If a new task emerges with a number of channels outside the predefined range of $C_0$ to $C_N$ in the model, how would this be addressed? This raises concerns about the generalizability and flexibility of the current method in handling unforeseen tasks with different channel configurations.
2. The purpose of using Mamba in this paper is to reduce computational complexity, conserve resources, and improve efficiency. Generally, the ultimate goal of improving efficiency in EEG models is to achieve real-time recognition. Given that EEG signals, like natural language, possess temporal characteristics, theoretically, a unidirectional Mamba would better meet this requirement, as it only requires past data rather than future information. The motivation for employing Mamba, especially Bidirectional Mamba, in this work is not sufficiently clear or logically aligned with this objective.
3. In Section 4.1, the authors present the performance of the single-task EEGMamba and other transformer-based models concerning memory usage and inference speed as the sequence length increases. However, it is not clearly evident from Figure 4 that EEGMamba demonstrates a significant advantage over the other methods. Additionally, while it is acknowledged that memory usage increases and inference speed decreases with longer sequence lengths for both single-channel and multi-channel scenarios, the authors do not specify the actual sequence lengths employed in the current eight EEG tasks. This omission lacks a reference point, making it difficult to ascertain whether EEGMamba exhibits superior performance. Furthermore, as indicated in Appendix I, EEGNet appears to perform better in terms of memory usage and inference speed, while also demonstrating commendable performance across various datasets. This further undermines the effectiveness of the proposed method in this paper.
4. In Section 4.2, the authors present the performance of EEGMamba in multi-task classification. However, I observe that EEGMamba does not demonstrate a significant advantage over the baseline models, and in many datasets, its performance is inferior to that of other single-task models, indicating that the multi-task approach does not facilitate mutual enhancement among tasks. Therefore, I question the necessity of employing a single model to address multiple tasks rather than utilizing several smaller models for different tasks, which might yield better results. The existing findings lack persuasiveness and do not adequately support the motivations for this work or the claims regarding the strong generalization capabilities of the proposed multi-task model.
5. Regarding Figure 5, I observe that, apart from the sleep stage task, where there is a considerable variation in the activation probabilities of different experts, the activation probabilities for the other tasks across the eight experts are generally quite uniform. This uniformity makes it challenging to demonstrate a particular preference for any specific expert. How can the effectiveness and necessity of the MoE approach be substantiated under these circumstances?
6. Figure 6 presents the ablation study results for the various modules of EEGMamba. However, the data indicate that these modules appear to have minimal discernible impact, as the experimental results across the Siena, CHB-MIT, and SHHS datasets show little variation. This raises concerns regarding the ability to substantiate the effectiveness of each module.

Meta-review: The work considers the problem of multitask decoding from EEG (to understand a model that can predict for different tasks without retraining). To so so it introduces a new architecture by using a bidirectional mamba state-space model and a MoE approach that can orient the model prediction based on a class token. Results are reported on 8 public datasets. To cope with the variability of input channels a different set of spatial filters are trained for each dataset.

Notable contributions:
- Mamba demonstrates memory efficiency and fast inference, making it advantageous for real-world applications.
- The model addresses the multi-task classification problem, showcasing how MoE enables prediction across a variety of downstream tasks.

Weaknesses:
- As pointed out by uVYA the motivation for multi-task training is not fully unclear. Effectively, the model cannot generalize to new datasets or subjects, particularly when new datasets vary in channel count.
- Comparison with SoTA including single tasks models are missing. While one can acknowledge the multi-task aspect of the work, as just pointed above taking a new dataset and a new task does effectively require retraining. There is no zero-shot transfer to new tasks here. So from a pure application point of view, if the objective is to have the best BCI model or the best sleep model it is not clear if the multitask approach is reasonable. Also as pointed out by uVYA results on very classical EEG BCI are pretty low. This questions the significance of the work.

While wapT is concerned about the claim of architectural novelty, the current decision is motivated by the 2 concerns above.

As a general comment to this community, it becomes very clear that benchmarks on such data need to be more standardized (test sets, validation strategy, etc.)

Decision: Reject