Omni Aggregation Networks for Lightweight Image Super-Resolution 【用于轻量级图像超分辨率的全聚合网络】

论文地址:https://openaccess.thecvf.com/content/CVPR2023/html/Wang_Omni_Aggregation_Networks_for_Lightweight_Image_Super-Resolution_CVPR_2023_paper.html
代码实现：https://github.com/francis0625/omni-sr

Abstract

​ 虽然轻量级 ViT 框架在图像超分辨率方面取得了巨大进步，但其单维自注意力建模以及同质聚合方案限制了其有效感受野 (ERF) 以包含来自空间和通道维度的更全面的交互。 为了解决这些缺点，这项工作在新的 Omni-SR 架构下提出了两个增强组件。 首先，基于密集交互原理提出了 Omni Self-Attention (OSA) 块，它可以同时从空间和通道维度对像素交互进行建模，挖掘跨全轴（即空间和通道）的潜在相关性。结合主流的窗口分区策略，OSA 可以通过令人信服的计算预算实现卓越的性能。 其次，提出了一种多尺度交互方案来减轻浅层模型中的次优 ERF（即过早饱和），这有助于局部传播和中观/全局尺度相互作用，呈现全尺度聚合构建块。 大量实验表明，Omni-SR 在轻量级超分辨率基准测试中实现了创纪录的性能（例如，26.95dB@Urban100 ×4，仅使用 792K 参数）。

1. Introduction

​ 图像超分辨率（SR）是一个长期存在的低级别问题，旨在从退化的低分辨率（LR）输入中恢复高分辨率（HR）图像。最近，出现了基于视觉转换器[14，51]（即基于ViT）的SR框架[5，30]，与以前占主导地位的卷积神经网络（CNNs）[66]相比，显示出显著的性能提升。然而，大多数尝试[30]都致力于改进基于大规模ViT的模型，而轻量级ViT（通常，小于1M的参数）的开发仍然充满了困难。本文的重点是提高基于ViT的轻量级框架的恢复性能。

​ 两个困难阻碍了基于轻量级ViT的模型的发展：1）一维聚合算子（即仅空间[30]或仅通道[59]）限制了自注意算子的全部潜力。当代的自我关注通常通过计算空间方向（即宽度和高度）的互协方差来实现像素之间的交互，并以通道分离的方式交换上下文信息。这种交互方案忽略了信道信息的显式使用。然而，最近的证据[59]和我们的实践表明，通道维度的自我注意（即计算上比空间自我注意更紧凑）在低级别任务中也至关重要。2） 同构聚合方案（即单算子的简单分层堆叠，如卷积、自注意）忽略了多尺度的丰富纹理模式，这是SR任务中迫切需要的。具体而言，单个操作员只对一个尺度的信息敏感[6，12]，例如，自我注意对长期信息敏感，而很少关注局部信息。此外，同质算子的叠加被证明是低效的，并且存在相互作用范围的过早饱和[8]，这反映为次优有效感受野。在轻量级模型中，由于轻量级模型无法堆叠足够的层，因此上述问题更加严重。

​ 为了解决上述问题并追求更高的性能，本工作提出了一种新的全维特征聚合方案，称为全方位自注意（OSA），同时利用空间和通道轴信息（即将交互扩展到三维空间），提供更高阶的感受野信息，如图1所示。与基于标量（一组重要系数）的信道交互[19]不同，OSA通过级联计算空间/信道维度之间的互协方差矩阵，实现了全面的信息传播和交互。所提出的OSA模块可以插入任何主流的自我注意变体（例如，Swin[34]、Halo[50]）中，这提供了更精细的重要编码粒度（与普通通道注意[19]相比），实现了上下文聚合能力的显著提高。此外，提出了一种多尺度分层聚合块，称为全尺度聚合组（Omni-scale aggregation Group，简称OSAG），以实现对不同尺度纹理图案的定制编码。具体而言，OSAG构建了三个级联聚合器：局部卷积（用于局部细节）、中尺度自注意（专注于中尺度模式处理）和全局自注意（追求全局上下文理解），提供了全尺度（即同时在局部/中尺度/全局尺度）的特征提取能力。与同质化特征提取方案[27，30]相比，我们的OSAG能够以更高的信息熵挖掘出更丰富的信息产生特征。结合上述两种设计，我们为轻量级超分辨率建立了一个新的基于ViT的框架，称为Omni-SR，它表现出优异的恢复性能，并覆盖了更大的交互范围，同时保持了有吸引力的模型尺寸，即792K。

