跳转至

Scaling Vision Transformers for Functional MRI with Flat Maps

会议: ICML 2026
arXiv: 2510.13768
代码: https://github.com/MedARC-AI/CortexMAE & https://github.com/MedARC-AI/Brainmarks (有)
领域: 医学图像 / 自监督学习 / 神经影像基础模型
关键词: fMRI 基础模型、Cortical Flat Map、MAE、Brainmarks 评测、Scaling Law

一句话总结

把 3D fMRI 体积按"皮层展平图"投影成 2D 视频后直接喂给标准 spacetime MAE-ViT,得到一个在 2.1K 小时 HCP 数据上训练的 CortexMAE:在认知状态解码上大幅超 SOTA,验证 flat map 是体素 (volume) 和脑区平均 (parcellation) 之间的"goldilocks zone";同时发布首个开源 fMRI 基础模型基准 Brainmarks,给出 fMRI 模型的第一份系统 scaling law 与一个"个体特质预测仍打不过简单功能连接 baseline"的诚实 null result。

研究背景与动机

领域现状:神经科学社区想用 fMRI + 大模型解码大脑活动 (诊断、行为预测、视觉重建)。已有一波 fMRI 自监督基础模型 (BrainLM, Brain-JEPA, NeuroSTORM, SwiFT 等),多数用 parcellation 表示(把 3D 脑体积按 100-400 个脑区平均,得到 1D 时间序列向量);少数用 volume 表示(直接处理 4D 时空 MRI 数据)。

现有痛点:(1) Parcellation 计算便宜但信息损失严重——一整个 cm 量级的脑区被压成单标量,丢掉 99% 维度;(2) Volume 信息全保留但 sequence length 巨大(一个 fMRI volume 切 patch 后 ~2000+ tokens),训练计算/IO 开销爆炸;(3) 整个 fMRI 基础模型领域几乎没有可复现 benchmark——各家用自家 dataset、自家预处理、自家评测设置,谁强谁弱不可比;(4) 之前的 trait prediction 论文经常报告"我们 beat baseline X%",但用的 baseline 太弱,没和"简单功能连接 (FC) + 逻辑回归"这种 30 年前的方法做严肃对比。

核心矛盾:fMRI 数据本质是 4D 时空体积,标准 ViT 假定 2D 输入。要么花大代价直接学 4D(信息全但贵),要么用强归纳偏置 (parcellation) 损失信息——存在一个 "既保留全皮层信号又给 ViT 友好的 2D 输入" 的中间表示吗?

本文目标:(i) 找到 fMRI 的 "goldilocks" 输入表示;(ii) 用标准 ViT + MAE 训练一组对比清晰的基础模型;(iii) 建立开源、可复现的 fMRI 基础模型评测 (Brainmarks);(iv) 第一次系统做 fMRI 自监督的 data/model scaling law。

切入角度:神经科学早就有 cortical flat map ——把大脑皮层这个 2D 流形(其实皮层就是一层 2-4mm 厚的褶皱片)展平到平面网格上。这样既保留了全皮层 BOLD 信号(不像 parcellation 那样平均掉细节),又得到了一张 224×560 的 2D "图像",可以直接被 spacetime ViT 当 video 处理。

核心 idea:用 cortical flat map 把 3D fMRI 投影成 2D 视频,套用现成 MAE-st 训练,不改 ViT 架构,只换 patch embedding——一个简单但被错过的好选择,最终带来 SOTA + 第一份 fMRI scaling law + 第一个开源 benchmark。

方法详解

整体框架

