跳转至

Efficient Visual State Space Model for Image Deblurring

会议: CVPR 2025
arXiv: 2405.14343
代码: https://github.com/kkkls/EVSSM
领域: 图像复原
关键词: 图像去模糊、状态空间模型、Mamba、频域前馈网络、高效扫描

一句话总结

本文提出一种高效视觉状态空间模型 EVSSM,通过几何变换替代多方向扫描策略捕获非局部信息,并设计高效频域前馈网络增强局部细节,在图像去模糊任务上以仅四分之一的计算代价超越现有 SSM 方法,达到 SOTA 效果。

研究背景与动机

领域现状:图像去模糊是低层视觉的经典任务。CNN 方法受限于局部感受野,无法有效建模空间变化特性和非局部信息。Transformer 通过自注意力机制可捕获全局依赖,但其二次复杂度在高分辨率图像上计算成本极高。为降低复杂度,现有 Transformer 方法采用局部窗口注意力、转置注意力等近似策略,但这些近似牺牲了对非局部信息的建模能力。

现有痛点:状态空间模型(SSM)如 Mamba 以线性复杂度建模长程依赖,但其设计面向 1D 序列。将 2D 图像展平为 1D 序列会破坏空间结构。现有视觉 SSM 方法(如 VMamba)采用多方向扫描策略来弥补这一缺陷,但多方向扫描使计算量成倍增加(VMamba 是原始 Mamba 的 4 倍),显著降低效率。

核心矛盾:想要用 SSM 的线性复杂度优势探索非局部信息,但多方向扫描又把这一优势消耗殆尽。同时,SSM 主要关注全局信息,对图像复原至关重要的局部细节则需要额外机制来增强。

本文目标 (1) 如何在最小计算开销下让 SSM 有效处理 2D 视觉数据,探索多方向非局部信息;(2) 如何高效增强局部频域信息以辅助清晰图像重建。

切入角度:作者观察到只需在每次扫描前对输入特征施加简单的几何变换(翻转、转置),就可以让单方向扫描在不同模块中自动覆盖不同方向的信息。几何变换本身几乎不增加计算成本,但能让后续的选择性扫描"看到"来自不同空间方向的上下文。

核心 idea:用交替几何变换+单向扫描代替多方向扫描,以极低开销实现多方向非局部信息探索。

方法详解

整体框架

EVSSM 采用 3 层对称的编码器-解码器架构。输入模糊图像经 3×3 卷积提取浅层特征后,进入多层级编码解码器。每层由若干高效视觉状态空间(EVSS)模块组成。编码器通过双线性插值+1×1 卷积实现下采样,解码器做上采样,层间有跳跃连接。最终输出残差图像,加回输入得到去模糊结果。通道数从 48 开始,随层级翻倍;各层 EVSS 模块数为 [6, 6, 12]。

关键设计

  1. 高效视觉扫描块(EVS Block):

    • 功能:以最小计算开销在单向扫描中探索多方向非局部信息
    • 核心思路:在每个 EVS 块开头,根据模块索引 \(i\) 对输入做不同的几何变换——偶数模块做转置(Transpose),奇数模块做翻转(Flip)。由于卷积具有平移不变性,几何变换不影响卷积本身,只影响后续选择性扫描的信息流向。每经过 4 个 EVSS 模块,特征自动恢复到原始空间结构。此外,在线性投影后对 SSM 的参数 \(B\)\(C\)\(\Delta\) 各施加核大小为 7 的 1D 深度可分离卷积,使这些参数能编码不同的空间信息,而非仅依赖同一线性变换输出。由于存在几何变换,1D 卷积实际上能聚合来自原始 2D 输入多方向的信息。
    • 设计动机:VMamba 的四方向扫描计算量是 Mamba 的 4 倍,而 EVS 仅用一次几何变换+单向扫描就达到类似效果,几何变换成本可忽略。对 \(B\)\(C\)\(\Delta\) 加 1D 深度卷积是为了打破它们仅从同一线性投影得到的"空间同质性",让模型捕获更丰富的局部-全局模式。
  2. 高效判别性频域前馈网络(EDFFN):

    • 功能:增强局部频率信息,辅助清晰图像重建
    • 核心思路:原始 FFTformer 的 DFFN 在 FFN 中间对扩展后的高维特征做 FFT 学习量化矩阵,计算量巨大(中间特征通道数为输入的 3 倍)。EDFFN 将频域筛选移到 FFN 末端、通道数已恢复原始大小处,用更小的可学习量化矩阵 \(W\) 在频域中筛选有用频率信息,大幅降低 FFT 的计算开销。
    • 设计动机:SSM 主要关注全局信息,局部纹理细节需要额外机制补充。在 FFN 末端而非中间做频域操作,可以在不损失性能的前提下显著减少运行时间。

