[论文解读] MSW-Transformer: Multi-Scale Shifted Windows Transformer Networks for 12-Lead ECG Classification

Renjie Cheng, Zhemin Zhuang+3 more|arXiv (Cornell University)|Jun 21, 2023

ECG Monitoring and Analysis被引用 7

一句话总结

本文提出 MSW-Transformer，一种轻量级、单层的 Transformer 架构，采用多尺度移位窗口与可学习特征融合机制，以提升 12 导联 ECG 分类性能。该方法在 PTBXL-2020 数据集上达到最先进性能，超级类别宏 F1 分数最高达 77.85%，样本宏 F1 分数达 81.26%，并通过注意力可视化提升了可解释性。

ABSTRACT

Automatic classification of electrocardiogram (ECG) signals plays a crucial role in the early prevention and diagnosis of cardiovascular diseases. While ECG signals can be used for the diagnosis of various diseases, their pathological characteristics exhibit minimal variations, posing a challenge to automatic classification models. Existing methods primarily utilize convolutional neural networks to extract ECG signal features for classification, which may not fully capture the pathological feature differences of different diseases. Transformer networks have advantages in feature extraction for sequence data, but the complete network is complex and relies on large-scale datasets. To address these challenges, we propose a single-layer Transformer network called Multi-Scale Shifted Windows Transformer Networks (MSW-Transformer), which uses a multi-window sliding attention mechanism at different scales to capture features in different dimensions. The self-attention is restricted to non-overlapping local windows via shifted windows, and different window scales have different receptive fields. A learnable feature fusion method is then proposed to integrate features from different windows to further enhance model performance. Furthermore, we visualize the attention mechanism of the multi-window shifted mechanism to achieve better clinical interpretation in the ECG classification task. The proposed model achieves state-of-the-art performance on five classification tasks of the PTBXL-2020 12-lead ECG dataset, which includes 5 diagnostic superclasses, 23 diagnostic subclasses, 12 rhythm classes, 17 morphology classes, and 44 diagnosis classes, with average macro-F1 scores of 77.85%, 47.57%, 66.13%, 34.60%, and 34.29%, and average sample-F1 scores of 81.26%, 68.27%, 91.32%, 50.07%, and 63.19%, respectively.

研究动机与目标

• 解决 ECG 信号中病理特征差异小的问题，该问题阻碍了自动分类。
• 克服 CNN 和标准 Transformer 在捕捉细微 ECG 变化以及处理长序列方面的局限性。
• 开发一种轻量级、高效的 Transformer 模型，降低复杂度与计算成本，适用于临床部署。
• 通过可视化注意力机制，使其与临床诊断规则对齐，提升模型可解释性。
• 在多类别 12 导联 ECG 分类任务中实现最先进性能，且参数量极少。

提出的方法

• 提出一种多尺度移位窗口机制，其中自注意力被限制在不同尺度的非重叠局部窗口内，以捕捉分层特征。
• 通过移位窗口实现跨窗口信息流动，同时保持局部注意力效率并降低计算成本。
• 引入可学习特征融合模块，将来自多个窗口尺度的注意力输出进行融合，以增强表征学习能力。
• 采用单层 Transformer 块并结合位置编码，以建模 12 导联 ECG 序列中的长距离依赖关系。
• 可视化不同窗口尺度下的注意力分数，并将其与临床诊断标准相关联，以增强可解释性。

实验结果

研究问题

• RQ1 单层 Transformer 搭配多尺度移位窗口是否能在 12 导联 ECG 分类任务中超越更深或更复杂的模型？
• RQ2 多尺度窗口机制如何提升对具有细微病理差异的 ECG 信号的特征提取能力？
• RQ3 与固定拼接或平均方法相比，可学习特征融合在多大程度上提升了模型性能？
• RQ4 MSW-Transformer 中的注意力机制是否能与临床诊断规则有意义地对齐？
• RQ5 该模型在类别不平衡的 ECG 分类任务中表现如何，特别是在罕见或临床相似性高的心律失常类别上？

主要发现

• 在 PTBXL-2020 数据集中，MSW-Transformer 在 5 个诊断超级类别上达到最先进宏 F1 分数 77.85%，在 23 个子类别上为 47.57%，在 12 个心律类别上为 66.13%，在 17 个形态类别上为 34.60%，在 44 个诊断类别上为 34.29%。
• 样本宏 F1 分数分别为 81.26%、68.27%、91.32%、50.07% 和 63.19%，表明模型在平衡与非平衡类别上均表现出色。
• 混淆矩阵显示，常见类别如 Norm（宏 F1 89.13%）、CLBBB（90.00%）、SR（95.96%）和 STACH（92.62%）的准确率较高，误分类率较低。
• 尽管存在数据不平衡，模型在区分常见 ECG 模式方面表现出鲁棒性，但对罕见标签如 ISCA、NST 和 LAFB/LPFB 的性能有所下降，原因在于类别间相似性高。
• 注意力可视化结果表明，模型聚焦于临床相关的 ECG 段（如 ST 段和 QRS 复合波），与专家诊断标准一致，显著增强了可解释性。
• 与标准 Transformer 相比，该模型参数量更少，推理速度更快，更适合实际临床部署。