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

from ilex.models.neurostorm import NeuroSTORM
model = 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.

Upstream license reference: https://opensource.org/licenses/MIT

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/NeuroSTORM

Provenance

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.

Downloads last month
27
Safetensors
Model size
7.71M params
Tensor type
F32
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support