Large Window-based Mamba UNet for Medical Image Segmentation: Beyond Convolution and Self-attention

摘要：

在 临 床 实 践 中 ， 医 学 图 像 分割 提 供 了 有 关 目 标 器 官 或 组织的轮 廓 和 尺 寸 的 有 用 信 息 ， 有 助 于 改 进 诊 断 、 分 析 和 治 疗 。 在 过 去 的 几年中 ， 卷 积神 经 网 络 (cnn) 和 transformer 在 这 一 领 域 占 据 了主 导 地 位 ， 但它们仍 然 受 到 有 限 的 感受野或昂贵的长程建模 的 影 响 。 状 态 空 间 序 列 模型(State Space Sequence Model ， 简 称 Mamba) 是 一 种 具 有 线性复杂度的长程依赖 建 模 方 法 。 本 文 提 出 一 种 基 于 大 窗 口 的 Mamba 型 网 络 (LMa-UNet)， 用 于 2D 和和和 3D 医 学 图 像 分割 。 LMa-UNet 的 一个 显 著 特 点 是 利用了大窗口，与基于小核的cnn和基于小窗口的transformer相比，在局部空间建模方面表现出色，同时与具有二次复杂度的自注意力相比，在全局建模方面保持了优越的效率 。 此 外 ， 设计 了一 种 新 颖 的 层次化和双向的Mamba块，进一步增强了Mamba的全局和邻域空间建模能力 。 综 合

实 验 证 明 了 所 提 出 方 法 的 有 效 性 和 效 率 ， 以 及 使 用 大 窗 口 尺 寸 来 实 现 大感受 野 的 可 行 性 。

代码：https://github.com/wjh892521292/LMa-UNet

一、Introduce

目前研究：

1. 自动分割模型利用流行的深度学习骨干，如卷积神经网络(cnn)和transformer，减少了人工注释要求，但可能会招致相当大的计算成本或忽视特定的细节。
2. 基于cnn的模型(例如UNet)通常寻求通过分层堆叠小内核来提取全局模式，在像素级特征提取方面表现出色，但由于其有限的感受野，在学习长程依赖关系方面无效。
3. 尽管最近有研究表明了大型卷积核的有效性，但它往往需要特定的优化策略和复杂的推理时间模型重构。
4. 基于transformer的算法提供了强大的长程建模，但牺牲了像素级空间建模。此外，一个关键组件，自注意力模块，产生了二次复杂度，不能处理过多的token，导致需要将像素打包到窗口中，从而牺牲分辨率信息。许多研究表明，transformer在7 × 7窗口大小上实现了最佳权衡，因为较小的窗口会导致更多的计算，而较大的窗口会导致性能的显著下降。
5. 后来的研究调查了CNN-Transformer混合模型或块内局部性建模方法。然而，由于医学图像通常比自然图像尺寸更大，如何降低全局图像块之间的交互复杂度以及如何扩大局部空间建模的感受野仍是有待解决的问题。

SSMs：

• 结构化状态空间序列模型(SSMs) (例如Mamba)已经成为一种强大的长序列建模方法，
• 在输入大小方面具有线性复杂度，揭示了局部和全局依赖的有效建模。
• 与传统的自注意力方法相比，ssm使一维数组(例如文本序列)中的每个元素都能通过压缩的隐藏状态与任何先前扫描的样本进行交互，有效地将二次复杂度降低为线性。
• 设计 ssm 是 为了 解 决 自 然 语 言 处 理 (NLP) 任 务 ，但也 显 示 了 计 算 机 视觉 的 有 效 性 。 对 于 医 学 图 像 分割 ， 例 如 ， U-Mamba和segamba 引 入 了 SSM-CNN 混 合 模 型 ， 直 接 应 用 Mamba 在 像 素 级 别 有 效地建 模 长 程 依 赖 性 。

Mamba潜 力 和 不 足 仍 未 充 分 挖掘 和 解 决：

