论文地址：https://arxiv.org/pdf/2503.06896
代码地址：https://github.com/EquationWalker/CATANet

摘要

基于 Transformer 的方法在低级视觉任务（如图像超分辨率）中表现出了令人瞩目的效果。然而，其计算复杂度会随着空间分辨率的提高而呈平方级增长。一系列研究试图通过将低分辨率图像划分为局部窗口、轴向条纹或膨胀窗口来缓解这一问题。超分辨率通常利用图像的冗余性来进行重建，这种冗余不仅存在于局部区域，还存在于长距离区域。然而，这些方法将注意力计算限制在与内容无关的局部区域，直接限制了注意力捕捉长距离依赖的能力。为了解决这些问题，我们提出了一种轻量级的内容感知令牌聚合网络（CATANet）。具体而言，我们提出了一个高效的内容感知令牌聚合模块，用于聚合长距离内容相似的令牌，该模块在所有图像令牌之间共享令牌中心，并仅在训练阶段更新它们。然后，我们利用组内自注意力来实现长距离信息交互。此外，我们还设计了一种组间交叉注意力机制，以进一步加强全局信息的交互。实验结果表明，与最先进的基于聚类的方法 SPIN 相比，我们的方法表现更优，最大峰值信噪比提高了 0.33 分贝，并且推理速度提高了近一倍。

引言

单图像超分辨率（SISR）是计算机视觉和图像处理领域中的一项经典任务。其目标是从低分辨率（LR）图像中恢复出高分辨率（HR）图像。SISR 在众多领域都有广泛应用，例如医学成像、数字摄影以及降低流媒体传输的服务器成本。自卷积神经网络（CNN）在 SISR 领域的开创性工作以来，众多基于 CNN 的方法被提出，以挑战从低分辨率图像重建高分辨率图像的任务。由于卷积的局部机制限制了对全局依赖关系的捕捉，一些基于 CNN 的方法使用非常深且复杂的网络架构来扩大感受野并实现更好的性能。然而，这些方法不可避免地增加了计算资源，这限制了它们的适用性。
最近，由于 Transformer 在自然语言处理领域中的成功应用，它已被应用于多个高级计算机视觉任务，并取得了显著成果。基于 Transformer 的方法之所以取得关键的成功，是因为其自注意力机制，能够有效地捕捉长距离依赖关系。由于 Transformer 的强大潜力，它们在低级计算机视觉任务中也受到了关注，包括图像超分辨率。HAT 通过实验表明，利用更多的全局信息可以有效地提高重建图像的质量。然而，这些方法将较大的图像分割成较小的局部区域进行单独处理，以减轻全局自注意力的高复杂性。管这种策略可以提高基于 Transformer 的模型的效率并提供更精细的局部信息，但它仍然存在一些局限性。SwinIR 将图像划分为与内容无关的局部窗口，限制了在长距离上使用相似标记的使用，并导致了不理想的结果。轴向条纹注意力机制以十字形模式扩展了感受域，但这种方法并不考虑内容，可能会引入一些无关的干扰信息。
为了缓解这些问题，一些方法探索了基于聚类的解决方案。例如，SPIN 采用基于软 k 均值的标记算法进行聚类，将聚类中心用作注意力机制中查询与键之间的中介，从而促进长距离信息的传播。然而，SPIN 仍然存在两个主要的局限性：（1）聚类中心提供了图像标记的稀疏表示。然而，仅依靠这些中心进行长距离信息传播会导致粗略的近似，这不足以捕捉和利用详细的长距离依赖关系。（2）由于在推理过程中需要对聚类中心进行迭代处理，SPIN 的推理速度受到限制，这制约了轻量级模型的部署。ATD 引入了一个辅助字典，从训练数据中学习先验知识，并使用该字典对标记进行分类，从而实现更准确的标记分组。然而，ATD 会同时运用多种注意力机制来提升性能，这显著增加了计算负担，并使得 ATD 不再适合于轻量级场景，如图 1 所示。

Performance and model complexity comparison

为解决这些挑战，我们提出了一种新颖高效的基于内容感知的令牌聚合网络（CATANet），其核心组件为令牌聚合块。令牌聚合块主要由高效的基于内容的令牌聚合模块、组内自注意力机制和组间交叉注意力机制组成。与之前讨论的聚类方法不同，我们在内容感知令牌聚合（CATA）模块中设计了一种高效的令牌中心更新策略。它在所有图像令牌之间共享令牌中心，并仅在训练阶段更新它们，从而消除了更新令牌中心对模型推理速度的影响。与 SPIN 不同，后者使用聚类中心进行长距离信息传播，而我们的组内自注意力机制则利用 CATA 模块将内容相似的令牌高效地聚合在一起，形成内容感知区域。在每个组内进行注意力操作，从而实现令牌信息之间更精细的长距离交互。此外，我们还引入了组间交叉注意力机制，它在每个组与令牌中心之间应用交叉注意力，进一步增强了全局信息交互。
总之，我们的主要贡献如下：
• 我们提出了一种新颖的轻量级图像超分辨率网络——内容感知令牌聚合网络（CATANet）。我们的 CATANet 将令牌聚合与注意力机制相结合，以捕捉长距离依赖关系，同时确保高推理效率。
• 为了减轻令牌中心更新对推理速度的影响，我们设计了高效的内容感知令牌聚合（CATA）算法，该算法仅在训练阶段更新令牌中心。
• 我们提出了组内自注意力和组间交叉注意力，它们在令牌之间运行，有效地捕捉长距离和全局依赖关系，有效地减少无关信息的干扰，并实现与局部区域基于方法相当的计算复杂度。
• 我们进行了广泛的实验，以证明我们的方法超越了当前最先进的基于集群的轻量级超分辨率方法 SPIN，最大 PSNR 提升达 0.33 分贝，推理速度提高近一倍。

