Decoding Matters: Efficient Mamba-Based Decoder with Distribution-Aware Deep Supervision for Medical Image Segmentation¶
会议: CVPR 2025
arXiv: 2603.12547
代码: 待发布(接收后公开)
领域: 医学图像
关键词: medical image segmentation, Mamba, decoder-centric, deep supervision, KL-divergence
一句话总结¶
提出 Deco-Mamba,一种以解码器为核心的混合 Transformer-CNN-Mamba 架构,通过 Co-Attention Gate、Vision State Space Module 和可变形卷积精炼块增强解码器能力,并引入基于窗口化 KL 散度的分布感知深度监督策略,在 7 个医学图像分割基准上取得 SOTA 性能,同时保持适中的模型复杂度。
研究背景与动机¶
- 领域现状:医学图像分割主流范式是 U-Net 及其变体(Att-UNet、UNet++等),后发展出 Transformer 架构(TransUNet、Swin-UNet)捕获长距离依赖,以及 Mamba/SSM 架构(U-Mamba、SliceMamba)以线性复杂度实现全局建模。
- 现有痛点:(a) 大多数方法是任务特定的,在单个数据集上表现好但跨模态泛化差;(b) 大多数工作聚焦编码器增强,使用大预训练骨干网络增加计算复杂度,但解码器设计被忽视;(c) Mamba 方法虽降低计算开销,但通常只在少量模态上验证,泛化能力不明。
- 核心矛盾:强编码器配弱解码器——即使编码器提取了丰富语义,设计不足的解码器在上采样过程中难以准确恢复目标边界和上下文结构,导致精细空间细节丢失。同时,现有级联解码器(如 MERIT)大幅增加计算复杂度。
- 本文要解决什么? (a) 如何在解码器中高效建模多尺度上下文和长距离依赖?(b) 如何在不增加过多参数的情况下提升解码器的重建精度?(c) 如何实现跨模态的泛化能力?
- 切入角度:以解码器为核心进行设计创新,编码器保持轻量,将 Mamba 的高效全局建模能力嵌入解码器,并提出分布感知深度监督避免传统深度监督中 resize 导致的信息损失。
- 核心idea一句话:轻编码器 + 重解码器(CAG+VSSMB+可变形卷积),配合基于窗口 KL 散度的分布感知深度监督。
方法详解¶
整体框架¶
Deco-Mamba 采用 U-Net 架构,编码器为双分支:CNN 分支提取高分辨率局部特征,PVT Transformer 分支提取四阶段全局特征。解码器为六阶段结构,核心模块包括 Co-Attention Gate (CAG)、Visual State Space Mamba Block (VSSMB) 和 Double Deformable Residual Block (DDConv)。提供 V0 (PVT-V2-B0, 9.67M) 和 V1 (PVT-V2-B2, 46.93M) 两个变体。
关键设计¶
- Co-Attention Gate (CAG):
- 做什么:自适应融合编码器跳跃连接特征和解码器特征
- 核心思路:不同于标准 AG 只用解码器特征门控编码器特征,CAG 让编码器和解码器特征互相门控——两个交叉注意力门的输出拼接后经 Channel Attention 精炼,\(D_i' = CA[AG(x=X_i, g=D_{i+1}), AG(x=D_{i+1}, g=X_i)]\)
-
设计动机:标准 AG 忽略解码器特征中的空间显著性,且只考虑空间注意力不考虑通道关系;双向门控 + 通道注意力补齐这两个缺陷
-
Vision State Space Mamba Block (VSSMB):
- 做什么:在解码器中以线性复杂度捕获长距离依赖
- 核心思路:将连续时间 SSM 应用于高度和宽度两个空间维度,通过四方向选择性扫描(水平、垂直及其反方向)传播上下文信息。瓶颈处用2个 VSSMB,2-5阶段各1个,最后一阶段不用
-
设计动机:Transformer 自注意力复杂度为 \(O(n^2)\),Mamba 的 SSM 为 \(O(n)\),在解码器中更适合处理逐步增大的空间分辨率
-
Deformable Residual Block (DRB):
- 做什么:恢复被 VSSMB 全局建模平滑掉的局部空间细节
- 核心思路:在残差框架内集成标准 3×3 卷积和可变形卷积(DCN)。DCN 预测像素级偏移和调制掩码,偏移分支估计采样位移、调制分支赋予像素重要性权重 \([0,2]\)
- 设计动机:SSM 擅长全局建模但可能忽略微妙的局部变化(如复杂组织边界),DRB 通过可变形卷积的几何自适应性恢复边界精度
损失函数 / 训练策略¶
\(\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{dice}} + \mathcal{L}_{\text{multi}}\)
Multi-Scale Distribution-Aware (MSDA) 深度监督: - 在每个解码器阶段的原生分辨率上操作,避免传统深度监督将中间输出 resize 到 GT 分辨率导致的信息损失 - 将 GT 标签用局部窗口平均得到类别频率分布 \(\tilde{P}^{(s)}\),将预测经分布头 + softmax 转为 \(Q^{(s)}\),计算 KL 散度 - 引入边界感知加权 \(W^{(s)}_{h,w} = (1 - \max_n \tilde{P}^{(s)}_{h,w,n})^\alpha\),对混合类别区域(边界附近)赋予更高权重 - 多尺度聚合时深层解码器阶段权重更大:\(\lambda_1 < \lambda_2 < \cdots < \lambda_S\)
实验关键数据¶
主实验(Synapse 多器官分割,8 类)¶
| 方法 | Dice↑ | HD95↓ | 参数量(M) | FLOPs(G) |
|---|---|---|---|---|
| Deco-Mamba-V1 | 85.07 | 14.72 | 46.93 | 17.24 |
| Deco-Mamba-V0 | 83.16 | 15.89 | 9.67 | 9.73 |
| Cascaded-MERIT | 83.59 | 15.99 | 147.86 | 33.31 |
| PAG-TransYnet | 83.43 | 15.82 | 144.22 | 33.65 |
| SliceMamba | 81.95 | 16.04 | - | - |
| Swin-UMamba | 80.34 | 21.51 | 59.88 | 31.35 |
消融实验(Synapse)¶
| 配置 | Dice↑ | HD95↓ | 说明 |
|---|---|---|---|
| Full Deco-Mamba-V1 | 85.07 | 14.72 | 完整模型 |
| w/o CNN 分支 | 84.07 | 18.92 | 掉1点Dice |
| w/o VSSMB | 83.51 | 15.96 | 去掉Mamba掉1.56% |
| CAG→AG | 82.98 | 15.69 | 标准AG掉2.09% |
| CAG→CBAM | 84.01 | 16.19 | CBAM次优 |
| DRB→标准Conv | 84.53 | 16.18 | 可变形卷积有帮助 |
| Dice only | 83.84 | 14.94 | 无深度监督 |
| Dice + 传统深度监督 | 84.24 | 15.89 | HD95反而变差 |
| Dice + MSDA (ours) | 85.07 | 14.72 | MSDA最优 |
关键发现¶
- Deco-Mamba-V0 仅 9.67M 参数就达到 83.16% Dice,接近 150M 参数的 Cascaded-MERIT (83.59%),效率提升 15×
- 在 7 个数据集(跨 CT、MRI、皮肤镜、腺体、细胞核)上均取得一致优势,验证跨模态泛化能力
- 传统深度监督提升 Dice 但恶化 HD95,因为低分辨率预测被强制放大破坏边界——MSDA 规避了这一问题
- MoNuSeg 上 Deco-Mamba-V1 比 Swin-UMamba 高 +4.46% Dice,比 U-Net 高 +8.69%
亮点与洞察¶
- 解码器中心设计是反直觉但有效的:大多数工作堆叠更强编码器,本文证明将创新放在解码器、编码器保持轻量同样甚至更好。这对计算受限场景(边缘部署)很有启发
- MSDA 深度监督解决了经典 resize 问题:在原生分辨率上用 KL 散度匹配类别分布而非像素级预测,避免信息损失且自然引入边界感知。这个监督策略可独立于架构使用
- 窗口化分布比像素级 GT 更适合多尺度监督:将小窗口内的 GT 类别频率作为软标签,比硬缩放更平滑,也自然反映了低分辨率下的语义分布
局限性 / 可改进方向¶
- 仅支持 2D 分割,未扩展到 3D 医学图像
- 代码未公开,可复现性待验证
- PVT 编码器仍是预训练模型,编码器的贡献与解码器的贡献难以完全解耦
- 各组件(CAG、VSSMB、DRB、MSDA)的组合较多,最优配置可能因数据集而异
相关工作与启发¶
- vs EMCAD:EMCAD 也关注解码器但用注意力+轻量卷积无长距离建模能力,Deco-Mamba 通过 VSSMB 弥补这一缺陷
- vs Cascaded-MERIT:MERIT 用级联编解码器(148M参数)实现 83.59% Synapse Dice,Deco-Mamba-V1 用 47M 参数达到 85.07%,效率更高性能更好
- vs SliceMamba/VM-UNet:同为 Mamba 方法但聚焦编码器,在 BTCV 等复杂多器官任务上表现不如 Deco-Mamba,印证解码器的重要性
评分¶
- 新颖性: ⭐⭐⭐⭐ 解码器中心+分布感知深度监督是清晰有价值的创新点,但各组件分别来看并非首创
- 实验充分度: ⭐⭐⭐⭐⭐ 7个数据集、多种模态、详细消融、多个backbone对比,非常充分
- 写作质量: ⭐⭐⭐⭐ 结构清晰,图表丰富,方法描述详细
- 价值: ⭐⭐⭐⭐ 关于解码器设计和分布感知监督的经验对后续工作有参考价值