当前位置：首页 > article >正文

[Mamba_4]LMa-UNet

article 2024/11/13 11:03:45

题目：LKM-UNet: Large Kernel Vision Mamba UNet for Medical Image Segmentation

中文题目：LMa-UNet: 探索大kernel Mamba在医学图像分割上的潜力

0摘要

在临床实践中，医学图像分割提供了有关目标器官或组织的轮廓和尺寸的有用信息，有助于改进诊断、分析和治疗。在过去的几年里，卷积神经网络(cnn)和Transformers占据了这一领域的主导地位，但它们仍然受到有限的接受域和昂贵的远程建模的困扰。Mamba，一种状态空间序列模型(SSM)，最近作为具有线性复杂性的远程依赖建模的一个有前途的范例出现。在本文中，我们引入了一种用于医学图像分割的大核视觉曼巴u形网络(LKM-UNet)。我们的LKM-UNet的一个显着特征是它利用了大型曼巴核，与基于小核的cnn和transformer相比，它在局部空间建模方面表现出色，同时与具有二次复杂度的自关注相比，它在全局建模方面保持了更高的效率。此外，我们设计了一个新的分层和双向曼巴块，以进一步提高曼巴的视觉输入的全局和邻域空间建模能力。综合实验证明了采用大尺寸曼巴核实现大接收野的可行性和有效性。代码可在https://github.com/wjh892521292/LKM-UNet上获得。

1介绍

在大尺寸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]和segamba[22]引入了SSM-CNN混合模型，直接应用Mamba在像素级高效建模远程依赖关系。虽然有效，但曼巴的潜力和不足仍未得到充分挖掘和解决，主要表现在三个方面。首先，得益于其线性复杂性，曼巴在空间分配上具有更大的灵活性。也就是说，与之前基于小卷积核或尺寸约束的自我注意窗口的方法的刻板印象不同，Mamba有希望赋予模型大感受野空间建模的能力，这在已知的基于Mamba的方法中被忽视了。其次，Mamba是一种单向序列建模方法，缺乏位置意识，更侧重于后验标记。第三，原始的Mamba是针对一维序列建模提出的，不适合需要空间感知理解的计算机视觉任务。主要由于局部相邻空间像素变得不连续，Mamba可能会出现遗忘问题和低效的局部建模能力。

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

我们的主要贡献如下(1) 我们提出了一种用于二维/三维医学图像分割的大核 Mamba UNet（LKM-UNet）。(2) 我们将大感受野的核分配给 SSM 层，使模型具备大空间建模能力。(3) 我们设计了用于位置感知序列建模的双向 Mamba。(4) 我们提出了一种由像素级 SSM（PiM）和斑块级 SSM（PaM）组成的新型分层 Mamba 模块，增强了局部邻域像素级和远距离全局斑块级建模能力。

2初步研究:SSM模型

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

其中 $A\in \mathbb{R}^{N*N}$ 为状态矩阵，B, C∈ $R^{N}$ 为投影参数。S4和Mamba是上述连续系统的离散版本，其中包含一个时间尺度参数∆，用于将连续参数A、B转换为离散参数 $\bar{A},\bar{B}$ 。通常采用零阶保持器(zero-order保持器，ZOH)作为离散化规则，其定义如下:

离散化后，Eq.(1)的离散化版本可以定义为:

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

式中，L为输入序列x的长度， $\bar{K}\in \mathbb{R}^{L}$ 为结构化卷积核。

3方法

在本节中，我们首先介绍 LKM-UNet 的整体架构。随后，我们将详细介绍 LM 模块这一核心组件。

3.1 LKM-UNet

图 1 是 LKM-UNet 的概览。具体来说，除了常见的深度卷积、带下采样层的编码器、带上采样层的解码器和跳跃连接的 UNet 组成外，LKM-UNet 改进了 UNet 的结构，在编码器中插入了提出的大核 Mamba（LM）块。给定分辨率为 C × D × H × W 的三维输入图像，深度卷积首先将输入图像编码为特征图 $F_{0}\in \mathbb{R}^{48*\frac{D}{2}*\frac{D}{2}*\frac{D}{2}}$ 。然后将特征图 $F_{0}$ 输入每个 LM 块和相应的下采样层，就得到了多尺度特征图。一个 LM 块包含两个 Mamba 模块：像素级 SSM（PiM）和斑块级 SSM（PaM）。对于第 l 层，处理过程可表述为:

其中，PiM 和 PaM 分别表示像素级 SSM 和斑块级 SSM。 Down-sampling 表示下采样层。每个阶段结束后，生成的特征图 $F_{l+1}$ 被编码为（ $2C_{l},\frac{D_{l}}{2},\frac{H_{l}}{2},\frac{W_{l}}{2},$ ），其中 $C_{l}$ ,( $D_{l},H_{l},W_{l},$ ) 表示特征图 $F_{l}$ 的信道和分辨率。至于解码器部分，我们采用 UNet 解码器和带跳接的残差块来进行上采样和预测最终的分割掩码。

