跳转至

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。

方法详解

整体框架

模型 = MAE-st (Feichtenhofer 等 2022) + 三种 patch embedding 互换的输入头。pipeline:(1) HCP-YA 数据预处理 (FreeSurfer / fMRIPrep 表面映射) → 拿到皮层表面网格上的时间序列;(2) 用 pycortex 把表面投到平面网格 → 得到 16 帧 × 224 × 560 的 fMRI flat map 视频;(3) 切 \(p_t \times 16 \times 16\) 时空 patch(默认 \(p_t = 4\)),mask 比例 0.9(tube masking 不跨时间插值);(4) ViT-B encoder 看到 sparse 观测 patch + [MASK] token,decoder 重建被 mask 的 patch,loss 用 MSE(只在非背景像素上);(5) 预训练完成后丢掉 decoder,encoder 输出作为特征,配合 linear probe / attentive probe 做下游 trait / state 预测。同步训练 parcellation MAE (用 Schaefer-400 脑区) 和 volume MAE (用 4D patch) 做严格对照。

关键设计

  1. Cortical Flat Map Patch Embedding (核心创新)

    • 功能:把 3D fMRI 体积转成 2D 图像视频,让标准 spacetime ViT 直接处理,又不丢失皮层全信号。
    • 核心思路:先用神经科学常用的 surface-based 流程 (FreeSurfer + fMRIPrep) 把每帧 fMRI 信号从 3D voxel 映射到皮层表面的 mesh vertex;再用 pycortex 的 flat map 把这张曲面 mesh 展平到 2D 平面 grid(左右半球分别展开拼成一张图);最后 resample 到固定 224×560 网格。每个时间步就是一帧 2D 图,16 帧拼起来就是 ViT 的 spacetime 输入。patch 大小 \(p_t \times 16 \times 16\),背景 (脑外区域) 全 0 patch 直接排除不参与计算,MSE loss 也只在非背景像素上算。
    • 设计动机:parcellation 把信号压成 ~400 维向量(dim 损失 ~100×),而 volume 直接处理需要 ~132K voxel 的稀疏 patch(计算量大且大部分是无效背景)。flat map 在二者之间——保留全皮层 ~77K 维信号但用 2D ViT 高效处理;序列长度 364(vs volume 的 465 和 parcellation 的 400)几乎持平但训练带宽与数据吞吐都更好。论文 Figure 1 把这个 trade-off 形象化为"光谱"。
  2. 三种表示的 head-to-head 对比设计

    • 功能:用同一架构、同一预训练数据、同一评测协议,把 parcellation/flat/volume 三种表示放在同一起跑线。
    • 核心思路:所有模型都用 ViT-B encoder、同样的 16 帧输入、同样的 0.9 mask ratio;只是 patch embedding 不同——parcel 用 \(p_t \times 1\) 仅时间维 patch、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)。
    • 设计动机:之前 fMRI 基础模型论文从不公平对比表示——通常自己只用一种然后宣称 SOTA。本文是第一份 multi-representation fMRI MAE family,结论可信度高得多。
  3. Brainmarks 开源评测套件

    • 功能:给 fMRI 基础模型领域提供一个所有方法都能跑、所有数据集都标准化处理的 benchmark。
    • 核心思路:纳入 6 个现有 fMRI 基础模型 (SwiFT, BrainLM, Brain-JEPA, BrainHarmonix-F, NeuroSTORM, Brain-Semantoks) 和 CortexMAE family;包含 7 个公开数据集——4 个临床诊断 (ABIDE/ADHD200/ADNI/PPMI) + HCP-A 年龄/性别 + HCP-YA Task21 + NSD COCO24;对小样本 trait prediction 用 linear probe + 100 次 train-test 随机分割,对大样本 state prediction 用 attentive probe 单 fixed split + 49 个学习率 grid 调优。所有方法用同一 probe 协议,避免每家自己 fine-tune 偷跑。
    • 设计动机:fMRI 模型评测的复现性灾难是社区公认问题;统一 protocol 才能让"谁真的更好"有意义。还专门设计了 NSD COCO24 这种"短 trial 重叠 + 不同 subject 测试集 + 高难度视觉解码"任务来真正区分模型质量。

损失函数 / 训练策略

预训练:MAE MSE loss on masked patches; data normalization 是关键——每个 voxel/ROI 时间序列 z-score (coordinate norm) + 每帧空间 z-score (frame norm),去除 BOLD 信号只有 1-2% 波动的静态噪声;temporal patch \(p_t = 4\);training schedule 默认 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 基础模型社区是教科书级贡献,会成为后续工作必引论文。