1. 得益于其线性复杂度，Mamba在空间分配方面具有更大的灵活性。也就是说，与之前基于小卷积核或大小约束的自注意力窗口的方法不同，Mamba在赋予模型大感受野空间建模能力方面很有希望，这在已知的基于Mamba的方法中被忽略了。
2. Mamba 是 一 种 单 向 的 序 列 建 模 方 法 ， 缺乏位置感知 ， 更 关 注 后 验 词 例 。
3. 原 始 的 Mamba 是 为一 维 序 列 建 模 而 提 出的， 不适合需要空间感知理解的计算机视觉任务 。 主 要 由 于 局 部 相 邻 空 间 像 素的不 连 续 ， Mamba 存 在 遗 忘 问 题 和 低 效 的 局 部 建 模 能 力 。

大窗口的曼巴型网络(LMa-UNet)：

1. 提出了一种基于大窗口的Mamba UNet (Large window-based Mamba UNet, LMa-UNet)模型，用于2D/3D医学图像分割。
2. 利用Mamba强大的序列建模能力和线性复杂度，通过为SSM模块分配大窗口来实现大的感受野。进一步，设计了一种新的基于层次和双向大窗口的Mamba块(LM块)，以增强ssm的表示建模能力。
3. 采 用 双 向 Mamba 设计 进行位 置 感 知 序 列 建 模 ， 降 低 输 入 顺 序 的 权 重 影 响 。
4. 层 次 式 的 Mamba 设计 由 两类SSM 操 作 组 成 :(i) 像 素 级 SSM (PiM) 和 (ii) patch 级 SSM (PaM) 。 PiM 通 过 大 窗口SSM 捕 获 邻 域 和 局 部 像 素 信 息 ， 避 免 SSM 中 因 标 记 化 后 相 邻 块 不 连 续 而 出 现

   的 遗 忘 问 题 。 PaM 处 理 长 程 依 赖 建 模 和 全 局 块 交互 。

 主要贡献：

• 提出了一种基于大窗口的Mamba UNet (Large window-based Mamba UNet, LMa-UNet)模型，用于2D/3D医学图像分割。
• 将大感受野窗口分配给SSM层，使模型具有大空间建模的能力。
• 设计了一种基于位置感知的双向Mamba序列建模方法。
• 提出了一种新的由像素级SSM (PiM)和块级SSM (PaM)组成的Mamba模块，增强了局部邻域像素级特征建模和长程全局块级建模。

二、基础知识:SSM模型

基 于 ssm 的 模 型 ， 即 结 构 化 状 态 空 间 序 列 模 型 (structured state space sequence models, S4)和 Mamba ，都是从连续系统演化而来的，该连续系统将一个一维函数或序列 x ( t ) → y ( t ) ∈ R映射到一个隐藏状态 。这个 过 程 可 以 表 示 为下 面 是 线 性 常 微 分 方 程 :

公式（1）

其中：

• 是状态矩阵
• 是投影参数

S4 和 Mamba 是 上 述连 续 系 统 的 离 散 版 本 ， 其 中 包 括 时 间 尺 度 参 数 ∆ 将连续 参 数 A , B 转 换 为 离 散参数 。通常 采 用 零阶 保 持 器 (zero-order hold,ZOH)作 为 离 散 化 规 则 ， 其 定 义 如 下 :

公式（2）

对A, B进行离散化后，公式一(1)的离散化版本(即SSM模型)可以写成: 

公式（3）

 然后通过全局卷积计算输出，定义如下:

公式（4）

其中

• L是输入序列的长度x,
• 是结构化卷积核。 

三、方法

3.1 LMa-UNet

概述LMa-UNet的解如图1所示： 

图1 LMa-UNet 结构图

具 体 来 说 ， 除 了 深 度 卷 积 、 具 有 下 采 样 层 的 编码器 、 具 有 上 采 样 层 的 解 码 器 和 跳跃 连 接 的 常 见 UNet 组 成 外 ， LMa-UNet 通 过在编 码 器 中 插 入 提 出 的 基 于 大 窗 口 的 曼 巴 (LM) 块 来 改 进 UNet 的 结 构 。 给 定 一个分 辨 率 为 C × D × H × W 的 3D 输 入 图 像 ， 深 度 卷 积 首 先 将 输 入 编 码 为特征图  。然 后 将 特 征 图 F 0 输 入 到 每 个 LM 块 和 对 应 的 下 采 样 层中， 得 到 多 尺 度 特 征 图 ; 一个 LM 块 包 含 两个 Mamba 模 块 : 像 素 级 SSM (PiM) 和 块级SSM (PaM)。对于 层，该过 程 可 以 表 述 为 :

