论文解读-SegMamba: Long-range Sequential ModelingMamba For 3D Medical Image Segmentation

SegMamba面向三维医学图像分割的长程序列建模 Mamba 模型MICCAI 2024Medical Image Computing and Computer Assisted Intervention医学图像计算与计算机辅助介入顶会代码链接GitHub - ge-xing/SegMamba: SegMamba: Long-range Sequential Modeling Mamba For 3D Medical Image Segmentation · GitHubhttps://github.com/ge-xing/SegMambaAbstractThe Transformer architecture has demonstrated remarkable results in 3D medical image segmentation due to its capability of modeling global relationships. However, it poses a significant computational burden when processing high‑dimensional medical images. Mamba, as a State Space Model (SSM), has recently emerged as a notable approach for modeling long‑range dependencies in sequential data, and has excelled in the field of natural language processing with its remarkable memory efficiency and computational speed. Inspired by this, we devise SegMamba, a novel 3D medical image Segmentation Mamba model, to effectively capture long‑range dependencies within whole‑volume features at every scale. Our SegMamba outperforms Transformer‑based methods in whole‑volume feature modeling, maintaining high efficiency even at a resolution of 64 × 64 × 64, where the sequential length is approximately 260k. Moreover, we collect and annotate a novel large‑scale dataset (named CRC‑500) to facilitate benchmarking evaluation in 3D colorectal cancer (CRC) segmentation. Experimental results on our CRC‑500 and two public benchmark datasets further demonstrate the effectiveness and universality of our method.翻译Transformer 架构凭借对全局关系的建模能力在三维医学图像分割任务上取得了优异效果。但在处理高维医学图像时该模型会带来巨大的计算开销。Mamba 作为状态空间模型SSM是序列数据长距离依赖建模的新兴方案凭借出色的内存利用率与运算速度在自然语言处理领域表现突出。受此启发本文提出全新的三维医学图像分割 Mamba 模型 SegMamba能够在各个尺度下高效捕获整个体数据特征内的长距离依赖。在整体体特征建模任务上SegMamba 性能优于基于 Transformer 的方法即便面对分辨率为 64*64*64、序列长度约 26 万的输入依旧可以保持很高的运算效率。此外本文收集并标注了全新大规模数据集 CRC‑500用于三维结直肠癌分割的基准评测。在 CRC‑500 数据集以及另外两个公开基准数据集上开展的实验进一步验证了该方法的有效性与通用性。创新点1提出面向三维医学体分割的TSMambaTri‑orientated Spatial Mamba基础模块模块内部包含 GSC 门控空间卷积与ToM 三向 Mamba 子模块。区别于 U‑Mamba 直接将 3D 特征整体展平为单方向一维序列的做法ToM 模块分别沿前向、反向、切片间三个不同方向对三维特征做 Mamba 序列扫描从多个维度捕获体数据全局依赖GSC 模块放置在 ToM 之前利用双分支卷积门控机制强化三维空间特征表达弥补 Mamba 将特征展平为一维序列带来的空间信息损失。TSMamba 块内部采用层归一化、残差连接与 MLP构成完整可堆叠的编码器单元能够在多个特征尺度上完成完整体数据的长程建模即便序列长度达到 260k 依然维持线性计算复杂度。2设计FUE 特征级不确定性估计模块嵌入 U 型网络的跳跃连接支路中。该模块对编码器多尺度特征计算特征不确定性权重对不确定性低、可靠性高的特征做增强抑制不确定性高的噪声特征优化编码器向解码器传递的跨层特征提升多尺度特征复用的鲁棒性。不同于普通直接拼接 / 相加的 skip‑connectionFUE 通过不确定性加权筛选特征缓解三维医学图像中背景、病灶等不同结构带来的特征混杂问题。3构建全新大规模 3D 结直肠癌 CT 数据集CRC‑500数据集包含 500 例专家标注的三维结直肠癌 CT 体数据填补了公开结直肠癌三维分割数据集规模偏小的短板为该任务提供新的基准评测资源。论文同时在 BraTS2023 脑肿瘤、AIIB2023 气道分割两个公开基准数据集开展充分对比实验在多模态、多器官、多病灶类型任务上验证整套框架的有效性与通用性。4搭建编码器使用 TSMamba 模块、解码器采用纯卷积的 U 型分割整体架构。编码器堆叠 TSMamba 模块完成多尺度全局上下文挖掘解码器仅依靠卷积残差块完成分辨率重建与局部细节恢复整套架构不需要在解码器引入开销较大的 SSM 运算在保证全局建模能力的前提下控制整体计算量。消融实验证明 GSC、ToM、FUE 三个子模块均带来性能增益与 CNN、Transformer 基线对比SegMamba 在大长度序列输入下不会出现 OOM 显存溢出兼顾分割精度与推理内存效率。SegMamba 整体架构图为 SegMamba 网络整体架构总览。论文原文描述SegMamba 的编码器由一个 Stem 层与多个 TSMamba Block 构成用于提取多尺度特征在每个 TSMamba Block 内部门控空间卷积 GSC 模块负责建模空间特征三向 MambaToM模块从不同方向表征全局信息除此之外本文设计特征级不确定性估计 FUE 模块对多尺度特征进行筛选实现更加鲁棒的特征复用。整个网络采用经典 U 型 3D 医学分割编解码结构。输入为三维医学体数据C*D*H*WC*D*H*W是三维医学体数据张量C为成像模态通道D、H、W分别对应体数据的深度、高度、宽度三个空间维度。首先送入Stem 层由深度卷积完成初步的底层特征提取与通道变换。Stem 之后堆叠多层TSMamba Block每一个 TSMamba 模块后紧跟下采样操作逐步压缩特征图的空间尺寸提升通道数逐层生成多尺度特征 z^0,z^1,z^2,z^3,z^4完成编码器的多尺度上下文挖掘。图左侧虚线框内是TSMamba Block内部细节。特征首先输入 GSC 门控空间卷积模块强化三维局部空间特征表达随后经过层归一化送入 ToM 三向 Mamba 模块沿三个不同空间方向执行 Mamba 序列扫描捕获三维体数据的全局长距离依赖ToM 输出经过残差相加再经过一层层归一化与 MLP再次使用残差连接形成完整可堆叠的编码器基础单元模块内部两层残差通路保障深层网络稳定训练。编码器每一层输出的多尺度特征 z^i会送入FUE 特征级不确定性估计模块也就是图中编码器和解码器之间的 FUE 单元。FUE 模块计算特征的不确定性权重筛选并优化跨层传递的特征增强可靠特征、抑制噪声特征经过 FUE 处理后的特征送入残差块作为跳跃连接特征传递到解码器。补充该公式用于计算 FUE 模块的不确定性权重图 u^i。首先对第i尺度的多通道编码器特征在通道维度上求取均值再经过 Sigmoid 函数归一化得到 bar z^i接着采用信息熵公式计算每个空间位置的不确定性。当 bar z^i 接近 0 或者 1 时熵值 u^i很小代表该位置特征响应明确、可信度高当 bar z^i 趋近 0.5 时熵达到最大值代表该位置特征模糊、不确定性高。得到不确定性图 u^i之后论文继续使用对原始特征做自适应调制增强高置信区域抑制噪声、模糊等高不确定区域对跳跃连接的特征进行筛选优化。解码器由转置卷积Transpose Conv图中向上箭头与残差块构成不使用 SSM 模块。解码器逐级上采样恢复空间分辨率每一级将上采样后的高层特征与 FUE 筛选后的编码器同层级跳跃特征融合依靠残差卷积块完成局部细节重建。多层解码器逐级融合多尺度信息最终特征送入 Seg-head 分割头输出最终三维分割预测结果。这套架构将局部空间建模、多向全局 SSM 建模与不确定性特征筛选结合在一起。GSC 负责局部空间信息ToM 从三维多方向捕获全局依赖FUE 优化跳跃连接的特征质量三者相互配合编码器使用 TSMamba 做全局上下文提取解码器仅使用卷积残差块完成重建在长序列三维体分割任务中在保证精度的同时控制显存开销避免 Transformer 类模型的二次复杂度带来的显存溢出问题。下面再讲解一下 TSMamba 的内部细节GSCGated Spatial Convolution门控空间卷积模块下图是 SegMamba 提出的门控空间卷积 GSC 模块它是 TSMamba Block 的首个子模块作用是在送入 ToM 三向 Mamba 之前提前强化三维局部空间特征弥补 Mamba 将 3D 特征展开为序列后容易丢失局部空间结构信息的缺陷。输入是三维特征图模块分为上下双分支门控结构。上分支采用 Norm 3×3×3 3D卷积 非线性激活利用 3×3×3 大感受野的三维卷积提取体数据局部空间邻域信息下分支采用 Norm 1×1×1 3D卷积 非线性激活负责通道维度的特征变换生成门控权重。两个分支输出逐元素相乘完成门控调制用 1×1×1 分支学习到的权重自适应筛选 3×3×3 卷积提取的空间特征抑制无效特征、保留有用局部空间信息。调制后的结果再送入右侧另一组 Norm3×3×3卷积非线性激活做进一步特征映射。最后模块采用残差连接将原始输入特征与处理后的特征相加输出。整体来看GSC 通过双分支门控卷积专门增强三维局部空间表达能力。在 TSMamba 里GSC 放在 ToM 三向 Mamba 模块的前面先把局部器官边界、组织纹理这类空间信息提取并强化再交给 ToM 去建模全局长距离依赖实现局部空间特征与全局序列建模的分工配合。ToMTri-orientated Mamba三向 Mamba模块下图为 TSMamba Block 中的三向 Mamba 模块 ToM是 SegMamba 的核心创新模块。输入为经过前置 GSC 处理后的三维特征体ToM 设计了三条独立的 Mamba 序列建模支路从三个不同方向对 3D 体数据做特征交互最后融合输出三维特征以此充分捕获三维医学图像的全局长程依赖。第一条支路是前向特征交互Forward Features Interaction沿着空间序列从前到后的顺序执行 Mamba 扫描学习序列靠前体素对后续体素的信息传递。第二条支路为反向特征交互Reverse Features Interaction将序列反转后再送入 Mamba捕捉从后向前的信息关联。前向 反向两个支路完成同一个空间维度上双向全局依赖建模。第三条支路是切片间特征交互Inter-slice Features Interaction沿着 CT/MRI 的切片深度 D 方向做 Mamba 序列扫描专门建模跨切片之间的上下文关联对应医学影像逐层切片的空间连续性。三条支路各自独立执行 Mamba 状态空间序列建模每条支路输出对应的三维特征最后将三路特征融合得到最终的 Fused 3D Features。总结与对比和 U-Mamba 直接把整个 3D 特征展平为单一方向一维序列的做法不同ToM 从三个不同空间维度分别进行序列建模更加贴合三维医学 Volume 的空间结构能够完整保留三维空间的上下文关系。ToM 依托 Mamba 线性复杂度的优势即使面对64*64*64、序列长度约 260k 的大尺度体特征也不会出现 Transformer 类模型的显存溢出问题高效完成完整体全局特征建模。