Code for the Brain2Vec paper, a 3D foundation model pre-trained on 74,425 T1-weighted brain MRI scans. In this repo, you can pre-train your own brain2vec model, fine-tune a Vision Transformer (ViT) using our pre-trained brain2vec encoder, train your own downstream ML models using brain2vec output embeddings or perform downstream classification inference using our pre-trained disease classification models trained on brain2vec embeddings.
Important
To use our model locally, you first need to apply for access to our hugging face model card. We follow a 2-step preprocessing with skull-stripping and linear registration, make sure to follow the steps here. Don't perform inference on raw images locally. You can always use Brain2Vec on NiChart where you can upload your raw T1w MRI scans via cloud and obtain Brain2Vec scores without preprocessing steps, installation or specialized infastructure (GPU).
Left: Masked input | Middle: Reconstruction (Z-normalized) | Right: Ground truth
By fine-tuning with our pre-trained model, we achieved both better accuracy and convergence speed on downstream deep learning tasks
We proved that using Brain2Vec embeddings we get superior or on-par classification performance across multiple disease groups as well as better accuracy than NeuroJEPA and DLMUSE ROI's for brain age regression.
We also proved that foundation models focus on limited but important set of brain regions. Manipulating a subset of important latent features produces monotonic, region-specific effects on the reconstructed anatomy.
-
Brain2Vec implementation follows the code of Masked Autoencoders from Meta.
-
Vision Transformers implementation follows the code of timm from huggingface.
Before using any of our models locally, make sure that the input is skull-stripped and linearly registered to a common atlas space. We provide a pre-processing pipeline that transforms your raw scans to Brain2Vec compatible scans ready for inference.
We use NiChart_DLMUSE for skull-stripping (s3 masked folder) and flirt for registration, though you can use any skull-stripping model and achieve comparable results.
NiChart_DLMUSE -i {dir with nifti files} -o ... --refaced-data
flirt -in {s3_masked NiChart_DLMUSE output} -ref registration/img_template.nii.gz -out ...You can use Brain2Vec for both feature extraction and disease/risk-factor classification or brain age prediction on NiChart by simply dragging and dropping your NIfTI files, we perform pre-processing for you, also no GPU or installation required.
You can install Brain2Vec locally using pip
pip install -e .We trained with two A100 GPUs for 500 epochs with a patience of 20 epochs. For more training details please read the Brain2Vec paper. Using our code, you can parallelize your training on SLURM cluster environments. Other environments are supported but not tested, so bugs may appear.
srun Brain2Vec brain2vec_pretrain --img_size 128 \
--masking_type random \
--mask_ratio 0.75 \
--model brain2vec_base \
--in_dir "input directory with train/test/eval
subdirectories" \
--batch_size 32 \
--out_dir "directory where the model will be
saved" \
--model_name "model name (.pth)" \
--device cuda \
--lr 0.001 \
--weight_decay 0.05 \
--optimizer AdamW \
--epochs 500 \
--patience 20 \
--distributed \
--verbose \One way of using Brain2Vec is for encoder fine-tuning. The pre-trained encoder weights can serve as initialization for Vision Transformer (ViT) models applied to downstream tasks like disease classification or brain age regression.
The only thing you need is a csv with MRID and target columns.
srun Brain2Vec vit_finetune --img_size 128 \
--mode binary \
--model vit_base \
--label_dict "csv with MRID and target columns" \
--target "target column" \
--pretrain "brain2vec.safetensors (model name from our
hugging face)" \
--in_dir "input directory with train/test/eval
subdirectories" \
--batch_size 32 \
--out_dir "directory where the model will be
saved" \
--model_name "model name (.pth)" \
--device cuda \
--lr 0.001 \
--weight_decay 0.05 \
--optimizer AdamW \
--epochs 200 \
--patience 20 \
--distributed \
--verboseOne other way of using Brain2Vec is for feature extraction. We can use the Brain2Vec frozen encoder to produce 768-dimensional embeddings from T1 scans. This compact numerical representation serves as input for lightweight downstream classifiers, such as SVM or MLP models.
python inference.py brain2vec_inference --model brain2vec_base \
--img_size 128 \
--in_dir "path to directory with nifti scans" \
--out_json "name of json file with output embeddings" \
--device cuda
python inference.py vit_inference --model vit_base \
--num_classes 1 \
--img_size 128 \
--mode regression \
--in_dir "path to directory with nifti scans" \
--out_csv "output CSV file with predictions" \
--device cudaWe provide a simple pipeline that enable users to train their own ML models using output embeddings from Brain2Vec.
python train_svm.py dx_classification --cn_embeddings "path to your CN embeddings (.json)" \
--dx_embeddings "path to your DX embeddings (.json)"
python train_svm.py ba_regression --embeddings "path to your embeddings (.json)" \
--label_dict "path to a CSV that contains MRID and Age columns"If you use our model, make sure to cite our work:
@misc{maggioros2026brain2vecselfsupervised3dfoundation,
title={Brain2Vec: A Self-Supervised 3D Foundation Model for Structural Brain MRI},
author={Spiros Maggioros and Guray Erus and Gareth Harman and George Aidinis and Alexander Getka and Aristeidis Sotiras and Christos Davatzikos},
year={2026},
eprint={2610.28411},
archivePrefix={arXiv},
primaryClass={cs.CV},
url={https://arxiv.org/abs/2610.28411},
}
We are open to any contributions about code organization, inference efficiency and quality, fixes, etc. Just open a PR with a full explanation of your code.