《Large Window-based Mamba UNet for Medical Image Segmentation:Beyond Convolution and Self-attention》

【Tiltle】

用于医学图像分割的基于大窗口的 Mamba UNet：超越卷积和自注意力

【Abstract】

        在临床实践中，医学图像分割提供了有关目标器官或组织的轮廓和尺寸的有用信息，有助于改进诊断、分析和治疗。 在过去的几年里，卷积神经网络（CNN）和 Transformer 占据了这一领域的主导地位，但它们仍然受到有限的感受野或昂贵的远程建模的困扰。 Mamba 是一种状态空间序列模型 (SSM)，最近成为具有线性复杂性的远程依赖建模的有前途的范例。 在本文中，我们介绍了一种基于大窗口的 Mamba U 形网络（LMa-UNet），用于 2D 和 3D 医学图像分割。

        我们的 LMa-UNet 的一个显着特点是它利用大窗口，与基于小型内核的 CNN 和基于小型窗口的 Transformer 相比，在局部空间建模方面表现出色，同时与具有二次复杂度的自注意力机制相比，在全局建模方面保持了卓越的效率。 此外，我们设计了一种新颖的分层和双向 Mamba 块，以进一步增强 Mamba 的全局和邻域空间建模能力。 综合实验证明了我们方法的有效性和效率以及使用大窗口尺寸实现大感受野的可行性。 代码可在 https://github.com/wjh892521292/LMa-UNet 获取。

【Introduction】

        在大尺寸 2D/3D 图像中有效分割感兴趣的生物医学对象（例如病变）可显着增强下游临床实践和生物医学研究。 目前，自动分割模型利用流行的深度学习主干，如卷积神经网络（CNN）和 Transformers [19]，减少了手动注释的要求，但可能会产生相当大的计算成本或忽略特定细节。

        基于 CNN 的模型（例如 UNet [18]）通常寻求通过分层堆叠小内核来提取全局模式，在像素级特征提取方面表现出色，但由于其有限的感受野，在学习远程依赖性方面效果不佳 [14]。 尽管最近的研究[3]已经证明了大卷积核的有效性，但它通常需要特定的优化策略和复杂的推理时间模型重构。

        相比之下，基于 Transformer 的算法提供了强大的远程建模，但牺牲了像素级空间建模 [21,20]。 此外，关键组件自注意力模块会产生二次复杂度，并且无法处理过多的标记[9]，导致需要将像素打包到窗口中，从而牺牲分辨率信息。 特别是，许多研究表明 Transformer 在 7 × 7 窗口大小下实现了最佳权衡，因为较小的窗口会导致更多的计算，而较大的窗口会导致性能显着下降 [13,1]。 后来的研究研究了用于块内局部性建模的 CNN-Transformer 混合模型或方法 [20,23]。 然而，由于与自然图像相比，医学图像的尺寸通常较大，如何降低全局斑块之间的交互复杂性以及如何扩大局部空间建模的感受野仍然是开放的。

        最近，结构化状态空间序列模型（SSM）[6]（例如，Mamba [5]）已经成为一种强大的长序列建模方法，在输入大小方面具有线性复杂性，为局部和全局的高效建模提供了线索。 依赖关系。 与传统的自注意力方法相比，SSM 使一维数组（例如文本序列）中的每个元素能够通过压缩隐藏状态与任何先前扫描的样本进行交互，从而有效地将二次复杂度降低为线性复杂度。 SSM 旨在解决自然语言处理 (NLP) 任务 [4]，但也显示出在计算机视觉方面的有效性 [25]。

        例如，对于医学图像分割，U-Mamba [5] 和 SegMamba [22] 引入了 SSM-CNN 混合模型，直接应用 Mamba 在像素级有效地建模远程依赖关系。 虽然有效，但Mamba的潜力和不足仍未得到充分挖掘和解决，主要表现在三个方面。

        首先，受益于其线性复杂度，Mamba 在空间分配上拥有更大的灵活性。 也就是说，与之前基于小卷积核或大小受限的自注意力窗口的刻板印象不同，Mamba 有望赋予模型大感受野空间建模的能力，而这一能力在已知的基于 Mamba 的方法中被忽视了。 其次，Mamba 是一种单向序列建模方法，缺乏位置意识，更注重后验标记。 第三，最初的 Mamba 是为一维序列建模而提出的，它不适合需要空间感知理解的计算机视觉任务。 主要是由于局部相邻空间像素变得不连续，Mamba 可能会遇到遗忘问题和低效的局部建模能力。

        在本文中，我们提出了一种基于大窗口的 Mamba U 形网络（LMa-UNet），用于 2D 和 3D 医学图像分割。 LMa-UNet利用Mamba强大的序列建模能力和线性复杂度，通过为SSM模块分配大窗口来实现大感受野。 此外，我们设计了一种新颖的分层和双向基于大窗口的 Mamba 块（LM 块）来增强 SSM 的表示建模能力。 双向 Mamba 设计用于位置感知序列建模，减少输入顺序的权重影响。 分层 Mamba 设计由两种类型的 SSM 操作组成：(i) 像素级 SSM (PiM) 和 (ii) 补丁级 SSM (PaM)。 PiM 通过大窗口 SSM 捕获邻域和局部范围的像素信息，避免了 SSM 中由于标记化后相邻块不连续而出现的遗忘问题。 PaM 处理远程依赖建模和全局补丁交互。 我们的主要贡献如下。

