{S$^3$-Mamba}: Small-Size-Sensitive Mamba for Lesion Segmentation

📄 arXiv: 2412.14546v1 📥 PDF

作者: Gui Wang, Yuexiang Li, Wenting Chen, Meidan Ding, Wooi Ping Cheah, Rong Qu, Jianfeng Ren, Linlin Shen

分类: cs.CV

发布日期: 2024-12-19

备注: Accept by AAAI 2025

🔗 代码/项目: GITHUB

💡 一句话要点

提出S$^3$-Mamba，提升Mamba模型对小病灶分割的敏感性，助力早期疾病诊断。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 医学图像分割 小病灶检测 Mamba模型 状态空间模型 注意力机制 课程学习 深度学习 计算机视觉

📋 核心要点

  1. 现有模型在分割小病灶时面临挑战，因为小病灶占比小，下采样易丢失局部特征。
  2. S$^3$-Mamba从通道、空间和训练策略三方面入手，提升模型对小病灶的敏感性。
  3. 实验表明，S$^3$-Mamba在医学图像分割任务上表现优异，尤其擅长分割小病灶。

📝 摘要（中文）

小病灶在早期疾病诊断和严重感染干预中起着关键作用。现有模型在分割小病灶时面临挑战，因为小病灶仅占据图像的一小部分，并且下采样操作可能不可避免地丢失小病灶的局部特征。为了应对这些挑战，我们提出了一种小尺寸敏感的Mamba模型（S$^3$-Mamba），它从通道、空间和训练策略三个维度提高对小病灶的敏感性。具体来说，我们设计了一个增强视觉状态空间块，通过多个残差连接来关注小病灶，以保留局部特征，并通过通道注意力选择性地放大重要细节，同时抑制不相关的细节。我们设计了一种基于张量的跨特征多尺度注意力，以整合输入图像特征和具有边缘特征的中间层特征，并利用跨多个尺度的特征的注意力支持，从而保留各种粒度的小病灶的空间细节。最后，我们引入了一种新的正则化课程学习，以自动评估病灶大小和样本难度，并逐渐从简单样本关注到困难样本，如小病灶。在三个医学图像分割数据集上的大量实验表明，我们的S$^3$-Mamba具有优越性，尤其是在分割小病灶方面。

🔬 方法详解

问题定义：论文旨在解决医学图像中小病灶分割的难题。现有方法，特别是基于深度学习的模型，在处理小病灶时表现不佳。这是因为小病灶在图像中占比很小，容易在下采样过程中丢失关键特征，导致分割精度下降。此外，模型往往难以区分小病灶与背景，造成假阳性。

核心思路：S$^3$-Mamba的核心思路是从三个维度提升模型对小病灶的敏感性：通道维度、空间维度和训练策略。通过增强局部特征的保留、选择性地放大重要细节、整合多尺度特征以及采用课程学习策略，使模型能够更好地关注和分割小病灶。

技术框架：S$^3$-Mamba的整体框架基于Mamba模型，并在此基础上进行了改进。主要包含以下几个模块：1) 增强视觉状态空间块（Enhanced Visual State Space block），用于保留局部特征并进行通道注意力加权；2) 基于张量的跨特征多尺度注意力（Tensor-based Cross-feature Multi-scale Attention），用于整合输入图像特征、中间层特征和边缘特征；3) 正则化课程学习（Regularized Curriculum Learning），用于优化训练过程，使模型逐渐关注小病灶等困难样本。

关键创新：S$^3$-Mamba的关键创新在于其三维敏感性设计，即通道敏感性、空间敏感性和训练敏感性。通道敏感性通过增强视觉状态空间块实现，空间敏感性通过跨特征多尺度注意力实现，训练敏感性通过正则化课程学习实现。与传统方法相比，S$^3$-Mamba能够更有效地提取和利用小病灶的特征，从而提高分割精度。

关键设计：增强视觉状态空间块中使用了多个残差连接，以保留局部特征。通道注意力机制用于选择性地放大重要细节，抑制不相关细节。跨特征多尺度注意力利用张量运算整合不同尺度的特征。正则化课程学习通过自动评估病灶大小和样本难度，动态调整训练样本的权重。损失函数未知，但推测使用了Dice Loss或Cross Entropy Loss的变体，以更好地处理类别不平衡问题。具体参数设置在论文中未详细说明，属于未知信息。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，S$^3$-Mamba在三个医学图像分割数据集上均取得了优异的性能，尤其是在分割小病灶方面。具体性能数据未知，但摘要中强调了其在小病灶分割上的优越性。与现有方法相比，S$^3$-Mamba能够更准确地分割小病灶，从而提高诊断的准确性和可靠性。代码已开源，方便研究者复现和进一步研究。

🎯 应用场景

S$^3$-Mamba在医学图像分析领域具有广泛的应用前景，可用于辅助医生进行早期疾病诊断、病灶检测和分割，以及治疗效果评估。该模型能够提高小病灶的检出率和分割精度，从而帮助医生更准确地判断病情，制定更有效的治疗方案，并最终改善患者的预后。

📄 摘要（原文）

Small lesions play a critical role in early disease diagnosis and intervention of severe infections. Popular models often face challenges in segmenting small lesions, as it occupies only a minor portion of an image, while down_sampling operations may inevitably lose focus on local features of small lesions. To tackle the challenges, we propose a {\bf S}mall-{\bf S}ize-{\bf S}ensitive {\bf Mamba} ({\bf S$^3$-Mamba}), which promotes the sensitivity to small lesions across three dimensions: channel, spatial, and training strategy. Specifically, an Enhanced Visual State Space block is designed to focus on small lesions through multiple residual connections to preserve local features, and selectively amplify important details while suppressing irrelevant ones through channel-wise attention. A Tensor-based Cross-feature Multi-scale Attention is designed to integrate input image features and intermediate-layer features with edge features and exploit the attentive support of features across multiple scales, thereby retaining spatial details of small lesions at various granularities. Finally, we introduce a novel regularized curriculum learning to automatically assess lesion size and sample difficulty, and gradually focus from easy samples to hard ones like small lesions. Extensive experiments on three medical image segmentation datasets show the superiority of our S$^3$-Mamba, especially in segmenting small lesions. Our code is available at https://github.com/ErinWang2023/S3-Mamba.