ECG-NAT: A Self-supervised Neighborhood Attention Transformer for Multi-lead Electrocardiogram Classification

Abstract: Electrocardiogram (ECG) arrhythmia classification remains challenging due to signal variability, noise, limited labeled data, and the difficulty in achieving both accuracy and efficiency in models. While self-supervised learning reduces label dependency, most methods target either global contextual features or local morphological patterns, but rarely implement hierarchical multi-scale feature extraction. ECG signals require architectures that simultaneously capture fine-grained beat-level morphology and broader rhythm-level dependencies with computational efficiency. To overcome this limitation, this paper proposes the Electrocardiogram Neighborhood Attention Transformer (ECG-NAT), a novel self-supervised learning approach tailored for multi-lead ECG classification. Our two-stage approach begins with generative pretraining, using a masked autoencoder to reconstruct partially masked ECG signals across multiple diverse datasets, enabling the model to learn robust, domain-invariant representations from unlabeled data. This is followed by discriminative fine-tuning with a dual-loss function that combines supervised contrastive and cross-entropy losses, aligning representation learning with label prediction. The hierarchical attention mechanism efficiently captures multi-scale temporal features from localized beat morphology to broader rhythm patterns at low computational cost. ECG-NAT achieves robust performance on benchmark datasets, with 88.1\% accuracy using only 1\% labeled data, demonstrating strong efficacy in low-resource settings. The framework combines superior classification performance with computational efficiency, making it practical for real-time ECG diagnosis. The code will be made available upon acceptance at: https://github.com/Mahsagazeran/ECG-NAT.

• The paper introduces ECG-NAT, a hierarchical self-supervised transformer using neighborhood attention for efficient multi-lead ECG classification.
• It combines masked autoencoding with dual-loss fine-tuning to enhance feature learning, yielding state-of-the-art accuracy and AUROC across benchmarks.
• The approach excels in low-resource environments by requiring minimal labeled data, making it ideal for resource-constrained clinical settings.

Self-Supervised Neighborhood Attention Transformers for Robust ECG Classification

Introduction and Motivation

Multi-lead electrocardiogram (ECG) analysis is a cornerstone of cardiovascular disease diagnosis, yet automated ECG classification confronts persistent obstacles: high inter- and intra-patient signal variability, limited access to expert-labeled data, and the computational burden of deploying representative models in clinical workflows. The prevailing deep learning approaches—CNNs, RNNs, and, more recently, transformers—are either biased toward local or global dependencies or demand extensive supervised datasets. Hierarchical modeling capable of capturing both morphological (beat-level) and rhythmic (cycle-level) signal properties, while being label-efficient and computationally tractable, remains underexplored.

The paper "ECG-NAT: A Self-supervised Neighborhood Attention Transformer for Multi-lead Electrocardiogram Classification" (2605.13194) addresses these challenges by proposing ECG-NAT, a transformer architecture based on neighborhood attention, specifically tailored for self-supervised learning and classification in multi-lead ECG analytics. The framework integrates generative masked autoencoding with discriminative contrastive and cross-entropy objectives, establishing a hierarchical multi-scale locality-aware representation learning pipeline.

ECG-NAT Framework and Architectural Advances

Locality in Attention: Neighborhood Attention Transformer

Traditional transformer designs, using global self-attention, scale quadratically with sequence length and can obfuscate diagnostically critical local ECG morphologies due to indiscriminate contextual mixing.

ECG-NAT introduces a hierarchical transformer architecture with neighborhood attention (NA), where each token interacts only with $k$-nearest neighbors. This drastically lowers attention complexity to $O(nk)$, allows progressive receptive field growth with depth, and better preserves physiologically meaningful local structures (e.g., QRS complex, ST segment).

Figure 1: Standard self-attention (SA) versus neighborhood attention (NA); NA restricts computations to the local neighborhood, improving both efficiency and diagnostic focus.

Self-Supervised Pretraining via Masked Autoencoding

ECG-NAT’s pretraining stage uses a masked autoencoder design: multi-lead ECG segments are tokenized, a subset of tokens is randomly masked with Gaussian noise, and the model reconstructs the missing regions using a convolutional decoder. Pretraining spans diverse, unlabeled datasets, producing robust, domain-invariant latent representations that are transferable across institutions and pathologies.

Fine-Tuning With Dual-Loss Optimization

For downstream classification, ECG-NAT employs a discriminative fine-tuning regime combining supervised contrastive loss (structuring the representation space so that intra-class distances are minimized and inter-class distances maximized) and standard cross-entropy loss. This joint objective enables both improved clustering of same-class signals and enhanced class predictivity.

