CVPR 2022 | Restormer：高分辨率图像恢复的高效Transformer

Restormer: Efficient Transformer for High-Resolution Image Restoration

一、Introduction

由于卷积神经网络（CNN）在从大规模数据中学习可概括的图像先验方面表现良好，因此这些模型已被广泛应用于图像恢复和相关任务。最近，另一类神经架构Transformers在自然语言和高级视觉任务上表现出显着的性能提升。虽然Transformer模型减轻了CNN的缺点（即，有限的接收场和对输入内容的不适应性），其计算复杂度随着空间分辨率二次增长，因此使得其不可行地应用于涉及高分辨率图像的大多数图像恢复任务。

在这项工作中，我们提出了一个有效的Transformer模型，通过在构建模块（多头注意力和前馈网络）中进行几个关键设计，使其可以捕获长距离像素交互，同时仍然适用于大图像。

我们的模型名为Restoration Transformer（Restormer），在多个图像恢复任务上实现了最先进的结果，包括图像去噪，单图像运动去模糊，散焦去模糊（单图像和双像素数据）和图像去噪（高斯灰度/彩色去噪和真实的图像去噪）。

二、Method

1. Network Architecture

在本文中，我们提出了一个有效的Transformer的图像恢复，是能够建模的全局连通性，仍然适用于大型图像。具体地，我们引入了多Dconv头“转置”注意力（MDTA）块（Sec.3.1)代替vanilla多头SA，其具有线性复杂度。它跨特征维度而不是空间维度应用SA，即，MDTA不是显式地对成对像素交互进行建模，而是计算跨特征通道的互协方差，以从（关键字和查询投影的）输入特征获得注意力图。

我们的MDTA块的一个重要特征是在特征协方差计算之前的局部上下文混合。这是通过使用1×1卷积的跨通道上下文的逐像素聚合和使用有效的逐深度卷积的局部上下文的逐通道聚合来实现的。该策略提供了两个关键优势。首先，它强调空间局部上下文，并在我们的流水线中引入卷积运算的互补优势。其次，它确保了像素之间的上下文全局关系被隐式地建模，同时计算基于协方差的注意力地图。

在这项工作中，我们用门控机制重新制定了规则FN的第一个线性变换层，以改善通过网络的信息流。该选通层被设计为两个线性投影层的逐元素乘积，其中一个用GELU非线性激活。我们的门控Dconv FN（GDFN）（Sec.3.2)也基于类似于MDTA模块的本地内容混合，以同样强调空间上下文。GDFN中的选通机制控制哪些互补特征应该向前流动，并允许网络层次结构中的后续层专门关注更精细的图像属性，从而产生高质量的输出。

除了上述架构的新颖性，我们显示了我们的Restormer渐进式学习策略的有效性。3.3）。在这个过程中，网络在早期的小块和大批量上进行训练，在后期的逐渐大的图像块和小批量上进行训练。这种训练策略有助于Restormer从大图像中学习上下文，并随后在测试时提供质量性能改进。

本文的主要贡献概述如下：

（1）我们提出了Restormer，一个编码器-解码器Transformer，用于在高分辨率图像上进行多尺度局部-全局表示学习，而不将其分解为局部窗口，从而利用远程图像上下文。

（2）我们提出了一个多Dconv头转置注意（MDTA）模块，能够聚合本地和非本地像素的相互作用，是足够有效的处理高分辨率图像。

（3）一种新的门控Dconv前馈网络（GDFN），其执行受控特征变换，即，抑制信息量较少的特征，并且仅允许有用的信息进一步通过网络分级结构。

我们的主要目标是开发一个高效的Transformer模型，可以处理高分辨率的图像恢复任务。为了缓解计算瓶颈，我们将关键设计引入多头SA层和多尺度分层模块，其计算要求低于单尺度网络。我们首先介绍了我们的Restormer架构的整体结构（见图2）的情况。然后，我们描述所提出的Transformer块的核心组件：（a）多Dconv头转置注意（MDTA）和（B）门控Dconv前馈网络（GDFN）。最后，我们提供了详细的渐进训练计划，有效地学习图像统计。

