💻 computer science

Primus: Enforcing Attention Usage for 3D Medical Image Segmentation

This paper introduces Primus and PrimusV2, the first competitive Transformer-centric architectures for 3D medical image segmentation that overcome the limitations of hybrid models by maximizing attention usage, thereby achieving state-of-the-art performance that surpasses or matches leading CNN-based methods across nine public datasets.

Original authors: Tassilo Wald, Saikat Roy, Fabian Isensee, Constantin Ulrich, Sebastian Ziegler, Dasha Trofimova, Raphael Stock, Michael Baumgartner, Gregor Köhler, Klaus Maier-Hein

Published 2026-05-01

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Tassilo Wald, Saikat Roy, Fabian Isensee, Constantin Ulrich, Sebastian Ziegler, Dasha Trofimova, Raphael Stock, Michael Baumgartner, Gregor Köhler, Klaus Maier-Hein

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Imposter" Transformers

Imagine you are trying to build a robot that can look at a 3D medical scan (like a CT scan of a kidney) and draw a perfect outline around the tumor.

For a long time, the best robots for this job were CNNs (Convolutional Neural Networks). Think of a CNN like a very skilled, local detective. It looks at a small neighborhood of pixels, gathers clues, and moves to the next neighborhood. It's great at seeing local details but sometimes misses the "big picture" of the whole organ.

Then, Transformers arrived. These are like "super-observers" that can look at the entire image at once and understand how the top of the kidney relates to the bottom. They became famous in language processing (like Chatbots) and natural photos. Everyone thought, "If Transformers can read a whole book or see a whole landscape, they must be perfect for medical scans!"

But here's the catch: When researchers tried to use Transformers for 3D medical scans, they didn't work very well. They were often slower and less accurate than the old CNN detectives.

The Investigation: Who is doing the work?

The authors of this paper decided to investigate why these Transformers were failing. They looked at nine popular "hybrid" models (models that try to use both CNNs and Transformers).

They discovered a shocking secret: The Transformers were mostly sleeping.

The Analogy: Imagine a construction crew where you hired a famous, expensive architect (the Transformer) to design a house, but you also hired 100 regular bricklayers (the CNNs). When the house was finished, the authors realized the architect didn't actually draw the plans. The bricklayers built the whole house themselves, and the architect just stood there nodding.
The Evidence: When the researchers removed the "architect" (the Transformer blocks) from these models, the models performed almost exactly the same. In some cases, removing the Transformer even made them faster and slightly better because the model stopped wasting energy on a useless component.

The paper calls this the "UNet Index." Most existing models had a high index, meaning they were actually just CNNs in disguise, carrying a heavy, useless Transformer backpack.

The Solution: Primus and PrimusV2

The authors asked: "What if we built a model where the Transformer has to do the work?" They created two new architectures named Primus (Latin for "First") and PrimusV2.

Here is how they forced the Transformer to wake up and do its job:

1. Don't Squash the Details (The Tokenizer)

Standard Transformers often chop the 3D image into huge chunks (like cutting a cake into giant, thick slices) to turn them into data tokens. This throws away tiny, important details, like a small tumor or a thin blood vessel.

The Fix: Primus uses a "high-resolution tokenizer." Instead of giant slices, it uses smaller, finer slices (8x8x8 voxels).
The Analogy: Instead of looking at a map of a city where every street is just a blurry line, Primus looks at a map where you can see individual houses. This gives the Transformer more detailed information to work with.

2. The "Iterative" Approach (PrimusV2)

Even with smaller slices, the Transformer sometimes struggled to understand the complex 3D shapes of organs right away.

The Fix: PrimusV2 uses an "Iterative Patch Embedding." Instead of trying to understand the whole chunk in one go, it processes the image in stages, like peeling an onion or refining a sketch. It uses a few small, smart convolutional steps to prepare the data before the Transformer sees it.
The Analogy: Imagine trying to solve a complex puzzle. Primus doesn't just dump all the pieces on the table. It first sorts them by color and edge pieces (the iterative steps), making it much easier for the "super-observer" (the Transformer) to see the final picture.

3. Giving it a GPS (3D RoPE)

Transformers are naturally "blind" to space; they don't know that "left" is different from "right" unless you tell them.

The Fix: The authors added a special 3D "Rotary Position Embedding" (RoPE).
The Analogy: It's like giving the Transformer a GPS system that understands depth, width, and height simultaneously. It knows exactly where a tumor is located in 3D space relative to the rest of the organ.

The Results: The Transformer Finally Wins

After these changes, the results were impressive:

Primus became the first Transformer-based model that could compete with the standard "gold standard" CNN (nnU-Net).
PrimusV2 didn't just compete; it beat the standard CNN on most tests and matched the very best, most complex CNNs available today.

The Key Takeaway:
The paper proves that Transformers can be the best tool for 3D medical imaging, but only if you stop treating them like a sidekick to a CNN. You have to design the system so the Transformer is the main character, fed with high-resolution data and given the right tools to understand 3D space.

Summary in One Sentence

The authors built a new AI model called Primus that finally forces the "super-observer" Transformer to do the heavy lifting in medical imaging, proving that when designed correctly, Transformers can outperform the traditional "local detective" CNNs.

1. Problem Statement

Despite the success of Transformers in natural language processing and natural image computer vision, their impact on 3D medical image segmentation has been limited. Current benchmarks are still dominated by Convolutional Neural Networks (CNNs), particularly the nnU-Net framework. The authors identify two critical shortcomings in existing Transformer-based approaches for medical imaging:

