STST-JEPA: 浅层目标时空联合嵌入预测架构用于EEG自监督学习

# 面向EEG自监督学习的浅层目标时空联合嵌入预测架构  
来源：https://arxiv.org/html/2607.06629  

###### 目录  
1. https://arxiv.org/html/2607.06629#Sx1Abstract  
2. https://arxiv.org/html/2607.06629#Sx2Introduction  
3. https://arxiv.org/html/2607.06629#Sx3Methods  
   1. https://arxiv.org/html/2607.06629#Sx3.SSx1Datasets and Cohort Partitioning  
   2. https://arxiv.org/html/2607.06629#Sx3.SSx2EEG Segmentation and Preprocessing  
   3. https://arxiv.org/html/2607.06629#Sx3.SSx3Tokenization and Input Embedding  
   4. https://arxiv.org/html/2607.06629#Sx3.SSx4Encoder, Target Tokenizer, and Predictor  
   5. https://arxiv.org/html/2607.06629#Sx3.SSx5Self\-Supervised Objective  
   6. https://arxiv.org/html/2607.06629#Sx3.SSx6Optimization  
   7. https://arxiv.org/html/2607.06629#Sx3.SSx7Downstream Age Prediction  
   8. https://arxiv.org/html/2607.06629#Sx3.SSx8Auxiliary Probe Protocol  
   9. https://arxiv.org/html/2607.06629#Sx3.SSx9External Benchmark Protocol  
   10. https://arxiv.org/html/2607.06629#Sx3.SSx10Brain\-Age Gap \(BAG\) Construction  
   11. https://arxiv.org/html/2607.06629#Sx3.SSx11Behavioural\-Capacity Targets and Covariates  
   12. https://arxiv.org/html/2607.06629#Sx3.SSx12Multiple\-Comparison Handling  
   13. https://arxiv.org/html/2607.06629#Sx3.SSx13Hyperparameter Summary  
4. https://arxiv.org/html/2607.06629#Sx4Results  
   1. https://arxiv.org/html/2607.06629#Sx4.SSx1Validation Age Regression  
   2. https://arxiv.org/html/2607.06629#Sx4.SSx2Pretraining Dynamics  
   3. https://arxiv.org/html/2607.06629#Sx4.SSx3Auxiliary Downstream Probes  
   4. https://arxiv.org/html/2607.06629#Sx4.SSx4External Benchmark Performance  
   5. https://arxiv.org/html/2607.06629#Sx4.SSx5Brain\-Age Gap and Behavioural Performance  
5. https://arxiv.org/html/2607.06629#Sx5Discussion  
6. https://arxiv.org/html/2607.06629#Sx6Acknowledgements  
7. https://arxiv.org/html/2607.06629#Sx7References  

## 摘要  

脑年龄——即从生理记录中推断出的年龄——是一种新兴的生物标志物，其与实际年龄的偏差可追踪神经和精神疾病负担。EEG因其成本低廉、便携且时间分辨率高，成为实现该生物标志物的理想载体。然而，EEG脑年龄模型必须应对跨站点导联布局异质性、小规模标注队列以及显著的个体级非平稳性。此外，目前很少有EEG基础模型能证明，在儿科到老年全年龄段（该生物标志物实际部署的区间）内，均能提供具有竞争力的年龄回归性能。  

我们提出了STST-JEPA，一种面向静息态和任务态EEG的自监督transformer，在来自brain.space和Healthy Brain Network（HBN）语料库的47,703个会话（年龄范围5–81岁）上进行了预训练。该模型结合了潜在预测目标（在EMA-of-tokenizer目标上预测掩码token表示）与辅助信号重建项，应用于时空块掩码下的30秒多通道窗口。基于冻结预训练嵌入训练的轻量级注意力探针，在3,367个会话的年龄回归任务上达到了3.06年的最佳保留验证平均绝对误差（MAE）（r = 0.924），而预测均值的基线MAE约为10年。通过对模型最后几层进行轻量级任务特定微调，相同的预训练编码器在公共NeuralBench × brain.space EEG排行榜上取得了性别分类（平衡准确率0.911）、年龄预测（r = 0.749）和精神病理学复合回归（r = 0.215）的排名第一，且使用模型原生的30秒窗口。我们进一步表明，模型年龄预测的残差与我们检查的多项任务中的认知效率呈负相关。  

## 引言  

