ICML 2026

Omni-fMRI

A universal atlas-free fMRI foundation model.

Omni-fMRI learns directly from voxel-level BOLD signals, replacing atlas-dependent regional summaries with dynamic patch tokenization and global self-attention over informative brain regions.

Mo Wang, Wenhao Ye, Junfeng Xia, Junxiang Zhang, Xuanye Pan, Minghao Xu, Haotian Deng, Hongkai Wen, Quanying Liu SUSTech, University of Warwick, Shenzhen University

Paper GitHub Checkpoint Docker Quick Start Data BibTeX

49,497 pretraining sessions

9 pretraining datasets

11 benchmark datasets

16 downstream tasks

Abstract

Voxel-level pretraining without an atlas bottleneck.

Self-supervised fMRI foundation models have improved transfer, but most still compress dense voxel signals into predefined ROI parcellations. Omni-fMRI removes that atlas dependency by operating directly on voxel-level time series while using dynamic patching to keep large-scale pretraining tractable.

Atlas-free inputs

The encoder works on voxel-level fMRI volumes rather than predefined regional summaries, preserving fine-grained spatial information and avoiding atlas-induced bias.

Dynamic patching

Background regions are pruned, simple foreground regions are merged into larger patches, and complex regions keep finer resolution for efficient global modeling.

Reproducible benchmark

The paper standardizes datasets, splits, and tasks for resting-state and task-based fMRI evaluation across diagnosis, demography, retrieval, emotion, and state prediction.

Quick Start

Process a raw fMRI volume and extract Omni-fMRI features.

The checkpoint is hosted on Hugging Face; keep the downloaded file at pretrain_checkpoint/checkpoint.pth. The smoke-test script and notebook below create a tiny synthetic NIfTI first, so a new user can verify the full preprocessing-to-feature path before preparing ADNI, HCP, ABIDE, or other benchmark data.

1. Install

Clone the repository and fetch the HF checkpoint.

git clone https://github.com/OneMore1/Omni-fMRI.git
cd Omni-fMRI
conda create -n omnifmri python=3.11 -y
conda activate omnifmri
pip install -r requirements.txt
pip install -U huggingface_hub

mkdir -p pretrain_checkpoint
huggingface-cli download OneMore1/Omni-fMRI checkpoint.pth \
  --local-dir pretrain_checkpoint

2. Preprocess

Convert raw NIfTI files into 96 x 96 x 96 x 40 NPZ segments.

python data_preparation/preprocessing.py \
  --input_dir examples/raw_nifti \
  --output_dir outputs/demo_npz \
  --pattern .nii.gz \
  --target_shape 96 96 96 \
  --segment_length 40

3. Extract

Run the frozen encoder and save feature tokens.

python extract_feat.py outputs/demo_npz \
  --checkpoint pretrain_checkpoint/checkpoint.pth \
  --output-dir outputs/demo_tokens \
  --device cuda:0 \
  --overwrite

Smoke test

Validated server path

Tested on volcengine2 with the released checkpoint: synthetic NIfTI -> official preprocessing -> official feature extraction.

Saved array	Shape	Dtype
`cls_token`	(768,)	float32
`patch_tokens`	(880, 768)	float32
`patch_coords`	(880, 3)	int64

Try it

Runnable starter assets

Notebook End-to-end cells for install, synthetic NIfTI, preprocessing, and extraction. Smoke script One-command validation for a server or fresh environment. HF checkpoint Download checkpoint.pth from the model repository.

Method

Dynamic patches make voxel-level ViT pretraining practical.

Instead of forcing a dense regular grid into a hierarchical backbone, Omni-fMRI tokenizes the volume according to local spatiotemporal complexity. This reduces the encoder sequence from roughly 14K tokens to about 4.3K tokens while keeping global self-attention over informative brain tissue.

Allocate content-adaptive patches

Prune non-brain background and assign patch scale by local signal variability.

Project mixed-scale tokens

Use dual-path multi-scale embedding with ZeroMLP for stable coarse-to-fine learning.

Reconstruct by scale

Route masked-autoencoding targets to scale-aware prediction heads.

Results

Strong transfer across resting-state and task-based fMRI.

95.54%

HCP sex classification ACC

Omni-fMRI leads the benchmark table while also achieving 95.05 F1 on HCP sex classification.

88.77%

ADNI AD diagnosis ACC

