Deep learning optimisation for cardiology: Neural Architecture Search-driven arrhythmia classification with electrocardiograms

Cardiovascular disease is the leading cause of death worldwide. Sudden cardiac death (SCD) accounts for roughly 50% of all cardiac deaths. The electrocardiogram (ECG) is widely used for early diagnosis of cardiac disease. However, the complexity of accurate interpretation limits the ECG’s efficacy. Modern deep learning methods have been applied to assist clinicians in diagnosis. We applied Neural Architecture Search (NAS), an automated machine learning technique, to identify optimal deep learning architectures for classifying cardiac arrhythmias from ECGs.

We applied the Differentiable Architecture Search strategy to an AutoFormer search space to identify optimal self-attention architectures for arrhythmia classification. We trained, validated, and tested the resulting model on the PhysioNet Challenge 2021 dataset (n = 88,253), comprising ECGs across three continents. We performed a hyperparameter optimisation on the NAS output, exploring input patch size, class weighting, and loss function. We evaluated performance using the PhysioNet Challenge metric and the area under the receiver operating characteristic curve (AUROC).

The NAS converged towards minimal architectural configurations (embedding dimension: 384, depth: 4, self-attention heads: 4, MLP ratio: 1) with a validation challenge metric of 0.66 (PhysioNet Challenge 21 Winner: 0.63). The NAS-created network achieved an AUROC of 0.97 and a challenge metric of 0.71 during testing. Normal Sinus Rhythm and Sinus Tachycardia achieved AUROCs of 0.99. Low-QRS Voltage and T-wave abnormality were the worst-performing arrhythmias, with AUROCs of 0.89 and 0.90, respectively.

We interpret that architectural simplicity drives performance in arrhythmia classification. Because SCD is unexpected, prevention strategies in free-living environments require lightweight computational resources suitable for wearable devices. Class imbalance fundamentally limits classification performance for rare arrhythmias such as Low-QRS Voltage and T-wave inversion, irrespective of hyperparameter choices. However, the self-attention mechanism can autonomously abstract clinical representations, simplifying clinical deployment by eliminating the need for an explicit feature-extraction pipeline.

Author summary

Heart diseases are the leading cause of death worldwide, of which around 50% happen suddenly. Analysis of the heart’s electrical signal can help detect heart disease early. However, it is a difficult and time-consuming process, even for experienced clinicians. Artificial intelligence (AI) can assist clinicians in detecting heart disease. We explore a specialised form of artificial intelligence that automatically finds its own optimal structure. To do this, we used a dataset from the general population containing 88,253 heart’s electrical activity records, commonly called electrocardiograms.

We found that simpler AI architectures perform better at identifying heart disease. This finding is contrary to the very prevalent thought “the bigger the better”, making this technology ideal for wearables, considering that almost 50% of the heart disease deaths are sudden. This finding is significant since simpler AI architectures require much less computing power and battery life, with the potential to make heart disease monitoring more accessible and equitable for everyone, especially in low-resource settings. Furthermore, we reduce human bias and the trial-and-error that often follows when creating optimised structures.