Adaptive Cross-Layer Attention for Image Restoration

1 INTRODUCTION . Image restoration algorithms aim to recover a high-quality image from a contaminated input image by solving an ill-posed image restoration problem . There are various image restoration tasks depending on the type of corruptions , such as image denoising ( Zhang et al. , 2019 ; Liu et al. , 2018b ) , demosaicing ( Zhang et al. , 2017b ; 2019 ) , single image super-resolution ( Fan et al. , 2019 ; Lai et al. , 2017 ; Tai et al. , 2017 ) , and image compression artifacts reduction ( Zhang et al. , 2017a ) . To restore corrupted information from the contaminated image , a variety of image priors ( Buades et al. , 2005 ; Zoran & Weiss , 2011 ; Zontak et al. , 2013 ) were proposed . Recently , image restoration methods based on deep neural networks have achieved great success . Inspired by the widely used non-local prior , most recent approaches based on neural networks ( Zhang et al. , 2019 ; Liu et al. , 2018b ) adapt non-local attention into their neural network to enhance the representation learning , following the non-local neural networks ( Wang et al. , 2018 ) . In a nonlocal block , a response is calculated as a weighted sum over all pixel-wise features on the feature map to account for long-range information . Such a module was initially designed for high-level recognition tasks such as image classification , and it has been proven to be beneficial for low-level vision tasks ( Zhang et al. , 2019 ; Liu et al. , 2018b ) . Though attention modules have been shown to be effective in boosting the performance . Most attention modules only explore the correlation among features at the same layer . Actually , features at different intermediate layers encode variant information at different scales , and might be helpful to augment the information used in recovering the high-quality image . Motivated by the potential benefit of exploring feature correlation across intermediate layers , Holistic Attention Network ( HAN ) ( Niu et al. , 2020 ) is proposed to find interrelationship among features at hierarchical levels with a Layer Attention Module ( LAM ) . However , LAM assigns a single importance weight to all features at the same layer and neglects the difference of spatial positions of these features . Recent research in omnidirectional representation ( Tay et al. , 2021 ) suggests that exploring the relationship among features at different layers can benefit the representation learning of neural networks . Nevertheless , calculating correlation among features at hierarchical layers is computationally expensive due to the quadratic complexity of dot product attention . The complexity of such cross-layer attention design is increased from ( HW ) 2L to ( HWL ) 2 , where H , W are the height and width of the feature map and L is the number of layers . 1.1 CONTRIBUTIONS . Our contributions are listed as follows . First , in order to address the problem caused by only referring to keys within the same layer in most attention modules , we propose a novel attention module termed Cross-Layer Attention ( CLA ) , which searches for keys across different layers for each query feature . With the help of the deformation mechanism , CLA only attends to a small set of keys at different layers for each query feature . Second , we propose an improved CLA termed Adaptive Cross-Layer Attention ( ACLA ) which selects an adaptive number of keys at each layer for each query , and searches for the optimal insert positions of ACLA modules . We deploy ACLA modules on commonly used neural network models , e.g . EDSR ( Lim et al. , 2017a ) , for image restoration . Extensive experiments on single image super-resolution , image denoising , image compression artifacts reduction , and image demosaicing demonstrate the effectiveness of our approach . 2 RELATED WORKS . 2.1 NEURAL NETWORKS FOR IMAGE RESTORATION . Adopting neural networks for image restoration has achieved great success by utilizing their power in representation . ARCNN ( Dong et al. , 2015a ) was first proposed to use CNN for compression artifacts reduction . Later , DnCNN ( Zhang et al. , 2017a ) uses residual learning and batch normalization to boost performance of CNN for image denoising . In IRCNN ( Zhang et al. , 2017b ) , a learned set of CNNs are used as denoising prior for other image restoration tasks . For single image superresolution ( Lai et al. , 2017 ; Zhang et al. , 2018d ; Haris et al. , 2018 ; Fan et al. , 2019 ) , even more efforts have been devoted to designing advanced architectures and learning methods . For example , RDN ( Zhang et al. , 2018d ) and CARN ( Ahn et al. , 2018 ) fuse low-level and high-level features with dense connection to provide richer information and details for reconstructing . Recently , non-local attention ( Liu et al. , 2018b ; Dai et al. , 2019 ; Mei et al. , 2020 ) is also used to further boost the peformance of CNN for image restoration . 2.2 ATTENTION MECHANISM . Attention mechanism has been applied to many computer vision tasks , such as image captioning ( Xu et al. , 2015 ; Chen et al. , 2017 ) and image classification ( Hu et al. , 2018 ; Wang et al. , 2017 ) . Non-local attention ( Wang et al. , 2018 ) was first proposed to capture long-range dependencies for high-level recognition tasks . Recently , several works propose to leverage non-local attention for low-level vision tasks . In NLRN ( Liu et al. , 2018b ) a recurrent neural network is proposed to incorporate non-local attention . RNAN ( Zhang et al. , 2019 ) proposed a residual local and non-local mask branch to obtain non-local mixed attention . RCAN ( Zhang et al. , 2018c ) exploits the interdependencies among feature channels by generating different attention for each channel-wise feature . HAN ( Niu et al. , 2020 ) is proposed to find interrelationship among features at hierarchical levels with a layer attention module . Besides , some recent works attempt to explore the benefits of transformer based models for image restoration . IPT ( Chen et al. , 2021 ) is proposed to solve various restoration problems in a multi-task learning framework based on visual Transformer . SwinIR ( Liang et al. , 2021 ) adopt the architecture of Swin Transformer . However , compared with methods using CNN architecture , transformer-based image restoration methods usually use large datasets for training . Specifically , IPT uses ImageNet to pretrain the model . SwinIR adapts a combination of four datasets consisting of over 8000 high-quality images as training set for the tasks of denoising and compression artifact reduction . 2.3 NEURAL ARCHITECTURE SEARCH . Neural Architecture Search ( NAS ) has attracted lots of attention recently . Early works of NAS adopt heuristic methods such as reinforcement learning ( Zoph & Le , 2016 ) and evolutionary algorithm ( Xie & Yuille , 2017 ) . The search process with such methods requires huge computational resources . Recently , various strategies are designed to reduce the expensive costs including weight sharing ( Pham et al. , 2018 ) , progressive search ( Liu et al. , 2018a ) and one-shot search ( Liu et al. , 2018c ; Xie et al. , 2018 ) . For example , DARTS ( Liu et al. , 2018c ) firstly relaxes the search space to be continuous and conducts differentiable search . The architecture parameters and network weights are trained simultaneously by gradient descent to reduce the search time . Despite the success of NAS methods for classification , dense prediction tasks such as semantic image segmentation and image restoration , usually demand more complicate network architectures . Some recent works have been devoted to explore hierarchical search space for dense prediction tasks . For example , Auto-DeepLab ( Liu et al. , 2019 ) introduces a hierarchical search space for semantic image segmentation . DCNAS ( Zhang et al. , 2021 ) build a densely connected search space to extract multi-level information . HNAS ( Guo et al. , 2020 ) also adopts a hierarchical search space for single image super-resolution . 3 METHOD . 3.1 CLA : CROSS-LAYER ATTENTION . Non-Local Attention . Non-local attention ( Wang et al. , 2018 ) is designed to integrate self-attention mechanism into convolutional neural networks for computer vision tasks . It is usually applied on an input feature map x ∈ RH×W×C to explore self similarities among all spatial positions of that feature map . For illustration purpose , we reshape x to N × C , where N = H ×W . A generic non-local attention can be formulated as yi = 1 C ( x ) N∑ n=1 f ( xi , xn ) g ( xn ) , ( 1 ) where i indexes the spatial position of feature maps . y is the output of non-local attention with the same size as x. f ( xi , xn ) outputs the pairwise affinity between the query feature xi and its key feature xn . g ( xn ) computes an embedding of feature xn . C ( x ) is a normalization term . Non-local attention in Equation ( 1 ) is usually wrapped into a non-local block ( Wang et al. , 2018 ) with a residual connection from the input feature x . The mathematical formulation is given as z = h ( y ) + x , ( 2 ) where h denotes a learnable feature transformation . Cross-Layer Attention . To search for keys from different layers for each query feature , we propose Cross-Layer Attention ( CLA ) . To derive the formulation of CLA , we start by adapting non-local attention in Equation ( 1 ) to a cross-layer design . Suppose that xi is the output of the i-th layer in a CNN backbone for image restoration , where i ∈ { 1 , . . . , L } and L is the number of layers . A non-local attention with cross-layer design can be formulated as yji = 1 C ( xj ) j∑ l=1 N∑ n=1 f ( xji , x l n ) g ( x l n ) , ( 3 ) where the superscripts j , l index the layer and the subscripts i , n index the spatial locations of features . y , x denote the output feature and input feature respectively . Query feature in this formulation can thus refer to features in previous layers . However , given the quadratic complexity of correlation computation , the complexity of such adaption of non-local attention is increased fromN2L to ( NL ) 2 . In order to mitigate the expensive computation cost , we leverage deformable mechanism proposed in deformable convolution ( DCN ) ( Dai et al. , 2017 ) to design our proposed CLA as yji = 1 C ( xj ) j∑ l=1 K∑ k=1 f ( xji , x l ( pi + ∆pik ) ) g ( x l ( pi + ∆pik ) ) , ( 4 ) where k indexes the sampled keys . pi represents the 2-d spatial position of the query in the feature map , and ∆pik is the 2-d offset from the position of query pi to the corresponding sampled key . As pi + ∆pik can be fractional , bilinear interpolation is used as in ( Dai et al. , 2017 ) to compute x ( pi+∆pik ) . To further reduce the computation complexity in computing attention weights , attention weights can be generated from query feature alone . Thus , our proposed CLA can be simplified as yji = 1 C ( xj ) j∑ l=1 K∑ k=1 f ( xji ) g ( x l ( pi + ∆pik ) ) , ( 5 ) where f is the function to generate attention weights from the query feature . Similar to the design for non-local attention module in ( Wang et al. , 2018 ) , when deploying CLA in CNN backbones , we also wrap it into a non-local block with residual connection . With our proposed CLA , each query feature now can refer to a fixed number of keys from each previous layer . However , query features at different spatial positions may have different preferences on keys sampled from different layers . Especially for image restoration tasks , the restoration process at different spatial position may vary significantly due to the diversity of textures in an image . Besides , when deploying CLA in CNN backbones for image restoration , we find that increasing the number of inserted CLA will not constantly improve the performance , but result in much heavier computation cost . This observation motivates us to search for an optimal configuration of deploying CLA in neural networks to efficiently explore its representation learning capacity .

Summary: This paper proposes a cross-layer attention module for the image restoration tasks. Unlike previous conventional non-local attention approaches that find correlated keys within the same layer, the proposed method selects keys across the layers. In order to prevent expensive computational costs, the authors propose adaptive cross-layer attention modules. The proposed method is validated on various image restoration tasks including SR, denoising, demosaicing, and compression.