跳转至

B[FM]\(^2\): Brain Foundation Model via Flow Matching with SplitUNet

会议: NeurIPS 2026
arXiv: 2606.20812
领域: 医学图像
关键词: 脑电基础模型, 流匹配, 时空解耦卷积, 连续时间生成预训练, 临床神经电生理

一句话总结

本文提出脑电基础模型 B[FM]\(^2\),摒弃传统离散分块与掩码重建范式,直接利用最优传输条件流匹配(OT-CFM)与专为时间-电极不对称性设计的 SplitUNet 在原始连续多通道 EEG 波形上预训练,仅用约 307 小时预训练数据便在 9 项下游任务中的 7 项取得 SOTA,且无条件生成的脑电波形在双盲测试中令认证神经科医师无法与真实临床脑电区分(\(\kappa = -0.096\))。

研究背景与动机

脑电(EEG)基础模型旨在从大规模无标注脑电语料库中学习通用神经表征,从而以单一预训练骨干网络迁移至癫痫发作检测、睡眠分期、脑机接口运动想象等多样化临床与神经工程任务。然而,现有主流脑基础模型(如 BIOT、LaBraM、CBraMod、REVE、CSBrain)几乎无一例外地直接照搬了计算机视觉与自然语言处理中的训练范式——将连续的多通道神经电信号强行切分为固定长度的时间块(patches)或量化为离散码本 token,再基于 Transformer 进行离散掩码重建或对比预测。

这种离散化处理在头皮脑电信号上引发了根本性的表征失配。真实的脑电节律是连续演化的动态生理过程,人为划定的分块边界不仅割裂了连续脑电波形的相位与时频连续性,还迫使自监督目标去建模离散片段之间的拼接过渡,而非底层的连续神经动力学演化;同时,随机块掩码会抹除对临床异常诊断至关重要的细粒度瞬态波形变化(如尖波、棘波等微秒至毫秒级放电特征)。如果转向连续时间生成式预训练,多通道 EEG 又给骨干网络带来了独特的结构性矛盾:时间维度采样密集且高度自相关(数千个时间步),而电极维度极短(标准 10–20 系统仅 19 个通道)且每个通道对应明确的解剖脑区位置。直接套用 Transformer 需要重新引入分块离散化,单通道 1D UNet 会丢失跨电极的空间耦合信息,而常规 2D UNet 则会在短促的电极轴上浪费大量跨轴卷积参数并在下采样时破坏脑区电极拓扑。

为了同时解决离散化表征失配与时间-电极维度的结构不对称性,本文跳出掩码 token 重建框架,直接将连续时间生成模型作为脑电基础模型的预训练引擎。核心 idea:抛弃分块、分词与掩码机制,利用最优传输条件流匹配(OT-CFM)直接在原始多通道连续脑电波形上学习从高斯噪声到真实脑电分布的确定性速度场,并设计将二维卷积解耦为“1D 时间卷积 + 1D 电极卷积”且仅沿时间轴下采样的 SplitUNet 骨干网络,以极高的数据效率同时实现可迁移通用表征学习与临床级逼真脑电生成。

方法详解

整体框架

B[FM]\(^2\) 的整体流程分为连续时间生成预训练、无条件临床脑电采样以及下游特征提取与微调三个环节。输入为多通道原始脑电波形 \(\boldsymbol{x} \in \mathbb{R}^{E \times T}\)(其中 \(E\) 为电极通道数,\(T\) 为时间采样点数)。在预训练阶段,模型首先通过最优传输条件流匹配(OT-CFM)在 mini-batch 内利用匈牙利算法将各向同性高斯噪声 \(\boldsymbol{x}_0\) 与真实脑电片段 \(\boldsymbol{x}_1\) 进行最优配对并构造直线插值轨迹 \(\boldsymbol{x}_t\);随后将插值信号与流时间步 \(t\) 输入 SplitUNet 速度场网络,依次经过 4 个仅压缩时间维度的编码阶段(每阶段内部由 (1+1)D 时空解耦卷积残差块与线性注意力组成)、时间轴压缩 \(8\times\) 后的全局自注意力瓶颈层以及镜像上采样解码器,回归预测目标速度场 \(\boldsymbol{x}_1 - \boldsymbol{x}_0\)。在下游迁移时,直接将目标任务脑电输入预训练好的 SplitUNet 并在固定时间步 \(t=0\) 处前向传播,提取倒数第二层(最终输出投影前)的激活张量(形状为 \(C \times E \times T\)),经时空池化或展平后接入单层线性分类头完成端到端微调。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["多通道原始脑电 x1 与高斯噪声 x0<br/>形状: E × T"] --> B["最优传输条件流匹配(OT-CFM)<br/>连续波形预训练"]
    B --> C["(1+1)D 时空解耦卷积算子<br/>1D 时间卷积 + 激活 + 1D 电极卷积"]
    C --> D["仅时间轴压缩与分级注意力架构<br/>4 级时间下采样 + 线性/瓶颈自注意力"]
    D -->|预训练回归 v_θ(x_t, t)| E["预测流场速度 x1 - x0<br/>ODE 积分生成逼真临床脑电"]
    D -->|下游迁移固定 t = 0| F["提取倒数第二层表征 (C × E × T)<br/>池化/展平 + 线性分类头微调"]

