上海大学提出Retinex-Mamba | 低光照下细节保持依旧完美，超越ViT家族！

在低光照图像增强领域，传统的Retinex方法与先进的深度学习方法（如Retinexformer）均显示出其独特的优势和局限性。传统Retinex方法旨在模仿人眼对亮度和颜色的感知，将图像分解为光照和反射成分，但在低光照条件下，它们在噪声管理和细节保持方面存在困难。

Retinexformer通过传统的自注意力机制增强光照估计，但在解释性和增强效果方面面临挑战。

为了克服这些局限性，本文提出了Retinex Mamba架构。RetinexMamba不仅捕捉到传统Retinex方法的物理直观性，还整合了Retinexformer的深度学习框架，利用状态空间模型（SSMs）的计算效率提高处理速度。

该架构特色创新的照明估计器和损伤恢复机制在增强过程中保持图像质量。

此外，RetinexMamba用融合注意力机制替换了Retinexformer中的IG-MSA（光照引导的多头注意力），提高了模型的解释性。在LOL数据集上的实验评估显示，RetinexMamba在基于Retinex理论的现有深度学习方法中，在定量和定性指标上均表现出色，证实了其在增强低光照图像方面的有效性和优越性。

代码可在https://github.com/YhuoyuH/RetinexMamba获取。

1 Introduction

低光照增强是指通过图像处理技术改善在光线不足或昏暗环境下捕捉到的图像质量和视觉外观的过程。这一领域在计算机视觉和图像处理中具有重要意义，因为低光照条件下的图像常常存在诸如昏暗、模糊和不清晰的细节等问题，这些问题影响了图像的质量和可用性。

传统的技术，如直方图均衡化和伽马校正，虽然是基础且重要的，但往往在处理复杂的照明动态和维护增强图像的自然性方面存在不足。受到人类视觉系统的启发，Retinex理论[21]为分离图像的照明和反射成分建立了一个概念框架，为处理多种多样且具有挑战性的低光照环境的能力更强的增强策略铺平了道路。

神经网络最近的发展，特别是在卷积神经网络（CNN）和Transformer模型的应用方面，为低光照图像增强设立了新的基准。CNNs通过有效地捕捉图像中的空间信息和局部特征，在低光照图像增强方面表现出强大的能力。然而，CNN在建模长距离依赖方面可能存在局限性，导致对整体图像信息的考虑不足，在有效解决低光照图像处理中的噪声放大和细节保留等问题时也面临挑战。另一方面，Transformer模型通过自注意力机制实现全局感知和长距离依赖的建模，有助于在低光照图像处理过程中更准确地恢复图像的细节和结构。然而，Transformer模型也面临诸如高计算复杂度和大参数规模等挑战，导致在实际应用中可能存在推理速度慢和资源消耗高的问题。

图1：上面的图像展示了注意机制与Mamba中的2D选择性扫描的视觉对比。在Mamba的2D选择性扫描中，扫描同时从图像的所有侧面开始，而注意机制从目标视图到全局视图分别计算注意力得分。

在这种背景下，作者的工作引入了 RetinexMamba 架构。最初遵循Retinexformer的方法，作者将整体架构分为照明估计器和损伤修复器。照明估计器用于初步照亮图像，而损伤修复器用于消除在照明过程中产生的放大伪影和噪声，以及色彩失真和过度曝光。为了解决Transformer中的高计算复杂性和Retinexformer中注意力机制解释性不足的问题，作者的损伤修复器的基本单元是Illumination Fusion State Space Model（IFSSM）。该模型使用2D Selective Scanning（SS2D）作为主干网络以实现线性计算效率，并采用Illumination Fusion Attention（IFA）替代Retinexformer中的Illumination-Guided Multi-head Self-Attention（IG-MSA）以增强注意力机制的解释性。