Figure2：用于高分辨率图像恢复的Restormer的体系结构。我们的Restormer由多尺度分层设计，结合高效的Transformer块。Transformer模块的核心模块有：（a）多Dconv头转置注意力（MDTA），其跨通道而不是空间维度执行（空间富集的）查询关键字特征交互，以及（b）门控Dconv前馈网络（GDFN），其执行受控特征变换，即，以允许有用信息进一步传播。

2、Multi-Dconv Head Transposed Attention

Transformers中的主要计算开销来自自我注意层。在传统的SA 中，键-查询点积交互的时间和存储器复杂度随着输入的空间分辨率二次增长，即，对于W×H像素的图像，

O(W^2H^2)

。因此，它是不可行的，以适用于大多数图像恢复任务，往往涉及高分辨率图像。

为了缓解这个问题，我们提出了MDTA，如图2（a）所示，具有线性复杂度。关键要素是跨通道而不是空间维度应用SA，即，计算跨信道的互协方差以生成隐式地编码全局上下文的注意力图。作为MDTA中的另一个重要组成部分，我们引入深度卷积，以强调在计算特征协方差以产生全局注意力图之前的局部上下文。

3、Gated-Dconv Feed-Forward Network

为了变换特征，常规前馈网络（FN）在每个像素位置上分别且相同地操作。它使用两个1×1卷积，一个用于扩展特征通道（通常是因子γ=4），第二个用于将通道减少回原始输入维度。在隐藏层中应用非线性。

在这项工作中，我们提出了FN的两个基本修改，以提高表示学习：（1）门控机制；（2）深度卷积。我们的GDFN的架构如图2（b）所示。门控机制被公式化为线性变换层的两个平行路径的元素式乘积，其中一个通过GELU非线性激活[27]。与MDTA一样，我们也在GDFN中包含深度卷积来编码信息。给定一个输入张量 X ∈ R ˆH× ˆW× ˆC，GDFN 的公式为：

其中

表示逐元素乘法，φ表示GELU非线性，LN是层归一化。总体而言，GDFN 控制着我们管道中各个层级的信息流，从而允许每个级别专注于与其他级别互补的细节。也就是说，与 MDTA（专注于利用上下文信息丰富特征）相比，GDFN 提供了独特的作用。由于与常规 FN相比，所提出的 GDFN 执行更多操作，因此我们降低了扩展率 γ，以便具有相似的参数和计算负担。

4、Progressive Learning

基于 CNN 的恢复模型通常在固定大小的图像块上进行训练。然而，在小裁剪补丁上训练 Transformer 模型可能不会对全局图像统计数据进行编码，从而在测试时在全分辨率图像上提供次优性能。为此，我们进行渐进式学习，其中网络在早期训练时期在较小的图像块上进行训练，在后期训练时期在逐渐增大的图像块上进行训练。

通过渐进式学习在混合大小的补丁上训练的模型在测试时表现出增强的性能，其中图像可以具有不同的分辨率（图像恢复中的常见情况）。渐进式学习策略的行为方式与课程学习过程类似，其中网络从更简单的任务开始，逐渐转向学习更复杂的任务（需要保留精细图像结构/纹理）。由于对大补丁的训练需要更长的时间，因此我们随着补丁大小的增加而减少批量大小，以保持每个优化步骤与固定补丁训练的时间相似。

三、Experiments

1、Image Deraining Results

2、Single-image Motion Deblurring Results

3、Defocus Deblurring Results

4、Image Denoising Results

四、Conclusion

我们提出了一种图像恢复 Transformer 模型 Restormer，它在处理高分辨率图像方面具有计算效率。

我们为 Transformer 模块的核心组件引入了关键设计，以改进特征聚合和转换。具体来说，我们的 Multi-Dconv 头部转置注意力（MDTA）模块通过跨通道而不是空间维度应用自注意力来隐式模拟全局上下文，因此具有线性复杂度而不是二次复杂度。此外，所提出的门控 DConv 前馈网络（GDFN）引入了门控机制来执行受控特征转换。为了将 CNN 的优势融入到 Transformer 模型中，MDTA 和 GDFN 模块都包含用于编码空间局部上下文的深度卷积。

对 16 个基准数据集的大量实验表明，Restormer 在众多图像恢复任务中实现了最先进的性能。