library_name: ilex
tags:
- jax
- equinox
- ilex
- neuroimaging
- 3d-mri
license: mit
license_link: https://github.com/wangshansong1/Triad/blob/main/LICENSE
Triad (3D-MRI self-supervised foundation backbone) -- Triad PlainConvUNet 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. The forward returns the per-stage encoder skip pyramid (channels 32, 64, 128, 256, 320, 320); a downstream consumer attaches its own decoder / classification / registration head. Single-channel input with each spatial dim a multiple of 32. Pretrained by masked autoencoding (MAE) on Triad-131K. Encoder-only: the SSL reconstruction decoder is not shipped.
Usage
from ilex.models.triad import TriadPlainConvUNet
model = TriadPlainConvUNet.from_pretrained('ilex-hub/triad.plainconvunet-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 nnU-Net backbone the encoder is taken from: Isensee F., Jaeger P. F., Kohl S. A. A., Petersen J., Maier-Hein K. H. (2021). nnU-Net: a self- configuring method for deep learning-based biomedical image segmentation. Nature Methods 18(2):203-211. doi:10.1038/s41592-020-01008-z
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.TriadPlainConvUNet and the weights have been converted
from their upstream format. See the upstream source above
for the canonical reference.