# Unleashing Diffusion and State Space Models for Medical Image Segmentation

arXiv: https://arxiv.org/abs/2506.12747 · HTML: https://mlsanity.com/paper/2506.12747/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 48 of its 68 references that are in the map.
- First posted on arXiv 15 Jun 2025.

## Builds on (most cited first, 30 of 48)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Segment Anything](https://mlsanity.com/paper/2304.02643/) (2023)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Masked-attention Mask Transformer for Universal Image Segmentation](https://mlsanity.com/paper/2112.01527/) (2021)
- [Large Language Models Encode Clinical Knowledge](https://mlsanity.com/paper/2212.13138/) (2022)
- [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://mlsanity.com/paper/2405.21060/) (2024)
- [Per-Pixel Classification is Not All You Need for Semantic Segmentation](https://mlsanity.com/paper/2107.06278/) (2021)
- [Sigmoid-Weighted Linear Units for Neural Network Function Approximation in Reinforcement Learning](https://mlsanity.com/paper/1702.03118/) (2017)
- [Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images](https://mlsanity.com/paper/2201.01266/) (2022)
- [Language-driven Semantic Segmentation](https://mlsanity.com/paper/2201.03546/) (2022)
- [The Medical Segmentation Decathlon](https://mlsanity.com/paper/2106.05735/) (2021)
- [TotalSegmentator: robust segmentation of 104 anatomical structures in CT images](https://mlsanity.com/paper/2208.05868/) (2022)
- [Open-Vocabulary Semantic Segmentation with Mask-adapted CLIP](https://mlsanity.com/paper/2210.04150/) (2022)
- [DenseCLIP: Language-Guided Dense Prediction with Context-Aware Prompting](https://mlsanity.com/paper/2112.01518/) (2021)
- [The Liver Tumor Segmentation Benchmark (LiTS)](https://mlsanity.com/paper/1901.04056/) (2019)
- [Scaling Open-Vocabulary Image Segmentation with Image-Level Labels](https://mlsanity.com/paper/2112.12143/) (2021)
- [Open-Vocabulary Object Detection Using Captions](https://mlsanity.com/paper/2011.10678/) (2020)
- [Self-Supervised Pre-Training of Swin Transformers for 3D Medical Image Analysis](https://mlsanity.com/paper/2111.14791/) (2021)
- [CRIS: CLIP-Driven Referring Image Segmentation](https://mlsanity.com/paper/2111.15174/) (2021)
- [AMOS: A Large-Scale Abdominal Multi-Organ Benchmark for Versatile Medical Image Segmentation](https://mlsanity.com/paper/2206.08023/) (2022)
- [CHAOS Challenge -- Combined (CT-MR) Healthy Abdominal Organ Segmentation](https://mlsanity.com/paper/2001.06535/) (2020)
- [A Tale of Two Features: Stable Diffusion Complements DINO for Zero-Shot Semantic Correspondence](https://mlsanity.com/paper/2305.15347/) (2023)
- [AbdomenCT-1K: Is Abdominal Organ Segmentation A Solved Problem?](https://mlsanity.com/paper/2010.14808/) (2020)
- [CLIP-Driven Universal Model for Organ Segmentation and Tumor Detection](https://mlsanity.com/paper/2301.00785/) (2023)

## Abstract (arXiv)

Existing segmentation models trained on a single medical imaging dataset often lack robustness when encountering unseen organs or tumors. Developing a robust model capable of identifying rare or novel tumor categories not present during training is crucial for advancing medical imaging applications. We propose DSM, a novel framework that leverages diffusion and state space models to segment unseen tumor categories beyond the training data. DSM utilizes two sets of object queries trained within modified attention decoders to enhance classification accuracy. Initially, the model learns organ queries using an object-aware feature grouping strategy to capture organ-level visual features. It then refines tumor queries by focusing on diffusion-based visual prompts, enabling precise segmentation of previously unseen tumors. Furthermore, we incorporate diffusion-guided feature fusion to improve semantic segmentation performance. By integrating CLIP text embeddings, DSM captures category-sensitive classes to improve linguistic transfer knowledge, thereby enhancing the model's robustness across diverse scenarios and multi-label tasks. Extensive experiments demonstrate the superior performance of DSM in various tumor segmentation tasks. Code is available at https://github.com/Rows21/k-Means_Mask_Mamba.