(1) 我们提出了一种基于大窗口的 Mamba UNet (LMa-UNet)，用于 2D/3D 医学图像分割。 (2) 将大感受野的窗口分配给SSM层，使模型具备大空间建模的能力。

(3) 我们设计了一个用于位置感知序列建模的双向 Mamba。

(4) 我们提出了一种由像素级SSM（PiM）和补丁级SSM（PaM）组成的新型分层Mamba模块，增强了局部邻域像素级特征建模和远程全局补丁级建模。

【Preliminaries: SSM Models】

基于SSM的模型，即结构化状态空间序列模型（S4）和Mamba [5]，都是从通过隐层映射一维函数或序列x(t) → y(t) ∈ R的连续系统演化而来的。 状态 h(t) ∈ RN。 这个过程可以表示为以下线性常微分方程：

其中 A ∈ RN×N 是状态矩阵，B,C ∈ RN 是投影参数。 S4和Mamba是上述连续系统的离散版本，其中包括时间尺度参数Δ，用于将连续参数A、B转换为离散参数A、B。 通常，采用零阶保持（ZOH）作为离散化规则，可以定义如下：

 对 A、B 进行离散化后，得到方程的离散化版本。 (1)（即SSM模型）可以写为：

 

图 1.我们提出的 LMa-UNet 的概述。

然后通过全局卷积计算输出，定义为：

其中 L 是输入序列 x 的长度，K ∈ RL 是结构化卷积核。 

 【Method】

3.1 LMa-UNet

图 1 给出了 LMa-UNet 的概述。具体来说，除了常见的 UNet 组成：深度卷积、具有下采样层的编码器、具有上采样层的解码器和跳跃连接之外，LMa-UNet 还改进了 UNet 的结构，其中提出了插入编码器的基于大窗口的 Mamba (LM) 块。 给定分辨率为 C×D×H×W 的 3D 输入图像，深度卷积首先将输入编码为特征图 F0 ∈ R48×D 2 ×H 2 ×W 2 。 然后将特征图F0输入到每个LM块和相应的下采样层中，获得多尺度特征图。 LM 块包含两个 Mamba 模块：像素级 SSM (PiM) 和补丁级 SSM (PaM)。 对于第l层，该过程可以表述为：

其中 PiM 和 PaM 分别表示像素级 SSM 和块级 SSM。 下采样表示下采样层。 在每个阶段之后，产生的特征图Fl+1被编码为(2Cl，Dl 2 ，Hl 2 ，Wl 2 )，其中Cl、(Dl，Hl，Wl)表示特征图Fl的通道和分辨率。 至于解码器部分，我们采用 UNet 解码器和带有跳跃连接的残差块来进行上采样和预测最终的分割掩模。