The model reaches the best reported ADNI Alzheimer's disease diagnosis accuracy in the paper's main comparison.

23.08%

Top-10 image retrieval

On 100-way image retrieval, Omni-fMRI improves over NeuroSTORM, SwiFT, and TFF in Top-10, Top-5, and Top-1 accuracy.

45.29%

HCP task accuracy at 10% data

Few-shot evaluation shows the model preserves strong task classification performance under severe data reduction.

Benchmark summary

Linear probing stays competitive.

With a frozen encoder and a lightweight linear head, Omni-fMRI often surpasses fine-tuned baselines and shows a smaller ADNI AD degradation than NeuroSTORM LP.

Efficiency summary

Scalability is part of the method.

The paper reports a 157M-parameter model pretrained for 35 epochs in about 32 hours on 4 NVIDIA A10G GPUs, enabled by dynamic token reduction.

Table 3

Main demography and disease diagnosis results.

Values are copied from Table 3. Accuracy/F1 use higher-is-better metrics; age regression reports MSE lower-is-better and Pearson R higher-is-better. Red cells mark the single best value in each column; stars follow the paper's reported effect-size markers.

Model	ABIDE MSE	ABIDE R	NKI age MSE	NKI age R	SALD MSE	SALD R	ABCD ACC	ABCD F1	HCP ACC	HCP F1
Brain-LM	0.907 +/- 0.007	0.239 +/- 0.023	0.503 +/- 0.015	0.678 +/- 0.014	0.677 +/- 0.055	0.634 +/- 0.061	59.24 +/- 1.03	58.74 +/- 0.70	62.71 +/- 4.43	61.74 +/- 4.72
Brain-JEPA	0.982 +/- 0.066	0.166 +/- 0.017	1.030 +/- 0.080	0.330 +/- 0.030	1.130 +/- 0.108	0.299 +/- 0.060	-	-	69.97 +/- 2.73	69.17 +/- 3.55
BrainMass	0.695 +/- 0.038	0.500 +/- 0.044	0.600 +/- 0.062	0.620 +/- 0.036	0.704 +/- 0.080	0.626 +/- 0.062	57.92 +/- 1.19	57.86 +/- 1.18	67.65 +/- 1.02	66.93 +/- 1.78
SwiFT	0.524 +/- 0.010	0.704 +/- 0.011	0.127 +/- 0.004	0.946 +/- 0.002	0.164 +/- 0.08*	0.926 +/- 0.003*	66.67 +/- 3.51	68.69 +/- 0.44	90.75 +/- 4.38	90.76 +/- 4.37
NeuroSTORM	0.537 +/- 0.016	0.648 +/- 0.013	0.152 +/- 0.033	0.835 +/- 0.035	0.194 +/- 0.019	0.892 +/- 0.010	76.51 +/- 1.88	76.24 +/- 1.70	94.17 +/- 2.65	94.17 +/- 2.15
Omni-fMRI	0.427 +/- 0.006*	0.734 +/- 0.002*	0.088 +/- 0.004*	0.959 +/- 0.003*	0.182 +/- 0.019	0.909 +/- 0.005	77.01 +/- 0.55	76.91 +/- 0.54	95.54 +/- 0.36	95.05 +/- 0.41

Model	BHRC ACC	BHRC F1	PPMI ACC	PPMI F1	ADNI MCI ACC	ADNI MCI F1	ADNI AD ACC	ADNI AD F1	NKI Edu ACC	NKI Edu F1
Brain-LM	57.09 +/- 1.33	54.69 +/- 1.03	54.86 +/- 1.20	49.03 +/- 2.64	64.53 +/- 2.01	61.35 +/- 0.66	80.96 +/- 0.46	72.43 +/- 0.64	60.46 +/- 0.73	58.93 +/- 0.70
Brain-JEPA	58.16 +/- 6.10	56.16 +/- 6.05	47.91 +/- 3.61	48.60 +/- 2.42	64.53 +/- 0.60	55.17 +/- 0.66	80.23 +/- 0.80	73.26 +/- 1.03	49.21 +/- 2.15	48.73 +/- 1.62
BrainMass	62.76 +/- 1.50	61.05 +/- 1.22	58.68 +/- 4.21	51.95 +/- 2.23	62.39 +/- 0.60	61.35 +/- 0.60	73.45 +/- 3.91	62.10 +/- 0.70	62.36 +/- 1.24	61.98 +/- 0.88
SwiFT	70.21 +/- 0.02	69.56 +/- 0.02	61.55 +/- 2.52	58.23 +/- 4.20	65.38 +/- 1.29	72.54 +/- 0.16	74.90 +/- 1.88	80.76 +/- 0.33	78.54 +/- 4.87	79.08 +/- 3.85
NeuroSTORM	73.02 +/- 2.19	72.49 +/- 1.96	68.25 +/- 1.18	65.18 +/- 2.55	66.67 +/- 1.06	62.17 +/- 8.93	84.26 +/- 0.80	83.91 +/- 0.75	89.35 +/- 4.72	89.34 +/- 4.49
Omni-fMRI	75.86 +/- 1.35*	75.96 +/- 1.31*	72.57 +/- 0.98*	70.19 +/- 0.84*	72.92 +/- 1.64*	72.87 +/- 1.74	88.77 +/- 0.71*	88.60 +/- 0.72*	92.05 +/- 2.85	91.64 +/- 2.25

