DAMamba-UNet3D: A Parameter-Efficient Mamba State Space U-Net with Dynamic Adaptive Scan for 3D Medical Image Segmentation

技术摘要：DAMamba-UNet3D

问题陈述
3D 医学图像分割需要在精确的局部边界描绘与连贯的全局空间上下文之间取得平衡，同时还必须遵守严格的计算和参数预算。现有架构面临着明显的权衡：

• 卷积 U-Net： 提供 $O(n)$ 的局部混合和参数效率，但缺乏显式的全局上下文，需要通过增加深度来捕捉长程依赖关系。
• 视觉 Transformer： 通过自注意力机制提供全局推理能力，但对于大型 3D 体素（例如，$128^3$ 个补丁包含约 210 万个 token）会产生难以承受的 $O(n^2)$ 内存和计算成本。
• 状态空间模型 (SSMs/Mamba)： 以 $O(n)$ 的时间和内存传播全局信息。然而，现有的医学 SSM 分割器（如 SegMamba）依赖于固定的扫描模式，且参数预算较大（约 70M 参数），缺乏数据依赖的适应性。

虽然动态自适应扫描 (Dynamic Adaptive Scan, DAS) 已被引入 2D 自然图像识别中，用于在选择性扫描前学习数据依赖的重排序，但尚未应用于医学成像或扩展到 3D 体素。

方法论
作者提出了 DAMamba-UNet3D，这是一种将 三平面 3D-DAS (Tri-Plane 3D-DAS) 模块集成到 U-Net 框架中的混合编码器-解码器架构。

1. 架构设计：

   • 混合结构： 模型保留了卷积块用于 Stem 层 (E1)、瓶颈层 (Bottleneck) 和整个解码器，以保留局部细节。
   • 编码器集成： 三平面 DAMamba 模块仅插入在编码器阶段 E2–E4。
   • 紧凑型变体： 主模型使用基础宽度 $C=16$，其参数量约为 530 万（比 SegMamba 少约 13 倍）。
   • 大规模变体 (DAMamba-L)： 一个更宽的变体（约 70M 参数），在编码器中堆叠了 DAMamba 模块，并配有卷积瓶颈层，旨在匹配 SegMamba 的容量以便进行直接比较。

2. 三平面动态自适应扫描 (Tri-Plane DAS)：

   • 机制： 与固定扫描（行优先、蛇形或多方向展平）不同，DAS 学习输入依赖的空间重排序。偏移预测网络 (OPN) 从特征图中预测 2D 偏移量，这些偏移量用于在选择性扫描前通过双线性插值对特征进行重采样。
   • 3D 扩展： 通过在三个解剖平面（轴向 $H \times W$、冠状 $D \times W$ 和矢状 $D \times H$）上应用 2D DAS 算子，将 DAS 扩展到体积数据。通过对输出求平均值来构建 3D DAS 模块：
     $$DAS3D(X) = \frac{1}{3}(DAS_{ax} + DAS_{cor} + DAS_{sag})$$
   • 选择性扫描： 重排序后的特征图由标准的 Mamba 选择性扫描块处理，从而实现 $O(N)$ 成本下的全局上下文传播，而无需二次方复杂度的注意力图。

3. 训练与评估：

   • 数据集： BraTS 2020 训练集（369 个受试者，4 种模态）。
   • 损失函数： 调制组合损失 ($L = (L_{CE} + \delta) \cdot (L_{mod}^{Dice} + \delta)$)，旨在聚焦于困难体素和边界。
   • 协议： 五折受试者级交叉验证，采用全卷推理（滑动窗口，$128^3$ 补丁）。指标包括增强肿瘤 (ET)、肿瘤核心 (TC) 和全肿瘤 (WT) 的平均 Dice 系数和 HD95。

核心贡献

1. 首次医学应用 DAS： 将动态自适应扫描从 2D 自然图像分类转移到稠密 3D 医学分割领域。
2. 2D 到 3D 的扩展： 通过三平面自适应扫描将 DAS 算子提升至体积数据，为手设计的多方向展平提供了一种可学习的替代方案。
3. 参数效率： 证明了仅在编码器中集成三平面 DAS 块的混合 U-Net，在实现近乎等同性能的同时，参数量仅为较大模型的约 1/13。
4. 双尺度验证： 在紧凑规模（$\sim 5.3$M）和大规模（$\sim 70$M）上均验证了该方法，表明当架构位置优化时，学习到的 DAS 可以优于固定的扫描模式。

结果

• 紧凑模型 (DAMamba-UNet3D, $\sim 5.3$M)： 在 BraTS 2020 上实现了 0.815 $\pm$ 0.013 的平均 Dice 系数。尽管参数量减少了约 13 倍，但其表现与重新训练的 SegMamba (0.824 $\pm$ 0.014) 近乎持平。
• 大规模模型 (DAMamba-L, $\sim 70$M)： 超越了重新训练的 SegMamba，平均 Dice 系数为 0.829 $\pm$ 0.012（提升了 0.5 个百分点），其中在肿瘤核心 (TC) 类别的提升最为显著（0.842 对比 0.827）。
• 消融实验：
  • 位置： 将 DAS 块仅放置在编码器（E2–E4）中效果最好（0.815）。将 DAS 添加到瓶颈层或解码器会降低性能（分别为 0.807 和 0.805），这表明仅在编码器中部署是至关重要的。
  • 扫描模式： 当编码器位置匹配时，学习到的三平面 DAS 优于固定扫描顺序（Raster, Snake, Cross, Tri-orientated）。

意义与主张
论文声称，在混合 U-Net 中使用学习到的三平面 DAS 是针对 3D 脑肿瘤分割的一种具有竞争力的固定三向 Mamba (ToM) 扫描替代方案。

• 效率： 紧凑型模型证明，通过利用自适应扫描而非单纯依靠增加容量，可以实现高水平的性能，同时大幅减少参数数量。
• 性能： 在大规模下，学习到的自适应机制 (DAS) 在相同的训练和评估协议下，可以改进固定扫描轨迹 (ToM) 的表现。
• 审慎性： 作者指出，虽然在 BraTS 2020 上的结果令人鼓舞，但本研究仅限于单一数据集。他们表示，“在对 DAS 与 ToM 做出普遍性结论之前，需要进行多数据集验证和统计检验”。其主要贡献在于将 DAS 扩展到 3D 医学体素，并证明了其在参数高效型混合架构中的可行性。