Brant-2：开启脑信号分析新篇章的基础模型

人工智能咨询培训老师叶梓 转载标明出处

脑信号，包括通过侵入性或非侵入性方式收集的脑电图（EEG）和立体脑电图（SEEG）等生物测量信息，为我们理解大脑的生理功能和相关疾病的机制提供了宝贵的洞见。然而脑信号的收集和分析面临着成本高、数据量大、标注困难等挑战。为了克服这些难题，浙江大学的研究团队提出了提出了Brant-2，这是一个基于大量未标记数据预训练的、能够适应多种脑信号应用场景的强大模型。

Brant-2模型是迄今为止最大的脑信号基础模型，它不仅在数据变化和建模尺度上表现出鲁棒性，还能够应用于更广泛的脑神经数据类型。与前一代模型Brant相比，Brant-2在预训练语料库的规模和多样性、对数据变化的鲁棒性以及在不同建模尺度上的适应性方面都有显著提升。

 Brant-2 模型的总体工作流程

图1展示了Brant-2模型的工作概览：研究团队首先利用来自超过15,000名受试者的约4TB脑神经数据构建了预训练语料库。接着通过两种预训练任务训练Brant-2，使其能够理解信号模式并预测未来变化。最后这个预训练好的模型可以通过微调，被应用于脑信号分析的各种应用场景，如癫痫检测、睡眠阶段分类等。这种方法不仅提高了模型的泛化能力，还减少了对大量标注数据的依赖。

架构

 Brant-2 的架构和预训练框架

Brant-2模型的设计旨在处理脑信号数据的多样性和复杂性。总体架构如图2所示，输入的原始脑信号xi​首先被处理成子序列级别的块pi​。接着，进行数据增强以增加训练数据的多样性，并对一部分块进行掩码处理。结合时域和频域的信息来获得输入嵌入这些嵌入随后被顺序送入时间编码器和空间编码器中。编码器的输出表示通过线性映射用于重建被掩码的块和预测未来的信号。架构主要包含四个模块：

1. Patching（分块）：将时间步长聚合成子序列级别的块，增强局部性并捕获全面的语义信息，同时降低计算成本。输入样本被划分为非重叠的块，生成一系列块。
2. Data Augmentation and Masking Module（数据增强与掩码模块）：数据的多样性对于训练基础模型至关重要。在预训练过程中，通过数据增强模块在时间和空间维度上扩展训练数据的多样性，增强模型对数据变化的鲁棒性。
3. Input Embedding Module（输入嵌入模块）：结合时域和频域的特征，生成输入嵌入。通过计算增强数据的功率谱密度（PSD），揭示信号在不同频带的谱功率分布，这与不同的脑功能状态相关。
4. Encoder（编码器）：为了适应不同的应用场景，Brant-2在预训练中整合了掩码预测和预测任务，学习具有丰富语义信息的表示。设计了一个多前馈（multi-FFN）Transformer块，包含两个前馈网络，一个用于信号重建，另一个用于预测。

预训练和微调

在Brant-2模型的预训练阶段，研究者们采用了掩码预测和预测任务，目的是充分提取信号的丰富语义信息，以适应不同的下游任务。这种方法使模型能够理解信号中特定片段的模式，并从当前观察到的序列中学习未来趋势的变化。预训练涉及两个关键步骤：

1. 掩码预测：在这一任务中，模型需要重建信号中被随机掩盖的部分。这有助于模型理解信号的局部模式和结构。通过这种方式，模型能够学习到信号的内在规律，这对于后续的任务至关重要。

2. 预测任务：与掩码预测不同，预测任务要求模型基于当前观察到的信号序列，预测未来的信号变化。这不仅考验了模型对信号动态变化的理解，也锻炼了其预测能力。

为了实现这些任务，模型使用了两个线性头（Wrec和Wfcst），它们将潜在的表示映射回原始信号。在预训练过程中，模型执行了块级别的重建和序列级别的预测。具体地，对于每个块，模型通过公式(4)计算重建的信号其中是重建的信号，zi​是经过编码器处理的潜在表示。同时，通过公式(5)进行序列级别的预测其中MeanPool操作对每个连续的块进行聚合，是部分观测的潜在表示。

预训练阶段的损失函数由两个均方误差(MSE)损失组成，分别对应掩码预测和预测任务。公式(6)和公式(7)分别表示这两个损失，其中N是训练样本的数量。通过联合优化这两个损失，模型能够在预训练阶段学习到有效的信号表示。

在微调阶段，模型首先使用均值池化操作聚合连续的块，然后通过加权和的方式聚合表示，公式(8)描述了这一过程其中λ是一个可学习的参数，用于平衡不同块的表示。聚合后的表示ri​将被送入下游任务的线性或非线性头部，以进行具体的任务处理。

这种方法不仅使Brant-2模型在预训练阶段能够学习到丰富的信号表示，而且通过微调，模型能够适应特定的下游任务，展现出强大的泛化能力和灵活性。

实验

Brant-2的预训练语料库结合了超过15,000名受试者的SEEG和EEG数据，总数据量接近4TB。SEEG数据集包含了26名受试者的颅内神经数据，经过去噪和过滤等预处理后，用于预训练的SEEG数据为2.3TB。EEG数据集使用了公开可用的TUEG数据集，包含来自14,987名个体的1,643GB临床记录。

在Brant-2的编码器块中，应用了RMSNorm和Swish激活函数。设置了上下文长度、掩码比率和预测长度等超参数，并选择了采样率调整因子和重组通道数。使用AdamW优化器进行训练，并采用了线性预热和余弦衰减的学习率调度策略。预训练过程中，模型参数更新了总共105k步。

评估设置