3.2 LM block

        LM block是我们的核心组件，用于对每个阶段不同尺度的特征图进行进一步的空间建模。 与之前使用 CNN 进行局部像素级建模、使用 Transformer 进行远程补丁级依赖建模的方法不同，LM 块可以同时完成像素级和补丁级建模，这得益于 Mamba 的线性复杂度。 更关键的是，较低的复杂度允许设置更大的窗口以获得更大的感受野，这将提高局部建模的效率，如图2（a）所示。 具体来说，LM块是一种分层设计，由像素级SSM（PiM）和补丁级SSM（PaM）组成； 前者用于局部邻域像素建模，后者用于全局远程依赖建模。 此外，LM 块中的每个 Mamba 层都是双向的，这是为位置感知序列建模而提出的。 

像素级 SSM (PiM)。 由于 Mamba 是连续模型，输入像素的离散性质会削弱局部相邻像素的相关建模。 因此，我们提出了一种像素级的SSM，将特征图分割成多个大的子窗口，并对子窗口执行SSM操作。 我们首先将整个特征图平均划分为不重叠的 2D 子窗口或 3D 子立方体。 以 2D 为例。 给定 H ×W 分辨率的输入，我们将特征图划分为每个大小为 m × n 的子窗口（m 和 n 最多可达 40）。 不失一般性，我们假设H/m和W/n都是整数。 然后我们有HW mn个子窗口，如图1的Pixel-level SSM所示。在该方案下，当这些子窗口被发送到Mamba层时，局部相邻像素将被连续输入到SSM中； 因此，可以更好地对局部邻域像素之间的关系进行建模。 此外，在大窗口划分策略下，感受野被放大，模型可以获得更多的局部像素细节。 然而，图像被分成不重叠的子窗口。 因此，我们需要一种在不同子窗口之间进行通信的机制，以进行远程依赖建模。

补丁级 SSM (PaM)。 我们引入了补丁级 SSM 层来在不同子窗口之间传递信息。 如图 1 的 Patch-level SSM 所示，分辨率为 H ×W 的特征图 F′ l 首先经过大小为 m × n 的池化层，以总结每个 HW mn 子窗口的重要信息 变成一个代表。 因此，我们获得了具有 HW mn 代表的聚合映射 Zl，然后使用它们通过 Mamba 在子窗口之间进行通信，以进行全局范围依赖建模。 在 Mamba 中交互后，我们将聚合映射恢复到相同的大小作为初始特征图F′l，并应用残差连接。 式中 F′′ l = PaM(F′ l) 的过程 (5) 可以执行为：

其中Pooling和Unpooling分别表示池化层和非池化层。 Bi-Mamba 表示提议的双向 Mamba 层。 双向曼巴 (BiM)。 与基于仅前向扫描方向 SSM 层的普通 Mamba 块不同，我们的 LM 块中的每个 SSM 层（包括 PiM 和 PaM）都是双向的。 图2（b）显示了差异。 在原始的 Mamba 中，作为一个连续模型，较早进入的元素会发生一些信息遗忘，而最新进入 Mamba 的元素会保留更多的信息。 因此，具有单一扫描方向的原始曼巴将更多地关注后部斑块，而不是通常具有更多器官和病变的特征图的中心区域。 为此，我们提出了一种双向 Mamba 结构，通过同时执行前向和后向扫描并将输出结果叠加。 详细结构如图1左侧所示。BiM有两个优点。 首先，该模型可以更多地关注图像中心区域的信息斑块，这些斑块可能具有更多的器官和病变，而不是角落区域。 其次，对于每个补丁，网络可以很好地建模其绝对位置信息和与其他补丁的相对位置信息。

