论文状态:已完成

EAMamba: Efficient All-Around Vision State Space Model for Image Restoration

发表:2025/06/27
原文链接PDF 下载
价格:0.100000
价格:0.100000
已有 1 人读过
本分析由 AI 生成,可能不完全准确,请以原文为准。

TL;DR 精炼摘要

本研究提出高效全方位Mamba(EAMamba),将多头选择性扫描模块和全方位扫描机制相结合,以解决视觉Mamba在图像恢复任务中的计算复杂度和局部像素遗忘问题。实验表明,EAMamba在保持相似性能的同时,FLOPs显著减少31-89%。

摘要

Image restoration is a key task in low-level computer vision that aims to reconstruct high-quality images from degraded inputs. The emergence of Vision Mamba, which draws inspiration from the advanced state space model Mamba, marks a significant advancement in this field. Vision Mamba demonstrates excellence in modeling long-range dependencies with linear complexity, a crucial advantage for image restoration tasks. Despite its strengths, Vision Mamba encounters challenges in low-level vision tasks, including computational complexity that scales with the number of scanning sequences and local pixel forgetting. To address these limitations, this study introduces Efficient All-Around Mamba (EAMamba), an enhanced framework that incorporates a Multi-Head Selective Scan Module (MHSSM) with an all-around scanning mechanism. MHSSM efficiently aggregates multiple scanning sequences, which avoids increases in computational complexity and parameter count. The all-around scanning strategy implements multiple patterns to capture holistic information and resolves the local pixel forgetting issue. Our experimental evaluations validate these innovations across several restoration tasks, including super resolution, denoising, deblurring, and dehazing. The results validate that EAMamba achieves a significant 31-89% reduction in FLOPs while maintaining favorable performance compared to existing low-level Vision Mamba methods.

思维导图

论文精读

中文精读

1. 论文基本信息

1.1. 标题

EAMamba: 高效全方位视觉状态空间模型用于图像恢复 (EAMamba: Efficient All-Around Vision State Space Model for Image Restoration)

1.2. 作者

Yu-Cheng Lin*, Yu-Syuan Xu*, Hao-Wei Chen, Hsien-Kai Kuo, Chun-Yi Lee *表示共同第一作者。 隶属机构包括:国立清华大学 (National Tsing Hua University)、国立台湾大学 (National Taiwan University) 和联发科技 (MediaTek Inc.)。

1.3. 发表期刊/会议

该论文作为预印本发表在 arXiv 上。截至目前,没有明确指出其已被正式接收或发表于特定期刊或会议,但其内容属于计算机视觉领域的顶级研究范畴,特别是图像处理和深度学习。

1.4. 发表年份

2025年

1.5. 摘要

图像恢复是低级计算机视觉 (low-level computer vision) 中的一项关键任务,旨在从降质输入中重建高质量图像。受先进状态空间模型 (State Space Model, SSM) Mamba 的启发,视觉 Mamba (Vision Mamba) 的出现标志着该领域的一个重大进展。视觉 Mamba 在以线性复杂度 (linear complexity) 建模长距离依赖 (long-range dependencies) 方面表现出色,这对于图像恢复任务具有至关重要的优势。尽管具有这些优点,视觉 Mamba 在低级视觉任务中仍面临挑战,包括计算复杂度随扫描序列数量增加而扩展的问题,以及局部像素遗忘 (local pixel forgetting) 问题。为解决这些限制,本研究引入了高效全方位 Mamba (Efficient All-Around Mamba, EAMamba),这是一个增强型框架,其集成了多头选择性扫描模块 (Multi-Head Selective Scan Module, MHSSM) 和全方位扫描机制 (all-around scanning mechanism)。MHSSM 有效地聚合多个扫描序列,避免了计算复杂度和参数数量的增加。全方位扫描策略采用多种模式来捕获整体信息,并解决了局部像素遗忘问题。实验评估验证了这些创新在多种恢复任务中的有效性,包括超分辨率 (super resolution)、去噪 (denoising)、去模糊 (deblurring) 和去雾 (dehazing)。结果表明,EAMamba 在保持与现有低级视觉 Mamba 方法相当的性能的同时,FLOPs (浮点运算) 显著减少了 31-89%。

1.6. 原文链接

原文链接: https://arxiv.org/abs/2506.22246 PDF 链接: https://arxiv.org/pdf/2506.22246v1.pdf 发布状态:预印本 (Preprint)。


2. 整体概括

2.1. 研究背景与动机

图像恢复 (Image Restoration) 是计算机视觉中的一项基础且关键的任务,旨在从各种降质(如噪声、模糊、下采样等)的低质量图像中重建出高质量图像。这项任务本质上是病态的 (ill-posed),具有相当大的技术挑战。

历史上,该领域的技术演进经历了几个阶段:

  • 基于卷积神经网络 (Convolutional Neural Network, CNN) 的方法: 早期的 CNN 模型在各种图像恢复基准测试中取得了成功。然而,由于其主要关注局部像素关系,在捕获全局信息方面存在固有限制。

  • 基于视觉 Transformer (Vision Transformer, ViT) 的方法: 随后的研究引入了基于 ViT 的架构,利用多头自注意力机制 (multi-head self-attention mechanisms) 来建模图像所有像素之间的关系。这些方法有效地捕获了全局依赖,并在多项恢复任务中取得了有希望的结果。然而,自注意力机制的计算复杂度与像素数量呈二次方增长,使得高分辨率图像处理变得不切实际。

  • 基于视觉 Mamba (Vision Mamba) 的方法: 为了应对计算复杂性挑战,最近引入的 Mamba 框架(一种先进的状态空间模型 (State Space Model, SSM))提供了一个有前景的解决方案。Mamba 具有线性计算缩放特性和高效的长距离依赖建模能力,在自然语言处理 (Natural Language Processing, NLP) 任务中展现出卓越潜力。受此启发,视觉 Mamba 模型将 Mamba 适应于视觉任务,旨在高效捕获全局图像信息,同时保持与像素数量呈线性关系的计算复杂度。

    尽管视觉 Mamba 具有优点,但现有的视觉 Mamba 方法在低级视觉任务中仍面临挑战:

  1. 计算复杂度与扫描序列数量成比例扩展: 现有的视觉 Mamba 通常采用双向扫描策略,将二维特征图转换为展平的一维序列。例如,二维选择性扫描 (2DSS) 会通过不同扫描模式生成多个一维序列(如四个),每个序列都需要独立的选择性扫描和参数,这不可避免地增加了计算开销。当需要更多扫描方向以捕获更全面的信息时,这一问题会更加突出。

  2. 局部像素遗忘 (Local Pixel Forgetting): 现有的二维扫描方法存在局部像素遗忘现象。当二维特征图中的空间相邻像素在转换为一维序列时变得距离遥远,会导致局部信息丢失。这对于图像恢复任务尤其不利,因为这些任务对保留局部空间关系至关重要。

    本文的切入点和创新思路在于解决上述两个限制,提出一种更高效、更全面的视觉 Mamba 架构。

2.2. 核心贡献/主要发现

本文的主要贡献体现在以下几个方面:

  1. 提出多头选择性扫描模块 (Multi-Head Selective Scan Module, MHSSM): 引入 MHSSM,通过通道分组策略进行选择性扫描,高效处理和聚合展平的一维序列,避免了计算复杂度和参数数量随扫描序列增加而带来的开销。这显著提高了视觉 Mamba 框架的可扩展性和效率。
  2. 引入全方位扫描策略 (All-Around Scanning Strategy): 受益于 MHSSM 的效率,EAMamba 能够整合水平、垂直、对角线和翻转对角线及其反向等多个扫描方向,捕获更全面的空间信息。这种多方向扫描方法有效解决了局部像素遗忘问题,增强了模型理解图像空间关系的能力。
  3. 全面的实验验证和性能提升: 在超分辨率、去噪、去模糊和去雾等多种图像恢复任务中进行了广泛的实验验证。结果表明,EAMamba 在保持甚至略微提升性能的同时,实现了显著的 FLOPs 减少(31-89%),同时参数数量也更少,设立了视觉 Mamba 图像恢复效率的新标杆。
  4. 深入的消融研究和分析: 对全方位扫描和 MHSSM 的性能影响进行了全面分析,并提供了关于这些创新优势的见解,包括不同扫描策略的有效感受野 (Effective Receptive Field, ERF) 可视化和通道 MLP (Multilayer Perceptron) 设计的比较。

2.3. 创新点