评估实验：在九个不同的SEEG和EEG数据集上进行了评估实验，涵盖了五个下游任务：癫痫发作检测、癫痫发作预测、睡眠阶段分类、情绪识别和运动想象分类。每个数据集被分成几个不重叠的组，并进行了n折交叉验证。

癫痫发作检测：使用MAYO和FNUSA两个SEEG数据集以及CHB-MIT和Siena两个EEG数据集来评估模型在癫痫发作检测上的性能。数据被分割成几秒长的数据片段，并使用了精确度、召回率、F1和F2分数作为评估指标。

癫痫发作预测：与癫痫发作检测不同，癫痫发作预测的任务是基于当前观察预测未来发作的可能性。使用了一家一流医院的临床SEEG数据集，数据集包含5名受试者，采样率为1000Hz，并采用了5折交叉验证。

睡眠阶段分类：选择了SleepEDF和Haaglanden Medisch Centrum睡眠分期数据库（HMC）两个EEG数据集来验证模型在睡眠阶段分类上的性能。评估指标包括准确率、敏感性、特异性、宏观F1分数和Cohen's kappa系数。

情绪识别：SEED数据集包含15名受试者在观看电影片段时的62通道EEG数据。评估指标包括准确率和宏观F1分数。

运动想象分类：选择了EEG Motor Movement/Imagery数据集，包含109名志愿者的超过1500个一分钟和两分钟的64通道EEG记录。评估指标包括准确率和F1分数。

基线

将Brant-2模型与12种先进方法进行了广泛的比较，这些方法分为三类：时间序列通用建模方法、基于脑信号自监督预训练的方法以及为每个任务特别设计的方法。这些方法在所有下游任务上进行了评估，而特定任务的方法仅在特定任务上进行了评估。

评估结果

图3总结了Brant-2与基线方法在所有下游任务上的整体结果。雷达图显示Brant-2在所有通用时间序列建模方法和脑信号预训练方法上都表现更好，超越了大多数特定任务的方法，表明Brant-2在各种脑信号场景中展现出强大的泛化能力。

表1显示了在MAYO和FNUSA数据集上的癫痫发作检测结果。Brant-2在这两个数据集上均取得了最佳的召回率和F2分数。

表2显示了在SleepEDFx和HMC数据集上的睡眠阶段分类结果。Brant-2和SleepHGNN表现出相当的性能。

表3显示了在临床数据集上的癫痫发作预测任务的平均性能。Brant-2在F1和F2分数上均取得了第一名，与Brant相比分别提高了37.97%和32.02%。

表4包含了SEED数据集上的情绪识别结果。Brant-2模型获得了最佳结果。

表5显示了运动想象分类的性能。Brant-2在准确率和F1分数上均取得了最高的成绩。

可扩展性分析

为了研究Brant-2模型的可扩展性，研究者预训练了两个较小版本的模型，分别含有2亿和4.6亿参数，并按照上文中描述的相同训练配置进行训练。然后在所有五个下游任务上评估这些模型，每个任务使用一个数据集。

图4(a)显示了两个预训练目标的训练损失（每5k步计算一次损失）。观察到随着训练的进行，不同大小的模型的训练损失持续下降；增加模型大小时，损失下降得更快。这表明Brant-2在预训练期间表现出可扩展性。另外如图4(b)所示，更大的模型在所有任务中都获得了更好的性能，表明模型的可扩展性在一系列下游任务中得到了转移。

消融研究

为了评估模型架构设计和预训练任务的有效性，研究者设置了三种模型变体来验证架构设计的有效性：1) 移除时间编码器的Brant-2；2) 移除空间编码器的Brant-2；3) 用普通的Transformer编码器块替换Brant-2中的多前馈网络（multi-FFN）Transformer编码器块。对于每种模型变体，控制模型的参数数量大致相同以确保公平比较。为了说明两种预训练任务的有用性，研究者进行了两组实验：4) 仅使用预测任务进行预训练的Brant-2；5) 仅使用掩码预测任务进行预训练的Brant-2。

图5显示了消融研究的结果，其中Brant-2在所有五个下游任务中都优于其他变体，证明了工作每个组件的有效性。没有时间编码器的Brant-2在这些下游任务中表现出整体较差的性能，突出了时间依赖性对脑信号的重要性。在某些任务（例如，癫痫发作检测、情绪识别）中，Brant-2 w/o mask优于Brant-2 w/o forecast，表明这些任务需要更好地理解信号段内的模式。另一方面，在某些任务（例如，癫痫发作预测）中，Brant-2 w/o forecast表现更好，表明这些任务更优先考虑基于当前观察序列预测未来变化。

标签稀缺场景探索

前文的结果表明Brant-2可以很好地泛化到各种任务。作为一个基础模型，研究者还旨在研究模型是否能解决对标签过度依赖的问题，并适用于标签稀缺的场景。选择临床数据集进行实验，这些数据集来自癫痫的真实世界临床场景，其中注释成本很高。通过选择这个数据集，旨在密切模拟真实世界场景，并解决临床设置中与昂贵注释相关的挑战。

图6显示了随着训练标签减少模型性能的变化。总体而言随着训练标签的减少，性能表现出一定程度的下降。当从100%过渡到1%标签时，Brant-2和Brant-2-460M的F1和F2分数下降了不到10%和15%，相比之下，Lopes等人[27]的F1和F2分数分别下降了50.6%和32.6%。结果表明，Brant-2可以减少对标签的依赖，从而在标签稀缺的场景中保持性能稳定。

通过深入分析，Brant-2模型不仅在脑信号分析的多个任务中展现出卓越的性能，还证明了其在标签稀缺场景下的鲁棒性和泛化能力，为未来的脑科学研究和临床应用提供了一个强大的基础工具。

论文链接：https://arxiv.org/abs/2402.10251