Primus: Enforcing Attention Usage for 3D Medical Image Segmentation
About
Transformers have achieved remarkable success across multiple fields, yet their impact on 3D medical image segmentation remains limited with convolutional networks still dominating major benchmarks. In this work, (A) we analyze current Transformer-based segmentation models and identify critical shortcomings, particularly their over-reliance on convolutional blocks. Further, we demonstrate that in some architectures, performance is unaffected by the absence of the Transformer, thereby demonstrating their limited effectiveness. To address these challenges, we move away from hybrid architectures and (B) introduce Transformer-centric segmentation architectures, termed Primus and PrimusV2. Primus leverages high-resolution tokens, combined with advances in positional embeddings and block design, to maximally leverage its Transformer blocks, while PrimusV2 expands on this through an iterative patch embedding. Through these adaptations, Primus surpasses current Transformer-based methods and competes with a default nnU-Net while PrimusV2 exceeds it and is on par with the state-of-the-art CNNs such as ResEnc-L and MedNeXt architectures across nine public datasets. In doing so, we introduce the first competitive Transformer-centric model, making Transformers state-of-the-art in 3D medical image segmentation. The code is available here: https://github.com/MIC-DKFZ/nnUNet/blob/master/documentation/primus.md.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Cardiac Segmentation | ACDC | -- | 68 | |
| 3D Image Classification | MedMNIST 3D v2 (test) | Organ Accuracy0.972 | 36 | |
| Medical Image Segmentation | MSD Pancreas (test) | DSC79.7 | 30 | |
| 3D Medical Image Segmentation | LIDC | Dice Coefficient74 | 24 | |
| 3D Medical Image Segmentation | MMWHS MRI | Dice66.2 | 24 | |
| 3D Medical Image Segmentation | MMWHS CT | Dice Score0.726 | 24 | |
| Medical Image Segmentation | WORD | Dice84.15 | 23 | |
| Medical Image Segmentation | LiTS (test) | Dice Score (Average)79.9 | 20 | |
| Cerebral Lesion Segmentation | SBM | Dice0.6636 | 20 | |
| Medical Image Segmentation | LiTS folds 1-4 (dev) | DSC81.73 | 19 |