Large Window-based Mamba UNet for Medical Image Segmentation: Beyond Convolution and Self-attention-CSDN博客

本文链接：https://blog.csdn.net/qq_46981910/article/details/138785267

摘要：

在临床实践中，医学图像分割提供了有关目标器官或组织的轮廓和尺寸的有用信息，有助于改进诊断、分析和治疗。在过去的几年中，卷积神经网络 (cnn) 和 transformer 在这一领域占据了主导地位，但它们仍然受到有限的感受野或昂贵的长程建模的影响。状态空间序列模型(State Space Sequence Model ，简称 Mamba) 是一种具有线性复杂度的长程依赖建模方法。本文提出一种基于大窗口的 Mamba 型网络 (LMa-UNet)，用于 2D 和和和 3D 医学图像分割。 LMa-UNet 的一个显著特点是利用了大窗口，与基于小核的cnn和基于小窗口的transformer相比，在局部空间建模方面表现出色，同时与具有二次复杂度的自注意力相比，在全局建模方面保持了优越的效率。此外，设计了一种新颖的 层次化和双向的Mamba块，进一步增强了Mamba的全局和邻域空间建模能力 。综合

实验证明了所提出方法的有效性和效率，以及使用大窗口尺寸来实现大感受野的可行性。

代码：https://github.com/wjh892521292/LMa-UNet

一、Introduce

目前研究：

自动分割模型利用流行的深度学习骨干，如卷积神经网络(cnn)和transformer，减少了人工注释要求，但可能会招致相当大的计算成本或忽视特定的细节。
基于cnn的模型(例如UNet)通常寻求通过分层堆叠小内核来提取全局模式，在像素级特征提取方面表现出色，但由于其有限的感受野，在学习长程依赖关系方面无效。
尽管最近有研究表明了大型卷积核的有效性，但它往往需要特定的优化策略和复杂的推理时间模型重构。
基于transformer的算法提供了强大的长程建模，但牺牲了像素级空间建模。此外，一个关键组件，自注意力模块，产生了二次复杂度，不能处理过多的token，导致需要将像素打包到窗口中，从而牺牲分辨率信息。许多研究表明，transformer在7 × 7窗口大小上实现了最佳权衡，因为较小的窗口会导致更多的计算，而较大的窗口会导致性能的显著下降。
后来的研究调查了CNN-Transformer混合模型或块内局部性建模方法。然而，由于医学图像通常比自然图像尺寸更大，如何降低全局图像块之间的交互复杂度以及如何扩大局部空间建模的感受野仍是有待解决的问题。

SSMs：

结构化状态空间序列模型(SSMs) (例如Mamba)已经成为一种强大的长序列建模方法，
在输入大小方面具有线性复杂度，揭示了局部和全局依赖的有效建模。
与传统的自注意力方法相比，ssm使一维数组(例如文本序列)中的每个元素都能通过压缩的隐藏状态与任何先前扫描的样本进行交互，有效地将二次复杂度降低为线性。
设计 ssm 是为了解决自然语言处理 (NLP) 任务，但也显示了计算机视觉的有效性。对于医学图像分割，例如， U-Mamba和segamba 引入了 SSM-CNN 混合模型，直接应用 Mamba 在像素级别有效地建模长程依赖性。

Mamba潜力和不足仍未充分挖掘和解决：

得益于其线性复杂度，Mamba在空间分配方面具有更大的灵活性。也就是说，与之前基于小卷积核或大小约束的自注意力窗口的方法不同，Mamba在赋予模型大感受野空间建模能力方面很有希望，这在已知的基于Mamba的方法中被忽略了。
Mamba 是一种单向的序列建模方法， 缺乏位置感知 ，更关注后验词例。
原始的 Mamba 是为一维序列建模而提出的， 不适合需要空间感知理解的计算机视觉任务 。主要由于局部相邻空间像素的不连续， Mamba 存在遗忘问题和低效的局部建模能力。

大窗口的曼巴型网络(LMa-UNet)：

