NeuroSTORM (4D Swin + Mamba MAE foundation model for fMRI) -- NeuroSTORM mask_ratio=0.8 pretraining checkpoint
Description
NeuroSTORM (Zhang et al., Nature Biomedical Engineering 2026)
is a general-purpose foundation model for fMRI analysis. The
architecture is a 4D Swin Transformer with Mamba selective-scan
state-space models as the per-window mixer (in place of the
conventional window attention) and a masked-autoencoder (MAE)
pretraining objective.
v0 ships two variants corresponding to the upstream's two
released mask-ratio pretraining checkpoints (0.5 and 0.8),
sharing parameter-identical architecture (7,712,892 trainable
scalars each); only the training-time mask ratio differs.
The 4D operates on the full (96, 96, 96, 20) voxel-time
volume (MNI152 brain-only resampling of fMRI BOLD); the
resulting embedding is the deepest-stage feature map
(288, 2, 2, 2, 20) -- consumed by downstream task heads
(sex/age regression, phenotype prediction, disease diagnosis)
per the upstream's benchmark.
Intended use
fMRI foundation backbone for representation learning. Input: (1, 96, 96, 96, 20) MNI152 BOLD clip (consumer handles registration + min-max normalisation + quantisation per the upstream pipeline). Output: deepest- stage encoder feature map (288, 2, 2, 2, 20) suitable for downstream task heads (linear probe, attention pool, MLP). The MAE decoder is loaded into the bundle for state-dict round-trip integrity but is not invoked at inference.
Usage
1from ilex.models.neurostorm import NeuroSTORM
2model = NeuroSTORM.from_pretrained('ilex-hub/neurostorm.ratio0.8.1')
Authors
Zhang Y. et al. (CUHK-AIM-Group, Chinese University of Hong Kong)
Citation
Zhang Y. et al. (2026). Towards a general-purpose foundation model for fMRI analysis. Nature Biomedical Engineering. doi:10.1038/s41551-026-01666-y.
References
- Zhang Y. et al. (2026). Towards a general-purpose foundation model for fMRI analysis. Nature Biomedical Engineering. doi 10.1038/s41551-026-01666-y.
- Gu A., Dao T. (2023). Mamba -- Linear-time sequence modeling with selective state spaces. arXiv 2312.00752.
- Liu Z. et al. (2021). Swin Transformer -- Hierarchical Vision Transformer using Shifted Windows. arXiv 2103.14030.
- Upstream code -- github.com/CUHK-AIM-Group/NeuroSTORM (Apache-2.0); upstream weights -- huggingface.co/zxcvb20001/NeuroSTORM (MIT).
License
HF Hub license tag: mit
Effective terms: MIT (on the released checkpoint weights, per the HF model card at huggingface.co/zxcvb20001/NeuroSTORM). The upstream code (github.com/CUHK-AIM-Group/NeuroSTORM) is separately Apache-2.0; both are fully permissive. The ilex JAX / Equinox port code is licensed under Apache-2.0 / GPL-3.0.
Copyright
NeuroSTORM is copyright (c) CUHK-AIM-Group 2026, Apache-2.0-
licensed on the code (github.com/CUHK-AIM-Group/NeuroSTORM,
LICENSE file) and MIT-licensed on the released checkpoints
(huggingface.co/zxcvb20001/NeuroSTORM, license tag mit). The
ilex JAX / Equinox port code is separately licensed under
Apache-2.0 / GPL-3.0; license restrictions on the upstream
artefacts are preserved through the canonical bundle's
_ilex.origin = 'pytorch' provenance.
Upstream source
Original weights / reference implementation:
https://github.com/CUHK-AIM-Group/NeuroSTORMProvenance
This artefact was produced by
ilex's
save/load pipeline. The architecture is implemented in
ilex.models.neurostorm.NeuroSTORM and the weights have been converted
from their upstream format. See the upstream source above
for the canonical reference.
