[论文解读] Cross Aggregation Transformer for Image Restoration

Zheng Chen, Yulun Zhang+4 more|arXiv (Cornell University)|Nov 24, 2022

Advanced Image Processing Techniques被引用 122

一句话总结

本论文提出 Cross Aggregation Transformer (CAT) 用于图像恢复，特征矩阵形窗口自注意力 (Rwin-SA) 结合轴向移位，以及 Locality Complementary Module 将全局注意力与局部 CNN 归纳偏差融合。

ABSTRACT

Recently, Transformer architecture has been introduced into image restoration to replace convolution neural network (CNN) with surprising results. Considering the high computational complexity of Transformer with global attention, some methods use the local square window to limit the scope of self-attention. However, these methods lack direct interaction among different windows, which limits the establishment of long-range dependencies. To address the above issue, we propose a new image restoration model, Cross Aggregation Transformer (CAT). The core of our CAT is the Rectangle-Window Self-Attention (Rwin-SA), which utilizes horizontal and vertical rectangle window attention in different heads parallelly to expand the attention area and aggregate the features cross different windows. We also introduce the Axial-Shift operation for different window interactions. Furthermore, we propose the Locality Complementary Module to complement the self-attention mechanism, which incorporates the inductive bias of CNN (e.g., translation invariance and locality) into Transformer, enabling global-local coupling. Extensive experiments demonstrate that our CAT outperforms recent state-of-the-art methods on several image restoration applications. The code and models are available at https://github.com/zhengchen1999/CAT.

研究动机与目标

• 在保持计算量可控的前提下，推动在图像恢复中改善长程依赖建模。
• 开发一种基于 transformer 的架构，在非方形窗口之间聚合特征以扩大感受野。
• 通过 Locality Complementary Module 融合 CNN 风格的局部性偏置，以耦合全局与局部信息。
• 将 CAT 应用于超分辨、JPEG 伪影抑制和真实图像去噪，并展示最先进的性能。

提出的方法

• 引入矩形窗口自注意力（Rwin-SA），在注意力头之间使用水平和垂直矩形窗口以扩大注意区域。
• 将一个矩形边固定为图像的高度或宽度，以形成轴向矩形窗口（axial-Rwin），实现更广泛的交互；在连续的 Rwin-SA 块之间加入轴向移位操作。
• 结合 Locality Complementary Module (LCM)，在自注意力并行的同时对值分支应用逐通道卷积，以融合局部与全局线索。
• 在受 RCAN 启发的骨干网络中嵌入 CATB（CAT 块），用 CATB 替换 RCAB 以形成 Cross Aggregation Transformer (CAT)。
• 使用由 CATB 组成的残差组，以及面向 SR、JPEG 伪影抑制和真实去噪任务定制的重建模块。

实验结果

研究问题

• RQ1 在图像恢复中，带轴向移位的矩形窗口自注意力是否比方形窗口自注意力更有效地捕获更长距离的依赖？
• RQ2 将具有局部偏置的 CNN 组件（LCM）与 Transformer 注意力结合，是否在不显著增加计算量的前提下提升恢复质量？
• RQ3 与最先进方法相比，CAT 在标准图像恢复任务（SR、JPEG 伪影抑制、真实去噪）上的表现如何？

主要发现

• 带轴向移位的矩形窗口自注意力优于矩形窗口自注意力，能够在图像恢复基准上带来更高的 PSNR/SSIM。
• LCM 通过耦合全局自注意力与局部卷积偏置带来额外提升，FLOPs 增长极小（约 0.26%–0.32%）。
• CAT-R（常规 Rwin）和 CAT-A（轴向 Rwin）在各个 SR 尺度上均取得显著提升，特别是在 Urban100 上；CAT-A 通常在所提出的变体中提供最佳性能。
• 基于 CAT 的模型在图像 SR、JPEG 伪影抑制和真实去噪任务中的 PSNR/SSIM 指标超过多种最先进方法。
• 消融结果表明，带有适当边长设置的 axial-Rwin 对于达到最佳性能至关重要，同时矩形窗口与轴向移位的组合可有效扩大感受野。