【源头活水】全新视觉 Transformer：高分辨率图像复原新SOTA——Restormer

“问渠那得清如许，为有源头活水来”，通过前沿领域知识的学习，从其他研究领域得到启发，对研究问题的本质有更清晰的认识和理解，是自我提高的不竭源泉。为此，我们特别精选论文阅读笔记，开辟“源头活水”专栏，帮助你广泛而深入的阅读科研文献，敬请关注。

来源：知乎—闪闪红星闪闪
地址：https://zhuanlan.zhihu.com/p/435090577

2021年11月19日，arXiv上发布了一篇全新的基于视觉 transformer 文章。该文章在 Transformer block 中进行了几个关键设计以提出一种新的高效的视觉 Transformer，称为 Restormer，再一次刷新了视觉 transformer 在高分辨率图像复原领域的性能并取得了新 SOTA。评估囊括了图像复原的主流任务，包括图像去雨、单图像运动去模糊、散焦去模糊和图像去噪。本文将完整详细地描述该论文，望大家喜欢。

概要

卷积神经网络 (CNN) 对于从大规模数据中学习泛化的图像先验方面表现良好，因此已广泛用于图像复原和相关任务中。而就在最近，另一类神经网络架构 Transformers 在自然语言和高级视觉任务上表现出了显着的性能提升。虽然 Transformer 模型减轻了 CNN 的缺点（即有限的感受野和对输入内容的不适应），但其计算复杂度随空间分辨率成二次方增长，因此无法应用于大多数涉及高分辨率图像的图像复原任务中。

在这项工作中，作者通过在构建块（多头注意力和前馈网络）中进行几个关键设计来提出一种高效的 Transformer 模型，以便它可以捕获远程像素交互，同时仍然适用于大图像。该模型取名为 Restoration Transformer (Restormer)，在多个图像复原任务上取得了SOTA结果，包括图像去雨、单图像运动去模糊、散焦去模糊（单图像和双像素数据）和图像去噪（高斯灰度/彩色去噪和真实图像去噪）。

01

论文和代码地址

论文题目：Restormer: Efficient Transformer for High-Resolution Image Restoration

发表单位：阿联酋起源人工智能研究院、穆罕默德·本·扎耶德人工智能大学、蒙纳士大学、加州大学默塞德分校、延世大学、谷歌研究院

论文地址：https://arxiv.org/abs/2111.09881

代码地址：https://github.com/swz30/Restor

提交时间：2021年11月19日

02

动机

2.1 CNN的局限性

图像复原是通过从退化（degraded）的图像输入中去除退化（例如，噪声、模糊、雨滴）来重建高质量图像的任务。由于不适定性质，这是一个极具挑战性的问题，通常需要强大的图像先验才能有效复原。而卷积神经网络 (CNN) 在从大规模数据中学习可泛化先验方面表现良好，因此与传统复原方法相比，它们已成为首选。

在 CNN 中，其基本操作是“卷积”，它为 CNN 提供了局部连接和平移不变性的特性。虽然这些特性为 CNN 带来了效率和泛化性，但它们也产生了两个主要问题：

1. 卷积算子的感受野有限，无法对长距离像素依赖性进行建模；

2. 卷积滤波器在推理时具有静态权重，不能灵活地适应输入内容。

为了解决上述缺点，一种更强大和动态的替代方案是自注意力 (SA) 机制，它通过所有其他位置的加权和来计算给定像素的输出特征。

2.2 使用 Transformer 替代 CNN

在本文中，作者提出了一种用于图像复原的高效 Transformer，它能够对全局连接进行建模，并且仍然适用于大图像。作者从自注意力（SA）、前馈网络（FN）对 transformer 进行改进，并使用一种渐进性学习策略来训练模型。详细地：

1. 引入多 Dconv 头“转置”注意力（multi-Dconv head "transposed" attention，MDTA）块来代替具有线性复杂性的普通多头 SA（详情见 3.1 节）。它跨特征维度而不是空间维度应用 SA，即，MDTA 不是显式建模成对像素交互，而是计算跨特征通道的交叉协方差，以从（关键和查询投影）输入特征获取注意力图。所提出的 MDTA 块的一个重要特征是特征协方差计算之前的局部上下文混合。这是通过使用 1×1 卷积的跨通道上下文的像素级聚合和使用有效的深度卷积（depthwise convolution）的局部上下文的通道级聚合来实现的。这种策略提供了两个关键优势。首先，它强调空间局部上下文，并在通道中引入卷积运算的互补强度。其次，它确保在计算基于协方差的注意力图时隐式建模像素之间的上下文全局关系。

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