图 2. (a) CNN、Transformer 和我们提出的 LMa-UNet 各自的领域比较。 CNN 通常使用小内核（如 3 × 3），而 Transformer 通常使用 7 × 7 大小的窗口。 (b) 普通 Mamba 与我们提出的双向 Mamba 的扫描顺序比较。

【Experiments】

4.1 Datasets 

我们在 2D 和 3D 分割任务的两个数据集上与最先进的方法进行实验比较，以验证 LMa-UNet 的有效性和可扩展性。

Table 1. Quantitative segmentation results on the 3D Abdomen CT dataset and 2D
Abdomem MR dataset. 

腹部CT。 腹部CT是一个公开的3D多器官分割数据集，包含来自MICCAI 2022 FLARE挑战赛的100个CT病例[16]，包括13种腹部器官（肝、脾、胰腺、右肾、左肾、胃、胆囊、食管） , 主动脉, 下腔静脉, 右肾上腺、左肾上腺和十二指肠）。 3D CT 图像的大小为 40 × 224 × 192。使用来自 MSD 胰腺数据集的 50 个带有 AbdomenCT-1K 注释的病例进行训练，另外使用来自不同医疗中心 [2] 的 50 个病例进行测试。 腹部先生。 腹部 MR 是一个公开的 2D 分割数据集，包含来自 MICCAI 2022 AMOS 挑战赛的 110 个 MRI 病例 [11]，包括 13 种腹部器官（与腹部 CT 数据集相同）。 2D MRI 图像的大小为 320 × 320。继之前的工作[15,22]之后，使用 60 个带注释的病例进行训练，另外 50 个病例用于测试。

4.2 Implementation Setup

我们的LMa-UNet是基于nnU-Net框架在PyTorch 1.9.0上实现的。 所有实验均在 NVIDIA GeForce RTX 3090 GPU 上进行。 3D 数据集（腹部 CT）的训练批量大小为 2，2D 数据集（腹部 MR）的训练批量大小为 24。 使用动量 = 0.99 的 Adam [12] 优化器。 初始学习率为 0.01，权重衰减为 3e-5。 最大训练epoch数为1000。对于腹部CT数据集，阶段为6，但维度不一致； 因此我们将三个维度的矩形窗口大小设置为 [20, 28, 24], [20, 28, 24], [10, 14, 12], [10, 14, 12], [5, 7, 6] ，以及每个阶段的[5,7,6]。 对于腹部 MR 数据集，阶段为 7，每个阶段的窗口大小为 40、20、20、10、10、5 和 5。

4.3 Overall Performances

基线模型包括三种类型的代表性医学图像分割网络：基于 CNN 的网络（nnU-Net [10] 和 SegResNet [17]）、基于 Transformer 的网络（UNETR [8]、SwinUNETR [7] 和 nnFormer [24] ]），以及最新的基于 Mamba 的网络（U-Mamba [15]）。 为了公平比较，我们还实现了nnU-Net框架中的所有模型，并使用默认的图像预处理。 表 1 列出了结果。 与基于 CNN 和基于 Transformer 的分割方法相比，我们提出的 LMa-UNet 实现了DSC 和 NSD 的性能都有所提高，这表明 Mamba 的全局建模能力对于医学图像分割至关重要。 请注意，与简单地将 Mamba 作为全局建模适配器的 U-Mamba 相比，LMa-UNet 展示了 U-Mamba 的改进，验证了我们双向和分层 Mamba 设计的有效性。 这些结果还证明了 Mamba 在具有基于大窗口的感受野的全局和局部特征建模中的潜力

表 2. LMa-UNet 在三种不同窗口大小设置下的性能。 窗口大小顺序表示每个阶段的窗口大小。 

 表 3.具有不同子模块的 LMa-UNet 的性能。 PiM = 像素级 SSM。 PaM = 补丁级 SSM。 BiM = 双向曼巴。

 4.4 Is the Window Size of LMa-UNet Important?

