MobileViTv2 with a normalised cosine-similarity classification head and synthetic data balancing for ECG arrhythmia classification on the MIT-BIH Arrhythmia Database

Technical Summary: MobileViTv2 with Normalised Cosine-Similarity Head for ECG Arrhythmia Classification

Problem Statement
The study addresses three persistent limitations in deep learning-based ECG arrhythmia classification: strong class imbalance in public databases (biasing models toward majority classes), the opacity of deep neural networks, and inconsistent evaluation protocols. Specifically, the authors aim to classify heartbeats from the MIT-BIH Arrhythmia Database into the five standard AAMI classes while mitigating imbalance and improving model efficiency for potential mobile deployment. The study explicitly notes that previous results often lack comparability due to the use of patient-independent versus beat-level evaluation splits.

Methodology
The proposed approach integrates data balancing, time-frequency representation, and a lightweight hybrid architecture with a modified classification head.

• Data Pre-processing: Beats were segmented around the R-peak (300 samples) and converted into 224×224 images using the Continuous Wavelet Transform (CWT) with a Gaussian-8 mother wavelet.
• Synthetic Data Balancing: To address class imbalance, the training set was balanced to 4,000 beats per class.
  • For classes with abundant data, a Wasserstein GAN with gradient penalty (WGAN-GP) was used.
  • For rare classes (approx. 200 beats), a 1D convolutional autoencoder with latent noise injection was employed.
  • Synthetic beats preserved the noise and artifacts of real signals. Fidelity was assessed via Pearson correlation of mean beats, spectral correlation, and Wasserstein distance.
• Model Architecture: The backbone is MobileViTv2_100, a hybrid convolution–transformer network designed for mobile devices, initialized with ImageNet weights.
• Classification Head (NormHead): The standard linear classifier was replaced by a Normalised Cosine-Similarity Head (NormHead). This head L2-normalizes the feature vector and class weight vectors, computing logits as the cosine of the angle between them. This approach, inspired by face recognition literature, removes dependence on the feature vector's norm.
• Training Strategy: The model was trained for 30 epochs using AdamW, cosine-annealing learning rate scheduling, label smoothing, and class weights. Data augmentation was limited to random horizontal flipping.
• Evaluation Protocol: A stratified beat-level hold-out split (80% training, 20% testing) was used. Crucially, the split was performed at the beat level, not the patient level, meaning beats from the same patient could appear in both training and test sets. The test set contained only real beats (7,712 total).
• Analysis: Grad-CAM was applied to misclassified beats for qualitative error analysis.

Key Contributions

1. NormHead for ECG: The application of a normalised cosine-similarity classification head to ECG beat classification, replacing the standard linear layer to potentially reduce sensitivity to amplitude and energy variations.
2. Hybrid Balancing Strategy: A dual-generator approach (WGAN-GP and Autoencoder) to synthesize data for balancing the MIT-BIH dataset while preserving signal fidelity.
3. Mobile-Optimized Pipeline: The evaluation of MobileViTv2 with CWT inputs for arrhythmia classification, targeting lightweight deployment.
4. Comparative Benchmarking: A rigorous comparison of the proposed model against various architectures (ConvNeXt V2, Swin Transformer V2, MobileNetV2 variants) on the exact same test set.

Results
On the 7,712-beat hold-out test set, the MobileViTv2 with NormHead achieved:

• Accuracy: 0.9793
• Macro F1-Score: 0.9538

Comparison with Baselines:

• The NormHead model outperformed the same backbone with a standard linear head (Accuracy: 0.9745; Macro F1: 0.9418), showing a modest improvement of +0.0048 in accuracy and +0.0120 in macro F1-score.
• The largest gains were observed in minority classes F (Fusion) and SVE (Supraventricular Ectopic).
• Class D (Unclassifiable/Paced) was nearly perfectly separated (F1: 0.9968), while F remained the most difficult class (F1: 0.8812).
• Most errors occurred between Normal (N) and SVE beats, likely due to the narrow QRS complex in SVE beats and the low energy of P-waves in CWT representations.
• In terms of inference time on an NVIDIA A100 GPU, the NormHead model (12.75 ms per beat) was slower than the standard head baseline (1.67 ms), though the measurement protocols differed (single-beat vs. batch average).

Significance and Limitations (Author's Claims)
The authors position the study as a demonstration of intra-patient performance. They explicitly state that the results describe how well the model performs when it has seen beats from the same patients during training.

• Modest Improvement: The paper claims the cosine-similarity head provides a "modest improvement" over standard heads in this specific setting, primarily benefiting minority classes.
• Generalization Caveat: The authors emphasize that because the split was not patient-independent, these results do not constitute evidence of generalization to unseen patients. They note that patient-independent validation (following the de Chazal et al. protocol) is still required.
• Clinical Readiness: The study concludes that before clinical or wearable deployment can be claimed, further work is needed regarding patient-independent splits, cross-database validation, and measuring the full pipeline (including CWT computation) on embedded hardware.
• Synthetic Data: While synthetic beats were highly similar to real beats for most classes, the authors acknowledge that fidelity was lower for classes with very few real samples.

In summary, the paper presents a technically sound, lightweight solution for ECG classification that improves upon standard baselines using a normalized head and synthetic balancing, but it maintains a cautious stance regarding its clinical applicability due to the evaluation protocol used.