1. Motivation and Background
1.1 CNNs及其局限性
图像修复在于去除低质量图片上的噪音以重建高质量的图片,而修复图像需要强大的图像先验。CNNs在学习广尺度数据的泛化先验上有着良好的性能,故近年出现了很多基于CNNs的方法,和传统修复方法相较,人们更偏爱于前者。“卷积”是CNN基本操作,它能够提供局部连接和平移等方差,这使得其带来在效率和泛化能力上提升的同时,也导致了两个问题:卷积操作的感受野有限,这使得它不能建立大范围的像素关系和卷积过滤的权重在理论推导时是静态的,因此不能随着输入内容的变化而变化。为了解决这个问题,动态、强大的自注意力机制(SA)是一个很好的选择,它能够通过给输入像素赋予不同权重再求和来计算对应输出。
1.2 自注意力机制(SA:self-attention)
自注意力机制是Transformer的核心组件,它的实现和CNN不同。例如,多头注意力对并行化和高效表征学习进行了优化。Transformer在自然语言处理、高层视觉等任务上有着state-of-the-art的性能。尽管SA在捕获广域像素交互上性能较好,但是它的复杂度会随着空间分辨率的提高呈平方式增长,故直接将其用于高分辨率图像处理(通常是图像修复)上是不可行的。
所以虽然Transformer模型减轻了cnn的缺点(即有限的接受域和对输入内容的不适应性),但其计算复杂度随着空间分辨率呈二次增长,因此无法应用于大多数涉及高分辨率图像的图像恢复任务。将transformer调整应用于图像修复上的工作,为了减小计算负荷,这些方法要么将SA用于每个像素周围8x8的小空间窗口,要么将输入图像分割成无重叠的48x48的块,再在单个块上计算SA。但是这对SA空间内容的限制和捕获真正的广域像素关系是相违背的,特别是对于高分辨率图像来说。
1.3 Restormer
为解决上述问题,本论文提出了一个有效的Transformer模型,通过在构建块(多头注意和前馈网络)中进行几个关键设计,使其可以捕获远程像素交互,同时仍然适用于大图像。模型会在几个图像恢复任务上实现了最先进的结果,包括图像去噪、单图像运动去模糊、离焦去模糊(单图像和双像素数据)和图像去噪(高斯灰度/彩色去噪和真实图像去噪)。
2. Method
2.1 Overall Pipline
Restormer的整体执行流程如下:首先输入图像大小为 I∈R^{H×W×3} ,利用一个卷积操作获得特征嵌入Fo∈R^{H×W×C} 。接着F0通过一个对称的4层编码-解码结构,得到高维特征Fd∈RH×W×2C,每一层编码/解码都包括多个Transformer 模块,从上到下,每一层中的Transformer模块数量逐渐递增,分辨率逐渐递减。其次在编码-解码器之间使用跳跃连接来传递低维特征信息。Fd进一步经过Refinement模块来提取细节特征。最后经过一个卷积层,并与输入图像进行叠加,得到最后的输出图像。
Restormer的整体结构如下图,包括四层对称的编码器-解码器。其核心组件包括Mutli-Dconv head Transposed Attention(MDTA)和Gated-Dconv Feed-forward Network(GDFN)
2.2 Mutli-Dconv head Transposed Attention
为了缓解传统SA中点积交互的时间和内存复杂度与输入空间分辨率成二次方增长,即对于W×H像素的图像复杂度为O(W2×H2) ,本文通过把HW×WH的attention优化成通道×通道的attention,使得MDTA具有线性复杂度,它有助于进行局部和非局部相关像素聚合,实现高效的高分辨率图像处理。Restormer模型在获得归一化张量Y后,通过1×1卷积聚合跨通道的上下文,再使用3×3深度卷积来编码通道级上下文,生成了query、key、value(即图中的)投影。接着应用softmax函数控制K和Q的点积大小,并通过SA计算生成转置注意图A。
MDTA的过程被定义为如下公式:
该模块结构示意图和过程定义如下:
2.3 Gated-Dconv Feed-forward Network
在传统的前馈网络(FN)中,特征转换是通过在每个像素位置独立且相同地操作来实现的。它使用两个 1×1 卷积,一个用于扩展特征通道(通常扩大 γ=4 倍),另一个用于将通道数减少回原始输入维度。在隐藏层应用非线性激活函数。在这项研究中,我们提出了两个基本修改以改善特征表示学习:(1)门控机制和(2)深度可分离卷积。我们的 GDFN 架构如图2(b)所示。门控机制被构建为线性变换层的两个并行路径的逐元素乘积,其中一个路径使用 GELU 非线性激活函数激活。与MDTA类似,我们还在GDFN中包含了深度可分离卷积,以编码来自空间相邻像素位置的信息,这对于学习有效恢复所需的本地图像结构非常有用。
给定输入张量,GDFN 的公式为:
其中 · 表示逐元素乘法,φ表示 GELU 非线性激活函数,LN 是层归一化操作。总体上,GDFN 控制着我们管道中各层次的信息流,从而使每个层次都能专注于与其他层次互补的精细细节。也就是说,与MDTA(专注于丰富上下文信息的特征)相比,GDFN 的作用更为独特。由于所提出的 GDFN 比常规 FN执行的操作更多,我们减小了扩张比例 γ,以保持相似的参数数量和计算负担。
2.4 Progressive Learning
在裁剪后的小 patch 上训练 Transformer 模型可能不会对全局的图像统计信息进行编码,从而在测试时的全分辨率图像上提供次优的性能。因此本问使用了渐进式学习,其中网络在早期阶段用较小的图像 patch 上进行训练,在后期的训练阶段中逐渐增大。通过渐进式学习在混合大小的 patch 上训练的模型在测试时表现出增强的性能,其中图像可以具备不同的分辨率。渐进式学习策略的行为方式类似于课程学习过程:网络从简单的任务开始,逐渐转向学习更复杂的任务,同时保留精细的图像结构。由于对大 patch 的训练需要花费更长的时间,因此会随着 patch 大小的增加而减小批量大小,以保持每个优化步骤的时间与固定 patch 训练相似。
3. Experiment and Analysis
3.1 Experiment Setting
Restormer 在四种图像处理任务上进行了实验评估:(a)图像去雨水、(b)单图像运动去模糊、(c)焦外模糊去除(针对单图像和双像素数据)、以及(d)图像去噪(针对合成和真实数据),同时为不同的图像恢复任务训练了单独的模型。
我们从模型架构、精化阶段、训练参数、渐进式学习和数据增强五个方面说明Restormer的实验设置。Restormer 使用了一个包含四个级别的编码器-解码器结构。从第一级到第四级,Transformer 块的数量分别为 [4, 6, 6, 8],MDTA 中的注意力头数为 [1, 2, 4, 8],通道数量分别为 [48, 96, 192, 384]。在精化阶段,Restormer包含 4 个块,并设置GDFN 中的通道扩展因子:γ=2.66。实验使用了 AdamW 优化器(参数为 β1=0.9,β2=0.999,权重衰减为 1e−4)和 L1 损失函数进行了 300K 次迭代的训练,初始学习率为 3e−4,并逐渐减少到 1e−6(采用余弦退火 [51])。为了渐进式学习,实验从 128×128 的图块大小和批处理大小为 64 开始训练,随着迭代次数的增加,图块大小和批处理大小的配对逐步更新为 [(1602, 40), (1922, 32), (2562, 16), (3202, 8), (3842, 8)] 分别在 [92K, 156K, 204K, 240K, 276K] 次迭代时更新。最后使用水平和垂直翻转技术进行数据增强。
3.2 Experiment Result
3.2.1 Image Deraining
5个去雨数据集上的图像去雨的结果如下表 1 和图 3 所示。
3.2.2 Single-image Motion Deblurring
GoPro数据集上的单图像运动去模糊的结果如下表 2 所示和图 4 所示。
3.2.3 Defocus Deblurring
DPDD测试集上的散焦去模糊的实验对比如下表 3 所示,包括单幅图像散焦去模糊和双像素散焦去模糊。双像素散焦去模糊对比案例如图 5 所示。
3.2.4 Image Denoising
高斯灰度图像去噪在两类方法(CNN 变体 和 Transformer 变体)上的结果比较如下表 4 所示。
高斯彩色图像去噪的实验结果如下表 5 所示。
真实图像去噪的实验结果下表 6 所示。
图像去噪的视觉结果如下图 6 所示。
3.2.5 Ablation Studies
Transformer 块对结果的消融实验如下表 7 所示。
在解码器层级1中的concat(w/o 1x1 conv)和细化阶段对实验结果的影响如下表 8 所示。
在固定 patch 大小和逐渐增大 patch 大小上训练 Restormer 的结果对比如下表 9 所示。
4. Conclusion
Restormer,它是计算效率高的处理高分辨率图像模型。我们介绍了Transformer模块核心组件的关键设计,以改进特征聚合和转换。具体来说,我们的multiDconv头部转置注意(MDTA)模块通过跨通道而不是空间维度应用自注意来隐式地模拟全局上下文,因此具有线性复杂性而不是二次复杂性。此外,所提出的门控-直流前馈网络(GDFN)引入了一种门控机制来进行受控特征转换。为了将cnn的强度整合到Transformer模型中,MDTA和GDFN模块都包括用于编码空间局部上下文的深度卷积。在16个基准数据集上的大量实验表明,Restormer在许多图像恢复任务中实现了最先进的性能。