CVPR 2023 | 即插即用的注意力模块 HAT: 激活更多有用的像素助力low-level任务显著涨点!

1,024 阅读9分钟

本文首发于微信公众号 CVHub,严禁私自转载或售卖到其他平台,违者必究。

Title: Activating More Pixels in Image Super-Resolution Transformer
Paper: arxiv.org/pdf/2205.04…
Code: github.com/XPixelGroup…

导读

本文提出了一种名为Hybrid Attention Transformer (HAT)的方法,旨在通过结合深度学习技术和注意力机制来改进图像超分辨率任务。

单图像超分辨率(SR)任务是计算机视觉和图像处理领域的一个经典问题,其目标是从给定的低分辨率输入中重建高分辨率图像。作者通过注意力分析(attribution analysis)的方法发现现有的 Transformer 方法在利用输入信息时只能使用有限的空间范围。这意味着目前的网络还没有充分发挥 Transformer 的潜力。

因此,为了激活更多的输入像素以实现更好的重建效果,本文构建了一种新颖的混合注意力 Transformer(HAT)方法。该方法结合了通道注意力和基于窗口的自注意力机制,充分利用它们在利用全局统计信息和强大的局部拟合能力方面的互补优势。此外,为了更好地聚合跨窗口信息,作者引入了一个重叠的交叉注意力模块,增强了相邻窗口特征之间的交互作用。在训练阶段,还采用了相同任务的预训练策略,以进一步挖掘模型的潜力以获得更好的性能。

最终,通过广泛的实验证明了所提出方法的有效性,并通过扩展模型规模进一步展示了任务性能的显著提升。整体而言,该方法在性能上显著优于当前最先进的方法,提升幅度超过 1dB。

动机

首先,我们先分析下Swin Transformer。它在图像超分辨率任务中表现出色,但很多时候我们并不清楚它相比基于 CNN 的方法有何优势。为了揭示该方法的工作机制,研究人员使用了一种名为LAM(Localization-Aware Mapping)的诊断工具,用于 SR 任务的归因分析。

LAM 能够显示出哪些输入像素对于重建贡献最大。通过 LAM 的分析,我们可以发现 SwinIR(基于Transformer的方法)在利用的信息范围上并不比基于 CNN 的方法(如RCAN)更广泛。这与常识相矛盾,但也给作者带来了额外的启示。

首先,这表明 SwinIR 比 CNN 具有更强的映射能力,因此可以利用较少的信息实现更好的性能。其次,由于利用的像素范围有限,SwinIR 可能会恢复错误的纹理,因此如果能够利用更多输入像素,它的性能可能进一步提升。因此,研究人员的目标是设计一种网络,既能够利用类似的自注意力机制,又能够激活更多像素进行重建。他们的 HAT 网络在图中几乎可以看到整个图像,并且能够恢复正确和清晰的纹理。

此外,通过上图我们可以观察到 SwinIR 的中间特征存在明显的阻塞伪影(blocking artifacts)。这些伪影是由窗口划分机制引起的,这表明移动窗口机制在建立窗口间的连接上效率较低。一些用于高层次视觉任务的研究也指出,增强窗口之间的连接可以改善基于窗口的自注意力方法。因此,当设计本文方法时,作者特意加强了窗口间的信息交互,从而显著减轻了 HAT 所得到的中间特征中的阻塞伪影。

方法

Framework

如上图所示,整体网络由三个部分组成,包括浅层特征提取深层特征提取图像重建。这种架构设计在之前的研究中得到了广泛应用。具体而言,对于给定的低分辨率(LR)输入,首先利用一个卷积层提取浅层特征。然后,采用一系列残差混合注意力组(RHAG)和一个3×3的卷积层进行深层特征提取。在此之后,图中添加了一个全局残差连接,将浅层特征和深层特征融合起来,然后通过重建模块重建高分辨率结果。

此外,每个 RHAG 包含多个混合注意力块(HAB),一个重叠的交叉注意力块(OCAB)和一个带有残差连接的3×3卷积层。对于重建模块,采用像素洗牌(pixel-shuffle)方法来上采样融合的特征。同样地,本文中简单地使用 L1 损失来优化网络参数。

Hybrid Attention Block (HAB)

HAB 用于结合不同类型的注意力机制来激活更多的像素,以实现更好的重建效果。该模块由两个关键组成部分组成:窗口自注意力机制(Window-based Self-Attention)和通道注意力机制(Channel Attention)。

