ilex-hub/triad.plainconvunet-mae.1

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.