Abstract
Cardiovascular disease remains one of the leading causes of mortality worldwide, underscoring the need for accurate as well as interpretable diagnostic machine learning tools. In this work, we investigate heart disease classification using electrocardiogram (ECG) data from two widely-utilized datasets: The MIT-BIH Arrhythmia and the PTB-XL datasets. We adapt a hierarchical attention network (HAN), originally developed for text classification, into an ECG-based heart-disease classification task. Our adapted HAN incorporates two attention layers that focus on ECG data segments of varying sizes. We conduct a comparative analysis between our adapted HAN and a more sophisticated state-of-the-art architecture, featuring a network with convolution, attention, and transformer layers (CAT-Net). Our empirical evaluation encompasses multiple aspects including test accuracy (quantified by 0-1 loss); model complexity (measured by the number of model parameters); and interpretability (through attention map visualization). Our adapted HAN demonstrates comparable test accuracy with significant reductions in model complexity and enhanced interpretability analysis: For the MIT-BIH dataset, our adapted HAN achieves 98.55\% test accuracy compared to 99.14\% for CAT-Net, while reducing the number of model parameters by a factor of 15.6. For the PTB-XL dataset, our adapted HAN achieves a 19.3-fold reduction in model complexity compared to CAT-Net, with only a 5\% lower test accuracy. From an interpretability perspective, the significantly simpler architecture and the hierarchical nature of our adapted HAN model facilitate a more straightforward interpretability analysis based on visualizing attention weights. Building on this advantage, we conduct an interpretability analysis of our HAN that highlights the regions of the ECG signal most relevant to the model's decisions.
Abstract (translated)
心血管疾病仍然是全球主要的死亡原因之一,这凸显了准确且易于解释的诊断机器学习工具的需求。在这项研究中,我们使用心电图(ECG)数据对心脏病进行分类,这些数据来自两个广泛使用的数据集:MIT-BIH心律失常数据集和PTB-XL数据集。我们将最初为文本分类开发的分层注意力网络(HAN)适应到基于ECG的心脏病分类任务中。我们的改进型HAN包含两个关注不同大小的ECG数据片段的关注层。 我们对改进型HAN与一种更复杂的最新架构——具有卷积、注意和变换器层的CAT-Net进行了比较分析,后者在模型复杂性和准确性方面表现出色。我们的实证评估包括多个方面:测试准确度(通过0-1损失衡量);模型复杂性(通过模型参数的数量衡量);以及可解释性(通过注意力图可视化)。改进型HAN展示了与CAT-Net相当的测试准确率,但显著降低了模型复杂性和提高了可解释性: - 对于MIT-BIH数据集,我们的改进型HAN达到了98.55%的测试准确度,而CAT-Net为99.14%,同时减少了15.6倍的模型参数数量。 - 在PTB-XL数据集中,与CAT-Net相比,我们的改进型HAN将模型复杂性降低了19.3倍,仅在测试准确性上低了5%。 从可解释性的角度来看,我们简化后的架构和分层特性使得基于可视化注意力权重的解释分析更加直接。在此优势的基础上,我们对HAN进行了详细的可解释性分析,突出了ECG信号中模型决策最相关的区域。
URL
https://arxiv.org/abs/2504.03703