​ 我们在主流开源图像超分辨率数据集上对所提出的框架进行了广泛的实验，并进行了定性和定量评估。研究表明，我们的框架在轻量级模型规模上实现了最先进的性能（例如，Urban100×4:26.95dB，Manga109×4: 31.50dB）。更重要的是，与现有的基于ViT的超级解决方案框架相比，我们的帧架显示出优越的优化特性（例如，收敛速度、更平滑的损失景观），这赋予了我们的模型更好的鲁棒性。

2. Related Works

​ 图像超分辨率。CNNs网络在图像SR任务中取得了显著的成功。SRCNN[13]是第一个将CNNs网络引入SR领域的工作。许多方法[25，48，66]采用跳过连接来加快网络收敛并提高重建质量。信道注意力[66]也被提出来增强SR模型的表示能力。为了在有限的计算资源下获得更好的重建质量，几种方法[23，38，42，47]探索了轻量级的架构设计。DRCN[26]利用递归运算来减少参数的数量。DRRN[47]在DRCN的基础上引入了全局和局部残差学习，以加速训练并提高细节质量。CARN[1]在残差网络上采用级联机制。IMDN[22]提出了一种信息多静态块，以归档更好的时间性能。另一个研究方向是利用模型压缩技术，例如知识提取[15，17，65]和神经架构搜索[11]）来降低计算成本。最近，出现了一系列具有卓越性能的基于transformer的SR模型[5，8，30，37]。Chenet等人[5]使用transformer架构为低级计算机视觉任务开发了一个预先训练的模型。在Swin-transformer[34]的基础上，SwinIR[30]提出了一个三阶段框架，刷新了SR任务的最新技术。最近，一些工作[5，29]探索了ImageNet预训练策略，以进一步提高SR性能。

​ 轻型视觉transformer。由于将网络应用于资源受限设备的迫切需求，轻量级视觉转换器[14，51]引起了广泛关注。已经进行了许多尝试[7，9，10，37，41，43，57，62]来开发具有类似性能的轻量级ViT。一系列方法专注于将卷积与变换器相结合，以学习局部和全局表示。例如，L VT[57]在自我注意中引入了卷积，以丰富低级特征。MobileViT[41]用变换器层代替卷积中的矩阵乘法来学习全局表示。类似地，EdgeViTs[43]为全空间交互采用了信息交换瓶颈。与将卷积解释为视觉转换器不同，LightViT[21]提出了聚合的自我注意力，以更好地聚合信息。 在这项工作中，我们采用 ViT 架构来实现轻量级和准确的 SR。

3. Methodology

3.1超分辨中的注意力机制

SR 中广泛采用两种注意范式来帮助分析和聚合综合模式。

​ 空间注意力。空间注意力可以看作是一个各向异性的选择过程。主要应用空间自我注意[37，51]和空间门[10，58]。如图2所示，空间自注意计算沿空间维度的互协方差，空间门生成通道分离的掩码。它们都不能在信道之间传输信息。

​ 通道注意力。有两类通道注意力，即基于标量的[19]和基于协方差的[59]，用于执行信道重新校准或信道之间的传输模式。如图2所示，前者预测一组重要性标量来加权不同的信道，而后者计算一个互协方差矩阵来同时实现信道重新加权和信息传输。与空间注意力相比，通道注意力各向同性地处理空间维度，因此，复杂性显著降低，这也损害了聚合的准确性。

