【医学分割】跨尺度全局状态建模和频率边界指导的分割架构
SkinMamba: A Precision Skin Lesion Segmentation Architecture with Cross-Scale Global State Modeling and Frequency Boundary Guidance
本文提出了一种基于 Mamba 和 CNN 的混合架构,称为 SkinMamba。它在保持线性复杂性的同时,提供了强大的长距离依赖关系建模和局部特征提取功能。具体来说,引入了尺度残差状态空间块(SRSSB),它能在宏观层面捕捉全局上下文关系和跨尺度信息交换,从而实现全局状态下交流。这有效地解决了皮损大小不一和目标区域不明显所带来的皮损分割难题。此外,为了减少模型下采样过程中的边界模糊和信息丢失,引入了频率边界引导模块(FBGM),提供足够的边界先验来引导精确的边界分割,同时还利用保留的信息来协助解码器进行解码。
Introduction
近年来,状态空间模型(SSM)吸引了众多研究人员的关注。S4 [15] 最初将 SSM 应用于深度学习领域,但其性能仍然落后于 CNN 和 Transformer。多亏了 Mamba,它通过高效的选择性扫描机制增强了 SSM,不仅建立了长程依赖关系,还展示了与输入大小相关的线性复杂性。受 Mamba 的启发,Vision Mamba和 Visual Mamba最先将 Mamba 引入视觉领域,并取得了显著成果。此外,许多研究人员已将其应用于医学图像分割,如 VM-UNet 、U-Mamba 、OCTAMamba 和 H-vmunet。然而,皮损分割仍面临挑战。例如,皮损区域的边界不清晰,皮损区域的大小不一,以及在下采样过程中存在严重的信息丢失。此外,由于无法捕捉局部细粒度特征,直接将 Mamba 应用于皮损分割会降低准确率。
为了解决上述问题,提出了 SkinMamba,它结合了 Mamba 在学习全局特征和 CNN 在提取局部特征方面的优势。具体来说,整体架构采用 5 级编码器-解码器结构。模型的核心组件是规模残差状态空间模块(SRSSB)和频率边界引导模块(FBGM)。在 SRSSB 中,引入了视觉状态空间块(VSSB)来实现高效的全局建模,而规模混合前馈层(SMFFL)则用于提取跨规模特征,促进多级信息的交互,实现全局背景下的多规模、多级特征表示。为了减少下采样造成的边界信息损失,引入了频率边界引导模块(FBGM),该模块从频率角度捕捉边界线索,并指导解码过程中的编码器。通过这种设计,开发出了 SkinMamba 模型,它既能保持计算效率,又能有效捕捉短程和长程依赖关系。它还解决了皮损图像中边界不清和皮损大小不一等难题,进一步探索了 Mamba 在皮损分割中的应用。
Method
Architecture Overview
提出的 SkinMamba 如图 1 所示。它包括初始 Conv 层、编码器、FBGM、解码器、最终投影层和跳转连接。给定输入图像 x∈R^H×W ×3,初始 Conv 层首先将 x 的通道映射到 C,C 默认为 16。然后,编码器执行深度特征提取。具体来说,编码器由五个阶段组成,每个阶段包括一个编码器块和一个降采样层。经过每个阶段后,输入被映射到更深的特征空间,同时图像的宽度和高度减半,通道数加倍。同样,解码器由五个阶段组成。每个阶段包括一个解码器块和一个上采样层,其中图像的宽度和高度增加一倍,通道数减少一半。值得注意的是,在编码器和解码器之间,引入了 FBGM,从频率的角度提取精确的边界线索,为皮损分割提供强有力的约束,并指导图像解码。在解码器之后,最终投影层会恢复通道数,使其与分割目标相匹配。对于跳转连接,使用简单的加法运算来避免不必要的参数,从而更容易验证我们架构的有效性。
Encoder and Decoder Blcoks
融合与平衡是架构的核心主题。编码器和解码器模块都由 CNN 和 SRSSB 组成,将 Mamba 学习全局特征的优势与 CNN 提取局部特征的能力结合起来。
如图 2 所示,在编码器区块中,输入特征首先通过 SRSSB,然后经过 3×3 卷积、批归一化(BN)层和 ReLU 激活函数,产生也用于跳转连接的输出特征。在解码器区块中,每个解码器区块接收来自 FBGM 或前一个解码器区块的特征,以及来自编码器区块跳转连接的相应输出特性。然后将这两组特征沿通道维度汇总,形成融合特征。融合后的特征经过 3×3 卷积、BN 层和 ReLU 激活函数,最后输入 SRSSB。
Scale Residual State Space Blcok
SRSSB 用于从全局流中提取跨尺度混合信息,并建立长程依赖关系模型。如图 2 所示,SRSSB 的核心组件是视觉状态空间块(VSSB)和尺度混合前馈层(SMFFL)。具体来说,它分为两个阶段。在第一阶段,输入特征首先经过层归一化处理 ,然后通过 VSSB 对全局信息进行建模。第二阶段采用第一阶段的特征,应用层归一化以防止模式崩溃,然后通过 SMFFL 进行处理,以提取多尺度特征。这样就得到了全局背景下的尺度特征,而残余连接有助于在 SRSSB 块之间保持一致的信息流,从而加速模型收敛。上述过程定义如下:
视觉状态空间块。图 2 展示了视觉状态空间块。输入特征 X 经过层归一化和线性层,生成两个信息流。这一过程可以表示为 :
其中,Linear(-) 表示通过线性层进行处理,F1 和 F2 分别是第一和第二信息流的输入。在第一信息流中,输入特征 F1 通过深度可分离卷积(DWConv)和 SiLU 激活函数进行处理,提取初步特征。然后通过二维选择性扫描(SS2D)模块对这些特征进行细化,并进行归一化处理,得到细化特征 F′1。在第二流中,输入特征 F2 由 SiLU 激活函数处理,生成补充特征 F ′ 2。然后,精炼特征和补充特征按元素顺序相乘、合并,并通过线性层,产生 VSSB 的最终输出,通过详细提取和高效的全局建模实现全面的特征表示。上述过程定义如下:
规模混合前馈层。如图 2 所示,SMFFL 是一种双分支多尺度架构。具体来说,输入特征 M ap 经过层归一化后被分为两个分支。由线性层和 3×3 或 5×5 卷积层组成。首先,通过线性层将特征映射到低维空间,然后使用不同大小的卷积从低维隐藏层捕捉特征。随后,将不同尺度的低维特征沿通道聚合,得到尺度混合特征。最后,应用 GELU 进行非线性变换,并使用线性层将特征重映射回原始高维,同时将残余特征添加到尺度混合特征中,以促进梯度流动和传播。该模块的整个计算过程定义如下:
其中,f x×y(-) 表示 X × Y 标准卷积运算,[-, -] 表示连接运算,GELU(-) 表示 GELU 激活函数,⊕ 表示元素相加。此外,还可以看到,虽然 VSSB 不需要位置编码,但 SMFFL 在促进多尺度信息流的互动方面发挥着至关重要的作用。SMFFL 加强了同步操作,使每个维度和尺度都能贡献自己的专业知识,从而实现精确和互补的细节提取。
Frequency Boundary Guided Module
一般来说,低频信息包含位置关系和空间信息,而高频信息主要包括具体的边界细节。然而,大多数方法在编码时都会通过降采样来降低图像分辨率,以减轻计算负荷。在这种情况下,解码只能依靠低分辨率图像,从而导致性能下降和高频信息丢失。高频信息对于准确分割皮肤病变区域至关重要。为此,从频率角度开发了 FBGM,以提供足够的边界先验,准确提取边界线索,并指导解码过程。如图 2 所示,频率特性引导调制层(FFGML)和紧凑型简单前馈层(CSFFL)是 FBGM 的核心组件。给定输入特征 Kl-1,拟议的 FBGM 可表示如下:
其中 K′ l 和 Kl 分别代表 FFGML 和 CSFFL 的输出特性。
频率特性引导调制层。图 2 展示了 FFGML。具体来说,给定输入特征 F0 后,使用快速傅立叶变换(FFT)将其转换到频域。然后,用两层点式卷积(PWConv)处理这些频域特征,中间穿插一个 ReLU 激活函数。最后,使用反向 FFT 将处理过的特征转换回原始空间域,然后使用 Sigmoid 激活函数进行二次调制。这种调制输出与输入特征相互作用,从而获得原始潜层隐藏空间的更新特征表示。上述过程定义如下 :
其中,θ(-) 表示 FFT,θ-1(-) 表示反向 FFT,PW(-) 表示点式卷积,Sigmoid(-) 表示 Sigmoid 激活函数。
紧凑型简单前馈层。受视觉变换器(ViT) 的启发,设计了一个紧凑型简单前馈层(CSFFL),用于非线性变换和降维。给定输入特征 X,首先应用 3×3 卷积来提高维度,同时增强通道数量和局部特征表示。然后通过 GELU 激活函数引入非线性。最后,使用 PWConv 恢复原始通道数。整个模块的定义如下 :