EAMamba 的创新点在于将两种协同的机制引入视觉 Mamba 架构:

  1. 高效的多头处理: MHSSM 避免了传统 2DSS 在增加扫描方向时导致的计算开销,使得更复杂的扫描策略在计算上可行。

  2. 全面的信息捕获: 全方位扫描 直接解决了 局部像素遗忘 问题,通过整合多方向扫描捕获更丰富的空间上下文,这对于需要精细局部细节的图像恢复任务至关重要。


3. 预备知识与相关工作

3.1. 基础概念

  • 图像恢复 (Image Restoration): 计算机视觉领域中的一类任务,旨在从受损(如噪声、模糊、下采样等)的低质量图像中恢复出高质量的原始图像。这是一个逆问题 (inverse problem),通常是病态的 (ill-posed),因为从降质图像到高质量图像可能存在多个解决方案。
  • 卷积神经网络 (Convolutional Neural Network, CNN): 一种深度学习模型,通过卷积层 (convolutional layers) 自动学习图像的层次化特征。CNN 在图像处理任务中表现出色,但其感受野 (receptive field) 有限,难以有效捕获全局依赖。
  • 视觉 Transformer (Vision Transformer, ViT): 受自然语言处理领域 Transformer 模型启发,将图像分割成固定大小的图像块 (patches),然后将这些图像块作为序列输入 Transformer 编码器。ViT 通过自注意力机制 (self-attention mechanism) 能够捕获全局范围的依赖关系,但在高分辨率图像上计算成本高昂。
  • 自注意力机制 (Self-Attention Mechanism): Transformer 的核心组件,允许模型在处理序列时,对序列中的每个元素都关注序列中的所有其他元素,并计算它们之间的关联强度。对于图像而言,这意味着每个像素可以“关注”图像中的所有其他像素。其计算复杂度通常与输入序列长度的平方成正比。
  • 状态空间模型 (State Space Model, SSM): 一类用于建模序列数据的数学模型,它通过一个隐藏状态来表示系统的内部状态,并根据当前输入和隐藏状态来预测下一个输出和更新隐藏状态。SSM 具有建模长距离依赖的能力,并且计算复杂度可以做到线性。
  • Mamba: 一种先进的具有选择性机制 (selective mechanism) 的状态空间模型。Mamba 通过其选择性状态空间 (selective state space) 特性,能够根据输入内容动态地过滤或保留信息,从而高效地处理长序列,并展现出线性计算复杂度。
  • 视觉 Mamba (Vision Mamba): 将 Mamba 框架应用于视觉任务的模型。通常将二维图像特征转换为一维序列,然后应用 Mamba 的选择性状态空间机制进行处理,以期在保持全局信息捕获能力的同时降低计算复杂度。
  • FLOPs (Floating Point Operations): 浮点运算次数,是衡量模型计算复杂度和推理速度的一个常用指标。FLOPs 越低通常意味着模型运行越快、计算资源消耗越少。
  • PSNR (Peak Signal-to-Noise Ratio): 峰值信噪比,一种衡量图像质量的客观指标。它通过计算原始图像与处理后图像之间的均方误差 (Mean Squared Error, MSE) 来评估图像失真程度。PSNR 值越高,表示图像质量越好,失真越小。
    • 概念定义: PSNR 用于量化处理后图像与原始(无噪声、无失真)图像之间的相似度。它基于像素值差异的均方误差,并以分贝 (dB) 为单位表示,值越大表示图像质量越好。
    • 数学公式: PSNR=10log10(MAXI2MSE) \mathrm{PSNR} = 10 \cdot \log_{10} \left( \frac{\mathrm{MAX}_I^2}{\mathrm{MSE}} \right) 其中, MSE=1mni=0m1j=0n1[I(i,j)K(i,j)]2 \mathrm{MSE} = \frac{1}{mn} \sum_{i=0}^{m-1} \sum_{j=0}^{n-1} [I(i,j) - K(i,j)]^2
    • 符号解释:
      • MAXI\mathrm{MAX}_I: 图像中像素的最大可能值。对于 8 位灰度图像,MAXI=255\mathrm{MAX}_I = 255
      • MSE\mathrm{MSE}: 均方误差。
      • I(i,j): 原始图像在坐标 (i,j) 处的像素值。
      • K(i,j): 降质或处理后图像在坐标 (i,j) 处的像素值。
      • m, n: 图像的行数和列数。
  • SSIM (Structural Similarity Index Measure): 结构相似性指数,另一种衡量图像质量的客观指标。它从亮度、对比度和结构三个方面衡量两幅图像的相似性,比 PSNR 更符合人类视觉感知。SSIM 值介于 -1 和 1 之间,1 表示两幅图像完全相同。
    • 概念定义: SSIM 旨在量化两幅图像之间的感知相似度,它考虑了人眼对图像亮度、对比度和结构变化的敏感性,从而提供比 PSNR 更符合视觉感知的质量评估。
    • 数学公式: SSIM(x,y)=(2μxμy+C1)(2σxy+C2)(μx2+μy2+C1)(σx2+σy2+C2) \mathrm{SSIM}(x,y) = \frac{(2\mu_x\mu_y + C_1)(2\sigma_{xy} + C_2)}{(\mu_x^2 + \mu_y^2 + C_1)(\sigma_x^2 + \sigma_y^2 + C_2)}
    • 符号解释:
      • x, y: 待比较的两幅图像的像素值集合(或局部窗口)。
      • μx,μy\mu_x, \mu_y: 图像 xxyy 的平均值。
      • σx2,σy2\sigma_x^2, \sigma_y^2: 图像 xxyy 的方差。
      • σxy\sigma_{xy}: 图像 xxyy 的协方差。
      • C1=(K1L)2,C2=(K2L)2C_1 = (K_1 L)^2, C_2 = (K_2 L)^2: 用于稳定除法的小常数,避免分母接近零。LL 是像素值的动态范围(例如 8 位灰度图像为 255)。K1=0.01,K2=0.03K_1 = 0.01, K_2 = 0.03 是默认值。
  • 有效感受野 (Effective Receptive Field, ERF): 指神经网络中输出层某个神经元对输入图像实际产生影响的区域。与理论感受野 (Theoretical Receptive Field) 不同,ERF 考虑了网络权重分布对输入影响强度的差异,更真实地反映了网络对局部信息的关注程度。
  • UNet 架构: 一种编码器-解码器 (encoder-decoder) 结构的卷积神经网络,最初用于图像分割。其特点是编码器和解码器之间存在跳跃连接 (skip connections),将编码器阶段的特征图直接传递给解码器阶段,有助于保留空间细节,对图像恢复任务非常有用。
  • Layer Normalization (LN): 层归一化,一种归一化技术,对单个样本的每个层的特征进行归一化,有助于稳定训练并加速收敛。
  • 深度可分离卷积 (Depth-wise Convolution, DWConv2D): 一种分解卷积操作,将标准卷积分解为深度卷积和点卷积 (pointwise convolution) 两步。深度卷积对每个输入通道独立进行卷积,点卷积则通过 1×11 \times 1 卷积聚合通道。这大大减少了参数量和计算量。
  • SiLU (Sigmoid Linear Unit): 一种激活函数,计算为 xσ(x)x \cdot \sigma(x),其中 σ(x)\sigma(x) 是 Sigmoid 函数。SiLU 具有非线性、无上界但有下界,且在负数区域具有平滑梯度的特性,有助于深度网络的训练。
  • AdamW 优化器: 一种带有解耦权重衰减 (decoupled weight decay) 的 Adam 优化器。它将权重衰减从梯度更新中分离出来,有助于提高模型的泛化能力。
  • L1 损失: 一种常用的损失函数,计算预测值与真实值之间绝对误差的和。在图像恢复中,L1 损失可以鼓励模型生成更锐利、更少模糊的图像。
  • Cosine Annealing (余弦退火): 一种学习率调度策略,学习率按照余弦函数的形式从最大值逐渐下降到最小值,有助于模型在训练后期进行更精细的优化。

3.2. 前人工作