​ 几次尝试[44，55]已经证明，空间注意力和通道注意力都有利于SR任务，并且它们的特征是互补的，因此以计算紧凑的方式将它们集成在一起将在表达能力方面带来显著的好处。

3.2.全方位自我关注块

​ 为了挖掘隐藏在潜在变量中的所有相关性，我们提出了一种新的自我注意范式，称为全自我注意（OSA）块。与只沉溺于一维处理的现有自我注意范式（例如，空间自我注意[5，37，51]）不同，OSA同时建立了空间和通道上下文。所获得的二维关系是非常必要和有益的，尤其是对于轻量级模型。一方面，随着网络的加深，重要信息分散到不同的渠道[19]，及时处理至关重要。另一方面，尽管空间自注意在计算协方差时利用了信道维度，但它不会在信道之间传输信息（参见第3.1节）。考虑到上述条件，我们的OSA旨在以紧凑的方式传输空间和维度信息。

​ 所提出的OSA通过顺序矩阵运算和旋转来计算与空间和信道方向相对应的分数矩阵，如图3所示。具体来说，假设X∈RHW×C表示输入特征，其中H和W是输入的宽度和高度，C是通道数。首先，通过线性投影将X嵌入到查询矩阵、键矩阵和值矩阵Qs、Ks、Vs∈RHW×C中。我们计算了查询和关键字的产生，以获得大小为RHW×HW的空间注意力图。然后，我们进行空间注意以获得中间聚合结果。请注意，窗口策略通常用于显著减少资源开销。下一阶段，我们旋转输入查询和关键矩阵，得到转置查询和关键阵Qc，Kc∈RC×HW，还旋转值阵，得到值阵VC∈RC X HW，用于后续的通道自注意。所获得的大小为RC×C的通道关注图对通道关系进行建模。最后，我们通过通道注意力输出Yc的反向旋转得到最终的聚合YOSA。整个OSA流程如下所示：

​ 其中Wq、Wk、Wv分别表示查询、关键字和值的线性投影矩阵。Q′，K′，V′是通道自注意的输入嵌入矩阵，它们是从前空间自注意嵌入的或直接从Qs，Ks，Vs复制的。R（·）表示绕空间轴的旋转操作，R−1（·）是逆旋转。为了简单起见，省略了一些归一化因子。特别是，这种设计显示出令人信服的特性，可以集成两个矩阵运算（即空间/通道矩阵运算）的元素结果，从而实现全轴交互。请注意，我们提出的OSA范式可以取代Swin[30，34]注意力块，以更少的参数获得更高的性能。得益于信道自注意的较小注意图大小，与Swin中的级联移位窗口自注意方案相比，所提出的OSA的计算密集度较低。

​ 与其他混合注意力范式的讨论。与之前的混合通道和空间注意力作品（如CBAM[55]和BAM[44]）相比，它们基于标量的注意力权重只反映了相对的重要性，没有进一步的像素间信息交换，导致关系建模能力有限。最近的几项工作[8]也将通道注意力与空间自我注意力结合在一起，但这些尝试仅采用标量权重进行通道重新校准，而我们的OSA范式使通道交互能够挖掘全轴中的潜在相关性。不同注意力范式的表现比较可以在第4.4节中找到。

3.3. 全方位聚合组

