BrainG3N: 3D Masked Autoencoder for Brain MRI (ViT-L)
Weights for BrainG3N: A Dual-Purpose Tokenizer for Controllable 3D Brain MRI Generation (arXiv:2606.19651).
Self-supervised ViT-Large encoder for 3D brain MRI. Pretrained with masked autoencoding on ~33k volumes spanning 18 public cohorts, covering healthy controls, neurodegenerative disease, neurodevelopmental cohorts, and glioma. The frozen encoder produces general-purpose features for clinical probing without fine-tuning.
Model
| Architecture | ViT-Large masked autoencoder (12 layers, d=1152, 16 heads) |
| Patch size | 16x16x16 voxels |
| Input | 160x192x160 voxels, single channel |
| Tokens | 1,200 (CLS prepended during encoding, stripped from output) |
| Masking ratio | 0.7 |
Quick start
git clone https://huggingface.co/gevaertlab/braing3n
cd braing3n && pip install -r requirements.txt
from modeling import BrainMAEEncoder
from preprocessing import load_volume
encoder = BrainMAEEncoder.from_pretrained("gevaertlab/braing3n", device="cuda")
x = load_volume("sub-0001_t1_preprocessed.nii.gz") # [1, 1, 160, 192, 160]
tokens = encoder.encode(x.cuda()) # [1, 1200, 1152]
features = tokens.mean(dim=1) # [1, 1152]
features is the representation used for the results reported in the paper.
Ready-made scripts
| Script | Does | Needs |
|---|---|---|
extract_features.py | Directory of NIfTIs to features.pt + features.csv (one 1152-d row per volume) | encoder |
reconstruct.py | Encode then decode one volume; writes original + reconstruction | encoder + decoder |
generate.py | Sample synthetic volumes from the conditional DiT | DiT + decoder |
# Batch feature extraction -- the common case
python extract_features.py --input_dir /path/to/scans --output_dir features/
# Non-BIDS filenames
python extract_features.py --input_dir scans/ --output_dir features/ \
--pattern '*_t1c_preprocessed.nii.gz' --id_regex 'CGGA_([A-Za-z0-9]+)_'
# Round-trip one volume through the d'=32 bottleneck
python reconstruct.py --input scan.nii.gz --output recon/
# Conditional generation (--list_conditions shows what this checkpoint accepts)
python generate.py --num_samples 4 --disease GBM --modality t1c --output samples/
Input requirements (read this)
The encoder expects volumes that are already skull-stripped and affinely registered to a common template. It applies only two on-the-fly steps: pad/crop to 160x192x160 and a nonzero-masked z-score. Feeding raw scanner output produces meaningless features -- this is the most common failure mode.
Data is z-scored, not scaled to [0, 1].
Companion weights
| Subfolder | Component | Use |
|---|---|---|
| (root) | MAE ViT-L encoder | Feature extraction |
decoder_d32/ | Projection (1152->32) + CNN decoder | Reconstruct volumes from latents |
dit_d32_6cond/ | Conditional DiT (flow matching) | Conditional generation in latent space |
Results
Evaluation and benchmark numbers are reported in the paper (arXiv:2606.19651). This repository hosts the weights only.
Training data
upenn, ucsf, abide, abide2, adhd200, adni, bgsp, corr, fcon1000, hbn, hcp_ya, ixi, nki, nki2, oasis1, oasis2, ppmi, schizo
Cohorts were used under their respective data use agreements. No individual scans are redistributed here -- only model weights.
License and intellectual property
Weights and accompanying code are released under CC BY-NC 4.0 for non-commercial research use only.
This is a copyright license. It grants no license, express or implied, to any patent or patent application covering the methods described in the paper or implemented here. All patent rights are reserved. Commercial use, or any use beyond non-commercial research, requires a separate agreement -- contact the authors and the Stanford Office of Technology Licensing.
Limitations
- Structural MRI only (T1, T1c, T2, FLAIR). Not validated on DWI, SWI, fMRI, or CT.
- Site is near-perfectly decodable from the features (multi-way site AUC ~1.0), so scanner and acquisition signal is present in the representation. Use site-aware cross-validation for any downstream clinical claim.
- Not a medical device. Research use only.
Citation
@article{vanpuyvelde2026braing3n,
title={{BrainG3N}: A Dual-Purpose Tokenizer for Controllable 3D Brain {MRI} Generation},
author={Van Puyvelde, Max and Gulluk, Ibrahim and Van Criekinge, Wim and Gevaert, Olivier},
journal={arXiv preprint arXiv:2606.19651},
year={2026}
}