WAS-Mamba: 3D Medical Image Segmentation via Windowed Attention State Space Model

Mamba, the state space model (SSM), has attracted significant attention for its ability to model long-range dependencies with linear complexity, achieving success in medical image segmentation. However, the previous cross-scanning approach in Mamba struggles to capture both long-range and short-range dependencies simultaneously and treats the features of each path equally. This imbalance between local and global modeling capabilities can adversely impact segmentation performance. To address these challenges, we propose WAS-Mamba, a novel method specifically designed for Mamba-based medical image segmentation. WAS-Mamba introduces a cross-channel window scanning strategy (CCWScan) that enables sequences to preserve original local image features during the transformation process. Furthermore, WAS-Mamba employs a weighted state space model (WSSM) to dynamically fuse spatial and frequency domain information, improving the capture of local details and global context for accurate segmentation. We validated the superior performance of WAS-Mamba across five datasets covering different anatomical regions, which include CT and MRI images: Synapse, BTCV, ACDC, BraTS, and Decathlon-Lung. In particular, we achieved a Dice coefficient of 88.09% on the Synapse dataset, with a 33% reduction in computational complexity and inference time compared to the second-best model. The code and model will be released at https://github.com/1605066114/WAS-Mamba.