​ 如何利用所提出的OSA范式来构建高性能、紧凑的网络是另一个关键课题。尽管基于窗口的自注意的分层堆叠（例如，swin[30，34]）已经成为主流，但各种工作已经发现，基于窗口的范式对于大范围的交互是非常低效的，尤其是对于浅层网络。值得指出的是，大范围交互可以提供令人愉悦的有效感受野，这对于提高图像恢复性能至关重要[37]。不幸的是，直接的全球交互会阻碍资源的使用，并降低本地聚合能力。考虑到这些点，我们提出了一个 Omni-Scale Aggregation Group（即简称 OSAG）来追求具有低计算复杂度的渐进式感受野特征聚合。 如图3所示，OSAG主要由三个阶段组成：local、meso和global aggregations。 具体来说，引入通道注意力 [19] 增强倒置瓶颈 [18] 来以有限的开销完成局部模式过程。 基于提出的 OSA 范式，我们派生了两个实例（即 Meso-OSA 和 Global-OSA），负责交互和聚合 meso 和全局信息。请注意，所提出的 omni self-attention 范式可用于不同的目的。 Meso-OSA 在一组非重叠补丁中执行注意力，这限制了 Meso-OSA 仅关注中尺度模式理解。 Global-OSA 以一种复杂的方式在整个特征中稀疏地采样数据点，赋予 GlobalOSA 以令人信服的成本实现全局交互的能力。

​ Meso-OSA和GlobalOSA之间的唯一区别是窗口分区策略，如图4所示。为了实现中尺度相互作用，中尺度OSA将输入特征X拆分为大小为P×P的非重叠块。

3.4.网络架构

​ 整体结构。基于Omni Self Attention范式和Omni-Scale Aggregation Group，我们进一步开发了一个轻量级的Omni-SR框架，以实现高性能的图像超分辨率。如图3所示，Omni-SR由三部分组成，即浅层特征提取、深层特征提取和图像重建。具体来说，给定LR输入ILR∈RH×W×Cin，我们首先使用3×3卷积HSF来提取浅层特征X0∈RH x W×C作为

其中Cin和C表示输入和浅特征的通道编号。卷积层提供了一种简单的方法，可以将图像空间的输入转换为高维特征空间。然后，我们以级联的方式使用K个堆叠的全尺度聚集群（OSAG）和一个3×3卷积层HCONV来提取深层特征FDF。这样的过程可以表示为

其中HOSAGi表示第i个OSAG，X1，X2。XK表示中间功能。在[30]之后，我们还在特征提取结束时使用卷积层，以获得更好的特征聚合。最后，我们通过将浅特征和深特征聚合为

其中，HRec（·）表示重建模块。详细地说，PixelShuffle[46]用于对融合的特征进行上采样。

​ 全规模聚合集团（OSAG）。如图3所示，每个OSAG包含一个局部卷积块（LCB）、一个meso-OSA块、一个全局OSA块和一个ESA块[27，33]。整个过程可以公式化为

其中，Xi−1和Xi表示第i个OSAG的输入和输出特性。在对卷积层进行映射后，我们插入了用于基于窗口的自我注意的Meso-OSB和用于扩大感受野以获得更好的信息聚合的Global-OSB。在OSAG结束时，我们在[27，66]之后保留卷积层和ESA块。

​ 具体而言，LCB被实现为逐点和深度卷积的堆栈，它们之间有CA模块[24]，以自适应地重新加权信道特征。该块旨在聚合本地上下文信息，并提高网络的可训练性[56]。然后遵循两种类型的OSA块（即，Meso-OSA块和Global OSA块），以获得来自不同区域的相互作用。基于不同的窗口划分策略，Meso-OSA块寻求内部块的相互作用，而Global OSA块旨在实现全局混合。OSA块遵循具有前馈网络（FFN）和LayerNorm[2]的典型变压器设计，唯一的区别是用我们提出的OSA算子取代了原点自注意操作。对于FFN，我们采用Restormer[59]提出的GDFN。将这些个体无缝地结合在一起，所设计的OSAG能够在特征图中的任何一对标记之间进行信息传播。我们使用[27，33]中提出的ESA模块来进一步细化融合特征。

​ 优化目标。根据先前的工作[30、31，53，67]，我们通过最小化模型预测和HR标签IHR之间的标准L1损失来训练模型，如下所示

4.Experiments

4.1实验设置

