Restormer: Efficient Transformer for High-Resolution Image Restoration解读

论文地址：Restormer: Efficient Transformer for High-Resolution Image Restoration。

摘要

由于卷积神经网络（CNN）在从大规模数据中学习可推广的图像先验方面表现出色，这些模型已被广泛应用于图像复原及相关任务。近年来，另一类神经架构——Transformer，在自然语言和高级视觉任务上取得了显著的性能提升。虽然Transformer模型缓解了CNN的缺点（即感受野有限和对输入内容的不适应性），但其计算复杂度随空间分辨率呈二次方增长，因此无法应用于大多数涉及高分辨率图像的图像复原任务。在这项工作中，我们通过在构建模块（多头注意力和前馈网络）中进行几个关键设计，提出了一种高效的Transformer模型，使其能够捕获长距离像素交互，同时仍适用于大图像。我们的模型名为复原Transformer（Restormer），在多个图像复原任务上取得了最先进的成果，包括图像去雨、单图像运动去模糊、散焦去模糊（单图像和双像素数据）以及图像去噪（高斯灰度/彩色去噪和真实图像去噪）。源代码和预训练模型可在https://github.com/swz30/Restormer获取。

1. Introduction

图像复原是一项通过去除受损输入图像中的退化因素（例如噪声、模糊、雨滴等）来重建高质量图像的任务。由于其不适定的性质，这是一个极具挑战性的问题，通常需要强有力的图像先验知识才能有效进行复原。由于卷积神经网络（CNN）在从大规模数据中学习可推广的先验知识方面表现出色，与传统的复原方法相比，它们已成为更优的选择。

CNN中的基本操作是“卷积”，它提供了局部连接性和平移等变性。虽然这些特性为CNN带来了效率和泛化能力，但也引发了两个主要问题。（a）卷积算子的感受野有限，因而无法对长距离像素依赖关系进行建模。（b）卷积滤波器在推理时具有静态权重，因此无法灵活地适应输入内容。为解决上述缺点，一种更强大且动态的替代方案是自注意力（SA）机制[17,77,79,95]，它通过对所有其他位置进行加权求和来计算给定像素处的响应。

自注意力是Transformer模型[34,77]中的核心组件，不过有着独特的实现方式，即多头自注意力（multi-head SA），它针对并行化和有效的表征学习进行了优化。Transformer在自然语言任务[10,19,49,62]以及高级视觉问题[11,17,76,78]上都展现出了最先进的性能。尽管自注意力在捕捉长距离像素交互方面非常有效，但其复杂度随空间分辨率呈二次方增长，因此无法应用于高分辨率图像（这在图像复原中是常见情况）。

近来，已有一些针对图像复原任务对Transformer进行定制的尝试[13,44,80]。为降低计算负荷，这些方法要么在每个像素周围大小为8×8的小空间窗口上应用自注意力[44,80]，要么将输入图像划分为大小为48×48的互不重叠的图像块，并在每个图像块上独立计算自注意力[13]。然而，限制自注意力的空间范围与捕捉真实长距离像素关系的目标相悖，尤其是在高分辨率图像上。

在本文中，我们提出了一种用于图像复原的高效Transformer，它能够对全局连接性进行建模，并且仍适用于大图像。具体而言，我们引入了一个多深度卷积头“转置”注意力（MDTA）模块（见第3.1节）来替代常规的多头自注意力[77]，该模块具有线性复杂度。它在特征维度而非空间维度上应用自注意力，也就是说，MDTA并非显式地对像素两两之间的交互进行建模，而是通过计算特征通道间的互协方差，从（经过键和查询投影的，key and query projected）输入特征中获取注意力图。我们的MDTA模块的一个重要特性是在计算特征协方差之前进行局部上下文混合。这是通过使用1×1卷积对跨通道上下文进行逐像素聚合以及使用高效的深度可分离卷积对局部上下文进行逐通道聚合来实现的。这一策略具有两个关键优势。首先，它强调空间局部上下文，并在我们的处理流程中引入了卷积操作的互补优势。其次，它确保在计算基于协方差的注意力图时，像素间的上下文全局关系能被隐式地建模。

