UNetMamba文献阅读笔记

2025 IEEE
论文地址：https://ieeexplore.ieee.org/document/10766630
代码地址：https://github.com/EnzeZhu2001/UNetMamba

一 背景

1. Transformer 的引入提升了语义分割精度，但它的二次计算复杂度和高参数数量限制了在高分辨率图像中的应用效率
2. Mamba 具有线性复杂度和长距离依赖建模能力，基于 Mamba 的模型在医学图像分割和遥感图像解释中展现出潜力，但未充分发挥其精度和效率的双重优势
3. 高分辨率遥感图像，信息复杂度，精度，效率难以兼得
4. 基于CNN的Unet建立了语义分割的基础架构：带有跳跃连接的U形“编码器-解码器”框架

二 目的

克服准确性和效率之间的困境

三 主要思想

基于Mamba的UNet类
采用类似于UNet的U形“编码器-解码器”框架
主要有三个部分：

1. 基于多阶段ResT主干的编码器
2. 基于Mamba的高效分割解码器（MSD）
3. 局部监督模块（LSM）： 旨在增强局部语义信息感知

以高效解码高分辨率图像中的复杂信息并增强局部语义信息的感知

四 方法介绍

4.1 整体架构

基于 U 形框架构建，跳跃连接，采用预训练的 ResT 作为编码器，MSD 作为解码器，并设计 LSM 增强局部语义信息感知

4.2 Res T编码器

• 多阶段特征提取
  • 以高效变压器块（Efficient Transformer Block, ETB）为核心，由四个阶段组成
  • 捕捉多尺度特征图，适应遥感图像中不同大小的目标物体和特征
• 采用EMSA
  • ETB 采用高效多头自注意力（EMSA）
  • 引入实例归一化（IN）和 1×1 卷积
  • 目的
    • 为解决传统多头自注意力机制的二次计算成本问题
• 在保证性能的同时提高计算效率，确保 UNetMamba 在处理高分辨率图像时能快速编码，避免二次计算成本

4.3 Mamba分割解码器（MSD）

• Mamba优势

  • 考虑到高分辨率遥感图像展开后类似长序列数据，Mamba 的线性缩放能力在处理时具有天然优势，将其引入模型

• 优化解码器结构

  • 引入 VMamba 的 VSS 块构建 MSD
    • 解码语义信息，具有全局感受野和线性复杂度
  • VSS 块先对 Patch 扩展后的 2-D 特征图进行层归一化，再经线性嵌入、3×3 深度卷积、SiLU 激活，通过 2-D 选择性扫描（SS2D）模块在四个方向扫描，以线性复杂度在全局感受野下解码语义信息。最后经层归一化、与旁路流元素相乘并通过线性层和残差连接输出
  • 在四个不同阶段解码多尺度特征图
  • 最终通过 1×1 卷积头输出语义分割结果

• 降低参数数量的同时利用长距离建模能力准确解码复杂信息

4.4 局部监督模块（LSM）

• 为解决 MSD 解码时对局部语义信息的部分忽视

  • VSS 块大感受野使 MSD 在解码时会忽略部分局部语义信息
  • 高分辨率图像的局部细节对分割精度至关重要

• 采用两个不同尺度卷积的并行分支

  • 对 MSD 输出的特征图处理后，经批归一化和 ReLU6 激活再合并

• 训练时在解码器 2 - 4 阶段添加 LSM 块

  • 计算辅助损失进行监督训练

• 以提高局部细节的感知，且训练阶段设计节省推理成本

4.5 损失函数

• 由用于整体优化的主损失Lp​和用于局部监督的辅助损失La​组成
  • Lp​由骰子损失Ldice​和交叉熵损失Lce​构成
  • La​采用交叉熵损失Lce​
  • 通过加权方程L=Lp​+αLa​平衡两者效果
    • α设为 0.4

五 实验

5.1 数据集

LoveDA和ISPRS Vaihingen数据集

5.2 评估指标

5.2.1 效率评估指标

模型参数量（Param）
内存占用（Memo）
浮点运算次数（FLOPs）

5.2.2 精度评估指标

平均F1分数（mF1）
平均交并比（mIoU）
总体准确率（OA）