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

科技   2024-10-30 10:34   中国香港  

点击下方深度学习爱好者”,选择加"星标"或“置顶


论文信息

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

论文创新点

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

2.4 解码器

与编码器相镜像,解码器使用两个连续的VSS块进行特征重建,使用补丁扩展层而不是合并层进行上采样深度特征[3]。这些层增强了分辨率(2倍上采样),同时将特征维度减半1/2,例如,一个初始层在重新组织和降低它们以增强分辨率之前将特征维度加倍。

2.5 瓶颈和跳跃连接

Mamba-UNet的瓶颈使用了两个VSS块。编码器和解码器的每个级别都采用了跳跃连接,将多尺度特征与上采样输出混合,通过合并浅层和深层来增强空间细节。随后的线性层保持了这个集成特征集的维度,确保与上采样分辨率的一致性。

3 实验和结果

声明

本文内容为论文学习收获分享,受限于知识能力,本文对原文的理解可能存在偏差,最终内容以原论文为准。本文信息旨在传播和学术交流,其内容由作者负责,不代表本号观点。文中作品文字、图片等如涉及内容、版权和其他问题,请及时与我们联系,我们将在第一时间回复并处理。
下载1:Pytoch常用函数手册
在「深度学习爱好者」公众号后台回复:Pytorch常用函数手册即可下载全网第一份Pytorch常用函数手册,涵盖Tensors介绍、基础函数介绍、数据处理函数、优化函数、CUDA编程、多线程处理等十四章章内容。

下载2:Python视觉实战项目52讲
小白学视觉公众号后台回复:Python视觉实战项目即可下载包括图像分割、口罩检测、车道线检测、车辆计数、添加眼线、车牌识别、字符识别、情绪检测、文本内容提取、面部识别等31个视觉实战项目,助力快速学校计算机视觉。

交流群


欢迎加入公众号读者群一起和同行交流,目前有SLAM、三维视觉驶、计算摄影、检测、分割、识别、医学影像、GAN算法竞赛等微信群(以后会逐渐细分),请扫描下面微信号加群,备注:”昵称+学校/公司+研究方向“,例如:”张三 + 上海交大 + 视觉SLAM“。请按照格式备注,否则不予通过。添加成功后会根据研究方向邀请进入相关微信群。请勿在群内发送广告,否则会请出群,谢谢理解~

深度学习爱好者
分享机器学习、深度学习和Python等知识与原理,每天分享深度学习与计算机视觉领域的经典和最新的论文总结,带读者一起跟踪前言科技!
 最新文章