Scaling Vision Transformers for Functional MRI with Flat Maps

AI summary

Overview

Research area: Self-supervised computer vision applied to neuroimaging — specifically, adapting Vision Transformers and masked autoencoders to functional MRI (fMRI) of human brain activity.

Technical level: Intermediate. The core idea is intuitive (turn 3D brain scans into 2D images so standard vision models work), but the paper assumes familiarity with Vision Transformers, masked autoencoders, self-supervised pretraining, and basic neuroimaging concepts like parcellations and cortical surfaces.

Scope: A study of how to best represent fMRI data for self-supervised foundation models, introducing a flattened-cortex representation, a model family (CortexMAE) trained on 2.1K hours of fMRI, the first open benchmark suite for fMRI foundation models (Brainmarks), and the first scaling laws for fMRI.

What This Paper Is About

Foundation models have transformed vision and language, but fMRI has been harder to adapt because brain activity is inherently 4D (3D space plus time) and noisy. Existing approaches either compress each brain volume into a few hundred region averages (lossy) or feed the raw 4D volume into a model (compute-hungry). This paper asks whether an intermediate representation — projecting the cortical surface onto a 2D "flat map" like an unfolded map of the Earth — gives models the best trade-off, and whether the usual scaling story (more data, bigger models) holds for brain data. It also addresses the field's reproducibility problem by releasing a standardized benchmark.

Key Contributions

1. CortexMAE, a multi-representation fMRI foundation model family. A spatiotemporal masked autoencoder trained on 2.1K hours of open fMRI data from the Human Connectome Project Young Adult dataset (980 subjects, 7.4M frames). Three variants are trained and directly compared: flat map, parcellation, and sparse cortical volume.

2. Flat map representation for fMRI. The first use of cortical flat map projections as the input format for a vision foundation model. Each 3D fMRI volume becomes a 2D image, making standard 2D spacetime ViT patch embeddings apply directly with no architectural surgery.

3. Brainmarks, the first open benchmark suite for fMRI foundation models. Seven public datasets covering subject-level trait prediction (ABIDE, ADHD-200, ADNI, PPMI, HCP-A age/sex) and dynamic cognitive state decoding (HCP-YA Task21, NSD COCO24), with a fixed evaluation protocol applied uniformly to six prior foundation models plus a functional connectivity baseline.

4. The first systematic scaling analysis for fMRI. Power laws fit for dataset size and model size, with downstream probe accuracy tracked as a function of both.

Main Findings

• Flat maps win on dynamic state decoding. On HCP-YA Task21 (21-way cognitive task classification), the flat map model reaches 98.9% accuracy versus 97.5% (parcellation) and 96.2% (volume). On NSD COCO24 (24-way object category decoding), flat maps reach 31.0% versus ~27.5% for both alternatives. This is the paper's main empirical support for a "goldilocks zone" of intermediate representations.

• Volume wins on age prediction. On HCP-A age quartile classification, the sparse cortical volume model reaches 53.4% versus 47.5% (flat) and 44.2% (parcellation). The authors suggest structural features like age-related cortical thinning may be leaking into the dense volume representation, which is a positive sign for future dense models but also a warning that the model may be reading structure rather than function.

• Parcellation is the most compute-efficient. Training takes 11 hours for the parcellation model, 28 hours for flat maps, and 50 hours for volume on a single H100, despite near-identical parameter counts (~85M encoder) and FLOP counts. Dense models are bottlenecked by data loading, not compute.

• Strict power-law data scaling, but weaker than language. Reconstruction loss on held-out HCP-YA data follows a clean power law with a data exponent of about -0.01, compared to roughly -0.1 for next-token prediction in language models. Scaling improves out-of-distribution (NSD) reconstruction too, but more slowly than the power law predicts, suggesting dataset diversity matters as much as dataset size.

• Model size saturates hard at 37M parameters. Scaling encoder depth from 3 to 15 layers improves both reconstruction and downstream accuracy, then flattens at depth 9. A relatively small model appears sufficient to model all of HCP-YA.

• Trait prediction is a null result. No single foundation model achieves clear state-of-the-art across the five subject-level trait datasets, model rankings vary by dataset, and all models struggle to beat a simple functional connectivity baseline. The authors call this the first benchmark to highlight this inconsistency.

• State prediction is robust and separates models cleanly. Rankings are consistent across both state datasets, most models beat the FC baseline (which sits at 82.4% on Task21 and 7.4% on COCO24), and the CortexMAE family outperforms all prior models by a large margin. State decoding is also more sample-efficient to benchmark because datasets are far larger.

• Input normalization is essential, and its absence explains competitors' poor results. Removing per-voxel coordinate normalization during pretraining drops Task21 accuracy to near chance. Several prior models (SwiFT, Brain-JEPA, NeuroSTORM) appear to have been pretrained without it.

• The brain's default mode network emerges on its own. The first principal component of the model's learned spatial position embedding closely matches the principal functional connectivity gradient, meaning functional network structure arises naturally during training rather than being hand-designed into the architecture.

• The model can denoise fMRI. Averaging 100 masked reconstructions over different random masks recovers large-scale spatiotemporal dynamics while leaving unstructured noise behind.

Methodology in Plain English

The researchers start with fMRI data that has already been processed through a standard surface-based pipeline, so brain activity is mapped onto a common cortical surface mesh. For the flat map variant, they copy that surface data onto a corresponding flat mesh and resample it to a regular 2D image grid. The result is a sequence of 2D images — effectively a video of the cortex unfolding over time — which a standard video Vision Transformer can consume without modification. Background pixels outside the cortex are excluded from the loss.

