VM-UNET模型简介

原创已于 2025-11-29 15:04:01 修改 · 1.1k 阅读

8 ·

CC 4.0 BY-SA版权

文章标签：

#语义分割 #深度学习 #医学影像 #学习记录

于 2025-11-26 10:21:29 首次发布

部署运行你感兴趣的模型镜像

简介

VM-UNet（Vision Mamba UNet）是首个完全基于状态空间模型（State Space Model, SSM）的医学图像分割架构，标志着 Mamba 系列模型在医学影像领域的开创性应用。由上海交通大学研究团队于 2024 年提出。

VM-UNet 构建了一种全新的 U 型编码器-解码器结构，其核心组件——VSS Block（Visual State Space Block）完全摒弃了传统的卷积（CNN）或自注意力（Transformer）操作，转而采用源自 VMamba 的纯 SSM 模块进行特征提取与重建。这一设计不仅继承了 Mamba 在建模长距离依赖方面的强大能力，还保持了对输入尺寸的线性计算复杂度，有效克服了 CNN 局部感受野有限和 Transformer 计算开销过大的固有缺陷。

作为首个纯 SSM 驱动的医学图像分割模型，VM-UNet 不仅为该领域提供了一个高效、简洁且性能强劲的新基线，也为未来探索更轻量、更全局感知的医学影像分析方法开辟了全新路径。

模型讲解

先看一下模型的论文摘要：

摘要

在医学图像分割领域，基于卷积神经网络（CNN）和基于 Transformer 的模型已得到广泛研究。然而，CNN 在长距离建模能力方面存在局限性，而 Transformer 则受限于其二次计算复杂度。近年来，以曼巴（Mamba）为代表的状态空间模型（SSMs）已成为一种极具潜力的方法。它们不仅在建模长距离交互方面表现出色，还保持了线性计算复杂度。本文利用状态空间模型，提出了一种用于医学图像分割的 U 型架构模型，名为视觉曼巴 U 型网络（VM-UNet）。具体而言，引入视觉状态空间（VSS）块作为基础模块以捕捉广泛的上下文信息，并构建了具有较少卷积层的非对称编码器 - 解码器结构，以节省计算成本。我们在 ISIC17、ISIC18 和 Synapse 数据集上进行了全面实验，结果表明 VM-UNet 在医学图像分割任务中表现出竞争力。据我们所知，这是首个基于纯状态空间模型构建的医学图像分割模型。我们旨在建立一个基准，并为未来开发更高效、更有效的基于状态空间模型的分割系统提供有价值的见解。我们的代码可在https://github.com/JCruan519/VM-UNet获取。

关键词：医学图像分割、状态空间模型、上下文信息

注：关于状态空间模型（State Space Model, SSM）以及Mamba的讲解可以看这位博主的文章：

https://kobebryant.blog.youkuaiyun.com/article/details/143676875?fromshare=blogdetail&sharetype=blogdetail&sharerId=143676875&sharerefer=PC&sharesource=qq_73038863&sharefrom=from_link

下面是博主博客里的内容：

状态空间模型（State Space Model, SSM）是一种源自控制理论和信号处理的序列建模方法，用于描述系统如何根据当前输入和内部状态逐步演化并产生输出。在深度学习中，现代 SSM（如 Mamba）通过将连续时间动态系统离散化，并引入数据依赖的选择性机制，使其能够高效地捕捉长距离依赖关系，同时保持线性计算复杂度（即计算量随序列长度线性增长，而非像 Transformer 那样呈平方增长）。

Mamba 是一种基于状态空间模型（State Space Model, SSM）的高效序列建模架构，于 2023 年提出，旨在克服传统深度学习模型在处理长序列时的计算瓶颈。SSM 最初源自控制理论，通过一个连续或离散的动态系统描述输入如何影响内部状态并产生输出，其经典形式使用固定的参数矩阵来传递信息，虽然具备线性计算复杂度的优势，但缺乏对输入内容的适应能力。Mamba 在此基础上进行了关键改进：它引入了数据依赖的选择性机制，使模型能够根据当前输入动态调整状态转移过程中的参数（如输入投影和输出观测矩阵），从而有选择地保留重要信息、过滤无关内容。这种“智能”的状态演化方式让 Mamba 在保持 O(N) 线性复杂度的同时，获得了媲美甚至超越 Transformer 的建模能力，尤其适合处理长距离依赖任务。

注：关于Mamba的讲解也可以看我的文章：

https://blog.youkuaiyun.com/qq_73038863/article/details/155375958?fromshare=blogdetail&sharetype=blogdetail&sharerId=155375958&sharerefer=PC&sharesource=qq_73038863&sharefrom=from_link

