Code: https://github.com/935963004/LaBraM
Data: 无
目录
- Abstract
- Introduction
- Method
- Neural tokenizer training:
- Pre-training LaBraM:
- Results
- Experimental setup:
- Pre-training result:
- Comparison with SOTA:
- Pre-training with/without downstream datasets:
- Scaling data size:
- Conclusion
- Related work
- Appendix
- Effectiveness of VQ:
- LaBraM without pre-training:
- Discussion
Abstract
本文介绍了一种新型的大型脑电图(EEG)模型,名为Large Brain Model(LaBraM),旨在克服传统基于EEG的深度学习模型在脑机接口(BCI)应用中的局限性,如模型规模有限、感知能力和泛化性不足。LaBraM通过无监督预训练来获得对EEG信号的通用感知能力,然后可以针对不同的下游任务进行微调。
LaBraM面临的挑战包括EEG数据集规模小、格式差异大、电极数量不匹配、数据样本长度不等、任务设计多样以及信噪比低。为了应对这些挑战,LaBraM通过将EEG信号分割成EEG通道片段,使用向量量化神经谱预测技术训练一个语义丰富的神经tokenizer,将连续的原始EEG通道片段编码成紧凑且离散的神经tokens。然后,通过预测掩蔽EEG通道片段的原始神经tokens来预训练神经Transformer。
LaBraM在约2500小时、来自约20个数据集的多种类型的EEG信号 上进行了预训练,并在多种下游任务上进行了验证,包括异常检测、事件类型分类、情感识别和步态预测。实验结果表明,LaBraM在各自领域的表现超过了所有比较的SOTA(State of the Art)方法。
Introduction
- 背景:当前的EEG模型缺乏跨任务的学习能力,泛化性较差。
- 动机:大语言模型(LLMs)的成功表明自监督掩码预训练的方式对于大规模数据应用的潜力,将重建思想应用于预训练神经Transformer可能对下游任务是有效的。
- 挑战:
- 缺乏足够的EEG数据:如何利用大量的未标记EEG数据集?以及多大的数据够用?
- 不同的EEG数据集有不同的采集配置:如何处理不同格式的EEG数据以匹配神经Transformer的输入层?
- 缺乏有效的EEG表征学习范式:EEG的低信噪比特点,以及如何平衡时间和空间特征?
- 贡献:
- 大规模EEG预训练:在超过2500小时的EEG数据上预训练
- 兼容各种EEG配置:LaBraM 是统一的模型,能够在灵活的辅助下处理各种通道和时间长度的 EEG 信号。预训练的 LaBraM 可以适应任何具有不同配置的下游数据集。
- 有效的EEG表征学习:神经 Transformer 使该模型能够有效捕获具有不同通道和长度的脑电图信号的时间和空间特征,使其适用于脑电图分析中的各种下游任务。并进一步定义了一个神经密码本(CodeBook),它提供了一种紧凑、通用且有意义的EEG信号表示。
- 下游数据集综合实验:作者在 BCI 的四个代表性下游任务上评估了 LaBraM,结果表明在很大程度上超越了所有为特定任务而开发的 SOTA 方法。
Method
Neural tokenizer training:
在通过掩码预测对 LaBraM 进行预训练之前,需要先将 EEG 标记为离散tokens。这一部分主要基于 Van Den Oord 的 VQVAE 架构实现。
- EEG数据token化:首先通过一个 tokenizer 模块将EEG根据时间和通道两个维度划分为patch矩阵
- tokens 的向量量化:基于VQVAE的思想,将上一步得到的patch通过最近邻的方式找到一个预定义的 CodeBook 上的向量来代替这个patch,更新patch矩阵
- 傅里叶谱预测:重建的目标被设定为EEG信号的傅里叶频谱的频率和相位分布(作者在预实验中发现直接重建EEG信号损失无法收敛,可能是SNR过低导致),作者认为这两种特征揭示了大脑潜在的神经生理活动(存在改进空间,这里的重建目标决定了后续预训练的下游任务的上限)。
- VQ的预测损失:包含两种预测目标的重建损失,以及 CodeBook 的更新损失
Pre-training LaBraM:
- 模型架构:时序 encoder → 叠加时序和空间 embedding → 大量Transformer模块 → Tokens预测头
- 掩码输入:同样地划分为patch矩阵,然后随机选取一定比例的patch进行掩蔽,同时为了提高训练效率和内存使用,将随机掩码的对称patch矩阵同时作为输入,也可以起到数据增广的作用。
- 预测目标:上一步得到的离散tokens。
Results
Experimental setup:
- Pre-training:在超过2500h的EEG数据上训练神经tokenizer和预训练 LaBram
- Fine-tuning:在四个下游数据集上进行全微调 LaBram
Pre-training result:
Comparison with SOTA:
表 1 和表 2 展示了最先进的baseline以及LaBraM在 TUAB 和 TUEV 的结果。结果表明,LaBraM-Base 模型在这两项任务的各种评估指标上均优于所有baseline。随着模型参数数量的增加,LaBraM-Huge 模型表现最好。作者认为有了足够的数据量,大规模脑电图模型可以学习更通用的脑电图模式,从而提高脑电图分析中各种下游任务的性能。
Pre-training with/without downstream datasets:
如图 4 所示,是否将下游任务数据集纳入模型的预训练过程不会显着影响模型在下游任务上的性能。这表明 LaBraM 模型具有学习通用脑电图表示的能力。
Scaling data size:
如图所示,随着数据规模的不断扩大,Huge模型的性能呈现出明显的上升趋势。这些结果基本遵循 scaling law,所以作者大胆推断,在数据量至少为一万小时的情况下,Huge 模型将继续表现得更好。
Conclusion
LaBraM是一个通过无监督预训练学习通用embedding的大型脑电图(EEG)模型,能够处理多样化的EEG数据集。LaBraM通过分割EEG信号和使用向量量化神经谱预测来生成丰富的语义tokenizer,并利用神经Transformer架构在大规模数据上预训练学习EEG信号的时空间特征,适用于多种EEG分析任务。在异常检测、事件分类、情感识别和步态预测等任务中,LaBraM的表现超越了当前最先进方法,有望推动EEG深度学习模型的发展。
Related work
- BCI 中的自监督学习:BrainBERT(Wang 等人,2023)掩蔽立体脑电图 (SEEG) 频谱图的随机部分,并使用 43.6 小时的数据生成原始嵌入。
Appendix
Effectiveness of VQ:
LaBraM without pre-training:
直接在下游数据集上从头开始训练 LaBraM,无需预训练,性能急剧下降证明了预训练的有用性。
Discussion
- Limitations:
- 虽然作者收集了有史以来最大的超过2500小时的脑电图数据集,并为BCI训练了有史以来最大的3.69亿参数的模型,但与当今的大型视觉模型和大型语言模型相比,仍然有很大的差距。
- LaBraM 需要进行全面微调以适应下游任务,这可能会耗费计算成本和内存成本。
- LaBraM 使用单峰脑电图数据进行训练。研究用其他方式训练大型脑电图模型是值得的。
- Outlook:
- 涌现能力:从各种BCI任务中收集更多的脑电数据,并训练更大的脑电模型,看看脑电模型中是否存在类似于大型语言模型的涌现能力
- 高效性:利用适配器、即时调优、LoRA等参数高效学习方法,减少微调开销,节省磁盘空间
- 多模态:将图像、语言、语音和其他生理信号等其他模态纳入大型脑电图模型训练中以构建新的范式,或将脑电图表示与语义空间中的其他模态对齐,这可能是未来工作的一个有意义且具有挑战性的方向。