Mamba-Based Prototypical Contrastive Learning With Augmented Feature Separation for Common and Rare Arrhythmia Classification.

Abstract

Early diagnosis of arrhythmia, a common cardiovascular condition, is crucial for improving prognosis. Electrocardiogram (ECG) is widely used as a non-invasive diagnostic tool. However, Computer-Aided Diagnosis of rare arrhythmias faces significant challenges due to the severe scarcity of samples for these rare disease classes. To tackle this, we propose a Mamba-based Prototypical Contrastive Learning framework, which can simultaneously identify both common and rare classes under the setting of generalized Few-Shot Learning (FSL). It primarily consists of: (1) the Mamba-based Spatio-Temporal Feature Fusion Network (MST), which integrates spatial features from multi-scale convolutions and temporal dynamics from bidirectional Mamba for ECG modeling; (2) the Prototypical Contrastive Learning framework with Augmented Feature Separation (PCAS), which employs a prototype augmentation strategy with an Augmented Prototype Consistency Loss to optimize prototype representations, and an Separation-Tuned Contrastive Loss to enhance intra-class compactness and inter-class distinctnessy, mitigating the risk of class collapse. Extensive experiments on publicly available datasets PTBXL and Chapman demonstrate the effectiveness of MST-PCAS, achieving superior rare-class recognition accuracies of 79.13% and 50.72%, respectively, for ECG arrhythmia classification.