3. 除了上述架构（上述架构本文称为 Restormer）新颖性之外，本文还展示了 Restormer 渐进式学习策略的有效性（详情见 3.3 节）。在这个过程中，网络在早期阶段对小的图像 patch 和大批量进行训练，在后期阶段对逐渐大的图像 patch 和小批量进行训练。这种训练策略帮助 Restormer 从大图像中学习上下文，并随后在测试时提供质量性能的改进。

本文进行了全面的实验，并在 16 个基准数据集上展示了 Restormer 在多个图像复原任务上的SOTA性能，包括图像去雨、单图像运动去模糊、散焦去模糊（在单图像和双像素数据上），和图像去噪（基于合成和真实数据）；参见下图 1。此外，本文提供了大量消融来展示架构设计和实验选择的有效性。

图1. Restormer 实现了 SOTA 性能在图像复原任务上

2.3 文章贡献

本文的主要贡献可以总结如下：

• 提出了 Restormer，这是一种编码器-解码器转换器，用于在高分辨率图像上进行多尺度局部-全局表示学习，而无需将它们分解为局部窗口，从而利用远处的图像上下文。

• 提出了一个多 Dconv 头部转置注意 (MDTA) 模块，该模块能够聚合局部和非局部像素交互，并且足够有效地处理高分辨率图像。

• 一种新的门控 Dconv 前馈网络 (GDFN)，它执行受控的特征转换，即抑制信息较少的特征，只允许有用的信息进一步通过网络层次结构。

03

方法

图2. Restormer 整体架构

本文所提方法的整体架构如上图 2 所示。其整体执行流程为：给定降质图像的图像  ，Restormer 首先使用一个卷积去获得低级特征嵌入  ，其中 H×W 表示空间维度而 C 表示通道数。接下来，这些浅层特征  通过一个 4 级对称的编码器-加码器并且转化为深层特征  。每一级编码器-解码器都包含多个 Transformer 块，块的数量从小到上逐渐增加，以保持效率。从高分辨率输入开始，编码器分层减少空间大小，同时扩展通道容量。解码器把低分辨率的潜在特征  作为输入并渐进式地恢复出高分辨率表示。

对于特征下采样和上采样，作者分别应用pixel-unshuffle和pixel-shuffle操作。为了帮助恢复过程，编码器的特征通过跳过连接与解码器的特征串联在一起。串联操作之后是 1×1 卷积，以减少所有层级的通道数量（减半），顶部的层级除外。在第1级，作者让 Transformer 块将编码器的低级图像特征与解码器的高级特征聚合在一起。这有助于在恢复的图像中保留精细的结构和纹理细节。

接下来，在高空间分辨率下操作的再精细阶段，深层特征  被进一步丰富。我们可以在实验部分中看到，这些设计选择提高了图像质量。最后，对细化后的特征采用卷积层生成残差图像 ，其中添加降质的图像以获得复原的图像。

3.1 多 Dconv 头部转置注意（Multi-Dconv Head Transposed Attention）

Transformers 中的主要计算开销来自于自注意力（SA）层。在传统的 SA 中，key-query点积交互的时间和内存复杂度与输入的空间分辨率成二次方增长，即对于 W×H 像素的图像，其复杂度为  。因此，将 SA 应用于经常涉及高分辨率图像的大多数图像复原任务是不可行的。为了缓解这个问题，作者提出了具有线性复杂度的 MDTA，如上图 2（a）所示。 

关键因素是跨通道而不是空间维度应用 SA，即计算跨通道的交叉协方差来生成隐式地编码全局上下文的注意力图。作为 MDTA 的另一个重要组成部分，作者引入了深度卷积，它在计算特征协方差来产生全局注意力图之前强调局部上下文。

在获得一个层归一化张量  之后，MDTA 首先前向传播生成 query（Q）、key（K）和 value（V），这丰富了局部上下文。它是通过应用 1×1 卷积来聚合像素别和跨通道的上下文，并应用 3×3 深度卷积来编码通道级空间上下文来实现的，产生  ，   。

其中，  是 1×1 逐点卷积（point-wise convolution），  是 3×3 的深度卷积。作者在网络中使用无偏置卷积层。接下来，作者重塑 query 和 key 投影，使得它们的点积交互生成一个大小为  的转置注意力图 A ，而不是大小为  的巨大的常规的注意力图。总的来说，MDTA 过程被定义为：

其中 X 和  是输入和输出特征图；矩阵 和  是从原始大小  对张量进行变形后获得的。这里 α 是一个可学习的缩放参数，用于在应用 softmax 函数之前控制  和  的点积的大小。与传统的多头 SA 相似，作者将通道数量划分为“头”并并行学习单独的注意力图。

3.2 门控 Dconv 前馈网络（Gated-Dconv Feed-Forward Network）

