Scaling Vision Transformers for Functional MRI with Flat Maps

TL;DR

This work investigates scaling fMRI foundation models by representing brain activity as flat maps on a cortical surface and training a spatiotemporal Vision Transformer via masked autoencoding (fm-MAE) on 2.3k hours of HCP data. It demonstrates that masked reconstruction follows a strict power-law scaling with pretraining data size, indicating efficient scaling in this domain. The learned representations support both fine-grained state decoding across subjects and cross-subject, zero-shot visual decoding on NSD, with competitive results on UKBB sex classification and promising performance in several benchmarks. Overall, the study positions flat-map fMRI as a viable input modality for foundation models and highlights the importance of data diversity and broader evaluation in realizing cross-subject generalization.

Abstract

A key question for adapting modern deep learning architectures to functional MRI (fMRI) is how to represent the data for model input. To bridge the modality gap between fMRI and natural images, we transform the 4D volumetric fMRI data into videos of 2D fMRI activity flat maps. We train Vision Transformers on 2.3K hours of fMRI flat map videos from the Human Connectome Project using the spatiotemporal masked autoencoder (MAE) framework. We observe that masked fMRI modeling performance improves with dataset size according to a strict power scaling law. Downstream classification benchmarks show that our model learns rich representations supporting both fine-grained state decoding across subjects, as well as subject-specific trait decoding across changes in brain state. This work is part of an ongoing open science project to build foundation models for fMRI data. Our code and datasets are available at https://github.com/MedARC-AI/fmri-fm.

Figures (8)

• Figure 1: Our flat map MAE (fm-MAE) architecture. Surface-mapped fMRI activity patterns are projected to a flattened cortical mesh Gao2015, resampled as 2D images, and tokenized into patches. We train a standard ViT Dosovitskiy2021 on temporal sequences of "patchified" flat maps using a spatiotemporal MAE He2022Feichtenhofer2022. A large fraction of the image patches are first masked. The encoder computes embeddings for the remaining observed patches, which are passed to the decoder. The model is trained to minimize the MSE loss between the decoder output and pixel values for masked patches.
• Figure 2: Visualization of MAE predictions. Within each panel of $3 \times 3$ images, we show the masked input (left), MAE prediction (middle), and target data (right). We show predictions for 3 frames spaced 4s apart from top to bottom. The model is a ViT-B with a spacetime patch size of $16 \times 16 \times 16$. RGB color mapping is for visualization only, model inputs and predictions are single channel.
• Figure 3: fMRI modeling performance scales with dataset size. The model is a ViT-B trained on varying size subsets of HCP from $N$ = 500K to 7.4M frames (59 to 979 subjects). Stars indicate epochs with lowest test loss selected for power law estimation. Power law parameters in (b) are fit using only the first 3 loss values to illustrate the deviation from prediction. In-distribution reconstruction obeys a strict power law, whereas OOD reconstruction shows signs of saturating.
• Figure 4: Downstream decoding performance as a function of dataset size (left column), model size (middle column), and temporal patch size $p_t$ (right column). Smaller temporal patch size corresponds to larger effective sequence length (tokens per input $= 364 \cdot 16 / p_t$). Black dashes indicate performance on independent validation sets used for classifier parameter tuning.
• Figure 5: 4D fMRI time series are first preprocessed using standard methods Smith2004. The cortical surface mesh is reconstructed using structural MRI and aligned to a standard surface template Dale1999Fischl2012. The fMRI data are then extracted for the cortical ribbon and resampled to the standard surface Glasser2013. This processing was performed by the dataset providers Van-Essen2013Alfaro-Almagro2018Allen2022. Middle figure adapted from Gopinath2023.