O-Mamba : 用于水下图像增强的 O 形状态空间模型

O-Mamba: O-shape State-Space Model for Underwater Image Enhancement

2408.12816 (arxiv.org)

Underwater image enhancement (UIE) face significant challenges due to complex underwater lighting conditions. Recently, mamba-based methods have achieved promising results in image enhancement tasks. However, these methods commonly rely on Vmamba, which focuses only on spatial information modeling and struggles to deal with the cross-color channel dependency problem in underwater images caused by the differential attenuation of light wavelengths, limiting the effective use of deep networks. In this paper, we propose a novel UIE framework called O-mamba. O-mamba employs an O-shaped dual-branch network to separately model spatial and cross-channel information, utilizing the efficient global receptive field of state-space models optimized for underwater images. To enhance information interaction between the two branches and effectively utilize multi-scale information, we design a Multi-scale Bimutual Promotion Module. This branch includes MS-MoE for fusing multi-scale information within branches, Mutual Promotion module for interaction between spatial and channel information across branches, and Cyclic Multi-scale optimization strategy to maximize the use of multi-scale information. Extensive experiments demonstrate that our method achieves state-of-the-art (SOTA) results.

由于复杂的水下光照条件，水下图像增强（UIE）面临着重大挑战。最近，基于Mamba的方法在图像增强任务中取得了显著成果。然而，这些方法通常依赖于Vmamba，它仅侧重于空间信息建模，并难以处理由于光波长差异衰减导致的水下图像跨颜色通道依赖性问题，从而限制了深度网络的有效利用。

本文提出了一种新颖的水下图像增强框架——O-mamba。O-mamba采用O形双分支网络分别建模空间和跨通道信息，利用针对水下图像优化的状态空间模型的有效全局感受野。为了增强两个分支之间的信息交互并有效利用多尺度信息，本文设计了一个多尺度双向促进模块。该模块包括MS-MoE（用于融合分支内的多尺度信息）、双向促进模块（用于跨分支的空间和通道信息交互）以及循环多尺度优化策略（以最大化利用多尺度信息）。大量实验表明，本文的方法达到了最先进（SOTA）水平。

Introduction

。。。。。。

为了应对复杂水下环境带来的挑战，传统上提出了基于水下图像物理特性的UIE方法。这些方法高度依赖于各种先验假设、物理模型和经典的、非数据驱动的图像处理技术。然而，由于这些基于物理的方法的表示能力有限以及水下场景的固有复杂性，很难使用物理模型为复杂的水下条件获得良好的增强效果

尽管基于CNN的方法取得了显著成果，但它们的感受野受限，因此无法有效捕获水下场景中的全局依赖。为了更好地捕获全局信息，Transformer基的方法被引入UIE领域。自注意力机制擅长建模长距离依赖并有效提取全局特征，但其在处理高分辨率图像时，计算复杂度与序列长度的平方成正比，从而带来了挑战。

最近，由于线性复杂度和对长序列的适用性，诸如Mamba之类的状态空间模型受到了越来越多的关注。这些模型通过有效平衡捕获长距离依赖和全局特征的需求，同时保持计算效率，提供了一种有吸引力的替代方案。现有的基于Mamba的方法在各种图像恢复领域取得了显著成果。这些方法通常依赖于标准的Mamba块，这些块使用具有四个空间方向的二维选择性扫描模块，从而限制了建模通道表示的能力。然而，在UIE（水下图像增强）任务中，由光波长差异衰减引起的独特颜色退化与跨通道信息密切相关，因此充分利用通道信息变得至关重要。这些现有模型缺乏对通道信息的足够关注，导致颜色校正能力不足。如图1所示，最近的SOTA（当前最佳）方法无法去除水下颜色干扰，表现出较弱的颜色校正能力。

本文提出了一种名为O-Mamba的新型框架用于水下图像增强，该框架采用O形双分支网络，充分利用由空间Mamba（SM）和通道Mamba（CM）块提供的有效全局感受野，分别建模空间和跨通道特征。为了在两个分支之间实现信息交互并有效利用多尺度信息，本文设计了一个多尺度双向促进（MSBMP）模块。

