用分层注意力模型分析心电图,准确率高且结果更易解释。
Hierarchical Attention Network for Interpretable ECG-based Heart Disease Classification
- 将文本分类的分层注意力机制改造用于心电图疾病分类。
- 在MIT-BIH数据集上达98.55%准确率,参数量减少15.6倍。
- 结构简单,注意力图清晰,适合医疗场景可解释性需求。
心血管疾病是全球主要致死原因,亟需高精度且可解释的机器学习诊断工具。本文利用MIT-BIH Arrhythmia和PTB-XL两个常用心电图(ECG)数据集,将原本用于文本分类的分层注意力网络(HAN)适配至心电图心脏病分类任务。该模型包含两层注意力机制,分别关注不同尺度的ECG信号片段。与包含卷积、注意力和变压器层的先进架构CAT-Net对比,本方法在多个维度表现优异:在MIT-BIH数据集上,测试准确率达98.55%,较CAT-Net的99.14%仅低0.59%,但模型参数量减少15.6倍;在PTB-XL数据集上,参数量减少19.3倍,准确率仅低5%。从可解释性看,因结构更简单且具备层级特性,注意力权重可视化更直观,能有效识别对模型决策最关键的ECG信号区域。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。