整篇工作的核心是一个赌注:fMRI 本质是 4D 时空数据,但只要把它投影成对的 2D 表示,就能原封不动地复用现成的 spacetime MAE-ViT,不必为它重新设计架构。于是 pipeline 只有两件事——先把 3D fMRI 体积压成 2D 视频,再把视频喂进标准 MAE。具体地,HCP-YA 数据先经 FreeSurfer / fMRIPrep 的 surface 流程把每帧信号从 3D voxel 映到皮层表面网格,再用 pycortex 展平成 16 帧 × 224 × 560 的 flat map 视频;视频切成 \(p_t \times 16 \times 16\) 时空 patch(默认 \(p_t=4\)),以 0.9 的比例 tube-mask 掉,ViT-B encoder 只看剩下的 sparse patch,decoder 重建被遮的部分。预训练完成后丢掉 decoder,encoder 输出当特征接 linear / attentive probe 做下游预测。为了让"flat map 是不是好表示"这个问题有可信答案,作者同时用同一套架构训了 parcellation MAE 和 volume MAE 做严格对照,并把这套对比固化成开源 benchmark。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400, 'subGraphTitleMargin': {'top': 8, 'bottom': 16}}}}%%
flowchart TD
    A["HCP-YA fMRI<br/>3D 体积 · 2.1K 小时"] --> B["surface 流程映到皮层表面网格<br/>+ 逐体素/逐帧 z-score 归一化"]
    subgraph EMB["三种表示 head-to-head 对比(只换 patch embedding)"]
        direction TB
        F["Cortical Flat Map(本文核心)<br/>pycortex 展平 224×560 · pt×16×16"]
        P["parcellation<br/>Schaefer-400 · pt×1(信息损失大)"]
        V["volume<br/>稀疏皮层体素 · pt×8×8×8(序列长 · 计算贵)"]
    end
    B --> EMB
    EMB -->|"0.9 tube-mask"| G["ViT-B encoder<br/>只看可见 sparse patch"]
    G --> H["decoder 重建被遮 patch<br/>MSE(仅非背景像素)"]
    H -->|"预训练完丢 decoder"| I["encoder 特征<br/>linear / attentive probe"]
    I --> J["Brainmarks 评测<br/>7 数据集 + 6 外部模型统一协议"]

关键设计

1. Cortical Flat Map Patch Embedding:在体素与脑区平均之间找皮层信号的"金发姑娘"中间点

fMRI 表示一直被两难夹住:parcellation 把整块 cm 量级脑区平均成单标量,~400 维向量丢掉了 99% 维度;volume 直接处理 4D 数据信息全保留,但一个体积切完 patch 是 ~132K voxel 的稀疏序列,计算和 IO 都爆炸,而且大半是脑外无效背景。这篇论文的做法是借神经科学几十年的老工具——皮层本质就是一层 2-4mm 厚的褶皱片,是个 2D 流形,可以无损展平。于是先用 surface-based 流程把信号从 3D voxel 映到皮层表面 mesh,再用 pycortex 的 flat map 把左右半球分别展开拼成一张 224×560 的 2D 图,每个时间步一帧、16 帧叠成 ViT 的 spacetime 输入,背景全 0 patch 直接剔除、MSE loss 也只在非背景像素上算。这样既保留了全皮层 ~77K 维信号(不像 parcellation 那样平均掉细节),序列长度 364 又和 volume 的 465、parcellation 的 400 几乎持平,却因为是规则 2D 网格而训练带宽和数据吞吐都更优——论文 Figure 1 把这个 trade-off 画成一条从"全压缩"到"全保留"的光谱,flat map 正落在甜点位置。

2. 三种表示的 head-to-head 对比:把 parcellation / flat / volume 放上同一条起跑线

以往 fMRI 基础模型的通病是各家只用自己那一种表示就宣称 SOTA,谁强谁弱根本不可比。这里作者刻意让三个变体共享几乎所有变量——同样的 ViT-B encoder、同样 16 帧输入、同样 0.9 的 mask ratio,唯一的差别就是 patch embedding:parcel 用 \(p_t \times 1\) 只切时间维、flat 用 \(p_t \times 16 \times 16\)、volume 用 \(p_t \times 8 \times 8 \times 8\) 的 4D patch。每个变体独立训 8 次取均值,再一起跑下游 8 个数据集(4 个临床诊断 + 性别 + 年龄 + HCP-YA Task21 + NSD COCO24)。因为除表示之外的所有东西都被钉死,最后得到的差异就能干净地归因到"表示"本身,这也是第一份真正成 family 的 multi-representation fMRI MAE,结论可信度远高于单点比较。

3. Brainmarks 开源评测套件:用统一 probe 协议终结 fMRI 模型的复现性灾难

