Enhanced ECG-Mamba2: Bidirectional State Space Model for ECG Classification
This repository contains an enhanced version of the ECG-Mamba model for 12-lead ECG arrhythmia classification. The Enhanced ECG-Mamba2 builds upon the original models_mamba_ecg.py implementation with significant architectural improvements.
Overview
Enhanced ECG-Mamba2 is a deep learning model for ECG classification that combines:
Original CNN feature extraction from models_mamba_ecg.py
Mamba-2 (State Space Duality) - 2-8x faster than the original Mamba/VisionMamba
Bidirectional scanning for better temporal context
Multi-branch architecture for lead-specific processing
Transformer attention for capturing short-term anomalies
Key Improvements over Original Implementation
Feature
Original (models_mamba_ecg.py)
Enhanced (Enhanced_ECG_Mamba_Test.ipynb)
State Space Model
VisionMamba (Mamba-1 based)
Mamba-2 (State Space Duality)
Scanning Direction
Unidirectional
Bidirectional (Forward + Backward)
Lead Processing
Single pathway
Multi-branch (4 lead groups)
Attention
None
Transformer attention layer
Training
Standard
Adversarial + Frequency Masking
Explainability
None
MambaLRP
Architecture
1. CNN Feature Extraction (Original)
The CNN layers from the original implementation are preserved:
Global average pooling followed by a linear classifier.
Training Features
Adversarial Training
FGSM-style perturbations are applied during training to improve model robustness.
Frequency Masking Augmentation
Random frequency bands are masked in the FFT domain to make the model robust to noise and artifacts.
Explainability: MambaLRP
MambaLRP (Layer-wise Relevance Propagation) provides interpretability by highlighting which parts of the ECG signal contribute most to the model's predictions.
Model Parameters
Total Parameters: ~29.3M
Embedding Dimension: 384
Number of Mamba-2 Layers: 4
Number of Attention Heads: 4
Dataset
The model is designed for the PhysioNet Challenge 2021 dataset with 5 arrhythmia classes:
1@software{enhanced_ecg_mamba2,
2 title={Enhanced ECG-Mamba2: Bidirectional State Space Model for ECG Classification},
3 year={2024},
4 note={Improvements over models_mamba_ecg.py with Mamba-2, bidirectional scanning, multi-branch architecture, and attention}
5}
License
This project is licensed under the MIT License - see the LICENSE file for details.