为了转换特征，常规的前馈网络 (FN) 分别且一致地对每个像素位置进行操作。它使用两个 1×1 卷积，一个是扩展特征通道（通常因子  ），第二个是将通道减少回原始的输入维度。在隐藏层中应用了非线性。在这项工作中，作者提出了 FN 中的两个基本修改来改进表示学习：（1）门控机制，以及（2）深度卷积。

所提出的 GDFN 的架构如上图 2（b）所示。门控机制被形式化为线性变换层的两条平行路径的元素乘积，其中之一被 GELU 非线性激活。与在 MDTA 中一样，作者还在 GDFN 中包含深度卷积来编码来自空间相邻像素位置的信息，这对于学习局部图像结构以进行有效的图像复原很有用。给定一个输入张量  ，GDFN 被形式化为如下形式：

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

3.3 渐进式学习（Progressive Learning）

在裁剪后的小 patch 上训练 Transformer 模型可能不会对全局的图像统计信息进行编码，从而在测试时的全分辨率图像上提供次优的性能。为此，作者使用渐进式学习，其中网络在早期阶段在较小的图像 patch 上进行训练，在后期的训练阶段中逐渐增大。通过渐进式学习在混合大小的 patch 上训练的模型在测试时表现出增强的性能，其中图像可以具备不同的分辨率（这是图像复原中的常见情况）。渐进式学习策略的行为方式类似于课程学习过程：网络从简单的任务开始，逐渐转向学习更复杂的任务（需要保留精细的图像结构/纹理）。由于对大 patch 的训练需要花费更长的时间，因此作者会随着 patch 大小的增加而减小批量大小，以保持每个优化步骤的时间与固定 patch 训练相似。

04

实验结果

实验有关的详细参数设置见原论文，这里将不赘述。

4.1 图像去雨

5个去雨数据集上的图像去雨的结果如下表 1 和图 3 所示。

表1. 图像去雨的结果

图3. 图像去雨的例子

4.2 单图像运动去模糊

GoPro数据集上的单图像运动去模糊的结果如下表 2 所示和图 4 所示。

表2. 单图像运动去模糊的结果

图4. 单图像运动去模糊的例子

4.3 散焦去模糊

DPDD测试集上的散焦去模糊的实验对比如下表 3 所示，包括单幅图像散焦去模糊和双像素散焦去模糊。双像素散焦去模糊对比案例如图 5 所示。

表3. DPDD测试集上的散焦去模糊的实验对比

图5. DPDD数据集上的双像素散焦去模糊比较

4.4 图像去噪

高斯灰度图像去噪在两类方法（CNN 变体 和 Transformer 变体）上的结果比较如下表 4 所示。

表4. 高斯灰度图像去噪在两类方法上的结果比较

高斯彩色图像去噪的实验结果如下表 5 所示。

表5. 高斯彩色图像去噪的实验结果

真实图像去噪的实验结果下表 6 所示。

表6. SIDD 和 DND 数据集上的真实图像去噪的实验结果

图像去噪的视觉结果如下图 6 所示。

图6. 图像去噪的视觉结果

4.5 消融实验

Transformer 块对结果的消融实验如下表 7 所示。

表7. Transformer 块的消融实验

在解码器层级1中的concat（w/o 1x1 conv）和细化阶段对实验结果的影响如下表 8 所示。

表8. 在编码器层级1中的concat和细化阶段对实验结果的影响

在固定 patch 大小和逐渐增大 patch 大小上训练 Restormer 的结果对比如下表 9 所示。

表9 在固定 patch 大小和逐渐增大 patch 大小上训练 Restormer 的结果对比

更深或更宽的模型对实验结果的影响如下表 10 所示。

表10. 更深或更宽

05

总结

本文提出了一种图像复原 Transformer 模型，Restormer，它在处理高分辨率图像方面具有计算效率。作者将关键设计引入到 Transformer 块的核心组件中，以改进特征聚合和转换。具体来说，提出的多 Dconv 头部转置注意力 (MDTA) 模块通过跨通道而不是空间维度应用自注意力来隐式地建模全局上下文，因此具有线性复杂性而不是二次方。

此外，所提出的门控 Dconv 前馈网络 (GDFN) 引入了门控机制来执行受控的特征转换。为了将 CNN 的强度整合到 Transformer 模型中，MDTA 和 GDFN 模块都包含用于编码空间局部上下文的深度卷积。 在 16 个基准数据集上的大量实验表明，Restormer 在众多图像复原任务中实现了SOTA性能。

本文目的在于学术交流，并不代表本公众号赞同其观点或对其内容真实性负责，版权归原作者所有，如有侵权请告知删除。

“源头活水”历史文章

更多源头活水专栏文章，

请点击文章底部“阅读原文”查看

分享、在看，给个三连击呗！