​ 数据集和度量。继先前的工作[30，31，38，49，66]之后，DIV2K[49]和Flickr2K[49]被用作训练数据集。为了进行公平的比较，我们使用了两种训练协议，即仅使用DIV2K进行训练和使用DF2K进行训练（DIV2K+Flickr2K）。请注意，使用DF2K训练的模型标有小†。为了进行测试，我们采用了五个标准的基准数据集：Set5[4]、Set14[60]、B100[39]、Urban100[20]和Manga109[40]。采用PSNR和SSIM[54]来评估变换的YCbCr空间的Y信道上的SR性能。

​ 实施细节。在训练过程中，我们通过随机水平翻转和90/270度旋转来增加数据。LR图像是通过对HR图像进行双三次下采样[63]生成的。OSAG编号设置为5，并且整个网络的信道编号设置为64。中尺度OSAB和全局OSAB的注意力头部数量和窗口大小都设置为4和8。我们使用AdamW[36]优化器来训练批量大小为64的模型，用于800K迭代。初始学习率设置为5×10−4，每200k次迭代减半。在每个训练批次中，我们随机裁剪大小为64×64的LR补丁作为输入。我们的方法是用PyTorch[45]实现的，所有实验都是在一个NVIDIA V100 GPU上进行的。注意，没有使用其他数据增强（例如，Mixup[61]、RGB通道混洗）或训练技能（例如，预训练[29]、余弦学习计划[35]）。需要指出的是，我们在消融研究中通过调整来保持模型参数的一致性。

4.2.与SOTA SR方法的比较

​ 为了评估Omni SR的有效性，我们将我们的模型与几种先进的轻量级SR方法在2/3/4的比例因子下进行了比较。特别是，介绍了以前的工作，VDSR[25]、CARN[1]、IMDN[22]、EDSR[31]、RFDN[32]、MemNet[48]、MAFFSRN[42]、LatticeNet[38]、RLFN[27]、ESRT[37]和SwinIR[30]进行比较。

​ 定量结果。在表1中，不同轻量级方法在五个基准数据集上进行了定量比较。在类似的模型大小下，我们的Omni-SR的性能超过了现有的方法，在所有基准测试中都有显著的优势。特别是，与具有类似参数的其他变压器架构（如SwinIR[30]和ESRT[37]）相比，所提出的Omni-SR获得了最佳性能。结果展示了OSA引入的全轴（即空间+信道）交互可以有效地提高模型的上下文聚合能力，这保证了卓越的SR性能。结合大型训练数据集DF2K，可以进一步提高性能，尤其是在Urban100上。我们认为，这种现象可以归因于Urban100中的图像具有许多相似的斑块，而OSAG引入的长期关系可以为细节恢复带来巨大的好处。更重要的是，在参数相似的情况下，我们的模型降低了28%的计算复杂度（Omni SR:36G FLOPs vs.SwinIR:50G FLOPs@1280×720），显示了它的有效性和效率。

​ 视觉对比。在图6中，我们还提供了不同轻量级SR方法在×4尺度上的视觉比较。我们可以观察到，Omni-SR构建的HR图像包含更多细粒度的细节，而其他方法在复杂区域生成模糊的边缘或伪影。例如，在第一排，我们的模型能够愉快地恢复墙壁的详细纹理，而所有其他方法都无法恢复。可视化结果也验证了所提出的OSA范式的有效性，该范式可以执行全轴像素交互建模，从而获得更强大的重建能力。

​ 在模型尺寸和性能之间进行权衡。在实验中，我们将OSAG的数量设置为5，使模型大小在800K左右，以便与其他方法进行公平比较。我们还通过减少OSAG数K来探索参数大小较小的模型性能。如图5（a）所示，与具有K=1时，增加OSAG的数量会带来稳定的性能改进。在图5（b）中，我们给出了PSNR与不同方法参数的对比。可以发现，Omni SR在各种设置下都能达到最佳效果，显示了其有效性和可扩展性。

4.3.全方位自我注意分析