论文在引言和相关工作部分回顾了图像恢复领域的主要技术路线:

  • 基于 CNN 的方法: 早期图像恢复研究主要依赖于 CNN,如 DnCNN [9]、IRCNN [10]、FFDNet [11]、DRUNet [13] 等。这些方法通过局部感受野捕获图像特征,在特定任务上表现良好。

    • 局限性: 固有的局部性限制了它们捕获全局信息的能力,这对于理解图像的整体结构至关重要。
  • 基于 ViT 的方法: 为了解决 CNN 的局部性限制,研究者引入了 ViT 架构,利用多头自注意力机制建模全局依赖。代表性工作包括 SwinIR [48]、Uformer [50]、Restormer [51]、IPT [49] 等。

    • 局限性: 自注意力机制的计算复杂度与像素数量呈二次方增长,导致在高分辨率图像处理时计算量巨大,难以扩展。为了缓解这一问题,SwinIRUformer 采用了基于窗口的注意力机制,Restormer 引入了多 Dconv 头转置注意力 (multi-Dconv head transposed attention) 和门控 Dconv 前馈网络 (gated-Dconv feed-forward)。
  • 基于 Vision Mamba 的方法: 近期 Mamba 模型的成功促使其被应用于视觉领域,旨在结合全局建模能力和线性计算复杂度。

    • VMamba [54] 和 Vision Mamba [55] 是先驱性工作,它们探索了如何将二维图像转换为一维序列并应用状态空间模型。VMamba 引入了 交叉扫描 (cross-scan)交叉合并 (cross-merge) 技术来聚合空间信息。
    • 在图像恢复领域,MambaIR [57] 将视觉状态空间模块与修改后的 MLP (多层感知机) 模块结合,以缓解 局部像素遗忘通道冗余 问题。
    • VMambaIR [58] 提出了 全方位选择性扫描 (Omni Selective Scan, OSS) 模块,执行四方向空间扫描和通道扫描,利用空间和通道信息。
    • MambaIR 的具体问题: MambaIR [57] 虽然尝试缓解局部像素遗忘,但其二维选择性扫描 (2DSS) 依然存在计算复杂度随扫描序列增加而线性增长的问题,且对角线方向的信息捕获能力不足(如论文 Figure 4(b) 所示)。

3.3. 技术演进

图像恢复技术从早期的基于信号处理方法,发展到以 CNN 为核心的深度学习方法,再到利用 Transformer 捕获长距离依赖。然而,Transformer 的高计算成本推动了对更高效架构的探索,Mamba 及其视觉变体应运而生。EAMamba 正是这一技术演进的最新成果,旨在解决 Vision Mamba 现有方法在效率和信息捕获方面的具体不足,特别是局部像素遗忘计算复杂度问题,以实现高效且全方位的图像恢复。

3.4. 差异化分析

EAMamba 与现有 Vision Mamba 方法的核心区别和创新点在于:

  1. 计算效率的提升: 现有 2DSS 方法(如 MambaIR [57] 中使用的)在增加扫描方向时会导致计算复杂度和参数数量的线性增长。EAMamba 引入的 Multi-Head Selective Scan (MHSS) 通过通道分组 (channel grouping) 策略,能够在不增加额外计算开销的情况下处理多个一维序列,从而显著提高了计算效率,如 Table 1 所示,EAMamba 的 FLOPs 远低于 MambaIR。

  2. 全面空间信息捕获: 现有 Vision Mamba 的二维扫描方法存在局部像素遗忘问题,尤其是在对角线方向上信息捕获不足。EAMamba 的全方位扫描策略通过结合水平、垂直、对角线和翻转对角线及其反向等多个扫描方向,实现了对图像更全面的空间信息捕获,有效缓解了局部像素遗忘问题,增强了模型对空间上下文的理解。Figure 4(b) 的 ERF 可视化结果证实了全方位扫描在保留局部信息方面的优势。

  3. 更好的性能与效率平衡: EAMamba 在各种图像恢复任务上,在实现显著 FLOPs 减少(31-89%)的同时,能保持甚至超越现有低级 Vision Mamba 方法的性能,在计算效率和图像恢复质量之间实现了更好的平衡(如 Figure 1 所示)。


4. 方法论

4.1. 方法原理

EAMamba 旨在通过创新的架构设计,解决现有视觉 Mamba 模型在低级视觉任务中面临的计算复杂度和局部像素遗忘问题。其核心思想是构建一个高效且能全面捕捉空间信息的视觉 Mamba 框架。这通过两个主要组件实现:

  1. 多头选择性扫描模块 (Multi-Head Selective Scan Module, MHSSM): 该模块通过将输入特征的通道分组,并对每个分组独立进行选择性扫描,从而在不增加计算开销和参数数量的前提下,高效聚合多个扫描序列。这使得 EAMamba 能够灵活地集成更复杂的扫描策略。

  2. 全方位扫描策略 (All-Around Scanning Strategy): 该策略利用 MHSSM 的效率优势,将多种扫描方向(如水平、垂直、对角线、翻转对角线及其反向)结合起来。这种多方向扫描能够捕获更全面的空间依赖关系,有效解决传统二维扫描中存在的局部像素遗忘问题。

    EAMamba 整体上采用 UNet-like 编码器-解码器架构,并通过自定义的 MambaFormer 模块作为其基础构建块,实现渐进式特征提取、融合和恢复。

4.2. 核心方法详解 (逐层深入)

4.2.1. EAMamba 框架概览

EAMamba 框架采用 UNet-like 架构,专门为图像恢复任务设计。其处理流程如下:

  1. 输入: 接收一张低质量图像 ILQRH×W×3I^{LQ} \in \mathbb{R}^{H \times W \times 3},其中 H, W 分别为图像的高度和宽度,3 表示 RGB 三通道。

  2. 编码器 (Encoder): 图像首先通过一个初始的线性层将输入图像 ILQI^{LQ} 投影到特征空间,然后经过三个 MambaFormer 编码器模块 (MambaFormer encoder modules)。这些模块在不同尺度上操作,逐步提取具有变化维度的特征嵌入 (feature embeddings)。编码器阶段旨在捕获图像的语义信息和长距离依赖。

  3. 瓶颈模块 (Bottleneck Module): 编码器之后是一个 瓶颈模块 (bottleneck module),通常包含更多的 MambaFormer 块,用于处理最抽象、最深层的特征。

  4. 解码器 (Decoder): 瓶颈模块的输出随后被送入三个 MambaFormer 解码器模块。解码器阶段负责将抽象特征逐步上采样并恢复空间细节。在每个解码器阶段,通过跳跃连接 (skip connections) 将来自对应编码器阶段的特征融合进来,以保留更多的图像细节信息。

  5. 精炼模块 (Refinement Module): 解码器之后是一个 精炼模块 (refinement module),通常包含两个 MambaFormer 块,用于对恢复的特征进行最终的细化处理。

  6. 输出: 精炼模块的输出通过一个最终的线性层生成一个残差图像 IHQrRH×W×3I^{HQ}r \in \mathbb{R}^{H \times W \times 3}

  7. 最终高质量图像: 最终的高质量图像 IHQRH×W×3I^{HQ} \in \mathbb{R}^{H \times W \times 3} 通过将残差图像 IHQrI^{HQ}r 与原始低质量输入图像 ILQI^{LQ} 进行逐元素相加得到:IHQ=ILQ+IHQrI^{HQ} = I^{LQ} + I^{HQ}r。这种残差学习 (residual learning) 有助于模型更好地学习降质与高质量图像之间的差异。

    该图像是一个示意图,展示了 EAMamba 框架及其关键组件,包括 MambaFormer 和 MHSSM。图中使用了 \(H \\times W \\times C\) 表示特征图形状,并展示了多头选择扫描模块的结构。这些模块通过高效的全方位扫描机制解决了低级视觉任务中的局部像素遗忘问题。 该图像是一个示意图,展示了 EAMamba 框架及其关键组件,包括 MambaFormer 和 MHSSM。图中使用了 H×W×CH \times W \times C 表示特征图形状,并展示了多头选择扫描模块的结构。这些模块通过高效的全方位扫描机制解决了低级视觉任务中的局部像素遗忘问题。

Figure 5 提供了 EAMamba 框架的整体架构图。

4.2.2. MambaFormer 模块

MambaFormer 模块是 EAMamba 框架(包括编码器、解码器、瓶颈和精炼阶段)的基本构建块。如 Figure 5 (a) 所示,每个 MambaFormer 模块包含两个主要组件:

  1. 多头选择性扫描模块 (Multi-Head Selective Scan Module, MHSSM): 用于令牌混合 (token mixing),旨在捕获长距离空间依赖。

  2. 通道多层感知机 (Channel Multilayer Perceptron, Channel MLP): 用于特征精炼,以增强特征表示。

    在每个组件之前,都应用了 层归一化 (Layer Normalization, LN) [66],并且通过残差连接 (residual connections) 将组件的输出与前一个输入集成。这个过程可以通过以下公式表示:

X=X+MHSSM(LN(X)),X=X+Channel MLP(LN(X)), \begin{array} { r l } & { X ^ { \prime } = X + \mathbf { M H S S M } ( \mathbf { L N } ( X ) ) , } \\ & { X ^ { \prime \prime } = X ^ { \prime } + \mathbf { C h a n n e l ~ M L P } ( \mathbf { L N } ( X ^ { \prime } ) ) , } \end{array}

其中:

  • XX: 输入特征。
  • LN()\mathbf{LN}(\cdot): 层归一化操作。
  • MHSSM()\mathbf{MHSSM}(\cdot): 多头选择性扫描模块。
  • Channel MLP()\mathbf{Channel~MLP}(\cdot): 通道多层感知机。
  • XX': 经过 MHSSM 处理后的特征。
  • XX'': 经过 Channel MLP 处理后的最终输出特征。

4.2.3. 多头选择性扫描模块 (MHSSM) 与全方位扫描

MHSSMMambaFormer 的关键组件,它通过将传统的 二维选择性扫描 (2DSS) 替换为 Multi-Head Selective Scan (MHSS) 来增强视觉选择性扫描模块 [54]。如 Figure 5 (b) 所示,MHSSM 的架构将输入特征 XRH×W×CX \in \mathbb{R}^{H \times W \times C}(其中 H, W 是特征图的高度和宽度,CC 是通道数)通过两个并行的分支进行处理:

  1. 左分支 (Gating Branch):
    • 首先,通过一个线性层 (Linear layer) 将特征通道扩展到 λC\lambda C,其中 λ\lambda 是预定义的通道扩展因子。
    • 扩展后的特征接着经过一系列操作:深度可分离卷积 (DWConv2D)SiLU 激活函数 [67]、Multi-Head Selective Scan (MHSS),以及 层归一化 (LN)
  2. 右分支 (Feature Branch):
    • 执行通道扩展 λC\lambda C,然后经过 SiLU 激活函数

      这两个分支的输出通过逐元素相乘 (\odot) 进行组合。最后,一个线性投影层将合并后的输出维度还原到原始的 CC 维度,得到 MHSSM 的最终输出。MHSSM 的完整过程可以用以下公式表示:

Y=LN(MHSS(SiLU(DWConv2D(Linear(X))))),Z=SiLU(Linear(X)),Xout=Linear(YZ), \begin{array} { r l } & { Y = \mathrm { L N } ( \mathrm { M H S S } ( \mathrm { S i L U } ( \mathrm { D W C o n v 2 D } ( \mathrm { L i n e a r } ( X ) ) ) ) ) , } \\ & { Z = \mathrm { S i L U } ( \mathrm { L i n e a r } ( X ) ) , \quad X _ { o u t } = \mathrm { L i n e a r } ( Y \otimes Z ) , } \end{array}

其中:

  • XX: 输入特征。
  • Linear()\mathrm{Linear}(\cdot): 线性投影层。
  • DWConv2D()\mathrm{DWConv2D}(\cdot): 深度可分离卷积操作。
  • SiLU()\mathrm{SiLU}(\cdot): SiLU 激活函数。
  • MHSS()\mathrm{MHSS}(\cdot): 多头选择性扫描操作。
  • LN()\mathrm{LN}(\cdot): 层归一化操作。
  • YY: 左分支(门控分支)的输出。
  • ZZ: 右分支(特征分支)的输出。
  • \odot: 逐元素相乘操作。
  • XoutRH×W×CX_{out} \in \mathbb{R}^{H \times W \times C}: MHSSM 模块的最终输出。

4.2.4. Multi-Head Selective Scan (MHSS)

MHSSMHSSM 中的核心组件,它通过多头方法和分组特征处理来捕获长距离空间信息,同时保持计算效率。如 Figure 6 所示,MHSS 的操作流程如下:

  1. 通道分组 (Channel Partitioning): MHSS 不像传统的 2DSS 那样在整个通道维度上进行扫描,而是将输入特征 MHSSinMHSS_{in} 沿通道维度分割成 nn 个组。

  2. 二维输入转换 (Transformation for 2D Input): 每个组的特征都会经过一个转换函数 Transform,将其二维输入(或其部分)转换成一维序列 SSiniSS_{in}^i。在本工作中,这里实现了全方位扫描策略,其细节在下一节阐述。

  3. 选择性扫描 (Selective Scan): 对于每个转换后的一维序列 SSiniSS_{in}^i,执行独立的 选择性扫描 (SelectiveScan) 操作,生成对应的输出 SSoutiSS_{out}^i。选择性扫描是 Mamba 模型的关键机制,它允许模型根据输入内容动态地选择要记住或遗忘的信息。

  4. 逆转换与拼接 (Inverse Transformation and Concatenation): 将每个组的输出 SSoutiSS_{out}^i 经过逆转换 InverseTransform,然后沿通道维度拼接 Concat 起来,形成 MHSS 的最终输出 MHSSoutMHSS_{out}

    MHSS 的操作可以表示为:

SSinN=Transform(Split(MHSSin)),SSoutN=SelectiveScan(SSinN),MHSSout=Concat(InverseTransform(SSoutN)), \begin{array} { r } { { S S _ { i n } ^ { N } = \mathrm { T r a n s f o r m } ( \mathrm { S p l i t } ( M H S S _ { i n } ) ) , } \qquad } \\ { { S S _ { o u t } ^ { N } = \mathrm { S e l e c t i v e S c a n } ( S S _ { i n } ^ { N } ) , \qquad } } \\ { { M H S S _ { o u t } = \mathrm { C o n c a t } ( \mathrm { I n v e r s e T r a n s f o r m } ( S S _ { o u t } ^ { N } ) ) , } } \end{array}

其中:

  • MHSSinMHSS_{in}: MHSS 模块的输入特征。

  • Split()\mathrm{Split}(\cdot): 沿通道维度将输入特征分割成 nn 个组的操作。

  • Transform()\mathrm{Transform}(\cdot): 将二维特征转换为一维序列的函数,实现了全方位扫描策略

  • SSinNSS_{in}^N: 转换后的一维序列集合,其中 N={iR:i={1,2,...,n}}N = \{i \in \mathbb{R} : i = \{1, 2, ..., n\}\} 表示各个组。

  • SelectiveScan()\mathrm{SelectiveScan}(\cdot): Mamba 模型的选择性扫描操作。

  • SSoutNSS_{out}^N: 选择性扫描后的输出序列集合。

  • InverseTransform()\mathrm{InverseTransform}(\cdot): 将一维序列逆转换为二维特征的函数。

  • Concat()\mathrm{Concat}(\cdot): 沿通道维度拼接操作。

  • MHSSoutMHSS_{out}: MHSS 模块的最终输出。

    MHSS 的主要优势在于其计算复杂度和参数数量与标准的选择性扫描 [53] 相当,但通过多头策略实现了更高的效率。与 2DSS 相比,MHSS 在增加扫描方向时不会导致复杂度的线性增加。

    Figure 6. Illustration of the Multi-Head Selective Scan (MHSS) with our proposed All-Around Scanning strategy. 该图像是一个示意图,展示了多头选择扫描(MHSS)与我们提出的全方位扫描策略的结合过程。图中描述了输入通道的分离、变换、选择扫描和逆变换等步骤,旨在高效聚合多条扫描序列以提升图像恢复性能。

Figure 6 展示了 Multi-Head Selective Scan (MHSS)全方位扫描策略的结合。

4.2.5. 全方位扫描 (All-Around Scanning)

受益于 MHSS 的高效率,EAMamba 引入了全方位扫描策略作为 MHSSM 中的 Transform 函数,以实现对空间依赖的全面理解。这种策略通过执行多方向的选择性扫描来解决传统二维扫描中固有的局部像素遗忘限制。

具体来说,全方位扫描策略包括以下扫描方向:

  • 水平扫描 (Horizontal Scan): 从左到右和从右到左。

  • 垂直扫描 (Vertical Scan): 从上到下和从下到上。

  • 对角线扫描 (Diagonal Scan): 从左上到右下和从右下到左上。

  • 翻转对角线扫描 (Flipped Diagonal Scan): 从右上到左下和从左下到右上。

    通过这种多方向扫描方法,全方位扫描策略能够纳入更广泛的邻域信息,从而加强空间上下文理解并减轻局部像素遗忘。例如,如图 Figure 3 所示,它结合了二维扫描对角线扫描。Figure 4 (b) 的 ERF 比较也支持了这一策略,显示 二维扫描(如 MambaIR [57] 中采用的)难以捕获对角线方向的信息,即使有局部卷积操作的补充。通过整合额外的扫描方向,EAMamba 实现了更大的感受野,特别是在保留目标像素周围的局部信息方面,这对于图像恢复任务至关重要。

    Figure 3. Illustration of an all-around scanning approach that combines two-dimensional scanning and diagonal scanning. 该图像是示意图,展示了全方位扫描的方法,该方法结合了二维扫描和对角线扫描。左侧部分显示了二维扫描的过程,右侧部分则展示了对角线扫描的方式,旨在捕捉图像的整体信息。