Figure 2: The ECG-NAT pipeline: a) masked autoencoder for robust generic feature acquisition; b) dual-loss supervised fine-tuning; c) local attention (NAT block) for context-specific feature extraction; d) contrastive loss structuring of latent spaces for inter/intra-class discrimination.

Experimental Validation and Comparative Results

Datasets and Evaluation Protocol

ECG-NAT is validated through extensive experiments using the PTB-XL and CPSC2018 multi-lead ECG datasets. Pretraining is performed on Chapman and Ningbo datasets (unlabeled), with downstream fine-tuning on the labeled datasets. Assessment metrics include Accuracy, Macro F1-score, and AUROC, with stratified cross-validation for robustness.

Comparative Performance

Linear Evaluation: Using only the pretrained encoder as a feature extractor (classification layer trained from scratch), ECG-NAT achieves 0.800 accuracy and 0.960 AUROC on PTB-XL, and 0.782 accuracy and 0.965 AUROC on CPSC2018—outperforming all prior self-supervised baselines.

Fine-Tuned Evaluation: With fully trainable encoders, ECG-NAT yields 0.902 accuracy and 0.977 AUROC (PTB-XL), and 0.903 accuracy and 0.986 AUROC (CPSC2018), surpassing prior state-of-the-art such as CRT, ST-MEM, and DLA. This demonstrates both the strength of self-supervised locality-aware pretraining and the discriminative adaptation during fine-tuning.

Low-Resource Settings

A key claim is the model’s strong efficacy in low-resource regimes. Using only 1% of the available labeled data, ECG-NAT attains 88.1% accuracy and 0.880 AUROC on PTB-XL and 0.931 AUROC on CPSC2018, closely matching or exceeding other models that require multiples more labeled data.

Ablation Studies and Representation Analysis

Ablations confirm that both Gaussian masking (vs. zero-masking) and the addition of contrastive loss during fine-tuning are critical to performance. Replacing NAT with ViT or Swin transformer under identical constraints leads to consistent drops in accuracy, F1, and AUROC.

The impact of contrastive objectives on representation geometry is validated by t-SNE visualizations, revealing substantial improvements in intra-class cluster compactness and inter-class separation with supervised contrastive loss.

Figure 3: t-SNE plot on CPSC2018; representations learned without contrastive loss display poor cluster separation, which is rectified by supervised contrastive fine-tuning.

Hyperparameter and Complexity Analysis

Exhaustive parameter sweeps solidify an optimal configuration of kernel size 7 (neighborhood size), batch size 32, and weight decay 4e-4, stable across both considered datasets.

Figure 4: Sensitivity of ECG-NAT accuracy to neighborhood size, batch size, and weight decay during fine-tuning on PTB-XL and CPSC2018.

Complexity assessment demonstrates the practical efficiency of ECG-NAT: 29.31M parameters, 1.10GB memory, and 1.57G FLOPs in the NAT-Tiny configuration—significantly lower than ViT and Swin baselines at comparable accuracy.

Figure 5: Model size (MB); ECG-NAT significantly reduces parameter count relative to standard transformer variants.

Figure 6: Convergence curves for PTB-XL; ECG-NAT attains stable minima within 25 epochs in both training and test loss, indicating swift and stable convergence.

Implications and Future Directions

Practical and Clinical Relevance

ECG-NAT offers an efficient, highly discriminative, and low-label-dependency approach for ECG analytics. These properties directly address deployment in resource-constrained or low-label clinical environments, such as primary healthcare or low-resource hospitals. The architecture’s modularity and computational efficiency further suggest suitability for point-of-care and edge-device integration, facilitating real-time or near real-time diagnostic workflows.

Theoretical Significance and Broader Impact

ECG-NAT’s design exemplifies the advantages of local attention mechanisms, especially where hierarchical multi-scale dependencies are physiologically salient (e.g., biomedical time series). The method demonstrates that locality-aware transformers, when combined with well-structured self-supervised objectives, can close the gap between representation robustness and practical efficiency. The approach is not limited to ECG but is extendable to other health-related time series (EEG, EMG, PPG) and possibly other domains encountering similar challenges.

Open challenges remain regarding interpretability of the attention mechanisms and alignment with clinical reasoning. Future work could exploit structured or low-rank attention variants, explicit motif or domain-informed regularization, and multi-modal extensions (e.g., via text reports or patient metadata).

Conclusion

The ECG-NAT framework advances ECG classification by integrating self-supervised masked autoencoding with hierarchical neighborhood attention and supervised contrastive adaptation. This methodology produces robust, computationally efficient, and highly discriminative representations for multi-lead ECG classification, validated by superior performance across standard benchmarks, notably in low-resource settings. The study substantiates the broader viability of locality-sensitive transformers for time-series biomedical analytics and paves the way for further developments in scalable, efficient, and label-efficient clinical AI.