【论文笔记】P-Mamba: Marrying Perona Malik Diffusion with Mamba for Efficient Pediatric Echocardiographic

P-Mamba: Marrying Perona Malik Diffusion with Mamba for Efficient Pediatric Echocardiographic Left Ventricular Segmentation
原文链接：https://arxiv.org/abs/2402.08506

1. 引言

超声波心动图是早期检测和治疗先天性心脏病的手段，准确分割其中的心脏结构（尤其是左心室）是关键。但其中的背景噪声会导致分割不精确，且需要考虑分割效率。

本文提出P-Mamba，可在减小噪声的同时保留局部目标边界细节，以达到最优的性能。此外，本文方法还有很高的效率。

P-Mamba的灵感来源为选择性状态空间模型Mamba和Perona–Malik扩散（PMD），前者可更好地建模长距离依赖性，且是硬件感知的、线性复杂度算法，后者则可减小图像噪声并保留边界细节。

2. 方法

如上图所示，P-Mamba分为3个组件：视觉Mamba编码器、基于DWT的PMD编码器、解码器。视觉Mamba编码器用于捕捉全局依赖并保证高效率，基于DWT的PMD编码器抑制背景噪声并保留边界细节。解码器分为分割头和全卷积头。

2.1 基于DWT的PMD块

给定输入特征图$u$，其PMD公式为：
$$\frac{\partial u}{\partial t}=div(g(|\nabla u|)\nabla u)$$

其中$g(|\nabla u|)=\frac{1}{1+(\frac{|\nabla u|}{k}{)}^2}$为扩散系数，$t$为扩散步数（可视为特征图的层数），$k$为控制扩散程度的正常数。该公式为各向异性扩散公式，在平坦或光滑区域，梯度值很小（$|\nabla u|\to 0$），扩散系数$g$很大，扩散较强，可视为高斯平滑。对于目标边界，梯度值很大，则$g$接近0，扩散较弱，可保留边界细节。可将上式重写为：
$$\frac{\partial u}{\partial t}=\frac{\partial }{\partial x}\left\{g\left(\sqrt{(\frac{\partial u_k}{\partial x}{)}^2+(\frac{\partial u_k}{\partial y}{)}^2}\right)\frac{\partial u_k}{\partial x}\right\}+\frac{\partial }{\partial y}\left\{g\left(\sqrt{(\frac{\partial u_k}{\partial x}{)}^2+(\frac{\partial u_k}{\partial y}{)}^2}\right)\frac{\partial u_k}{\partial y}\right\}$$

其中$\frac{\partial u}{\partial x}$和$\frac{\partial u}{\partial y}$为特征图的水平和垂直梯度。特征图的离散小波变换（DWT）可表达为
u i = D W T ( u ) , i ∈ { u L L , u L H , u H L , u H H } u_i=DWT(u),i\in\{u_{LL},u_{LH},u_{HL},u_{HH}\} ui​=DWT(u),i∈{uLL​,uLH​,uHL​,uHH​}

其中$u_{LL}$为特征图的低频部分，$u_{LH},u_{HL},u_{HH}$分别为水平、垂直和对角方向的高频部分，主要包括边缘细节。

分别使用$u_{LH}$和$u_{HL}$近似$\frac{\partial u}{\partial x}$和$\frac{\partial u}{\partial y}$，并把扩散步长$\delta t$视为1，则离散形式的PMD公式为：
$$u_k=u_{k-1}+g{\left[(\sqrt{u_{LH}^2+u_{HL}^2})\cdot u_{LH}\right]}_{LH}+g{\left[(\sqrt{u_{LH}^2+u_{HL}^2})\cdot u_{HL}\right]}_{HL}$$

使用PMD增强特征图后，将输出送入ResNet块中。堆叠多个基于DWT的PMD块，组成基于DWT的PMD编码器。

2.2 视觉Mamba块

初始阶段，2D输入${\mathbb{R}}^{H\times W\times C}$被变换为拉长的patch $P_N$，尺寸为$M\times (N^2\cdot C)$，其中$N$为patch的大小，$M$为patch的数量。然后，$P_N$通过线性投影变为$D$维向量，并添加位置编码$E_{\text{pos}}\in {\mathbb{R}}^{M\times D}$，得到
$$X_0=[x^{1}W;x^{2}W;\cdots ;x^{M}W]+E_{\text{pos}}$$

其中$x^m$为$P_N$的第$m$个patch，$W\in {\mathbb{R}}^{(N^2\cdot C)\times D}$为可学习投影矩阵。第$l-1$层的输出$X_{l-1}$会通过第$l$层得到$X_l$：
$$X_l=Vim(X_{l-1})+X_{l-1}$$

2.3 损失函数

本文使用交叉熵损失，包含主要损失$L_{prim}$（融合+分割头）和3个辅助损失$L_{fcn}$（融合+FCN头），$L_{pmd}$（基于DWT的PMD+分割头），$L_{vim}$（视觉Mamba+分割头）。

3. 实验结果与分析

3.3 结果比较

与基于CNN和ViT的方法相比，P-Mamba能达到最优性能。此外，PMD的设计使得本文方法受到噪声的影响最小。

3.4 消融研究

将基于DWT的PMD块替换为只有边缘保留作用的Sobel算子，性能有所下降。这是因为前者还有噪声抑制作用。

将视觉Mamba替换为ViT，性能同样下降。

3.5 模型效率比较

实验表明，本文的P-Mamba比其余方法有更低的计算量和参数。PMD块并未带来过多的参数。