标题：Efficient and Explicit Modelling of Image Hierarchies for Image Restoration
会议：CVPR2023
论文地址：http://arxiv.org/abs/2303.00748
官方代码：https://github.com/ofsoundof/GRL-Image-Restoration
作者单位：苏黎世联邦理工学院、Meta现实实验室、维尔茨堡大学、鲁汶大学

Abstract

本文的目的是提出一种机制，在全局、区域和局部范围内高效、显式地建模图像层次结构，以进行图像复原。为此，我们从分析自然图像的两个重要性质入手，包括跨尺度相似性和各向异性图像特征。受此启发，我们提出了锚条纹自注意力，在自注意力的空间和时间复杂度以及超出区域范围的建模能力之间取得了很好的平衡。然后，我们提出了一种新的网络结构，称为GRL，通过锚条纹自注意力、窗口自注意力和通道注意力增强卷积来显式地建模全局、区域和局部范围内的图像层次结构。最后，将所提出的网络应用于7种图像复原类型，涵盖真实场景和合成场景。所提出的方法在其中几个达到了新的SOTA。

1. Introduction

图像复原旨在从低质量图像恢复高质量图像，低质量图像是由于图像退化过程导致的，例如模糊、下采样、噪声破坏和JPEG压缩等。图像复原是一个不可逆问题，因为在图像退化过程中丢失了图像的重要内容信息。因此，为了恢复出高质量的图像，需要充分挖掘退化图像中所蕴含的丰富信息。
自然图像在全局、区域和局部包含特征层次结构，其可以被深度神经网络用于图像复原。首先，局部范围覆盖几个像素的跨度，典型特征是边缘和局部颜色。为了建模这些局部特征，使用小核（例如3×3）的卷积神经网络（CNNs）。其次，区域范围由数十个像素的窗口来表征。该范围的像素可以覆盖小目标和大目标的组成部份（图1中粉色方框）。由于范围较大，使用大核的CNNs显式建模区域特征（一致性、相似性）在参数和计算量上都会使效率低下。作为代替，具有窗口注意力机制的transformers非常适合该任务。第三，超出局部和区域范围，一些特征具有全局跨度（图1中的青色矩形框），包括但不限于对称性、多尺度的模式重复（图1a）、同尺度的纹理相似性（图1b）以及大目标和内容中的结构相似性和一致性（图1c）。为了对此范围的特征进行建模，需要全局的图像理解。

不同于局部和区域范围特征，对全局范围特征进行建模存在两大挑战。首先，现有的基于卷积和窗口注意力的图像复原网络无法通过使用单个计算模块显式地捕获长距离依赖关系。尽管在一些工作中使用了非局部操作，但它们要么在网络中稀疏地使用，要么应用于小的图像剪裁。因此，全局的图像理解仍主要借助于特征的逐步传播，通过重复的计算模块来实现。其次，如今图像分辨率的不断提高对长距离依赖建模提出了挑战。高图像分辨率会导致与成对像素比较以及相似性搜索相关的计算负担。
上述讨论引出了一系列研究问题：①如何高效地建模高维图像中的全局范围特征以进行图像复原；②如何通过单个计算模块来显式地建模图像层次结构（局部性、区域性、全局性）以进行高维图像复原；③以及这种联合建模如何能够对不同的图像复原任务带来统一的性能提升。
首先，我们提出了锚条纹自注意力用于超出区域范围的高效依赖建模。提出的自注意力受自然图像两个性质的启发，包括跨尺度相似性和各向异性图像特征。跨尺度相似性是指自然图像中的结构在不同尺度下被重复。受此启发，我们提出使用锚作为中介来近似自注意力中query和key之间精确的注意力图。由于锚将图像信息汇总到了一个更低维的空间，因此可以显著降低自注意力的空间和时间复杂度。此外，基于对各向异性图像特征的观察，我们提出在垂直和水平条纹内进行锚自注意力。由于注意力范围的各向异性收缩，实现了复杂度的进一步降低。并且轴向条纹的组合也保证了图像内容的全局视图。当配备条纹移动操作时，四种条纹自注意力模式（水平、垂直、移动水平、移动垂直）在计算复杂度和全局范围依赖建模能力之间取得了良好的平衡。会面会进一步从低秩性和相似性传播的角度对提出的锚条纹自注意力进行分析。
其次，提出了一种新的transformer网络，在单个计算模块中显式地建模全局、区域和局部范围的依赖关系。通过提出的锚条纹自注意力、窗口自注意力和通道注意力增强卷积的并行计算实现图像的层次建模。这个transformer结构称名为GRL。
最后，将提出的GRL transformer应用于各种图像复原任务。根据数据的可得性，这些任务可以分为3类：真实图像复原、合成图像复原和基于数据合成的真实图像复原。总共针对所提出的网络探讨了7项任务，包括图像超分辨率、图像去噪、JPEG压缩伪影去除、去马赛克、真实图像超分辨率、单幅图像去除运动模糊和去除失焦模糊。如图2所示，所提出的网络在所有被研究的图像复原任务上都显示出良好的结果。