Figure 3 展示了全方位扫描方法如何结合二维扫描和对角线扫描。

Figure 4. (a) Ilustration of the local pixel forgetting phenomenon, where spatially adjacent pixels become distantly separated in the one-dimensional token sequence during scanning. The target pixel (highlighted in red square) and its adjacent pixels demonstrate how different scanning patterns affect spatial relationships. (b) The ERF visualization results averaged across the SIDD dataset \[59\], which depict improved spatial dependency preservation with the proposed all-around scanning approach. 该图像是图示,通过对比2D扫描和全方位扫描展示了不同扫描方式对空间关系的影响。左侧展示了相邻扫描参考图,右侧为有效感受野的可视化结果,表明全方位扫描策略更好地保持空间依赖性。

Figure 4 (a) 描述了局部像素遗忘现象,(b) 通过 ERF 可视化结果表明全方位扫描在保留空间依赖性方面的改进。

这种设计使得 MHSSM 能够同时实现两个目标:通过通道分割实现计算效率,以及通过全方位扫描捕获全面的视觉信息。这种策略能够在不带来过度计算开销的情况下,有效地进行图像恢复。


5. 实验设置

5.1. 数据集

EAMamba 在多种图像恢复任务上进行了广泛的实验,涉及以下数据集:

  • 图像去噪 (Image Denoising):
    • 合成高斯彩色图像去噪 (Synthetic Gaussian Color Denoising):
      • 训练数据: DIV2K [82]、Flickr2K [83]、WED [84] 和 BSD [71] 中的图像,用于训练模型处理 σ[050]\sigma [0-50] 范围内的噪声水平。
      • 评估基准: CBSD68 [71]、Kodak24 [73] 和 McMaster [72],在多个噪声水平 σ=[15,25,50]\sigma = [15, 25, 50] 下进行评估。
    • 真实世界去噪 (Real-World Denoising):
      • 训练和评估数据: SIDD [59] 数据集。
  • 图像超分辨率 (Image Super-Resolution, SR):
    • 训练和评估数据: RealSR [61] 基准数据集,用于 ×2,×3,×4\times 2, \times 3, \times 4 缩放因子。

    • 示例样本 (Super-Resolution ×4\times 4):

      该图像是展示EAMamba在图像恢复任务中比较不同方法性能的示意图,包括真实图像、低质量图像与其他恢复方法以及EAMamba的结果,最后一列显示了EAMamba获得的最高PSNR值35.04。 该图像是展示EAMamba在图像恢复任务中比较不同方法性能的示意图,包括真实图像、低质量图像与其他恢复方法以及EAMamba的结果,最后一列显示了EAMamba获得的最高PSNR值35.04。

      Figure 8 展示了 RealSR 数据集在 ×4\times 4 缩放因子下的超分辨率结果,裁剪区域表示真实图像与生成结果之间的归一化差异。

  • 图像去模糊 (Image Deblurring):
    • 训练数据: GoPro [62] 数据集。

    • 评估基准: GoPro [62] 和 HIDE [86] 基准数据集。

    • 示例样本 (Image Deblurring):

      该图像是图表,展示了EAMamba在图像重建任务中与其他方法(包括MPRNet、MAXIM-3S、Restormer和SFNet)的PSNR(峰值信噪比)性能对比。在各方法中,EAMamba的PSNR达到36.52,显示了其在低级视觉任务中的优越性。 该图像是图表,展示了EAMamba在图像重建任务中与其他方法(包括MPRNet、MAXIM-3S、Restormer和SFNet)的PSNR(峰值信噪比)性能对比。在各方法中,EAMamba的PSNR达到36.52,显示了其在低级视觉任务中的优越性。

      Figure 9 展示了 GoPro 数据集上的图像去模糊任务的定性结果。

  • 图像去雾 (Image Dehazing):
    • 训练和评估数据: RESIDE [63] 合成基准数据集。

    • 示例样本 (Image Dehazing):

      该图像是一个比较展示,左侧为真实图像,右侧展示了低质量图像及多种图像恢复方法的结果,包括Dehamer、MAXIM-2S、DehazeFormer-L和我们的EAMamba方法,显示出EAMamba在图像恢复中取得的最大PSNR值46.28,显著高于其他方法。 该图像是一个比较展示,左侧为真实图像,右侧展示了低质量图像及多种图像恢复方法的结果,包括Dehamer、MAXIM-2S、DehazeFormer-L和我们的EAMamba方法,显示出EAMamba在图像恢复中取得的最大PSNR值46.28,显著高于其他方法。

      Figure 10 展示了 SOTS-Indoor 子集上的图像去雾定性结果。

5.2. 评估指标

性能评估主要采用以下指标:

  • 峰值信噪比 (Peak Signal-to-Noise Ratio, PSNR):
    1. 概念定义: PSNR 用于量化处理后图像与原始(无噪声、无失真)图像之间的相似度。它基于像素值差异的均方误差,并以分贝 (dB) 为单位表示,值越大表示图像质量越好。
    2. 数学公式: PSNR=10log10(MAXI2MSE) \mathrm{PSNR} = 10 \cdot \log_{10} \left( \frac{\mathrm{MAX}_I^2}{\mathrm{MSE}} \right) 其中, MSE=1mni=0m1j=0n1[I(i,j)K(i,j)]2 \mathrm{MSE} = \frac{1}{mn} \sum_{i=0}^{m-1} \sum_{j=0}^{n-1} [I(i,j) - K(i,j)]^2
    3. 符号解释:
      • MAXI\mathrm{MAX}_I: 图像中像素的最大可能值。对于 8 位灰度图像,MAXI=255\mathrm{MAX}_I = 255
      • MSE\mathrm{MSE}: 均方误差。
      • I(i,j): 原始图像在坐标 (i,j) 处的像素值。
      • K(i,j): 降质或处理后图像在坐标 (i,j) 处的像素值。
      • m, n: 图像的行数和列数。
  • 结构相似性指数 (Structural Similarity Index Measure, SSIM):
    1. 概念定义: SSIM 旨在量化两幅图像之间的感知相似度,它考虑了人眼对图像亮度、对比度和结构变化的敏感性,从而提供比 PSNR 更符合视觉感知的质量评估。
    2. 数学公式: SSIM(x,y)=(2μxμy+C1)(2σxy+C2)(μx2+μy2+C1)(σx2+σy2+C2) \mathrm{SSIM}(x,y) = \frac{(2\mu_x\mu_y + C_1)(2\sigma_{xy} + C_2)}{(\mu_x^2 + \mu_y^2 + C_1)(\sigma_x^2 + \sigma_y^2 + C_2)}
    3. 符号解释:
      • x, y: 待比较的两幅图像的像素值集合(或局部窗口)。
      • μx,μy\mu_x, \mu_y: 图像 xxyy 的平均值。
      • σx2,σy2\sigma_x^2, \sigma_y^2: 图像 xxyy 的方差。
      • σxy\sigma_{xy}: 图像 xxyy 的协方差。
      • C1=(K1L)2,C2=(K2L)2C_1 = (K_1 L)^2, C_2 = (K_2 L)^2: 用于稳定除法的小常数,避免分母接近零。LL 是像素值的动态范围(例如 8 位灰度图像为 255)。K1=0.01,K2=0.03K_1 = 0.01, K_2 = 0.03 是默认值。
  • 浮点运算次数 (Floating Point Operations, FLOPs):
    1. 概念定义: FLOPs 用于衡量模型执行一次前向传播所需的计算量。它反映了模型的计算复杂度,较低的 FLOPs 通常意味着更高的计算效率和更快的推理速度。

    2. 符号解释: 论文使用 fvcore [70] 工具计算 FLOPs,并在 256×256256 \times 256 分辨率下进行测量,以确保与其他 Vision Mamba 方法的公平比较。

      在评估时,PSNRSSIM 均在 RGB 通道上计算。对于 RealSR 数据集,PSNRSSIM 在 YCbCr 颜色空间的 Y 通道上计算。

5.3. 对比基线