关键设计

1. 最优传输条件流匹配(OT-CFM)连续波形预训练:直接在原始连续多通道脑电波形上学习流场,消除离散分块与掩码带来的表征失配

针对现有方法将连续脑电切分为离散 patch 或码本 token 导致细粒度时序动态丢失的痛点,B[FM]\(^2\) 直接在原始信号空间 \(\mathbb{R}^{E \times T}\) 中定义连续时间流匹配目标。模型学习一个时变速度场网络 \(v_\theta(\boldsymbol{x}, t): \mathbb{R}^{E \times T} \times [0, 1] \rightarrow \mathbb{R}^{E \times T}\),将 \(t=0\) 处的标准高斯先验分布 \(q_0\) 传输至 \(t=1\) 处的真实脑电分布 \(q_1\)。若采用独立随机配对,不同噪声-数据对在 \(\mathbb{R}^{E \times T}\) 空间中的直线插值轨迹会发生交叉,导致交叉点处的边缘速度场被迫平均相互冲突的条件速度方向,从而推高梯度方差并使学习到的流轨迹弯曲。为此,本文引入小批量最优传输耦合 \(\pi\)(OT-CFM):在每个 mini-batch 内以平方欧氏距离为代价矩阵,使用匈牙利算法(Hungarian algorithm)精确求解最优分配,将每个噪声样本 \(\boldsymbol{x}_0\) 与批次内距离最近的真实脑电片段 \(\boldsymbol{x}_1\) 配对。在此基础上构造确定性直线插值 \(\boldsymbol{x}_t = (1-t)\boldsymbol{x}_0 + t\boldsymbol{x}_1\),对应的条件速度目标简化为常数位移向量 \(u_t(\boldsymbol{x} \mid \boldsymbol{x}_0, \boldsymbol{x}_1) = \boldsymbol{x}_1 - \boldsymbol{x}_0\),条件流匹配损失定义为:

\[\mathcal{L}_{\mathrm{CFM}}(\theta) = \mathbb{E}_{t \sim \mathcal{U}[0,1],\; (\boldsymbol{x}_0, \boldsymbol{x}_1) \sim \pi} \left\lVert v_\theta\bigl((1-t)\boldsymbol{x}_0 + t\boldsymbol{x}_1, t\bigr) - (\boldsymbol{x}_1 - \boldsymbol{x}_0) \right\rVert^2\]

这种设计使模型在所有电极和所有时间步上同步学习从噪声到清晰波形的去噪流场,不仅倒逼骨干网络捕获跨尺度的连续神经节律,而且生成的速度场更加平直,无需额外蒸馏即可通过常规 ODE 求解器实现少步高质量采样。

2. (1+1)D 时空解耦卷积算子:将二维时空卷积拆分为时间维与电极维一维卷积串联,以更少参数高效捕获跨通道耦合