相关工作

深度网络在超分辨率处理中的应用。近年来，由于其强大的表示学习能力，深度神经网络已成为图像超分辨率处理的主流解决方案。能力。自 SRCNN 首次通过三层卷积神经网络（CNN）成功将 CNN 应用到超分辨率（SR）领域以来，大量基于 CNN 的方法利用更复杂或更高效的结构（如具有残差连接和 U 形架构）取得了最先进的性能。与 CNN 相比，注意力机制具有更好的长距离建模能力，因此注意力机制也被引入到超分辨率中，以提取长距离范围内的最重要特征。例如，RACN 利用通道注意力机制，而 CSFM 和 DAT 结合了空间和通道注意力机制。最近，一系列基于 Transformer 的方法被提出，并刷新了最先进的水平（SOTA），展示了 Transformer 强大的表示学习能力。例如，SwinIR 通过将整个图像划分为大小为 8×8 的小窗口，并在应用多头注意力机制时移动窗口，应用了 Swin Transformer 框架来进行超分辨率处理。尽管上述方法能够有效地提取有用特征，但它们通常需要大量的参数。
轻量级超分辨率方法。由于将网络应用于资源受限设备的需求十分迫切，轻量级超分辨率技术受到了广泛关注。例如，ESPCN 和 FSRCNN 利用后置上采样技术来降低计算成本。IDN 和 IMDN 采用信息提炼模块来提取有用信息，通过切片操作将输入特征分为两路流，其中一路流通过卷积进一步增强特征表示。随后，这两路特征被合并以获取更丰富的信息。对于基于转换器的轻量级超分辨率技术，大多数方法侧重于稀疏注意力机制，例如基于窗口的注意力机制和分组式注意力机制，以降低注意力计算的高计算复杂度。尽管这些方法有效地降低了计算复杂度，但它们的稀疏特性与内容无关，无法保证超分辨率重建的质量。
计算机视觉中的标记聚类。标记聚类在计算机视觉领域是一个已被深入研究的任务，而深度学习的最新进展在此方面取得了显著成果。常见的方法是使用传统的聚类算法，如 k 均值算法，将相似的标记聚集在一起。例如，BOAT 采用层次聚类将相似的标记聚集在一起，并在每个聚类内计算注意力。SPIN 利用基于软 k 均值的标记算法对标记进行聚合，将聚合后的标记用作查询和键之间的中介，以利用长距离信息。尽管这些标记聚类方法在计算机视觉中已显示出有希望的结果，但聚类过程的计算成本很高，这阻碍了标记聚类在轻量级超分辨率中的应用。

方法

1. 网络架构

CATANet, TAB and LRSA

所提出的 CATANet 的整体网络由三个模块组成：浅层特征提取、深层特征提取和图像重建，如图 2(a) 所示。浅层特征提取由一个 3×3 的卷积层完成，该层将原始输入空间中的图像映射到一个高维特征空间。设 I_LR 和 I_SR 分别代表低分辨率（LR）输入图像和超分辨率（SR）图像。我们首先通过以下方式获得浅层特征 X0 ：
X0 = F_S (I_LR),
其中，FS 表示浅层特征提取的功能。然后，通过 K 个顺序的残差组（RG）来提取更深层的特征。每个 RG 包括三个组成部分：令牌聚合块（TAB）、局部区域自注意力（LRSA）和 3×3 卷积（Conv）。
每个 RG 的输入特征首先通过一个 TAB 模块进行处理，以实现词元聚合，从而捕捉长距离的依赖关系。然后，我们使用一个 LRSA 来增强局部区域中词元之间的依赖关系。此外，在每个 RG 的末尾，我们使用一个 3×3 的卷积来进一步细化局部特征，并隐式地学习位置嵌入。具体而言，对于第 i 个 RG，整个过程可以表述为：
X_i = X_i−1 + F_Conv (F_LRSA (F_TAB (X_i−1))),
其中 Xi 表示第 i 个 RG 中的输出特征，F(·) 表示每个单独组件的函数。参考先前的研究，采用残差连接来稳定训练过程。
经过 K 个 RG 后，我们使用图像重建模块来获取全局残差信息，然后将其添加到 I_LR 的上采样图像中，以获得高分辨率图像 I_SR。
I_SR = F_Up(I_LR) + F_IR(X_K),
其中，F_Up 表示上采样操作的函数，而 F_IR 表示图像重建函数，该函数包括一个 3×3 的卷积和一个pixel shuffle操作。