论文将 EAMamba 与以下基线模型进行了比较:

  • 图像去噪:

    • 合成高斯去噪: IRCNN [10]、FFDNet [11]、DnCNN [9]、BRDNetBRDNet* [12]、DRUNet [13]、SwinIRSwinIR* [48]、Restormer [51]、MambaIRMambaIR* [57]。
    • 真实世界去噪: DnCNN [9]、BM3D [74]、CBDNetCBDNet* [8]、RIDNetRIDNet* [75]、VDN [76]、SADNetSADNet* [77]、DANetDANet* [78]、CycleISPCycleISP* [79]、MIRNet [35]、DeamNetDeamNet* [80]、MPRNet [36]、DAGL [81]、HINet [37]、IPTIPT* [49]、MAXIM-3S [85]、UFormer-B [50]、Restormer [51]、MambaIR-UNet [57]。
  • 图像超分辨率: Restormer [51]、MambaIR-UNet [57]、VMambaIR [58]。

  • 图像去模糊: DeblurGAN-v2 [16]、SRN [15]、DBGAN [17]、DMPHN [18]、SPAIR [38]、MIMO-UNet+ [19]、MPRNet [36]、HINet [37]、IPT [49]、MAXIM-3S [85]、UFormer-B [50]、Restormer [51]、Stripformer [43]、SFNet [39]。

  • 图像去雾: DehazeNet [28]、AOD-Net [29]、GridDehazeNet [30]、MSBDN [31]、FFA-Net [32]、PFFNet [33]、Dehamer [46]、MAXIM-2S [85]、DehazeFormer-S [47]、DehazeFormer-B [47]、DehazeFormer-L [47]。

    这些基线模型涵盖了 CNN、Transformer 和 Vision Mamba 等不同架构,具有代表性,有助于全面评估 EAMamba 的性能和效率。

5.4. 架构与训练细节

  • 架构细节:

    • EAMamba 采用四级 UNet 架构,在不同级别分别包含 [4, 6, 6, 7]MambaFormer 模块。
    • 精炼阶段包含两个 MambaFormer 模块。
    • 通道维度 CC 保持恒定值 64。
    • 默认的通道 MLP 采用简单的前馈网络 (Feed-Forward Network, FFN) [68]。
  • 训练细节:

    • 总训练迭代次数:450,000 次。

    • 初始学习率:3×1043 \times 10^{-4},通过余弦退火 (cosine annealing) 衰减到 1×1061 \times 10^{-6}

    • 优化器:AdamW [69],参数 β1=0.9,β2=0.999\beta_1 = 0.9, \beta_2 = 0.999,权重衰减 10410^{-4}

    • 损失函数:L1 损失。

    • 渐进式训练策略 [51]:

      • 训练开始时使用 128×128128 \times 128 像素的图像块和批大小 64。
      • 在不同迭代次数逐步调整图像块大小和批大小:
        • 138K 次迭代时:(160,40)(160, 40)
        • 234K 次迭代时:(192,32)(192, 32)
        • 306K 次迭代时:(256,16)(256, 16)
        • 360K 次迭代时:(320,8)(320, 8)
        • 414K 次迭代时:(384,8)(384, 8)
    • 数据增强:随机水平翻转、垂直翻转和 90 度旋转。


6. 实验结果与分析

6.1. 核心结果分析

6.1.1. 图像去噪 (Image Denoising)

EAMamba 在图像去噪任务上的表现,包括合成高斯彩色去噪和真实世界去噪:

合成高斯彩色图像去噪:

  • 量化结果 (Table 1): EAMamba 在 CBSD68Kodak24McMaster 数据集上,在 σ=[15,25,50]\sigma = [15, 25, 50] 不同噪声水平下,实现了与 MambaIRMambaIR* [57] 相当甚至略优的 PSNR 性能。最显著的是,EAMamba 的 FLOPs 仅为 137 G,而 MambaIRMambaIR* [57] 为 1290 G,EAMamba 仅需 MambaIR 计算量的 11%11\%

  • 效率优势: EAMamba 展现了显著的计算效率优势,参数量也较少 (25.3M vs 15.8M for MambaIR*, but MambaIR* uses separate models for each noise level).

    以下是原文 Table 1 的结果:

    Method Param. (M) ↓ FLOPs (G) ↓ CBSD68 [71] Kodak24 [73] McMaster [72]
    σ = 15 σ = 25 σ = 50 σ = 15 σ = 25 σ = 50 σ = 15 σ = 25 σ = 50
    IRCNN [10] - - 33.86 31.16 27.86 34.69 32.18 28.93 34.58 32.18 28.91
    FFDNet [11] - - 33.87 31.21 27.96 34.63 32.13 28.98 34.66 32.35 29.18
    DnCNN [9] - - 33.90 31.24 27.95 34.60 32.14 28.95 33.45 31.52 28.62
    BRDNet* [12] - - 34.10 31.43 28.16 34.88 32.41 29.22 35.08 32.75 29.52
    DRUNet [13] 32.6 144 34.30 31.69 28.51 35.31 32.89 29.86 35.40 33.14 30.08
    SwinIR* [48] 11.5 788 34.42 31.78 28.56 35.34 32.89 29.79 35.61 33.20 30.22
    Restormer [51] 26.1 141 34.39 31.78 28.59 35.44 33.02 30.00 35.55 33.31 30.29
    MambaIR* [57] 15.8 1290 34.43 31.80 31.80 35.34 32.91 29.85 35.62 33.35 30.31
    EAMamba (Ours) 25.3 137 34.43 31.81 28.62 35.36 32.95 29.91 35.59 33.34 30.31

真实世界去噪:

  • 量化结果 (Table 2):SIDD 数据集上,EAMamba 的 FLOPs 为 137 G,相较于 MambaIR-UNet [57] 的 230 G 减少了 41%41\%PSNR 仅有 0.02 dB 的轻微下降 (39.87 dB vs 39.89 dB),SSIM 保持一致。

  • 定性结果 (Figure 7): EAMamba 生成的图像在视觉上更接近真实标注数据 (Ground Truth),显示出其在细节保留和噪声抑制方面的有效性。

    以下是原文 Table 2 的结果:

    Method Param. (M) ↓ FLOPs (G) ↓ SIDD [59]
    PSNR ↑ SSIM↑
    DnCNN [9] - 23.66 0.583
    BM3D [74] 25.65 0.685
    CBDNet* [8] 30.78 0.801
    RIDNet* [75] 1.5 98 38.71 0.951
    VDN [76] 7.8 44 39.28 0.956
    SADNet* [77] - - 39.46 0.956
    DANet* [78] 63.0 30 39.47 0.957
    CycleISP* [79] 2.8 184 39.52 0.957
    MIRNet [35] 31.8 785 39.72 0.959
    DeamNet* [80] 2.3 147 39.47 0.957
    MPRNet [36] 15.7 588 39.71 0.958
    DAGL [81] 5.7 273 38.94 0.953
    HINet [37] 88.7 171 39.99 0.958
    IPT* [49] 115.3 380 39.10 0.954
    MAXIM-3S [85] 22.2 339 39.96 0.960
    UFormer-B [50] 50.9 89 39.89 0.960
    Restormer [51] 26.1 141 40.02 0.960
    MambaIR-UNet [57] 26.8 230 39.89 0.960
    EAMamba (Ours) 25.3 137 39.87 0.960

    该图像是图表,展示了不同图像恢复方法的PSNR值。左侧为真实图像PSNR值为33.11,接下来是低质量图像以及MPRNet(39.62)、UFormer-B(39.73)、Restormer(40.47)、MambaIR-UNet(37.45)等方法的恢复结果,最终是我们的EAMamba方法,PSNR值达到40.99。 该图像是图表,展示了不同图像恢复方法的PSNR值。左侧为真实图像PSNR值为33.11,接下来是低质量图像以及MPRNet(39.62)、UFormer-B(39.73)、Restormer(40.47)、MambaIR-UNet(37.45)等方法的恢复结果,最终是我们的EAMamba方法,PSNR值达到40.99。

Figure 7 展示了 SIDD 数据集上的真实世界去噪定性结果。

6.1.2. 图像超分辨率 (Image Super-Resolution)

