[ArXiv (fm-MAE)]
[Discord]
[Notion]
[Intro video]
Our goal is to train state-of-the-art foundation models for fMRI and unlock novel clinical applications for improving brain and mental health. This is an open community driven project where we share ideas, data, and compute. Join our Discord to get involved!
A (growing) list of our past and current contributors: @Atom-101, @ckadirt, @Dojo2024, @har5h1l, @ks9249, @Leema-Krishna-Murali, @LeoYu, @Manishram-ai, @mihirneal, @MohammedSB, @sagarigrandhi, @souvik3333, @ss8319, @Sue0515, @utxrsh, @yuxiangwei0808. Thank you! 🤗
A key question for building fMRI foundation models is how to format or "tokenize" the native 4D volumetric fMRI data for model input. Most current approaches use either parcellation based averaging [Brain-JEPA, BrainLM] or 4D volume convolution [SwiFT, NeuroStorm]. Our current work [fm-MAE] uses a novel flat map based representation that is in between the two along the "bitter lesson continuum": more information rich than parcellation approaches, yet more structured than native 4D volume approaches.
Another important axis for fMRI foundation models is the self-supervised learning objective. For our initial work, we picked the MAE objective for its simplicity.
Beyond flat maps, there are several other tokenization strategies we are interested in exploring:
- "vector" parcellation embeddings (cf [Feng et al., 2025])
- native surface mesh tokenization (cf [Yu et al., 2025])
- sparse gray matter only 4D volume patch embedding
Beyond MAE, there are many different types of SSL objectives we could try:
- Contrastive: [SimCLR]
- Self-distillation: [SimSiam, DINO(v2,v3), JEPA]
- Reconstructive: [AIM, l-DAE]
- Masked image modeling: [CAPI]
There is also a lot of work to do to set up a solid infrastructure for training and evaluating SOTA fMRI foundation models:
- assembling diverse large-scale pretraining datasets
- establishing reproducible benchmarks
- developing methods to interpret models' learned representations
- optimizing model training and inference
Clone the repo, install uv, and run
uv syncThis will create a new virtual environment for the project with all the required dependencies. Activate the environment with
source .venv/bin/activateor use uv run. See the uv docs for more details.
See our quickstart notebook for a demo of loading an fm-MAE model and plotting masked predictions.
For datasets that allow open sharing, we release our preprocessed flat map datasets to huggingface.
-
hcp-flat-wds: HCP-YA minimally preprocessed fMRI data converted to flat maps (webdataset format). -
nsd-flat-wds: NSD preprocessed surface-mapped fMRI data converted to flat maps (webdataset format).
The datasets can be downloaded and used locally:
from huggingface_hub import snapshot_download
from data.flat_data import make_flat_wds_dataset
# download first shard only
snapshot_download(
repo_id="medarc/hcp-flat-wds", repo_type="dataset", allow_patterns="*_0000.tar", local_dir="datasets/hcp-flat"
)
url = "datasets/hcp-flat/hcp-flat_0000.tar"
dataset = make_flat_wds_dataset(url)or streamed directly from huggingface:
from huggingface_hub import get_token
from data.flat_data import make_flat_wds_dataset
# first ten shards
url = "https://huggingface.co/datasets/medarc/hcp-flat-wds/resolve/main/hcp-flat_{0000..0009}.tar"
hf_token = get_token()
url = f"pipe:curl -s -L {url} -H 'Authorization:Bearer {hf_token}'"
dataset = make_flat_wds_dataset(url)We also share some misc data in a shared working data repo: medarc/fmri-fm.
| model | arch | patch size (p-pt) | # params | pretrain data | weights |
|---|---|---|---|---|---|
| fm-MAE | ViT-B | 16-16 | 89M | HCP-YA | download |
| fm-MAE | ViT-B | 16-2 | 89M | HCP-YA | download |
We have uploaded an initial set of pretrained models to huggingface. To load a pretrained model, you can run:
from flat_mae.models_mae import MaskedAutoencoderViT
model = MaskedAutoencoderViT.from_pretrained("medarc/fm_mae_vit_base_patch16-16.hcp")See the quickstart notebook for an example of running masked prediction inference.
To launch a debug pretraining run on one gpu, run:
uv run python src/flat_mae/main_pretrain.py \
--cfg-path src/flat_mae/config/debug_pretrain.yamlOr to run on cpu:
uv run python src/flat_mae/main_pretrain.py \
--cfg-path src/flat_mae/config/debug_pretrain.yaml \
--overrides device=cpuTo launch a full pretraining run on one gpu:
uv run python src/flat_mae/main_pretrain.pyRead the default_pretrain.yaml config for a description of all the options.
The codebase follows a flat organization for (hopefully) easier forking and hacking. It contains the following modules which each have a flat layout and are independent from one another.
Additional documentation for each module can be found in the module README.
As the project progresses, the codebase will grow width-wise, with additional modules added to try out different approaches.
This is a community-driven open science project. We welcome all contributions. To get started contributing, see our contributing guide and code of conduct. Then take a look at our open issues.
By default, code and model weights are released under a CC-BY-NC 4.0 license. See LICENSE for details. However, many of the individual source files are derived from third-party sources licensed under an Apache 2.0 or MIT license. In these cases, we preserve the original license. The relevant license for a particular source file (if different from the default) is indicated in the file header.