ilex-hub/triad.swinb-mae.1

Triad-Swin (3D-MRI self-supervised Swin-B backbone) -- Triad Swin-B encoder (MAE-pretrained)

Description

Triad vision foundation model for 3D MRI, ported to JAX / Equinox from the upstream PyTorch release. Triad is an nnUNet PlainConvEncoder pretrained self-supervised on Triad-131K (131,170 3D MRI volumes spanning brain, breast, and prostate; T1/T2/FLAIR/DWI/fMRI/DCE) and serves as a transfer-learning backbone for downstream MRI segmentation, classification, and registration. The published checkpoints are encoder-only (the self-supervised decoder / mask token are stripped); this port exposes the pretrained encoder, whose multi-scale features are the transfer representation. Two backbone families are ported: the nnUNet PlainConvUNet encoder (TriadPlainConvUNet) and the 3D Swin Transformer encoder (TriadSwinViT, the Swin-B variant, via the shared nimox SwinViT primitive). Each is released under two self-supervised objectives -- masked autoencoding (MAE) and SimMIM -- as separate bundles (four in total).

Intended use

Transfer-learning backbone for 3D MRI: contrast / task-agnostic multi-scale feature extraction via a 3D Swin Transformer. The forward returns the 5-level SwinViT feature pyramid (channels 48, 96, 192, 384, 768 at strides 2, 4, 8, 16, 32); a downstream consumer attaches its own decoder / head. Single-channel input with each spatial dim a multiple of 16. Pretrained by masked autoencoding (MAE) on Triad-131K. Encoder-only: the SSL reconstruction decoder is not shipped.

Usage

from ilex.models.triad import TriadSwinViT
model = TriadSwinViT.from_pretrained('ilex-hub/triad.swinb-mae.1')

Authors

Wang S., et al.

Citation

Wang S., Safari M., Li Q., Chang C.-W., Qiu R. L. J., Roper J., Yu D. S., Yang X. (2025). Triad: Vision Foundation Model for 3D Magnetic Resonance Imaging. arXiv:2502.14064. The Swin backbone is MONAI's SwinUNETR swinViT (use_v2): Hatamizadeh A., Nath V., Tang Y., Yang D., Roth H., Xu D. (2022). Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images. BrainLes 2021. arXiv:2201.01266. Tang Y., et al. (2022). Self-Supervised Pre-Training of Swin Transformers for 3D Medical Image Analysis. CVPR. arXiv:2111.14791.

References

• Wang S., et al. (2025). Triad: Vision Foundation Model for 3D Magnetic Resonance Imaging. arXiv:2502.14064. https://arxiv.org/abs/2502.14064
• Codebase: https://github.com/wangshansong1/Triad

License

HF Hub license tag: mit

Effective terms: MIT (Shansong Wang et al.) on both the Triad code (https://github.com/wangshansong1/Triad) and the released pretrained checkpoints. No commercial restrictions; no gating required. The arXiv preprint (2502.14064) is separately distributed under CC BY 4.0, but the code and weights the ilex bundle re-expresses are MIT. The ilex JAX / Equinox port code is separately licensed under Apache-2.0 / GPL-3.0; that does not alter the upstream MIT terms governing the weights.

Upstream license reference: https://github.com/wangshansong1/Triad/blob/main/LICENSE

Copyright

Network architecture and pretrained weights: copyright (c) the Triad authors, released under the MIT License. JAX / Equinox port: copyright (c) the ilex authors, released under the Apache-2.0 / GPL-3.0 dual license used by ilex itself.

Upstream source

Original weights / reference implementation: https://github.com/wangshansong1/Triad

Provenance

This artefact was produced by ilex's save/load pipeline. The architecture is implemented in ilex.models.triad.TriadSwinViT and the weights have been converted from their upstream format. See the upstream source above for the canonical reference.