基于多级轴向加性网络的轻量级单图超分辨率

Lightweight Single Image Super-Resolution with Multi-Level Axial Additive Network

摘要: 信息技术发展日新月异，视觉信息的质量广受重视，图像超分辨率技术正因此经过了长久的迭代。但作为一个不适定问题，这项技术仍将是一个长久的难题。随着自注意力机制的出现及引入，传统卷积神经网络方法逐渐在性能上落后。然而，包含自注意力的方法通常计算成本高昂，或是只能为节约计算成本在性能上妥协。因此，本文提出了一种多级轴向加性网络，很好地平衡了性能与成本。具体来说，我们首先设计了一种多级轴向注意力模块，在注意力机制内实现了轴向窗口的模式。然后，我们提出了一种高效的加性注意力，使注意力计算免于矩阵乘法运算。同时，我们还构建了一个轻量级的超分辨率网络MLAAN。最后，我们在五个基准数据集上评估了所提出的MLAAN的效果。在与SOTA方法的对比中，MLAAN在参数量较少的前提下体现了优越的超分辨率性能。

Abstract: The importance of visual data has been increasingly emphasized due to the swift advancement of information technology nowadays. As an ill-posed problem, Single Image Super-Resolution continues to present an enduring challenge even after years of progression. Massive self-attention based methods proposed have shown performance exceeding traditional Convolutional Neural Networks based methods. However, methods including self-attention either suffer from large computational cost, or have to compromise on the weakened ability on capturing information thanks to modification on attention. We propose a Multi-Level Axial Additive Network with well-balanced trade-off in this work. Specifically, we first elaborate a Multi-Level Axial Attention Block enabling axial window patterns within attention. Then we present an effective additive attention that eliminates the need for expensive matrix multiplication operations in attention. We also construct a Feature Extraction Module base on shift-convolution to extract local features. We evaluate the efficacy of our proposed MLAAN on five benchmark datasets and show that it significantly enhances the super-resolution performance of the network. Our experimental results demonstrate state-of-the-art performance in lightweight SISR while using a low number of parameters.

1. 引言

单图像超分辨(SISR)的目标是从相应的低分辨率(LR)图像中推算出还原的高分辨率(HR)图像。在监控成像、自动驾驶和医疗成像等各种计算机视觉应用中，该技术都发挥着重要作用。作为一个不适定问题，即使经过多年的发展，SISR仍然是一个持久的挑战。人们提出了许多基于卷积神经网络(CNN)的方法 [1] [2] ，以直接学习LR和HR图像对之间的映射。例如，Dong [3] 等率先提出了基于CNN的初始模型SRCNN。SRCNN仅用三个卷积层，就成功超越了传统方法。之后，研究人员开始设计更多具有更复杂架构的网络，基于CNN的SISR方法取得了长足进步 [4] 。

尽管这些模型已经取得了显著的成就，但其庞大的计算成本却一直阻碍着它们的传播与实际应用投产。为了扩大SISR的应用范围，必须在性能和计算成本的平衡中求进。因此，学界与业界都把目光投向了轻量级的超分辨率方法，这些方法既有不俗的效果，又能将计算开支降到最低。DRRN [5] 利用循环网络结构在不增加参数的情况下增加了网络深度。然而，由于牺牲了性能和计算量，这种方法的相对实际效率并未提升。随着人们对SISR研究的钻研，高效的网络结构设计分化出了多种路径，包括神经结构搜索(NAS)、多尺度结构和通道分组策略。CARN [6] 通过级联局部和全局特征来恢复HR图像，速度和精度都很高。IMDN [7] 通过引入多重信息蒸馏模块来聚合和提炼特征。RFDN [8] 在IMDN的基础上进一步改进，加入了特征蒸馏连接(FDC)。BSRN [9] 引入了蓝图可分离卷积，以更小的模型实现了更好的重建效果。但是，卷积核通常用于提取局部特征，这意味着对图像中长距离依赖关系彻的彻底忽视。