针对单通道 1D 卷积忽略电极间空间相关性、而常规 2D 卷积在非对称的(电极 \(\times\) 时间)网格上参数冗余且易过拟合的痛点,SplitUNet 将标准 DDPM UNet 中的每一个 \(k_t \times k_e\) 二维卷积替换为解耦的 \(\mathrm{conv}_{(1+1)\mathrm{D}}\) 算子。对于特征张量 \(\boldsymbol{z} \in \mathbb{R}^{C \times E \times T}\),该算子先沿时间轴滑动核长为 \(k_t\) 的一维时间卷积(权重在所有电极间共享),经过非线性激活函数 \(\sigma\) 后,再沿按国际 10–20 系统顺序排列的电极轴滑动核长为 \(k_e\) 的一维电极混合卷积(权重在所有时间步间共享):

\[\mathrm{conv}_{(1+1)\mathrm{D}}(\boldsymbol{z}) = \mathrm{conv}_{\text{electrode}}\bigl(\sigma\bigl(\mathrm{conv}_{\text{time}}(\boldsymbol{z})\bigr)\bigr)\]

具体实现中,输入首层采用核大小为 \((1, 7)\)、填充为 \((0, 3)\) 的宽视野时间卷积以扩大初始时间感受野,后续残差块内电极混合卷积采用 \((3, 1)\) 核与 \((1, 0)\) 填充。相比未解耦的 2D 卷积核,这种 (1+1)D 解耦具有双重优势:一是在时间滤波与空间电极混合之间额外插入了一次非线性激活 \(\sigma\),在不增加任何参数的前提下将非线性网络深度翻倍;二是将单个滤波器的权重数量从 \(k_t \cdot k_e\) 压缩至 \(k_t + k_e\),把节省下来的参数预算重新分配给特征通道宽度与网络深度,恰好契合脑电信号中局部时间动态与跨电极空间耦合高度可分离的物理特性。

3. 仅时间轴压缩与分级注意力架构:全程保留电极拓扑维度,结合阶段内线性注意力与瓶颈层全局自注意力兼顾长程时序与空间交互

针对多通道脑电中电极轴极短(19 通道)且每个通道具有明确脑区解剖语义、一旦沿电极维下采样就会混叠不同脑区特征的痛点,SplitUNet 设计了严格的不对称多尺度层级结构。编码器包含 4 个阶段,基础通道宽度为 32,通道倍增系数为 \((1, 2, 4, 8)\),每个子块通过自适应层归一化(AdaLN)注入正弦流时间嵌入 \(t\) 的缩放与偏移调制。在空间-时间分辨率控制上,编码器与解码器仅沿时间轴进行下采样与上采样(采用核大小 \((1, 4)\)、步长 \((1, 2)\)、填充 \((0, 1)\) 的操作),在网络每一层都完整保留电极维度 \(E\) 不变;这使得模型天然具备通道数无关性(electrode-count-invariant),无需重初始化首层即可直接适配不同电极数量(如 4 至 64 通道)的下游脑电数据集。在注意力排布上,由于高分辨率阶段时间步 \(T\) 极长,各编码/解码阶段内部采用具有线性复杂度的多头线性注意力模块(4 头,每头 32 维)进行高效混合;仅在时间轴已被压缩 \(8\times\) 的网络最深瓶颈层(bottleneck)放置单个全局全自注意力模块(full self-attention),此时时间维度与电极维度尺度相当,能够以可控计算开销完成全局时空联合交互。

损失函数 / 训练策略

  • 预训练数据集与预处理:在 Temple University Hospital EEG Corpus(TUEG)上进行预训练。原始脑电重采样至 \(200\text{ Hz}\),经 $0.3\text}75\text{ Hz}$ 带通滤波与 \(60\text{ Hz}\) 陷波滤波后切分为无重叠的 \(30\text{ s}\) 片段(保留国际 10–20 系统除耳电极 A1/A2 外的 19 个标准通道),剔除绝对幅值超过 \(100\,\mu\text{V}\) 的伪影片段并除以 \(100\,\mu\text{V}\) 归一化至 \([-1, +1]\)。从预处理池中仅随机抽取 \(36,895\) 个片段(约 \(307\text{ h}\),仅占预处理池的 \(3.3\%\)、TUEG 总时长的 \(1\%\))作为预训练集。
  • 预训练超参数:模型总参数量仅 \(6\text{ M}\)。使用 Adam 优化器(\(\beta_1=0.9, \beta_2=0.99\),无权重衰减),恒定学习率 \(8 \times 10^{-5}\),梯度裁剪阈值 \(\lVert g \rVert_2 \le 1.0\)。在单张 NVIDIA H100 80GB GPU 上采用 fp16 混合精度训练 70 个 epoch,batch size 为 4,每 10 步更新一次衰减率为 \(0.995\) 的指数移动平均(EMA)权重用于生成采样与下游特征提取。
  • 下游微调策略:固定流时间输入 \(t=0\),取倒数第二层特征 \(C \times E \times T\) 经全局平均池化(mean pooling)或展平(flattening)后送入线性分类头,使用交叉熵损失与 AdamW 优化器(\(\beta_1=0.9, \beta_2=0.999\),权重衰减 \(0.01\))端到端微调 40 个 epoch,batch size 为 32,在 \(\{5\times 10^{-5}, 10^{-4}, 2\times 10^{-4}\}\) 中按验证集平衡准确率(Balanced Accuracy)选择最优学习率与聚合方式。

