心电自监督论文分享 （6）—— ST-MEM：时空掩码心电图建模
GUIDING MASKED REPRESENTATION LEARNING TO CAPTURE SPATIO-TEMPORAL RELATIONSHIP OF ELECTROCARDIOGRAM（ICLR-2024）
研究背景与动机
心电图（ECG）是监测心脏电活动、诊断心血管疾病的重要无创工具。标准 12 导联 ECG 可被视为一个多变量时间序列：12 个导联从不同空间位置采集信号（肢体导联位于额面，胸导联位于水平面），同时每个导联记录了随时间变化的电压波形。
现有 ECG 自监督学习方法主要分为两类：
对比学习（如 MoCo v3、CMSC/CLOCS）：通过数据增强前后的表征一致性进行学习。但简单的增强操作（裁剪、翻转、平移）可能剧烈改变 ECG 信号的语义含义（如翻转可能将 ST 段抬高变成压低）。
生成式学习（如 MaeFE 的 MTAE / MLAE）：借鉴 MAE 范式，通过重建被遮蔽的 ECG 片段来学习表征。但现有方法要么仅做时间维度的 patchify（MTAE，忽略导联间空间关系），要么仅做空间维度的 patchify（MLAE，忽略导联内时序关系），未能同时捕获时空信息。
针对上述问题，作者提出了ST-MEM（Spatio-Temporal Masked Electrocardiogram Modeling）：一种简单但有效的 ECG 自监督生成式学习框架，通过时空联合 patchify + 导联指示模块，显式捕获 ECG 信号中的时空关系。
数据集与实验任务
- 预训练阶段：合并三个公开 12 导联 ECG 数据集，共188,480 条无标注 ECG，将他们统一重采样至250Hz
| 数据集名称 | 原始采样率 | 长度 |
|---|---|---|
| Chapman | 500Hz | 10s |
| Ningbo | 500Hz | 10s |
| CODE-15 | 400Hz | 10s |
下游任务：心律失常 + 心肌梗死分类
- PTB-XL：21,837 条 12 导联，5 类（NORM、CD、MI、HYP、STTC）
- CPSC2018：6,877 条 12 导联，9 类心律失常
- PhysioNet2017：8,528 条单导联（Lead I），4 类心律失常
评估范式：线性评估（Linear Evaluation）+ 微调（Fine-tuning），指标为Accuracy、F1、AUROC；还包含低资源（1%/5%标注）和任意导联数（1/6/12 导联）实验
架构图与核心方法
整体框架如图 3 所示：编码器-解码器架构，编码器接受 patchify 后的 ECG 信号（附加位置嵌入和导联嵌入），解码器逐导联重建被遮蔽的 patch（MAE 损失）。
- 时空 Patchify（Spatio-Temporal Patchifying）
这是 ST-MEM 最核心的设计选择。给定 ECG 信号X ∈ R L × T X \in \mathbb{R}^{L \times T}X∈RL×T（L LL个导联，T TT个时间步）： 对每个导联独立地沿时间轴分割为n = T / p n = T/pn=T/p个非重叠 patch（patch 大小p pp），总计产生L × n L \times nL×n个 patch
如图所示，与仅时间 patchify（图 2a，3 个 patch）和仅空间 patchify（图 2b，12 个 patch）不同，时空 patchify（图 2c）产生细粒度的12 × n 12 \times n12×n个 patch，使自注意力机制同时在时间维度和导联维度上运行，显式捕获时空关系。
- 导联感知共享解码器（Lead-wise Shared Decoder）
上面时空 patchify 虽然有利于学习时空表征，但也降低了重建任务的难度——解码器可以轻松地从其他导联的未遮蔽 patch 中获取同一时刻的信息来重建，这样模型就不会学到真正有效的特征。为了解决这个问题：限制解码器每次只处理单一导联的 patch 序列，各导联共享同一个解码器参数。这一设计：增加了重建任务难度，迫使编码器学习更丰富的时空表征，提升训练效率（所有导联共享解码器）
- 导联指示模块（Lead Indicating Modules）
仅增加重建难度还不够——编码器和解码器需要能够区分 patch 来自哪个导联，才能真正学习空间关系。为此引入两个模块：
| 模块 | 作用 |
|---|---|
| 导联嵌入（Lead Embedding） | 每个 patch 嵌入添加一个可学习的导联特定嵌入L e a d ∈ R D Lead \in \mathbb{R}^DLead∈RD，同一导联的所有 patch 共享相同的导联嵌入 |
| 分离嵌入（SEP Embedding） | 在每个导联的 patch 序列前后插入共享的[SEP]嵌入，帮助模型区分不同导联的 patch 序列边界 |
- 预训练与下游任务
预训练损失：对随机遮蔽的 patch（掩码率m mm），最小化原始信号与重建信号的 MAE：
L S S L = 1 L n m ∑ i ∈ M ∥ P a t c h ^ i − P a t c h i ∥ L_{SSL} = \frac{1}{Lnm} \sum_{i \in M} \lVert \widehat{Patch}_i - Patch_i \rVertLSSL=Lnm1i∈M∑∥Patchi−Patchi∥
下游任务：下游任务分为线性评估和微调两种任务：其中线性评估丢弃解码器，冻结编码器，加一个可训练的线性分类头。微调丢弃解码器，编码器与线性分类头参数共同更新。两者都采用交叉熵损失：
L C E = 1 K ∑ c = 1 K ( − y c log ( y ^ c ) ) L_{CE} = \frac{1}{K} \sum_{c=1}^{K} (-y_c \log(\hat{y}_c))LCE=K1c=1∑K(−yclog(y^c))
任意导联数量微调：ST-MEM 天然支持任意导联组合——patch 的形状不受导联数影响，导联嵌入模块使模型理解不同导联的特性。因此预训练于 12 导联后可无缝微调于 6 导联（仅肢体导联）或单导联数据。
实验结果与结论
12 导联线性评估与微调
低资源场景（仅用 1%/5% 标注数据）
- ST-MEM 在所有低资源配置下均最优
- CPSC2018 仅用 5% 数据即可接近 100% 数据的性能
任意导联数量鲁棒性
- ST-MEM 在所有导联组合下均显著优于其他方法（p < 0.05）
- 证明了预训练于 12 导联再微调于降导联场景的可行性
空间关系可视化分析
- t-SNE + GMM 聚类：ST-MEM 学习到的表征在嵌入空间中自然地按额面（肢体导联 I、II）和水平面（胸导联 V2、V3）分离，聚类准确率高达 89.5%，而 MTAE 仅 60.7%、MoCo v3 仅 50.6%（≈随机）
- 注意力图分析：以 Lead III（肢体导联）的 patch 为 query，注意力高分区域集中在肢体导联范围，同时时间上相似波形（周期性心搏）的 patch 也获得高注意力——验证了 ST-MEM 同时捕获了空间和时间关系
消融实验
每个导联指示模块的加入都持续提升性能，验证了空间信息建模对 ECG 表征学习的重要性。
局限性与发展
优势
- 时空联合建模：通过时空 patchify 使自注意力同时跨越时间和导联维度，显式捕获 ECG 的时空关系，这是该方法的最核心创新
- 生成式无增强：基于 MAE 重建范式，不依赖可能有语义破坏风险的数据增强，更适用于 ECG 信号
- 任意导联数量适应：导联嵌入 + patch 形状不受导联数影响，预训练于 12 导联后可无缝迁移到任意导联组合（12/6/1 导联）
- 低资源优势显著：在极端标注稀缺（1%）场景下大幅领先 Supervised 和各类 SSL Baseline
- 可解释性强：t-SNE 聚类和注意力图直观展示了模型学到的空间（额面 vs 水平面）和时间（心搏周期）关系
局限性
- 仅验证 ECG 任务：虽然声称可扩展到通用多变量时间序列，但仅在附录中展示了 HAR 任务的初步结果
- 编码器骨干固定为 ViT-B：未探索不同 backbone 规模（ViT-S/L）的影响
- 下游微调策略简单：仅采用全量微调（naive fine-tuning），未探索部分微调（partial fine-tuning）等更适合小样本的迁移策略
未来拓展方向
- 可变形 patch 或层次化 patch：替代固定大小的非重叠 patch，更好地对齐 ECG 信号中的 P-QRS-T 波边界
- 多模态扩展：将 ST-MEM 扩展到 ECG + 临床文本 / EHR 等多模态数据
- 部分微调策略：探索仅微调部分层或添加 adapter 的方式，进一步提升低资源场景的迁移效率
- 通用时序扩展：系统性验证 ST-MEM 在其他多元时间序列（如 EEG、fNIRS、HAR）上的泛化能力
代码地址：https://github.com/bakqui/ST-MEM