2. 令牌聚合块（TAB）

在以往基于标记聚类的方法中，BOAT 提出了基于层次聚类的注意力机制，以实现精细的长距离信息交互。然而，层次聚类需要对聚类中心进行迭代更新，这会减慢模型的推理速度。SPIN 将聚类中心用作查询和键之间的中间代理，在自注意力机制中促进长距离信息的传播。不过，这种方法会导致信息传输的粗糙性，不可避免地引入无关信息。相比之下，我们提出了“组内自注意力”，以实现内容相似标记之间更精细的信息交互。此外，与 BOAT 类似，SPIN 也存在推理速度慢的问题。
为了更有效地解决这些问题，我们提出了“令牌聚合块”（TAB）这一概念。如图 2（b）所示，TAB 主要由四个部分组成：内容感知令牌聚合模块、组内自注意力模块、组间交叉注意力模块以及一个 1×1 卷积层。通过利用这四个模块，我们能够高效地实现细粒度的长距离信息交互，同时还能保持与局部区域注意力类似的计算复杂度。

2.1. 内容感知令牌聚合（CATA）

Figure 3

Visualization of token grouping results of TAB

Parameter, Multi-Adds, and Running Time comparison

2.2. 组内自注意力机制（IASA）

给定子群 S = {Sj}，我们将它投影为矩阵 {Qj}、{Kj} 和 {Vj}，具体如下：
Qj 、Kj 、Vj = SjWQ 、SjWK 、SjWV，
其中 WQ、WK 和 WV ∈ R^(d×d) 为权重矩阵。为了提高并行效率，我们将这些组 G 分成子群 S。然而，这可能会导致内容相似的标记被拆分成相邻的子群。为了缓解这个问题，我们允许每个子群的 Qj 与两个连续子群的 Kj 进行交互，如图 3(c) 所示。此外，如表 1 所示，与相邻子群进行交互有效地提高了性能，而不会增加计算开销。该过程可以表述为：
公式详见原论文
这里，MSA(·) 表示多头自注意力运算，而 PushBack 运算会将每个标记放回其在特征图中的原始位置，从而形成 Xo1 。

2.3. 群体间交叉注意力（IRCA）

如算法 1 所示，在训练阶段，令牌中心 C 会聚集内容相似的令牌，从而总结出全局先验信息。为了利用这一全局先验信息，我们对每个子组 Sj ∈ R gs×d 和 C ∈ RM×d 进行交叉注意力计算，具体如下：
公式详见原论文
其中，Wq、Wk 和 Wv 分别为位于 R^d×d 空间中的权重矩阵。我们令 M ≪ N 以保持较低的计算成本。最后，通过卷积将式（1）中的 Xo2 和 Xo1 融合起来，得到 Xo:
Xo = F_Conv (Xo1 + Xo2)

3. 局部区域自注意力机制（LRSA）

如图 2(c) 所示，LASA 采用了 HPINet [28] 的方法，使用重叠的块来增强特征交互。LRSA 模块负责学习更精细的局部细节。给定输入 Xo ∈ R N×d，此过程可以表示为：
X_out = MSA(X_oW_Q, X_oW_K, X_o*W_V )，
其中 Xout ∈ R N×d。WQ、WK 和 WV 是共享在各个块之间的权重矩阵。
ConvFNN。与 Transformer 层类似，在 TAB 和 LRSA 之后，ConvFNN 被用于沿着通道维度进行特征交互，如图 2 所示。此外，在 TAB、LRSA 和 ConvFFN 之前添加了层归一化，并且在两个模块之后都添加了残差连接。

实验

训练集：DIV2K
测试集：Set5, Set14, B100, Urban100 and Manga109
Baseline：

1. CNN-based model: CRAN, LatticeNet, IMDN, RFDN, AWSRAN-M and OSFFNe；
2. Transformer-based model: ESRT, ELAN light, A-CubeNet, SwinIR and SPIN；
   评价指标：PSNR，MS-SSIM
   实验结果：详见原论文。

结论

在本文中，我们提出了一种名为 CATANet 的新型轻量级图像超分辨率网络，该网络利用令牌中心将内容相似的令牌聚合到内容感知区域中。具体而言，CATANet 的核心组件是令牌聚合块。令牌聚合块主要由内容感知令牌聚合（CATA）、组内自注意力和组间交叉注意力组成。CATA 模块有效地将相似的令牌聚合在一起，形成内容感知区域。组内自注意力负责在内容感知区域内的精细长距离信息交互。同时，我们引入了组间交叉注意力，它通过每个组与令牌中心之间的交叉注意力来进一步增强全局信息交互。我们在各种基准数据集上进行了广泛的实验结果展示，我们的方法在轻量级图像超分辨率中实现了卓越的性能，同时保持了较高的推理效率。