Scaling Vision Transformers for Functional MRI with Flat Maps

Connor Lane; Mihir Tripathy; Leema Krishna Murali; Ratna Sagari Grandhi; Shamus Sim Zi Yang; Sam Gijsen; Debojyoti Das; Manish Ram

cs.CVOct 15, 2025

Scaling Vision Transformers for Functional MRI with Flat Maps

Authors: Connor Lane, Mihir Tripathy, Leema Krishna Murali, Ratna Sagari Grandhi, Shamus Sim Zi Yang, Sam Gijsen, Debojyoti Das, Manish Ram, +10 more

Organizations: MedARC · Sophont · Baylor College of Medicine · University of Tübingen · Georgia Tech · University of Florida

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.

Figures & tables

Figure 1 : Spectrum of fMRI data representations: volume is the native 3D MRI format, flat map is our proposed representation based on cortical flat map projection, and parcellation data results from averaging signal within a set of regions (parcels). Volume patches are 3D cubes, flat map patches are 2D squares, and parcellation “patches” are single scalars. We hypothesize there is a “goldilocks zone” of representations that are neither too lossy nor too dense.

Figure 2 : MAE model architecture ( He et al., 2022 ) adapted to fMRI flat maps. Single frame input and prediction shown for convenience. Model inputs are temporal sequences of 16 frames.

Dataset	Target	Subjects	Samples	Seq length	TR	#Classes	Majority %
ABIDE ( Di Martino et al., 2014 )	ASD Dx	578:124:124	578:124:124	150	2.0s	2	55%
ADHD200 ( ADHD-200, 2012 )	ADHD Dx	301:64:65	301:64:65	150	2.0s	2	57%
ADNI ( Jack Jr et al., 2008 )	AD Dx	328:41:41	328:41:41	100	3.0s	2	77%
PPMI ( Marek et al., 2011 )	PD Dx	463:99:100	463:99:100	120	2.5s	2	62%
HCP-A ( Bookheimer et al., 2019 )	Age	455:53:52	455:53:52	500	0.7s	4	27%
HCP-A ( Bookheimer et al., 2019 )	Sex	471:58:55	471:58:55	500	0.7s	2	58%

Table 1: Summary of trait prediction (top) and state prediction (bottom) evaluation datasets. Dx = diagnosis classification, Age = quartile classification, Task21 = cognitive task state decoding, COCO24 = object category decoding. Subject and sample counts are train:validation:test. Trait prediction datasets include one sample per subject. For diagnosis datasets, controls are the majority class.

Figure 5 : CortexMAE learns the brain’s default mode network from scratch. (left) Map of default mode network (principal gradient, Margulies et al. 2016 ). (right) First principal component of the model’s learned spatial position embedding.

space	ABIDE	ADHD200	ADNI	PPMI	HCP-A Age	HCP-A Sex	HCP-YA Task21	NSD COCO24
parcel	62.0 $\pm$ 0.8	56.8 $\pm$ 0.6	61.6 $\pm$ 1.2	61.4 $\pm$ 1.3	44.2 $\pm$ 0.5	71.2 $\pm$ 1.0	97.5 $\pm$ 0.2	27.5 $\pm$ 0.5
flat	61.4 $\pm$ 1.3	59.2 $\pm$ 1.0	62.4 $\pm$ 1.4	58.8 $\pm$ 1.1	47.5 $\pm$ 1.6	87.4 $\pm$ 0.7	98.9 $\pm$ 0.1	31.0 $\pm$ 0.7
volume	60.4 $\pm$ 0.8	58.8 $\pm$ 1.1	64.3 $\pm$ 1.6	59.1 $\pm$ 1.2	53.4 $\pm$ 0.5	86.3 $\pm$ 0.7	96.2 $\pm$ 0.3	27.7 $\pm$ 0.7
connectome	59.8	57.0	58.6	58.0	45.6	81.9	82.4	7.4

Table 2 : Downstream probe performance across fMRI representations. Values are mean accuracy

\pm

standard deviation over 8 random repeats. Bold indicates robust improvement over non-bold (

p<0.0001

). The flat map MAE performs best on dynamic state classification. The volume MAE performs best on age classification. Both volume and flat outperform parcel on sex classification.

Figure 6 : MAE predictions across different fMRI representations for a single example. We project the parcel and volume data to the cortical flat map for consistent visualization. This projection is lossy for the volume model, capturing only voxels intersecting the cortical surface.

