Adaptive Cross-Layer Attention for Image Restoration

Yancheng Wang, Yingzhen Yang, Chong Chen, Ning Xu

Keywords: image restoration, neural architecture search, non-local attention

Abstract: Non-local attention module has been proven to be crucial for image restoration. Conventional non-local attention processes features of each layer separately, so it risks missing correlation between features among different layers. To address this problem, we propose Cross-Layer Attention (CLA) module in this paper. Instead of finding correlated key pixels within the same layer, each query pixel is allowed to attend to key pixels at previous layers of the network. In order to mitigate the expensive computational cost of such hierarchical attention design, only a small fixed number of keys can be selected for each query from a previous layer. We further propose a variant of CLA termed Adaptive Cross-Layer Attention (ACLA). In ACLA, the number of keys to be aggregated for each query is dynamically selected. A neural architecture search method is used to find the insert positions of ACLA modules to render a compact neural network with compelling performance. Extensive experiments on image restoration tasks including single image super-resolution, image denoising, image demosaicing, and image compression artifacts reduction validate the effectiveness and efficiency of ACLA.

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2203.03619/code)