2. Related Works

卷积用于局部范围建模。 基于学习的图像复原的一个基本假设是，重复的模式可以存在于相同或不同的图像中，它们携带的冗余信息有助于恢复局部块。因此，如果能够检测到重复模式并对其建模，就会有所帮助。这种直觉很好地匹配了卷积的计算过程，即在图像上滑动kernel，并检测出与可学习kernel相似的局部模式。通过堆叠多个卷积层，CNN的感受野被逐步扩大，从而捕获丰富的图像特征。自深度学习问世以来，人们一直致力于设计用于图像复原的卷积神经网络。
非局部和全局先验。 除了局部特征外，对图像的非局部先验和全局先验建模也很重要。早期工作如非局部均值就是为了达到这个目的，它计算一个输出像素，即图像中所有像素的加权和。受此启发，后来的工作发展到利用非局部范围内的重复模式进行图像去噪和超分辨率。除了传统的方法，非局部操作也被引入到深度神经网络中用于视频分类和图像超分辨率。
除了非局部操作，自注意力已经被开发用于建模全局范围的依赖。受自然语言处理领域的成功启发，transformers在视觉感知任务和图像复原任务中均有应用。然而，全局自注意力的计算复杂度随token数量呈二次方增长。因此，一些工作对全局自注意力的效率提升进行了研究。
区域自注意力。 在加速transformers的方法中，区域自注意力似乎很有前途。该思想在开创性的工作中提出，并改进为移动窗口注意力。受移动窗口注意力在视觉识别和感知上成功的启发，该方法也被用于图像复原。尽管窗口注意力机制具有良好的性能，但最近的工作指出，更宽范围的像素参与可以导致更好的图像复原。因此，在本文中，我们尝试提出一种高效地将自注意力的建模能力带到区域范围之外的方法。

3. Motivation

3.1. Self-attention for dependency modelling