该模块由三个部分组成：

1）为了在每个分支内充分利用多尺度信息，本文引入了多尺度混合专家（MS-MoE），它使用多个Mamba专家来学习不同尺度的特征表示。

2）设计了相互促进（MP）模块，旨在实现空间和通道信息的融合。

3）此外，本文引入了循环多尺度（CMS）优化策略，以减轻同时优化多个尺度时的干扰。

通过利用O形双分支网络和MSBMP模块，O-Mamba实现了对空间和跨通道信息的全面关注，通过充分利用多尺度信息有效地恢复了水下图像。

Preliminaries

状态空间模型（SSMs），受连续线性时不变（LTI）系统的启发，为序列到序列的建模提供了一个强大的框架。这些模型使用一个隐含的潜在状态h(t) ∈ R^N来将一维函数或序列输入x(t) ∈ R映射到输出y(t) ∈ R。数学上，该系统可以由以下线性常微分方程（ODE）表示：

在此上下文中，N是隐藏状态的大小，A ∈ R^(N×N)、B ∈ R^(N×1)、C ∈ R^(1×N)是状态大小的参数，D ∈ R表示跳跃连接。

为了将方程（1）集成到深度学习算法中，通常使用零阶保持（ZOH）规则进行离散化。该规则涉及时间尺度参数∆，将连续参数A和B转换为它们的离散对应物A¯和B¯。定义如下：

其中，离散化参数定义为：

最近，Mamba的引入进一步提升了SSMs的能力，通过引入选择机制，Mamba能够有效地处理非常长的序列，同时利用并行扫描算法进行更快的并行计算，从而实现高效的训练。

Methodology

Overall Framework

如图2所示，本文的O-Mamba是一个O形网络，由两个主要分支组成：空间Mamba分支和通道Mamba分支。空间Mamba分支是一个类似UNet的网络，由堆叠的空间Mamba块组成。它通过在不同尺度上提取空间信息来增强图像的空间结构，从而改善图像的整体空间完整性。相比之下，通道Mamba分支由堆叠的通道Mamba块组成，专注于不同尺度上的跨通道退化信息。为了在两个分支之间实现信息交互并充分利用多尺度信息，本文设计了一个多尺度双向促进（MSBMP）模块。图2的中间部分表示MSBMP模块，它由用于分支内特征融合的多尺度混合专家（MS-MoE）、用于分支间交互的相互促进（MP）模块和一个循环多尺度（CMS）优化策略组成。

Spatial Mamba Branch

    空间Mamba块

空间Mamba块是在Transformer基本架构的基础上构建的。如图2所示，本文利用了一个四向选择性扫描，即空间选择性扫描模块（S-SSM），来替代自注意力机制，并在空间维度上全面捕获长期依赖性。简单地将注意力机制替换为SSM并不能达到最优效果。

本文设计了一个前馈混合专家（FF-MoE）来替代Transformer中的多层感知机（MLP），具体细节将在后文详细解释。SM块的输入由两部分组成：一部分是前一块的输出Y^(n-1)，另一部分是当前尺度的多尺度输入X_n，当前块的处理过程可以定义如下：

其中，LN代表LayerNorm，n ∈ {1, 2, 3}代表当前Mamba块的编号。

    空间选择性扫描模块

在空间信息建模方面，注意力机制往往需要在计算效率和有效的全局感受野之间取得平衡。

本文采用了Spatial-SSM，它以线性复杂度提供全局建模能力，以处理水下图像中的全局空间信息。如图所示，Spatial-SSM的输入X通过两个分支进行处理。在第一分支中，特征通道通过线性层扩展到λC，其中λ是预定义的通道扩展因子。随后是卷积层、SiLU激活函数、SSM层和LayerNorm。具体来说，这一分支可以表示如下：

与第一分支类似，第二分支也将特征通道扩展到λC，并应用SiLU激活函数：

随后，两个分支的特征通过Hadamard乘积进行聚合。将通道维度重新投影回C，以产生与输入形状相匹配的输出X_out：

其中，⊙代表Hadamard乘积。

    前馈混合专家