fMRI 领域长期是"各家用自家 dataset、自家预处理、自家评测设置",加上 trait 论文常拿过弱的 baseline 自夸,谁真的更好无从判断。Brainmarks 把这件事标准化:一边纳入 6 个现有基础模型(SwiFT、BrainLM、Brain-JEPA、BrainHarmonix-F、NeuroSTORM、Brain-Semantoks)和 CortexMAE family,一边统一 7 个公开数据集(4 个临床诊断 ABIDE/ADHD200/ADNI/PPMI + HCP-A 年龄/性别 + HCP-YA Task21 + NSD COCO24)。关键是 probe 协议对所有方法一视同仁——小样本 trait prediction 用 linear probe + 100 次随机 train-test 分割,大样本 state prediction 用 attentive probe + 单 fixed split + 49 个学习率 grid,谁都不许自己 fine-tune 偷跑。作者还特意设计了 NSD COCO24 这种"短 trial 重叠 + 测试集换不同 subject + 高难度视觉解码"的任务,专门用来把真正强的模型从弱的里区分出来。

损失函数 / 训练策略

预训练目标就是 masked patch 上的 MAE MSE,但有两步 normalization 是成败关键:每个 voxel/ROI 时间序列做 z-score(coordinate norm)压住静态体素差异,每帧再做空间 z-score(frame norm)去掉全局漂移——因为 BOLD 信号本身只有 1-2% 的波动,不归一化静态噪声就会淹没有用信号。其余超参:temporal patch \(p_t=4\)、默认 625K steps、batch 32(= 512 帧),并用 repeated sampling 缓解 IO 瓶颈。下游 trait prediction 用 average-pooled embedding + logistic regression 5-fold CV,state prediction 用 attentive probe + early stopping。

实验关键数据

主实验

三种表示在 8 个下游任务上的探针准确率(8 次预训练随机种子均值):

Dataset parcel flat volume FC 基线
ABIDE (ASD 诊断) 62.0 61.4 60.4 59.8
ADHD200 56.8 59.2 58.8 57.0
ADNI (AD) 61.6 62.4 64.3 58.6
PPMI (PD) 61.4 58.8 59.1 58.0
HCP-A Age 44.2 47.5 53.4 45.6
HCP-A Sex 71.2 87.4 86.3 81.9
HCP-YA Task21 (状态) 97.5 98.9 96.2 82.4
NSD COCO24 (视觉解码) 27.5 31.0 27.7 7.4

总结:(1) flat map 在动态状态解码上全面胜出(Task21、COCO24、性别);(2) volume 在年龄预测上有优势(可能跟 dense 体素能捕捉皮层厚度等结构线索有关);(3) parcel 最高效但状态解码弱;(4) 临床诊断 4 个数据集上所有方法几乎打平,且勉强超过 FC baseline——揭示样本太少时 fMRI 基础模型还看不出优势。

跨模型 controlled benchmark(Figure 8):trait prediction 上没有任何模型显著超过简单 FC 基线(包括 BrainLM、Brain-JEPA、NeuroSTORM 等 SOTA);state decoding 上 CortexMAE flat map 全面领先,比 NeuroSTORM 等 volume 模型在 NSD COCO24 上高 3-5 个百分点。

消融实验

配置 现象
Full flat map MAE baseline
不做 frame normalization 全局信号漂移污染,下游精度掉
不做 coordinate normalization 静态体素差异主导特征,状态解码崩
tube masking → random masking 时间内插泄露信息,重建任务变 trivial
mask ratio 0.5 → 0.9 高 mask ratio 强迫学结构性表示,下游更强
增大 encoder depth depth ~9 (37M 参数) 后 saturate
增大 pretrain data 在 HCP 内部数据严格遵循 power law (指数 -0.01),OOD NSD 上 saturate

关键发现

  • fMRI 严格遵循 data scaling law,但指数比语言模型弱十倍 (-0.01 vs Kaplan 2020 的 -0.1),意味着 fMRI 边际收益小,scaling 不会自动救场。
  • Model scaling 在 depth 9 (37M 参数) 就饱和——HCP-YA 这种 2K 小时数据集就只能撑住这点容量。
  • 模型自发学到大脑的默认模式网络 (DMN):position embedding 的第一主成分跟 Margulies 2016 的 FC principal gradient 几乎一致,证明 MAE 真的学到神经生物学有意义的结构。
  • 诚实 null result:所有 fMRI 基础模型在 individual trait prediction 上都打不过简单 FC + linear——这对整个领域是个警钟。
  • state decoding 上基础模型有 robust 优势,CortexMAE flat 是其中最强。