自注意力擅长显式地建模长距离依赖关系，并且它有利于信息在所建模的依赖关系之间传播。这个操作允许一个token与所有其它token进行比较。输出token基于相似性比较，由所有token的加权和计算，即：
$$\mathbf{Y}=\mathrm{Softmax}(\mathbf{Q}\cdot {\mathbf{K}}^T/\sqrt{d})\cdot \mathbf{V}\text{\hspace{1em}(1)}$$其中，$\mathbf{Q}={\mathbf{W}}_Q\cdot \mathbf{X},\mathbf{K}={\mathbf{W}}_K\cdot \mathbf{X},\mathbf{V}={\mathbf{W}}_V\cdot \mathbf{X}$，${\mathbf{W}}_Q,{\mathbf{W}}_K,{\mathbf{W}}_V\in {\mathbb{R}}^{d\times d}$以及$\mathbf{X},\mathbf{Y}\in {\mathbb{R}}^{d\times N}$。$N$和$d$分别表示token的数量和一个token的维度。此外，$\mathbf{M}$表示注意力图，即$\mathbf{M}=\mathrm{Softmax}(\mathbf{Q}\cdot {\mathbf{K}}^T/\sqrt{d})$。
自注意力的时间复杂度为$\mathcal{O}(N^{2}d)$，空间复杂度由注意力图$\mathbf{M}$的项$\mathcal{O}(N^2)$主导。自注意力的计算复杂度和内存占用随token数量呈二次方增长。因此，对于那些token数量是特征图两个维度乘积的图像，自注意力很容易成为计算瓶颈。为了克服这一问题，提出了在窗口内应用自注意力。通过这种方式，参与自注意力的token数量显著减少，计算负担也有所减轻。
窗口自注意力的问题在于，由于窗口尺寸较小（8×8），该操作的建模能力被限制在区域范围内。另一方面，最近的工作表明，即使窗口大小稍微增加也能导致更好的图像复原。因此可以推测，超出区域范围的依赖关系建模对图像复原仍然很重要。因此，如何在控制计算预算的情况下保持长距离依赖建模的能力还有待研究。

3.2. Motivation I: cross-scale similarity

注意力图$\mathbf{M}$在自注意力中起着至关重要的作用，因为它记录了图像中每一成对像素之间的相似性。因此，提高公式1中自注意力的效率需要分析注意力图的性质。而我们是受图像一个性质的启发，即跨尺度相似性。即，在具有不同缩放因子的图像的不同版本中，保留了图像的线条、边缘等基本结构。图3显示了图像中像素之间的注意力图。特别地，一个像素与整幅图像之间的注意力图被可视化为灰度热图。如图所示，无论这个像素来自高分辨率图像还是降尺度的版本，其与高分辨率图像之间的热图都显示了图像的基本结构。图3c与图3d的热图非常相似。

锚自注意力。 受图3所示的跨尺度相似性的启发，我们尝试降低公式1中全局自注意力的复杂度，通过对不同分辨率的图像进行操作，并控制token的数量，即$\mathcal{O}(N^{2}d)$中的$N^2$项。为了实现这个目标，我们在query、key和value的三元组之外，引入了一个称为锚（anchor）的额外概念。锚集是对图像特征图中信息的汇总，并且具有较低的维度。不同于直接进行query和key之间的相似性比较，锚是作为相似性比较的中介。形式上，锚自注意力由下式提出：
$$\begin{array}{rl}\mathbf{Y}& ={\mathbf{M}}_e\cdot \mathbf{Z}={\mathbf{M}}_e\cdot ({\mathbf{M}}_d\cdot \mathbf{V})\\ {\mathbf{M}}_d& =\mathrm{Softmax}(\mathbf{A}\cdot {\mathbf{K}}^T/\sqrt{d})\\ {\mathbf{M}}_e& =\mathrm{Softmax}(\mathbf{Q}\cdot {\mathbf{A}}^T/\sqrt{d})\end{array}\text{\hspace{1em}(2)(3)(4)}$$其中，$M\ll N$，$A\in {\mathbb{R}}^{M\times d}$是锚，${\mathbf{M}}_e\in {\mathbb{R}}^{N\times M}$和${\mathbf{M}}_d\in {\mathbb{R}}^{M\times N}$表示query-anchor对和anchor-key对之间的注意力图。
由于锚的数量远小于token的数量，因此得到的两个注意力图${\mathbf{M}}_e$和${\mathbf{M}}_d$的尺寸远小于公式1中的原始注意力图$\mathbf{M}$。然后从右边计算公式2中的矩阵乘法。首先对anchor和key进行自注意力。注意力图${\mathbf{M}}_d$将token$\mathbf{V}$提取为中间特征$\mathbf{Z}$。然后在query和anchor之间进行自注意力。第二个注意力图${\mathbf{M}}_e$扩大了特征$\mathbf{Z}$的尺寸，恢复了$\mathbf{V}$中的信息。锚自注意力的计算复杂度降低到了$\mathcal{O}(NMd)$。空间复杂度降低到了$\mathcal{O}(NM)$。