公式（5）

其中：

• PiM和PaM分别表示像素级SSM和块级SSM。
• Down-sampling表示下采样层。 

每 个 阶 段 之 后 ， 生 成 的 特 征 图 被 编码为 ;其中 表示特征图 的通道和分辨率。在解码器部分，采用UNet解码器和带有跳跃连接的残差块进行上采样并预测最终的分割掩码。

3.2 LM块

LM 块 是 核 心 组 件 ， 用 于 在 每 个 阶 段 对 不 同 尺 度 的 特 征 图 进 行 进 一 步 的空间 建 模 。 与以 往 利 用 CNN 进 行 局 部 像 素 级 建 模 、 Transformer 进 行 长 程 块 级 依赖建 模 不 同 ， 利 用 Mamba 的 线性复杂度 ， LM 块 可 以 同时完成像素级和块级 建模。 更 重 要 的 是 ， 较 低 的 复 杂 度 允 许设 置 更 大 的 窗 口 来 获 得 更大的感受野 ， 这将提 高 局 部 建 模 的 效 率 ， 如 图 2(a) 所 示 。 具 体 来 说 ， LM 块 是 一个 分 层 设计 ， 由像素 级 SSM (PiM) 和 块 级 SSM (PaM) 组 成 ; 前 者 用 于 局 部邻 域 像 素 建 模 ， 后 者 用于全 局 长 程 依 赖 建 模 。 此 外 ， LM 块 中 每 个 Mamba 层 都 是 双 向 的 ， 用 于 位 置 感知序 列 建 模 。

图2 (a) CNN、Transformer 和我们提出的 LMa-UNet 之间的视野比较。CNN 通常使用小内核（如 3 × 3），Transformers 通常使用 7 × 7 大小的窗口。(b) vanilla Mamba 与我们提出的双向 Mamba 的扫描顺序比较。

1、像素级SSM (PiM) ：

由 于 Mamba 是 连 续 模 型 ， 输 入 像 素 的 离 散 性 会 削 弱 局 部相邻 像 素 的相 关 性 建 模 。 因 此 ， 本 文 提 出 一 种 像 素 级 SSM ， 将 特 征 图 分割 为 多个大 的 子 窗 口 ， 并 在 子 窗 口 上 执 行 SSM 操 作 。 首 先 将 整 个 特 征 图 平 均 划分 为互不重 叠 的 2D 子 窗 口 或 3D 子 立 方 体 ; 以 2D 为 例 。 给 定 H × W 分 辨 率 的 输 入 ， 将特 征 图 划分 为 大 小 为 m × n 的 子 窗 口 ， 每 个 子 窗 口 ( m 和 n 可 以 达 到 40) 。 为了不失 一 般 性 ， 假 设 H/m 和 W/n 都 是 整数 。 然 后 我 们 有 HW/ mn 子 窗 口 ， 如 图 1 的像素 级 SSM 所 示 。 在 该 方 案 下 ， 当 这 些 子 窗 口 被 送 入 Mamba 层 时 ， 局 部 相 邻 像素将 连 续 输 入 SSM; 因 此 ， 可 以 更 好 地 建 模 局 部邻 域 像 素 之 间 的 关 系 。 此 外 ， 在大窗 口 划分 策 略 下 ， 增 大 了 感 受 野 ， 使 模 型 能 够 获 得 更 多 的 局 部 像 素 细 节 。 然而， 图 像 被 划分 为不 重 叠 的 子 窗 口 。 因 此 ， 需 要 一 种 机 制 来 进 行 不 同 子 窗口之 间 的 通 信 ， 以 进 行 长 程 依 赖 关 系 建 模 。

2、补丁级SSM (PaM)：

引 入 补 丁 级 SSM 层 在 不 同 子 窗 口 之 间 传 递 信 息 。 如 图 1 的补丁 级 SSM 所 示 ， 分 辨 率 为 H × W 的 特 征 图 首 先 通 过 大 小 为 m × n 的 池 化 层 ，以允 许 每 个 的 重 要 信 息 HW/ mn将子 窗 口 汇 总 为 单 个代 表 。 因 此 ， 获 得 聚 合 映射 与 HW/ mn 代 表 ， 然 后 用 于 通 过 Mamba 进 行 全 局 范 围 依 赖 建 模 的 子 窗 口 之 间 的通信。在Mamba中的交互之后，将聚合映射解池回初始特征映射 的相同 大小 ， 并应 用 残 差 接，公式 (5)中 的过程 可 以 这 样 进 行 :