前馈网络（FN）是Transformer模型[77]的另一个构建模块，它由两个全连接层以及中间的非线性层构成。在这项工作中，我们利用门控机制[16]对常规前馈网络[77]的第一个线性变换层进行了重新设计，以改善网络中的信息流。这个门控层被设计为两个线性投影层的逐元素乘积，其中一个线性投影层通过高斯误差线性单元（GELU）非线性激活函数[27]进行激活。我们的门控深度卷积前馈网络（GDFN）（见第3.2节）也基于与MDTA模块类似的局部内容混合，同样对空间上下文予以重视。GDFN中的门控机制控制着哪些互补特征应该向前传递，并允许网络层级中的后续层专门聚焦于更精细的图像属性，从而生成高质量的输出。

除了上述架构上的创新之外，我们还展示了针对Restormer的渐进式学习策略的有效性（见第3.2节）。在此过程中，网络在训练早期阶段基于小图像块和大批次进行训练，而在后期阶段则基于逐渐增大的图像块和小批次进行训练。这种训练策略有助于Restormer从大图像中学习上下文信息，进而在测试时提升性能表现。我们进行了全面的实验，并展示了Restormer在16个基准数据集上针对多个图像复原任务（包括图像去雨、单图像运动去模糊、散焦去模糊（针对单图像和双像素数据）以及图像去噪（针对合成数据和真实数据））所取得的最先进的性能；见图1。此外，我们还进行了大量的消融实验，以展示架构设计和实验选择的有效性。

这项工作的主要贡献总结如下：

• 我们提出了Restormer，它是一种编解码器Transformer，用于在高分辨率图像上进行多尺度局部 - 全局表征学习，而无需将图像分解为局部窗口，从而能够利用远距离的图像上下文信息。
• 我们提出了一个多深度卷积头转置注意力（MDTA）模块，该模块能够聚合局部和非局部像素交互，并且效率足够高，可用于处理高分辨率图像。
• 一种新的门控深度卷积前馈网络（GDFN），它能执行可控的特征变换，即抑制信息量较少的特征，只允许有用信息在网络层级中进一步传递。

2. 背景

图像复原。近年来，数据驱动的卷积神经网络（CNN）架构[7, 18, 92, 93, 105, 107]已被证明优于传统的复原方法[26, 36, 53, 75]。在卷积设计中，基于编解码器的U-Net架构[3, 14, 39, 80, 90, 93, 99]由于其具有分层多尺度表示且计算效率较高的特点，在图像复原方面得到了大量研究。同样，基于跳跃连接的方法由于特别侧重于学习残差信号，也已被证明对图像复原有效[24, 48, 92, 106]。空间和通道注意力模块也已被引入，用于有选择性地关注相关信息[43, 92, 93]。我们建议读者参考NTIRE挑战赛报告[2, 5, 30, 57]以及近期的文献综述[8, 42, 73]，这些资料总结了图像复原的主要设计选择。

视觉Transformer。Transformer模型最初是为自然语言任务中的序列处理而开发的[77]。它已被应用于众多视觉任务中，如图像识别[17, 76, 88]、图像分割[78, 83, 108]、目标检测[11, 50, 109]等。视觉Transformer（Vision Transformers）[17, 76]将一幅图像分解为一系列图像块（局部窗口），并学习它们之间的相互关系。这些模型的显著特点是具有很强的学习图像块序列之间长距离依赖关系的能力，以及对给定输入内容的适应性[34]。由于这些特性，Transformer模型也已被用于研究诸如超分辨率[44, 85]、图像彩色化[37]、去噪[13, 80]以及去雨[80]等低级视觉问题。然而，Transformer中自注意力（SA）的计算复杂度会随着图像块的数量呈二次方增长，从而限制了其在高分辨率图像上的应用。因此，在需要生成高分辨率输出的低级图像处理应用中，近期的方法通常会采用不同策略来降低复杂度。一种可能的解决办法是利用Swin Transformer设计[44]在局部图像区域内应用自注意力[44, 80]。然而，这种设计选择限制了局部邻域内的上下文聚合，违背了使用自注意力取代卷积的主要初衷，因此并不十分适合图像复原任务。相比之下，我们提出了一种Transformer模型，它能够学习长距离依赖关系，同时保持计算高效性。

3.方法

