Transformer模型在图像复原领域取得令人瞩目的成绩。不过考虑到Transformer过高的计算复杂度，部分工作尝试使用局部窗口注意力来限制自注意力的作用范围。但是这些方法缺少窗口之间的直接交互，限制了模型长程依赖的建模。针对这个问题，本文提出交叉聚合Transformer（Cross Aggregation Transformer，CAT），使用新颖的Rwin-SA与Axial-Shift操作，在不增加计算复杂度的情况下，提高窗口间的信息聚合。同时提出LCM模块，实现全局与局部信息的耦合。我们将CAT模型应用于3个经典的图像复原任务：图像超分辨率、压缩伪影去除、真实图像去噪，均取得最佳的性能。
论文连接：
https://arxiv.org/abs/2211.13654
代码连接：
https://github.com/zhengchen1999/CAT
一、研究问题
图像复原是一个长期的低级视觉任务，旨在从低质量的图片（LQ），恢复出对应的高质量图片（HQ）。根据退化类型的不同，可以分为多个子任务，例如图像超分辨率，图像去噪，图像去模糊等。目前主流的复原算分是通过深度神经网络模型进行端到端的图像重构。
二、研究背景
2.1 图像复原
基于卷积神经网络（CNN）的模型，在图像复原任务中取得出色的表现，并取代传统方法（例如基于滤波），成为该领域的主流。大量的模型被提出用于改进复原效果 [1, 2, 3, 4]，这些方法借助残差学习、密集连接等技术，来尽可能的提高CNN模型的深度，从而更多的高频信息。同时空间与通道注意力机制也被引入到模型中，使网络能够更加关注特征图中的特定信息。
2.2 视觉Transformer
Transformer模型首先在自然语言处理（NLP）中被提出，其核心为自注意力机制（self-attention，SA）。与卷积相比，自注意力具有动态权重，并且能够建模长距离依赖关系。考虑到Transformer在NLP中取得的性能，部分工作将Transformer引入到视觉领域 [6, 7]，并且在多个高级视觉任务中取得最佳的性能。这些工作主要关注于降低Transformer的计算复杂度并提升模型性能。例如使用窗口注意力机制并通过Shift操作增加窗口交互，或者提出交叉形式的窗口。
同时，在低级视觉任务，如图像复原领域，目前也提出部分工作，尝试用Transformer替换传统的CNN模型。SwinIR [8] 使用窗口划分图像，并在每个图像上分别执行注意力操作，同时也采用Shift操作增加窗口交互。Restormer [9] 则在特征图的通道维度计算交叉协方差。这些方法都能实现与图像大小成线性的计算复杂度，并在一定程度上保留Transformer长距离建模能力，从而优于基于CNN的模型。
三、方法介绍
在本节中，我们首先介绍本文模型的整体架构。然后是核心模块的三个创新点：矩形窗口注意力机制（rectangle- window self-attention，Rwin-SA）、轴向位移操作（axial-shift）以及局部补充模块（locality complementary module，LCM）。
3.1 模型架构
3.2 交叉聚合Transformer块（CATB）
CATB是本文的核心模型，使用一种新颖的矩形窗口注意力机制（Rwin-SA）和轴向位移操作（Axial-Shift），以及通过局部补充模块（LCM）耦合全局与局部特征。
矩形窗口注意力（Rwin-SA） 如图片2 (a)，本文提出的Rwin-SA，同样适用窗口对特征图进行划分，这里我们用蓝色与绿色区分相邻窗口。但是与之前的工作不同，我们使用的是矩形窗口（长为sh，宽为sw ）而不是方形窗口，并考虑到矩形窗口长-宽特性，我们将其分为水平矩形与垂直矩形两种形式。通过将特征图沿通道维度拆分为两组，并分别应用两种矩形窗口进行划分，执行注意力操作后，再沿着通道维度进行合并，从而能够在不增加计算复杂度的情况下，有效提升感受野，聚合不同窗口的信息特征。可以看到，在中间特征图中，对于红色点，其感受区域为粉色方格组成，明显大于蓝色与绿色窗口。对于每一个窗口，注意力的计算公式如下：
四、实验结果
4.1 实验设置
本文依据传统实验设置，对于图像超分辨率任务，使用DIV2K和Flickr2K进行训练，并在Set5、Set14、B100、、Urban100和Manga109上进行测试，放大比例为2、3、4。在此任务中，我们提出两个模型CAT-A和CAT-R，分别使用上述介绍的axial-Rwin和regual-Rwin。对于压缩伪影去除，使用DIV2K、Flickr2K、BSD500、WED进行训练，在Classic5和LIVE1进行测试，压缩等级为10、20、30、40。在此任务中，我们使用CAT（CAT-A）。对于真实图像去噪，本文在SIDD上训练，并在SIDD和DND上测试。在此任务中，我们同样适用CAT（CAT-A）。以上所有测试，均适用PSNR和SSIM来测量复原效果。
4.2 消融实验
通过实验对比，我们可以发现，我们提出的Rwin-SA和Axial-Shift，与传统的方向窗口和Shift操作相比，在不增加计算复杂度（FLOPs）的情况下，性能明显提升。同时LCM模型，在CAT-A和CAT-R模型上均能提升性能，同时对复杂度的影响很小。此外，增大窗口大小也能明显提升模型性能，但是考虑到计算复杂度的影响，我们也不能无限制的提高窗口大小。
4.3 结果对比
图像超分辨率 由于篇幅问题，这里只展示部分结果，更多详细对比数据，可以参考论文。可以看到，本文提出的CAT-A与CAT-R，与目前主流的模型对比，均取得最佳性能。
压缩伪影去除 与目前主流模型对比，本文提出的CAT模型，在除了Classic5-q=10的测试数据上略低于SwinIR，其他结果均取得最佳性能。
真实图像去噪 在去噪任务上，我们提出的模型，与目前最有性能模型Restormer相比，性能相似，同时模型复杂度更低。
以上这些实验，均表明本文提出方法的有效性。
五、 结论
在本文中，我们提出了一个新的Transformer模型（Cross Aggregation Transformer，CAT）用于图像复原任务。CAT的核心为矩形窗口注意力（Rectangle-Window Self-Attention，Rwin-SA），具体而言Rwin-SA在不同组中分别执行H-Rwin和V-Rwin，从而聚合不同窗口的特征来增加感受区域。我们还提出轴向位移操作（Axial-Shift）来进一步增加窗口之间的交互。此外，局部补充模型（Locality Complementary Module，LCM）被提出，以将CNN的归纳偏置引入到Transformer中，实现局部与全局特征的耦合。大量的实验证明，我们提出的CAT在多个图像复原任务：图像超分辨率、压缩伪影去除、真实图像去噪，均取得最佳的性能。
参考文献
[1] Chao Dong, Chen Change Loy, Kaiming He, and Xiaoou Tang. Learning a deep convolutional network for image super-resolution. In ECCV, 2014.
[2] Chao Dong, Yubin Deng, Chen Change Loy, and Xiaoou Tang. Compression artifacts reduction by a deep convolutional network.
[3] Yulun Zhang, Kunpeng Li, Kai Li, Lichen Wang, Bineng Zhong, and Yun Fu. Image super-resolution using very deep residual channel attention networks. In ECCV, 2018.
[4] Yulun Zhang, Kunpeng Li, Kai Li, Bineng Zhong, and Yun Fu. Residual non-local attention networks for image restoration.
[5] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swintransformer: Hierarchical vision transformer using shifted windows. In ICCV, 2021.
[6] Xiaoyi Dong, Jianmin Bao, Dongdong Chen, Weiming Zhang, Nenghai Yu, Lu Yuan, Dong Chen, and Baining Guo. Cswin transformer: A general vision transformer backbone with cross-shaped windows. In CVPR, 2022.
[7] Jingyun Liang, Jiezhang Cao, Guolei Sun, Kai Zhang, Luc Van Gool, and Radu Timofte. Swinir: Image restoration using swin transformer. In ICCV, 2021.
[8] Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Shahbaz Khan,and Ming-Hsuan Yang. Restormer: Efficient transformer for high-resolution image restoration. In CVPR, 2022.
作者：陈铮Illustration by WOOBRO LTD from IconScout
快抢沙发！