3.3. Motivation II: anisotropic image features

通过去除二次项$N^2$，锚自注意力可以显著降低公式1中的自注意力的空间和时间复杂度。然而，对于图像复原任务，剩余项$N$是图像宽度和高度的乘积。因此，公式2中锚自注意力的复杂度仍然可能由于较大的项$N$而无法承受。因此，我们希望进一步降低锚自注意力的复杂度。
为了达到这个目的，我们借助了自然图像的另一个性质，即各向异性图像特征。如图4所示，图4c和d中的单个目标，图4h中的多尺度相似性，图4e和g中的对称性等自然图像特征以各向异性的方式跨越。因此，整个图像上的各向同性全局范围的注意力对于捕获各向异性的图像特征是冗余的。对此，我们提出在图4所示的各向异性条纹内进行注意力。

条纹注意机制。 所提出的条纹注意力机制由四种模式组成，包括水平条纹、垂直条纹、移动水平条纹和移动垂直条纹。水平和垂直的条纹注意力机制可以在transformer网络中交替使用。通过这种方式，在保持全局范围建模能力和控制全局自注意力的计算复杂度之间进行了权衡。因此，结合锚的概念，我们提出了锚条纹自注意力。对于该注意力机制，通过引入的锚，在垂直和水平条纹的内部进行高效的自注意力。

3.4. Discussion

所提出的锚条纹自注意力机制与低秩性和相似性传播两个概念密切相关。
注意力图的低秩性。 通过比较公式1和公式2中的自注意力机制，不难发现，原始的注意力图$\mathbf{M}$被分解为秩不大于$M$的小注意力图${\mathbf{M}}_d$和${\mathbf{M}}_e$。这里的本质是在不事先计算原始注意力图的情况下提供低秩近似。对于锚自注意力的成功，重要的是确保以锚为中介，近似的注意力图要与原始的注意力图相似。因此，图5提供了一个额外的分析。

首先，通过观察query、key和value，我们可以推断出锚与query和key具有非常相似的结构。因此，锚很好地汇总了query和key中的信息。用锚作为中介来近似自注意力似乎是合理的。此外，近似的注意力图${\mathbf{M}}_e\cdot {\mathbf{M}}_d$和精确的注意力图$\mathbf{M}$也在图5中进行了比较。如图所示，近似的注意力图保持了精确的注意力图中的主要结构，两个注意力图之间大的皮尔逊相关系数（0.9505）证实了这一点。因此，锚自注意力的质量被保证了。
度量和相似性传播。 从另一个角度来看，在提出的锚自注意力机制中，首先将query和key与锚进行比较，然后计算query和key的相似性。因此，该计算过程需要将query-anchor和key-anchor的相似性传播到query-key对。而相似性传播与度量空间中的三角不等式有关。
一个数学度量需要满足几个条件，包括基本的三角不等式：
$$d(\mathbf{q},\mathbf{k})\le d(\mathbf{a},\mathbf{q})+d(\mathbf{a},\mathbf{k})\text{\hspace{1em}(5)}$$其中，$d(\cdot ,\cdot )$定义了两个实体之间的度量。因此，$\mathbf{q}$/$\mathbf{k}$距离的上界是$\mathbf{a}$/$\mathbf{q}$距离和$\mathbf{a}$/$\mathbf{k}$距离之和。这意味着如果$\mathbf{a}$与$\mathbf{q}$和$\mathbf{k}$都相似（接近），那么$\mathbf{q}$与$\mathbf{k}$也应该相似（接近）。然而，公式1和公式2中的相似性度量是用点积来定义的，而不是token之间的距离，不满足三角不等式。因此，相似性传播在理论上无法得到保证。为了研究相似性度量的影响，进行了消融研究，结果如第5节所示。点积和距离作为相似性度量进行了比较。从结果来看，虽然点积不是严格遵守三角不等式，但仍然保证了较好的图像复原效果。因此，我们可以从经验上得出结论，在当前场景下，点积足以进行相似性传播。