通过全面的实验设置，作者提供了定性和定量的结果，证明了作者的模型在LOL数据集等标准基准上的优越性。根据表1中的数据，采用这种方法，作者已经在LOL数据集上超过了基于Retinex理论的深度学习方法的状态-of-the-art（SOTA）。

作者的贡献可以总结如下：

1. 首次将Mamba用于低光照增强，使用SS2D替代Transformer捕捉长距离依赖。
2. 作者提出了一种融合模块，更好地实现了与Retinex理论一致的照明特征的嵌入。
3. 大量的定性和定量实验证明，作者的方法超过了所有之前基于Retinex理论的深度学习方法。

2 Related Work

In recent years, significant progress has been made in the field of AI, particularly in the realm of computer vision. This section reviews the related work that is most relevant to our research.

2.1.1 Distribution Mapping Method.

在早期对低光照图像增强的研究中，最直观的方法之一是将低光照输入的分布进行映射，以放大较小的值（这些值看起来较暗）。这种方法代表性的技术包括直方图均衡化和基于S曲线的方法，如伽马校正。然而，现有的基于分布映射的方法常常由于在分布映射过程中缺乏对语义信息的识别和利用，而导致颜色失真和其他影响增强结果视觉质量的艺术效果。

2.1.2 Traditional Model Method.

Retinex理论[21]为增强弱光图像的过程提供了直观的物理解释。该理论假设通过从低光照输入中移除光照分量，可以得到所需的正常图像（即反射率图）。Jobson等人[19, 20]基于Retinex模型进行了探索性研究。随着研究的进展，很明显，使用Retinex方法实现亮度增强的关键在于估计光照层。这些方法依赖于手工制作的先验知识，并且通常需要仔细的参数调整。不准确的先验知识或正则化可能导致增强图像中出现伪影和颜色偏差，显示出较差的泛化能力和耗时的优化过程。此外，这些研究常常忽略了噪声的存在，这可能导致在增强图像中保留或放大噪声。

基于深度学习的低光照图像增强方法起源于2017年[26]，并随后成为该领域的主导方法。在传统的Retinex理论[21]作为模型架构基石的基础上，一系列工作沿着这些思路被设计出来[46, 24]。卷积神经网络（CNN）已被广泛应用于低光照图像增强。例如，Wei等人及其后续工作将Retinex分解与深度学习相结合。然而，这些基于CNN的方法在捕捉不同区域间的长距离依赖方面存在限制。Star将 Transformer 架构引入低光照增强领域，解决了捕捉长距离依赖的问题。同时，Retinexformer[2]将Retinex理论与一阶段 Transformer 的设计相结合，进一步改进和优化了这种方法。尽管如此，由于自注意力机制， Transformer 模型在处理长序列时带来了巨大的计算负担和复杂性。

State Space models

近期，状态空间模型（SSMs）越来越被认为是一个有前景的研究方向。在中提出了一种结构化的状态空间序列模型，称为S4，作为CNN或Transformer的一种新颖替代方案，用于建模长距离依赖关系。后续的发展见证了各种结构化状态空间模型的涌现，这些模型具有复杂的对角线、多输入多输出支持[41]、对角分解和低秩操作[16]，增强了它们的表达能力。现代SSMs，如Mamba[10]，不仅建立了长距离依赖关系，而且对于输入大小而言，其复杂性呈现线性。基于SSM架构的模型在各个领域都引起了广泛的研究兴趣[33]。Mamba[10]引入的选择性扫描机制与视觉领域中普遍的基础模型的性能相匹配。Vision Mamba[60]表明，纯SSM模型可以作为一个通用的视觉 Backbone 网络。提供了实证验证，证明了在医学图像分割任务中具有潜力，而在低级视觉任务中，应用如也展示了良好的成果。受到这一研究的启发，作者的工作利用Mamba对长序列进行线性分析的能力，处理与Retinex理论[21]融合的特征。对低光照图像的增强效果证实了基于Mamba模型的在低光照图像增强领域的潜力。

3 Method