实际年龄是静息态电生理学的一个强大锚点。EEG频谱从儿童期到老年期系统性地重组——α峰迁移，1/f斜率变平，睡眠相关节律减弱——而这些特征中预测年龄与实际年龄之间的差距，已成为神经和精神疾病负担的候选生物标志物，在认知衰退人群中观察到系统性残差[Cole & Franke, 2017; Franke & Gaser, 2019]，以及睡眠EEG的纵向死亡率追踪[Sun et al., 2019]。EEG是这类生物标志物的天然载体：成本低、便携、时间分辨率丰富。另一方面，它难以适应现代深度学习架构的假设（特别是transformer [Vaswani et al., 2017]）：不同站点的EEG导联布局不同（仅在本工作中我们就面临一个115通道和一个128通道的布局），标注队列相对现代transformer的期望而言规模较小，且个体级非平稳性主导了条件内方差——表示容易与个体记录绑定，而非记录所要表达的构造。因此，我们追求一种表示，它能在已经抽象了这些干扰的空间中预测记录中缺失的内容，并通过重建项对其源自的波形负责。  

基于M/EEG的脑年龄预测已有成熟的经典文献和快速发展的深度学习文献。使用岭回归或随机森林回归器的手工提取频谱和协方差特征仍是强基线：Engemann等人[2022]的基准测试跨越四个M/EEG队列和超过2,500名参与者，报告MAE在7–8年范围内，R²高达约0.75，在清洁的成人临床EEG（TUAB）上最佳MAE为7.75年。任务特定的卷积模型，如EEGNet [Lawhern et al., 2018]以及Schirrmeister等人[2017]的浅/深卷积网络，在BCI范式上缩小了与经典管道的差距，同时保持紧凑和任务定制。大型睡眠EEG回归器在成年队列上达到接近7–8年的MAE [Sun et al., 2019]；儿科或窄年龄研究不可避免地报告更小的绝对误差，无法直接与成人基准比较。  

与此同时，EEG已涌现第一代基础模型，试图通过无标注预训练语料库摊销导联布局和队列异质性的成本：BENDR [Kostas et al., 2021]将wav2vec风格的对比学习应用于原始EEG；BrainBERT [Wang et al., 2023]将掩码谱图建模引入颅内记录；BIOT [Yang et al., 2023]独立将每个通道分词化以吸收不匹配的导联布局和缺失通道；LaBraM [Jiang et al., 2024]通过向量量化神经频谱预测扩展到约2,500小时和数亿参数；Neuro-GPT [Cui et al., 2024]在堆叠的EEG上训练GPT风格的掩码块预测器；EEG2Rep [Foumani et al., 2024]通过预测掩码时空区域的潜在表示，将联合嵌入预测训练引入EEG。其中一些模型确实报告了年龄回归，部分在共享基准（如NeuralBench）上进行了评估；然而，相对罕见的是单一基础模型既能跨异质、多站点语料库（从儿科到老年受试者）提供有竞争力的年龄回归，又能从一个冻结骨干网络实现排行榜级别的跨任务迁移——这正是脑年龄生物标志物实际部署的场景。  

视觉领域的自监督方法涵盖广泛设计空间——匹配掩码区域表示与EMA目标的潜在预测方法[iBOT, Zhou et al., 2022; LeCun, 2022; I-JEPA, Assran et al., 2023; V-JEPA, Bardes et al., 2024; data2vec, Baevski et al., 2022]、预测原始输入或其变分潜在变量的信号重建方法[VAE, Kingma & Welling, 2013; MAE, He et al., 2021]、将同一输入的不同增强视图拉近的对比方法[SimCLR, Chen et al., 2020; BYOL, Grill et al., 2020]，以及稳态/保方差正则化器[VICReg, Bardes et al., 2022; LeJEPA, Balestriero & LeCun, 2025]。STST-JEPA位于掩码预测象限，并带有辅助重建项。每个家族都有其特征性的失败模式。纯潜在预测可能在缺乏EMA目标、停止梯度或显式正则化器时崩溃——目标流漂移向平凡解，学生跟踪该漂移。纯信号重建过度投资于表面统计信息：在EEG上，噪声、肌肉伪影和眨眼携带大量能量，训练用于重建波形的解码器自然会分配容量来重建这些干扰。我们实例化了一个联合目标，结合两者——这种配对松散地映射到皮层的预测处理理论[Rao & Ballard, 1999; Friston, 2010; Clark, 2013]，其中潜在生成模型不断受到信号级后果的约束——但我们将其视为架构直觉而非对皮层保真度的主张。联合公式是否优于其在该语料库上的任意单一消融，留待未来工作，本文不做实证主张。  

STST-JEPA（本文）实现了这一设计。模型接收30秒、256Hz的窗口，统一128通道预算，通过一个学习的池化多头注意力（PMA）块在每个时间索引处折叠每个通道的时间补丁，该块通过坐标感知、掩码感知的通道池化允许任意导联布局。预训练在掩码位置应用潜在预测损失（预测器输出与对应EMA分词器输出之间，λ_lat=1.0），以及一个平滑L1重建损失，将每个掩码补丁解码回其16样本波形（λ_rec=0.35），应用于平均覆盖约24%网格的时空块掩码。重建项被有意降低权重：并非要同等最小化，而是作为一个软下限，使学得的潜在空间足够接近信号，以便一个短的线性–GeLU–线性头部能够解码。  