由于MAMBA的线性复杂性，LMA-UNET可以实现CNN和Transformers无法实现的大窗口建模。为了验证Mamba的巨大空间建模潜力，我们考察了LMA-UNET在不同窗口大小设置下的性能。表2显示了在级数为7的腹部MR数据集上的结果。比较三种窗口大小设置的性能，可以看到窗口大小越大的LMA-UNET获得了更好的性能。这表明大的感受野对于医学图像分割是至关重要的，由于Mamba的线性复杂性，可以使用Mamba来实现医学图像分割。这将为未来的大型内核或窗口建模提供更多可能性。

4.5 Ablation Study

为了验证 LMa-UNet 中每个关键组件的效果，我们在腹部 MR 数据集上进行了消融实验。 结果如表3所示。 可以看到，PiM 和 PaM 都为 LMa-UNet 提供了相对于基线模型的改进，分别验证了 PiM 和 PaM 在局部像素级建模和全局建模中的有效性。 值得注意的是，使用 PiM 的模型比使用 PaM 的模型获得了更多的改进，这表明扩大局部特征建模的感受野是提高模型性能的关键。 引入BiM后，LMa-UNet的性能进一步提升，由此可见双向Mamba对于位置感知序列建模的重要性。最后，具有所有组件的 LMa-UNet 达到了最佳性能，这进一步证明了我们的方法及其组件的有效性。

【Conclusions】

在本文中，我们引入了一种新的基于 Mamba 的 UNet 医学图像分割模型，实现了大窗口空间建模。 此外，我们设计了双向分层 SSM，以增强 Mamba 局部和全局特征建模的能力。 对多器官分割数据集的综合实验证明了我们提出的方法的有效性。 未来，我们将探索SSM在其他医学成像任务中的扩展和应用，如检测、配准、重建等。

【References】

1. Chu, X., Tian, Z., Wang, Y., Zhang, B., Ren, H., Wei, X., Xia, H., Shen, C.:
Twins: Revisiting the design of spatial attention in vision Transformers. Advances
in Neural Information Processing Systems 34, 9355–9366 (2021)
2. Clark, K., Vendt, B., Smith, K., Freymann, J., Kirby, J., Koppel, P., Moore, S.,
Phillips, S., Maffitt, D., Pringle, M., et al.: The cancer imaging archive (TCIA):
Maintaining and operating a public information repository. Journal of Digital Imag-
ing 26, 1045–1057 (2013)
3. Ding, X., Zhang, X., Ma, N., Han, J., Ding, G., Sun, J.: RepVGG: Making VGG-
style convnets great again. In: Proceedings of the IEEE/CVF Conference on Com-
puter Vision and Pattern Recognition. pp. 13733–13742 (2021)
4. Fu, D.Y., Dao, T., Saab, K.K., Thomas, A.W., Rudra, A., Ré, C.: Hungry hun-
gry hippos: Towards language modeling with state space models. arXiv preprint
arXiv:2212.14052 (2022)
5. Gu, A., Dao, T.: Mamba: Linear-time sequence modeling with selective state
spaces. arXiv preprint arXiv:2312.00752 (2023)
6. Gu, A., Goel, K., Ré, C.: Efficiently modeling long sequences with structured state
spaces. arXiv preprint arXiv:2111.00396 (2021)
7. Hatamizadeh, A., Nath, V., Tang, Y., Yang, D., Roth, H.R., Xu, D.: Swin UNETR:
Swin Transformers for semantic segmentation of brain tumors in MRI images. In:
International MICCAI Brainlesion Workshop. pp. 272–284. Springer (2021)
8. Hatamizadeh, A., Tang, Y., Nath, V., Yang, D., Myronenko, A., Landman, B.,
Roth, H.R., Xu, D.: UNETR: Transformers for 3D medical image segmentation.
In: Proceedings of the IEEE/CVF Winter Conference on Applications of Computer
Vision. pp. 574–584 (2022)
9. Hua, W., Dai, Z., Liu, H., Le, Q.: Transformer quality in linear time. In: Interna-
tional Conference on Machine Learning. pp. 9099–9117. PMLR (2022)
10. Isensee, F., Jaeger, P.F., Kohl, S.A., Petersen, J., Maier-Hein, K.H.: nnU-Net: A
self-configuring method for deep learning-based biomedical image segmentation.
Nature Methods 18(2), 203–211 (2021)
11. Ji, Y., Bai, H., Ge, C., Yang, J., Zhu, Y., Zhang, R., Li, Z., Zhanng, L., Ma,
W., Wan, X., et al.: AMOS: A large-scale abdominal multi-organ benchmark for
versatile medical image segmentation. Advances in Neural Information Processing
Systems 35, 36722–36732 (2022)
12. Kingma, D.P., Ba, J.: Adam: A method for stochastic optimization. arXiv preprint
arXiv:1412.6980 (2014)
13. Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., Guo, B.: Swin Trans-
former: Hierarchical vision Transformer using shifted windows. In: Proceedings of
the IEEE/CVF International Conference on Computer Vision. pp. 10012–10022
(2021)
14. Luo, W., Li, Y., Urtasun, R., Zemel, R.: Understanding the effective receptive field
in deep convolutional neural networks. Advances in Neural Information Processing
Systems 29 (2016)
15. Ma, J., Li, F., Wang, B.: U-Mamba: Enhancing long-range dependency for biomed-
ical image segmentation. arXiv preprint arXiv:2401.04722 (2024)
16. Ma, J., Zhang, Y., Gu, S., Ge, C., Ma, S., Young, A., Zhu, C., Meng, K.,
Yang, X., Huang, Z., et al.: Unleashing the strengths of unlabeled data in pan-
cancer abdominal organ quantification: The FLARE22 challenge. arXiv preprint
arXiv:2308.05862 (2023)

