README.md

---
license: cc-by-nc-4.0
tags:
  - medical-imaging
  - brain-mri
  - masked-autoencoder
  - vision-transformer
  - 3d
  - self-supervised
  - foundation-model
library_name: pytorch
pipeline_tag: feature-extraction
extra_gated_prompt: >-
  This model was pretrained on brain MRI from consortia that impose Data Use
  Agreements (including ADNI and PPMI). By requesting access you confirm that
  you will use these weights for non-commercial research only, will not attempt
  to reconstruct or re-identify any individual scan or subject, and will comply
  with the terms of the source datasets. No patent rights are granted; see the
  License and intellectual property section.
extra_gated_fields:
  Name: text
  Institution: text
  Intended use: text
  I agree to use these weights for non-commercial research only: checkbox
---
[![arXiv](https://img.shields.io/badge/arXiv-2606.19651-b31b1b.svg)](https://arxiv.org/abs/2606.19651)

# 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](https://arxiv.org/abs/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

```bash
git clone https://huggingface.co/gevaertlab/braing3n
cd braing3n && pip install -r requirements.txt
```

```python
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 |

```bash
# 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](https://arxiv.org/abs/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

```bibtex
@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}
}
```