[Paper Review] Cross Aggregation Transformer for Image Restoration

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

Advanced Image Processing Techniques122 citations

TL;DR

This paper introduces Cross Aggregation Transformer (CAT) for image restoration, featuring rectangle-window self-attention (Rwin-SA) with axial-shift, and a Locality Complementary Module to fuse global attention with local CNN inductive biases.

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.

Motivation & Objective

• Motivate improving long-range dependency modeling in image restoration while keeping computation manageable.
• Develop a transformer-based architecture that aggregates features across non-square windows to enlarge receptive fields.
• Integrate CNN-like locality biases via a Locality Complementary Module to couple global and local information.
• Apply CAT to super-resolution, JPEG artifact reduction, and real image denoising and demonstrate state-of-the-art performance.

Proposed method

• Introduce rectangle-window self-attention (Rwin-SA) that uses horizontal and vertical rectangle windows across attention heads to enlarge the attention region.
• Fix one rectangle side to image height or width to form axial rectangle windows (axial-Rwin) for broader interaction; include an axial-shift operation between consecutive Rwin-SA blocks.
• Incorporate a Locality Complementary Module (LCM) that applies depthwise convolution on the value branch in parallel with self-attention to fuse local and global cues.
• Embed CAT blocks (CATB) in a backbone inspired by RCAN, replacing RCAB with CATB to form the Cross Aggregation Transformer (CAT).
• Use residual groups of CATBs and a reconstruction module tailored to SR, JPEG artifact reduction, and real denoising tasks.

Experimental results

Research questions

• RQ1 Can rectangle-window self-attention with axial-shift capture longer-range dependencies more effectively than square-window self-attention in image restoration?
• RQ2 Does integrating a locality-biased CNN component (LCM) with Transformer attention improve restoration quality without significant computational overhead?
• RQ3 How does CAT perform on standard image restoration tasks (SR, JPEG artifact reduction, real denoising) compared with state-of-the-art methods?

Key findings

• Rectangle-window self-attention with axial-shift outperforms square-window attention, yielding higher PSNR/SSIM on image restoration benchmarks.
• LCM provides additional gains by coupling global self-attention with local convolutional biases, with minimal FLOPs increase (~0.26%–0.32%).
• CAT-R (regular-Rwin) and CAT-A (axial-Rwin) achieve strong improvements across SR scales, particularly on Urban100; CAT-A generally offers the best performance among the proposed variants.
• CAT-based models outperform several state-of-the-art methods in image SR, JPEG artifact reduction, and real denoising in PSNR/SSIM metrics.
• The ablations show axial-Rwin with appropriate side-length settings is crucial for optimal performance, and the combination of rectangle windows and axial-shift expands receptive fields effectively.