PixMamba- Leveraging State Space Models in a Dual-Level Architecture for Underwater Image Enhancement

发表于 更新于 分类于

论文地址：arxiv

摘要

背景

水下图像增强（UIE）对海洋的研究与探索很重要，但是会受到复杂的颜色失真与严重的模糊的阻碍。

问题

虽然已经在深度学习方面有了成果，但是那些方法在高计算成本与不足的全局建模方面依然存在，从而导致局部区域调整不足或过度调整。

改进

提出 PixMamba，通过状态空间模型（SSMs）进行高效的全局依赖建模来克服那些挑战。

正文

背景

利用卷积神经网络（CNN）的方法由于可以端到端的学习视觉表示而被广泛用于 UIE，然而这种方法也有局限性：小的感受野限制了长距离像素依赖关系的建模，而固定的卷积核无法适应各种水下场景中的图像。

Transformer 模型可以克服 CNN 的局限性，然而，Transformer 在序列上的二次复杂度造成了性能的问题。

而状态空间模型及其改进版本 （Mamba 与 Mamba-2）通过专门设计的结构化重新参数化来建模非常长的依赖关系，保持了良好的性能。

改进

PixMamba 是一种利用状态空间模型的线性复杂性与长距离建模能力的新方法。由两部分组成：Mamba 网络（EMNet）与 PixMamba 网络（PixNet）。EMNet 结合了高效 Mamba 块（EMB）用于高效的补丁(patch)级特征提取与依赖关系建模，以及 Mamba 上采样块（MUB）用于细节保留的上采样。PixNet 是一种并重像素级网络，可以在像素级处理整个图像，从而捕捉全局视角下的详细特征。

EMNet 在小块中提取丰富的局部特征以重建高质量图像（保留了复杂的纹理和结构），而 PixNet 则在整个图像的像素级高效建模交互和依赖关系（确保了全局一致性、清晰度和自然的视觉效果）。通过这些互补的层次的结合使得 PixMamba 可以同时增强微观细节和宏观图像质量。

贡献

  • 引入了一种高效且详细的图像修复新方法，通过无缝集成本地 patch 级处理的全局像素级处理，PixMamba 提供了精细且增强的高质量水下图像。
  • EMNet 结合了 Efficient Mamba Block（EMB）和 Mamba Upsampling Block（MUB），这种组合显著提高了修复图像的质量并增强了整体处理效率。
  • 通过 PixNet 与 EMNet 的强大 patch 级处理的协同，使得 PixMamba 在水下图像自得过程中实现了更精细和增强的效果。

预备知识

结构化状态空间模型（S 4）是一种新的序列模型，该栈受连续系统的启发，主要通过隐藏状态 $h(t)$ 将一维序列 $x(t)$ 映射为输出序列 $y(t)$，连续系统可以用线性常微分议程 （ODE）表示：
$$
h’(t) = Ah(t)+Bx(t)
$$
$$
y(t) = Ch(t)+Dx(t)
$$
其中，$h(t)$ 为隐藏状态，$A$，$B$，$C$ 是当状态大小等于 $N$ 时的参数。$D$ 表示跳跃连接。

之后对以上两个方程离散化，使用零阶保持（ZOH）规则可以获得，离散化可以定义如下：
$$
h’(t)=\overline{A}h_{t-1}+\overline{B}x_{t}
$$
$$
y(t)=Ch_{t}+Dx_{t}
$$
$$
\overline{A}=e^{\Delta A}
$$
$$
\overline{B} = (\Delta A)^{-1}(e^{\Delta A}-I)
$$
其中， $\Delta A$ 是时间缩放参数。

整体架构

PixMamba 的整体架构依赖于 EMNet，并集成了一个并行的 PixNet。

EMNet 的工程流程：

  1. 输入一个要处理的水下图像
  2. 图像通过 PatchEmbed 进行编码，生成特征特征图。
  3. 经过三阶段的 EMB 编码，每个阶段后进行下采样，逐渐减少特征图的尺寸
  4. 通过 MUB 和 EMB 进行三阶段解码，恢复图像特征，并最终扩展和投影生成空间丰富的块级特征图（IFD）