损失函数 / 训练策略

采用像素空间 L1 损失与频域 L1 损失的加权组合:\(\mathcal{L} = \|I_{deblur} - I_{gt}\|_1 + \lambda \|\mathcal{F}(I_{deblur}) - \mathcal{F}(I_{gt})\|_1\),其中 \(\lambda = 0.1\)。训练采用渐进策略:先在 128×128 patch 上以 batch 64 训练 300K 步,再在 256×256 patch 上以 batch 16 训练 300K 步。使用 AdamW 优化器和余弦退火调度。

实验关键数据

主实验

数据集 指标 EVSSM FFTformer NAFNet GRL 提升
GoPro PSNR 34.51 34.21 33.71 33.93 +0.30
GoPro SSIM 0.9713 0.9692 0.9668 0.9680 +0.0021
HIDE PSNR 31.99 31.62 31.31 31.65 +0.34
RealBlur-R PSNR 41.27 40.11 - - +1.16
RealBlur-J PSNR 34.34 32.62 - - +1.72

EVSSM 在所有基准上均取得最佳 PSNR 和 SSIM。在 RealBlur 真实数据集上优势尤其明显,PSNR 提升超过 1 dB。

消融实验

配置 GoPro PSNR 说明
完整 EVSSM 34.51 完整模型
无几何变换 ~33.8 退化为纯单向扫描,丢失多方向信息
无 1D DWConv on B/C/Δ ~34.1 参数空间表达力下降
DFFN 替代 EDFFN ~34.3 性能接近但计算量大幅增加

关键发现

  • EVSSM 的计算量仅为 VMamba 等多方向扫描 SSM 的约四分之一,但去模糊性能更优
  • 几何变换策略是最关键的设计——零成本实现多方向信息探索
  • EDFFN 在保持频域判别能力的同时,相比 DFFN 显著降低运行时间
  • 模型在真实场景(RealBlur)提升特别大,说明非局部信息对真实模糊的恢复尤其重要
  • EVSSM 还展现出良好的泛化性,可直接迁移到去雨(+0.46 dB)和去雾(+0.60 dB)任务

亮点与洞察

  • 几何变换替代多方向扫描是本文最精彩的设计。翻转和转置几乎零成本,却让单向 SSM 等效获得多方向信息——这个思路优雅且通用,可以迁移到任何基于 SSM 的视觉任务
  • 在 FFN 末端做频域筛选的改进虽然简单,但抓住了计算瓶颈的关键——中间层通道数膨胀导致 FFT 成本激增,挪到末端就解决了
  • 1D 卷积+几何变换的组合使得 1D 操作实际上覆盖了 2D 空间的多方向信息,这个insight值得注意

局限与展望

  • 几何变换仅限于翻转和转置两种,是否可以引入更复杂的可学习空间变换进一步提升信息多样性
  • 未在更大分辨率(如 4K/8K)上验证效率优势的实际表现
  • 训练策略的渐进式增大 patch 对最终性能的贡献未充分消融
  • 频域 FFN 的量化矩阵 \(W\) 的学习机制和作用可以进一步分析

相关工作与启发

  • vs VMamba: VMamba 四方向扫描计算量为 4 倍 Mamba,EVSSM 用几何变换+单向扫描实现类似效果但仅约 1 倍开销,效率优势显著
  • vs FFTformer: 同样使用频域信息,但 EVSSM 将频域筛选放在 FFN 末端而非中间,计算效率更高且性能不降
  • vs NAFNet: NAFNet 追求极简 CNN 设计,EVSSM 引入 SSM 实现了更强的长程信息建模

评分

  • 新颖性: ⭐⭐⭐⭐ 几何变换替代多方向扫描的idea简洁有效但不算颠覆性
  • 实验充分度: ⭐⭐⭐⭐⭐ 多个去模糊基准+跨任务泛化+消融分析,非常全面
  • 写作质量: ⭐⭐⭐⭐ 结构清晰,动机和方法描述明确
  • 价值: ⭐⭐⭐⭐ 为 SSM 在图像复原中的高效使用提供了实用方案

Efficient Visual State Space Model for Image Deblurring

会议: CVPR 2025
arXiv: 2405.14343
代码: https://github.com/kkkls/EVSSM
领域: 图像复原
关键词: 图像去模糊、状态空间模型、Mamba、频域FFN、几何变换扫描

一句话总结

