PHMamba: Preheating State Space Models with Context-Augmented Features for Medical Image Segmentation

Abstract

The recent Mamba model has demonstrated the competitive potential of State Space Models (SSMs) on various image benchmarks, particularly in modeling long-range sequences. However, most of the improvements in Mamba methods focus on scanning strategies, and lack an effective means of aggregating context with the structure of SSMs. In this paper, we design a framework that applies coarse segmentation map to preheat the SSMs, enabling effective and efficient information aggregation. The preheating method for SSMs is proposed to enhance its ability to learn contextual information, which plays a crucial role in medical image segmentation tasks where regional features are particularly prominent. To reduce the model’s potential bias towards coarse context, the compensation module is introduced, which allows the output sequence to partially revert to the original feature map. The results show that our method outperforms various competing approaches on different medical image tasks including cardiac segmentation and multi-organ segmentation. The code is available at https://github.com/nchennnn/PH-Mamba.