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

📄 arXiv: 2407.20254v2 📥 PDF

作者: Yiyu Gui, MingZhi Chen, Yuqi Su, Guibo Luo, Yuchao Yang

分类: eess.SP, cs.LG

发布日期: 2024-07-20 (更新: 2024-10-06)

💡 一句话要点

提出EEGMamba，首个基于双向状态空间模型和混合专家模型的通用脑电多任务分类网络。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 脑电信号 多任务学习 状态空间模型 混合专家模型 癫痫检测 情绪识别 睡眠分期 运动想象

📋 核心要点

  1. 现有脑电分类模型计算复杂度高，且难以泛化到不同信号长度和通道数的多个任务。
  2. EEGMamba通过ST-Adaptive模块、双向Mamba和混合专家模型，实现脑电信号的统一特征提取和多任务学习。
  3. 在八个公开数据集上的实验表明，EEGMamba在癫痫检测、情绪识别等任务中表现优异。

📝 摘要（中文）

近年来，随着深度学习的发展，脑电图(EEG)分类网络取得了一定的进展。基于Transformer的模型在捕获脑电信号的长期依赖性方面表现良好。然而，它们二次方的计算复杂度带来了巨大的计算挑战。此外，大多数脑电分类模型只适用于单任务，难以推广到不同的任务，尤其是在信号长度和通道数量变化的情况下。本文提出了EEGMamba，这是第一个真正实现脑电应用多任务学习的通用脑电分类网络。EEGMamba将时空自适应(ST-Adaptive)模块、双向Mamba和混合专家(MoE)无缝集成到一个统一的框架中。所提出的ST-Adaptive模块通过空间自适应卷积对不同长度和通道数的脑电信号进行统一的特征提取，并结合类令牌来实现时间自适应性。此外，我们设计了一种特别适用于脑电信号的双向Mamba，以进一步提取特征，在处理长脑电信号时平衡了高精度、快速推理速度和高效的内存使用。为了加强跨多个任务的脑电数据处理，我们引入了具有通用专家的任务感知MoE，有效地捕捉了来自不同任务的脑电数据之间的差异和共性。我们在八个公开的脑电数据集上评估了我们的模型，实验结果表明它在四种类型的任务中表现优异：癫痫检测、情绪识别、睡眠阶段分类和运动想象。

🔬 方法详解

问题定义：现有的脑电(EEG)分类模型，特别是基于Transformer的模型，在处理长序列脑电信号时面临计算复杂度高的问题，并且大多数模型是为单任务设计的，难以适应不同长度和通道数的脑电信号，限制了其在多任务场景下的应用。

核心思路：EEGMamba的核心思路是构建一个通用的脑电分类网络，能够同时处理多种任务，并有效提取不同长度和通道数脑电信号的特征。通过结合时空自适应模块、双向Mamba模型和混合专家模型，实现对脑电信号的全面分析和高效处理。

技术框架：EEGMamba的整体架构包含三个主要模块：1) ST-Adaptive模块：用于统一提取不同长度和通道数脑电信号的特征，并实现时间自适应性。2) 双向Mamba模型：用于进一步提取脑电信号的特征，平衡精度、速度和内存使用。3) 任务感知混合专家模型(MoE)：用于处理跨多个任务的脑电数据，捕捉不同任务之间的差异和共性。整个框架通过端到端的方式进行训练。

关键创新：EEGMamba的关键创新在于：1) 首个通用脑电多任务学习框架：能够同时处理多种脑电分类任务。2) 双向Mamba模型：专门为脑电信号设计的，能够有效处理长序列数据，并具有高效的计算性能。3) ST-Adaptive模块：能够处理不同长度和通道数的脑电信号，提高了模型的泛化能力。4) 任务感知MoE：能够有效捕捉不同任务之间的差异和共性，提高多任务学习的性能。

关键设计：ST-Adaptive模块采用空间自适应卷积，根据输入信号的通道数动态调整卷积核的大小。类令牌被添加到输入序列中，以实现时间自适应性。双向Mamba模型通过分别从前向和后向处理脑电信号，捕捉信号中的双向依赖关系。任务感知MoE包含一个通用专家和多个特定任务的专家，通过门控机制动态选择合适的专家进行处理。损失函数采用交叉熵损失，并针对多任务学习进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

EEGMamba在八个公开脑电数据集上进行了评估，涵盖癫痫检测、情绪识别、睡眠阶段分类和运动想象四种任务。实验结果表明，EEGMamba在所有任务上均取得了优异的性能，显著优于现有的单任务模型。具体性能数据和与基线模型的详细对比将在代码发布后公开。

🎯 应用场景

EEGMamba具有广泛的应用前景，可用于开发智能脑电分析系统，例如用于癫痫预测、情绪识别、睡眠质量监测和脑机接口等。该研究有助于推动脑电信号处理技术的发展，为医疗诊断、人机交互等领域提供更高效、更准确的解决方案，并有望促进相关产业的智能化升级。

📄 摘要（原文）

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.