在这方面，Transformer是一个出色的替代品。近年来，凭借它强大的全局建模能力，视觉Transformer (ViT)在许多视觉任务中崭露头角。近期有研究将Transformer引入了SISR领域。Swin IR [10] 利用移位窗口方案对长距离依赖关系进行建模，证明了ViT在SISR领域同样潜能巨大。ESRT [11] 将CNN和Transformer结合，构建了高效的轻量级模型。ELAN [12] 进一步简化了网络，避免了上游视觉任务庞大冗杂的网络。然而，现有的大多数基于Transformer的方法都使用密集注意力策略或移动窗口策略。本质上，感受野依然受到限制，需要堆叠大量模块才可能有效提取全局信息。为此，我们采用了轴向窗口注意力策略。由于相邻区域间的像素元素通常比远距离的像素元素相互作用更强，我们将注意力划分，分配到局部窗口和两个轴上。通过调节这三个部分的权重，实现了对局部与全局注意力的粒度调整。值得注意的是，在我们的架构中，局部窗口、水平轴向和垂直轴向的自注意力是以一种免于额外计算开支的并行模式计算的。

同时，我们还提出了一个用于SISR的多级轴向加性网络(MLAAN)，以基于ViT的架构组成。我们首先设计了了多级轴向注意力模块(MLAAB)，使提出的轴向窗口模式能够以轻量级的方式收集全局特征。然后，我们提出了一个高效的加性注意力模块(EAA)，在注意力计算中将繁杂的矩阵乘法运算实现取代。我们还构建了一个作用于局部的特征提取模块(FEM)，其中引入了移位卷积层和GELU激活函数。总之，本文的主要贡献可以概括为以下三个方面：

(1) 为SISR任务提出了一种结构简洁但功能强大的网络MLAAN。在轻量化的网络中巧妙地利用ViT的全局建模能力，在抑制计算成本的同时显著地提高了性能。并且通过标准数据集上的实验，定性和定量地论述了网络的优越性。

(2) 设计了一种多级轴向注意力模块MLAAB，在注意力机制内实现了轴向窗口模式，使全局依赖提取的轻量化成为可能。

(3) 设计了一种高效的加性注意力模块EAA，让注意力中繁杂的矩阵乘法运算得以被取代，进一步降低了计算复杂度。

2. 网络结构设计

以MLAAB为骨干单元，我们提出了MLAAN，网络结构如图1所示。该网络有着清晰简洁的结构，先后包括了浅层特征提取模块和深层特征提取模块，两部分的输出之间进行残差连接，合并后的输出再输入到上采样模块进行重建。

Figure 1. The overall architecture of the proposed MLAAN

图1. 所提出的MLAAN的主体网络架构

2.1. 注意力机制

在Transformer中，核心单元是多头自注意力(MHSA)。所谓自注意，是指将输入 $X\in {ℝ}^{N\times C}$ 分别线性投影为查询矩阵 $Q\in {ℝ}^{N_q\times C}$ ，键矩阵 $K\in {ℝ}^{N_{kv}\times C}$ ，以及值矩阵 $V\in {ℝ}^{N_{kv}\times C}$ ，注意力函数会将每个 $Q$ 矩阵转换为 $V$ 矩阵的求和权重。该权重是通过 $Q$ 和 $K$ 之间的归一化点积确定的。上述计算可表示为如下矩阵运算：

$\text{Attention}\left(Q,K,V\right)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{C}}\right)V.$

这里引入了可变标量 $\sqrt{C}$ 以避免权重集中和梯度消失，通常根据输入的维度确定。对于视觉Transformer来说， $X$ 是一个二维空间特征图 $N=H\times W$ ，其中 $H$ 和 $W$ 分别是特征图的高度和宽度。多头是指需要沿通道维度将输出划分为 $h$ 段。不同注意力头的投射权重不同。上述计算可表述为：

$\text{MHSA}\left(X\right)=\text{Concat}\left(hea{d}_0,hea{d}_1,\dots ,hea{d}_h\right)W^o,$

$hea{d}_i=\text{Attention}\left(X{W}_i^Q,X{W}_i^K,X{W}_i^V\right),$

其中 $hea{d}_i\in {ℝ}^{N\times \frac{C}{h}}$ 是第 $i^{th}$ 注意力头的输出。 $W_i^Q,W_i^K,W_i^V\in {ℝ}^{C\times \frac{C}{h}}$ 矩阵用于输入映射。通过对输入进行额外的线性变换得到的权重 $W^o\in {ℝ}^{C\times C}$ ，用于各输出进行合并求和。