提出了一种基于大窗口的Mamba UNet (Large window-based Mamba UNet, LMa-UNet)模型，用于2D/3D医学图像分割。
利用Mamba强大的序列建模能力和线性复杂度，通过为SSM模块分配大窗口来实现大的感受野。进一步，设计了一种新的基于层次和双向大窗口的Mamba块(LM块)，以增强ssm的表示建模能力。
采用双向 Mamba 设计进行位置感知序列建模，降低输入顺序的权重影响。
层次式的 Mamba 设计由两类SSM 操作组成 :(i) 像素级 SSM (PiM) 和 (ii) patch 级 SSM (PaM) 。 PiM 通过大窗口SSM 捕获邻域和局部像素信息，避免 SSM 中因标记化后相邻块不连续而出现

的遗忘问题。 PaM 处理长程依赖建模和全局块交互。

主要贡献：

提出了一种基于大窗口的Mamba UNet (Large window-based Mamba UNet, LMa-UNet)模型，用于2D/3D医学图像分割。
将大感受野窗口分配给SSM层，使模型具有大空间建模的能力。
设计了一种基于位置感知的双向Mamba序列建模方法。
提出了一种新的由像素级SSM (PiM)和块级SSM (PaM)组成的Mamba模块，增强了局部邻域像素级特征建模和长程全局块级建模。

二、基础知识:SSM模型

基于 ssm 的模型，即结构化状态空间序列模型 (structured state space sequence models, S4)和 Mamba ，都是从连续系统演化而来的，该连续系统将一个一维函数或序列 x ( t ) → y ( t ) ∈ R映射到一个隐藏状态

。这个过程可以表示为下面是线性常微分方程 :

其中：

是状态矩阵
是投影参数

S4 和 Mamba 是上述连续系统的离散版本，其中包括时间尺度参数 ∆ 将连续参数 A , B 转换为离散参数

。通常采用零阶保持器 (zero-order hold,ZOH)作为离散化规则，其定义如下 :

公式（2）

对A, B进行离散化后，公式一(1)的离散化版本(即SSM模型)可以写成:

然后通过全局卷积计算输出，定义如下:

其中

L是输入序列的长度x,
是结构化卷积核。

三、方法

3.1 LMa-UNet

概述LMa-UNet的解如图1所示：

具体来说，除了深度卷积、具有下采样层的编码器、具有上采样层的解码器和跳跃连接的常见 UNet 组成外， LMa-UNet 通过在编码器中插入提出的基于大窗口的曼巴 (LM) 块来改进 UNet 的结构。给定一个分辨率为 C × D × H × W 的 3D 输入图像，深度卷积首先将输入编码为特征图

。然后将特征图 F 0 输入到每个 LM 块和对应的下采样层中，得到多尺度特征图 ; 一个 LM 块包含两个 Mamba 模块 : 像素级 SSM (PiM) 和块级SSM (PaM)。对于 $l^{th}$ 层，该过程可以表述为 :

其中：

PiM和PaM分别表示像素级SSM和块级SSM。
Down-sampling表示下采样层。

每个阶段之后，生成的特征图 $F_{l+1}$ 被编码为

;其中

表示特征图 $F_{l}$ 的通道和分辨率。在解码器部分，采用UNet解码器和带有跳跃连接的残差块进行上采样并预测最终的分割掩码。

3.2 LM块

LM 块是核心组件，用于在每个阶段对不同尺度的特征图进行进一步的空间建模。与以往利用 CNN 进行局部像素级建模、 Transformer 进行长程块级依赖建模不同，利用 Mamba 的 线性复杂度 ， LM 块可以 同时完成像素级和块级 建模。更重要的是，较低的复杂度允许设置更大的窗口来获得 更大的感受野 ，这将提高局部建模的效率，如图 2(a) 所示。具体来说， LM 块是一个分层设计，由像素级 SSM (PiM) 和块级 SSM (PaM) 组成 ; 前者用于局部邻域像素建模，后者用于全局长程依赖建模。此外， LM 块中每个 Mamba 层都是双向的，用于位置感知序列建模。

图2 (a) CNN、Transformer 和我们提出的 LMa-UNet 之间的视野比较。CNN 通常使用小内核（如 3 × 3），Transformers 通常使用 7 × 7 大小的窗口。(b) vanilla Mamba 与我们提出的双向 Mamba 的扫描顺序比较。

1、像素级SSM (PiM) ：

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

2、补丁级SSM (PaM)：