我们的主要目标是开发一种高效的Transformer模型，使其能够处理用于复原任务的高分辨率图像。为缓解计算瓶颈问题，我们针对多头自注意力（SA）层以及一个多尺度分层模块引入了关键设计，该多尺度分层模块相较于单尺度网络[44]具有更低的计算需求。我们首先展示Restormer架构的整体流程（见图2）。然后，我们将描述所提出的Transformer模块的核心组件：（a）多深度卷积头转置注意力（MDTA）和（b）门控深度卷积前馈网络（GDFN）。最后，我们会详细介绍用于有效学习图像统计信息的渐进式训练方案。

整体流程。给定一幅受损图像(I \in R^{H×W×3})，Restormer首先应用卷积操作来获取低级特征嵌入(F_0 \in R^{H×W×C})；其中(H×W)表示空间维度，(C)为通道数量。接下来，这些浅层特征(F_0)会通过一个四级对称的编解码器，并被转换为深层特征(F_d \in R^{H×W×2C})。编解码器的每一层都包含多个Transformer模块，且模块数量从顶层到底层逐渐增加，以保持效率。从高分辨率输入开始，编码器按层级逐步缩小空间尺寸，同时扩展通道容量。解码器以低分辨率的潜在特征(F_l \in R^{\frac{H}{8}×\frac{W}{8}×8C})作为输入，并逐步恢复高分辨率表示。对于特征下采样和上采样操作，我们分别应用像素重排（pixel-unshuffle）和像素洗牌（pixel-shuffle）操作[69]。为辅助恢复过程，编码器特征通过跳跃连接[66]与解码器特征进行拼接。拼接操作之后会紧跟一个(1×1)的卷积操作，以在除顶层之外的所有层级将通道数量减半。在第一层，我们让Transformer模块将编码器的低级图像特征与解码器的高级特征进行聚合。这有助于在恢复后的图像中保留精细的结构和纹理细节。接下来，深层特征(F_d)会在以高空间分辨率运行的细化阶段进一步丰富。正如我们将在实验部分（第4节）看到的那样，这些设计选择能够提升图像质量。最后，对细化后的特征应用一个卷积层，生成残差图像(R \in R^{H×W×3})，将其与受损图像相加，便可得到复原后的图像：(\hat{I} = I + R)。接下来，我们将介绍Transformer模块的各个组件。

图2. 用于高分辨率图像复原的Restormer架构。我们的Restormer包含多尺度分层设计，并融入了高效的Transformer模块。Transformer模块的核心组件为：（a）多深度卷积头转置注意力（MDTA），它跨通道（而非空间维度）执行（空间丰富的）查询 - 键特征交互；（b）门控深度卷积前馈网络（GDFN），其执行可控的特征变换，即允许有用信息进一步传播。

3.1. 多深度卷积头转置注意力

Transformer中的主要计算开销来自自注意力层。在传统的自注意力（SA）[17, 77]中，键 - 查询点积交互的时间和内存复杂度随输入的空间分辨率呈二次方增长，也就是说，对于尺寸为(W×H)像素的图像，其复杂度为(O(W^{2}H{2}))。因此，将自注意力应用于大多数通常涉及高分辨率图像的图像复原任务是不可行的。为缓解这一问题，我们提出了如图2（a）所示的多深度卷积头转置注意力（MDTA），它具有线性复杂度。关键思路是跨通道而非空间维度应用自注意力，即通过计算通道间的互协方差来生成一个注意力图，从而隐式地对全局上下文进行编码。作为MDTA中的另一个重要组成部分，我们引入了深度可分离卷积，以便在计算特征协方差以生成全局注意力图之前强调局部上下文。

从一个经过层归一化的张量(Y \in R^{{\hat{H}×\hat{W}×\hat{C}})开始，我们的MDTA首先生成融入了局部上下文的查询（(Q)）、键（(K)）和值（(V)）投影。这是通过应用(1×1)卷积来聚合逐像素的跨通道上下文，随后应用(3×3)深度可分离卷积对逐通道的空间上下文进行编码来实现的，从而得到。其中，(W_{(·)}}{p})是(1×1)逐点卷积，(W_{(·)}^{{d})是(3×3)深度可分离卷积。我们在网络中使用无偏置的卷积层。接下来，我们对查询和键投影进行重塑，使得它们的点积交互生成一个尺寸为(R}{\hat{C}×\hat{C}})的转置注意力图(A)，而非尺寸为(R^{\hat{H}\hat{W}×\hat{H}\hat{W}})的庞大常规注意力图[17, 77]。

