本文基于大kernel或者说大window提出了基于分层和双向的大核 Mamba 模块，以增强 SSM 的表示建模能力。通过大window的SSM同时建模全局和局部关系，并结合UNet进行高效的医学图像分割
。
>>
加入极市CV技术交流群，走在计算机视觉的最前沿
论文
：
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操作。
首先
, 我们将一个完整的特征图划分为不重叠的
window或
cube。这里以
为例给定一个分辨率为
的输入, 我们将特征图划分为每个大小为
的子窗口(
和
最高可达到40
! )。在不失一般性的前提下, 我们假设
和
都是整数。然后我们有
子窗口。
在该方案下，当这些大窗口被送入Mamba层时，局部相邻像素将连续输入到SSM中，因此局部邻域像素之间的关系可以更好地建模。此外，在大窗口分割策略下，感受野被扩大，模型可以获得更多的局部像素细节。
图2(a)展示了LMa-UNet与CNN、Transformer的感受野对比。CNNs的kernel感受野通常在
, Transformer系列的 window感受野通常在
, 而本文提出的LMa-UNet的感受野则可扩展到
.
Patch-level SSM (PaM)
由于图像被分成不重叠的子窗口。因此, 我们需要一种在不同子窗口之间进行通信的机制, 以进行远程依赖建模。我们引入了一个补丁级patch-level SSM层来在不同的子窗口之间传递信息。分辨率为
的特征图首先通过大小为
的池化层, 以便将每个
子窗口的重要信息汇总为单个代表。
因此, 我们获得了具有
