💻 computer science

ENSAM: an efficient foundation model for interactive segmentation of 3D medical images

ENSAM is a lightweight, foundation model for interactive 3D medical image segmentation that achieves state-of-the-art performance under limited data and computational budgets by combining a SegResNet-based architecture with advanced training techniques like normalized attention and the Muon optimizer, outperforming several pretrained baselines in the CVPR 2025 Foundation Models for Interactive 3D Biomedical Image Segmentation Challenge.

Original authors: Elias Stenhede, Agnar Martin Bjørnstad, Arian Ranjbar

Published 2026-07-02

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Elias Stenhede, Agnar Martin Bjørnstad, Arian Ranjbar

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

Imagine you are trying to teach a computer to find specific organs inside a 3D medical scan (like a CT or MRI). Usually, this requires a massive, expensive supercomputer and thousands of hours of training.

The authors of this paper, Elias, Agnar, and Arian, built a new tool called ENSAM. Think of ENSAM as a "smart, lightweight assistant" that can learn to segment (outline) any part of a 3D medical image just by looking at a few examples and taking a few hints from a human.

Here is the story of how they built it, explained simply:

1. The Goal: A "One-Man Army"

The team wanted to create a model that could do the job of a giant, heavy-duty foundation model but run on a single, standard computer graphics card (a 32 GB GPU). They wanted to prove you don't need a supercomputer to get great results.

The Name: They named it ENSAM (Equivariant, Normalized, Segment Anything Model). The name is a subtle joke: in Swedish and Germanic languages, "ensam" means "alone" or "lonely." This reflects their goal of training a powerful model all by itself, without help from other pre-trained giants.

2. The Recipe: How They Built It

They didn't just copy existing models; they cooked up a new recipe with three special ingredients:

The Brain (Image Encoder): Instead of using the newest, most complex "Transformer" brain (which is like a heavy, fuel-guzzling engine), they used a SegResNet. Think of this as a reliable, efficient workhorse engine that is proven to handle 3D medical data very well.
The GPS (Relative Positional Encoding): When you tell a computer "click here," it needs to know where "here" is. Old models use "Absolute Positioning" (like giving a street address: "123 Main St"). If you move the house, the address is wrong. ENSAM uses Relative Positioning (like saying "two blocks north of the park"). This is like giving the computer a flexible map that works no matter where the organ is in the body. This made the model learn much faster.
The Coach (Muon Optimizer): Training AI is like teaching a student. Usually, you use a standard teacher (an optimizer called Adam). The authors swapped this for a new, faster coach called Muon. Imagine a coach who doesn't just tell the student "try harder," but reorganizes the entire lesson plan instantly to find the best path forward. This helped the model learn in record time.

3. The Challenge: The "Coreset" Test

The paper describes a competition (the CVPR 2025 Challenge) where teams had to build these models.

The Twist: There was a special track called the "Coreset Track." Teams were only allowed to use 10% of the available training data. It was like being asked to learn a whole language using only a pocket dictionary instead of a library.
The Constraint: They also had to train on a single GPU in just 6 hours.

4. The Results: Punching Above Its Weight

Despite using only a fraction of the data and a single computer, ENSAM performed incredibly well:

The Score: It finished 5th out of 10 teams in the Coreset track.
The Big Win: It was the best team that didn't use any pre-trained weights. (Many other teams started with a model that had already been trained on massive datasets; ENSAM started from scratch, like a blank slate).
Comparison: It beat two other famous models (VISTA3D and SAM-Med3D) and came very close to the third (SegVol), even though ENSAM was trained with much less computing power.

5. How It Works in Practice (The "Interactive" Part)

Imagine a doctor looking at a 3D scan of a liver.

The doctor draws a box around the liver.
The model guesses the outline.
The doctor clicks "Yes" on the liver and "No" on a nearby kidney to correct the model.
The model instantly updates the outline.
ENSAM is designed to handle this back-and-forth conversation very quickly, refining its answer with every click.