实验关键数据

主实验

在涵盖精神疾病诊断(Mumtaz)、心理压力检测(MAT)、癫痫发作检测(Siena)、睡眠分期(ISRUC、HMC)、癫痫样事件分类(TUEV)、异常脑电检测(TUAB)、运动想象(BCIC-IV-2a)与情绪识别(SEED-V)的 9 个标准下游基准上,与 7 种从头训练的专用监督模型及 5 种主流预训练脑电基础模型进行对比(指标为 5 个随机种子的平衡准确率 Balanced Accuracy \(\pm\) 标准差):

方法 类型 / 预训练时长 Mumtaz MAT Siena ISRUC HMC TUEV TUAB BCIC-IV-2a SEED-V 9 任务平均
EEGConformer 监督从头训练 0.931 0.680 0.756 0.740 0.715 0.407 0.776 0.470 0.354 0.648
BIOT 脑电基座模型 0.936 0.688 0.735 0.753 0.686 0.528 0.796 0.475 0.384 0.664
LaBraM-Base 脑电基座 (\(2,500\text{ h}\)) 0.941 0.691 0.708 0.763 0.728 0.641 0.814 0.487 0.398 0.686
CBraMod 脑电基座 (\(9,000\text{ h}\)) 0.956 0.726 0.732 0.786 0.727 0.667 0.829 0.514 0.409 0.705
REVE 脑电基座 (\(60,000\text{ h}\)) 0.964 0.766 0.740 0.782 0.740 0.676 0.832 0.640 0.405 0.727
CSBrain 脑电基座模型 0.964 0.756 0.766 0.792 0.735 0.690 0.817 0.566 0.420 0.723
B[FM]\(^2\) (本文) 流匹配基座 (\(\approx 307\text{ h}\)) 1.000 0.840 0.776 0.806 0.764 0.715 0.819 \(\underline{0.570}\) 0.492 0.754

此外,在双盲临床神经科医师评估实验(Table 3)中,两位执业认证神经科医师对随机混排的 25 段真实 TUEG 留出片段与 25 段 B[FM]\(^2\) 无条件生成片段进行 1–5 分真实度评分(\(1=\) 确定真实,\(5=\) 确定伪造):

评估医师 真实脑电均值 (\(\downarrow\) 更真) B[FM]\(^2\) 生成均值 (\(\downarrow\) 更真) Mann-Whitney \(U\) \(p\) 值 医师间 Spearman \(\rho\) 真实/合成判别 Cohen's \(\kappa\)
医师 R1 \(2.96 \pm 1.02\) \(2.60 \pm 1.04\) 251.5 0.204 0.052 -0.096
医师 R2 \(3.08 \pm 1.12\) \(2.32 \pm 0.80\) 188.5 0.011 - -
汇总 (Pooled) \(3.02 \pm 1.06\) \(2.46 \pm 0.93\) 876.0 0.007 0.052 -0.096

消融实验

固定相同的 OT-CFM 预训练目标与 \(307\text{ h}\) 数据集,对比 4 种不同空间耦合方式的 UNet 速度场骨干网络在无条件生成质量(5 个标准频段对数功率谱分布的平均 Wasserstein-1 距离 \(W_1 \downarrow\),\(n=5,000\))与 9 项下游分类任务上的表现(Table 4 & Table 5):

