[论文解读] MambaUIE&SR: Unraveling the Ocean's Secrets with Only 2.8 GFLOPs

Zhihao Chen, Yiyuan Ge|arXiv (Cornell University)|Apr 22, 2024

Oceanographic and Atmospheric ProcessesEarth and Planetary Sciences被引用 3

一句话总结

该论文提出MambaUIE&SR，一种基于状态空间模型（SSM）架构Mamba的新型水下图像增强模型，在仅2.8 GFLOPs的计算量下实现了最先进（SOTA）性能。通过集成视觉状态空间（VSS）模块、动态交互模块（DIB）以及空间门控前馈网络（SGFN），该方法高效捕捉全局上下文与局部细粒度特征，在计算成本上相比SOTA方法降低了67.4%，同时保持高精度。

ABSTRACT

Underwater Image Enhancement (UIE) techniques aim to address the problem of underwater image degradation due to light absorption and scattering. In recent years, both Convolution Neural Network (CNN)-based and Transformer-based methods have been widely explored. In addition, combining CNN and Transformer can effectively combine global and local information for enhancement. However, this approach is still affected by the secondary complexity of the Transformer and cannot maximize the performance. Recently, the state-space model (SSM) based architecture Mamba has been proposed, which excels in modeling long distances while maintaining linear complexity. This paper explores the potential of this SSM-based model for UIE from both efficiency and effectiveness perspectives. However, the performance of directly applying Mamba is poor because local fine-grained features, which are crucial for image enhancement, cannot be fully utilized. Specifically, we customize the MambaUIE architecture for efficient UIE. Specifically, we introduce visual state space (VSS) blocks to capture global contextual information at the macro level while mining local information at the micro level. Also, for these two kinds of information, we propose a Dynamic Interaction Block (DIB) and Spatial feed-forward Network (SGFN) for intra-block feature aggregation. MambaUIE is able to efficiently synthesize global and local information and maintains a very small number of parameters with high accuracy. Experiments on UIEB datasets show that our method reduces GFLOPs by 67.4% (2.715G) relative to the SOTA method. To the best of our knowledge, this is the first UIE model constructed based on SSM that breaks the limitation of FLOPs on accuracy in UIE. The official repository of MambaUIE at https://github.com/1024AILab/MambaUIE.

研究动机与目标

• 解决基于CNN和Transformer的现有水下图像增强（UIE）方法存在的高计算成本与性能有限的问题。
• 探索状态空间模型（SSM），特别是Mamba，在高效且有效的UIE中的潜力。
• 通过保留局部细粒度特征，克服Mamba在UIE中直接应用时表现不佳的问题。
• 设计一种轻量化、参数高效的架构，以平衡全局上下文建模与局部细节恢复。
• 在大幅降低FLOPs的前提下实现UIE的SOTA性能，打破效率与精度之间的权衡。

提出的方法

• 提出一种定制的基于Mamba的架构MambaUIE，集成视觉状态空间（VSS）模块以建模长距离依赖关系并捕捉全局上下文。
• 引入动态交互模块（DIB），实现在模块内部的自适应特征交互，提升模块内特征聚合能力。
• 采用空间门控前馈网络（SGFN）以增强局部特征学习并保留细粒度细节。
• 在分层编码器-解码器结构中结合VSS模块、DIB与SGFN，以平衡全局与局部信息处理。
• 通过利用Mamba的SSM机制所具有的线性复杂度，优化模型以实现低FLOP运行。
• 在UIEB数据集上端到端训练模型，以学习联合增强与超分辨率（SR）能力。

实验结果

研究问题

• RQ1 Mamba状态空间模型能否被有效适配于水下图像增强，以实现高效率？
• RQ2 为何Mamba的直接应用在UIE中表现不佳？如何更好地保留局部细粒度特征？
• RQ3 结合SSM与局部特征学习机制的混合架构能否在低FLOP成本下实现SOTA性能？
• RQ4 所提出的动态交互模块（DIB）是否相比标准注意力机制或前馈网络机制能改善特征聚合？
• RQ5 能否通过利用基于SSM的模型打破UIE中精度与FLOPs之间的权衡？

主要发现

• MambaUIE&SR的FLOP仅达2.8 GFLOPs，相比SOTA方法计算成本降低67.4%。
• 该模型在UIEB数据集上的定量指标（如PSNR、SSIM）与定性视觉结果方面均优于现有SOTA方法。
• DIB与SGFN的集成显著提升了特征聚合效果，从而更好地保留了纹理与色彩细节。
• 所提出的架构在保持少量参数的同时实现高精度，展现出强大的参数效率。
• 这是首个基于SSM的UIE模型，成功打破了水下图像增强中传统FLOPs与精度之间的权衡。
• 该方法在多种多样的水下图像退化模式下表现出强大的泛化能力与鲁棒性。