3.2 LM block

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

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

补丁级 SSM (PaM)。我们引入了一个斑块级 SSM 层来传递不同子核之间的信息。如图 1 中的片段级 SSM 所示，分辨率为 H × W 的特征图 $F_{l}^{'}$ 首先要经过一个大小为 m × n 的汇集层，以便将 $\frac{HW}{mn}$ 个子内核中每个子内核的重要信息汇总为一个代表。这样，我们就得到了具有 $\frac{HW}{mn}$ 代表的集合地图 $Z_{l}$ ，然后通过 Mamba 在子内核之间进行通信，以建立全局范围的依赖关系模型。在 Mamba 中进行交互后，我们将聚合图解压缩到与初始特征图 $F_{l}^{'}$ 相同的大小，并应用残差连接。公式（5）中的 $F_{l}^{''}=PaM\left ( F_{l}^{'} \right )$ 过程可按以下方式进行：

其中Pooling和Unpooling分别表示池化层和非池化层。Bi-Mamba表示提议的双向曼巴层。

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

图二所示：(a) CNN、Transformer 和我们提出的 LKM-UNet 的各自领域比较。CNN 通常使用较小的内核（如 3 × 3），而 Transformer 通常使用 7 × 7 大小的内核（窗口）。我们的 LKM-UNet 可以将内核大小扩展到 40 × 40。(b) Vanilla Mamba 与我们提出的双向 Mamba 的扫描阶数比较。

4实验

4.1数据集

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

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

腹部MR。腹部MR是一个公开的二维分割数据集，包含来自MICCAI 2022 AMOS挑战赛[11]的110个MRI病例，包括13种腹部器官(与腹部CT数据集相同)。二维MRI图像大小为320 × 320。根据之前的工作[15,22]，使用60个带注释的用例进行训练，另外使用50个用例进行测试。

4.2实现设置

我们的LKM-UNet是在PyTorch 1.9.0上基于nnU-Net框架实现的。所有实验均在NVIDIA GeForce RTX 3090 GPU上进行。训练中的批大小对于3D数据集(腹部CT)为2，对于2D数据集(腹部MR)为24。使用动量= 0.99的Adam[12]优化器。初始学习率为0.01，权值衰减为3e-5。最大训练历元数为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整体性能

基线模型包括三种代表性网络:基于cnn的网络(nnU-Net[10]和SegResNet[17])，基于transformer的网络(UNETR [8]， SwinUNETR [7]， nnFormer[24])，以及最新的基于mamba的网络(U-Mamba[15])。为了公平比较，我们还在nnU-Net框架中实现了所有模型，并使用默认的图像预处理。表1给出了结果。与基于 CNN 和基于变换器的分割方法相比，我们提出的 LKM-UNet 在 DSC 和 NSD 方面的性能都有所提高，这表明 Mamba 的全局建模能力对于医学图像分割至关重要。需要注意的是，与仅仅将 Mamba 用作全局建模适配器的 U-Mamba 相比，LKM-UNet 的性能比 U-Mamba 有所提高，从而验证了我们的双向和分层 Mamba 设计的有效性。这些结果还证明了 Mamba 在更大感受野的全局和局部特征建模方面的潜力。

4.4 LKM-UNet的核大小重要吗?

为了验证Mamba的大空间建模潜力，我们探索了LKM-UNet在不同内核大小设置下的性能。表2显示了在腹部MR数据集上的结果。比较三种内核大小设置的性能，可以看到内核大小较大的LKM-UNet可以获得更好的性能。这表明，由于曼巴的线性复杂性，大的接受野对于医学图像分割是至关重要的。

表2。LKM-UNet在三种不同内核大小设置下的性能。内核大小序列表示每个阶段的内核大小(阶段总数为7)。

4.5消融研究

我们对腹部MR数据集进行了消融实验，以验证LKM-UNet中每个关键成分的效果，如表3所示。PiM和PaM都对LKM-UNet基线模型进行了改进，分别验证了PiM和PaM在局部像素级建模和全局建模方面的优越性。值得注意的是，使用PiM的模型得到了更多的改进，这表明扩大局部建模的接受域是提高模型性能的关键。引入BiM后，LKM-UNet的性能进一步提升，可见双向曼巴对于位置感知建模的重要性。最后，所有组件的LKM-UNet都达到了最佳性能，进一步证明了我们的方法及其组件的有效性。

表3。LKM-UNet不同子模块的性能分析。PiM =像素级SSM。PaM =补丁级SSM。BiM =双向曼巴。

4.6有效的感受野可视化

为了展示感受野的更多细节，我们在图 3 中展示了其他方法和 LKM-UNet 的有效感受野（ERF）[14]。基于 CNN 的方法更侧重于局部特征提取，而基于 Transformer 的方法的 ERF 范围更广。虽然 U-Mamba 利用 Mamba 获得了全局 ERF，但却削弱了对局部的关注。相比之下，我们提出的带有大核 Mamba 的 LKM-UNet 在全局和局部方面都获得了更大的 ERF。