6. Where It Stumbles (Limitations)

The authors are honest about where their model isn't perfect yet:

Tiny Details: In very small, dense images (like microscopy), the model sometimes misses tiny details or gets confused by noise.
The "Simulation" Gap: The model was tested by simulating clicks from a computer. The authors admit that real human doctors might click differently, and they haven't tested it with real people yet.
One at a Time: Currently, the model handles one object (like one kidney) at a time. It doesn't naturally understand that if you click "kidney," you definitely don't mean "liver."

Summary

In short, ENSAM is a proof-of-concept that you can build a highly effective, interactive 3D medical segmentation tool without needing a massive budget or a supercomputer. By using a "relative map" system and a "fast-track coach" for training, they created a model that learns quickly, works well on a single computer, and competes with much larger, more expensive systems.

Technical Summary: ENSAM – An Efficient Foundation Model for Interactive 3D Medical Image Segmentation

Problem Statement

Accurate segmentation of 3D medical images is critical for clinical tasks such as diagnosis and treatment planning. While deep learning models have automated these tasks, recent foundation models like SAM (Segment Anything Model) and SAM2, designed for natural images, struggle to generalize to medical modalities without significant adaptation. Existing medical-specific models often face trade-offs: some lack support for iterative refinement (e.g., MedSAM), others require massive computational resources and pretrained weights, and many fail to generalize across diverse modalities (CT, MRI, PET, Ultrasound, Microscopy). Furthermore, a unified benchmark for comparing these models on a standardized, hidden test set was previously lacking.

The specific challenge addressed in this work is the CVPR 2025 Foundation Models for Interactive 3D Biomedical Image Segmentation Challenge. The goal is to develop a universal, promptable model that can segment any anatomical or pathological structure given a user prompt (bounding box or clicks) within a simulated interactive setting. A critical constraint of the "coreset track" is the limitation to 10% of the total training data and the requirement to train from scratch without leveraging external pretrained weights, all while operating under strict computational budgets (single 32 GB GPU).

Methodology: ENSAM Architecture

The authors propose ENSAM (Equivariant, Normalized, Segment Anything Model), a lightweight, promptable model built on a U-Net-style architecture. It consists of three primary components: an image encoder, a prompt encoder, and a mask decoder.

1. Image Encoder

Instead of Vision Transformers, which often underperform CNNs in specific segmentation benchmarks, ENSAM utilizes a SegResNet-based encoder. This architecture comprises a cascade of residual blocks interleaved with downsampling layers (StridedConv3D). It transforms single-channel inputs into multi-channel feature maps, maintaining volumetric consistency.

2. Prompt Encoder and Interaction Mechanism

The prompt encoder handles 3D bounding boxes, foreground clicks, and background clicks. Key innovations include:

Lie Rotational Positional Encoding (LieRE): To preserve equivariance and improve efficiency, the model replaces absolute positional encoding with relative positional encoding. This is achieved by applying position-dependent rotation matrices to key and query vectors. The rotation matrix $R(p)$ is derived from the Lie algebra $so(n)$, ensuring attention scores depend only on the relative position $p_j - p_i$ .
Normalized Attention: Inspired by recent work on normalized transformers, ENSAM replaces standard LayerNorm and weight decay with $\ell_2$ normalization applied to all weight matrices and activations after each optimization step. This constrains outputs to a unit hypersphere, utilizing a modified residual update strategy:
$x \leftarrow \text{Norm}(\text{Norm}(x) + \lambda(\text{Norm}(\text{nBlock}(x)) - \text{Norm}(x)))$
where $\lambda$ is a learnable eigen learning rate. This aims to accelerate convergence and improve numerical stability.
Interaction Flow: User prompts and previous segmentation logits (downsampled via strided convolutions) interact with image embeddings through a four-step process: (1) normalized self-attention on prompts, (2) prompt-to-image attention, (3) MLP processing, and (4) image-to-prompt attention. This cycle repeats twice.

