[论文解读] Guiding Masked Representation Learning to Capture Spatio-Temporal Relationship of Electrocardiogram

Yeongyeon Na, Minje Park+2 more|arXiv (Cornell University)|Feb 2, 2024

ECG Monitoring and AnalysisMedicine被引用 3

一句话总结

该论文提出 ST-MEM，一种自监督掩码自编码器框架，通过重建12导联心电图（ECG）数据的时空块来学习通用的心电图表征。通过联合建模时间与空间关系，采用分块和掩码重建策略，ST-MEM 在心律失常分类任务中达到最先进性能，优于对比学习和生成式自监督基线方法，并在减少导联和低资源设置下展现出良好的泛化能力。

ABSTRACT

Electrocardiograms (ECG) are widely employed as a diagnostic tool for monitoring electrical signals originating from a heart. Recent machine learning research efforts have focused on the application of screening various diseases using ECG signals. However, adapting to the application of screening disease is challenging in that labeled ECG data are limited. Achieving general representation through self-supervised learning (SSL) is a well-known approach to overcome the scarcity of labeled data; however, a naive application of SSL to ECG data, without considering the spatial-temporal relationships inherent in ECG signals, may yield suboptimal results. In this paper, we introduce ST-MEM (Spatio-Temporal Masked Electrocardiogram Modeling), designed to learn spatio-temporal features by reconstructing masked 12-lead ECG data. ST-MEM outperforms other SSL baseline methods in various experimental settings for arrhythmia classification tasks. Moreover, we demonstrate that ST-MEM is adaptable to various lead combinations. Through quantitative and qualitative analysis, we show a spatio-temporal relationship within ECG data. Our code is available at https://github.com/bakqui/ST-MEM.

研究动机与目标

• 为解决在训练用于多种心脏疾病检测的模型时标注 ECG 数据有限的挑战。
• 通过显式建模 12 导联 ECG 信号中固有的时间与空间关系，改进心电图中的自监督表征学习。
• 开发一种灵活的框架，可在不同导联配置下泛化，包括减少导联和单导联 ECG。
• 通过定量和定性分析验证所学习的表征是否捕捉了有意义的时空模式。

提出的方法

• 对 12 导联 ECG 信号应用时空分块，将其划分为时间块和空间块以供模型输入。
• 采用掩码自编码器（MAE）架构，其中在预训练过程中随机掩码 60% 的块，并使用共享解码器和可学习导联嵌入进行重建。
• 模型采用双分支编码策略，结合独立的时间注意力和空间注意力机制，以捕捉动态和解剖依赖关系。
• 引入导联共享解码器和分离嵌入，以在重建过程中保持导联间的空间关系。
• 预训练持续 200 个周期，固定初始学习率为 0.001，最小化原始信号与重建信号之间的均方误差。
• 微调在下游分类任务中使用交叉熵损失进行 30 个周期，固定初始学习率为 0.001。

实验结果

研究问题

• RQ1 掩码自编码器框架是否能通过捕捉 12 导联 ECG 数据中的时间与空间依赖关系，有效学习通用的心电图表征？
• RQ2 在标准和低资源设置下，ST-MEM 在心律失常分类任务中与对比学习和生成式自监督基线相比表现如何？
• RQ3 ST-MEM 在减少导联 ECG 配置（如肢体导联或单导联记录）中能多大程度实现泛化？
• RQ4 有哪些证据表明所学习的表征编码了心电图信号中有意义的时空关系？

主要发现

• 在 CPSC2018 数据集上，ST-MEM 在仅使用 1% 训练数据的情况下达到 0.897 ± 0.025 的 AUROC，PTB-XL 数据集上达到 0.815 ± 0.012，优于所有复现的基线方法，包括 CPC 和 MLAE。
• 在人体活动识别（HAR）基准上，ST-MEM 达到 98.4% 的 F1 分数，表明其在通用多变量时间序列任务中具有强大的迁移能力。
• 在仅使用 1% 训练数据的低资源设置下，ST-MEM 仍保持优越性能，在 PTB-XL 上达到 0.815 AUROC，在 CPSC2018 上达到 0.897 AUROC。
• 定性分析表明，模型能够重建出有意义的心电图形态，表明其有效捕捉了时空模式。
• ST-MEM 在减少导联配置下表现出色，即使在肢体导联或单导联上微调，仍能保持高性能。
• 在某些设置下，模型性能与监督学习相当，表明自监督预训练具有强大的泛化能力。