[论文解读] MambaCapsule: Towards Transparent Cardiac Disease Diagnosis with Electrocardiography Using Mamba Capsule Network

Yinlong Xu, Xiaoqiang Liu+7 more|arXiv (Cornell University)|Jul 30, 2024

ECG Monitoring and AnalysisMedicine被引用 3

一句话总结

MambaCapsule 提出了一种新颖的、可解释的深度学习框架，用于基于心电图（ECG）的心律失常分类，通过整合 Mamba 以提取长序列特征，并结合胶囊网络实现分层特征表示与重建。该方法在 MIT-BIH 和 PTB 数据集上分别实现了 99.54% 和 99.59% 的准确率，同时通过信号重建实现透明的决策过程，揭示了模型所学习的特征及其相互关系。

ABSTRACT

Cardiac arrhythmia, a condition characterized by irregular heartbeats, often serves as an early indication of various heart ailments. With the advent of deep learning, numerous innovative models have been introduced for diagnosing arrhythmias using Electrocardiogram (ECG) signals. However, recent studies solely focus on the performance of models, neglecting the interpretation of their results. This leads to a considerable lack of transparency, posing a significant risk in the actual diagnostic process. To solve this problem, this paper introduces MambaCapsule, a deep neural networks for ECG arrhythmias classification, which increases the explainability of the model while enhancing the accuracy.Our model utilizes Mamba for feature extraction and Capsule networks for prediction, providing not only a confidence score but also signal features. Akin to the processing mechanism of human brain, the model learns signal features and their relationship between them by reconstructing ECG signals in the predicted selection. The model evaluation was conducted on MIT-BIH and PTB dataset, following the AAMI standard. MambaCapsule has achieved a total accuracy of 99.54% and 99.59% on the test sets respectively. These results demonstrate the promising performance of under the standard test protocol.

研究动机与目标

• 为解决基于深度学习的心电图心律失常诊断中可解释性不足的问题，该问题限制了临床信任与可靠性。
• 开发一种不仅能够以高准确率分类心律失常，还能通过信号重建解释其预测结果的模型。
• 利用胶囊网络的类脑不变性与动态路由机制，提升心电图信号中特征表示与关系建模的能力。
• 整合 Mamba 架构，实现对心电图数据中长上下文序列的高效建模，增强对长时序模式的特征提取能力。

提出的方法

• 该模型采用基于 Mamba 的编码器，从原始心电图信号中提取长程、分层的表示，利用状态空间模型实现计算效率与长序列建模的平衡。
• 胶囊网络解码器处理编码后的特征，输出胶囊向量，编码类别概率以及信号各组成部分之间的空间关系。
• 模型包含一个重建头，能够从预测的胶囊状态重建输入心电图信号，从而通过可视化模型关注的特征实现可解释性。
• 通过分析输入扰动下重建的 P 波与 R 波形态变化，利用重建过程验证特征学习效果，并解释模型决策。
• 架构采用端到端联合训练策略，编码器-解码器联合训练，重建损失引导特征学习过程。
• 胶囊网络中的路由机制动态分配注意力至相关特征，提升预测的鲁棒性与可解释性。

实验结果

研究问题

• RQ1 与 CNN、Transformer 或 ResNet 架构相比，基于 Mamba 的编码器是否能显著提升长序列心电图表征学习能力？
• RQ2 胶囊网络的动态路由机制在心电图信号中对特征层次结构与关系建模的增强作用有多大？
• RQ3 从胶囊状态重建信号在多大程度上能揭示模型的内部知识与决策逻辑？
• RQ4 Mamba 与胶囊网络的融合是否在保持或提升模型透明性的同时，带来更高的分类准确率？
• RQ5 输入扰动（如时间偏移）如何影响重建的心电图信号？这揭示了模型在鲁棒性与特征理解方面的哪些特性？

主要发现

• MambaCapsule 在 MIT-BIH 心律失常数据集上达到 99.54% 的测试准确率，在 PTB 数据集上达到 99.59%，在标准 AAMI 评估协议下展现出最先进性能。
• 在输入扰动下，从胶囊输出重建的心电图信号表现出高度一致性，表明模型学习到了稳定且有意义的特征及其相互依赖关系。
• 对正常（N）与房性早搏（S）标签的重建信号清晰显示：S 类信号中缺失 P 波，直观解释了模型的分类依据。
• Mamba 编码器在准确率与特征表征方面均优于基于 CNN、ResNet 和 Transformer 的编码器，表现为更高的 F1 分数与更优的重建保真度。
• 尽管训练时间更长，胶囊网络解码器相比标准 MLP 解码器产生了更具可解释性与鲁棒性的预测，因其能有效保持特征间的关系。
• 模型的重建能力表明，它不仅捕捉了单个波形形态（如 P 波、R 波），还捕捉了其相对配置关系，支持了其可解释性主张。