鉴于水下环境的多样性和复杂性，包括但不限于蓝色、绿色、低光和浑浊等场景，单个网络很难有效地建模如此广泛的水下图像。本文在前馈网络中引入了混合专家（MoE），如图 3 所示，使Mamba能够在不增加计算成本的情况下处理各种场景。在此设置中，FFMoE作为当前尺度的专家，通过门控网络动态选择最适合当前水下条件的专家。为了定义FF-MoE过程，本文从输入X开始，过程如下：

Channel Mamba Branch

    通道Mamba块

与空间Mamba块类似，通道Mamba块也采用了另一种SSM来替代自注意力机制，本文称之为通道选择性扫描模块（C-SSM）。C-SSM是一个跨通道的双向选择性扫描模块，专注于提取跨通道依赖性。与SM块一样，CM块的输入也由两部分组成：一部分是前一块的输出Y^(n-1)，另一部分是当前尺度的多尺度输入X_n，当前块的处理过程可以定义如下：

    通道选择性扫描模块

在通道信息建模方面，现有的Mamba架构主要从四个空间方向扫描图像，而没有考虑跨通道的依赖性。为了克服这一局限性，本文将2D-SSM扩展到通道选择性扫描模块中，以实现跨通道扫描。如图所示，Channel-SSM首先对输入X应用全局平均池化（GAP）操作以获得池化特征。然后，池化特征通过SiLU激活函数和SSM层进行处理。最后，使用sigmoid激活函数生成通道注意力，并将其与原始输入X相乘，以产生Channel-SSM的输出，定义如下：

Multi-scale Bi-mutual Promotion

    多尺度Mamba混合专家（MS-MoE）

与处理不同水下条件的FF-MoE相比，多尺度Mamba混合专家（MS-MoE）的作用是融合由不同Mamba块提取的多尺度特征。为了应对不同水下场景中不同尺度特征的复杂性，MS-MoE使用当前分支中的Mamba块作为其专家。以空间Mamba分支为例，假设编码器中S-SSM的输出为Y1、Y2和Y3。在接收到这些特征后，MS-MoE首先通过卷积层统一它们的维度，并将它们拼接为Y_cat以供后续处理。然后，将合并结果输入到多尺度Mamba专家模块以获得Y_s。处理过程如下：

..

其中，Expert代表S-SSM块。通道分支的结果Y_c与上述过程类似。

    相互促进模块

为了实现空间信息和通道信息之间的交互，本文引入了空间-通道分支MS-MoE相互促进机制。对于来自空间分支MS-MoE（Ys）和通道分支MS-MoE（Yc）的输出，本文通过逐元素相加的方式将它们组合起来，以获得融合特征F。然后，这个融合特征F通过三个不同的卷积MoE进行处理，以调整其维度与Unet的跳跃连接相同，从而产生Zi，其中i ∈ {1, 2, 3}。这些特征随后被添加到相应的跳跃连接上。具体过程如下：

其中，MoEi代表用于处理和调整的卷积MoE，Skipi代表跳跃连接，其中i ∈ {1, 2, 3}。

    循环多尺度优化策略

本文将S1, S_{1/2}, S_{1/4}表示为空间Mamba分支的输出，C1, C_{1/2}, C_{1/4}表示为通道Mamba分支的输出。相应的真实标签分别表示为G1, G_{1/2}, G_{1/4}。多尺度损失可以表示为：

其中i ∈ {1, 1/2, 1/4}。如算法 1 所示，本文在每次迭代中仅优化三个损失中的一个，以减少不同尺度之间的干扰，并在后续迭代中循环进行。这种方法被称为循环多尺度优化策略，它通过防止同时优化多个尺度损失时可能出现的潜在冲突，从而提高了模型的性能。

Datasets

本文使用UIEBD数据集和LSUI数据集来训练和评估本文的模型。UIEBD数据集包含890张带有相应标签的水下图像。本文使用800张图像进行训练，其余90张图像用于测试。LSUI数据集被随机划分为4604张训练图像和400张测试图像。此外，为了验证O-Mamba的泛化能力，本文还使用了参考基准EUVP，它包含515对测试样本，以及非参考基准U45和Challenge-60，它们分别包含45张和60张用于测试的水下图像。