图2展示了作者方法的综合结构。如图2所示，作者的RetinexMamba包括一个光照估计器（a）和一个损伤恢复器（b）。光照估计器（IE）受到传统Retinex理论的启发。损伤恢复器的设计基于光照融合视觉Mamba（IFVM）。如图2(b)所示，IFVM的核心组件是光照融合状态空间模型（IFSSM），其特点包括层归一化（LN）、光照融合注意力机制（IFA）、2D选择性扫描（SS2D）和前馈网络（FFN）。IFA的具体细节在图2(c)中详细描述。

3.3.1 Retinex-based Framework

传统的Retinex图像增强算法模拟了人类对亮度和颜色的视觉感知。它将图像 分解为光照分量 和反射分量 。这个结论可以用以下公式表示：

其中 表示逐元素乘法。反射分量 由物体的固有属性决定，而光照分量 表示光照条件。然而，在公式表达下，传统Retinex算法没有考虑到不平衡光照分布或低光照条件下暗场景产生的噪声和伪影，这种质量的损失在图像增强过程中被进一步放大。因此，受到Retinex算法的启发，作者采用了[2]提出的扰动建模，为原始公式中的光照分量和反射分量 和 引入了扰动项，如下方程所示：

其中 和 分别表示反射分量和光照分量的扰动项。简化后，作者可以将光照图像 表示如下：

其中 表示光照映射，通过卷积进行特征提取获得， 表示之前提到的所有损失。因此，作者的RetinexMamba可以表示为：

(6)

其中IE表示光照估计器，IFVM表示损伤修复器。IE接收低光照图像 和光照先验 作为输入，输出光照图像 和光照特征图 。 通过计算图像每个通道的平均值获得，用于评估图像的整体亮度或光照水平；因此作者使用 作为光照先验，为图像提供光照信息。然后，这两个结果被送入损伤修复器（IFVM）以修复图像光照过程中放大的质量损失，并最终生成修复结果 。

3.2.2 Illumination Estimator.

图2(a)展示了光照估计器(IE)的结构。作者将低光照原始图像 与通过计算得到的先验光照 合并，并增加通道维度作为输入。随后通过三个卷积层提取特征。第一个 卷积将之前合并的输入融合在一起，即将在低光照图像中应用先验光照融合。第二个深度可分离 卷积对输入进行上采样，进一步提取特征以生成光照特征图 ，其中特征维度 设置为40。最后，另一个 用于下采样以恢复3通道光照映射 ，然后将其与低光照图像 逐元素相乘得到光照后的图像 。

3.2.3 Illumination-Fused Vison Mamba.

损伤恢复器（IFVM）的结构如图2(b)所示，它基于光照融合视觉曼巴构建了一个编码器和解码器。编码器代表下采样过程，而解码器代表上采样过程。上采样和下采样过程是对称的，并分为两个 Level 。首先，从光照估计器IE获得的照亮图像 通过一个 （步长=2）下采样以匹配光照特征图 的尺寸，便于后续操作。接下来，作者执行下采样以降低计算复杂度并提取深层特征。下采样过程分为两个 Level ，每个 Level 包含一个光照融合状态空间模型（IFSSM）和一个步长为2，核大小为 的卷积层。在每次卷积层之后，图像的宽度和高度减半，而特征维度加倍。因此，在两个 Level 的下采样之后，最深层特征维度应为 。在提取图像特征之后，作者需要执行上采样以恢复图像。与下采样类似，上采样也分为两个 Level ，每个 Level 包含一个 （步长=2）和一个 ，以及一个光照融合状态空间模型（IFSSM）。在每次反卷积层之后，图像的宽度和高度加倍，而特征维度减半。反卷积层的输出随后与相应 Level 的下采样光照融合状态空间模型（IFSSM）的输出连接，以减轻下采样过程中图像信息的损失。最后，对图像应用一个 （步长=2）以降低维度并将其恢复为具有三个通道的RGB格式。通过将恢复的图像和 进行残差连接，得到增强的图像 。

