[论文解读] ECGMamba: Towards Efficient ECG Classification with BiSSM

Yupeng Qiang, Xunde Dong+4 more|arXiv (Cornell University)|Jun 14, 2024

ECG Monitoring and AnalysisMedicine被引用 3

一句话总结

ECGMamba 提出了一种新颖且高效的 ECG 分类模型，采用基于 Mamba 架构的双向状态空间模型（BiSSM），以克服 Transformer 模型在计算复杂度上的二次方问题。该模型在两个公开的 ECG 数据集上实现了最先进性能，同时具备出色的推理效率，展现出准确率与速度之间的良好平衡。

ABSTRACT

Electrocardiogram (ECG) signal analysis represents a pivotal technique in the diagnosis of cardiovascular diseases. Although transformer-based models have made significant progress in ECG classification, they exhibit inefficiencies in the inference phase. The issue is primarily attributable to the secondary computational complexity of Transformer's self-attention mechanism. particularly when processing lengthy sequences. To address this issue, we propose a novel model, ECGMamba, which employs a bidirectional state-space model (BiSSM) to enhance classification efficiency. ECGMamba is based on the innovative Mamba-based block, which incorporates a range of time series modeling techniques to enhance performance while maintaining the efficiency of inference. The experimental results on two publicly available ECG datasets demonstrate that ECGMamba effectively balances the effectiveness and efficiency of classification, achieving competitive performance. This study not only contributes to the body of knowledge in the field of ECG classification but also provides a new research path for efficient and accurate ECG signal analysis. This is of guiding significance for the development of diagnostic models for cardiovascular diseases.

研究动机与目标

• 为解决基于 Transformer 的模型在 ECG 分类中因自注意力机制的二次方复杂度导致的高推理成本问题。
• 探索双向状态空间模型（BiSSM）在高效且有效的 ECG 信号建模中的潜力。
• 开发一种轻量化但强大的架构，在不牺牲推理速度的前提下，保持对长 ECG 序列的高性能。
• 为临床 ECG 诊断提供一种新型、高效的深度学习框架，具备更高的可扩展性和计算效率。

提出的方法

• ECGMamba 以基于 Mamba 的模块作为核心组件，利用其选择性扫描机制实现对长序列的高效建模。
• 模型集成残差连接和层归一化，以提升训练稳定性与梯度流动。
• 在每个 Mamba 模块后应用前馈网络（FFN），以增强非线性表示学习能力。
• 采用基于卷积层的 ECG 编码器进行初始特征提取，随后输入至 Mamba 模块。
• 架构采用双向设计（BiSSM），以捕捉 ECG 序列中的过去与未来上下文信息。
• 模型采用端到端训练，使用交叉熵损失函数，并通过标准深度学习训练协议进行优化。

实验结果

研究问题

• RQ1 基于 BiSSM 的架构是否能在保持低推理成本的同时，优于 Transformer 在 ECG 分类中的表现？
• RQ2 将 Mamba 模块与残差连接及归一化机制结合，对 ECG 分类性能有何影响？
• RQ3 在 ECG 分类中，为实现准确率与效率的平衡，Mamba 层的最优数量是多少？
• RQ4 各个组件（ECG 编码器、层归一化、FFN）对整体模型性能的贡献如何？

主要发现

• 在 CPSC2018 数据集上，ECGMamba 的 AUC 达到 0.9526，F1 分数为 0.7909，优于基线模型。
• 在 PTB-XL 数据集上，ECGMamba 的 AUC 达到 0.9167，F1 分数为 0.7528，展现出强大的泛化能力。
• 消融实验证实，ECG 编码器、层归一化和 FFN 均显著提升了模型性能。
• 使用 8 个 Mamba 层时达到最佳性能，CPSC2018 数据集上的 AUC 为 0.9539，PTB-XL 数据集上的准确率为 0.8881。
• 由于 Mamba 机制具有线性复杂度，模型保持了高效的推理速度，而自注意力机制的复杂度为二次方。
• 结果表明，基于 BiSSM 的模型是 Transformer 在 ECG 分类中的可行且高效的替代方案。