​ 在本节中，我们说明了OSA的优化特征，并进一步揭示了其潜在机制。自注意是一种低偏差算子，这使得其优化变得困难，并且需要更多的训练时间。为此，我们引入了额外的通道交互来缓解它。在图7（a）中，我们展示了DIV2K训练集上不同自我注意范式的损失曲线，包括空间自我注意、通道自我注意和所提出的全方位自我注意。我们可以看到，我们的OSA呈现出明显优越的收敛速度。更重要的是，在最后阶段的表现也明显领先于他们。上述现象清楚地表明，我们的OSA具有优越的良好优化特性。此外，我们深入研究了为什么渠道交互会导致这些改进。我们计算由上述三个计算基元组成的网络的隐藏层特征的归一化熵[52]。我们在图7（c）中说明了熵的结果。如图所示，在所有传出层中，我们的OSA编码特征显示出更高的熵，这表明我们的OSA编码的信息更丰富。更多的信息可能来自不同的尺度，这些信息可以帮助操作员更快地重建确切的细节。我们推测这就是为什么我们的OSA显示出更好的优化性能的潜在原因。此外，继之前的工作[8，16]之后，我们还采用了LAM分析。DI[16]度量可以测量模型的最远交互距离。从图8中，我们可以观察到Omni-SR通常比其他方法具有最高的最大扩散指数，这表明我们的OSA范式可以有效地捕捉长程相互作用。

4.4.消融研究

​ 全方位自我注意的效果。我们框架的核心思想是通过渠道关系扩展香草自我关注，以建立全轴像素交互。基于Omni-SR框架，我们设计了几个变体模型，其SR结果如表2所示。我们首先简单地去除信道分量以形成仅限空间的变体（Omni-SRsp），与完整模型相比，其性能降低了0.13dB。如此显著的退化证明了通道相互作用的重要性。请注意，Omni SRsp仍然优于SwinIR0.04dB@Urban100×4，这得益于网格窗口分区引入的全局交互。类似地，我们去除了空间自注意分量，以导出通道自注意变体OmniSRca，这种修改也会导致不希望的性能下降。此外，我们使用最广泛采用的信道和空间注意力配置（即SE[19]和CBAM[55]）作为信道和空间聚合的替代算子。与完整型号相比，这两种替换（OmniSRSE、Omni SRCBAM）都损害了PNSR的性能。上述结果表明，特定的交互模式（例如，基于标量的、基于协方差的）同样重要，并且我们基于协方差矩阵的信道交互显示出很大的优势。

​ 全规模聚合组的影响。在OmniSR中，我们提出了一种局部-中全局交互方案（即OSAG）来追求渐进的特征聚合。为了研究其有效性，我们基于Omni-SR框架设计了三种不同的交互方案：分离方案、混合方案和我们完全设计的Omni方案（即我们提出的OSAG），消融研究结果如图7（a）所示。在图中，我们使用不同的词（例如，“Local”、“Meso+Global”）来表示特定的方案，例如，“Local-”表示使用Local-Conv块来代替Meso-OSA和GlobalOSA；“局部+全局”表示用级联的局部Conv和全局OSA取代原来级联的Meso-OSA和全局OSA。我们可以观察到，单个交互方案（例如“本地”）的性能最差。有趣的是，“全局”方案不如“Meso”方案，因为其全局自注意的优化性能较差[3，34，50]。一旦两个交互操作符结合在一起，性能就会稳步提高。其中，“Meso+Global”设置的性能排名第二。此外，将所有三种交互方案组合在一起，我们获得了性能最好的方案，即“Omni”。从上述实验中，我们可以推断，通过引入各种尺度的相互作用可以获得明显的性能增益，这也说明了我们提出的OSAG的可行性和有效性。

5. Conclusion

​ 在这项工作中，我们提出了Omni-SR，这是一个用于图像SR的轻量级框架。此外，我们提出了一种全尺度聚合方案，以低计算复杂度有效地扩大感受野，该方案以渐进的分层方式编码上下文关系。在公共基准数据集上进行的大量实验和全面的分析研究验证了其显著的SR性能。