PixNet 的工作流程：

  1. 输入一个要处理的水下图像（与 EMNet 的相同）
  2. 图像逐像素嵌入特征，并通过块级位置嵌入 （BPE）丰富局部信息
  3. 经过 L 个阶段的 Mamba Block 表层增强图像的深层特征
  4. 生成像素级详细特征图（IFP）

最终的图像生成：

  1. 将 IFD 与 IFP 结合，生成最终增强的水下图像。

高效 Mamba 模块（EMB）

该模块使用 Efficient SS2D（ESS2D）处理图像特征 patch。这个模块是 VMamba 引入的一个计算效率更高的变体。该模块在每个扫描以独立处理特征后，使用窨与通道注意力模块结合撮的信息并去除通道冗余。

该注意力模块通过在通道与空间两个维度上进行调整来优化特征表示。它包含两个分支：一个捕捉广泛特征表示并在通道间分配权重的通道注意力分支，一个评估特征中单个 token 重要性的空间注意力分支，从而实现详细和重要性加权的特征提取。最后，在基于注意力的过滤之后，来自不同扫描分支的特征会被合并，以创建最终的深度特征表示。

Mamba 上采样模块

该模块可以克服细节丢失或噪声引入的问题。该模块在上采样过程之前结合了 SSM 机制。SSM 可以有效捕捉通道和标记之间的依赖关系。通过在上采样阶段之间整合 SSM，可以智能的调整需要上采样的特征，有选择地维护和优化保留特征图中最相关的信息。因此，该模块可以显著提高上采样过程中的细节保留，从而增强恢复图像的整体性能与质量。

这个过程可以定义如下：

$$
I^{s-1}{D} = Norm(TransposeConv2D(EMB(I^s_D W)))
$$
其中，$W$ 为一个可以学习的投影矩阵，$I{D}^{s}$ 是 EMNet 在第 s 阶段解码的特征。

PixMamba Net

PixNet 专门设计用于像素级别的处理，这会产生细节更清晰，噪声更不的图像，其核心是 Mamba Block，其中在像素级别使用 SSM。通过 SSM，可以利用所有的像素，还可以获得更细粒度的特征与全局一致性。

为了让 PixNet 可以同时具有局部空间信息与全局像素信息，使用了块级学习位置嵌入 （BPE）。空间信息对于 SSM 建模很重要。BPE 使用块级设计，将位置嵌入分割成 $\frac{HW}{B^2}$ 并双线性采样到输入序列大小。之后，BPE 在 PixNet 处理之前会被添加到像素级序列特征中，来提供空间信息。

上述过程定义如下：

$$
\begin{aligned}
\text{PE} &= \textbf{Upsample}(\text{BPE}) \
I^0_P &= [I^0 W ; I^1 W ; \ldots ; I^{HW} W ] + \text{PE} \
I^l_P &= \textbf{MambaBlock}(I^{l-1}P) + I^{l-1}P \
I{\text{FP}} &= \textbf{Project}(I^L_P) \quad (11)
\end{aligned}
$$
$I{i}$ 为输入图像 $I$ 的第 i 个像素。

性能评估

使用以下 5 个评估指标：

  1. 均方误差
  2. 峰值信噪比
  3. 结构相似性指数
  4. 水下图像质量测量（UIQM）
  5. 水下彩色图像质量评估（UCIQE）

UIQM 与 UCIQE 无需参考图像即可评估图像质量。

数据的对比

PixMamba 在图像质量、颜色增强和信噪比等方面均有显著提升

消融研究

通过对每个组件的逐步添加，可以发现模型的准确率也在逐步上升，因此可以证明每个模块都是必须的。

知识点

零阶保持

零阶保持的基本思想是将连续时间信号在每个采样周期内保持恒定值，直到下一个采样时刻。这意味着在每个采样点，信号的值会保持不变，直到下一个采样点到来。

如果有一个连续时间信号 $x(t)$，如果希望将其离散化，则零阶保持的具体步骤如下：

  1. 采样：在特定的时间间隔 T 处对连续信号 $x(t)$ 进行采样，得到离散时间信号 $x[n]=x(nT)$，其中，$n$ 为采样点的索引。
  2. 保持：在每个采样周期内，信号的值保持恒定，直到下一个采样时刻。因此，在时间区间 [nT, (n+1)T] 内，信号的值为 $x(nT)$ 不变。