Illumination-Fused State Space Models.

在低光照增强研究中，卷积神经网络（CNNs）在处理整体图像信息方面存在限制。由于计算需求高， Transformer （Transformers）可能会影响实际应用的效率。为了解决这个问题，作者设计的照明融合状态空间模型（IFSSM）包括了融合注意力（IFA），一个SS2D模块，一个LN层，一个前馈网络（FFN）以及一个卷积层，以匹配照明特征图与输入的维度，如图3所示。

3.2.1 Illumination-Fused Attention.

如图3所示，由IE生成的光照特征图 与经过特征提取后的增强图像一起输入到IFA中。在[2]中，将光照特征图 作为一个token，经过转换并与来自 的 (值)相乘以计算注意力得分。然而，这种方法导致 (键和值)并非来自同一输入，违反了Transformer中确保所有信息处理都集中在输入数据内容上的原则，而不是与当前任务无关的外部或无关信息，破坏了数据一致性和对齐，额外的乘法操作也增加了参数需求和计算复杂性。因此，作者的IFA采用了交叉注意力机制，首先通过改变维度将 和输入 调整为适合多头注意力的形式：

其中 , , 是头的数量， 表示每个头的维度， = ， 表示输入特征的维度。作者注意到，使用光照先验图作为 Query 向量可以使模型更具体地处理图像中的暗区。因此，模型的注意力机制可以专注于需要增强的低光区域，而不是统一处理整个图像。因此，作者将 视为 ，将输入 视为 ，以融合两个输入特征，使 指导 的自注意力计算。

其中 ， ， 是由卷积层构建的可学习参数矩阵。因此，每个头的自注意力可以用以下公式表示：

其中 是一个作为缩放因子的可学习参数，用于调整注意力得分，从而控制注意力权重的锐度， 个头随后重新调整为标准图像格式（ ），并通过卷积层进行聚合，得到与原始匹配维度的输出。

2D-Selective Scan.

图3：作者的光照融合状态空间模型（IFSSM）通过融合块集成光照特征和输入向量 ，并使用线性的2D选择性扫描模型（SS2D）进行特征提取。在IFA中，作者将光照特征视为 ，输入向量视为 以计算注意力分数。

受到[10]中SSM模型的启发，作者采用了将SSM集成到视觉任务中的方法，如[47]所示。在[47]中的SS2D模块包括扫描扩展、S6块和扫描合并操作。对于作者所输入的处理过的图像，它们首先经历扫描扩展操作，如图1所示，图像从四个角展开。然后，图像被展平，即将高度（H）和宽度（W）合并为一个标记长度（L）。来自扫描的每个序列随后输入到S6模块中进行特征提取。S6的计算公式可以表示为：

其中 是输入变量， 是输出， ， 和 都是可学习的参数。之后，来自四个方向提取的特征输出 被求和并合并，合并输出的维度被重新调整以匹配输入大小。此外，为了提取更深层次的潜在特征，作者设置SS2D中的隐藏层数量随着IFSSM的每个 Level 增加。作者默认SS2D中的隐藏层数量 为16，每个采样 Level 翻倍层数，因此在最深采样时达到64层。这个设置允许从集成光照特征的向量中逐步提取更深层次的特征。

4 Experiment

Datasets and Implementation details

4.1.1 LOL

LOL数据集分为两个版本，v1和v2。在LOL_v1中，训练数据与测试数据对的比率为485:15。每对数据包括一个输入的低光照图像和一个目标参考图像。LOL_v2数据集进一步划分为LOLv2_real和LOLv2_synthetic。LOLv2_real和LOLv2_synthetic中的训练数据与测试数据对的比例分别为689:100和900:100。数据对的分布与LOL_v1中的相同。

4.1.2 Implementation Details.