17. Myronenko, A.: 3D MRI brain tumor segmentation using autoencoder regular-
ization. In: Brainlesion: Glioma, Multiple Sclerosis, Stroke and Traumatic Brain
Injuries: 4th International Workshop, BrainLes 2018, Held in Conjunction with
MICCAI 2018, Revised Selected Papers, Part II 4. pp. 311–320. Springer (2019)
18. Ronneberger, O., Fischer, P., Brox, T.: U-Net: Convolutional networks for biomed-
ical image segmentation. In: Medical Image Computing and Computer-Assisted
Intervention. pp. 234–241. Springer (2015)
19. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser,
Ł., Polosukhin, I.: Attention is all you need. Advances in Neural Information Pro-
cessing Systems 30 (2017)
20. Wu, H., Xiao, B., Codella, N., Liu, M., Dai, X., Yuan, L., Zhang, L.: CvT: In-
troducing convolutions to vision Transformers. In: Proceedings of the IEEE/CVF
International Conference on Computer Vision. pp. 22–31 (2021)
21. Wu, Y., Liao, K., Chen, J., Wang, J., Chen, D.Z., Gao, H., Wu, J.: D-Former: A U-
shaped dilated transformer for 3D medical image segmentation. Neural Computing
and Applications 35(2), 1931–1944 (2023)
22. Xing, Z., Ye, T., Yang, Y., Liu, G., Zhu, L.: SegMamba: Long-range sequential mod-
eling mamba for 3D medical image segmentation. arXiv preprint arXiv:2401.13560
(2024)
23. Yuan, K., Guo, S., Liu, Z., Zhou, A., Yu, F., Wu, W.: Incorporating convolution
designs into visual Transformers. In: Proceedings of the IEEE/CVF International
Conference on Computer Vision. pp. 579–588 (2021)
24. Zhou, H.Y., Guo, J., Zhang, Y., Han, X., Yu, L., Wang, L., Yu, Y.: nnFormer:
Volumetric medical image segmentation via a 3D Transformer. IEEE Transactions
on Image Processing (2023)
25. Zhu, L., Liao, B., Zhang, Q., Wang, X., Liu, W., Wang, X.: Vision Mamba: Efficient
visual representation learning with bidirectional state space model. arXiv preprint
arXiv:2401.09417 (2024)