[Paper Review] 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 AnalysisMedicine3 citations

TL;DR

MambaCapsule proposes a novel, explainable deep learning framework for ECG-based cardiac arrhythmia classification by integrating Mamba for long-sequence feature extraction and Capsule networks for hierarchical feature representation and reconstruction. It achieves 99.54% and 99.59% accuracy on MIT-BIH and PTB datasets respectively, while enabling transparent decision-making through signal reconstruction that reveals learned features and their relationships.

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.

Motivation & Objective

• To address the lack of interpretability in deep learning-based ECG arrhythmia diagnosis, which limits clinical trust and reliability.
• To develop a model that not only classifies arrhythmias with high accuracy but also explains its predictions through signal reconstruction.
• To leverage the brain-inspired invariance and dynamic routing of Capsule networks for improved feature representation and relationship modeling in ECG signals.
• To integrate the Mamba architecture for efficient, long-context sequence modeling in ECG data, enhancing feature extraction over long temporal patterns.

Proposed method

• The model uses a Mamba-based encoder to extract long-range, hierarchical representations from raw ECG signals, leveraging state space models for computational efficiency and long-sequence modeling.
• A Capsule network decoder processes the encoded features, outputting capsule vectors that encode both class probabilities and spatial relationships between signal components.
• The model includes a reconstruction head that reconstructs input ECG signals from the predicted capsule states, enabling interpretability by visualizing what features the model focuses on.
• The reconstruction process is used to validate feature learning and explain model decisions by analyzing changes in reconstructed P-wave and R-wave morphology under input perturbations.
• The architecture employs a joint training scheme where the encoder-decoder is trained end-to-end, with reconstruction loss guiding feature learning.
• The routing mechanism in the Capsule network dynamically assigns attention to relevant features, improving robustness and interpretability of predictions.

Experimental results

Research questions

• RQ1 Can a Mamba-based encoder improve long-sequence ECG representation learning compared to CNN, Transformer, or ResNet architectures?
• RQ2 To what extent does the Capsule network’s dynamic routing mechanism enhance feature hierarchy and relationship modeling in ECG signals?
• RQ3 How effectively can signal reconstruction from capsule states reveal the model’s internal knowledge and decision rationale?
• RQ4 Does the integration of Mamba and Capsule networks lead to higher classification accuracy while maintaining or improving model transparency?
• RQ5 How do input perturbations (e.g., shifts) affect the reconstructed ECG signals, and what does this reveal about the model’s robustness and feature understanding?

Key findings

• MambaCapsule achieved a test accuracy of 99.54% on the MIT-BIH arrhythmia dataset and 99.59% on the PTB dataset, demonstrating state-of-the-art performance under standard AAMI evaluation protocols.
• The reconstruction of ECG signals from capsule outputs showed high consistency under input perturbations, indicating that the model learns stable, meaningful features and their interdependencies.
• Reconstructed signals for normal (N) and supraventricular premature beat (S) labels clearly showed the absence of P-wave in S-classified signals, visually explaining the model’s classification rationale.
• The Mamba encoder outperformed CNN, ResNet, and Transformer-based encoders in both accuracy and feature representation, as evidenced by superior F1-scores and reconstruction fidelity.
• The capsule-based decoder produced more interpretable and robust predictions than a standard MLP decoder, despite longer training time, due to its ability to preserve feature relationships.
• The model’s reconstruction capability revealed that it captures not only individual wave morphologies (P, R) but also their relative configurations, supporting its interpretability claims.