总体而言，MDTA过程定义如下：

在此，(\alpha)是一个可学习的缩放参数，用于在应用Softmax函数之前控制(\hat{K})和(\hat{Q})点积的大小。与传统的多头自注意力（SA）[17]类似，我们将通道数量划分为多个“头”，并并行学习独立的注意力图。

3.2. 门控深度卷积前馈网络

为了对特征进行变换，常规的前馈网络（FN）[17, 77]会对每个像素位置分别且相同地进行操作。它使用两个1×1卷积，第一个用于扩展特征通道（通常扩展倍数γ = 4），第二个用于将通道数量减少回原始输入维度。在隐藏层中会应用非线性激活函数。在这项工作中，我们对前馈网络提出了两项根本性的改进，以提升表征学习能力：（1）门控机制；（2）深度可分离卷积。我们的门控深度卷积前馈网络（GDFN）的架构如图2（b）所示。门控机制被构建为两个并行的线性变换层路径的逐元素乘积形式，其中一条路径通过高斯误差线性单元（GELU）非线性激活函数[27]进行激活。与多深度卷积头转置注意力（MDTA）一样，我们也在门控深度卷积前馈网络（GDFN）中纳入了深度可分离卷积，以便对来自空间相邻像素位置的信息进行编码，这对于学习局部图像结构以实现有效复原很有帮助。给定一个输入张量(X \in R^{\hat{H}×\hat{W}×\hat{C}})，门控深度卷积前馈网络（GDFN）的计算公式如下：

总体而言，门控深度卷积前馈网络（GDFN）控制着我们整个流程中各层级的信息流，从而使每一层都能聚焦于与其他层级互补的细节信息。也就是说，与多深度卷积头转置注意力（MDTA）（侧重于利用上下文信息丰富特征）相比，门控深度卷积前馈网络（GDFN）发挥着不同的作用。由于所提出的门控深度卷积前馈网络（GDFN）相较于常规的前馈网络[17]执行了更多的操作，我们降低了扩展比率(\gamma)，以使二者具有相近的参数数量和计算负担。

3.3. 渐进式学习

基于卷积神经网络（CNN）的复原模型通常是在固定尺寸的图像块上进行训练的。然而，在裁剪后的小图像块上训练Transformer模型可能无法对全局图像统计信息进行编码，因此在测试时针对全分辨率图像会产生次优的性能表现。为此，我们采用渐进式学习方法，即在训练的早期阶段让网络基于较小的图像块进行训练，而在后续的训练阶段则基于逐渐增大的图像块进行训练。通过渐进式学习在混合尺寸图像块上训练得到的模型，在测试时（图像可能具有不同分辨率，这在图像复原中是常见情况）展现出了性能提升。渐进式学习策略与课程学习过程类似，网络从较简单的任务开始，逐渐过渡到学习更复杂的任务（即需要保留精细图像结构/纹理的任务）。由于基于大图像块进行训练会耗费更长时间，随着图像块尺寸的增加，我们会减小批次大小，以便在每次优化步骤中保持与固定图像块训练相近的时间。

4. 实验与分析

我们在基准数据集以及针对四项图像处理任务的实验设置下，对所提出的Restormer进行评估，这四项任务分别是：（a）图像去雨，（b）单图像运动去模糊，（c）散焦去模糊（针对单图像及双像素数据），（d）图像去噪（针对合成数据和真实数据）。关于数据集、训练协议以及更多可视化结果的详细信息在补充材料中给出。在表格中，所评估方法的最佳和次佳质量得分均已突出显示并加下划线标注。

实现细节。我们针对不同的图像复原任务分别训练模型。在所有实验中，除非另有说明，我们均使用以下训练参数。我们的Restormer采用一个四级编解码器结构。从第1层到第4层，Transformer模块的数量分别为[4, 6, 6, 8]，多深度卷积头转置注意力（MDTA）中的注意力头数量分别为[1, 2, 4, 8]，通道数量分别为[48, 96, 192, 384]。细化阶段包含4个模块。门控深度卷积前馈网络（GDFN）中的通道扩展因子γ = 2.66。我们使用AdamW优化器（β₁ = 0.9，β₂ = 0.999，权重衰减为1e⁻⁴）以及L1损失对模型进行30万次迭代训练，初始学习率为3e⁻⁴，并通过余弦退火[51]逐渐降至1e⁻⁶。对于渐进式学习，我们以图像块尺寸为128×128、批次大小为64开始训练。在迭代次数分别达到[92000, 156000, 204000, 240000, 276000]时，图像块尺寸和批次大小的组合依次更新为[(160², 40), (192², 32), (256², 16), (320², 8), (384², 8)]。对于数据增强，我们使用水平翻转和垂直翻转操作。