Over-reliance on Convolutions: Most "Transformer" architectures are actually hybrid models where the majority of parameters and computational capacity reside in convolutional blocks (encoders/decoders), rendering the Transformer components merely auxiliary.
Ineffectiveness of Attention: In many hybrid models, removing the Transformer blocks results in negligible performance drops, suggesting the models are learning representations via CNNs alone. Furthermore, existing pure or near-pure Transformer models often underperform strong CNN baselines, likely due to the difficulty of training Transformers from scratch on small, sparsely labeled medical datasets.

2. Methodology

The authors propose a two-pronged approach: first, a rigorous deconstruction of existing models to prove their limitations, and second, the design of new, truly Transformer-centric architectures.

A. Analysis of Existing Architectures

The authors analyzed nine popular Transformer-based models (e.g., TransUNet, SwinUNETR, UNETR) using a metric called the UNet Index.

UNet Index Definition: A ratio measuring the amount of convolutional parameters outside the Transformer blocks relative to a standard U-Net.
Findings: Most models have a high UNet Index (heavy CNN reliance). When the Transformer blocks were replaced with identity mappings, 6 out of 9 models showed less than a 3% drop in Dice Similarity Coefficient (DSC), proving the Transformer was not the primary driver of performance.

B. The Primus and PrimusV2 Architectures

To address these issues, the authors designed Primus and PrimusV2, the first competitive Transformer-centric architectures for 3D medical imaging. These models minimize convolutional parameters, forcing the network to rely on attention mechanisms.

Key Design Innovations:

High-Resolution Tokenization:
- Unlike standard ViTs that use large $16\times16\times16$ patches (which lose local detail), Primus uses a $8\times8\times8$ patch size to preserve fine-grained anatomical structures.
- PrimusV2 introduces an Iterative Patch Embedding. Instead of a single projection, it uses a chain of residual convolutional blocks to iteratively downsample and extract semantic features before tokenization. This creates richer tokens without relying on a heavy CNN backbone.
3D Rotary Positional Embeddings (RoPE):
- Standard absolute positional embeddings struggle with the spatial shifts common in 3D medical data. The authors adapted Axial Rotary Position Embeddings (RoPE) to 3D, encoding relative distances and orientations directly into the attention mechanism to better capture volumetric spatial correlations.
Modern Transformer Components:
- MLP Blocks: Adoption of SwiGLU (Swish-activated Gated Linear Units) with LayerNorm, proven effective in natural image Transformers (e.g., EVA-02).
- Stability Mechanisms: Use of LayerScale and Post-Attention Normalization to stabilize training on large models (Primus-L) and prevent divergence.
- Lightweight Decoder: A decoder composed of transposed convolutions that upsamples tokens independently, minimizing convolutional influence and keeping the UNet Index low.

3. Key Contributions

Empirical Deconstruction: Provided quantitative evidence (via UNet Index and ablation studies) that current "Transformer" medical models are effectively CNNs with minor attention additions.
First Competitive Transformer-Centric Models: Introduced Primus and PrimusV2, which achieve state-of-the-art performance while keeping the vast majority of parameters and FLOPs within the Transformer blocks.
Architectural Innovations: Demonstrated that combining high-resolution tokenization ( $8\times8\times8$ ), iterative patch embedding, and 3D RoPE allows Transformers to overcome the data scarcity and local feature requirements of medical imaging.
Efficiency in Self-Supervised Learning: Highlighted that Transformer-centric designs are more efficient for Masked Image Modeling (MIM) pre-training, using significantly less VRAM (40% less) compared to CNN-heavy architectures when masking tokens.

4. Experimental Results

The models were evaluated on nine public datasets (including AMOS22, KiTS23, TotalSegmentator-BTCV, SBM, etc.) covering various modalities (CT, MRI) and organs.

Performance vs. Baselines:
- Primus: Surpassed current Transformer-based methods and competed directly with the default nnU-Net, exceeding it on 4 out of 9 datasets.
- PrimusV2: Significantly outperformed all other Transformer baselines (including UNETR++, SwinUNETRv2) and reached parity with state-of-the-art CNNs like ResEnc-L and MedNeXt.
- Gap Reduction: PrimusV2 reduced the performance gap to the best CNNs to within 0.05–0.29 DSC points, a massive improvement over previous Transformer attempts which lagged by 8–11 DSC points.
Ablation Insights:
- Removing the Transformer from PrimusV2 caused a ~5.0 DSC drop, confirming the model relies heavily on its attention mechanisms (unlike previous hybrids where the drop was near 0).
- The Iterative Tokenizer in PrimusV2 provided a ~2 DSC gain over the standard projection, particularly improving performance on small lesions (e.g., brain metastases in the SBM dataset).
Scaling: While scaling up model size (S to L) showed diminishing returns on very small datasets (due to overfitting), PrimusV2-L achieved the highest overall scores on larger datasets.

5. Significance

This work marks a paradigm shift in 3D medical image segmentation:

Viability of Pure Transformers: It proves that Transformers can be the primary driver of performance in medical imaging, not just an add-on to CNNs.
Future-Proofing: By establishing a strong Transformer-centric baseline, the paper enables the adoption of self-supervised learning (SSL) and multi-modal integration (e.g., combining images with text reports) in medical AI, areas where Transformers excel but CNNs struggle.
Clinical Applicability: Inference time analysis shows that even the largest PrimusV2 models fit within clinical decision windows (e.g., stroke detection), making them viable for real-world deployment.

In conclusion, Primus and PrimusV2 demonstrate that with the right architectural adaptations (high-res tokens, 3D RoPE, iterative embedding), Transformers can match or exceed the best CNNs in 3D medical segmentation, opening the door for next-generation medical AI pipelines.