在 HAB 模块中,首先将输入特征进行归一化处理,然后利用窗口自注意力机制对特征进行处理。窗口自注意力机制将输入特征划分为局部窗口,并在每个窗口内计算自注意力。这样可以捕捉到局部区域的关联信息。接下来,通过通道注意力机制,全局信息被引入,用于计算通道注意力权重。通道注意力机制能够利用全局信息对特征进行加权,从而激活更多的像素。

HAB 模块的输出是窗口自注意力机制和通道注意力机制的加权和,并通过残差连接与输入特征相加。这种设计使得网络能够同时利用局部和全局信息,从而实现更好的重建效果。

Overlapping Cross-Attention Block (OCAB)

Overlapping Cross-Attention Block (OCAB) 模块则通过引入重叠交叉注意力层,在窗口自注意力中建立了窗口之间的交叉连接,以增强网络的表征能力。该模块的设计可以更好地利用窗口内部的像素信息进行查询,从而提高重建任务的性能。

The Same-task Pre-training

预训练在高级视觉任务中已经证明是有效的,而最近的研究也表明它对低级视觉任务有益处。有些方法强调使用多个低级任务进行预训练,例如去噪、去雨、超分辨率等,而另一些方法则利用了特定任务的不同降级水平进行预训练。然而,与这些方法不同的是,本研究直接在一个更大规模的数据集(如ImageNet)上使用相同任务进行预训练,结果显示预训练的有效性更多地取决于数据的规模和多样性。例如,在进行×4超分辨率模型训练时,首先在ImageNet上训练一个×4超分辨率模型,然后在特定数据集(如DF2K)上进行微调。这种相同任务预训练的策略更加简单,但却能够带来更好的性能提升。

不过大家需要注意的是,预训练的有效性取决于充分的训练迭代次数以及在微调阶段采用适当的小学习率,这是非常重要的。Transformer 模型需要更多的数据和迭代次数来学习任务的通用知识,但在微调时需要小的学习率以避免对特定数据集的过拟合。因此,预训练阶段需要足够的时间和数据来学习通用特征,而微调阶段则需要细致调整以适应特定任务的要求。

实验

正如我们上面所讨论的,激活更多的输入像素有助于实现更好的超分辨率性能。扩大窗口尺寸是实现这一目标的直观方法。在先前的一些相关工作中,有人研究了不同窗口尺寸的影响。然而,这些实验是基于移位交叉局部注意力的,并且只探索了最大为12×12的窗口尺寸。

本文进一步探究了自注意力的窗口尺寸如何影响表示能力。为了消除新引入的模块的影响,作者直接在SwinIR上进行以下实验。如表1所示,窗口尺寸为16×16的模型在性能上表现更好,尤其是在Urban100数据集上。此外,本文还在图6中提供了定性比较。对于标记为红色的补丁区域,窗口尺寸为16的模型利用了比窗口尺寸为8的模型更多的输入像素。重建结果的定量性能也证明了大窗口尺寸的有效性。基于这个结论,作者将窗口尺寸16直接作为默认设置。

从表6中可以看出,通过比较HAT和HAT†的性能,我们可以看到HAT可以极大地受益于预训练策略。为了展示所提出的同任务预训练的优越性,作者还将多相关任务预训练方法应用于HAT进行比较,并在完整的ImageNet数据集上使用与相同的训练设置,实验做得还是蛮充分的。

总结

本文提出了一种名为 HAT 的新型超分辨率 Transformer 方法,通过结合不同类型的注意力机制和大规模数据预训练,实现了更好的图像重建效果。该方法在实验证明了其在超分辨率任务中的优越性能,并超过了当前最先进方法。这项研究拓展了Transformer在计算机视觉任务中的应用,并提供了一种改进低级视觉任务的方法。


CVHub是一家专注于计算机视觉领域的高质量知识分享平台,全站技术文章原创率达99%,每日为您呈献全方位、多领域、有深度的前沿AI论文解决及配套的行业级应用解决方案,提供科研 | 技术 | 就业一站式服务,涵盖有监督/半监督/无监督/自监督的各类2D/3D的检测/分类/分割/跟踪/姿态/超分/重建等全栈领域以及最新的AIGC等生成式模型。关注微信公众号,欢迎参与实时的学术&技术互动交流,领取CV学习大礼包,及时订阅最新的国内外大厂校招&社招资讯!