速度场骨干网络配置 频段功率 \(W_1 \downarrow\) PSNR \(\uparrow\) Mumtaz MAT TUEV BCIC-IV-2a SEED-V 9 任务平均准确率 \(\uparrow\) 说明
Real-vs-Real (真实下界) \(0.156 \pm 0.028\) 16.32 - - - - - - 不同真实数据分片间的噪声基线
UNet-1D (indep.) \(1.903 \pm 0.098\) 6.74 0.896 0.588 0.675 0.458 0.485 \(0.674 \pm 0.153\) 通道完全独立处理,无跨电极混合
UNet-1D (joint) \(0.256 \pm 0.002\) 16.88 0.895 0.711 0.559 0.420 0.438 \(0.682 \pm 0.170\) 仅在首层卷积将电极视为输入通道混合一次
UNet-2D \(0.248 \pm 0.001\) 17.14 0.909 0.715 0.519 0.468 0.470 \(0.686 \pm 0.161\) 每层使用未解耦的 2D 时空联合卷积
SplitUNet (本文) \(0.195 \pm 0.002\) 17.16 1.000 0.840 0.715 0.570 0.492 \(0.754 \pm 0.150\) 每层交替 1D 时间与 1D 电极卷积,平均提升 6.8 pp

关键发现

  • 极高的预训练样本效率:B[FM]\(^2\) 仅使用约 \(307\text{ h}\)(\(36,895\) 个片段)的预训练数据,规模比 LaBraM(\(2,500\text{ h}\))、CBraMod(\(9,000\text{ h}\))和 REVE(\(60,000\text{ h}\))小 1 至 2 个数量级(约 \(30\times\) 至 \(200\times\)),却在 9 个下游任务中的 7 个刷新 SOTA,并在 MAT(\(+7.4\text{ pp}\))和 SEED-V(\(+7.2\text{ pp}\))上取得大幅领先,证明连续波形流匹配比离散掩码重建具有显著更高的数据利用效率。
  • (1+1)D 解耦显著优于纯 1D 与联合 2D 卷积:完全不混合电极的 UNet-1D (indep.) 在生成频谱距离(\(W_1 = 1.903\))和运动想象/压力检测任务上严重崩塌;仅首层混合的 UNet-1D (joint) 与全层二维卷积的 UNet-2D 下游平均准确率仅分别为 \(0.682\) 和 \(0.686\),在依赖精细时空耦合的 TUEV、MAT、BCIC-IV-2a 上比 SplitUNet 低 $10\text{20\text{ pp}$,验证了逐层解耦交替混合与保留电极拓扑的必要性。
  • 跨领域迁移边界与临床级生成保真度:在唯一未夺冠的运动想象数据集 BCIC-IV-2a 上,B[FM]\(^2\)(\(0.570\))仅次于预训练中包含 4 个运动想象数据集(相当于域内微调)的 REVE(\(0.640\)),且显著优于同样仅在临床静息/病理脑电 TUEG 上预训练的其他基线。同时,两位神经科专家在双盲评估中的判别一致性 \(\kappa = -0.096\),甚至给生成脑电打出了比真实留出片段更逼真的平均分(\(2.46\) vs. \(3.02\)),且无条件生成池中可自动涌现出尖波瞬态、\(\delta\) 慢波睡眠活动和闭眼枕叶 \(\alpha\) 节律等典型生理模式。

亮点与洞察

  • 用连续流匹配替代离散掩码重建的范式转换:本文最令人眼前一亮的洞察在于指出“将连续脑电切块 token 化再做 MAE 掩码重建”是借用视觉/语言模型时的权宜之计而非最优归纳偏置。通过在原始多通道波形上直接做 OT-CFM 速度场回归,仅用 \(6\text{ M}\) 参数和 \(307\text{ h}\) 数据就击败了使用上万小时数据预训练的离散 Transformer 基座。
  • 针对非对称传感器阵列的 SplitUNet 架构设计:多通道生理信号普遍存在“时间轴极长、传感器通道数少且具特定拓扑含义”的不对称特性。SplitUNet 通过“仅沿时间下采样 + 保持通道维恒定 + (1+1)D 时空卷积解耦”,既避免了空间脑区信息在池化中混叠,又让同一个预训练骨干无需修改网络结构即可无缝迁移到任意通道数(4 到 62 通道)的下游设备。
  • 兼具强判别表征与可解释生成能力的统一模型:以往基于掩码 token 的脑电基础模型无法在原始信号空间自由生成完整的新波形,难以直观检验模型究竟学到了什么神经动力学知识;而 B[FM]\(^2\) 不仅能提取 \(t=0\) 特征用于下游分类,还能直接采样生成经受住临床神经科医生双盲检验的多样化生理/病理波形。这一思路可直接迁移至脑磁图(MEG)、皮层脑电图(ECoG)、多导联心电图(ECG)等结构化传感器时序数据。

局限与展望

  • 预训练语料规模与领域覆盖面单一:尽管 \(307\text{ h}\) 的 TUEG 子集证明了流匹配的高样本效率,但预训练仅局限于单一临床脑电语料库和标准 19 通道 10–20 导联,未纳入任务态(如运动想象、认知任务)或高密度脑电数据,导致在 BCIC-IV-2a 等运动想象任务上仍落后于多语料库预训练的 REVE。
  • 电极排列顺序的隐式依赖与下游不平衡指标波动:SplitUNet 的一维电极卷积(核大小 \(3 \times 1\))沿电极轴滑动,隐式依赖电极在一维列表中的排列顺序(如标准 10–20 顺序),当下游数据集采用完全不同的通道子集或导联配置(如 4 通道 HMC 或 62 通道 SEED-V)时,一维滑动邻域对应的物理空间相邻关系可能发生偏移;此外在类别极度不平衡的二分类任务(如 Siena 和 MAT)中,虽然平衡准确率达 SOTA,但附录显示其 AUC-PR 低于部分基线,说明分类头阈值校准与概率输出仍有优化空间。
  • 未来方向:一是引入显式的三维电极空间坐标编码或图卷积来替代固定顺序的 1D 电极卷积;二是将无条件流匹配扩展为类别条件或受试者条件生成模型,用于合成罕见癫痫发作或长尾精神疾病脑电样本以增广临床训练集。

相关工作与启发

  • vs LaBraM / CBraMod / CSBrain:这些主流脑电基础模型均将连续 EEG 切分为固定时间片并通过神经分词器或块掩码自编码器训练 Transformer,预训练需要数千小时数据且丢失了块内亚秒级连续动态;本文完全摒弃分块与掩码,直接在原始连续波形上用流匹配训练轻量级 SplitUNet,以约 \(1/30\) 的预训练数据量取得了显著更高的下游平均准确率。
  • vs REVE:REVE 通过 4D 电极位置编码与块掩码在涵盖 \(25,000\) 名受试者、\(60,000\text{ h}\) 的多数据集上预训练,在包含同源预训练数据的运动想象任务上表现优异;B[FM]\(^2\) 则证明了无需海量语料和复杂位置掩码,仅靠连续时间生成目标与保留电极维度的时空解耦架构即可在绝大多数临床与情感解码任务上实现反超,并额外具备真实波形生成能力。
  • vs 图像/时间序列扩散与流匹配表征学习(如 OT-CFM、CSDI):以往扩散/流匹配表征学习多集中于二维自然图像,或在时间序列中仅做特定任务的插补与预测;B[FM]\(^2\) 首次将 OT-CFM 扩展为多通道连续生理信号的基础模型预训练范式,并通过 SplitUNet 解决了多变量传感器信号与二维自然图像在空间-时间尺度上的结构差异。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次将最优传输条件流匹配与时空解耦 SplitUNet 引入多通道原始脑电基础模型预训练,彻底摒弃分块、分词与掩码。
  • 实验充分度: ⭐⭐⭐⭐⭐ 覆盖 9 个多样化下游脑电基准、4 种骨干消融对比、频谱 Wasserstein 距离评估以及两位认证神经科医师的双盲临床图灵测试。
  • 写作质量: ⭐⭐⭐⭐⭐ 动机阐述切中离散化脑电表征的要害,架构设计与脑电物理特性的对应关系清晰,图表与附录细节极为详实。
  • 价值: ⭐⭐⭐⭐⭐ 以极低的预训练算力与数据成本(单卡 H100、307 小时数据、6M 参数)实现 SOTA 迁移与临床级脑电合成,对神经科学与生理信号基座模型具有很强的范式启发意义。