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

TL;DR

Electroencephalogram classification network EEGMamba is introduced, the first universal EEG classification network to truly implement multi-task learning for EEG applications and introduces task-aware MoE with a universal expert, effectively capturing both differences and commonalities among EEG data from different tasks.

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.

Figures (12)

• 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.
• Figure 2: Overall structure of EEGMamba. The model consists of ST-Adaptive module, Bidirectional Mamba (BiMamba) blocks and Task-aware MoE modules.
• Figure 3: Overall structure of ST-Adaptive module.
• Figure 4: Memory-usage and inference speed of Single-task EEGMamba compared with Transformer-based models. OOM indicates out of memory.
• Figure 5: Activation probabilities of MoE experts in the final layer.