亮点与洞察

  • "换 patch embedding 就完事"是非常优雅的工程选择:不重新设计架构、不重写 attention,只是把输入从 3D 体积投到 2D 流形——任何 ViT 论文未来想入局 fMRI 都可以这样做。
  • goldilocks zone 的概念可迁移:在表示学习里"全保留 vs 大压缩"是经典 trade-off,cortical flat map 是一个利用领域几何 (cortex 本质是 2D) 找到完美中间点的例子,类似的可以套到 EEG (1D 时间 + 电极几何)、心电、显微图像等。
  • 诚实 null result + 开源 benchmark:神经影像社区长期有"小数据集 + 各家自评"的复现性灾难,作者直接发布 Brainmarks 并公开承认 trait prediction 打不过 FC baseline,这是社区急需的诚实声音。
  • 第一份 fMRI scaling law:让大家看清"fMRI 不是 NLP"——指数小十倍意味着堆数据收益有限,真正瓶颈可能是数据多样性而非规模。
  • DMN 自然涌现:自监督表征对应到神经生物学已知结构,是 fMRI MAE 的优秀诠释性证据。

局限与展望

  • HCP-YA 全是 22-35 岁年轻人 + 健康人群,预训练分布狭窄,OOD 泛化弱(论文图 7 已经展示了在 NSD 上 scaling 失效)。
  • 临床诊断结果(ABIDE、ADHD200 等)全 60% 上下徘徊,几乎没用——基础模型在小样本临床数据上无法 transfer,是社区共性难题,但本文没提供解决方案。
  • depth 9 就 saturate,说明现有数据规模不够——但 fMRI 数据采集极贵 (一小时成本几百美元),要 10× 数据需要全社区合作。
  • flat map 投影本身丢掉了皮层下结构 (subcortical regions like thalamus, basal ganglia),这些区域对很多临床任务很重要;volume 模型在这块有结构性优势。
  • 没探讨 multi-modal fMRI (task + rest + diffusion) 联合预训练,未来工作空间大。
  • 评测只在英语母语、北美人群数据上做,存在 demographic bias。

相关工作与启发

  • vs BrainLM (Caro et al. 2024) / Brain-JEPA (Dong et al. 2024):这些是 parcellation-based fMRI 基础模型,dim 损失大;CortexMAE flat 全皮层信号保留,在状态解码上显著更强。
  • vs SwiFT (Kim et al. 2023) / NeuroSTORM (Wang et al. 2025a):volume-based 模型,计算贵但年龄预测有优势;本文 volume MAE 也复现了这一点,说明 dense 表示有 niche,但状态解码上 flat 仍占优。
  • vs functional connectivity baselines:从 Finn et al. 2015 起,FC + linear 一直是 trait prediction 标准 baseline;本文证明深度 fMRI 模型至今没真正超越它,是对整个深度 fMRI 领域的一次"打脸"。
  • vs vision MAE (He et al. 2022):直接搬运 MAE-st,主要贡献在于"找到合适的 2D 投影把 fMRI 装进现成框架"——示范了"领域几何 + 通用架构"是个高效组合。

评分

  • 新颖性: ⭐⭐⭐⭐ 严格说 cortical flat map 是神经科学几十年的工具,作者首次系统用它做 ViT-friendly representation;技术上小但战略上聪明。
  • 实验充分度: ⭐⭐⭐⭐⭐ 三表示严格对比 + 6 个外部模型 + Brainmarks 7 数据集 + scaling law + 解释性分析,几乎是 fMRI MAE 的"实验白皮书"。
  • 写作质量: ⭐⭐⭐⭐⭐ 思路清晰、动机和结论都直白,关键 figure (光谱图、DMN 涌现) 极其有说服力。
  • 价值: ⭐⭐⭐⭐⭐ Brainmarks + null result + flat map 三件套对 fMRI 基础模型社区是教科书级贡献,会成为后续工作必引论文。