💻 computer science

USEMA: a Scalable Efficient Mamba Like Attention for Medical Image Segmentation

This paper introduces USEMA, a hybrid UNet architecture that integrates a novel Scalable and Efficient Mamba-like Attention (SEMA) mechanism to achieve superior medical image segmentation performance and computational efficiency by effectively combining local feature extraction with global context modeling.

Original authors: Elisha Dayag, Nhat Thanh Tran, Jack Xin

Published 2026-05-13

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Elisha Dayag, Nhat Thanh Tran, Jack Xin

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 solve a giant jigsaw puzzle of a human body, but the pieces are tiny, blurry, and there are thousands of them. This is what doctors face when they look at medical images like MRIs or microscope slides to find tumors or organs. To help them, computer scientists build "AI detectives" that can automatically draw lines around these body parts. This process is called medical image segmentation.

The paper you shared introduces a new AI detective called USEMA. Here is how it works, explained simply:

The Problem: The "Too Many Neighbors" Dilemma

To understand an image, an AI needs to look at two things:

The Details: What is happening right next to a specific pixel? (Is this a liver cell or a kidney cell?)
The Big Picture: How does this pixel relate to the rest of the image? (Is this a tumor on the left side of the body?)

Older AI models (called Transformers) were great at looking at the "Big Picture." They could connect any two pixels in the image instantly. However, they had a major flaw: they were incredibly slow and expensive to run. It was like trying to introduce every single person in a stadium of 100,000 people to every other person individually. The math gets so heavy (quadratic complexity) that it becomes impossible to do quickly on large medical scans.

Newer models (called Mamba) are faster because they look at the image in a line, like reading a book. But sometimes, they get a bit "myopic," focusing too much on the immediate neighbors and missing the global context.

The Solution: USEMA (The "Smart Hybrid")

The authors created USEMA (a hybrid of a classic "U-Net" shape and a new attention mechanism called SEMA). They solved the speed vs. accuracy problem with a clever two-step strategy, using a Window and a Whistle analogy:

The Window (Local Focus):
Imagine you are in a crowded room. To understand your immediate surroundings, you only look at the people standing within a 5-foot circle around you. This is Window Attention. It's fast and efficient because you aren't trying to talk to everyone in the stadium, just your neighbors. This handles the fine details.
The Whistle (Global Focus):
But what if you need to know if someone is shouting from the other side of the room? The paper noticed that when AI models look at too many things at once, the importance of any single item gets diluted (like a whisper in a hurricane). To fix this, they added a simple, mathematically proven trick: Arithmetic Averaging.

Think of this as the AI taking a quick "average" of everything it sees. It's not a deep, complex conversation with every pixel; it's a quick summary. The paper argues that this simple average is actually enough to capture the "global" feeling of the image without the heavy math cost.

USEMA combines these two: It uses the "Window" to see the details and the "Average" to get the general vibe of the whole image.

How They Tested It

The team didn't just guess; they put USEMA to the test in three very different "puzzle rooms":

The Abdomen MRI: A 3D scan of internal organs (liver, kidneys, etc.).
The Endoscopy: A video camera inside the body looking at surgical tools.
The Microscope: Tiny images of individual cells.

The Results

In every single test, USEMA won:

Better Accuracy: It drew the outlines of organs and cells more correctly than the previous best models (both the slow Transformers and the fast Mamba models).
Smaller Size: It achieved these results with fewer "brain cells" (parameters) than the heavy Transformer models, making it lighter and faster to run.
Efficiency: It proved that you don't need to do the heavy math of connecting every pixel to every other pixel to get great results. A smart mix of "local windows" and a "simple average" works better.

The Bottom Line

The paper claims that USEMA is a new, faster, and more accurate way for computers to understand medical images. By mixing the "local focus" of looking at neighbors with a "global average" of the whole scene, it avoids the computational bottleneck that has held back AI in medicine for years. It's like finding a way to understand a whole city by knowing your street and having a quick map of the whole area, rather than trying to memorize every single building in the city at once.

Technical Summary: USEMA – A Scalable Efficient Mamba Like Attention for Medical Image Segmentation

Problem Statement
Accurate medical image segmentation is critical for diagnostic pipelines but faces unique challenges, including limited training data and complex, overlapping object geometries. While Vision Transformers (ViTs) offer the ability to capture global dependencies through self-attention, their application in medical imaging is hindered by two primary issues:

Quadratic Computational Complexity: Vanilla self-attention incurs an $O(n^2)$ cost relative to sequence length, making it computationally prohibitive for high-resolution medical images.
Attention Dispersion: As sequence length increases, the scores in the attention matrix (softmax output) tend to disperse uniformly toward zero. This phenomenon, theoretically grounded in the weak law of large numbers, hinders the model's ability to select important keys and values, effectively reducing the attention mechanism to a uniform averaging process.

Methodology
The authors propose USEMA, a hybrid UNet architecture that integrates a novel attention mechanism called Scalable Efficient Mamba-like Attention (SEMA). The methodology addresses the limitations of both pure Transformers and existing Mamba-based models through the following components:

SEMA Mechanism: SEMA combines local and global approximations to mimic the behavior of full self-attention while maintaining linear complexity.
- Local Component: It utilizes window attention to maintain token localization and focus, avoiding the dispersion issue inherent in long sequences.
- Global Component: To capture global context without quadratic cost, SEMA employs theoretically consistent arithmetic averaging. This term approximates the global aspect of attention, reflecting the $O(1/n)$ bounds observed in the dispersion inequality. The global term is defined as the average of all value vectors, broadcasted across the sequence.
- Integration: The SEMA block is inserted into a Mamba macro structure, utilizing conditional positional embeddings, layer normalization, and gating mechanisms similar to Mamba (e.g., SiLU activation and depthwise convolutions).
USEMA Architecture: The model adopts a symmetric U-shaped structure (UNet).
- Encoder: Consists of residual blocks followed by SEMA blocks. Features are reshaped from $(B, C, H, W)$ to $(B, C, HW)$ before entering the attention mechanism. Rotary Positional Embeddings and locally-enhanced position encodings are used to augment spatial awareness.
- Bottleneck: Due to spatial compression, the bottleneck utilizes full self-attention rather than windowed attention, as the sequence length is sufficiently small to be computationally manageable.
- Decoder: Composed of residual blocks and transposed convolutions, utilizing skip connections to fuse hierarchical features from the encoder.

Key Contributions

Hybrid Attention Design: The paper introduces SEMA, which theoretically justifies the use of arithmetic averaging as a global approximation for self-attention, complementing local window attention to prevent score dispersion.
Architectural Integration: USEMA successfully merges the local feature extraction efficiency of CNNs with the global modeling capabilities of Mamba-like attention within a UNet framework.
Comprehensive Benchmarking: The model is evaluated across three diverse medical imaging modalities (Abdomen MRI, Endoscopy, and Microscopy) with varying resolutions and segmentation targets (organ, instrument, and cell segmentation).

Experimental Results
USEMA was benchmarked against state-of-the-art Transformer-based models (UNETR, Swin-UNETR, nnFormer) and Mamba-based/inspired models (U-Mamba, Swin-UMamba, Mamba UNet, MLLA-UNet) on three datasets:

Abdomen MRI (MICCAI 2022 AMOS): USEMA achieved a Dice Similarity Coefficient (DSC) of 0.7704 and Normalized Surface Distance (NSD) of 0.8345. This outperformed the best Transformer baseline (nnFormer) by 5.83% in DSC and the best Mamba baseline (U-Mamba Enc) by 10.53% in DSC. Notably, USEMA achieved these results with 52M parameters, which is 13% fewer than the best-performing Transformer model.
Endoscopy (MICCAI 2017): USEMA achieved a DSC of 0.6463 and NSD of 0.6621, outperforming Transformer baselines by ~5.6% and Mamba baselines by ~13.5% in DSC.
Microscopy (NeurIPS 2022): In this instance segmentation task, USEMA achieved an F1 score of 0.5791, surpassing the best Transformer baseline by 8.61% and the Mamba-inspired MLLA-UNet by 19.23%. It also utilized 43% fewer parameters than U-Mamba Enc.

Ablation Study
An ablation study confirmed the necessity of the global attention approximation. Removing the arithmetic averaging term (using only window attention) resulted in performance drops across all datasets (e.g., a drop of 0.013 in DSC on the Abdomen MRI dataset), validating the contribution of the global term.

Significance and Claims
The authors claim that USEMA effectively balances computational efficiency and segmentation accuracy. By leveraging the dispersion property of vanilla self-attention and the linear complexity of Mamba, USEMA provides a scalable solution that:

Outperforms purely convolutional, Transformer-based, and Mamba-based models across diverse 2D medical imaging benchmarks.
Reduces computational costs compared to full self-attention Transformers.
Demonstrates that a theoretically grounded, efficient global approximation (arithmetic averaging) can successfully complement local attention mechanisms in medical image analysis.

The paper concludes with plans to extend the averaging mechanism to learned or weighted averaging and to apply USEMA to large-scale pathology images, but does not claim current success in these specific future domains.