Efficient Visual State Space Model for Image Deblurring
论文信息
题目: Efficient Visual State Space Model for Image Deblurring
用于图像去模糊的高效视觉状态空间模型
源码:https://github.com/kkkls/EVSSM
创新点
- 提出了高效视觉状态空间模型EVSSM,解决了Transformer模型计算复杂度与图像分别率呈现二次增长的问题
- 提出了高效扫描块(EVS),使用几何变换,影响扫描过程,减小复杂度。
- 提出了 基于高效判别频域的前馈网络(EDFFN)
摘要
卷积神经网络(CNNs)
和视觉Transformer(ViTs)
在图像恢复任务中取得了优异的性能。由于ViTs能够捕捉长距离依赖关系和输入相关特征,因此在图像恢复方面通常比CNNs表现更出色
。然而,基于Transformer的模型计算复杂度会随着图像分辨率的增加呈二次增长
,这限制了它们在高分辨率图像恢复任务中的实际应用。在本文中,我们提出了一种简单而有效的视觉状态空间模型(EVSSM)用于图像去模糊
,该模型将状态空间模型(SSMs)的优势应用于视觉数据。与现有方法采用多个固定方向扫描进行特征提取,从而显著增加计算成本不同,我们开发了一种高效的视觉扫描模块,在每个基于SSM的模块之前应用各种几何变换,能够在保持高效率的同时捕捉有用的非局部信息。大量实验结果表明,在基准数据集和实际拍摄的图像上,所提出的EVSSM与最先进的图像去模糊方法相比表现出色
。
1 引言
1.1 研究背景
图像去模糊旨在从模糊图像中恢复出清晰图像,随着各种相机和手持成像设备的普及,这一任务受到了广泛关注。由于在实际情况中,只有模糊图像可用,而模糊核和清晰图像的信息缺失,使得该任务极具挑战性。
1.2 研究现状与挑战
深度学习的发展推动了基于深度卷积神经网络(CNNs)的图像去模糊方法取得显著进展。然而,卷积操作在空间上具有不变性和局部性,无法捕捉图像内容的空间可变属性,也难以探索对去模糊有益的非局部信息
。
相比之下,Transformer中的自注意力机制能够通过计算每个标记与所有其他标记之间的相关性来捕捉全局信息
,为图像去模糊提取更好的特征。但自注意力机制(即缩放点积注意力)的空间和时间复杂度与标记数量呈二次关系,在处理高分辨率图像时,计算成本变得难以接受。虽然基于局部窗口的方法、转置注意力和基于频域的近似方法被提出以降低计算成本,但这些方法牺牲了对非局部信息和空间信息的建模能力,进而影响了恢复图像的质量。
近年来,状态空间模型(SSMs)
在自然语言处理(NLP)任务中展现出强大的长距离依赖建模能力,且具有线性或接近线性的计算复杂度
。改进后的SSM,特别是Mamba,开发了一种选择性扫描机制
(S6),能够在保持线性计算复杂度的同时记住相关信息并忽略无关内容。这启发我们利用Mamba高效地探索有用的非局部信息,以实现更好的图像去模糊效果。然而,Mamba是为处理一维(1D)序列而设计的,如果直接将其应用于视觉任务,需要先将图像数据展平为一维图像序列,这会破坏图像的空间结构,难以从相邻像素中捕捉局部信息。一些方法采用多方向扫描机制将状态空间模型应用于视觉领域,但多方向扫描机制显著增加了计算成本。
1.3 研究内容与贡献
本文提出了一种高效的视觉状态空间模型用于图像去模糊。我们发现现有视觉状态空间模型大多采用固定方向扫描进行特征提取,这种方式可能无法自适应地探索非局部信息,并且会导致较高的计算成本。因此,我们开发了一种简单有效的扫描策略,在保持低计算成本的同时捕捉非局部空间信息。具体而言,我们仅在一个方向上扫描输入特征,但在每次扫描前应用简单的几何变换,以最小的计算成本增加有效地自适应探索有用信息。
本文的主要贡献如下:
- 第一,提出了一种简单而有效的视觉状态空间模型,能够高效恢复高质量图像;
- 第二,开发了一种高效的扫描策略,在保持低计算成本的同时捕捉非局部空间信息;
- 第三,在基准数据集和实际图像上对所提方法进行了定量和定性评估,结果表明其在性能和效率上优于现有方法。
2 相关工作
2.1 基于深度卷积神经网络的图像去模糊方法
近年来,基于深度卷积神经网络的图像去模糊方法取得了显著进展。一些研究提出了基于多尺度框架的深度CNN,直接从模糊图像中估计清晰图像;引入了有效的尺度循环网络,以增强多尺度框架中各尺度信息的利用;提出了选择性网络参数共享方法,进一步改进了之前的方法。此外,生成对抗网络(GANs)也被广泛应用于图像去模糊,旨在通过生成逼真清晰的图像来提高去模糊结果的质量。
然而,由于利用额外尺度带来的性能提升有限,一些研究提出了采用多补丁策略的网络,使去模糊过程逐步执行,网络能够逐步优化输出;引入了跨阶段特征融合技术,以增强去模糊方法的整体性能;提出了多输入多输出网络架构,在保持去模糊性能的同时降低计算负担;通过分析基线模块并提出简化版本,提高了图像恢复的效率。但由于卷积操作的局限性,这些方法无法有效建模全局和空间可变信息,限制了其实现更好图像恢复的能力。
2.2 基于Transformer的图像去模糊方法
Transformer能够建立长距离依赖关系并有效建模全局信息,在图像分类、目标检测和语义分割等高级视觉任务中取得了显著进展。研究人员将其应用扩展到图像超分辨率、图像去模糊和图像去噪等领域。但Transformer的自注意力机制计算复杂度为二次方,对于高分辨率图像恢复任务来说难以接受。
为降低Transformer的计算复杂度,一些研究提出了在特征深度域计算缩放点积注意力的高效Transformer模型;通过构建条带内和条带间标记简化自注意力计算;基于UNet架构引入Transformer,应用非重叠窗口自注意力进行单图像去模糊;提出了基于频域的Transformer并取得了领先的结果。尽管这些方法采用了各种策略来降低计算复杂度,但它们无法有效表征长距离依赖关系和非局部信息。相比之下,本文开发了一种高效的视觉状态空间模型,能够以较低的计算成本探索有用的非局部信息。
2.3 状态空间模型
状态空间模型长期以来一直是时间序列分析和建模的基石。近年来,一些方法采用状态空间模型来捕捉序列建模中的长距离依赖关系,基于状态空间模型的方法可以通过递归或卷积高效计算,具有线性或接近线性的计算复杂度。一些研究提出了相关框架和方法,以解决状态表示带来的高计算和内存需求问题。Mamba引入了具有动态权重的选择性扫描层,在自然语言处理中展现出巨大潜力。
为将状态空间模型应用于视觉任务,一些方法采用多方向扫描策略,但这会增加计算成本。相比之下,本文提出了一种高效的视觉扫描模块,在每次扫描前应用几何变换,以高效地探索非局部信息。
3 高效视觉状态空间模型
3.1 整体架构
本文提出的高效视觉状态空间模型(EVSSM)的整体架构基于分层编码器 - 解码器框架。给定模糊图像 I b l u r ∈ R H × W × 3 I_{blur} \in \mathbb{R}^{H×W×3} Iblur∈RH×W×3,首先使用一个 3 × 3 3×3 3×3卷积层获得浅层特征 F s ∈ R H × W × C F_{s} \in \mathbb{R}^{H×W×C} Fs∈RH×W×C,其中 H × W H×W H×W表示空间维度, C C C是特征通道数。
然后,浅层特征 F s F_{s} Fs被输入到一个3级对称编码器 - 解码器网络中。每个级别的编码器/解码器由多个高效视觉状态空间(EVSS)模块组成。对于第(l)级的编码器/解码器,输入特征通过每个EVSS模块逐步处理,生成中间特征 F e n c l / F d e c l ∈ R H 2 l − 1 × W 2 l − 1 × 2 l − 1 C F_{enc}^{l} / F_{dec}^{l} \in \mathbb{R}^{\frac{H}{2^{l - 1}}×\frac{W}{2^{l - 1}}×2^{l - 1}C} Fencl/Fdecl∈R2l−1H×2l−1W×2l−1C(在本文中 l = 1 , 2 , 3 l = 1,2,3 l=1,2,3)。接着,使用双线性插值和 1 × 1 1×1 1×1卷积进行上采样和下采样,并在每个级别编码器和解码器之间添加跳跃连接。
最后,对特征 F d e c 3 F_{dec}^{3} Fdec