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

标题：《Large Window-based Mamba UNet for Medical Image Segmentation: Beyond Convolution and Self-attention》
论文：https://arxiv.org/pdf/2403.07332.pdf
代码：https://github.com/wjh892521292/LMa-UNet

导读

目前，基于 CNN 和 Transformers 等 backbone 的自动分割模型减少了人工标注需求，大多基于小核进行局部特征提取，大核或全局的建模依赖则需要大量的计算成本。例如，基于CNN的模型通常通过分层堆叠小核来提取全局模式，在像素级特征提取方面表现出色，但由于其有限的接受域，在学习远程依赖关系方面效果不佳。

虽然最近的一些研究已经证明了大卷积核的有效性，但它通常需要特定的优化策略和复杂的推理时间模型重构。相比之下，基于 Transformer 的算法提供了强大的远程建模，但牺牲了像素级的空间建模。

此外，自注意力机制这一关键组件的复杂性较高，因此像 swin transformer 这样的方法通过将像素打包到 window 中，从而降低计算成本，但牺牲了更多空间信息。

特别是，许多研究表明，transformer在7 × 7的窗口尺寸下实现了最佳权衡，因为较小的窗口会导致更多的计算，而较大的窗口会导致性能显著下降。后来的研究研究了CNN-Transformer混合模型或局域内建模方法。

然而，由于医学图像通常比自然图像尺寸更大，如何降低全局补丁之间相互作用的复杂性以及如何扩大局部空间建模的接受场仍然是一个开放的问题。

最近，以 Mamba 为代表的状态空间模型（SSMs）已然成为一种有前景的方法，在长序列建模上具有线性复杂度。得益于其较低的复杂度，mamba在空间分配上具有更大的灵活性。也就是说，与之前基于小kernel或小window的自注意力的方法不同，Mamba有希望赋予模型大感受野空间建模的能力，这在已知的基于Mamba的方法中被忽视了。

因此，本文基于大kernel或者说大window提出了基于分层和双向的大核 Mamba 模块，以增强 SSM 的表示建模能力。通过大window的SSM同时建模全局和局部关系，并结合UNet进行高效的医学图像分割。

方法

网络架构

上图为LMa-UNet的模型结构图。 包括一个depth-wise conv层、编码器、解码器、投影层以及经典的跳跃连接。其中编码器由所提出的Large Mamba Block(LM Block)组成。 我们重点介绍所提出的LM Block。

LM Block

LM block的核心由三个组件组成，分别是Pixel-level SSM (PiM)层，Patch-level SSM (PaM)层以及bidirectional Mamba (BiM)。如图1所示。 通过划分给SSM大window以充分进行局部和全局建模。

Pixel-level SSM (PiM)

由于Mamba是一个连续模型，输入像素的离散性会削弱局部邻接像素的相关性建模。因此，我们提出了一个像素级的Pixel-level SSM，将特征映射分割成多个大window（子窗口），并在子窗口上执行SSM操作。

首先，我们将一个完整的特征图划分为不重叠的2D window或3D cube。这里以2D 为例给定一个分辨率为$H\times W$的输入，我们将特征图划分为每个大小为$m\times n$的子窗口($m$和$n$最高可达到40！)。在不失一般性的前提下，我们假设$H/m$和$W/n$都是整数。然后我们有$\frac{hw}{mn}$子窗口。

在该方案下，当这些大窗口被送入Mamba层时，局部相邻像素将连续输入到SSM中，因此局部邻域像素之间的关系可以更好地建模。此外，在大窗口分割策略下，感受野被扩大，模型可以获得更多的局部像素细节。

图2(a)展示了LMa-UNet与CNN、Transformer的感受野对比。CNNs的kernel感受野通常在$3\times 3$, Transformer系列的window感受野通常在$7\times 7$，而本文提出的LMa-UNet的感受野则可扩展到$40\times 40$.

Patch-level SSM (PaM)

由于图像被分成不重叠的子窗口。因此，我们需要一种在不同子窗口之间进行通信的机制，以进行远程依赖建模。我们引入了一个补丁级patch-level SSM层来在不同的子窗口之间传递信息。分辨率为$H\times W$的特征图首先通过大小为$m\times n$的池化层，以便将每个$\frac{hw}{mn}$子窗口的重要信息汇总为单个代表。

因此，我们获得了具有$\frac{hw}{mn}$个子窗口代表的聚合序列，然后使用它通过Mamba在子窗口之间进行通信，以进行全局范围的依赖关系建模。在Mamba中进行子窗口信息交互之后，我们将聚合图反池化到原始特征图大小。

Bidirectional Mamba (BiM)

与基于仅向前扫描方向SSM层的原始Mamba块不同，LM块中的每个SSM层(包括PiM和PaM)都是双向的。如图2
2(b)所示。原始的Mamba中作为一个连续模型，较早输入的元素会出现一定遗忘问题，而最近输入到Mamba的元素会保留更多的信息。

因此，原始的单一扫描方向的曼巴将更多地关注后面的patch。而在医学中，图像的中心区域，往往有更多的器官和病变。为此，我们提出了一种双向曼巴结构，通过同时进行前向和后向扫描并叠加输出结果。

详细结构如图1左侧所示，其有两个优点：

1. 该模型可以更多地关注图像中心区域可能有更多器官和病变的信息斑块，而不是角落区域。
2. 对于每个patch，该网络可以很好地建模出其绝对位置信息和与其他patch的相对位置信息。

实验

实验在两个器官分割数据集(Abdomen CT和Abdomen MR)上进行，包括3D任务和2D任务。 从论文提供的实验结果可以看出这种大kernel或者大window的Mamba模型在医学图像分割任务上表现出了一定的竞争力。

消融实验部分在Abdomen MR上进行，进一步证明了大window的Mamba表现出强劲的性能提升。

总结

本文首次探索了大kernel(window) Mamba块在医学图像分割上的潜力，为了充分利用大kernel Mamba的能力，我们设计了一种基于分层和双向的大window的Mamba块，增强SSM的表示建模能力。通过分层设计即PiM和PaM同时建模全局和局部关系，并结合UNet进行高效的医学图像分割。在多个器官数据集上的实验表明，大kernel(window) Mamba在医学图像分割任务中具有竞争力，并值得未来深入探索。

最后，如果你对本文有任何的观点或疑问，欢迎评论区下方留言讨论。同时也欢迎对多模态等前沿相关技术感兴趣的同学扫描屏幕下方二维码添加微信好友，备注“交流学习”即可。