--- license: cc-by-nc-sa-4.0 tags: - fMRI - neuroscience - foundation_model --- # Flexibrain Flexibrain is a voxel-level fMRI representation learning framework for pretraining and downstream classification. It keeps fMRI volumes in a fixed 96 x 96 x 96 input grid, reads each sample's voxel spacing and TR from the NIfTI header, and resizes patch embedding kernels in physical spatial and temporal units before learning with a Mamba-JEPA backbone.

FlexiBrain framework pipeline

## Installation The code was tested on l40 with Python 3.10, PyTorch 2.1.2, CUDA 12.1, `causal-conv1d`, `mamba-ssm`, and `flash-attn`. ```bash conda create -n flexibrain python=3.10 conda activate flexibrain pip install -r requirements.txt pip install -e . ``` Check the CLI: ```bash python -m flexibrain --help python -m flexibrain pretrain --help python -m flexibrain downstream --help ``` ## Data Preparation Each sample should be a 4D NIfTI file shaped as: ```text 96 x 96 x 96 x T ``` Flexibrain uses the NIfTI header to read voxel spacing and TR. If a dataset has missing TR metadata, fix the header before training or pass an explicit fallback with `--default-tr` / `data.default_tr`. `T_prime` and `tau_seconds` control the selected temporal length: ```text kt = round(tau_seconds / TR) T_selected = T_prime * kt ``` The preprocessing script can convert native/T1/MNI-space inputs to 96 x 96 x 96, apply sample-wise global z-score normalization over foreground voxels, and write 4D NIfTI outputs: ```bash python data_process.py \ --input-root /path/to/input_root \ --output-root /path/to/output_root \ --spaces all \ --groups class0,class1,class2 ``` Expected grouped input layout: ```text input_root/ |-- nativespace/class0/*.nii.gz |-- t1space/class0/*.nii.gz `-- mnispace/class0/*.nii.gz ``` If files are not organized by group subfolders, omit `--groups`. For MNI-space inputs, provide `--template-mask` when the default mask is not available. Pretraining list files contain one NIfTI path per line: ```text /path/to/sub-0001_bold.nii.gz /path/to/sub-0002_bold.nii.gz ``` Downstream classification uses the same list format plus a CSV label table: ```csv Subject,Group_idx 003_S_0908,2 011_S_0002,1 1001,0 ``` Default label fields are `Subject` and `Group_idx`. `path_id_mode=auto` supports ADNI-style IDs such as `003_S_0908`, ADHD-style filenames, and fallback digit IDs. ## Pretraining Run from a config: ```bash python -m flexibrain pretrain --config configs/pretrain_example.yaml ``` Or use CLI arguments: ```bash python -m flexibrain pretrain \ --train-list /path/to/pretrain_train.txt \ --val-list /path/to/pretrain_val.txt \ --checkpoint-dir ./checkpoints/pretrain/example \ --log-dir ./logs/pretrain/example \ --embed-dim 512 \ --depth 24 \ --predictor-depth 2 \ --bimamba-type v2 \ --if-devide-out \ --batch-size 4 \ --epochs 30 \ --lr 5e-4 \ --weight-decay 0.05 \ --warmup-epochs 3 \ --mask-ratio 0.65 \ --grad-accumulation-steps 4 \ --t-prime 30 \ --tau-seconds 6.0 \ --use-amp ``` Outputs: ```text checkpoint_latest.pt checkpoint_best.pt pretrain_*.log ``` ## Downstream Classification Run from a config: ```bash python -m flexibrain downstream --config configs/downstream_example.yaml ``` Or use CLI arguments: ```bash python -m flexibrain downstream \ --train-list /path/to/downstream_train.txt \ --val-list /path/to/downstream_val.txt \ --test-list /path/to/downstream_test.txt \ --csv /path/to/labels.csv \ --pretrain-checkpoint /path/to/checkpoint_best.pt \ --num-classes 3 \ --head-type transformer \ --batch-size 8 \ --epochs 30 \ --lr 1e-5 \ --lr-backbone 6e-6 \ --lr-head 6e-5 \ --checkpoint-dir ./checkpoints/downstream/example \ --log-dir ./logs/downstream/example \ --use-amp ``` During downstream training, validation metrics select `downstream_best.pt`. The test set is evaluated once at the end after loading that best validation checkpoint, and the final metrics are written to `test_metrics.json`. ## Configuration YAML config mirrors the CLI options. Keep private paths in local config files and leave shared configs as portable examples. The provided examples use placeholder paths under `data/`: ```text configs/pretrain_example.yaml configs/downstream_example.yaml ``` ## Checkpoint Compatibility The downstream loader can initialize from the original pretraining checkpoint path format: ```text /path/to/checkpoint_best.pt ``` When `use_checkpoint_config: true`, model-shape settings stored in the checkpoint are applied before loading the backbone. ## License This repository is provided for non-commercial research use under CC BY-NC-SA 4.0. See `LICENSE` and `NOTICE` for license boundaries and preserved notices.