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

TL;DR

EEGMamba is a universal multi-task EEG classification network that combines adaptive feature extraction, bidirectional Mamba, and mixture of experts to handle various EEG tasks efficiently and accurately.

Contribution

This paper introduces EEGMamba, the first universal EEG classification model capable of multi-task learning across different EEG applications with a unified framework.

Findings

Outperforms existing models on eight EEG datasets.

Effective in seizure detection, emotion recognition, sleep staging, and motor imagery.

Balances accuracy, speed, and memory efficiency.

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.