3. Mask Decoder

The decoder mirrors the encoder structure but uses single residual blocks per upsampling layer. It concatenates skip connections, processes them through a ResBlock3D, and applies trilinear upsampling to output segmentation logits.

4. Training Strategy

Optimizer: The authors employ the Muon optimizer, which operates on 2D weight matrices using an efficient approximation of Singular Value Decomposition (SVD) for updates, rather than standard gradient descent. For 1D parameters, AdamW is used.
Data Handling: Training involves simulated interactions where initial bounding boxes are generated from ground truth with random offsets, followed by five iterative clicks placed in the largest error regions (foreground for under-segmentation, background for over-segmentation).
Loss Function: A compound loss of Generalized Dice Loss and Cross-Entropy (weighted 1:2) is used.
Coreset Selection: For the coreset track, the training set was limited to ~4,471 samples (10% of total), selected to ensure uniform representation across modalities, excluding datasets with normalization or annotation issues.

Key Contributions

Efficient Foundation Model: ENSAM achieves competitive performance in universal 3D medical segmentation while being trained from scratch on a single 32 GB GPU in just 6 hours, using less than 5,000 volumes.
Architectural Innovations: The integration of Lie Rotational Positional Encoding and Normalized Attention within a SegResNet backbone significantly accelerates convergence and improves segmentation quality compared to absolute encoding and standard attention.
Optimizer Application: The paper demonstrates the efficacy of the Muon optimizer for 3D medical segmentation tasks, showing it outperforms AdamW in training speed and final loss.
Benchmarking: The work provides a rigorous evaluation on a hidden, multimodal test set, offering a fair comparison against state-of-the-art baselines (VISTA3D, SAM-Med3D, SegVol) under constrained data and compute conditions.

Results

Quantitative Performance

On the hidden test set (coreset track), ENSAM achieved:

DSC AUC: 2.404
NSD AUC: 2.266
Final DSC: 0.627
Final NSD: 0.597

Comparisons:

ENSAM outperformed VISTA3D and SAM-Med3D across all four metrics.
It matched SegVol in DSC AUC and NSD AUC but surpassed SegVol in Final DSC (0.627 vs. 0.622), though it trailed slightly in Final NSD.
In the overall challenge ranking, ENSAM placed 5th out of 10 teams and was the best-performing model that did not utilize pretrained weights.

Ablation Studies

Experiments confirmed that switching from absolute to relative positional encoding (LieRE) improved training efficiency. Furthermore, combining LieRE with the Muon optimizer resulted in the fastest convergence and the best final loss, outperforming the baseline configuration (Absolute PE + Adam).

Qualitative Analysis

Successes: ENSAM demonstrated accurate delineation of large structures (e.g., liver, spleen, kidneys in CT; vertebral bodies in MRI) and handled 2D+t ultrasound sequences effectively.
Failures: The model struggled with fine details in microscopy (thin vessels) and sometimes exhibited over-segmentation in PET images, likely due to intensity clipping during preprocessing. It also showed jagged edges in ultrasound due to motion artifacts.

Significance and Claims

The authors claim that ENSAM demonstrates that high-performance foundation models for medical imaging can be trained efficiently from scratch without relying on massive computational resources or external pretrained weights. The paper emphasizes that the combination of relative positional encoding and the Muon optimizer is a key driver for this efficiency.

The work highlights that while the model is not yet superior to all baselines on every metric (particularly when those baselines use significantly more data and compute), it achieves a strong balance of performance and resource efficiency. The authors modestly note that the current architecture is likely not optimal and that future work could benefit from exploring deeper networks, handling anisotropic spacing via physical coordinates, and integrating more diverse prompt types (e.g., scribbles, lassos).

The paper concludes that ENSAM represents a significant step toward democratizing access to foundation models for medical imaging, proving that effective interactive segmentation is achievable under strict constraints of data and hardware.