cortex_mae 0.1.1

CortexMAE fMRI foundation model

CortexMAE

CortexMAE is an fMRI foundation model trained on 2.1K hours of fMRI data from the Human Connectome Project using the masked autoencoder framework. We release a family of models trained with different fMRI input representations:

• CortexMAE-P: a computationally efficient model based on the Schaefer-400 parcellation.
• CortexMAE-F: our flagship model based on fMRI flat maps.
• CortexMAE-V: a dense volume model based on an efficient cortex-only representation.

Installation

uv pip install cortex_mae

Or install the latest version from github

uv pip install "cortex_mae @ git+https://github.com/MedARC-AI/CortexMAE.git"

Or clone the repo and install locally

git clone https://github.com/MedARC-AI/CortexMAE.git
cd CortexMAE
uv sync --python 3.11

Quickstart

Load a pretrained model and compute embeddings on a preprocessed fMRI time series from OpenNeuro:

from cortex_mae import CortexMAE, resolve_file

model = CortexMAE.from_pretrained("cortex_mae_flat")

path = resolve_file(
  "s3://openneuro.org/ds006072/NON_BIDS/ciftis/sub-1_Drug2_rsfMRI_uout_bpss_sr_noGSR_sm4.dtseries.nii",
  anon=True,
)
embeds = model.run_embedding(path)
print(embeds.patch_embeds.shape)  # (clips, tokens, dim)

See notebooks/quickstart.ipynb for the full demo.

Pretrained models

Pretrained checkpoints and training logs are available on HuggingFace. We release default models for each input space:

name	input space	shape	size
cortex_mae_flat	flat map	224×560	ViT-B
cortex_mae_parcel	Schaefer-400	400×1	ViT-B
cortex_mae_volume	MNI cortex	465×512	ViT-B

as well as >50 ablation variants covering data scale, model scale, alternative parcellations, etc. List all the available models with cortex_mae.list_models().

model = CortexMAE.from_pretrained("cortex_mae_flat")     # default
model = CortexMAE.from_pretrained("cortex_mae_flat_r2")  # repeat with new seed
model = CortexMAE.from_pretrained("cortex_mae_flat_d6")  # depth-6 model

Datasets

Benchmark datasets are distributed in HuggingFace Arrow format on the MedARC R2 bucket, maintained by Brainmarks. To request access, fill out this form, then configure credentials:

export AWS_ACCESS_KEY_ID=...
export AWS_SECRET_ACCESS_KEY=...
export AWS_ENDPOINT_URL_S3=...   # Cloudflare R2 endpoint

The HCP-YA pretraining data are also available as webdataset shards. The data can be streamed from R2 during pretraining or downloaded locally.

Pretraining

To reproduce pretraining of the default CortexMAE-F model, run

uv run python src/cortex_mae/main_pretrain.py

You can also override defaults

uv run python src/cortex_mae/main_pretrain.py \
  --config config.yaml \
  --overrides \
  input_space=schaefer400 \
  base_lr=3e-4

See the default config src/cortex_mae/config/default_pretrain.yaml for all available options. To reproduce specific model variants, use the original configs on HuggingFace.

Downstream evaluation

Probe evaluation uses Brainmarks. The CortexMAE encoders are registered as cortex_mae_{parcel,flat,volume}:

uv run python -m brainmarks.main_probe cortex_mae_flat patch attn nsd_cococlip

To evaluate a different model variant:

uv run python -m brainmarks.main_probe cortex_mae_flat patch attn nsd_cococlip \
    --overrides model_kwargs.variant=d6

To see a list of all variants:

from brainmarks.models.cortex_mae_wrapper import list_variants

print(list_variants("cortex_mae_flat"))

Support

For help with any issues, reach out to us on MedARC Discord in the #neuro-fm channel.

License

Code is released under the Apache License 2.0 (LICENSE). Model weights are relased under CC-BY-NC 4.0 (LICENSE.models).

Citation

@inproceedings{lane2026scaling,
    title={Scaling Vision Transformers for Functional MRI with Flat Maps},
    author={Connor Lane and Mihir Tripathy and Leema Krishna Murali and
            Ratna Sagari Grandhi and Shamus Sim Zi Yang and Sam Gijsen and
            Debojyoti Das and Manish Ram and Utkarsh Kumar Singh and
            Cesar Kadir Torrico Villanueva and Yuxiang Wei and Will Beddow and
            Gianfranco Cortés and Suin Cho and Daniel Z. Kaplan
            and Benjamin Warner and Tanishq Mathew Abraham and Paul S. Scotti},
    booktitle={ICML},
    year={2026},
}