Source-linked AI summary

Cross Aggregation Transformer for Image Restoration

Zheng Chen, Yulun Zhang, Jinjin Gu, Yongbing Zhang, Linghe Kong, Xin Yuan

arXiv:2211.13654v2cs.CV

TL;DR

Image restoration needs efficient modeling of long-range dependencies because global Transformer attention is quadratic, while local square windows limit interaction between windows. CAT addresses this with rectangle-window attention, axial shifts, and a locality complementary module, and reports state-of-the-art results across several restoration tasks.

  • Problem

    Global self-attention has quadratic complexity for image resolution, while local square-window methods lack direct inter-window interaction and restrict long-range dependency modeling.

  • Method

    CAT combines horizontal and vertical rectangle-window self-attention, Axial-Shift for window interaction, and a convolutional Locality Complementary Module.

  • Results

    CAT achieves state-of-the-art quantitative and visual performance on image super-resolution, JPEG compression artifact reduction, and real image denoising.

  • Takeaways & Limitations

    CAT provides a restoration Transformer with a large receptive field and cross-window feature aggregation while retaining linear complexity.

Abstract

from arXiv · show

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.

1 Introduction

Image restoration seeks to recover high-quality images from degraded inputs, but global Transformer attention is computationally expensive and local windows limit cross-window dependencies. CAT addresses this with rectangle-window attention, axial shifts, and a locality complement that combines global and local information.

  • Image restoration recovers high-quality images from low-quality counterparts across tasks including super-resolution, denoising, and JPEG artifact reduction.
  • Global Transformer self-attention models long-range dependencies but has quadratic complexity O((HW)2), limiting high-resolution image restoration.
  • Local square-window or channel-wise attention reduces complexity but restricts inter-window interaction, receptive-field growth, or spatial information.
  • CAT uses horizontal and vertical rectangle-window attention across heads to expand the attention area and aggregate features across windows without increasing computational complexity.
  • Axial-Shift explicitly links same-direction windows while implicitly enabling interaction between horizontal and vertical windows.
  • The Locality Complementary Module applies convolution to self-attention values in parallel, coupling Transformer-based global dependencies with CNN locality and translation invariance.
  • CAT is evaluated on image super-resolution, JPEG compression artifact reduction, and real image denoising, with experiments reporting state-of-the-art quantitative and visual performance.

2 Related Work

Prior work established CNNs as effective image-restoration models and explored efficient Transformer attention through local, shifted, cross-shaped, and combined global-local designs. CAT builds on this progression with a Transformer block tailored to restoration.

  • CNNs became mainstream in image restoration after strong results from SRCNN and ARCNN, with later work using residual, dense, dropout, and UNet strategies.
  • Vision Transformer research introduced efficient attention variants including Swin Transformer's shifted local windows and CSwin's cross-shaped windows.
  • Twins combines global and local attention, while restoration methods such as SwinIR use 8×8 windows to avoid the quadratic cost of full self-attention.
  • Figure 1 presents CAT's overall architecture and illustrates its cross aggregation Transformer block.

3 Method

