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

TL;DR

USEMA introduces a hybrid UNet architecture combining local CNN features with a scalable, efficient Mamba-like attention mechanism to improve medical image segmentation performance and efficiency.

Contribution

The paper proposes USEMA, a novel hybrid architecture that integrates local CNN features with a scalable attention mechanism for better segmentation.

Findings

USEMA achieves higher segmentation accuracy than pure CNN and Mamba models.

USEMA demonstrates improved computational efficiency over transformer-based models.

USEMA performs well across various modalities and image sizes.

Abstract

Accurate medical image segmentation is an integral part of the medical image analysis pipeline that requires the ability to merge local and global information. While vision transformers are able to capture global interactions using vanilla self-attention, their quadratic computational complexity in the input size remains a struggle for medical image segmentation tasks. Motivated by the dispersion property of vanilla self-attention and recent development of Mamba form of attention, Scalable and Efficient Mamba like Attention (SEMA) utilizes token localization via local window attention to avoid dispersion and maintain focusing, complemented by theoretically consistent arithmetic averaging to capture global aspect of attention. In this work, we present USEMA, a hybrid UNet architecture that merges the local feature extraction ability of convolutional neural networks (CNNs) with SEMA attention. We conduct experiments with USEMA across a variety of modalities and image sizes, demonstrating improved computational efficiency compared to transformer based models using full self-attention, and superior segmentation performance relative to purely convolution and Mamba-based models.