4.1. 图像去雨结果

我们采用与现有方法[32, 61, 93]类似的方式，利用YCbCr色彩空间中的Y通道来计算峰值信噪比（PSNR）/结构相似性指数（SSIM）得分。表1显示，我们的Restormer在所有五个数据集上相较于现有方法都取得了持续且显著的性能提升。与近期最佳方法SPAIR [61]相比，Restormer在所有数据集上取平均值时实现了1.05分贝的提升。在单个数据集上，提升幅度可高达2.06分贝，例如在Rain100L数据集上。图3展示了一个具有挑战性的可视化示例。我们的Restormer复原出了一幅无雨滴的图像，同时有效地保留了图像的结构内容。

4.2. 单图像运动去模糊结果

我们在合成数据集（GoPro [56]、HIDE [67]）以及真实世界数据集（RealBlur-R [65]、RealBlur-J [65]）上对去模糊方法进行评估。表2显示，我们的Restormer在所有四个基准数据集上均优于其他方法。在所有数据集上取平均值时，我们的方法相较于近期算法MIMO - UNet + [14]性能提升了0.47分贝，相较于此前最佳方法MPRNet [93]性能提升了0.26分贝。与MPRNet [93]相比，Restormer的浮点运算次数（FLOPs）减少了81%（见图1）。此外，我们的方法相较于Transformer模型IPT [13]性能提升了0.4分贝，同时参数数量少了4.4倍，运行速度快了29倍。值得注意的是，我们的Restormer仅在GoPro [56]数据集上进行训练，但它通过创造新的最先进水平，展现出了对其他数据集很强的泛化能力。图4显示，与其他算法相比，我们方法生成的图像更加清晰，在视觉上更接近真实图像。

4.3. 散焦去模糊（Defocus Deblurring）结果

表3展示了传统散焦去模糊方法（增强双边深度图法（EBDB）[33]和联合非局部贝叶斯法（JNB）[68]）以及基于学习的方法在双像素散焦去模糊数据集（DPDD）[3]上的图像保真度得分。我们的Restormer在所有场景类别下的单图像和双像素散焦去模糊任务中，都显著优于现有的最先进方案。特别是在综合场景类别中，Restormer相较于此前最佳方法（IFAN）[41]实现了约0.6分贝的提升。与Transformer模型Uformer [80]相比，我们的方法在峰值信噪比（PSNR）方面有1.01分贝的显著增益。图5表明，相较于其他方法，我们的方法在去除空间变化的散焦模糊方面更为有效。

4.4. 图像去噪结果

我们在添加了加性高斯白噪声生成的合成基准数据集（Set12 [101]、BSD68 [52]、Urban100 [29]、Kodak24 [20]以及麦克马斯特（McMaster）[104]）以及真实世界数据集（智能手机图像去噪数据集（SIDD）[1]和数字自然图像去噪数据集（DND）[60]）上进行去噪实验。按照[54, 93, 99]中的做法，我们使用无偏置的Restormer进行去噪。

高斯去噪。表4和表5分别展示了不同方法在若干用于灰度图像和彩色图像去噪的基准数据集上的峰值信噪比（PSNR）得分。与现有方法[44, 99]一致，我们在测试中包含了噪声水平为15、25和50的情况。被评估的方法分为两类实验情况：（1）学习单个模型来处理不同噪声水平；（2）针对每个噪声水平学习单独的模型。我们的Restormer在不同数据集和不同噪声水平下的这两种实验设置中都取得了最先进的性能。具体而言，对于高分辨率的Urban100数据集[29]中具有挑战性的噪声水平50的情况，如表4所示，Restormer相较于此前最佳的基于卷积神经网络（CNN）的方法深度残差去噪网络（DRUNet）[99]实现了0.37分贝的增益，相较于近期基于Transformer的网络SwinIR [44]实现了0.31分贝的提升。在表5所示的高斯彩色去噪中也能观察到类似的性能提升。值得一提的是，深度残差去噪网络（DRUNet）[99]需要噪声水平图作为额外输入，而我们的方法仅需输入含噪图像即可。此外，与SwinIR [44]相比，我们的Restormer的浮点运算次数（FLOPs）少了3.14倍，运行速度快了13倍。图6展示了不同方法进行灰度去噪（顶行）和彩色去噪（中行）的去噪结果。我们的Restormer能复原出清晰、明快的图像。

