ICCV 2023 | 南开程明明团队提出新颖注意力机制用于图像超分辨率任务

导读

TL;DR: 本文提出了一种新的置换自注意力（PSA）用于图像超分辨率任务，可以在保持计算成本较低的同时，处理大窗口自注意力，并基于PSA构建了一个新的超分辨率网络，SRFormer，在各项基准测试中取得了最先进的性能。

先简单介绍一下图像超分辨率（SR）的背景知识。SR 是计算机视觉中的一个重要任务，目的是从低分辨率的图像恢复出高质量的版本。早期，以 CNNs 为主的方法主导了 SR 领域，通过利用残差学习，密集连接或者通道注意力来构建网络结构。然而，近期的研究显示，基于Transformer的模型在 SR 任务上展现出了更好的表现，其中一个典型的例子就是SwinIR，它采用了Swin Transformer的技术，在多个基准测试中大大提升了SR模型的表现。【其实视觉任务的性能大部分取决于预训练骨干网络】

然而，这些基于 Transformer 的模型在提高窗口大小以优化性能的同时，也带来了相应的计算负担。这就提出了一个问题："如果我们在减少通道数量的同时，增大窗口大小，会产生怎样的效果？"

针对这个问题，本文提出了一种新的方法，即置换自注意力。PSA 的设计目标是在大窗口（例如24x24）内建立有效的两两关系，而不增加额外的计算负担。为此，作者在键和值矩阵上减少了通道的维度（压缩），并采用了一种置换操作，将部分空间信息传递到通道维度（置换）。这种方法即使在通道减少的情况下，也不会损失空间信息，并且每个注意力头可以保持适当数量的通道以产生表达力强的注意力图。

最后，基于PSA，论文最终构建了一个新的SR网络，称为SRFormer，并在五个广泛使用的数据集上对SRFormer进行了评估。结果显示，SRFormer在几乎所有数据集上的性能都有所提高。尤其在2倍的超分辨率任务中，仅用DIV2K数据集训练的SRFormer在Urban100数据集上取得了 33.86 的 PSNR 得分，比 SOTA 模型 SwinIR（33.40）和 ELAN（33.44）的结果都要高。

方法

如上图所示，SRFormer 的架构分为三部分：

像素嵌入层：将低分辨率的RGB图像转换为特征嵌入，这些嵌入会被送入特征编码器进行处理。
特征编码器：包含了 N 个排列自注意力组，每个组由 M 个 PAB 和一个 3×3 卷积层构成。
高分辨率图像重建层：最后，特征编码器的输出和特征嵌入的残差，会被送入高分辨率图像重建层，从而得到高分辨率图像。

此处，排列自注意块（PAB）是整个模型的核心，它包含了置换注意力层(PSA)和卷积前馈网络(ConvFFN)。这里我们可以简单看下 PSA 和常规的 MHSA 的区别：

PSA 层的主要思想是，它将输入特征图划分成 N 个非重叠的方块窗口，然后通过三个线性层得到 Q，K，V 三组值。这里的 $Q$ 保持与输入 $X$ 同样的通道维度，而 $K$ 和 $V$ 的通道维度被压缩。随后，作者提出了一种将空间信息排列到通道维度的方法，使得每个注意力头能生成更具表现力的注意力图。最后，利用 $Q$ ， $K_{p}$ ， $V_{p}$ 来进行自注意力操作。

扫描二维码关注公众号，回复： 16172967 查看本文章

ConvFFN 则是一个辅助部分，它在自注意之后添加一个 3x3 卷积层以恢复高频信息。作者发现，这样的操作几乎不增加计算量，但能够补偿自注意操作过程中造成的高频信息的损失。

此外，文中还介绍了两种基于大窗口自注意力的变体，一种是减少 token 的数量，另一种是对 token 进行随机采样。尽管这两种方法都可以在保持计算成本不变的前提下实现大窗口自注意力，但是它们都有一些缺点，例如，token 采样会丢失图像的结构信息，这对于图像超分辨率来说影响还是蛮大滴。

