U‑Mamba增强生物医学图像分割的长距离依赖建模能力MIDL 2024Medical Imaging with Deep LearningAbstractConvolutional Neural Networks (CNNs) and Transformers have been the most popular architectures for biomedical image segmentation, but both of them have limited ability to handle long‑range dependencies because of inherent locality or computational complexity. To address this challenge, we introduce U‑Mamba, a general‑purpose network for biomedical image segmentation. Inspired by the State Space Sequence Models (SSMs), a new family of deep sequence models known for their strong capability in handling long sequences, we design a hybrid CNN‑SSM block that integrates the local feature extraction power of convolutional layers with the abilities of SSMs for capturing the long‑range dependency. Moreover, U‑Mamba enjoys a self‑configuring mechanism, allowing it to automatically adapt to various datasets without manual intervention. We conduct extensive experiments on four diverse tasks, including the 3D abdominal organ segmentation in CT and MR images, instrument segmentation in endoscopy images, and cell segmentation in microscopy images. The results reveal that U‑Mamba outperforms state‑of‑the‑art CNN‑based and Transformer‑based segmentation networks across all tasks. This opens new avenues for efficient long‑range dependency modeling in biomedical image analysis.翻译卷积神经网络CNN与 Transformer 是生物医学图像分割领域最主流的网络架构。但受限于 CNN 固有的局部感受野以及 Transformer 巨大的计算开销二者在长距离依赖建模上均存在短板。为解决该问题本文提出面向生物医学图像分割的通用网络 U‑Mamba。受擅长处理长序列的深度序列模型 —— 状态空间序列模型SSM启发本文设计CNN‑SSM 混合模块将卷积层强大的局部特征提取能力与 SSM 捕获长距离依赖的能力相结合。此外U‑Mamba 具备自配置机制无需人工调参即可自动适配各类数据集。本文在四项不同的分割任务上开展大量实验包括 CT 与 MR 影像的三维腹部器官分割、内窥镜图像器械分割、显微图像细胞分割。实验结果表明在全部测试任务上U‑Mamba 性能均优于当前 SOTA 的基于 CNN 与 Transformer 的分割网络。该工作为生物医学图像分析中高效长距离依赖建模开辟了新的研究方向。创新点1提出CNN‑SSM 混合 U‑Mamba 基础块将残差卷积与 Mamba 选择性 SSM 模块结合。卷积负责提取图像局部细粒度特征把空间特征展平为一维序列后送入 Mamba 模块捕获全局长距离依赖整体沿用 U‑Net 编解码器架构以线性复杂度完成长距离建模规避 Transformer 自注意力平方复杂度带来的显存与算力消耗同时保留卷积对医学图像边界、细小结构的局部建模能力。设计两个可选用变体U‑Mamba_Enc编码器全部采用混合块U‑Mamba_Bot仅瓶颈层使用 Mamba 混合块其余模块使用普通残差块方便在性能与计算开销之间做权衡。2继承 nnU‑Net 框架的自配置self‑configuring自动化机制。网络的 patch 尺寸、batch size、下采样阶段数、各轴池化数量等超参数由数据集统计特性自动推导无需人工手动调参同一套网络能够无缝适配 2D、3D支持 CT、MRI、内窥镜、显微图像多种模态的分割任务提升模型的通用性与落地便捷性。3较早把原生一维 MambaS6 选择性状态空间模型迁移到生物医学图像分割领域开展大规模多任务、多模态基准验证。在 4 个公开数据集上与 nnU‑Net、SegResNet、UNETR、Swin‑UNETR 开展完备对比实验。实验结果表明该模型对灰度分布异质性大的软组织鲁棒性更强能够减少漏分割、过分割与分割离群噪声在腹部器官、手术器械、细胞分割任务上整体优于 CNN、Transformer 类 SOTA证实 SSM/Mamba 范式在医学分割的可行性为后续大量 Mamba 变体分割工作提供基线参考。4解码器保留纯残差卷积设计分工明确。编码器使用 U‑Mamba 混合块兼顾局部特征与全局长距离上下文解码器完全使用普通残差块 转置卷积只聚焦局部细节恢复与分辨率重建不引入 Mamba 模块。编码器负责挖掘全局语义解码器专注还原分割细节既降低整体计算量又充分发挥卷积在分辨率恢复上的优势配合 U‑Net 原始跳跃连接完成多尺度特征融合。U-Mamba 编码器模块这张图是 U-Mamba 编码器模块U-Mamba (Enc)的整体结构也就是 U-Mamba 的基础构建块。根据论文描述U-Mamba 的基础模块由两个连续的残差卷积块再接一个基于 SSM 的 Mamba 模块串联而成以此兼顾局部特征提取能力与长距离依赖建模能力。输入数据维度为(B, C, H, W, D)代表 batch 数量、通道数、三维医学图像的高、宽、深度首先送入下方的卷积残差分支该分支重复 2 次内部由普通 3D 卷积、实例归一化 IN 与 Leaky ReLU 激活函数构成最后使用残差相加这部分负责提取医学图像的局部空间特征和传统 UNet 的卷积残差单元思路相近保证模型能捕捉边缘、纹理这类短距离局部信息。经过两层残差卷积处理后的三维特征图继续向上送入 Mamba 分支。第一步先执行 Flatten 操作将三维空间维度 H×W×D 展平得到序列长度 LH×W×D特征维度为 C 的序列特征张量形状转变为(B, L, C)随后经过 Layer Norm 做层归一化稳定序列建模的数值分布。归一化后的特征会分流成两条并行支路右侧支路经过线性层映射、SiLU 激活得到门控权重左侧支路先经过线性层、一维卷积、SiLU 激活再送入核心的 SSM 模块SSM 也就是 Mamba 的选择性状态空间模型用来对展平后的长序列建模高效捕获三维医学图像中跨空间位置的长距离依赖规避 Transformer 二次复杂度的开销。两条支路输出完成逐元素相乘融合门控信息与 SSM 提取的全局上下文特征再经过一层线性层做特征映射。最后通过 Reshape 操作把一维序列(B, L, C)还原回三维特征张量(B, C, H, W, D)完成整个 Mamba 块的计算。补充Mamba 块内部那条门控分支只是门控权重生成支路执行相乘操作属于 Gated‑MLP 门控机制用来调制 SSM 输出特征它不等于残差连接。这种“卷积残差块 Mamba 模块”的组合是 U-Mamba 编码器的核心设计。两层连续的残差卷积负责挖掘精细的局部图像特征而后面的 SSM-Mamba 块专门弥补 CNN 固有的长程建模短板传统 CNN 感受野有限Transformer 计算开销巨大U-Mamba 通过这种混合 CNN-SSM 的模块设计将卷积强大的局部提取能力和 SSM 线性复杂度的长序列建模能力结合在一起。整个模块使用残差结构贯穿缓解深度网络训练的梯度消失问题同时模块输出依然保持三维特征图格式可以直接嵌入 UNet 型编码器 - 解码器架构适配各类三维生物医学图像分割任务不需要大幅改动基础的 U 型网络范式。U-Mamba 整体 U 型编解码架构下图为 U-Mamba 整体 U 型编码器 - 解码器框架示意图论文原文描述为U-Mamba 采用编码器 - 解码器架构编码器中堆叠 U-Mamba Block解码器中使用残差块同时搭配跳跃连接Skip Connection该图属于概念示意图侧重表达整体网络逻辑并非严格的参数级绘制。整个网络保持 UNet 经典的对称 U 型结构左侧为编码器右侧为解码器。编码器部分由多层 U-Mamba Block 串联构成图中绿色方块即为 U-Mamba Block也就是刚才提到的 “2 个连续卷积残差块 Mamba 子模块” 的混合 CNN-SSM 基础单元。编码器依靠黑色箭头代表的步幅卷积Strided Convolution实现下采样逐层缩小特征图的空间尺寸不断扩大感受野每一层 U-Mamba Block 同时结合卷积提取局部医学图像细节以及 SSM 模块建模三维图像全局长距离依赖在不断下采样的过程中学习多尺度上下文特征。编码器每一层输出的特征会通过图中虚线代表的跳跃连接Skip Connection直接传递到解码器对应的层级。网络右侧是解码器分支每一级解码器的数据流顺序为紫色箭头的转置卷积Convolution Transposed先完成上采样恢复特征图分辨率上采样得到的特征送入绿色 U-Mamba Block借助 SSM 继续建模全局上下文融合编码器跳跃连接传来的浅层细节特征之后再送入橙色 Residual Block 残差块完成局部特征的精细化重建。这对应论文原文描述编码器堆叠 U-Mamba Block解码器引入残差块但解码器并不是完全移除 U-Mamba 模块而是在 U-Mamba Block 之后增加残差卷积单元做局部优化。上采样后的高层全局特征与编码器跳跃连接传递过来的浅层轮廓、边缘特征完成融合跳跃连接弥补上采样过程中丢失的细节这也是 UNet 系列分割网络的经典设计思路。论文的差异化架构设计思路可以理解为编码器依靠步幅卷积不断压缩特征用 U-Mamba Block 挖掘多尺度全局上下文解码器在进行分辨率恢复时依然保留 U-Mamba 块用于全局信息融合再搭配残差块做局部细节的微调重建平衡全局建模能力与局部细节还原能力。补充残差块的作用局部特征精细化加工。U‑Mamba Block 内部的 SSM 模块是把特征展平成一维序列运算擅长捕获全局长距离依赖但展平扫描会一定程度弱化三维空间局部纹理、器官边界、细小组织的局部建模。残差块由 3D 卷积、实例归一化 IN、Leaky ReLU搭配残差加法构成依靠卷积的局部感受野对经过全局建模之后的特征进一步打磨边缘、纹理、组织交界这些细粒度局部信息修正全局序列处理带来的局部特征偏差。稳定解码器深层训练缓解梯度消失。解码器网络层数很深上采样、特征融合之后特征分布容易震荡。残差块内部自带残差加法通路梯度可以绕过卷积层直接回传改善深层解码器的梯度流通让很深的 U‑Net 型解码器更容易收敛这也是残差结构最基础的能力。完成融合后特征的通道与空间特征变换。跳跃连接会把编码器低层特征和解码器高层特征拼接融合融合之后特征维度、信息混杂。残差卷积块承担特征变换任务对混合后的多尺度特征做进一步的提取与筛选把融合后的特征整理为适合下一步上采样的表示。架构上的分工互补全局交给 U‑Mamba Block局部交给残差卷积块。U‑Mamba Block 负责挖掘跨空间的全局上下文残差块承接其后专门负责局部重建。二者形成分工一个管全局语义一个管局部细节不用全部模块都做开销更大的 SSM 序列运算在保证分割精度的同时平衡计算开销。意义U-Mamba 这篇工作的核心意义是将新兴的选择性状态空间模型Mamba/SSM成功嵌入经典 UNet 编码器 - 解码器范式为三维医学图像分割开辟了一条CNN SSM 混合建模的新路线解决传统医学分割模型长期存在的固有矛盾。在 U-Mamba 出现之前3D 医学分割领域主要有两类主流方案以 3D UNet 为代表的纯 CNN 模型擅长提取局部纹理、边界特征但受限于卷积的局部感受野很难高效捕获三维体积内远距离的全局依赖而基于 Transformer 的医学分割模型依靠自注意力建模长程上下文但自注意力O(L^2)的时间复杂度在高分辨率三维医学体数据上会带来极高显存与算力开销很难落地到大尺寸三维影像上。对于医学图像分割的意义作者发现目前心脏分割相关比赛中Mamba架构模型表现较好1. 心脏是“大范围耦合”的器官如WHS 要分 7 个子结构——四腔室、心肌、大血管空间上分散但解剖关系强耦合心肌是薄壁结构几个体素厚边界消歧、跨切片的拓扑连续性都需要远处上下文。CNN 要堆很深才有全局感受野Transformer 有全局但代价是 O(N²)——3D 体积 token 数爆炸128³ 体积就是百万级体素、数万 token注意力矩阵根本算不动。Mamba 扫描是 O(N)让“3D 全局建模”第一次变得便宜。心脏分割非常依赖 anatomical context解剖上下文。Mamba 的长程依赖建模可以帮助网络获得这种上下文。心脏结构具有很强的全局空间关系比如你要分左心室 LV、左心室心肌 myocardium、右心室 RV、左心房 LA、心肌瘢痕、水肿区域这些结构不是独立存在的。一个像素究竟属于血池 / 心肌 / 背景 / 病灶很多时候不能只看这个像素附近几层 voxel。它需要结合局部纹理 整个心室形状 相邻切片 其他心脏结构之间的关系这就是 Mamba 第一个比较契合的地方。2. 选择性机制匹配医学数据的构成Mamba 的 selectivity 是数据依赖的门控按内容决定记什么、忘什么。心脏数据里背景占比高、目标沿扫描方向跨度大多方向扫描你跑的 SegMamba-V2 的 TSMamba 三向扫描、LKM-UNet 的 BiPixel/BiWindow 就是实例 选择性遗忘天然过滤无关体素、保留长程解剖关联。3. 真正的赢家模式是“混合架构”不是纯 Mamba比赛中的管线几乎都是 nnU-Net 式 conv 编解码 Mamba 瓶颈/逐 stage 混合——conv 保局部细节和薄边界Mamba 供全局先验。这正是 UMamba、LKM-UNet、EM-Net 的结构谱系纯 Mamba 通常反而不如混合。挑战赛按 Dice HD95 评分HD95 考边界精度conv 的局部归纳偏置在小数据上仍然关键。