在 VM-UNet 中，Mamba 的视觉变体（VMamba）被用作核心构建模块，其 VSS Block（Visual State Space Block）完全替代了传统 U-Net 中的卷积或注意力操作，构成了首个纯 SSM 驱动的医学图像分割架构，既实现了全局上下文感知，又显著提升了计算效率。

模型整体架构

VM-UNET的整体架构如下图所示：

VM-UNet 包括补丁嵌入（Patch Embedding）层、编码器、解码器、最终投影（Final Projection）层和跳跃连接。与先前的方法不同，没有采用对称结构，而是使用了非对称设计。

补丁嵌入层将输入图像划分为 4×4 大小的非重叠补丁，随后将图像的维度映射到 C（C 默认值为 96）。此过程产生嵌入图像。最后，在将输入编码器进行特征提取之前，使用层归一化（Layer Normalization）对其进行归一化。

编码器由四个阶段组成，前三个阶段的末尾应用补丁合并操作，以减少输入特征的高度和宽度，同时增加通道数。在四个阶段中采用 [2, 2, 2, 2] 个 VSS 块，每个阶段的通道数分别为 [C, 2C, 4C, 8C]。

解码器也分为四个阶段。后三个阶段的开始采用补丁扩展操作，以减少特征通道数并增加高度和宽度。在四个阶段中，使用 [2, 2, 2, 1] 个 VSS 块，每个阶段的通道数分别为 [8C, 4C, 2C, C]。

解码器之后，采用最终投影层将特征大小恢复为与分割目标匹配。具体而言，通过补丁扩展进行 4 倍上采样以恢复特征的高度和宽度，然后通过投影层恢复通道数。

对于跳跃连接，采用了简单直接的加法操作，没有额外的复杂设计，因此不会引入任何额外参数。

VSS block

源自 VMamaba 的 VSS 块是 VM-UNet 的核心模块，如下图所示：

输入经过层归一化后，被分为两个分支：

（1）在第一个分支中，输入经过线性层和激活函数。

（2）在第二个分支中，输入经过线性层、深度可分离卷积和激活函数处理后，送入二维选择性扫描（2D-Selective-Scan，SS2D）模块进行进一步特征提取。随后，对特征进行层归一化，然后与第一个分支的输出进行逐元素乘法（element-wise production），以融合两个路径的特征。最后，使用线性层混合特征，并将结果与残差连接相结合，形成 VSS 块的输出。默认采用 SiLU 作为激活函数。

其中，SS2D 由三个组件组成：扫描扩展（scan expanding）操作、S6 块和扫描合并（scan merging）操作。

如图 3（a）所示，扫描扩展操作将输入图像沿四个不同方向（左上到右下、右下到左上、右上到左下、左下到右上）展开为序列。然后，这些序列由 S6 块进行特征提取，确保来自各个方向的信息都被充分扫描，从而捕捉多样化的特征。

如图 3（b）所示，扫描合并操作将四个方向的序列求和并合并，将输出图像恢复到与输入相同的大小。

S6 块源自 Mamba ，通过基于输入调整 SSM 的参数，在 S4 的基础上引入了选择性机制。这使模型能够区分并保留相关信息，同时过滤掉无关信息。

S4（Structured State Space Model）是一种将经典控制理论中的连续状态空间模型引入深度学习的序列建模方法。它通过一个线性动态系统描述输入如何随时间演化为输出，其核心由固定的参数矩阵 A、B 和 C 控制，具有严格的数学结构和高效的并行化能力。S4 的关键优势在于能够以线性计算复杂度（O(N)）建模极长距离依赖，并支持全局感受野，适用于语音、时间序列等任务。然而，由于其参数是静态的、与输入无关的，S4 缺乏对不同内容的适应能力，在处理语义复杂的数据（如自然语言或医学图像）时表现受限。

S6 是 Mamba 模型中对 S4 的重要改进，全称为“Selective State Space Model”。它保留了 S4 的高效线性结构，但引入了输入依赖的选择性机制：S6 动态地根据当前输入生成状态转移参数 Bt、Ct和时间尺度 Δt，使模型能够有选择地关注或忽略某些信息。这种“内容感知”能力让 S6 在保持 O(N) 复杂度的同时，显著提升了表达能力和泛化性能，尤其适合需要精细上下文理解的任务。在 VM-UNet 中，S6 作为 SS2D 模块的核心组件，被用于高效提取医学图像中的全局语义特征，成为纯 SSM 架构实现高性能分割的关键技术基础。

S6 块的伪代码如下所示：