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

Yiyu Gui,MingZhi Chen,Yuqi Su,Guibo Luo,Yuchao Yang

2024-10-06

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.

Signal Processing,Machine Learning

What problem does this paper attempt to address?

The main problem this paper attempts to address is the difficulty existing EEG classification models face in balancing high accuracy, fast inference speed, and efficient memory usage when handling long sequence signals. Additionally, these models often struggle with different EEG classification tasks, exhibiting poor generalization capabilities. To solve these issues, the paper proposes a new EEG classification network—EEGMamba.

Specifically, EEGMamba addresses the problems through the following aspects:
1. **Bidirectional Mamba Design**: Introduces a bidirectional Mamba module specifically designed for EEG signals to balance high-speed inference, efficient memory usage, and good global perception capabilities.
2. **Multi-task Learning Implementation**: EEGMamba is the first EEG classification network to truly achieve multi-task learning, capable of handling data from various tasks in a single training session.
3. **Spatio-Temporal Adaptive Module (ST-Adaptive)**: Proposes a module that can automatically adapt to EEG signals of different lengths and channel numbers, enabling the model to handle multiple task data in a single training session.
4. **Task-aware MoE**: Designs a task-aware MoE module with general experts that can capture the differences and commonalities in EEG data across different tasks.

Through these innovations, EEGMamba not only improves model performance but also demonstrates excellent performance on four different types of tasks, including epilepsy detection, emotion recognition, sleep staging, and motor imagery. Experimental results show that EEGMamba outperforms existing methods on multiple public datasets and has significant advantages in memory usage and inference speed when processing long sequences.