MHSA中有 $N$ 个 $Q$ ，每个 $Q$ 将处理 $N$ 个键值对，因此复杂度为 $O\left(N^2\right)$ 。MHSA的高复杂度给视觉任务的输入分辨率带来了不少限制。

2.2. 多级轴向窗口注意力(MLAW)

Figure 2. The basic scheme of MLAW. The input scale is set as (9, 9) for better understanding. H, W, C stand for input height, input width, and input channel respectively.

图2. MLAW的流程结构。为了清晰可视化将输入尺寸设为9 × 9，H，W，C分别为输入的高、宽和通道

与MHSA中相同，输入特征 $X\in {ℝ}^{\left(H\times W\right)\times C}$ 将首先线性投射到 $K$ 个注意力头上，而此后每个注意力头在局部窗口或横轴或纵轴内进行自注意力运算，如图2。

2.2.1. 轴向窗口注意力

在提出的水平轴向注意力中， $X$ 被平均分割成不重叠的水平条状窗口 $\left[X^1,\cdots ,X^H\right]$ ，每个窗口包含 $1\times W$ 个元素。形式上，假设第 $k^{th}$ 注意力头的 $Q,K,V$ 的维度都是 $d_k$ ，那么 $k^{th}$ 注意力头的水平轴向注意力输出定义为：

$X=\left[X^1,X^2,\cdots ,X^H\right],$

$Y_k^i=\text{MSA}\left(X^i{W}_k^Q,X^i{W}_k^K,X^i{W}_k^V\right),$

${\text{H-MSA}}_k\left(X\right)=\left[Y_k^1,Y_k^2,\cdots ,Y_k^H\right],$

其中， $X^i\in {ℝ}^{\left(1\times W\right)\times C},i\in \left\{1,2,\cdots ,H\right\}$ ， $\text{MSA}$ 表示多头注意力。 $W_k^Q\in {ℝ}^{C\times d_k},W_k^K\in {ℝ}^{C\times d_k},W_k^V\in {ℝ}^{C\times d_k}$ 分别代表 $k^{th}$ 注意力头的 $Q,K,V$ 的投影矩阵， $d_k=C/K$ 。类似可得垂直轴向注意力，其对第 $k^{th}$ 注意力头的输出表示为 ${\text{V-MSA}}_k\left(X\right)$ 。

对于局部窗口注意力， $X$ 被平均分割成高度和宽度等于 $M$ 的非重叠局部窗口 $\left[X_m^1,\cdots ,X_m^N\right]$ ，每个窗口包含 $M\times M$ 元素。类似地，第 $k^{th}$ 注意力头的局部窗口注意力输出定义为：

$X_m=\left[X_m^1,X_m^2,\cdots ,X_m^N\right],$

$Y_k^i=\text{MSA}\left(X_m^i{W}_k^Q,X_m^i{W}_k^K,X_m^i{W}_k^V\right),$

${\text{W-MSA}}_k\left(X\right)=\left[Y_k^1,Y_k^2,\cdots ,Y_k^N\right],$

其中 $N=\left(H\times W\right)/\left(M\times M\right)$ 。

2.2.2. 粒度差异的构建

我们将 $K$ 个注意力头分为三部分，给两个轴向窗口各分配 $K/4$ 个注意力头，给局部窗口分配 $K/2$ 个注意力头。通过这种刻画粒度的方式，我们为局部和稀疏全局特征设置了权重。第一组注意力头执行水平轴向注意力，第二组注意力头执行垂直轴向注意力，第三组注意力头执行局部窗口注意力。并行计算后，输出将被重新合并：