实验

消融实验

:::block-1

PSA 提供了一种有效且高效的方法来增大窗口大小。为了研究不同窗口大小对模型性能的影响，文中进行了三组实验：

第一组实验是 vanilla SwinIR 模型，窗口大小为 8 × 8，12 × 12，和 16 × 16。
第二组实验是不使用 ConvFFN，只在 SRFormer 中使用 PSA，窗口大小设置为 12 × 12，16 × 16，和 24 × 24，以观察性能差异。
第三组实验使用完整的 SRFormer，并设置窗口大小为 12×12，16×16，和 24×24，以探索性能的变化。

结果表明，对于所有三组实验，更大的窗口大小都会带来更好的性能提升。此外，使用24 × 24窗口的 SRFormer 参数和 MACs（乘法和累加运算）甚至比使用 8 × 8 窗口的原始 SwinIR 还要少。为了平衡性能和 MACs，他们将 SRFormer 的窗口大小设置为 24 × 24，而 SRFormer-light 的窗口大小设置为 16 × 16。

:::

:::block-1

如上所述，我们提到文中在每个块的最后添加了 ConvFFN，其目的是在不增加过多计算量的情况下编码更多的局部信息。为了探索哪种内核大小可以带来最好的性能提升，论文尝试使用 3×3 深度卷积和 5×5 深度卷积。由于深度卷积对参数数量和MACs（乘法和累加运算）的影响很小，因此在表中并未列出。显然，5 × 5深度卷积带来了最好的结果。因此，SRFormer 中使用 5 × 5 深度卷积。

:::

:::block-1

尽管 token 减少可以在使用大窗口时略微提高 SwinIR 的性能，但参数数量并未减少，且性能增益低于本文方法。作者认为，这是因为直接对键和值应用下采样操作会导致空间信息的丢失。对于 token 采样，其性能甚至比原始的 SwinIR 还要差，这大概率是因为丢弃一些 token 严重破坏了原始图像内容结构。

:::

Quantitative comparison

:::block-1

表4中展示了针对经典图像超分辨率（SR）方法的量化比较结果。为了公平比较，SRFormer 的参数数量和 MACs（Multiply-Accumulate操作）都低于SwinIR。结果清晰地显示出，SRFormer 在几乎所有的五个基准数据集上的所有尺度因子中都实现了最佳性能。因为在大窗口内计算自注意力可以聚合更大区域的信息，所以 SRFormer 在高分辨率测试集上的表现更好。特别是，对于使用DIV2K进行的2倍SR训练，SRFormer在Urban100数据集上实现了33.86dB的PSNR分数，这比SwinIR高了0.46dB，但使用的参数和计算更少。当引入SRFormer+ 的集成策略时，性能提升更大。

:::

Qualitative comparison

:::block-1

从上图第一个例子中，可以清楚地看到，SRFormer 能够恢复清晰且细节丰富的纹理和边缘。相比之下，其他模型恢复的纹理模糊或质量较低。对于第二个例子，SRFormer 是唯一能清晰恢复每个字母的模型。这些定性比较显示出，本文方法可以有效地从低分辨率图像中恢复出更好的高分辨率图像。

:::

总结

本文提出了一种新颖且高效的单图像超分辨率模型SRFormer，其核心是一个新的自注意力机制——PSA。PSA 能够在大窗口内有效地构建配对关联，而不会引入过多的计算成本，使得更多的 token 能参与自注意力计算。此外，SRFormer引入 ConvFFN 的模块，通过添加一个局部深度卷积分支来增强模型对高频信息的捕获能力。实验证明，SRFormer 在多个基准数据集上均优于现有技术，并且在参数数量和计算成本方面也表现出了优越性。

写在最后

如果有对深度学习在Low-level的应用研究感兴趣的童鞋，非常欢迎扫描屏幕下方二维码或者直接搜索微信号 cv_huber 添加小编好友，备注：学校/公司-研究方向-昵称，与更多小伙伴一起交流学习！