space	time	params	FLOPs	compute	data
parcel	11 hr	85M	89G	10K fps	60K fps
flat	28 hr	86M	92G	9K fps	4K fps
volume	50 hr	87M	116G	8K fps	2K fps

Table 3 : Training time comparison across representations for default setting (ViT-B) on a single H100. Parameter count is for encoder only. FLOP count is for forward pass on a single sample. Compute and data loading are in fMRI frames per second (fps).

Panel of Figure 7 — Figure 7 : Scaling analysis. (a) Models are trained on subsets of HCP-YA (split by subject). Lines indicate rolling median over 5 epochs. Vertical lines indicate best epochs for each dataset. Power laws are estimated from best losses using the three smallest data splits. (b) Encoder depth is scaled from 3 to 15 while fixing model proportions. Horizontal lines indicate best loss for each model. Power laws are estimated from best losses using the three smallest models. (c-d) Error bars indicate $\pm$ $2\times$ stdev, using the standard deviations in Table 2 .

Table 4 : Ablations on NSD COCO24. The model is our default flat map CortexMAE (ViT-B). Default settings in gray . red indicates

{>3}\sigma

below baseline. No scores

{>3}\sigma

above baseline. (a) All masking strategies perform well. Uniform requires higher ratio. tube

2\times

( Yang et al., 2025 ) requires lower ratio. (b) No clear benefit from any augmentation. (c) No clear benefit from target patch normalization ( He et al., 2022 ) or global PCA normalization. (d) Better performance for smaller temporal patch size.

Table 11

Figure 8 : fMRI foundation model probe comparison for subject-level trait prediction (left; ABIDE, ADHD200, ADNI, PPMI, HCP-A) and dynamic state prediction (right; HCP-YA Task21, NSD COCO24). To handle small sample size and class imbalance, trait prediction uses balanced accuracy and logistic probe over 100 random splits. State prediction uses raw accuracy and attentive probe over a single fixed split. Confidence intervals indicate

\pm

2\times

stdev over pretraining repeats (only available for CortexMAE and Brain-Semantoks).

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 9 : Schaefer 400 parcellation ( Schaefer et al., 2018 ) with Yeo resting-state networks ( Yeo et al., 2011 ) on the cortical surface and flat map. Relaxation cuts required for flat map transformation ( Gao et al., 2015 ) are marked in white.

config	value
optimizer	AdamW
momentum	$\beta_{1},\beta_{2}{=}0.9,0.95$
weight decay	0.05
learning rate	1.25e-4 (flat, vol), 3.75e-5 (parcel)
lr schedule	cosine decay
warmup steps	31K

Appendix

Table 7 : Pretraining setting on HCP-YA.

config	value
optimizer	AdamW
momentum	$\beta_{1},\beta_{2}{=}0.9,0.999$
base learning rate	3e-4
base weight decay	0.05
lr scale grid	[0.02, 0.023, 0.028, 0.033, 0.038, 0.045, 0.053, 0.062, 0.074, 0.087, 0.1, 0.12, 0.14, 0.17, 0.2, 0.23, 0.27, 0.32, 0.38, 0.44, 0.52, 0.61, 0.72, 0.85, 1, 1.2, 1.4, 1.6, 1.9, 2.3, 2.7, 3.1, 3.7, 4.3, 5.1, 6, 7.1, 8.3, 9.8, 12, 14, 16, 19, 22, 26, 31, 36, 43, 50]
wd scale grid	[1.0]

Appendix

Table 8 : Attentive probe setting

Figure 10 : Attentive probe performance across learning rate for best and last checkpoints. The dense learning rate grid (trained in parallel) allows precise tuning of all models.

Table 17

subset	hours	Age	Sex	Task21	COCO24
all	2058	47.5	87.4	98.9	31.0
rest	1072	47.7	85.5	97.5	30.8
task	987	47.4	85.4	98.5	28.7

Appendix

Table 12 : Pretraining on all HCP-YA data outperforms resting-state or task-based only pretraining. The model is a flat map CortexMAE. Baseline performance is from Table 2 .

Figure 11 : Comparing pretraining datasets across data scales. Downstream performance as a function of pretraining steps (in millions) for models trained on subsets of HCP-YA (left) and UKBB (right). Dots are individual checkpoints; lines show a rolling median. All models within each family are trained for a matched number of effective epochs, with total steps scaled proportionally to dataset size.

Figure 13 : Patch embedding $16\times 16$ filters for the official MAE ( He et al., 2022 ) showing edge artifacts. Maps show the red channel only for consistency with Figure 12 .