${\text{head}}_k=\{\begin{array}{c}{\text{H-MSA}}_k\left(X\right),\\ {\text{V-MSA}}_k\left(X\right),\\ {\text{W-MSA}}_k\left(X\right),\end{array}\begin{array}{c}k=1,\cdots ,K/4,\\ k=K/4+1,\cdots ,K/2,\\ k=K/2+1,\cdots ,K,\end{array}$

$\text{MLAW}\left(X\right)\text{=Concat}\left({\text{head}}_1\text{,}\cdots {\text{,head}}_K\right)W^O,$

其中， $W^O\in {ℝ}^{C\times C}$ 为投影矩阵，用于融合各注意力头的输出。与分别逐步实现轴向和窗口注意力相比，这种并行机制的计算复杂度更低。并且，可以通过改变各组注意力头的数量来实现不同的粒度。

2.3. 高效的加性注意力

此前，加性注意力机制在NLP中已被应用，通过元素乘法取代了点积运算，利用成对标记间的交互来获得全局信息。它与惯例方法相同，仍用 $Q,K,V$ 的交互作用来编码输入序列上下文信息的相关性分数。在此之上，本文提出EAA只需通过加入线性投影层以聚焦于 $Q$ 、 $K$ 间有效交互，就足以学习表征之间的关系(见图1)。进一步简化了运算，在提高推理速度的条件下仍能稳健地提取特征。具体来说，输入 $X$ 通过两个映射矩阵 $W^Q,W^K$ 转换成 $Q$ 和 $K$ ，其中 $Q,K\in {ℝ}^{n\times d},W^Q,W^K\in {ℝ}^{d\times d}$ ， $n$ 是输入长度， $d$ 是输入向量的维数。接下来， $Q$ 矩阵可学习参数向量 $w_a\in {ℝ}^d$ 相乘，学习 $Q$ 的注意力权重，然后通过Softmax运算产生全局注意力查询向量 $\alpha \in {ℝ}^n$ 如下：

然后，根据学习到的注意力权重对 $Q$ 矩阵进行池化，最终得到一个单一的全局查询向量 $q\in {ℝ}^d$ ，如下所示：

接下来，通过元素乘积对全局查询向量 $q$ 和键矩阵 $K\in {ℝ}^{n\times d}$ 的交互进行编码，从而形成全局上下文( ${ℝ}^{n\times d}$ )。该矩阵与MHSA中的注意力矩阵相似，能捕捉全局的元素信息，在学习输入间的相关性上足够敏锐。不过，与MHSA相比，它的计算成本相对较低，复杂度与输入长度呈线性关系。受Transformer架构的启发，我们在 $Q$ 、 $K$ 交互中加入线性变换，从而学习输入的隐藏层表征。高效加性注意力的输出 $\hat{X}$ 可以描述为：

$\hat{X}=\hat{Q}+T\left(K\ast q\right),$

其中， $\hat{Q}$ 表示归一化查询矩阵， $T$ 表示线性变换。

2.4. 复杂度分析

对于MLAW，输入特征的大小为 $H\times W\times C$ ，窗口大小为 $\left(M,M\right)$ ，标准的全局自注意力的计算复杂度为：

$\text{Ω}\left(\text{Global}\right)=4HW{C}^2+2{\left(HW\right)}^{2}C,$

而并行计算的MLAW的计算复杂度为：

$\text{Ω}\left(\text{MLAW}\right)=4HW{C}^2+HWC\ast \left(\frac{1}{2}H+\frac{1}{2}W+M^2\right),$

与全局计算相比，这显然可以减轻计算和内存负担，因为 $2HW>>\left(\frac{1}{2}H+\frac{1}{2}W+M^2\right)$ 总是成立的。

至于EAA，其中学习 $Q$ 、 $K$ 矩阵的映射与交互的参数部分的计算复杂度为 $O\left(n\cdot d\right)$ ，而此后的逐元素乘积复杂度也同为 $O\left(n\cdot d\right)$ 。与传统自注意力的 $O\left(n^2\cdot d\right)$ 复杂度相比，EAA在节约计算上显出优越性，且在图像的大输入尺寸的背景下更为显著。

2.5. 多级轴向注意力模块

在上述内容的基础上，构建了FEM作为MLAAB的局部特征提取基础单元，我们在FEM中使用了两个移位卷积，并采用了GELU激活函数。这个基本卷积单元的结构如图1所示。至此，MLAAB可定义为：

$\widehat{X^l}=\text{FEM}\left(\text{LN}\left(X^i\right)\right)+X^i,X^l=\text{MLAW}\left(\text{LN}\left(\widehat{X^l}\right)\right)+\widehat{X^l},$

其中 $\widehat{X^l},X^l$ 分别表示第 $l$ 个MLAAB模块中的 $\text{FEM}$ 模块和 $\text{MLAW}$ 模块的输出特征。

2.6. 损失函数

在训练阶段，给定一个LR-HR训练集 ${\left\{I_i^{LR},I_i^{HR}\right\}}_{i=1}^N$ ，MLAAN的损失函数可以表示为：

其中， $F_{MLAAN}(\cdot )$ 表示我们提出的MLAAN， $\theta$ 表示MLAAN的参数集，N表示训练集中LR-HR图像对的数量。

3. 实验结果与讨论

3.1. 数据集和评估标准

在训练阶段，我们采用了DF2K数据集，包含DIV2K和Flickr2K数据集，共3450张图像。至于模型的评估，我们使用5个SISR标准数据集进行，分别为Set5、Set14、BSD100、Urban100和Manga109。为了呈现明晰的量化结果，我们采用了峰值信噪比(PSNR)和结构相似度(SSIM)作为数值指标。具体来说，SR图像的PSNR和SSIM是在YCbCr色彩空间的亮度Y通道上计算的。

3.2. 训练细节

为了得到符合SISR条件的训练数据，我们首先应用双三次插值法进行HR图像下采样。每个训练批次由16个随机抽取的 $48\times 48$ 像素块组成，每个输入随机进行旋转或水平翻转的数据增强。我们通过Adam优化器对模型进行了1000次训练，动量参数为0.9，损失函数为 $L1$ 。初始学习率为 $2\times {10}^{-4}$ ，每经200个epoch减半。该网络的训练使用了英伟达RTX3080Ti GPU和PyTorch框架。

3.3. 实验结果对比

本节，我们通过与最先进的轻量级SR模型的对比，展示了所们提出的模型的有效性。首先，我们直接呈现了量化结果(PSNR和SSIM)和计算成本，这是轻量级网络通常的关注重点。其次，我们从SR任务的实际目的出发，展示了视觉效果并进行了定性评估。

3.3.1. 定量比较

在表1和表2中，我们将我们的网络与其他先进的SISR模型进行了不同缩放尺度上的比较，其中包括VDSR、EDSR、SRMDNF、CARN、IMDN、ESRT、SwinIR-light。最好的结果均用粗体标出。可见，我们的MLAAN在各标准测试集上取得了可观的结果，在PSNR和SSIM上可与这些最先进的模型相媲美。此外，值得注意的是，我们提出的方法在实验的各尺度上基本优于众多类似的基于或融合了ViT的模型：ESRT和SwinIR-light。这主要得益于我们的模型在整合局部和长距离依赖方面取得了很好的平衡。因此，重建的图像包含丰富的结构细节，局部细节连贯，看起来比其他模型的图像更自然。

Table 1. Average PSNR/SSIM comparison with SISR models on ×3 scale

表1. 与其他SISR方法在×3缩放尺度上的PSNR/SSIM均值比较

Table 2. Average PSNR/SSIM comparison with SISR models on ×4 scale

表2. 与其他SISR方法在×4缩放尺度上的PSNR/SSIM均值比较

Table 3. Overall capacity comparison with lightweight SISR models

表3. 与其他轻量级SISR方法的综合性能比较

我们还进一步探究了所提出模型的计算成本，并与其他方法进行了比较。出于提高性能，并尽可能削减参数量、复杂度和推理速度等项的考量，我们的MLAAN实现了适当的平衡，如表3所示。从表中可以看出，首先MLAAN在性能上显著超出上述方法。同时作为一种基于ViT的模型，MLAAN在参数量方面与多数CNN和混合模型相当，性能则是有着巨大提升。每秒浮点运算次数(FLOPs)上，明显优于传统CNN方法，和较为新进的CNN和混合网络相比作为ViT模型仍有优化空间。最重要的是，MLAAN在参数量、FLOPs和性能方面全面优于同样基于ViT的SwinIR。综合而言，MLAAN比SOTA方法实现了更好的性能_开支平衡。

3.3.2. 定性比较

此外，我们还在图3中提供了我们的MLAAN与其他SISR方法的直观视觉对比。就图3中的上图而言，多数比较方法重建的SR图像都含有严重的伪影，图像上的线条模糊不清。相比之下，MLAAN重建的SR图像更加贴近原始图像，线条和色块更加清晰分明。不仅如此，其他方法存在过度平滑的问题，丢失了许多图像中的高频细节。基于混合的方法和其他基于ViT的方法可以缓解但不能完全克服这些现象，而MLAAN呈现出了最好的效果。在此之上MLAAN如实地重建了图像的结构，其他效果较清晰的图像则都在图像的右部出现了失真。对于下图来说，与上图的情况一样，其他方法要么存在过度平滑的问题，要么无法稳健地整合全局信息，导致视觉上的失调，图像出现了不同程度的扭曲。然而，我们的MLAAN一方面克服了过度平滑的问题，呈现出相对更清晰的边缘，另一方面还能重建出符合原始图像结构的图像。很明显，我们的MLAAN可以重建具有更精确纹理细节和边缘的高质量图像。这进一步证明了所提出的MLAAN的有效性。

Figure 3. The Qualitative comparison of ×4 image SR on the Urban100 dataset

图3. 在Urban100数据集×4缩放尺度上的定性比较

Figure 4. An example of LAM results of MLAAN

图4. MLAAN的LAM结果样本

不仅如此，我们还在图4中展示了MLAAN的局部归因图(LAM)结果。图中红色高亮部分代表了最左侧图像中，红框部分输入重建所建立的依赖范围。可见，网络对全局依赖建模的范围相当广泛。由此，LAM进一步证明了MLAAN能够在大范围内聚集像素信息，从而重构图像细节。

3.4. 消融实验

在本节中，我们进行了一系列研究来探究各模块的作用，以进一步展示模型的有效性。

3.4.1. 表征之间的相似性

首先，我们引入了中心核对齐(CKA)，以直观展示MLAW学习到的表征模式。具体来说，CKA结果显示了神经元之间的相似性。CKA得分越高，表示相似神经元越多，神经元间传递的信息就越少。如图5所示，MLAW的轴向窗口内部神经元相似性明显更低，这表明我们的MLAW的模式能有效提取这些范围内的信息。

Figure 5. An example of LAM results of MLAAN

图5. MLAAN的LAM结果样本

3.4.2. MLAW和EAA的有效性

为了验证所提出的MLAW与EAA的有效性，我们在实验中进行了两模块的消融对比，如表4。从表中我们可以清楚地看到，MLAW的引入使重建质量显著提升，参数的膨胀也远不及上文中对比的基于ViT的网络严重。这表明，MLAW可以提高模型的学习能力，从而在有效抑制计算成本升高的情况下提高模型性能。而EAA的引入，相比传统自注意力，显然使参数量和推理延时下降了，性能上则在Set5的SSIM上有微小的下降，但在PSNR上略有提高。该结果完美符合了轻量级网络的要求，在把握性能和开支平衡的基础上向前跃进。

Table 4. Evaluate the effectiveness of MLAW and EAA

表4. MLAW和EAA的有效性验证

3.4.3. MLAAB数量的影响

为了整体调整模型大小与性能间的平衡，我们考察了MLAAB数量对模型的影响，在五个数据集上比较了了PSNR和SSIM。从表5中可见，当MLAAB数量为12时，模型拟合程度明显不足；而与18的数量相比，模型在24个MLAAN块下依然有显著的性能提升，而参数量恰好符合轻量级的条件。当数量进一步增加时，模型性能的提高成本将难以承受。因此，我们将模型中的MLAAB数量设定为24。

Table 5. Evaluate the model capacity under different numbers of MLAAB

表5. MLAAB数量的影响评估

4. 结论

近年，图像超分辨率技术经过了快速的迭代。但作为一个不适定问题，这项技术仍将是一个长久的难题。随着自注意力机制的出现及引入，关键点逐渐落在了性能与计算成本的平衡上。本文中，我们提出了一种轻量级的多级轴向加性网络(MLAAN)。具体来说，我们首先设计了多级轴向注意力模块(MLAAB)，在注意力机制内实现了轴向窗口的模式以整合全局特征。然后，我们提出了一种高效的加性注意力(EAA)，注意力计算免于繁杂的矩阵乘法运算。同时，我们还构建了一个轻量级的超分辨率网络MLAAN。最后，我们在五个基准数据集上评估了所提出的MLAAN的效果。在与SOTA方法的对比中，MLAAN在参数量较少的前提下体现了优越的超分辨率性能。

NOTES

^*通讯作者。