EAMamba 在 RealSR 数据集上的超分辨率任务表现:

  • 量化结果 (Table 3): EAMamba 在 FLOPs (137 G) 和参数量 (25.3M) 方面均优于所有 Vision Mamba 基线 (MambaIR-UNet [57], VMambaIR [58]) 和 Restormer [51]。在 ×4\times 4 缩放因子下,EAMamba 取得了最高的 PSNR 性能 (29.60 dB)。在 ×2\times 2×3\times 3 缩放因子下,EAMamba 保持了与 Vision Mamba 方法非常小的 PSNR 差距 (小于 0.05 dB)。

  • 定性结果 (Figure 8): 裁剪的差异结果表明,EAMamba 在结构保留和细节重建方面优于 Vision Mamba 基线。

    以下是原文 Table 3 的结果:

    Method Param. (M) ↓ FLOPs (G) ↓ x2 x3 x4
    PSNR ↑ SSIM ↑ PSNR ↑ SSIM ↑ PSNR ↑ SSIM ↑
    Restormer [51] 26.1 155 34.33 0.929 31.16 0.874 29.54 0.836
    MambaIR-UNet [57] 26.8 230 34.20 0.927 31.16 0.872 29.53 0.835
    VMambaIR [58] 26.3 200 34.16 0.927 31.14 0.872 29.56 0.836
    EAMamba (Ours) 25.3 137 34.18 0.927 31.11 0.872 29.60 0.835

6.1.3. 图像去模糊 (Image Deblurring)

EAMamba 在 GoProHIDE 数据集上的图像去模糊任务表现:

  • 量化结果 (Table 4): EAMamba 在 GoPro 基准上取得了最佳性能,PSNR 为 33.58 dB,超过次优结果 0.31 dB。在 HIDE 基准上排名第二,PSNR 为 31.42 dB。参数量和 FLOPs 也非常有竞争力。

  • 定性结果 (Figure 9): EAMamba 在保留细节方面表现出色,无论从近处物体(如汽车后保险杠)到更广阔的环境(如砖铺路面),其预测结果几乎与真实标注数据无法区分。

    以下是原文 Table 4 的结果:

    Method Param. (M) ↓ FLOPs (G) ↓ GoPro [62] HIDE [86]
    PSNR ↑ SSIM ↑ PSNR ↑ SSIM↑
    DeblurGAN-v2 [16] - - 29.55 0.934 26.61 0.875
    SRN [15] . - 30.26 0.934 28.36 0.915
    DBGAN [17] - - 31.10 0.942 28.94 0.915
    DMPHN [18] 21.7 195 31.20 0.940 29.09 0.924
    SPAIR [38] - - 32.06 0.953 30.29 0.931
    MIMO-UNet+ [19] 16.1 151 32.45 0.957 29.99 0.930
    MPRNet [36] 20.1 760 32.66 0.959 30.96 0.939
    HINet [37] 88.7 171 32.71 0.959 30.32 0.932
    IPT [49] 115.3 380 32.52 - -
    MAXIM-3S [85] 22.2 339 32.86 0.961 32.83 0.956
    UFormer-B [50] 50.9 89 33.06 0.967 30.90 0.953
    Restormer [51] 26.1 141 32.92 0.961 31.22 0.942
    Stripformer [43] 19.7 155 33.08 0.962 31.03 0.940
    SFNet [39] 13.3 125 33.27 0.963 31.10 0.941
    EAMamba (Ours) 25.3 137 33.58 0.966 31.42 0.944

6.1.4. 图像去雾 (Image Dehazing)

EAMamba 在 RESIDE 数据集上的图像去雾任务表现:

  • 量化结果 (Table 5): EAMamba 在 PSNRSSIM 方面取得了领先性能,同时在参数量 (25.3M) 和 FLOPs (137 G) 上保持了良好的平衡,显示出其高效性。

  • 定性结果 (Figure 10): 评估结果表明,EAMamba 在细节保留方面表现优越,与真实标注数据之间的偏差最小。

    以下是原文 Table 5 的结果:

    Method Param. (M) ↓ FLOPs (G) ↓ SOTS-Indoor [63]
    PSNR ↑ SSIM ↑
    DehazeNet [28] - - 29.28 0.923
    AOD-Net [29] - - 29.62 0.929
    GridDehazeNet [30] - - 31.25 0.946
    MSBDN [31] - - 31.69 0.950
    FFA-Net [32] - - 32.06 0.954
    PFFNet [33] - - 32.74 0.960
    Dehamer [46] 34.1 177 32.93 0.962
    MAXIM-2S [85] 10.3 68 33.09 0.964
    DehazeFormer-S [47] 12.5 78 33.15 0.964
    DehazeFormer-B [47] 24.3 157 33.22 0.965
    DehazeFormer-L [47] 40.3 246 33.32 0.966
    EAMamba (Ours) 25.3 137 33.51 0.967

6.1.5. 整体效率对比

  • Figure 1 提供了计算效率与图像质量 (PSNR) 的对比图。EAMamba (红色星号) 在多个图像恢复任务中均表现出优于其他 Vision Mamba 方法 (蓝色圆点) 和现有方法 (黑色圆点) 的计算效率,同时保持了有竞争力的性能。这表明 EAMamba 为基于 Vision Mamba 的图像恢复设定了新的效率前沿。

    Figure 1. Computational efficiency versus image quality across model architectures. Our method (denoted by \(\\cdot\) )demonstrates superior efficiency compared to other Vision Mamba-based methods \(( \\bullet )\) and existing approaches \(( \\bullet )\) . EAMamba establishes a new efficiency frontier for Vision Mamba-based image restoration. 该图像是图表,展示了EAMamba在不同图像恢复任务中的PSNR与计算效率的对比。EAMamba在多项任务中表现出色,取得了优于其他方法的效果,并在计算复杂度和效率之间设置了新的平衡。图中以红星标识EAMamba的位置。

Figure 1. 计算效率与图像质量跨模型架构的对比。

6.2. 效率与各种扫描策略的有效性

6.2.1. 各种扫描策略的有效性可视化 (ERF)

为了验证全方位扫描策略的有效性,论文可视化了不同扫描策略的 ERF 结果 (Figure 11)。

  • 2D 扫描 (Figure 11 (a)): 能够捕获水平和垂直方向上的长距离依赖。

  • 对角线扫描 (Diagonal Scan) 和 锯齿形扫描 (Zigzag Scan) (Figure 11 (b) 和 (c)): 擅长捕获对角线路径上的全局信息。

  • Z-order 扫描 (Figure 11 (d)): 捕获全局信息但存在不连续性。

  • Hilbert 扫描 (Figure 11 (e)): 缺乏全局信息捕获能力。

  • 全方位扫描 (All-around Scan) (Figure 11 (f)): 结合了 2D 扫描对角线扫描,捕获了图像的全局和局部上下文信息,其 ERF 覆盖范围更广,更均匀,特别是在对角线方向上。

    这些结果表明,单一扫描策略只能捕获特定的空间信息。全方位扫描通过结合多种扫描策略(如 2D 扫描对角线扫描)来捕获互补的空间信息,从而实现更大的感受野和更好的局部信息保留。

    Figure 11. Illustration of the ERF results for different scanning strategies, including two-dimensional scan, diagonal scan, zigzag scan, Z-order scan, Hilbert scan, and our all-around scan with reversing and flipping. 该图像是图表,展示了不同扫描策略(如2D扫描、对角线扫描、锯齿形扫描、Z字形扫描、Hilbert扫描和全方位扫描)的ERF结果。ERF结果展示了各方法在不同指标上的表现,以验证EAMamba的优越性。

Figure 11. 不同扫描策略的 ERF 结果,包括二维扫描、对角线扫描、锯齿形扫描、Z 字形扫描、Hilbert 扫描以及我们的全方位扫描(带有反向和翻转)。

6.2.2. 扫描策略的量化结果

  • Table 6 比较了不同扫描策略在 SIDD 数据集上的 PSNR 表现。全方位扫描策略实现了 39.87 dB 的 PSNR,比其他单一扫描策略高出 0.070.15 dB0.07 - 0.15 \ \mathrm{dB}。这进一步证明了其在图像恢复任务中的优势。

    以下是原文 Table 6 的结果:

    2D Diagonal Zigzag Z-order Hilbert All-around
    39.80 39.79 39.77 39.74 39.74 39.87
  • Table 7 评估了不同扫描策略组合在多个图像恢复数据集上的 PSNR 表现。结果显示,2D + Diagonal 组合通常能提供良好的性能,并被设为默认配置。例如,在 SIDDRealSRx4 上,2D + Diagonal 取得了最佳或接近最佳的性能。在 SOTS-Indoor 上,2D + Diagonal + Z-order 取得了最佳性能。这表明全方位扫描机制不仅通过整合多种扫描策略解决了图像恢复任务,还提供了灵活性,可以通过 MHSSM 无缝集成新的扫描策略。

    以下是原文 Table 7 的结果:

    Dataset 2D + Diagonal 2D + Z-order 2D + Hilbert 2D + Diagonal + Z-order
    SIDD [59] 39.87 39.82 39.83 39.83
    RealSRx4 [61] 29.60 29.58 29.51 29.57
    GoPro [62] 33.58 33.51 33.66 33.56
    SOTS-Indoor [63] 43.19 43.20 43.07 43.37

