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 work presents cross-layer attention (CLA) modules to find informative keys across different CNN layers for each query feature. Furthermore, an adaptive cross-layer attention (ACLA) is also formulated to dynamically select keys from different CNN layers by using a NAS method. After that, the authors embedded the presented CLA or ACLA modules into EDSR to formulate a deep model for image restoration. Experimental results on several image restoration tasks show that the developed deep model outperforms state-of-the-art methods.