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} Fdec3应用一个 3 × 3 3×3 3×3卷积层，生成残差图像 R ∈ R H × W × 3 R \in \mathbb{R}^{H×W×3} R∈RH×W×3。恢复后的图像 I d e b l u r I_{deblur} Ideblur通过 I d e b l u r = R + I b l u r = N ( I b l u r ) + I b l u r I_{deblur}=R + I_{blur}=N(I_{blur}) + I_{blur} Ideblur=R+Iblur=N(Iblur)+Iblur获得，其中 N N N表示通过最小化以下损失函数进行正则化的编码器 - 解码器网络：
L = ∥ I d e b l u r − I g t ∥ 1 + λ ∥ F ( I d e b l u r ) − F ( I g t ) ∥ 1 \mathcal{L}=\left\| I_{deblur }-I_{gt}\right\| _{1}+\lambda\left\| \mathcal{F}\left(I_{deblur }\right)-\mathcal{F}\left(I_{gt}\right)\right\| _{1} L=∥Ideblur−Igt∥1+λ∥F(Ideblur)−F(Igt)∥1
其中 m a t h c a l F mathcal{F} mathcalF表示离散傅里叶变换，权重参数 l a m b d a lambda lambda经验性地设置为0.1。
3.2 高效视觉状态空间模块
3.2.1 状态空间模型
状态空间模型是一种常用于时间序列分析和控制系统的数学框架。状态方程描述了一个潜在系统随时间的演变，代表了系统隐藏状态与其时间动态之间的关系。输入信号 x ( t ) x(t) x(t)通过隐藏状态 h ( t ) h(t) h(t)映射到输出响应 y ( t ) y(t) y(t)，通常被建模为一组一阶差分或微分方程：
h ′ ( t ) = A h ( t ) + B x ( t ) , y ( t ) = C h ( t ) + D x ( t ) h'(t)=Ah(t)+Bx(t), y(t)=Ch(t)+Dx(t) h′(t)=Ah(t)+Bx(t),y(t)=Ch(t)+Dx(t)
其中 A A A、 B B B、 C C C和 D D D是可学习的权重矩阵。
为此，状态方程可以使用零阶保持（ZOH）技术进行离散化：
h t = A ‾ h t − 1 + B ‾ x t , y t = C h t + D x t , A ‾ = e Δ A , B ‾ = ( Δ A ) − 1 ( e Δ A − I ) ⋅ Δ B . \begin{aligned} & h_{t}=\overline{A} h_{t - 1}+\overline{B} x_{t}, y_{t}=C h_{t}+D x_{t}, \\ & \overline{A}=e^{\Delta A}, \overline{B}=(\Delta A)^{-1}\left(e^{\Delta A}-I\right) \cdot \Delta B . \end{aligned} ht=Aht−1+Bxt,yt=Cht+Dxt,A=eΔA,B=(ΔA)−1(eΔA−I)⋅ΔB.
基于上述公式，Mamba提出了一种选择性扫描（S6）机制，以同时实现依赖于输入的权重和线性计算复杂度。在自然语言处理任务中使用状态空间模型没有问题，因为自然语言数据本质上是因果序列。但视觉任务带来了重大挑战，因为视觉数据本质上不是顺序的，并且包含局部纹理和全局结构等空间信息。由于S6是一个递归过程，在处理当前时间步(t)的输入时，它只能利用先前时间步的信息，而不能考虑未来时间步的信息。
3.2.2 高效视觉扫描（EVS）块
一种直接的方法是在不同方向（例如向前和向后）扫描视觉数据来解决上述问题，但这种策略会显著增加计算成本。例如，VMamba中的双向扫描策略使其计算成本比Mamba高4倍。
本文开发了一种高效的视觉状态空间模型，以最小的计算成本增加探索更多有用信息。关键在于提出的EVS块，我们仅在一个方向上扫描，并且在每次扫描前对输入应用一种几何变换（例如翻转和转置）。由于卷积的平移不变性，几何变换不会影响卷积本身，只会影响选择性扫描的过程。
具体来说，对于每个EVS块，假设它位于整个网络的第 i i i个EVSS模块中，首先对输入特征 F i n F_{in} Fin进行转置或翻转：
G = { T r a n s p o s e ( F i n ) i f i % 2 = 0 , F l i p ( F i n ) i f i % 2 = 1. G= \begin{cases}Transpose\left(F_{in}\right) & if\ i \% 2=0, \\ Flip\left(F_{in}\right) & if\ i \% 2=1 .\end{cases} G={Transpose(Fin)Flip(Fin)if i%2=0,if i%2=1.
这里， F l i p ( ⋅ ) Flip(·) Flip(⋅)表示沿特征的水平和垂直轴进行翻转。根据上述公式，每经过4个EVSS模块，图像特征将自动恢复到原始空间结构。特别地，如果网络中EVSS模块的总数不能被4整除，可以通过应用相应的逆变换来恢复原始空间结构，因为翻转和转置都是可逆的。这样，我们的EVSS模块有效地解决了上述问题，除了高效的几何变换外，避免了任何额外的计算负担。然后，选择性扫描可以表述如下：
X 1 , X 2 = s p l i t ( C o n v 1 × 1 ( G ) ) X ^ 1 = S 6 ( R e s h a p e ( σ ( D c o n v 3 × 3 ( X 1 ) ) ) ) X ^ 2 = σ ( X 2 ) F o u t = C o n v 1 × 1 ( R e s h a p e ( L ( X ^ 1 ) ) ⋅ X ^ 2 ) , \begin{aligned} & X_{1}, X_{2}=split\left(Conv_{1 × 1}(G)\right) \\ & \hat{X}_{1}=S6\left(Reshape\left(\sigma\left(Dconv_{3 × 3}\left(X_{1}\right)\right)\right)\right) \\ & \hat{X}_{2}=\sigma\left(X_{2}\right) \\ & F_{out }=Conv_{1 × 1}\left(Reshape\left(\mathcal{L}\left(\hat{X}_{1}\right)\right) \cdot \hat{X}_{2}\right), \end{aligned} X1,X2=split(Conv1×1(G))X^1=S6(Reshape(σ(Dconv3×3(X1))))X^2=σ(X2)Fout=Conv1×1(Reshape(L(X^1))⋅X^2),
其中 C o n v 1 × 1 ( ⋅ ) Conv_{1 ×1}(\cdot) Conv1×1(⋅)表示滤波器大小为 1 × 1 1×1 1×1像素的卷积层， D C o n v 3 × 3 ( ⋅ ) DConv _{3 ×3}(\cdot) DConv3×3(⋅)表示滤波器大小为 3 × 3 3×3 3×3像素的深度卷积层， m a t h c a l L ( ⋅ ) mathcal{L}(\cdot) mathcalL(⋅)表示归一化层， s p l i t ( ) split() split()在通道维度上分割图像特征， s i g m a sigma sigma表示GeLU激活函数， S 6 S6 S6表示Mamba提出的选择性扫描机制。
3.2.3 基于高效判别频域的前馈网络（EDFFN）
为了有效且高效地转换来自EVSS模块的特征，我们开发了一种基于高效判别频域的前馈网络。FFN部分通常是深度学习模型的核心组件，有助于潜在清晰图像的重建。FFTformer开发了一种基于判别频域的FFN，自适应地确定应保留哪些频率信息，但这在执行基于频域的操作时增加了计算成本。与在FFN网络中间应用基于频域操作的DFFN不同，我们的方法是在FFN网络的最后阶段对特征进行频域筛选。
4 实验结果
4.1 数据集和实现
4.1.1 数据集
遵循现有先进方法，我们在常用的GoPro数据集、HIDE数据集和RealBlur数据集上评估我们的方法。GoPro数据集包含2103张训练图像和1111张测试图像；HIDE数据集包括2025张主要关于人物的测试图像；RealBlur数据集包含由不同后处理策略生成的RealBlur-J和RealBlur-R子集，使用182个场景进行训练，50个场景进行测试。为了进行公平比较，我们遵循这些数据集的协议来评估我们的方法。
4.1.2 实现细节
浅层特征 F s F_{s} Fs的通道数为48，编码器/解码器中从第1级到第3级的EVSS模块数量分别为[6, 6, 12]。在训练过程中，我们使用默认参数的ADAM优化器。我们采用翻转和旋转操作的数据增强方法来生成训练数据。我们应用类似于但比Restormer更简单的渐进式训练：训练从补丁大小为 128 × 128 128×128 128×128像素、批量大小为64开始，进行300,000次迭代，学习率从 1 × 1 0 − 3 1×10^{-3} 1×10−3逐渐降低到 1 × 1 0 − 7 1×10^{-7} 1×10−7；然后将补丁大小扩大到(256×256)像素，批量大小为16，进行300,000次迭代，学习率初始化为 5 × 1 0 − 4 5×10^{-4} 5×10−4
并降低到 1 × 1 0 − 7 1×10^{-7} 1×10−7。学习率根据余弦退火策略进行更新。除非另有说明，所有实验均在NVIDIA RTX 4090 GPU上使用PyTorch框架进行。训练代码和测试模型可在https://github.com/kkkls/EVSSM上获取。
4.2 与现有方法的比较
4.2.1 在GoPro数据集上的评估
我们首先在GoPro数据集上评估所提方法的性能。将我们的方法与包括基于CNN的方法、基于Transformer的方法、基于MLP的方法和基于SSM的方法等现有先进方法进行比较。如果基于深度学习的方法未在基准数据集上训练，我们对其进行重新训练或微调以进行公平比较。由于CU - mamba未提供训练和测试代码，为了公平比较，我们将我们的结果与文献中报告的结果进行比较。我们使用峰值信噪比（PSNR）和结构相似性指数（SSIM）作为评估指标来衡量每个恢复图像的质量。
实验结果表明，我们的方法在PSNR和SSIM值上均优于其他方法。由于基于CNN的方法无法有效利用全局信息，其恢复的图像仍包含严重的模糊残留；虽然基于Transformer的方法能够建模全局上下文，但为降低计算成本采用的各种近似方法影响了它们对全局信息的建模能力，导致一些主要结构恢复不佳。相比之下，我们提出的视觉状态空间模型能够以较低的计算成本有效探索非局部信息，恢复的图像效果更好
4.2.2 在RealBlur数据集上的评估
使用相同的协议，我们在Rim等人提供的真实世界模糊数据集RealBlur上评估我们的方法。从表格数据可知，所提方法在该数据集上显著优于先前的工作，在RealBlur-J和RealBlur-R数据集上，PSNR分别提高了至少1.53dB和0.93dB。从视觉对比结果图中可以看出，我们的方法生成的结果中字符更清晰，结构细节更精细。
4.2.3 在HIDE数据集上的评估
我们在HIDE数据集上检验我们的方法。与现有先进方法类似，我们直接使用在GoPro数据集上训练的模型进行测试。我们方法的PSNR值比竞争方法高0.35dB，这表明我们的方法具有良好的泛化能力，因为模型并未在该数据集上进行训练。
4.2.4 在其他图像恢复任务数据集上的评估
我们还在用于图像去雨的真实世界数据集和用于图像去雾的RESIDE-6K数据集上评估我们的方法。实验结果表明，所提方法在去雨和去雾数据集上与现有先进方法相比表现出色，这证明了我们的方法在其他图像恢复任务上的泛化能力。
5 分析与讨论
5.1 EVS块的有效性
我们提出的EVS块的核心是在块开始时应用的几何变换。为了证明EVS块的效果，我们首先去除几何变换（简称为单向扫描，即像Mamba那样在一个方向上进行扫描），并使用与我们相同的设置训练这个基线模型。实验数据显示，我们的方法的PSNR比这个基线方法高0.14dB。状态空间模型需要将图像特征展平为一维序列，这损害了视觉数据的空间结构信息。与基线相比，我们带有几何变换的方法可以更好地探索非局部信息。同时，我们方法的参数数量和浮点运算次数（FLOPs）与基线相同，运行时间也几乎相同，这证明了所提出的EVS块在以最小的计算成本增加来提高状态空间模型处理视觉数据能力方面的有效性。
此外，我们还与分别在两个方向和四个方向进行扫描的基线进行比较。实验数据显示，虽然多方向扫描可以缓解状态空间模型在处理视觉数据方面的限制，但会导致网络参数数量和计算复杂度增加，从而使运行时间显著变长。需要注意的是，四个方向扫描的结果略低于两个方向扫描的结果，这是因为多方向扫描的方法只是简单地通过求和并归一化来融合从不同扫描方向提取的特征，因此多方向信息没有得到有效和充分的利用。与同时在多个方向进行扫描不同，我们的方法在每个EVS块开始时对输入特征应用几何变换，这使得每次扫描都能从不同方向捕捉上下文信息，并减轻了计算复杂度和运行时间的增加。从对比图可以看出，我们的方法比其他方法能更好地对图像进行去模糊处理，其中窗户的结构恢复得很好。
5.2 几何变换的有效性
在EVS块中，我们采用了翻转和转置这两种经典的图像几何变换。为了证明它们的有效性，我们分别去除翻转变换、转置变换以及同时去除翻转和转置变换。对比结果表明，应用翻转或转置变换可以取得更好的结果，PSNR至少提高0.04dB。我们同时使用翻转和转置变换的方法在没有显著增加计算成本或运行时间的情况下，优于所有这些基线方法。视觉对比图进一步证明了我们方法的有效性，我们方法恢复的车轮更清晰。
5.3 模型复杂度
我们进一步从模型参数、浮点运算次数（FLOPs）和平均运行时间方面研究了所提方法和其他高性能方法的模型复杂度。实验数据显示，所提方法的FLOPs更少，运行速度比评估的方法更快。
5.4 局限性
我们开发了一种有效且高效的方法，利用状态空间模型的特性进行高质量的图像恢复。然而，目前我们只考虑了翻转和转置等简单变换。在未来的工作中，我们将考虑更强大的变换方法，如极坐标变换，以便更好地利用状态空间模型来表征视觉数据的空间信息。
6 结论
本文提出了一种用于图像去模糊的高效视觉状态空间模型。具体而言，我们开发了一种高效的视觉扫描模块，在每次扫描前应用各种几何变换，使状态空间模型适用于视觉数据。与现有同时沿多个方向扫描的方法相比，我们的方法在不显著增加计算成本的情况下，更有效地探索非局部信息。与现有先进方法的大量评估和比较表明，我们的方法在实现良好性能的同时效率更高。