公式（6）

其 中：

• Pooling和Unpooling分别表示池化层和非池化层。
• Bi-Mamba表示提出的双向Mamba层。

3、Bidirectional Mamba (BiM)：

与 基 于 前 向 扫 描 方 向 SSM 层 的 普 通 Mamba 块不同 ， 我 们 的 LM 块 中 的 每 个 SSM 层 ( 包 括 PiM 和 PaM) 是 双 向 的 。 图 2 (b) 显 示 了差异 。 在 原 始 的 曼 巴 (Mamba) 中 ， 作 为一个 连 续 模 型 ， 一些 信 息 遗 忘 发 生 在 较早输 入 的 元 素 上 ， 而 最 新 输 入 的 元 素 将 保 留 更 多 的 信 息 。 因 此 ， 原 始 的 单 一 扫描方 向 的 Mamba 将 更 多 地 关 注 后 片 ， 而 不 是 特 征 图 的 中 心 区 域 ， 通 常 具 有 更 多的器 官 和 病 变 。 通 过 同 时 进 行 前 向和后向 扫 描 并 叠 加 输 出 结 果 ， 提 出 了一 种 双向Mamba 结 构 。 详 细结 构 如 图 1 左 侧 所 示 。 BiM 有 两个 优 点 。 首 先 ， 该 模 型 可以更 多 地 关 注 图 像 中 心 区 域 可 能 具 有 更 多 器 官 和 病 变 的 信 息 块 ， 而 不 是 角 落 区域。 其 次 ， 对 于 每 个 块 ， 网 络 都 可 以 很 好 地 建 模 其 与 其 他 块 的 绝 对 位 置 信 息 和相对 位 置 信 息 ;

四、实验

4.1数据集

在 2D 和 3D 分割 任 务 的 两个 数 据 集 上与 最 先 进 的 方 法 进 行 了 实 验 比 较 ， 以 验证LMa-UNet 的 有 效 性 和 可 扩 展 性 。

腹部CT 。 腹 部 CT 是 一个 公 开 可 用 的 3D 多 器 官 分割 数 据 集 ， 包 含 来 自 MICCAI 2022 FLARE Challenge [16]的 100 个 CT 病 例 ， 包 括 13 种 腹 部 器 官 ( 肝 、 脾 、 胰腺、 右 肾 、 左 肾 、 胃 、 胆 囊 、 食 管 、 主 动 脉 、 下 腔 静 脉 、 右 肾 上 腺 、 左 肾 上腺和 十 二 指 肠 ) 。 三 维 CT 图 像 的 大 小 为 40 × 224 × 192 。 来 自 MSD 胰 腺 数 据 集的50 例 与 来 自腹腔 的 注 释 一 起 用 于 训 练 ， 另 外 来 自 不 同 医 疗 中 心 的 50 例 用 于测试 。

腹部 MR。Abdomen MR 是一个公开可用的 2D 分割数据集，包含来自 MICCAI 2022 AMOS Challenge [11] 的 110 个 MRI 病例，包括 13 种腹部器官（与腹部 CT 数据集相同）。2D MRI 图像的大小为 320 × 320。60 个带注释的案例用于训练，另外 50 个案例用于测试。

4.2实现设置

我们的 LMa-UNet 基于 nnU-Net 框架在 PyTorch 1.9.0 上实现。所有实验均在 NVIDIA GeForce RTX 3090 GPU 上进行。3D 数据集（Abdomen CT）的训练批量大小为 2，2D 数据集（Abdomen MR）的批量大小为 24。使用动量 = 0.99 的 Adam [12] 优化器。初始学习率为 0.01，权重衰减为 3e-5。最大训练 epoch 数为 1000。对于Abdomen CT 数据集，阶段为 6，但维度不一致；因此，我们将每个阶段的矩形窗口大小设置为 [20, 28, 24]、[20, 28, 24]、[10, 14, 12]、[10, 14, 12]、[5, 7, 6] 和 [5, 7, 6]。对于 Abdomen MR 数据集，阶段为 7，每个阶段窗口大小为 40、20、10、10、5 和 5。

