摘要
基于Transformer的方法在低级视觉任务(如图像超分辨率)上表现出令人印象深刻的性能。然而,我们发现这些网络只能通过归因分析利用有限的输入信息空间范围。这意味着Transformer的潜力在现有网络中仍未得到充分利用。为了激活更多输入像素以获得更好的重构效果,我们提出了一种新颖的混合注意力Transformer(Hybrid Attention Transformer,HAT)。它结合了通道注意力和基于窗口的自注意力方案,从而利用全局统计和强大的局部拟合能力。此外,为了更好地聚合跨窗口信息,我们引入了一个重叠的跨注意力模块,以增强相邻窗口特征之间的交互。在训练阶段,我们还采用了同任务预训练策略来开发模型的潜力以进一步提高性能。大量实验表明了所提模块的有效性,并且我们进一步扩大了模型规模以证明该任务的可性能得到极大改进。我们的整体方法比最先进的方法高出1dB以上,从而证明了其有效性。
1. Introduction
单一图像超分辨率(SR)是计算机视觉和图像处理中的经典问题。它的目标是从给定的低分辨率输入中重建出高分辨率图像。自从深度学习被成功地应用于SR任务以来,已经提出了许多基于卷积神经网络(CNN)的方法[8, 11, 12, 24, 29, 32, 68, 70],这些方法在过去的几年中几乎主导了这一领域。最近,由于在自然语言处理方面的成功,Transformer[53]已经引起了计算机视觉界的关注。在高级视觉任务上取得快速进步后[14, 39, 54],基于Transformer的方法也被开发用于低级视觉任务[6, 57, 65],以及SR[27, 31]。特别是,新设计的网络SwinIR[31]在这个任务上取得了突破性的改进。
尽管取得了成功,但“为什么Transformer比CNN更好”仍然是一个谜。一个直观的解释是,这种网络可以从自注意力机制中受益,并利用长期信息。因此,我们采用归因分析方法LAM[15]来检查SwinIR中用于重建的信息所涉及的范围。有趣的是,我们发现SwinIR在超分辨率方面并没有比基于CNN的方法(例如RCAN[68])利用更多的输入像素,如图2所示。此外,尽管SwinIR平均获得了更高的定量性能,但由于利用信息的范围有限,它在一些样本上产生了不如RCAN的结果。这些现象表明,Transformer具有更强的建模局部信息的能力,但其利用信息的范围需要扩大。此外,我们还发现SwinIR的中间特征中会出现阻塞伪影,如图3所示。这表明移位窗口机制( shift window mechanism)不能完全实现跨窗口信息交互。
为了解决上述限制并进一步开发用于SR的Transformer的潜力,我们提出了一种混合注意力Transformer(Hybrid Attention Transformer),即HAT。我们的HAT结合了通道注意力和自注意力方案,以利用前者的使用全局信息的能力和后者的强大代表性能力。此外,我们引入了一个重叠的交叉注意力模块,以实现相邻窗口特征之间的更直接交互。得益于这些设计,我们的模型可以激活更多的像素进行重建,从而获得显着的性能改进。
由于Transformers不像CNN那样具有归纳偏差,因此大规模数据预训练对于解锁这种模型的潜力是重要的。在这项工作中,我们提供了一种有效的相同任务的预训练策略。与使用多个恢复任务的IPT [6]和使用多个降级级别的EDT [27]不同,我们直接使用大型数据集对相同任务进行预训练。我们认为,大规模数据是预训练真正重要的因素,实验结果也显示了我们策略的优越性。装备上述设计后,HAT可以以巨大的优势超越最先进的方法(0.3dB~1.2dB)。
Contributions
1)我们设计了一种新型的混合注意力变压器(HAT),它结合了自注意力、通道注意力以及一个新的重叠交叉注意力,以激活更多的像素进行更好的重建。
2)我们提出了一种有效的相同任务的预训练策略,以进一步利用SR变压器的潜力,并展示了大规模数据预训练对该任务的重要性。
3)我们的方法达到了先进的性能水平。通过进一步扩展HAT来构建一个大型模型,我们大大提高了该SR任务的性能上限。
3. Methodology
3.1. Motivation
Swin Transformer [39]在图像超分辨率[31]领域已经表现出了出色的性能。那么我们很想知道是什么让它比基于CNN的方法表现更好。为了揭示其工作机制,我们求助于一种诊断工具-LAM [15],这是一种针对SR的归属方法。通过LAM,我们可以判断哪些输入像素对选定区域的贡献最大。如图2所示,红色标记的点是有信息的像素,它们对重建有贡献。
直观地说,利用的信息越多,就可以获得更好的性能。这对基于CNN的方法来说是正确的,比如比较RCAN [68]和EDSR [32]。然而,对于基于Transformer的方法-SwinIR,其LAM并没有显示出比RCAN更大的范围。这与我们的常识相矛盾,但也可以为我们提供额外的见解。首先,这暗示了SwinIR具有比CNN更强的映射能力,因此可以使用更少的信息来获得更好的性能。其次,SwinIR可能会因为所利用的像素范围有限而恢复错误的纹理,我们认为如果它可以利用更多的输入像素,它可以进一步改进。因此,我们的目标是设计一个网络,它可以在激活更多像素进行重建的同时利用类似的自注意力。如图2所示,我们的HAT几乎可以看到图像的所有像素并恢复正确和清晰的纹理。
此外,我们可以观察到SwinIR的中间特征中存在明显的块效应伪影,如图3所示。这些伪影是由窗口划分机制引起的,这表明移位窗口机制在建立跨窗口连接时效率低下。一些针对高级视觉任务的作品[13,20,44,60]也指出,增强窗口之间的连接可以提高基于窗口的自注意力方法。因此,在设计我们的方法时,我们增强了跨窗口的信息交互,HAT获得的中间特征中的阻塞伪影得到了显著缓解。
3.2. Network Architecture
3.2.1 The Overall Structure
如图4所示,整个网络由三个部分组成,包括浅层特征提取、深层特征提取和图像重建。这种架构设计在以前的工作中得到了广泛的应用[31,68]。具体来说,对于给定的低分辨率(LR)输入 I L R ∈ R H × W × C i n I_{LR}∈R^{H×W×Cin} ILR∈RH×W×Cin,我们首先利用一个卷积层提取浅层特征 F 0 ∈ R H × W × C F_0∈R^{H×W×C} F0∈RH×W×C,其中 C i n C_{in} Cin和C表示输入和中间特征的通道数。然后,一系列残差混合注意力组(RHAG)和一个3×3卷积层HConv(·)被用来进行深层特征提取。之后,我们添加一个全局残差连接来融合浅层特征 F 0 F_0 F0和深层特征 F D ∈ R H × W × C F_D∈R^{H×W×C} FD∈RH×W×C,然后通过重建模块重建高分辨率结果。如图4所示,每个RHAG包含多个混合注意力块(HAB)、一个重叠交叉注意力块(OCAB)和一个具有残差连接的3×3卷积层。对于重建模块,采用了像素洗牌方法[47]来上采样融合的特征。我们简单地使用L1损失来优化网络参数。
3.2.2 Hybrid Attention Block (HAB)
如图2所示,采用通道注意力时,更多的像素被激活,因为全局信息被用于计算通道注意力的权重。此外,许多作品表明卷积可以帮助Transformer获得更好的视觉表示或更容易优化[26,59,61,63,71]。因此,我们将一个基于通道注意力的卷积块集成到标准Transformer块中,以增强网络的表示能力。如图4所示,我们在标准Swin Transformer块中第一个LayerNorm(LN)层后面并联插入一个通道注意力块(CAB),与基于窗口的多头自注意力(W-MSA)模块并联。需要注意的是,类似于[31,39],我们在连续的HAB之间采用了移位的窗口自注意力(SW-MSA)。为了避免CAB和MSA在优化和视觉表示上可能存在的冲突,我们将一个小常数α乘以CAB的输出。对于给定的输入特征X,HAB的整个计算过程如下所示:
其中 X N X_N XN和 X M X_M XM表示中间特征。Y表示HAB的输出。特别是,我们将每个像素视为一个token进行嵌入(即,将patch大小设置为1以进行patch embedding,见[31])。MLP表示多层感知器。对于自注意力模块的计算,给定大小为H×W×C的输入特征,它首先被划分为 H W / M M HW/MM HW/MM个大小为M×M的局部窗口,然后在每个窗口内计算自注意力。对于局部窗口特征 X W ∈ R M 2 × C X_W∈R^{M^2×C} XW∈RM2×C,query、key和value矩阵通过线性映射计算得到Q、K和V。然后基于窗口的自注意力可以表示为
其中d表示query/key的维度。B表示相对位置编码,计算方法见[53]。需要注意的是,我们使用大窗口大小计算自注意力,因为发现它显着扩大了所用像素的范围,如第4.2节所述。此外,为了在相邻的非重叠窗口之间建立连接,我们还采用了移位窗口划分方法[39],并将移位大小设置为窗口大小的一半。CAB由两个具有GELU激活函数的标准卷积层和一个通道注意力(CA)模块组成,如图4所示。由于基于Transformer的结构通常需要大量的通道来进行token嵌入,直接使用具有恒定宽度的卷积会导致计算量大。因此,我们将两个卷积层的通道数压缩一个常数β。对于具有C个通道的输入特征,第一层卷积之后输出特征的通道数被压缩到Cβ,然后通过第二层将特征扩展到C个通道。接下来,我们利用标准CA模块[68]自适应地重新调整通道特征的权重。
3.2.3 Overlapping Cross-Attention Block (OCAB)
我们引入了OCAB来直接建立跨窗口连接并增强窗口自注意力的表示能力。我们的OCAB由一个重叠交叉注意力(OCA)层和一个与标准Swin Transformer块类似的MLP层组成[39]。但对于OCA,
如图5所示,我们使用不同的窗口大小来对投影的特征进行分区。具体来说,对于输入特征X的 X Q 、 X K 、 X V ∈ R H × W × C X_Q、X_K、X_V∈R^{H×W×C} XQ、XK、XV∈RH×W×C, X Q X_Q XQ被分成 H W / M M HW/MM HW/MM个大小为 M × M M×M M×M的非重叠窗口,而 X K 、 X V X_K、X_V XK、XV被展开成HW个大小为Mo×Mo的重叠窗口。它的计算方式如下:
其中γ是一个常数,用于控制重叠的大小。为了更好地理解这一操作,标准窗口划分可以被认为是使用大小和步长都等于窗口大小M的核进行滑动划分。相比之下,重叠窗口划分可以看作是使用大小等于Mo的核进行滑动划分,而步长等于M。使用大小为γM^2的零填充来确保重叠窗口的大小一致性。注意力矩阵按式2计算,相对位置偏置B∈RM×Mo也被采用。不同于WSA将查询、键和值从相同的窗口特征计算出来,OCA从更大的领域计算键/值,其中可以利用更多有用的信息用于查询。需要注意的是,尽管[44]中的多分辨率重叠注意力(MOA)模块执行类似的重叠窗口划分,但我们的OCA与MOA根本不同,因为MOA使用窗口特征作为token来计算全局注意力,而OCA在每个窗口特征中使用像素token来计算交叉注意力。
3.3. The Same-task Pre-training
预训练被证明对许多高级视觉任务是有效的【1、14、16】。最近的工作【6、27】还表明,预训练对低级视觉任务也是有益的。IPT【6】强调使用各种低级任务,如去噪、去雨、超分辨率等,而EDT【27】利用特定任务的不同退化程度进行预训练。这些工作侧重于探讨目标任务的多元任务预训练的效果。相比之下,我们直接在大规模数据集(即ImageNet【9】)上基于相同任务进行预训练,表明预训练的有效性更多地取决于数据的规模和多样性。例如,当我们想要训练一个×4 SR模型时,我们首先在ImageNet上训练一个×4 SR模型,然后在其特定数据集(如DF2K)上进行微调。所提出的策略称为相同任务的预训练,它更简单,同时带来了更多的性能改进。值得一提的是,预训练的充分训练迭代和微调的适当小学习率对预训练策略的有效性非常重要。我们认为这是因为Transformer需要更多数据和迭代来学习任务的一般知识,但需要小学习率进行微调以避免过度拟合特定数据集。