【魔改UNet系列】Mamba-UNet: 医学图像分割的UNet类纯视觉Mamba

小白学视觉

发布于 2024-10-29 12:59:43

9220

文章被收录于专栏：深度学习和计算机视觉深度学习和计算机视觉

论文信息

Mamba-UNet: UNet-Like Pure Visual Mamba for Medical Image Segmentation

Mamba-UNet: 医学图像分割的UNet类纯视觉Mamba

作者：Ziyang Wang, Jian-Qing Zheng, Yichi Zhang, Ge Cui, Lei Li

源码：https://github.com/ziyangwang007/MambaUNet

论文创新点

新颖的架构融合：Mamba-UNet结合了UNet的对称编码器-解码器风格架构和Mamba架构的能力，特别擅长处理长序列和全局上下文信息。这种融合在医学图像分割领域是创新的，旨在改进长距离依赖的建模。
纯视觉Mamba（VMamba）基础的编码器-解码器结构：该网络采用了基于纯视觉Mamba的编码器-解码器结构，并注入了跳跃连接以保留不同尺度上的空间信息。这种设计促进了全面的特征学习过程，能够在医学图像中捕捉复杂的细节和更广泛的语义上下文。
视觉状态空间（VSS）块的引入：Mamba-UNet在编码器和解码器中使用了视觉状态空间（VSS）块，这些块通过交叉扫描模块（CSM）将非因果视觉图像转换为有序的补丁序列，增强了模型在计算机视觉任务中的适用性。
针对医学图像分割优化的集成机制：作者在VMamba块内引入了一种新颖的集成机制，确保了编码器和解码器路径之间的无缝连接和信息流动，从而增强了分割性能。这一机制针对医学图像分割任务进行了特别优化。

摘要

在医学图像分析的最新进展中，卷积神经网络（CNN）和视觉变换器（ViT）树立了重要的基准。前者通过其卷积操作擅长捕捉局部特征，后者通过自注意力机制实现了显著的全局上下文理解。然而，这两种架构在有效建模医学图像中的长距离依赖方面都存在局限性，这对于精确分割至关重要。受Mamba架构的启发，该架构以其在处理长序列和全局上下文信息方面的专业性而闻名，并且作为状态空间模型（SSM），我们提出了Mamba-UNet，这是一种新颖的架构，它将UNet在医学图像分割中的能力与Mamba的能力相结合。Mamba-UNet采用了基于纯视觉Mamba（VMamba）的编码器-解码器结构，并注入了跳跃连接以保留网络不同尺度上的空间信息。这种设计促进了全面的特征学习过程，在医学图像中捕捉复杂的细节和更广泛的语义上下文。我们引入了一种新颖的集成机制，在VMamba块内确保编码器和解码器路径之间的无缝连接和信息流动，从而增强了分割性能。我们在公开可用的ACDC MRI心脏分割数据集和Synapse CT腹部分割数据集上进行了实验。结果表明，在相同的超参数设置下，Mamba-UNet在医学图像分割方面优于几种类型的UNet。

关键字

医学图像分割 · 卷积 · 变换器 · Mamba · 状态空间模型

2 方法

2.1 架构概述

所提出的Mamba-UNet的架构在图2中进行了概述，其灵感来自UNet和Swin-UNet。输入的2D灰度图像大小为H × W × 1，首先被分割成类似于ViT和VMamba的补丁，然后转换为1-D序列，尺寸为H/4 × W/4 × 16。一个初始的线性嵌入层将特征维度调整为任意大小，表示为C。这些补丁令牌随后通过多个VSS块和补丁合并层进行处理，创建层次化特征。每个编码器阶段的输出分辨率分别为H/4 × W/4 × C、H/8 × W/8 × 2C、H/16 × W/16 × 4C和H/32 × W/32 × 8C。解码器由VSS块和补丁扩展层组成，遵循编码器的风格，使得输出的特征大小完全相同，从而通过跳跃连接增强了在下采样过程中丢失的空间细节。在编码器和解码器中，每个阶段都使用了2个VSS块，并且在编码器中加载了预训练的VMamba-Tiny，遵循与Swin-UNet加载预训练的SwinViT-Tiny相同的过程。VSS块、编码器的补丁合并和解码器的补丁扩展的细节在以下小节中讨论。

2.2 VSS块

VSS网络块在图3中进行了说明，主要基于视觉Mamba。具体来说，传统的SSM作为线性时不变系统，通过隐藏状态h(t) ∈ RN，将x(t) ∈ R映射到y(t) ∈ R，给定A ∈ CN×N作为演化参数，B、C ∈ CN作为状态大小N的投影参数，以及跳跃连接D ∈ C1。该模型可以被公式化为线性常微分方程（ODEs），如方程1所示：

这个线性模型的离散版本可以通过零阶保持在一个时间尺度参数∆ ∈ RD下进行转换。

其中A = e^∆A, B = (e^∆A − I)A−1B, C = C, B, C ∈ RD×N。使用一阶泰勒级数对B进行精细化的近似B = ∆A(∆A)−1 ∆B = ∆B。视觉Mamba进一步引入了交叉扫描模块（CSM），然后将卷积操作集成到块中，详细内容见[7,16]。在VSS块中，输入特征首先通过一个线性嵌入层，然后分成两个路径。一个分支经过深度卷积和SiLU激活，然后进入SS2D模块，并在层归一化后与另一个经过SiLU激活的分支合并。这个VSS块不像典型的视觉变换器那样使用位置嵌入，而是选择一个没有MLP阶段的流线型结构，使得在相同的深度预算内可以堆叠更密集的块。

2.3 编码器

在编码器中，C维的标记化输入在降低分辨率的情况下经历了两个连续的VSS块进行特征学习，保持维度和分辨率。补丁合并作为下采样过程在Mamba-UNet的编码器中使用了三次，通过将输入分割成1/4的象限，将它们连接起来，然后每次通过layernorm归一化维度，将令牌数量减少1/2，并将特征维度加倍2倍。