4. Modelling Image Hierarchies

在这一节，我们回答引言中描述的第二个研究问题，即如何通过单个计算模块显式地建模图像层次结构。为此，我们提出了包含全局范围、区域范围和局部范围图像建模能力的GRL网络结构。

网络结构。 所提出的网络的总体结构如图6所示。该网络将退化的低质量图像作为输入，在网络内部对图像进行处理，输出恢复的高质量图像。具体而言，该网络包含三个部分。①用简单的卷积实现的特征提取层，将输入图像转换为特征图。②表示学习组件丰富了前面操作中提取到的信息。transformer stage由若干个transformer层组成，并以卷积层结束。特征图的维度在整个表示学习模块中保持不变。跳跃连接同时应用于transformer stage和表示学习模块。③图像重建模块利用前面操作计算出的丰富特征，估计出一幅恢复图像。
transformer层。 图6b所示的transformer层是提供全局、区域和局部范围的图像层次结构建模能力的关键组件。该层首先通过并行的自注意力模块和通道注意力增强卷积来对输入特征图进行处理。卷积分支用于捕获输入特征图中的局部结构。另一方面，自注意力模块包含Swin Transformer V2中提出的窗口注意力和本文提出的锚条纹注意力。特征图沿着通道维度被等分，在两个注意力模块内并行处理之后再沿着通道维度concat。窗口注意力提供了机制来捕获区域范围依赖。然后，将卷积模块和注意力模块的输出特征图加到输入特征图中，由后面的MLP模块进行处理。
锚条纹自注意力。 所提出的锚条纹自注意力的操作根据公式2来进行，可视化如图6c所示。图中还显示了不同特征的维度。$\mathbf{Q}$、$\mathbf{V}$、$\mathbf{K}$的三元组由普通的线性投射得到。为了将信息汇总到锚，锚投射被实现为一个平均池化层接一个线性投射。经过锚投射后，图像特征图的分辨率沿着两个方向均按因子$s$降尺度。如图6所示，两个注意力图${\mathbf{M}}_d$和${\mathbf{M}}_e$发挥了与原始注意力图$\mathbf{M}$相似的作用，但具有较小的空间和时间复杂度。

5. Experimental Results

我们通过考察所提出的网络在不同图像复原任务上的性能来回答引言中提出的第三个研究问题。根据数据类型，将研究任务分为三种常用的场景，包括①真实图像复原（单幅图像去除运动模糊和去除失焦模糊），②基于合成数据的图像复原（图像去噪、单幅图像SR、JPEG压缩伪影去除、去马赛克）和③基于数据合成的真实图像复原。我们提供了三种不同模型大小的网络，包括tiny、small和base版本（GRL-T、GRL-S、GRL-B）。对于真实图像复原和合成图像复原，使用Adam优化器和$L_1$损失，以$2\times 10^{-4}$的初始学习率来训练网络。

5.1. Image deblurring

我们首先考察了所提出的网络在单幅图像去除运动模糊和去除失焦模糊两个真实图像复原任务上的性能。
单幅图像去除运动模糊。 表1和表2分别在合成数据集（GoPro 、HIDE）和真实数据集（RealBlur-R）上给出了单幅图像去除运动模糊的实验结果。在GoPro数据集上，与之前最先进的Restromer相比，本文提出的GRL取得了1.01 dB的显著PSNR提升。在HIDE数据集上，PSNR提升了0.43 dB。请注意，这些提升是在更少的参数预算下取得的。如表4所示，GRL-B比Restromer减少了24%的参数。如表2所示，GRL-B在RealBlur-R数据集上取得了40.20 PSNR的SOTA性能。