Datasets

Pretraining and downstream benchmark coverage.

Table 2 in the paper groups datasets into pretraining-only data, internal downstream data, and external downstream evaluation data.

ID	Dataset	Participants	Sessions	Age	Sex (M/F)	Downstream task
Type I: only for pretraining
1	UK Biobank	38,301	38,372	68.37 +/- 7.48	17879/20493	-
2	AOMIC (PIOP1)	168	672	22.15 +/- 1.83	94/67	-
3	AOMIC (PIOP2)	180	720	21.99 +/- 1.83	98/81	-
4	CHCP	244	1,224	18-79 range	117/127	-
5	ISYB	130	520	22.52 +/- 2.61	98/32	-
Type II: internal downstream tasks
6	ABCD	1,680	6,720	9.93 +/- 0.63	813/866	Sex classification
7	ABIDE	609	2,436	6-58 range	506/103	Age regression
8	HCP rest	606	2,424	28.79 +/- 3.53	321/281	Sex classification
9	PPMI	331	1,324	35-84 range	181/150	Disease classification
Type III: external downstream tasks
10	ADNI	497	1,988	55-90 range	-	Disease classification
11	SALD	493	493	19-80 range	197/306	Age regression
12	BHRC	465	465	9.93 +/- 1.79	272/193	Sex classification
13	NKI	717	717	7-86 range	311/406	Education / age regression
14	NSD	6	70,566	19-32 range	-	Image retrieval
15	HCP task	118	1,647	23-36 range	54/64	State prediction
16	StudyForrest	20	71,980	26.6 mean	12/8	Emotion prediction

Resources

Project links for release, reproduction, and deployment.

Code

GitHub repository

Training, evaluation, preprocessing, and benchmark scripts.

Runtime

Docker image

Containerized environment for reproducing Omni-fMRI experiments.

Checkpoint

Hugging Face model

Released Omni-fMRI checkpoint and model repository entry.

Notebook

Feature extraction demo

Runnable walkthrough from synthetic NIfTI to saved feature tokens.

Smoke

Server validation script

Minimal script for checking preprocessing and extraction end to end.

Paper

arXiv entry

Public paper page for discovery, citation, and external linking.

Performance Map

One model leads across heterogeneous fMRI tasks.

Figure 1 summarizes the headline comparison: Omni-fMRI is reported as the top performer across a broad collection of resting-state and task-based benchmarks, including linear probing variants.

Interpretability

Voxel representations align with disease-related brain maps.

SSIM

0.5639 vs 0.4636

Omni-fMRI shows higher similarity to the Neurosynth Alzheimer reference map than the NeuroSTORM baseline.

RMSE

0.0562 vs 0.1386

Lower error supports the paper's claim that the learned voxel-level representation captures task-relevant patterns.

Paper Figures

Original figure crops from the manuscript PDF.

Citation

BibTeX

@inproceedings{wang2026omnifmri,
  title = {Omni-fMRI: A Universal Atlas-Free fMRI Foundation Model},
  author = {Wang, Mo and Ye, Wenhao and Xia, Junfeng and Zhang, Junxiang and Pan, Xuanye and Xu, Minghao and Deng, Haotian and Wen, Hongkai and Liu, Quanying},
  booktitle = {Proceedings of the International Conference on Machine Learning},
  year = {2026}
}