Multi-Scale Representation Learning for Image Restoration with State-Space Model

📄 arXiv: 2408.10145v1 📥 PDF

作者: Yuhong He, Long Peng, Qiaosi Yi, Chen Wu, Lu Wang

分类: cs.CV

发布日期: 2024-08-19

💡 一句话要点

提出基于状态空间模型的多尺度图像复原网络MS-Mamba，实现高效高质量图像重建。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 图像复原 状态空间模型 多尺度表示学习 图像去噪 图像增强

📋 核心要点

  1. 现有图像复原方法受限于Transformer的高计算复杂度和CNN的有限感受野，难以兼顾性能与效率。
  2. 提出基于状态空间模型的多尺度图像复原网络MS-Mamba，利用全局和区域SSM模块增强多尺度表示学习能力。
  3. 实验表明，该方法在图像去雨、去雾、去噪和低光增强等任务上取得了SOTA性能，并保持较低的计算复杂度。

📝 摘要（中文）

图像复原旨在从退化的图像中重建高质量、细节丰富的图像，这在摄影和各种计算机视觉系统中至关重要。在实际场景中，不同类型的退化会导致图像细节在各个尺度上的损失，并降低图像对比度。现有方法主要依赖CNN和Transformer来捕获多尺度表示，但这些方法通常受到Transformer的高计算复杂度和CNN的有限感受野的限制，阻碍了它们在图像复原中实现卓越的性能和效率。为了解决这些挑战，我们提出了一种新颖的基于多尺度状态空间模型（MS-Mamba）的高效图像复原方法，该方法通过我们提出的全局和区域SSM模块增强了多尺度表示学习的能力。此外，还提出了自适应梯度块（AGB）和残差傅里叶块（RFB），通过捕获各个方向的梯度并促进频域中的细节学习，来提高网络的细节提取能力。在图像去雨、去雾、去噪和低光增强等四个经典图像复原任务的九个公共基准上进行的大量实验表明，我们提出的方法在保持低计算复杂度的同时，实现了新的最先进的性能。源代码将公开提供。

🔬 方法详解

问题定义：论文旨在解决图像复原任务中，现有方法无法有效捕获多尺度信息，且计算复杂度高的问题。现有方法如CNN感受野有限，Transformer计算量大，难以在实际应用中达到理想的性能和效率。

核心思路：论文的核心思路是利用状态空间模型（SSM）来高效地捕获图像的多尺度信息。通过设计全局和区域SSM模块，增强网络对不同尺度特征的建模能力，同时保持较低的计算复杂度。此外，还引入自适应梯度块（AGB）和残差傅里叶块（RFB）来增强细节提取能力。

技术框架：MS-Mamba的整体架构包含多尺度状态空间模型（MS-SSM）、自适应梯度块（AGB）和残差傅里叶块（RFB）。MS-SSM负责捕获全局和区域的多尺度特征，AGB用于提取各个方向的梯度信息，RFB则在频域中学习图像细节。整个网络通过残差连接进行信息传递，以提高训练的稳定性。

关键创新：论文的关键创新在于将状态空间模型引入图像复原任务，并设计了多尺度SSM模块。与传统的CNN和Transformer相比，SSM具有更强的长距离依赖建模能力和更低的计算复杂度。此外，AGB和RFB的设计也进一步提升了网络的细节提取能力。

关键设计：MS-SSM模块包含全局SSM和区域SSM，分别用于捕获全局上下文信息和局部细节信息。AGB通过多个方向的梯度卷积来提取图像的梯度特征。RFB利用傅里叶变换将图像转换到频域，并在频域中进行细节学习。损失函数方面，论文可能采用了L1损失或L2损失，以及感知损失等，具体细节需要在代码中进一步确认（未知）。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

MS-Mamba在图像去雨、去雾、去噪和低光增强等四个经典图像复原任务的九个公共基准上进行了广泛的实验，结果表明该方法在保持低计算复杂度的同时，实现了新的state-of-the-art性能。具体的性能提升幅度需要在论文中进一步查看（未知）。

🎯 应用场景

该研究成果可广泛应用于各种图像复原场景，例如手机摄影中的图像增强、监控视频的清晰化处理、医学图像的重建以及遥感图像的修复等。通过提高图像的质量和细节，可以提升后续计算机视觉任务的性能，例如目标检测、图像分割和人脸识别等，具有重要的实际应用价值和潜在的商业前景。

📄 摘要（原文）

Image restoration endeavors to reconstruct a high-quality, detail-rich image from a degraded counterpart, which is a pivotal process in photography and various computer vision systems. In real-world scenarios, different types of degradation can cause the loss of image details at various scales and degrade image contrast. Existing methods predominantly rely on CNN and Transformer to capture multi-scale representations. However, these methods are often limited by the high computational complexity of Transformers and the constrained receptive field of CNN, which hinder them from achieving superior performance and efficiency in image restoration. To address these challenges, we propose a novel Multi-Scale State-Space Model-based (MS-Mamba) for efficient image restoration that enhances the capacity for multi-scale representation learning through our proposed global and regional SSM modules. Additionally, an Adaptive Gradient Block (AGB) and a Residual Fourier Block (RFB) are proposed to improve the network's detail extraction capabilities by capturing gradients in various directions and facilitating learning details in the frequency domain. Extensive experiments on nine public benchmarks across four classic image restoration tasks, image deraining, dehazing, denoising, and low-light enhancement, demonstrate that our proposed method achieves new state-of-the-art performance while maintaining low computational complexity. The source code will be publicly available.