引入补丁级 SSM 层在不同子窗口之间传递信息。如图 1 的补丁级 SSM 所示，分辨率为 H × W 的特征图 $F_{l}^{'}$ 首先通过大小为 m × n 的池化层，以允许每个的重要信息 HW/ mn将子窗口汇总为单个代表。因此，获得聚合映射 $Z_{l}$ 与 HW/ mn 代表，然后用于通过 Mamba 进行全局范围依赖建模的子窗口之间的通信。在Mamba中的交互之后，将聚合映射解池回初始特征映射 $F_{l}^{'}$ 的相同大小，并应用残差接，公式 (5)中

的过程可以这样进行 :

其中：

Pooling和Unpooling分别表示池化层和非池化层。
Bi-Mamba表示提出的双向Mamba层。

3、Bidirectional Mamba (BiM)：

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

四、实验

4.1数据集

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

腹部CT 。腹部 CT 是一个公开可用的 3D 多器官分割数据集，包含来自 MICCAI 2022 FLARE Challenge [16]的 100 个 CT 病例，包括 13 种腹部器官 ( 肝、脾、胰腺、右肾、左肾、胃、胆囊、食管、主动脉、下腔静脉、右肾上腺、左肾上腺和十二指肠 ) 。三维 CT 图像的大小为 40 × 224 × 192 。来自 MSD 胰腺数据集的50 例与来自腹腔的注释一起用于训练，另外来自不同医疗中心的 50 例用于测试。

腹部 MR。Abdomen MR 是一个公开可用的 2D 分割数据集，包含来自 MICCAI 2022 AMOS Challenge [11] 的 110 个 MRI 病例，包括 13 种腹部器官（与腹部 CT 数据集相同）。2D MRI 图像的大小为 320 × 320。60 个带注释的案例用于训练，另外 50 个案例用于测试。

4.2实现设置

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

4.3综合性能

基线模型包括三种类型的代表性医学图像分割网络 : 基于 cnn 的网络 (nnU-Net和 SegResNet) ，基于 transformer 的网络 (UNETR, SwinUNETR 和 nnFormer)，以及最新的基于 mamba 的网络 (U-Mamba [15]) 。为了公平比较，我们还在nnU-Net 框架中实现了所有模型，并使用默认的图像预处理。表 1 显示了结果。与基于 cnn 和基于 transformer 的分割方法相比，提出的 LMa-UNet 在以下方面取得了更好的性能 DSC 和 NSD 均表明， Mamba 的全局建模能力对医学图像分割至关重要。请注意，与 U-Mamba 相比， U-Mamba 只是将 Mamba 作为全局建模适配器， LMa-UNet 展示了对 U-Mamba 的改进，验证了双向和分层 Mamba 设计的有效性。这些结果还表明了 Mamba 在基于大窗口的感受野的全局和局部特征建模方面的潜力。

4.4LMa-UNet的窗口大小

由于 Mamba 的线性复杂度， LMa-UNet 可以实现 cnn 和 transformer 无法实现的大窗口建模。为了验证 Mamba 的巨大空间建模潜力，探索了 LMa-UNet 在不同窗口大小设置下的性能。表2显示了腹部MR数据集的结果，该数据集的分期为7。比较三种窗口大小设置的性能，可以发现窗口大小较大的LMa-UNet具有更好的性能。这表明大感受野对于医学图像分割是至关重要的，而由于Mamba的线性复杂性，它可以实现医学图像分割。这将为未来的大型内核或窗口建模提供更多可能性。

4.5消融实验

为了验证 LMa-UNet 中每个关键成分的效果，在腹部 MR 数据集上进行了消融实验。表 3 显示了结果。实验结果表明， PiM 和 PaM 对 LMa-UNet 的建模效果都优于基线模型，分别验证了 PiM 和 PaM 在局部像素级建模和全局建模方面的有效性。值得注意的是， PiM 模型比 PaM 模型获得了更多的改进，这表明扩大局部特征建模的感受野是提高模型性能的关键。引入 BiM 后， LMa-UNet 的性能进一步提升，说明了双向 Mamba 对位置感知序列建模的重要性。最后，所有组件的LMa-UNet 均取得了最好的性能，进一步证明了算法的有效性该方法及其组成部分的有效性。

五、结论

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