6.3. 消融实验 (Ablation Studies)

6.3.1. MHSS 和全方位扫描的有效性

  • Table 8 量化了 EAMamba 各个设计元素的贡献。
    • Baseline: 使用 2DSSM [57] 和 2D 扫描

    • + MHSSM2DSSM 替换为 MHSSM 后,FLOPs 减少了一半 (从 286 G 降至 137 G),但 PSNR 仅有 0.1% 的轻微下降。这证明了 MHSSM 在显著节省计算成本的同时,对图像恢复质量的影响微乎其微。

    • + all-around scanMHSSM 的基础上,进一步引入全方位扫描PSNR 得到提升,表明全方位扫描相较于传统的 2D 扫描能够提供更高的图像质量。

      以下是原文 Table 8 的结果:

      Method Param. (M) ↓ FLOPs (G) ↓ Urban100 [87]
      σ = 15 σ = 25 σ = 50
      Baseline 31.1 286 35.15 33.00 30.08
      + MHSSM 25.3 137 35.06 32.89 29.95
      + all-around scan 25.3 137 35.10 32.93 30.01

这些结果有力地支持了 EAMamba 的设计选择,突显了其在效率和恢复性能之间取得平衡的能力。

6.3.2. 不同通道 MLP 的比较

  • Table 9 评估了 MambaFormer Block 中各种通道 MLP 模块的性能,包括 vanilla FFN [88]、Gated-Dconv FFN (GDFN) [51]、Simple FFN [68] 和 Channel Attention (CA) [6]。
    • 无 MLP: 移除通道 MLP 会导致 PSNR 下降超过 0.1 dB,表明 Channel MLP 对于性能提升是必要的。

    • GDFN: 表现出最佳性能,但参数量和 FLOPs 相对较高。

    • Simple FFN: 在性能和计算效率之间实现了最佳平衡,因此被选作默认配置。

      以下是原文 Table 9 的结果:

      Channel MLP Param. (M) ↓ FLOPs (G) ↓ Urban100 [87]
      σ = 15 σ = 25 σ = 50
      None 16.5 90 34.98 32.79 29.82
      FFN 28.3 153 35.10 32.93 30.01
      GDFN 34.5 189 35.15 32.98 30.08
      Simple FFN 25.3 137 35.10 32.93 30.01
      CA 28.3 123 35.05 32.88 29.95

7. 总结与思考

7.1. 结论总结

本文提出了 EAMamba,一个增强型视觉 Mamba 框架,旨在解决现有视觉 Mamba 方法在低级视觉任务中面临的计算复杂度和局部像素遗忘问题。EAMamba 引入了两项关键的架构创新:

  1. 多头选择性扫描模块 (MHSSM): 通过通道分割处理和高效的序列聚合,MHSSM 显著提高了视觉 Mamba 框架的可扩展性和计算效率,避免了计算复杂度和参数数量随扫描序列增加而线性增长的弊端。

  2. 全方位扫描机制 (All-Around Scanning Mechanism): 这一机制超越了传统的二维扫描策略,整合了水平、垂直、对角线和翻转对角线及其反向等多个扫描方向,从而能够捕获更全面的空间信息,有效解决了局部像素遗忘问题,增强了模型对图像空间关系的理解能力。

    通过在超分辨率、去噪、去模糊和去雾等多种图像恢复任务上的广泛实验评估,EAMamba 的这些架构创新得到了验证。定性和定量结果均表明,EAMamba 在实现与现有低级视觉 Mamba 方法相当的性能的同时,显著减少了参数数量和浮点运算 (FLOPs)(降低了 31-89%)。这使得 EAMamba 在图像恢复领域取得了性能与效率的良好平衡,为未来的研究奠定了坚实基础。

7.2. 局限性与未来工作

论文中未明确指出 EAMamba 自身的具体局限性。然而,从其设计和实验结果可以推断一些潜在方面:

  • 扫描策略选择的泛化性: 论文提到 2D + Diagonal 组合通常表现良好并作为默认配置,但其他组合在特定数据集上可能表现更好(如 SOTS-Indoor2D + Diagonal + Z-order 最佳)。这可能意味着对于不同的图像恢复任务或数据集,最佳的全方位扫描组合可能需要手动调整或通过更复杂的自适应机制来确定,增加了部署的复杂性。

  • 计算资源需求: 尽管 EAMamba 显著降低了 FLOPs,但作为深度学习模型,特别是在处理高分辨率图像时,仍然需要相当的计算资源(如 GPU)进行训练和推理。对于极端资源受限的边缘设备,可能仍需进一步优化。

  • 理论分析的深度: 论文主要侧重于实证验证,对 MHSS全方位扫描为何能有效解决局部像素遗忘计算效率问题的更深层理论分析(例如,它们如何改变特征表示或信息流)可以进一步加强。

    至于未来工作,论文中提到全方位扫描机制的灵活性,可以无缝集成新的扫描策略。这暗示了未来可以探索:

  • 自适应扫描策略: 开发能够根据图像内容或任务需求自适应选择或组合扫描方向的机制。

  • 更复杂的扫描模式: 研究除了文中提及的几种基本扫描模式之外,更复杂、更高效的扫描路径,以进一步提升信息捕获能力。

  • 与其他先进技术的融合: 将 EAMamba 的高效架构与最新的自监督学习、生成模型或其他图像处理技术相结合,以探索更高质量或更通用的图像恢复方案。

7.3. 个人启发与批判

EAMamba 的工作给我带来了几点启发和思考:

  1. 效率与性能的协同: 这篇论文再次强调了在深度学习模型设计中,效率和性能并非总是相互排斥的。通过巧妙的架构创新(如 MHSSM),可以在大幅降低计算成本的同时,维持甚至提升性能。这对于将先进模型部署到实际应用和资源受限环境中至关重要。

  2. 局部与全局信息融合的重要性: 图像恢复任务对局部细节和全局上下文都高度敏感。EAMamba 强调并有效解决了 Vision Mamba 中局部像素遗忘的问题,通过全方位扫描同时捕获局部和全局信息,这对于提升图像质量具有决定性作用。这启发我们在设计其他视觉任务模型时,也应充分考虑如何更全面、更细致地处理多尺度和多方向的信息流。

  3. Mamba 模型在视觉领域的潜力: 论文进一步证明了 Mamba 架构在视觉任务中的巨大潜力,尤其是在处理长距离依赖和保持计算效率方面。EAMamba 的成功可能会激励更多研究者深入探索 SSM 在图像生成、图像理解等更广泛视觉任务中的应用。

  4. 模块化设计的优势: EAMamba 的 MambaFormer 模块和其中的 MHSSM全方位扫描等设计体现了良好的模块化思想。这种模块化使得模型易于理解、修改和扩展,也便于进行消融实验来验证各组件的贡献。

    批判性思考方面:

  • 默认配置的解释: 论文将 Simple FFN 作为默认的 Channel MLP 配置,理由是其在性能和效率之间达到最佳平衡。虽然定量结果支持这一选择,但如果能更深入地分析 GDFN 等更优性能 MLP 为什么会带来更高的计算成本,以及这种成本在实际应用中的可接受程度,将更有助于读者理解权衡。

  • 超参数敏感性: 论文中对学习率调度、批大小和图像块大小的渐进式训练策略进行了详细描述。虽然这能带来好的结果,但也可能意味着模型对这些超参数比较敏感。如果能提供更多关于模型对这些超参数鲁棒性的分析,将更有价值。

  • 更复杂的退化场景: 论文主要在噪声、模糊、下采样和雾等常见退化类型上进行实验。对于更复杂、更现实、更混合的图像退化场景,EAMamba 的性能和鲁棒性如何,值得进一步探索。例如,同时包含多种退化的真实世界图像恢复,或者泛化到未见过的退化类型。

    总的来说,EAMamba 为高效图像恢复提供了一个强大且设计精巧的 Vision Mamba 框架,其提出的 MHSSM全方位扫描创新点,为 Mamba 架构在低级视觉任务中的应用树立了新的标杆。

相似论文推荐

基于向量语义检索推荐的相关论文。

暂时没有找到相似论文。