本文提出 EVSSM,通过在单方向 SSM 扫描前施加交替的几何变换(转置/翻转)来高效捕获非局部信息,并设计高效判别性频域 FFN (EDFFN) 增强局部细节,在图像去模糊任务上以仅 1/4 的计算量超越了现有 SSM 方法并达到 SOTA。

研究背景与动机

领域现状:图像去模糊旨在从模糊图像中恢复清晰图像,主流方法分为 CNN 和 Transformer 两大类。CNN 方法受限于卷积操作的空间不变性和有限感受野,难以捕获空间变化特性和非局部信息。Transformer 通过自注意力机制建模全局依赖,效果更好,但其计算复杂度与 token 数量呈二次方关系,处理高分辨率图像时成本不可接受。

现有痛点:为降低 Transformer 的计算开销,现有方法采用局部窗口注意力、转置注意力、频域近似等策略,但这些方法在降低计算量的同时牺牲了对非局部信息或空间信息的建模能力,限制了恢复质量。近期状态空间模型 (SSM/Mamba) 展示了以线性复杂度建模长程依赖的潜力,但现有视觉 SSM 方法(如 VMamba)采用多方向扫描机制,计算成本是单方向的 4 倍,效率增益被大幅抵消。

核心矛盾:如何在保持线性计算复杂度的前提下,让 SSM 有效地探索二维图像中的非局部信息?多方向扫描虽然覆盖全面但计算代价过高,而单方向扫描又无法充分利用空间结构。

本文目标 (1) 设计高效的视觉扫描策略,在不显著增加计算量的情况下捕获多方向的非局部信息;(2) 解决 SSM 参数 B/C/Δ 从相同线性变换导出、空间信息单一的问题;(3) 降低频域 FFN 的计算开销,同时保持局部细节增强能力。

切入角度:作者观察到,与其在多个方向重复扫描,不如在每次扫描前对输入特征进行简单的几何变换(转置或翻转),这样单方向扫描就能自动覆盖不同方向的信息。由于卷积具有平移不变性,几何变换不影响卷积本身,仅改变选择性扫描的行为。

核心 idea:用交替几何变换 + 单方向扫描代替多方向扫描,以近乎零开销实现多方向非局部信息探索。

方法详解

整体框架

EVSSM 采用经典的三级对称编码器-解码器架构。输入模糊图像 \(I_{blur} \in \mathbb{R}^{H \times W \times 3}\) 首先通过 3×3 卷积提取浅层特征 \(F_s \in \mathbb{R}^{H \times W \times C}\)(C=48),送入三级编码器-解码器。每级编码器/解码器由若干个 EVSS 模块堆叠而成(各级数量为 [6, 6, 12]),级间通过双线性插值和 1×1 卷积实现上下采样,并加入跳跃连接。最终通过 3×3 卷积输出残差图像 R,加上输入得到去模糊结果 \(I_{deblur} = R + I_{blur}\)

关键设计

  1. 高效视觉扫描块 (EVS Block):

    • 功能:以最小计算代价实现多方向非局部信息探索
    • 核心思路:对于第 \(i\) 个 EVSS 模块,在扫描前根据模块索引交替施加几何变换:\(i \% 2 = 0\) 时做特征转置,\(i \% 2 = 1\) 时做水平+垂直翻转。这样每 4 个 EVSS 模块自动恢复到原始空间结构。变换后,通过线性层拆分为两个分支 \(X_1, X_2\)\(X_1\) 经 3×3 深度卷积和 SiLU 激活后送入 S6 选择性扫描,\(X_2\) 作为门控信号与扫描结果相乘输出
    • 设计动机:几何变换几乎零开销(仅涉及内存重排),却能让单方向扫描在不同模块中"看到"不同方向的信息排列,等效实现了多方向扫描的效果
  2. 1D 深度卷积增强 SSM 参数多样性:

    • 功能:使 SSM 的 B、C、Δ 参数编码不同的空间信息
    • 核心思路:在线性投影导出 B、C、Δ 后,分别对每个参数施加 kernel size=7 的 1D 深度卷积。由于前序有几何变换,1D 卷积实际上在原始 2D 输入上聚合了多方向信息,使各参数具有差异化的空间表征
    • 设计动机:原始 Mamba 中 B、C、Δ 都由相同输入的线性变换得到,编码了相同的空间信息,限制了模型捕获多样化空间模式的能力。加入 1D 深度卷积后各参数有了独立的局部关注能力
  3. 高效判别性频域 FFN (EDFFN):

    • 功能:增强 SSM 未充分覆盖的局部细节信息
    • 核心思路:在 FFN 末端(而非中间)对特征执行 FFT 并学习一个量化矩阵 W,自适应筛选需要保留的频率信息。由于 FFN 末端特征通道数远小于中间层(原 DFFN 在 3 倍通道扩展的中间层做 FFT),计算开销大幅降低
    • 设计动机:FFTformer 的 DFFN 在 FFN 中间做 FFT,此时通道数是输入的 3 倍,导致 FFT 计算量巨大。将频域筛选移至 FFN 末端,在不影响性能的前提下显著降低计算时间