For comparison, they build two other variants of the same model. The parcellation variant averages the signal within each of 400 cortical regions, giving a 400-dimensional time series per frame. The volume variant keeps the 3D voxel grid but discards voxels outside a cortical gray matter mask, cutting the sequence length roughly fourfold versus a naive full-volume approach.

All three use the same masked autoencoder recipe: split the input into space-time patches, hide 90% of them using tube masking (masking the same spatial location across all time frames to prevent the model from interpolating between neighbors in time), and train encoder and decoder transformers jointly to reconstruct the missing patches. Pretraining uses 16-frame clips (16 seconds) with a temporal patch size of 4 frames, a ViT-B encoder, and 625K training steps.

For evaluation, the decoder is discarded and the frozen encoder is probed. Trait prediction uses a logistic regression probe on average-pooled embeddings with 100 random train-test splits, chosen because sample sizes are tiny (400-900 subjects, one scan each). State prediction uses a more sensitive attentive probe on unpooled embeddings with a single fixed split, since those datasets have tens of thousands of samples. The same protocol is applied to six prior foundation models and a classical functional connectivity baseline.

Why This Matters

Impact on research. This is the first work to systematically compare fMRI input representations within a single controlled setup, and the first to establish scaling laws for the modality. It also directly challenges a wide body of prior work: several published fMRI foundation models fail to beat a simple connectivity baseline on trait prediction, and the inclusion of state decoding benchmarks reveals that some models may be sensitive mainly to structural rather than functional signal. The open benchmark and open models should make future comparisons reproducible, which the field has lacked.

Real-world applications:

• Clinical diagnostics. Classification of autism, ADHD, Alzheimer's, and Parkinson's from resting-state scans, though the paper's own results suggest current foundation models are not yet reliable here.
• fMRI denoising. The mask-averaging reconstruction technique can clean up noisy scans using population priors learned from thousands of subjects, potentially improving any downstream analysis.
• Brain-computer interfaces and neurofeedback. Accurate cognitive state decoding from short 16-second clips is a prerequisite for real-time applications like detecting mental workload or attention states.
• Cognitive neuroscience tooling. A model that natively "sees" fMRI could act as a perceptual prosthesis for researchers, surfacing representations that are hard to inspect by hand.

Industry relevance. The paper is a case study in adapting a mature vision architecture to a new modality cheaply. The finding that a 37M-parameter model saturates on this dataset size suggests that fMRI foundation models may not need enormous compute budgets, which lowers the barrier for smaller labs and companies. The compute comparison across representations is directly actionable for anyone building neuroimaging pipelines, and the demonstration that standard augmentations (cropping, color jitter) do not transfer to fMRI is useful negative knowledge for practitioners.

Future Directions

1. Better intermediate representations. Flat maps do not universally win, and they cannot represent subcortical structures at all. The authors call for new representations that combine the efficiency of flat maps with subcortical coverage, and the paper's own ablation shows that adding subcortex to the parcellation model did not help — so the right way to include it remains open.

2. Scaling beyond single-source datasets. Data scaling follows a power law in-distribution but generalizes poorly out-of-distribution, implying that dataset diversity rather than sheer volume is the bottleneck. The authors report preliminary experiments with UKBB in the appendix but position multi-source pretraining as unfinished work.

3. fMRI-specific data augmentation. None of the standard image or video augmentations tested helped. Temporal TR scaling showed a modest effect, but the paper treats useful fMRI augmentation as an unsolved problem, alongside better techniques for focusing the model on fine-grained detail rather than low-frequency structure.

4. Making trait prediction work, or explaining why it does not. The gap between foundation models and the functional connectivity baseline on clinical diagnosis is unexplained. Closing it, or demonstrating that connectivity is genuinely the right inductive bias for subject-level traits, would clarify the value proposition of the whole foundation model approach in this domain.

Target Audience

Machine learning researchers working on foundation models or self-supervised learning for scientific and medical data will find the representation comparison and scaling analysis directly relevant. Neuroimaging methodologists and computational neuroscientists building or benchmarking fMRI models are the primary intended audience, since the paper's main deliverables are an open model family and an open benchmark that they can build on. Clinicians and translational researchers interested in functional neuroradiology will benefit from the honest framing of where current models fall short. Readers new to the area can follow the core idea without deep background, since the central innovation is a change of input format rather than a new architecture.

Authors’ abstract

We study the problem of training self-supervised foundation models for functional MRI. Our main contributions are: (1) we introduce a new model family (CortexMAE) trained using the masked autoencoder framework on 2.1K hours of open fMRI data, and (2) we release the first open evaluation suite (Brainmarks) for fMRI foundation models. Our core innovation is simple: we adapt the Vision Transformer to fMRI by first converting each 3D fMRI volume to a 2D map using a cortical flat map projection. We directly compare flat maps to both parcellation and volume-based representations. While each has its advantages, flat maps generally perform best. We perform the first systematic scaling analysis for fMRI and observe strict power law scaling, albeit with limits. Finally, we use Brainmarks to do controlled benchmark comparisons. On subject-level trait prediction, we report a challenging null result: no single model achieves clear state-of-the-art performance. Moreover, all models struggle to outperform a simple functional connectivity baseline. On cognitive state decoding, we observe more robust performance, and in this setting our CortexMAE family outperforms prior models by a large margin. Code, models, and datasets are available at https://github.com/MedARC-AI/CortexMAE and https://github.com/MedARC-AI/Brainmarks.