本篇分享 NeurIPS2024 论文
ECMamba: Consolidating Selective State Space Model with Retinex Guidance for Efficient Multiple Exposure Correction
，ECMamba 结合 Retinex Theory 和 Mamba 实现高效曝光矫正！
论文链接：
https://arxiv.org/pdf/2410.21535
代码链接：
https://github.com/LowlevelAI/ECMamba
亮点直击
提出一种名为ECMamba的新型曝光校正模型，该模型采用双支路结构并且引入Retinex 信息作为指导ECMamba去分离光照信息和物体的内在反射特性。
为了使Mamba有效处理图像数据，我们将图像数据扫描视为一个“特征敏感型”问题，并提出了一个具有Retinex信息指导的特征感知的2D扫描机制
大量实验和消融研究证明了ECMamba在多重曝光校正和低光照图像增强任务上有优异表现。
在当前基于深度学习的曝光校正研究中，很少有模型完全将Retinex理论嵌入到其架构中。此外，如何平衡高性能与高效率仍然值得探索。
本文提出一种基于Retinex theory 和 Mamba的双通道高效曝光校正模型ECMamba。首先通过理论分析，本文将输入分别映射到与目标reflectance和 illumination map接近的两个过渡空间。
此外在Mamba架构中，为实现精细曝光校正，本文开发了一种基于可变形特征聚合的新型2D扫描策略。大量实验结果表明ECMamba的优异性能。
方法
本文提出的方法的概览如图2所示，可见本文提出的ECMamba是基于Mamba和Retinex理论设计的。
首先，通过分析Retinex 理论，作者提出了双支路的曝光校正架构，每个支路分别旨在修复reflectance 和illumination map。
其次，在每个支路中，本文利用ECMM模块实现对reflectance 或者illumination map的精确恢复。
此外，为了进一步增强效率和效果，本文开发了一种新的2D SSM 层以及一种基于可变形特征聚合的新型2D扫描策略。
基于Retinex 理论的双支路曝光校正架构
Retinex 理论可以表示为
, 其中
表示 Hadamard 乘积,
是一个理想的无退化图像，
和
分别代表 reflectance 和 illumination map。
然而，一个在非理想照明条件下（过曝光或欠曝光场景）拍摄的低质量图像
不可避免地会受到噪声、颜色失真和对比度不足的影响。
因此, 可以分别对
和
引入了扰动（ (
和
来模拟这些退化图像,公式（1）如下：
当前一些基于 Retinex 的方法
将反射分量
视为最终的增强结果，从而忽略了公式（1）中的最后三项，只专注于使用网络
去模拟此映射：
。
但是，这些模型在多重曝光校正任务中只能达到次优性能，因为我们很难获得准确映射的困难，特别是多个欠曝光 （UE）和过曝光（OE）输入对应一个正常曝光（NE）图像时。
因此, 为了获得令人满意的结果, 本文选择同时恢复
和
。
具体来说, 将公式 (1) 的两边分别与
和
进行逐元素相乘公式 (2):
其中
和
是满足
和
的矩阵, 并且本文使用 Retinex 预测器
去估计这两个矩阵。
和
分别表示
和
'中剩余的退化, 因此可以采用深度学习网络来获得良好曝光的图片公式 (3) :
其中
和
是用来预测
和
中负退化的网络，
是从
中得到的 Retinex 指导信息。
如图2 所示, Retinex 估计器
将
及其沿通道维度的均值矩阵作为输入 (为清晰起见，图2中没有显示该矩阵）。
首先使用
卷积和一个
的 depthwise 卷积来提取特征，然后分别通过一个
卷积生成
、
和
。
更为关键的是，
和
被输入到
和
进行进一步的复原。除了优化
使其接近
，本文的训练目标还包括了对
和
的限制。
讨论：
（i）许多基于Retinex的方法 [37]旨在学习输入与反射图像和照明图之间的映射，然后通过 Hadamard 乘积获得最终结果。然而，这种策略不适用于多重曝光校正任务。图1(a)展示了过曝光（OE）和欠曝光（UE）图像的分布很复杂且距离正常曝光图像很远。这种复杂的分布使得准确建议从输入到输出的映射是极其困难的。然而，通过仔细分析Retinex理论，本文构建了一个中间空间如图1(b)所示。该空间显著减少了与优化目标的距离并有利于后续的微调恢复过程。
(ii) 一些方法
将
视为最终增强结果, 这与 Retinex 理论的原始解释偏离。因此, 本文采用了一个双分支框架, 使用独立的深度学习网络去分别重建
和
。本文在消融研究中讨论了该框架中各个部分的重要性。
ECMM 以及RMB模块
为了使提出的曝光校正架构可以应用在资源有限的设备上并且实现高性能，本文提出了一个新颖的由Retinex信息指导的ECMM模块，它成功继承了Mamba的强大建模能力。
如图2所示, ECMM采用了双尺度U-Net架构。在编码过程中, 输入
首先通过一个
的卷积和一个RMB模块以获得初始特征
。然后, 通过一个步长为 2 的
卷积实现下采样, 并且将下采样后的特征输入到另一个RMB模块中, 以获得中间特征
。在解码阶段，
首先被一个步长为 2 的
反卷积上采样得到
。为了降低下采样过程中造成的信息丢失, 本文采用了一种自适应混合特征融合策略[46] 将编码信息传递到解码阶段公式（4）:
其中
是一个可学习的参数,
代表 sigmoid 方程。融合后的特征
依次进入
和卷积层，按照公式（3）通过残差叠加的方式来获得修复后的反射分量
。
作为 ECMM 模块的核心, RMB 块 采用了与 Transformer 相似的结构。然而，需要大量计算的自注意力和交叉注意力机制影响了 Transformer 的效率，使其无法应用于实时或资源受限的环境中。为此，本文去除了注意力过程，并引入了一种新型的由 Retinex 信息指导的 2D SSM 层来捕捉远距离依赖并促进动态特征聚合。因此, RMB 块可以这样表示公式（5）：
其中, LN 表示 LayerNorm,
是由 Retinex 估计器
提取的 Retinex 指导信息。此外，受到 ConvNext
的启发，本文移除了门控机制和深度卷积，引入了一个效率更高的前向反馈层（EFF）。该层遵循
卷积
卷积的流程, 其操作类似于 Transformers 中的 MLPs 但只需要更少的参数。
Retinex-SS2D 层和基于可变形特征聚合的新型 2D 扫描策略
图 3 详细展示了 Retinex-SS2D 层的结构。首先通过线性运算、depth-wise 卷积、逐元素相乘及 SiLU操作对输入特征
和
进行融合。然后, 融合后的特征
被输入到本文提出的 FA-SS2D 机制中，该机制旨在捕捉动态的远距离依赖并实现自适应空间聚合。此外, 还运用了门控信号
和线性运算来得到最终的聚合特征
。
基于可变形特征聚合的新型2D扫描策略
选择性状态空间模型（S6）在处理涉及时间序列的自然语言处理等任务时表现出色。但是，将其应用于二维图像时面临不小挑战。为了更好地处理二维图像中的空间信息，当前研究提出多种扫描策略将图像块转换为一维序列。举例来说，文献[27]提出了一种交叉扫描策略。该策略沿四个不同的扫描路径生成四个序列，每个序列都单独进行S6处理。然而，这种策略显著增加了计算负担，这与S6本身高效、低计算需求的特性相悖。此外，这些方法只是简单地在不同方向进行图像扫描，导致在某些序列中局部纹理和全局结构发生分离。这种分离在一定程度上影响了S6模型对图像的建模效果。
现行扫描方法的局限性促使作者重新考虑如何更有效地应用S6于二维图像。在S6 过程中， 对于一维序列中的每一个元素，其输出y(t)取决于它的输入x(t)和之前的输入{x(1), x(2), ..., x(t − 1)}。这种机制要求从二维图像转化而来的一维序列必须满足以下两个条件才能确保优良的表现：(1) 序列应在起始部分优先处理最关键的特征区域，而将次要信息置于末尾。(2) 空间相邻的特征应当被紧密地排序，以避免序列中出现显著的间隙。然而，现有的二维扫描策略未能满足这两个要求，这促使作者提出新的解决方案来弥补这一差距。
基于这些观察，本文引入了一种高效的FA-SS2D机制如图3所示。本文首先开发了一种由Retinex信息指导的可变形特征聚合操作。具体而言，采用可变形卷积（DCN）[48, 9]来捕捉融合特征
的动态远程依赖。
例如，当DCN应用于图3中
内的红框时，此时激活区域是一个不规则的并以蓝色标记。更为重要的是，当红框
在上滑动时，这个不规则的区域会发生变化。最终可以得到每个小框的平均激活频率，进而得到如图3中
所示的激活响应图，其中激活频率高的区域代表重要的特征。
具体而言，对于欠曝光图像的相对较亮区域或过曝光图像的相对正常曝光区域，这些区域包含重要特征从而具有比较大的激活响应。
基于获得的激活响应图，本文提出了一种新的扫描策略，与“方向敏感”的扫描方法[27]不同，本文提出的策略根据不同区域的激活频率进行排序，并将频率高的标记放在序列的前端，从而使得生成的序列有效地满足Mamba的要求，显著地提升了其对图像数据建模的能力。
损失函数
在本文中, 本文选择了一种端对端训练策略去同时优化
、
和
。本文的最终训练目标是使
接近
，同时我们也在
、
和
上施加了若干约束以实现稳定训练。因此, 完整的损失函数如下所示公式（6）：
其中