损失函数 / 训练策略

训练损失由像素域 L1 损失和频域 L1 损失组成:\(\mathcal{L} = \|I_{deblur} - I_{gt}\|_1 + 0.1 \|\mathcal{F}(I_{deblur}) - \mathcal{F}(I_{gt})\|_1\)。采用渐进式训练:先用 128×128 patch + batch 64 训练 300K 迭代,再切换到 256×256 patch + batch 16 继续 300K 迭代,均使用 AdamW 优化器和余弦退火策略。

实验关键数据

主实验

数据集 方法 PSNR (dB) SSIM
GoPro FFTformer 34.21 0.9692
GoPro GRL 33.93 0.9680
GoPro EVSSM 34.51 0.9713
HIDE FFTformer 31.62 0.9455
HIDE GRL 31.65 0.9470
HIDE EVSSM 31.99 0.9503
RealBlur-R FFTformer 40.11 0.9753
RealBlur-R EVSSM 41.27 0.9776
RealBlur-J FFTformer 32.62 0.9326
RealBlur-J EVSSM 34.34 0.9456

消融实验

配置 计算特点 说明
VMamba 四方向扫描 4× 计算量 现有 SSM 视觉方法的通用做法
EVSSM 几何变换+单方向 1× 计算量 仅增加几何变换的微小开销
CU-Mamba (另一 SSM 方法) GoPro PSNR 33.53 明显低于 EVSSM 的 34.51

此外在去雨(PSNR 49.00 vs Restormer 47.98)和去雾(PSNR 32.05 vs DehazeFormer 31.45)任务上也验证了 EVSSM 的泛化性。

关键发现

  • 几何变换策略是核心贡献:以几乎为零的计算代价,实现了等效于多方向扫描的信息覆盖
  • EDFFN 将频域筛选移至 FFN 末端,在保持相同性能的前提下显著减少运行时间
  • 1D 深度卷积对 B、C、Δ 参数的增强也有独立贡献,使得扫描过程中参数具有空间差异性
  • 在真实模糊数据集 RealBlur 上提升特别明显(+1.16 dB / +1.72 dB),体现了对真实退化的鲁棒性

亮点与洞察

  • 几何变换替代多方向扫描是本文最精妙的设计——利用转置/翻转这种零成本操作改变了信息在序列中的相对位置,让单方向扫描在不同层中等效扫描了不同空向,是一种非常优雅的工程化思路
  • 频域 FFN 的位置优化(从中间移至末端)是一个简单但高效的改进思路,具有广泛的可迁移性——任何在特征中间层做昂贵变换的网络都可以考虑将变换后移到通道数更小的位置
  • EVSSM 在去模糊之外还成功应用于去雨和去雾,说明该架构具有通用的图像复原能力

局限与展望

  • 几何变换策略虽然简洁,但变换模式固定(交替转置和翻转),未探索自适应选择变换类型的可能性
  • 文中未充分分析几何变换对不同尺度特征的影响差异,也未与可学习的扫描方向进行对比
  • EDFFN 将频域筛选移至末端虽然降低了成本,但可能也损失了一些在高通道特征上的表达能力,这一点缺乏消融验证

相关工作与启发

  • vs FFTformer: FFTformer 使用频域 Transformer,EVSSM 用 SSM 替代注意力机制实现线性复杂度,同时改进了 FFTformer 的 DFFN 组件。EVSSM 在所有基准上均超过 FFTformer
  • vs VMamba/CU-Mamba: 这些方法采用多方向扫描,EVSSM 通过几何变换实现等效效果但仅需 1/4 计算量
  • vs NAFNet/Restormer: CNN/Transformer 方法难以高效建模非局部信息,EVSSM 的线性复杂度优势在高分辨率场景下尤为突出

评分

  • 新颖性: ⭐⭐⭐⭐ 几何变换替代多方向扫描的思路新颖且优雅,但整体框架仍沿用标准编码器-解码器
  • 实验充分度: ⭐⭐⭐⭐ 在多个去模糊基准和其他复原任务上验证,但消融实验可以更详细
  • 写作质量: ⭐⭐⭐⭐ 逻辑清晰,动机推导合理,图表展示清楚
  • 价值: ⭐⭐⭐⭐ 为视觉 SSM 的高效化提供了简洁实用的方案,对相关领域有参考价值

相关论文