CAT replaces square-window attention with cross-window aggregation through rectangle windows, axial shifts, and a convolutional locality complement. These designs combine directional and long-range feature interactions with local image structures in an efficient restoration architecture.

  • CAT Architecture: CAT combines shallow extraction, deep feature extraction, and task-specific reconstruction modules for image restoration.The deep feature extractor uses residual groups; reconstruction differs across super-resolution, JPEG artifact reduction, and real-image denoising.
  • Rectangle-Window Self-Attention: Rwin-SA divides rectangle windows into horizontal and vertical forms assigned to different attention heads in parallel.H-Rwin uses sh<sw, while V-Rwin uses sh>sw, enabling directional feature aggregation.
  • Rectangle-Window Self-Attention: H-Rwin and V-Rwin aggregate features across different windows and expand the attention area without increasing computational complexity.Rectangle windows capture horizontal and vertical features that square windows may not handle comparably, especially for directional and repetitive textures.
  • Rectangle-Window Self-Attention: Axial-Rwin fixes one rectangle side to the image height or width, creating flexible stripe-like windows whose other side has length sl.Its complexity is O(axial-Rwin) = HWC × (4C + sl × H + sl × W).
  • Axial-Shift Operation: Axial-shift shifts horizontal and vertical window partitions between consecutive Rwin-SA blocks to increase interactions among different windows.It supports same-direction window interaction and implicitly enables interaction between horizontal and vertical windows.
  • Locality Complementary Module: The Locality Complementary Module applies a 3×3 depth-wise convolution to value V in parallel with attention, coupling global dependencies with local image structures.The convolution operates in the same feature domain as attention and complements Rwin-SA with CNN inductive biases such as translation invariance and locality.

4 Experiments

Experiments evaluate CAT across image super-resolution, JPEG artifact reduction, and real image denoising, including ablations of its attention and locality designs. CAT generally outperforms competing methods while retaining favorable efficiency, with gains in both quantitative and visual restoration quality.

  • Experimental settings: CAT is trained and evaluated on standard datasets for super-resolution, JPEG artifact reduction, and real image denoising, using multiple scale factors and JPEG qualities.Super-resolution uses ×2, ×3, and ×4 settings; JPEG evaluation uses Classic5 and LIVE1.
  • Ablation study: 0.16 dB: rectangle-window attention improves PSNR from 32.50 dB to 32.66 dB without increasing FLOPs, while axial shift reaches 32.91 dB.The comparison is conducted in the ablation study against square-window and shifted square-window attention.
  • Ablation study: 0.07 dB and 0.10 dB: LCM improves CAT-R and CAT-A, respectively, while increasing FLOPs by only 0.26%–0.32%.LCM is embedded in self-attention without changing the attention implementation.
  • Image super-resolution: 0.45 dB: CAT-A improves PSNR over SwinIR for ×2 super-resolution, with CAT-R and CAT-A outperforming other methods across datasets and scale factors.The largest improvements occur on Urban100, which contains many repetitive texture structures.
  • Visual results: CAT recovers sharper edges, explicit textures, and more high-frequency structural details than compared methods in challenging super-resolution and JPEG artifact-reduction examples.The visual comparisons particularly report reduced blurring and blocking artifacts.
  • JPEG compression artifact reduction: 29.89 dB: CAT achieves this LIVE1 result at JPEG quality q=10, while CAT+ performs best among compared methods across all tested JPEG qualities and datasets.At q=40, CAT improves over SwinIR by 0.05 dB on one comparison and 0.06 dB on another reported comparison.
  • Real image denoising: CAT performs better than compared methods on SIDD and DND except Restormer, while achieving 0.02 dB higher PSNR than Restormer on DND with fewer parameters.The paper describes CAT as comparable to Restormer with fewer parameters.
  • Efficiency: CAT has less computational complexity and fewer parameters than EDSR, and substantially lower computational complexity than CSNLN in the reported super-resolution comparison.FLOPs are measured for a 3×512×512 output and PSNR on Urban100 (×4).

5 Conclusion

The paper concludes that CAT combines rectangle-window attention, axial shifting, and locality-aware convolution to aggregate information across windows. Experiments across three restoration tasks show that CAT outperforms current state-of-the-art methods.

  • Conclusion: CAT combines Rwin-SA, axial shift, and LCM to model cross-window dependencies while coupling global self-attention with local convolutional information.Rwin-SA uses parallel horizontal and vertical rectangle-window attention, while LCM adds CNN inductive bias.
  • Conclusion: Extensive experiments on image super-resolution, JPEG compression artifact reduction, and real image denoising demonstrate that CAT outperforms current state-of-the-art methods.The conclusion presents this result across the three listed image restoration applications.
Loading 2211.13654v2…