Table 22

model	params	FLOP	data/s	fwd/s	FLOP/s
BrainLM	113M	382G	267K	96K	73T
Brain-JEPA	87M	1511G	338K	86K	261T
BrainHarmonix-F	89M	3135G	310K	20K	122T
Brain-Semantoks	63M	6G	179K	1079K	12.2T
CortexMAE-P	85M	8062G	380K	19K	303T
SwiFT	4M	183G	0.5K	15K	5.5T

Appendix

Table 16 : Compute comparison of fMRI foundation models. FLOP count is for forward pass on a single sample (500 frames). Data loading and forward pass throughput is in frames per second. FLOP/s compute throughput is for model forward pass. Models are evaluated with AMP enabled at bfloat16. CortexMAE-{P,F,V} refer to the parcel, flat map, and volume models respectively.

Figure 14 : Attentive probe performance on subject-level trait prediction. Markers indicate each model’s performance (balanced accuracy) on a fixed train/test split using the attentive probe protocol. Confidence intervals indicate the

\text{mean}\pm 2\times\text{stdev}

for the logistic probe over the 100 randomized splits. Attentive probe performance is variable, but within logistic probe CIs.

Panel of Figure 15 — Figure 15 : Uncurated examples of MAE predictions on fMRI flat maps.

Explore similar work

Natural Image Autoencoder-Based fMRI Representations for Trait and State Prediction
Sep 28, 2026Juhyeon Park, Yeonwoo Kim, Peter Yongho Kim +5Functional Magnetic Resonance Imaging Electroencephalography Foundation Models
IPAI, Seoul National University · Microsoft Research · ECE, Seoul National University +1
The Variance Brain Foundation Models Forgot: Third-Order Statistics Predict Cognition Where Billion-Parameter Models Fail
May 29, 2026Giovanni Marraffini, Gabriel Mahuas, Trinidad Borrell +2Functional Magnetic Resonance Imaging Wireless Foundation Models
Inria Saclay Île-de-France, CEA, Université Paris-Saclay, Palaiseau, France · Sigma Nova · Sorbonne Université, Institut du Cerveau - Paris Brain Institute - ICM +1
A Scaling Study for fMRI Foundation Models
Sep 23, 2026Wenhao Ye, Xuanye Pan, Junfeng Xia +3Model Size Functional Magnetic Resonance Imaging
Southern University of Science and Technology · Shenzhen University · SLAI +1

Sep 28, 2026, cs.CVJ/K move · Enter open · S save

Natural Image Autoencoder-Based fMRI Representations for Trait and State Prediction

Juhyeon Park, Yeonwoo Kim, Peter Yongho Kim, Yansen Wang+4

IPAI, Seoul National University · Microsoft Research · ECE, Seoul National University+1

Functional Magnetic Resonance Imaging Electroencephalography Foundation Models Frozen Encoder Autoencoder Architectures

Open paper AskPDF arXiv

Foundation models pre-trained on large-scale fMRI datasets have shown strong downstream performance, but at substantial data and computation cost. To investigate how much fMRI-specific pre-training is actually needed for such performance, we introduce FReD, which derives fMRI representations from a frozen Deep Compression AutoEncoder (DCAE) pre-trained exclusively on natural images and pairs them with a task specific readout. For trait prediction, FReD summarizes frame-wise representations by their temporal mean and log-standard deviation and applies linear probing, with late fusion across two normalization schemes. For state prediction, it represents each frame as a single token and models temporal dependencies with a shallow Transformer. Across four resting-state datasets spanning six trait-prediction targets, linear probes on frozen DCAE features generally outperform those on fMRI foundation model representations and remain competitive with fully fine-tuned fMRI foundation models. On three task-fMRI state-prediction tasks, a temporal readout on DCAE features performs comparably to the strongest foundation models evaluated. A Gaussian injection analysis further shows that localized signal changes are recovered more accurately from the frozen DCAE features than from the evaluated foundation-model representations. Together, these results show that strong performance on current fMRI benchmarks is possible without fMRI-specific representation pre-training, making frozen natural-image features as a useful baseline for assessing its added value.

Abstract

Explore similar work

Natural Image Autoencoder-Based fMRI Representations for Trait and State Prediction

The Variance Brain Foundation Models Forgot: Third-Order Statistics Predict Cognition Where Billion-Parameter Models Fail

A Scaling Study for fMRI Foundation Models