Humans struggle to "see" the structure in functional MRI (BOLD) brain maps. Our goal is to train a GPT that understands brain maps better than humans. This kind of "foundation" model should be useful for things like phenotype prediction and brain activity decoding. Plus it will hopefully generate neat fake brain maps.

• Prepare the natural scenes dataset for GPT consumption
  • NSD "beta" activity vectors to flat maps (NSD-Flat, 🤗)
  • Flat maps to tokenized patch sequences (BoldTokenizer) supporting raster, radial, and random order.
• Implement model
  • BoldGPT supporting causal/cross attention, masking, and shuffled prediction. Borrows elements from timm, nanoGPT, and BEiT.
• Implement training script
  • train_gpt.py for training models on shuffled next patch prediction.
• Train models
  • Initial training run with ViT-small, shuffled next patch prediction with discrete token targets, 1000 epochs.
• Implement benchmarks
  • Masked patch prediction
  • Patch-wise object classification (i.e. segmentation)
  • image-to-activity
  • activity-to-image
• Evaluation
• Generate pretty brain maps

In our first training run, we trained a ViT-small decoder to auto-regressively predict the next patch (like GPT). However, the patch order is shuffled for each training example. So it's actually a bit like training a masked prediction (BERT/MAE) style model on all possible masking ratios simultaneously. The targets are discrete patch tokens from k-means (like HUBERT). We trained the model on NSD-Flat for 1000 epochs (90K steps) with a batch size of 2048 on 4 V100 GPUs. The figure shows the patch order, activity data, target tokens, and predictions for a few training samples at the end of training.

Clone the repo and install the package in a new virtual environment

git clone https://github.com/clane9/boldGPT.git
cd boldGPT

python3 -m venv --prompt boldgpt .venv
source .venv/bin/activate
pip install -U pip
pip install -r requirements.txt
pip install -e .

If you'd like to contribute, please feel free fork the repo and join (or start!) a conversation in our issues.

If you find this repository helpful, please consider citing:

@article{lane2023boldgpt,
  author       = {Connor Lane},
  title        = {boldGPT: A GPT foundation model for brain activity maps},
  howpublished = {\url{https://github.com/clane9/boldGPT}},
  year         = {2023},
}