[论文解读] Masked Transformer for Electrocardiogram Classification

Zhou Ya, Xiaolin Diao+4 more|arXiv (Cornell University)|Aug 31, 2023

ECG Monitoring and AnalysisMedicine被引用 3

一句话总结

该论文提出 MTECG，一种用于心电图（ECG）分类的掩码 Transformer 模型，其在大规模、专家标注的 ECG 数据集（共 220,251 个记录）上进行了掩码自编码预训练。通过将 ECG 信号分割为非重叠的片段，并采用轻量级解码器的波动重建目标，MTECG-T 在私有和公共 ECG 数据集上均实现了最先进性能，优于近期的 SOTA 方法，同时保持了仅 5.7M 参数的轻量化架构。

ABSTRACT

Electrocardiogram (ECG) is one of the most important diagnostic tools in clinical applications. With the advent of advanced algorithms, various deep learning models have been adopted for ECG tasks. However, the potential of Transformer for ECG data has not been fully realized, despite their widespread success in computer vision and natural language processing. In this work, we present Masked Transformer for ECG classification (MTECG), a simple yet effective method which significantly outperforms recent state-of-the-art algorithms in ECG classification. Our approach adapts the image-based masked autoencoders to self-supervised representation learning from ECG time series. We utilize a lightweight Transformer for the encoder and a 1-layer Transformer for the decoder. The ECG signal is split into a sequence of non-overlapping segments along the time dimension, and learnable positional embeddings are added to preserve the sequential information. We construct the Fuwai dataset comprising 220,251 ECG recordings with a broad range of diagnoses, annotated by medical experts, to explore the potential of Transformer. A strong pre-training and fine-tuning recipe is proposed from the empirical study. The experiments demonstrate that the proposed method increases the macro F1 scores by 3.4%-27.5% on the Fuwai dataset, 9.9%-32.0% on the PTB-XL dataset, and 9.4%-39.1% on a multicenter dataset, compared to the alternative methods. We hope that this study could direct future research on the application of Transformer to more ECG tasks.

研究动机与目标

• 解决尽管在视觉和自然语言处理中表现优异，但 Transformers 在 ECG 分类中性能有限的问题。
• 通过构建包含 220,251 个 ECG 记录的大规模、专家标注数据集，克服小规模、标注 ECG 数据集的挑战。
• 将已在视觉和自然语言处理中验证有效的掩码自编码（MAE）技术适配至时间序列 ECG 数据，实现自监督预训练。
• 研究有效的训练策略，包括重建目标设计、学习率调度和正则化方法，以优化轻量化 Transformer 模型在 ECG 任务中的表现。
• 评估原始视觉 Transformer 架构是否能在结合适当的预训练与微调策略后，超越专用 CNN 模型。

提出的方法

• 将 ECG 信号分割为非重叠的固定长度片段，以减少序列长度并保留结构信息。
• 使用可学习的位置嵌入以维持序列顺序并区分片段位置。
• 在未标注的 ECG 数据上应用掩码自编码（MAE）预训练策略，随机掩码部分片段，并通过轻量级解码器重建。
• 引入波动的重建目标，以促使模型学习有意义的波形特征，而非平凡模式。
• 在微调阶段，采用逐层学习率衰减和 DropPath 正则化，以减少过拟合并提升泛化能力。
• 最终模型 MTECG-T 是基于轻量化视觉 Transformer 架构，通过完整的预训练与微调流程训练而成。

实验结果

研究问题

• RQ1 掩码自编码能否有效适配至 ECG 时间序列数据，以提升自监督表征学习效果？
• RQ2 如重建目标设计、训练周期长度和正则化技术等训练组件，如何影响 ECG 分类中的模型性能？
• RQ3 当结合适当的预训练与微调策略时，轻量化、原始的视觉 Transformer 架构能否在 ECG 分类中实现最先进性能？
• RQ4 所提出的掩码预训练方法是否能有效利用未标注 ECG 数据，尤其在其他模态（如超声心动图）标签稀缺的临床场景中？
• RQ5 Transformer 模型在性能上超越 CNN 的优势，是源于架构本身的优越性，还是预训练策略的有效性？

主要发现

• MTECG-T 在私有和公共 ECG 数据集上均实现了最先进性能，分类准确率优于近期 SOTA 方法。
• 该模型仅含 570 万个参数，在 5% 至 75% 的广泛掩码比率范围内表现稳定，展现出对预训练超参数的强鲁棒性。
• 消融实验证实，波动重建目标、更长的训练周期、逐层学习率衰减以及最优的 DropPath 比率对性能至关重要。
• 当模型规模超过某一临界值后，性能出现下降，表明数据集规模限制了更大架构的收益。
• 结果挑战了认为 Transformers 不适合轻量化 ECG 分类的假设，表明结合恰当训练策略的原始 ViT 可超越专用 CNN。
• MTECG 的自监督特性使其能有效利用未标注 ECG 数据，适用于标签数据稀缺的新临床场景。