作者在Pytorch中实现了作者的RetinexMamba模型，并在一台装有A10和V100 GPU的PC上，在Linux系统下（CUDA 11.7，Python 3.8，Pytorch 1.13）进行了训练和测试。作者将图像的分辨率设置为128x128。LOL_v1和LOLv2-synthetic的批处理大小设置为8，LOLv2_real设置为4。作者使用了标准的增强方法，如随机旋转和翻转，来增强训练数据。为了最小化损失，作者采用了Adam优化器，动量项 设置为0.9，RMSprop控制参数 设置为0.999，旨在最小化增强图像与 GT 之间的平均绝对误差（MAE）。此外，还使用了余弦退火调度来防止损失陷入局部最小值。

Low-light Image Enhancement 2024-05-07-03-53-13

作者在表1中将在监督和无监督领域内的多种最新方法（SOTA）与作者的方法进行了比较。用于比较的数据集是来自LOLv1的合成数据以及来自LOLv2的真实数据。表中的所有数据都是在相同条件下获得的，使用公开可用的代码进行训练和测试，或者直接源自原始论文。结果表明，在PSNR和RMSE方面，作者的方法优于上述SOTA，而SSIM略低于Retinexformer。

4.2.2 Quantitative Results.

作者的方法有效减少了颜色失真并增强了光照效果。

作者用于比较的指标包括PSNR、SSIM和RMSE。更高的PSNR表明图像增强效果更好，而更高的SSIM表明结果中保留了更多的高频细节和结构。较低的RMSE值意味着预测模型的性能更好，因为它表示误差较小。与基准方法和最佳现有技术方法[2]相比，作者的方法在LOL_v1和LOLv2_real数据集上分别使PSNR提高了0.093和0.77。在LOLv2_real数据集上，RMSE降低了0.39，这是作者所期望的，因为值越低越好。

4.2.2 Qualitative Results.

定性比较RetinexMamba与其他SOTA算法的结果显示在图4和图5中。请放大以获得更好的视觉清晰度。图3比较了LOLv1数据集，而图5比较了LOLv2_real数据集。如图3所示，先前的方法表现出噪声放大，如在Retinex-Net中，以及欠曝光的情况，如KinD和过曝光的情况，如Uretiinex-Net。同样，在图5的合成数据集中，KinD在顶部空调区域显示过曝光，底部楼梯区域显示欠曝光，而ZeroRCE则显示出大量的噪声和伪影。同样，在Retinexformer中，图4顶部保龄球周围出现欠曝光，底部体育场出现颜色失真。相比之下，作者的RetinexMamba有效地控制了曝光强度，减少了颜色失真，并最小化了噪声。

与所有消融实验结果相比，作者的消融设置取得了最高的PSNR和SSIM。"作者的FixedHS"展示了在SS2D模型中具有固定数量的隐藏层在特征提取不足和无法捕获长序列方面的缺点。"作者的NoFB"使用直接逐元素乘法来融合光照特征，缺乏逻辑解释。同时，"作者的NoSS2D"和"作者的IGMAa"分别突显了仅使用Transformer架构和来自[2]的注意力计算可解释性差的问题。

5 Conclusion

在本文中，作者介绍了RetinexMamba架构，该架构基于Retinexformer和Mamba以增强低光照图像。最初，在Retinexformer模型的基础上，作者将它划分为一个光照估计器和一个损伤恢复器，并且受到VMDA的启发，作者融入了SS2D模型以解决Transformer中视觉数据的固有位置敏感性。此外，作者用更可解释的融合注意力模块替换了IG-MSA模块，以融合光照特征和输入向量。广泛的定量和定性实验表明，作者的RetinexMamba在LOL数据集上超越了当前的最先进水平。尽管SS2D的计算复杂度降低了，但参数的总数有所增加，消耗了更多的计算资源。因此，作者未来的工作将集中在减少总参数数量同时保持计算复杂度。

参考

[1].Retinex Mamba: Retinex-based Mamba for Low-light Image Enhancement.