去除失焦模糊。 表3给出了使用单幅图像和双像素图像来去除失焦模糊的实验结果。我们的GRL在三种场景类型上都优于之前的方法。在组合场景下，GRL的单幅图像和双像素图像去除失焦模糊性能相比Restromer分别提升了0.20 dB和0.38 dB。与Uformer和IFAN相比，GRL在双像素场景下取得了1.39 dB和1.05 dB的PSNR增益。

5.2. Image restoration based on synthetic data

利用合成数据进行图像复原研究对于揭示复原方法网络的能力也是有价值的。除了在真实数据上的实验，我们还研究了网络在合成数据上的表现。
图像去噪。 首先，高斯图像去噪实验结果如表4所示。为了公平地比较不同模型之间的差异，网络复杂度和准确率都显示在表中。并且观察到几个关键的发现。①tiny版本的GRL-T极其高效，在不牺牲网络精度的同时，将模型复杂度降低了两个数量级（只有DRUNet的0.76%和2.7%）。②small版本的GRL-S与此前SOTA的SwinIR和Restromer具有竞争力。③在Urban100上，base版本的表现远远优于Restromer。

图像SR。 对经典图像的实验结果如表5所示。对轻量级模型和准确的SR模型都进行了总结。从结果可以得出类似的结论。①在轻量级网络中，GRL-T优于包括DBPN、SwinIR和EDT在内的基于卷积和基于自注意力的网络。与EDT相比，在Urban100和Manga109数据集上取得了明显的改进（×4 SR分别提升0.44 dB和0.22 dB）。②GRL-B为准确的图像超分辨率取得了新的SOTA。③GRL-S在网络复杂度和SR图像质量之间取得了很好的平衡。

JPEG压缩伪影去除。 针对彩色和灰度图像的实验结果如表6和表7所示。研究了JPEG压缩的四个图像质量因子，范围从10到40。如表所示，在不同的数据集和质量因子下，所提出的GRL-S网络都优于之前SOTA的方法。值得注意的是，GRL-S具有比FBCNN小得多的模型复杂度。

去马赛克。 图像去马赛克的结果如表8所示。所提出的方法明显优于之前的方法RNAN和DRUNet。

5.3. Real image restoration based on data synthesis

最后，我们还考察了该网络在真实世界图像复原中的性能。目标是按照4的升尺度因子对低质量图像进行超分辨率。由于该任务没有真值图像，因此图7只给出了视觉对比。与其它方法相比，本文提出的GRL能够去除更多低分辨率图像中的伪影。

5.4. Ablation study

在本节中，对所提方法的各种设计选择进行了消融。
相似性比较方法的影响。 如3.4节所述，为保证相似性传播的理论性，应使用数学度量而非点积。为了研究两者之间的差异，对两种操作下的图像复原进行了比较，结果如表9所示。从表中可以看出，点积与距离度量相比是非常有竞争力的，它在几个场景中都优于距离度量。考虑到这一点，在其它实验中仍采用点积。

锚投射的影响。 锚投射操作有助于汇总特征图中的信息。消融研究见表10。综合考虑精度性能和参数预算，最终使用平均池化+线性投射。此外，窗口和条纹尺寸在图像SR中的影响也在表11中消融。在Urban100和Manga109上，增加窗口和锚的尺寸可以带来显著的性能提升。

6. Conclusion

在本文中，我们提出了GRL，一个用于图像复原的具有高效和显式层次结构建模能力的网络。所提出的网络主要受到两个图像性质的启发，包括跨尺度相似性和各向异性图像特征。基于此，我们提出了高效的锚条纹自注意力模块用于长距离依赖建模。然后提出了一种用于图像复原的通用网络结构。所提出的网络可以在全局、区域和局部范围内建模图像层次结构。由于先进的计算机制，所提出的网络结构取得了各种图像复原任务的SOTA。