真实图像去噪。表6显示，我们的方法是唯一一种在两个数据集上峰值信噪比（PSNR）都超过40分贝的方法。值得注意的是，在智能手机图像去噪数据集（SIDD）上，我们的Restormer相较于此前最佳的基于卷积神经网络（CNN）的方法多尺度图像复原网络（MIRNet）[92]以及Transformer模型Uformer [80]，分别获得了0.3分贝和0.25分贝的峰值信噪比（PSNR）增益。图6（底行）显示，我们的Restormer在生成清晰图像的同时不会损害精细纹理。

4.5. 消融研究

在消融实验中，我们仅针对尺寸为128×128的图像块训练高斯彩色去噪模型，迭代次数为10万次。在Urban100数据集[29]上进行测试，并针对具有挑战性的噪声水平σ = 50进行分析。浮点运算次数（FLOPs）和推理时间是基于尺寸为256×256的图像来计算的。表7 - 10表明我们所做的改进提升了性能表现。接下来，我们将分别描述各个组件的影响。

多头注意力方面的改进。表7c显示，我们的多深度卷积头转置注意力（MDTA）相较于基线（表7a）带来了0.32分贝的良好增益。此外，通过深度可分离卷积为MDTA引入局部性可提高鲁棒性，因为去除它会导致峰值信噪比（PSNR）下降（见表7b）。

前馈网络（FN）方面的改进。表7d表明，前馈网络中用于控制信息流的门控机制相较于常规的前馈网络[77]带来了0.12分贝的增益。与多头注意力情况一样，在前馈网络中引入局部机制也带来了性能优势（见表7e）。我们通过纳入门控深度可分离卷积进一步强化了前馈网络。对于噪声水平为50的情况，我们的门控深度卷积前馈网络（GDFN）（见表7f）相较于标准的前馈网络[77]实现了0.26分贝的峰值信噪比（PSNR）增益。总体而言，我们Transformer模块所做的改进相较于基线带来了0.51分贝的显著增益。

第1层解码器的设计选择。为了在第1层将编码器特征与解码器特征进行聚合，我们在拼接操作之后未使用1×1卷积（该卷积会将通道数量减半）。如表8所示，这有助于保留来自编码器的精细纹理细节。这些结果进一步证明了在细化阶段添加Transformer模块的有效性。

渐进式学习的影响。表9显示，渐进式学习相较于固定图像块训练能取得更好的结果，且训练时间相近。

更深还是更宽的Restormer？表10表明，在参数/浮点运算次数预算相近的情况下，一个深度较深但宽度较窄的模型比宽度较宽但深度较浅的对应模型表现得更准确。不过，宽度较宽的模型由于可并行化，运行速度更快。在本文中，我们使用深度较深且宽度较窄的Restormer模型。

5. 结论

我们提出了一种用于图像复原的Transformer模型——Restormer，它在处理高分辨率图像时计算效率较高。我们针对Transformer模块的核心组件引入了关键设计，以改进特征聚合和变换。具体而言，我们的多深度卷积头转置注意力（MDTA）模块通过跨通道（而非空间维度）应用自注意力来隐式地对全局上下文进行建模，因此其复杂度是线性的，而非二次方的。此外，所提出的门控深度卷积前馈网络（GDFN）引入了门控机制来执行可控的特征变换。为了将卷积神经网络（CNN）的优势融入Transformer模型中，多深度卷积头转置注意力（MDTA）和门控深度卷积前馈网络（GDFN）模块都包含了深度可分离卷积，用于对空间局部上下文进行编码。在16个基准数据集上进行的大量实验表明，Restormer在众多图像复原任务中都取得了最先进的性能表现。