我们在年龄回归上检查得到的表示，使用一个组合的内部brain.space加上Healthy Brain Network（HBN; Alexander et al., 2017）预训练语料库，包含47,703个会话，年龄范围5至81岁。一个基于冻结预训练嵌入的轻量级注意力探针在保留验证集上达到3.06年MAE和5.11年RMSE，而预测训练集均值基线的MAE为10.09年，RMSE为13.27年——MAE降低了69.7%。该结果不能直接与在清洁成人队列（如TUAB）上报告的7–8年MAE相比较：我们的语料库偏向儿科（HBN贡献25,115个会话，年龄集中在儿童和青少年期），该队列中年龄的条件方差小于成人临床EEG，这机械性地降低了可实现MAE。因此，我们将年龄作为表示的压力测试——判断预训练是否捕捉到受试者级结构的最清晰标尺——而非本文的主要终点。  

在相同冻结嵌入上的辅助探针讲述了互补的故事：性别（平衡准确率0.89）和范式身份（睁眼/闭眼0.88，其他三个任务分类器高于随机）也可恢复。为了检查这并非我们内部划分或标签集的偶然现象，我们还在公共NeuralBench × brain.space EEG排行榜上评估了相同的预训练检查点（根据排行榜协议进行最后一层微调；方法§*外部基准协议*），它在三个不同的下游任务（性别、年龄和精神病理学复合指标）上均从单个共享骨干取得排名第一——这最具体地体现了STST-JEPA的基础模型定位，以脑年龄作为主要应用。  

我们对范围保持明确。NeuralBench × brain.space排行榜结果（表4）是在NeuralBench保留测试分区上的固定协议评估。对于内部brain.space队列，我们报告的年龄回归数字（表2）是训练轨迹中的最佳验证值，而非单一固定协议测试抽取；内部测试分区保留用于未来的报告，本文唯一触及该分区的部分是结果中的探索性脑年龄差距×行为能力分析，该分析合并验证和测试分区以增强统计功效，并在全文标注为探索性。联合目标的受控消融超出我们的计算预算，尽管在开发过程中探索了广泛的架构和目标。我们不提出临床有效性声明。以下工作应作为一项架构和实证研究来阅读，聚焦于一个设计选择——联合潜在+信号监督——在一个大型异质语料库上，针对年龄回归，并辅以探针面板和外部多任务基准进行报告。  

贡献。  
- • STST-JEPA，一种24层transformer基础模型，具有坐标感知PMA通道池化，在来自两个不同导联布局和年龄分布的队列（brain.space和HBN）的47,703个EEG会话上，以联合潜在预测+重建目标（预测器输出受EMA-of-tokenizer目标监督）进行预训练，通过共享128通道预算和缺失通道掩码统一。  
- • 一个可重用的注意力探针协议，基于冻结预训练嵌入，支持年龄回归和一组辅助标签及范式探针，无需重新训练骨干网络。  
- • 从单个预训练检查点的跨任务泛化：在公共NeuralBench × brain.space EEG排行榜上，性别（平衡准确率0.911）、年龄（Pearson r 0.749）和精神病理学复合指标（Pearson r 0.215）均获得排名第一，使用模型原生的30秒窗口。在基准的2秒标准协议下，性别和精神病理学仍保持第一，年龄为第四（§外部基准性能）。这直接证明了模型的基础模型地位——一个骨干网络，三个不同的下游任务。  
- • 脑年龄作为主要应用：在涵盖5–81岁的受试者不相交分区上，最佳保留验证MAE为3.06年，RMSE为5.11年，辅以辅助探针面板显示相同的冻结嵌入也远高于随机水平地恢复性别和范式身份。  
- • 一项探索性的BAG×行为能力分析，基于保留验证和测试分区的并集：21个认知效率目标中有7个通过了Benjamini–Hochberg FDR校正（q = 0.05），所有目标均符合预期负方向（老化大脑↔效率降低）；效应量小（|r| < 0.10），关键观察是效应方向的一致性而非效应量的发现。  
- • 一种受大脑启发的框架，通过松散类比皮层的预测处理理论（架构直觉，非对皮层保真度的主张）来激励联合目标。  

## 方法  

我们使用一种潜在预测目标（在EMA-of-tokenizer目标上预测掩码token表示）结合每补丁信号重建项，预训练了一个EEG基础模型——下文称为STST-JEPA——并评估学得的表示在下游年龄回归上的表现。本节其余部分描述了预训练语料库和队列划分、EEG预处理管道、分词器和编码器架构、自监督目标、优化以及注意力年龄探针。以下所有值均指代