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

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

导读

现在，根据 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 为例给定一个分辨率为$HW$的输入，咱们将特征图划分为每个巨细为$mn$的子窗口($m$和$n$最高可达到40！)。在不失一般性的前提下，咱们假定$H/m$和$W/n$都是整数。然后咱们有hwmnfrac{hw}{mn}mnhw​子窗口。

在该计划下，当这些大窗口被送入Mamba层时，部分相邻像素将接连输入到SSM中，因而部分邻域像素之间的联系能够更好地建模。此外，在大窗口切割战略下，感触野被扩大，模型能够获得更多的部分像素细节。

图2(a)展示了LMa-UNet与CNN、Transformer的感触野比照。CNNs的kernel感触野一般在$33$, Transformer系列的window感触野一般在$77$，而本文提出的LMa-UNet的感触野则可扩展到$4040$.

Patch-level SSM (PaM)

因为图画被分成不堆叠的子窗口。因而，咱们需求一种在不同子窗口之间进行通讯的机制，以进行长途依靠建模。咱们引入了一个补丁级patch-level SSM层来在不同的子窗口之间传递信息。分辨率为$HW$的特征图首先经过巨细为$mn$的池化层，以便将每个hwmnfrac{hw}{mn}mnhw​子窗口的重要信息汇总为单个代表。

因而，咱们获得了具有hwmnfrac{hw}{mn}mnhw​个子窗口代表的聚合序列，然后使用它经过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在医学图画切割使命中具有竞争力，并值得未来深化探究。

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