4.3综合性能

基 线 模 型 包 括 三 种 类 型 的 代 表 性 医 学 图 像 分割 网 络 : 基 于 cnn 的 网 络 (nnU-Net和 SegResNet) ， 基 于 transformer 的 网 络 (UNETR, SwinUNETR 和 nnFormer)， 以 及 最 新 的 基 于 mamba 的 网 络 (U-Mamba [15]) 。 为了 公 平 比 较 ， 我 们 还在nnU-Net 框 架 中 实 现 了 所 有 模 型 ， 并 使 用 默 认 的 图 像 预 处 理 。 表 1 显 示 了 结果。 与 基 于 cnn 和 基 于 transformer 的 分割 方 法 相 比 ， 提 出 的 LMa-UNet 在 以下 方面取 得 了 更 好 的 性 能 DSC 和 NSD 均 表 明 ， Mamba 的 全 局 建 模 能 力 对 医 学 图 像 分割至 关 重 要 。 请 注 意 ， 与 U-Mamba 相 比 ， U-Mamba 只 是 将 Mamba 作 为 全 局 建模适 配 器 ， LMa-UNet 展 示 了 对 U-Mamba 的 改 进 ， 验 证 了 双 向和 分 层 Mamba 设计的 有 效 性 。 这 些 结 果 还 表 明 了 Mamba 在基 于 大 窗 口 的 感 受 野 的 全 局 和 局 部 特征建 模 方 面 的 潜 力 。

4.4LMa-UNet的窗口大小 

由 于 Mamba 的 线 性 复 杂 度 ， LMa-UNet 可 以 实 现 cnn 和 transformer 无 法 实 现 的 大窗口 建 模 。 为了 验 证 Mamba 的 巨 大 空 间 建 模 潜 力 ， 探 索 了 LMa-UNet 在 不 同 窗口大小设置下的性能。表2显示了腹部MR数据集的结果，该数据集的分期为7。比较三种窗口大小设置的性能，可以发现窗口大小较大的LMa-UNet具有更好的性能。这表明大感受野对于医学图像分割是至关重要的，而由于Mamba的线性复杂性，它可以实现医学图像分割。这将为未来的大型内核或窗口建模提供更多可能性。

4.5消融实验

为了 验 证 LMa-UNet 中 每 个 关 键 成 分 的 效 果 ， 在 腹 部 MR 数 据 集 上 进 行 了 消 融 实验。 表 3 显 示 了 结 果 。 实 验 结 果 表 明 ， PiM 和 PaM 对 LMa-UNet 的 建 模 效 果 都 优于基 线 模 型 ， 分别 验 证 了 PiM 和 PaM 在 局 部 像 素 级 建 模 和 全 局 建 模 方 面 的 有 效性。 值 得 注 意 的 是 ， PiM 模 型 比 PaM 模 型 获 得 了 更 多 的 改 进 ， 这 表 明 扩 大 局 部特征 建 模 的 感 受 野 是 提 高 模 型 性 能 的 关 键 。 引 入 BiM 后 ， LMa-UNet 的 性 能 进 一步提 升 ， 说 明 了 双 向 Mamba 对 位 置 感 知 序 列 建 模 的 重 要 性 。 最 后 ， 所 有 组 件的LMa-UNet 均 取 得 了 最 好 的 性 能 ， 进 一 步 证 明 了 算 法 的 有 效 性 该 方 法 及 其 组 成部分 的 有 效 性 。

 五、结论

提 出 了一 种 新 的 基 于 mamba 的 UNet 医 学 图 像 分割 模 型 ， 实 现 了 大 窗 口 空 间 建模。 此 外 ， 为了 增 强 Mamba 的 局 部 和 全 局 特 征 建 模 能 力 ， 设计 了一 种 双 向 分层SSM 。 在 多 器 官 分割 数 据 集 上 的 综 合 实 验验 证 了 所 提 方 法 的 有 效 性 。 在 未来， 我 们 将 探 索 SSMs 在 其 他 医 学 成 像 任 务 中 的 可 扩 展 性 和 应 用 ， 如 检 测 、 配准、 重 建 等 。