Compacter: A Lightweight Transformer for Image Restoration

Authors: zhijian  wu
jun  li
yang  hu
dingjiang  huang
Venue: acmmm_org_ACMMM_2024_Conference
Published: 2024-03-14 08:40:07

Abstract: Although deep learning-based methods have made significant advances in the field of image restoration (IR), they often suffer from excessive model parameters. To tackle this problem, this work proposes a compact Transformer (Compacter) for lightweight image restoration by making several key designs. We employ the concepts of projection sharing, adaptive interaction, and heterogeneous aggregation to develop a novel Compact Adaptive Self-Attention (CASA). Specifically, CASA utilizes shared projection to generate Query, Key, and Value to simultaneously model spatial and channel-wise self-attention. The adaptive interaction process is then used to propagate and integrate global information from two different dimensions, thus enabling omnidirectional relational interaction. Finally, a depth-wise convolution is incorporated on Value to complement heterogeneous local information, enabling global-local coupling. Moreover, we propose a Dual Selective Gated Module (DSGM) to dynamically encapsulate the globality into each pixel for context-adaptive aggregation. Extensive experiments demonstrate that our Compacter achieves state-of-the-art performance for a variety of lightweight IR tasks with approximately 400K parameters.

Main paper: 
 introduction: 
 As a classical task in low-level computer vision, image restoration (IR) aims to recover high-quality counterparts from degraded images by removing degraded content. With the rapid development of deep learning, convolutional neural networks (CNNs) have become the de facto method for modern IR algorithms. Various advanced CNN designs were introduced into the image restoration task leading to further performance progress [55, 61, 62] . However, CNNs suffer from limitations in long-range dependencies modeling because convolution operators are more advantageous in extracting local information [21] . Recently, the dominance of CNNs has been challenged by vision Transformer-based models, which exhibit superior performance due to the advantage of self-attention in long-range modeling [15] . However, the quadratic complexity of self-attention makes it difficult to apply to image restoration tasks involving high-resolution images [6] . Many attempts have been made to develop more efficient self-attention mechanisms for image restoration. SwinIR [26] uses the shifted window mechanism to limit the scope of attention to the window, such that linear complexity is achieved. Similar strategies were adopted by Uformer [42] for building a hierarchical network architecture. Restormer [53] proposes to use channel-wise self-attention instead of spatial one to achieve linear complexity.Despite significant advances, most attempts have been devoted to improving large-scale image restoration models while neglecting the development of lightweight models. Lightweight image restoration networks are still fraught with great challenges. Specifically, most existing methods stack a single spatial or channel selfattention, making single-dimensional information modeling fail to achieve comprehensive feature interactions. Although some recent methods [7, 37] have been proposed to use the two self-attention mechanisms alternatively, it may lead to a large number of model parameters. In addition, the simple sequential combination is not fully capable of effectively modeling omnidirectional higher-order relational interactions. Besides, self-attention focuses more on global long-range information and is inferior at modeling local information [33] , which is detrimental to pixel-level image restoration tasks. These problems are more severe in lightweight models, which cannot stack enough layers due to limited capacity.In this paper, we propose a novel compact Transformer called Compacter for lightweight image restoration. Specifically, we propose a compact adaptive self-attention (CASA) mechanism through the design of projection sharing, adaptive interaction, and heterogeneous aggregation. CASA first utilizes shared projections for generating Query, Key, and Value to simultaneously establish spatial and channel self-attention. Subsequently, the bidirectional interaction process is utilized to adaptively propagate and integrate the two global information from two different axes, thereby realizing omnidirectional relational interactions. Finally, a depth-wise convolution is performed in parallel on Value to complement the locality to achieve heterogeneous aggregation of global and local information. With these designs, CASA enables comprehensive information dissemination and interaction within a compact computational unit. Furthermore, we propose a dual selective gated module (DSGM) to further calibrate the aggregated features and obtain per-pixel global dependencies. DSGM utilizes parallel branches to encode local information and gated to each other to adaptively promote favorable pixels and suppress detrimental pixels, thus producing high-quality restoration results. Thanks to two complementary components, our Compacter enables comprehensive pixel-relational interaction and maintains a desirable model size. We conducted comprehensive experiments on four lightweight image restoration tasks, our method achieved superior performance with significantly fewer parameters. Figure 1 illustrates the comparison of the image superresolution and low-light enhancement tasks. We further provide extensive ablation experiments to demonstrate the effectiveness of the architectural design.To sum up, we summarize the contributions of this paper as follows:• We propose a compact Transformer-based network, called Compacter, for lightweight image restoration.  
related work 2.1 image restoration: 
 The purpose of image restoration is to recover a clean version from a degraded image. In recent years, deep learning-based image restoration methods have achieved unprecedented success. In particular, CNN models have achieved state-of-the-art performance on various tasks such as super-resolution [16, 29, 38] , denoising [54] , deraining [23, 49] , and low-light enhancement [24, 59] . Image restoration algorithms are driven by the design of advanced CNN network architectures. For example, the residual connection has been introduced for tasks such as super-resolution [62] , image denoising [20] , and has become a necessary component of image restoration networks. Multi-scale representation learning is also widely used in various restoration tasks [55] due to its superior performance. Subsequently, the emerging attention mechanism has been explored for further performance improvements in low-level vision tasks. RCAN [61] coupled channel attention within the SR network, remarkably improving the representational capability of the model. To better learn feature correlation, SAN [13] further proposed a second-order channel attention module. However, most attempts have been devoted to improving the performance of large restoration models, while neglecting the development of lightweight models. Although some lightweight methods [27, 29, 45] have emerged recently, they tend to be dedicated to specific tasks such as super-resolution. In contrast, our Compacter is a general-purpose lightweight image restoration network that can be used for various restoration-related tasks such as image super-resolution, denoising, deraining, and low-light enhancement.Vision Transformer [15] has achieved impressive performance due to the self-attention mechanism, which naturally incorporates powerful dynamic weights and global dependency capture. However, the quadratic computational complexity of self-attention limits the application to vision tasks, which typically involve high-resolution images. IPT [6] pioneered introducing vision Transformer to image restoration tasks, which reduces the computational cost by decomposing the image into small patches and using a sequence of small patches as input. After that, SwinIR [26] used shifted-window-based self-attention to reduce the computational cost by restricting the scope of self-attention to the local window. A similar strategy is also used by Uformer [42] , which constructs a U-shaped hierarchical network. In addition, Restormer [53] adopts channel attention instead of the original spatial attention, making the computational complexity linear in spatial resolution. Despite achieving promising performance, they restrict attention to a local scope and may not fully exploit the potential of Transformers in capturing global dependencies [64] . Some recent methods [7, 37] suggest alternating spatial and channel attention to achieve further performance improvements. However, these serial models do not adequately represent the process of fusion of spatial and channel information, due to the lack of ability to model their interactions. In addition, these efforts remain committed to developing large vision Transformer models. In contrast, this paper is dedicated to designing a lightweight image restoration network. Loss Function. Following prior works [12, 35] , we optimize our model using dual-domain loss:EQUATIONEQUATIONwhere 𝐼 𝑃 and 𝐼 𝐺 represent the predicted image and ground truth. F (•) is the 2D Fast Fourier Transform. 𝜆 is set to 0.1 for balancing dual-domain training. 
compact transformer layer: 
 As shown in Figure 2 (b), our Compact Transformer Layer (CTL) consists of a compact adaptive self-attention (CASA) and a dual selective gated module (DSGM). It can be formulated as:EQUATIONEQUATIONwhere 𝑋 𝑖𝑛 and 𝑋 𝑜𝑢𝑡 are the input and output features, 𝐿𝑁 (•) denotes the LayerNorm operation [3] .Comprehensive and dense feature interactions help reconstruct finer results in restoration tasks that require dense per-pixel predictions. To fully mine the potential full correlation of features within a compact computational unit, we propose compact adaptive self-attention (CASA), as shown in Figure 3 . Our CASA is based on the philosophy of projection sharing, adaptive interaction and heterogeneous aggregation. It not only adaptively models spatial and channel-wise global information and their interactions but also achieves global-local coupling through the aggregation of heterogeneous operators. Projection Sharing. Given an input feature map 𝑋 ∈ R 𝐶 ×𝐻 ×𝑊 , we use linear layers to project it into query, key, and value 𝑄, 𝐾, 𝑉 ∈ R 𝐶 ×𝐻 ×𝑊 . Subsequently, the generated 𝑄, 𝐾, 𝑉 are fed to spatial window self-attention (SWSA) and channel-wise self-attention (CWSA), respectively. Based on this design, SWSA and CWSA share 𝑄, 𝐾, 𝑉 generated by the same projection, allowing to modelling of both spatial and channel information in a more compact manner.𝑄, 𝐾, 𝑉 = Linear 𝑖 (𝑋 ), 𝑖 ∈ {0, 1, 2}𝑌 𝑠 = SWSA(𝑄, 𝐾, 𝑉 ), 𝑌 𝑐 = CWSA(𝑄, 𝐾, 𝑉 )where Linear(•) denotes linear layer. The computational processes of SWSA and CWSA are elaborated as follows. Spatial Window Self-Attention (SWSA). Given the query, key and value 𝑄, 𝐾, 𝑉 , we split them into non-overlapping windows with the resolution of 𝑁 𝑤 and flatten them inEQUATIONEQUATION)where 𝐵 is a dynamic position bias, whose value is learned through an MLP [40] .Channel-Wise Self-Attention (CWSA). Given the query, key, and value 𝑄, 𝐾, 𝑉 , we reshape them into Q, K, V ∈ R 𝐶 ×𝐻𝑊 and compute the transposed-attention map across channels. The CWSA process is defined as:EQUATION)where 𝛼 is a learnable scaling factor. Adaptive Interaction. After establishing spatial and channel longrange dependencies, we use a bidirectional interaction process to adaptively propagate and integrate the two global information from two different axes. We adopt the cross-modulation method [17] to realize the bidirectional interaction. The computational procedure is formulated as follows:EQUATION)where ⊙ is element-wise multiplication. With bidirectional interaction, the spatial and channel self-attentions complement and augment each other, thereby enabling omni-axial global representations.Heterogeneous Aggregation. Although self-attention excels at modeling long-range dependencies, it lacks an inductive bias that favors extracting high-frequency local information [21] . To address this issue, we propose to add a depth-wise convolution operation on the value 𝑉 to complement the heterogeneous local information. Finally, all these representations are combined followed by aggregation using a linear layer:EQUATIONBased on this design, our CASA not only realizes comprehensive dense feature interactions but also provides heterogeneous feature aggregation capability at different granularities within a compact computational unit.We propose a dual selective gated module (DSGM) to further calibrate the aggregated features generated by CASA and encapsulate the comprehensive feature correlation to each pixel. Similar to the gating mechanism [14] , features are fed to two parallel linear layers followed by depth-wise convolution to encode local pixel information.EQUATIONThe local-aware representations are converted into context-aware weights via Sigmoid and GELU activation functions, and adaptively filter the extracted pixel features using dual selective gating. Finally, a linear layer is used to aggregate these representations.EQUATIONEQUATIONThe two different activation functions allow generating weights with different properties to adaptively promote beneficial pixels and suppress detrimental pixels when performing information aggregation. With the above design, our DSGM can more efficiently aggregate global information and local contexts while introducing less computation. 
experiment 4.1 experimental setup: 
 4.1.1 Datasets and Evaluation Metrics. We evaluate our Compacter on four image processing tasks. For the sake of fairness, the datasets used for all experiments follow exactly previous works [9, 10, 53, 65] .Training and testing datasets for each task are listed below. Image Super-Resolution. Our model is trained on the DIV2K [1] dataset. The evaluation is implemented on five public datasets: Set5 [5] , Set14 [48] , BSD100 [31] , Urban100 [22] and Manga109 [32] . Image Denoising. The image denoising task is trained on DIV2K, Flickr2K [36] , BSD400 [2] , and WED [30] , and tested on CBSD68 [31] , Kodak [18] , McMaster [58] , and Urban100 [22] , Low-Light Image Enhancement. The training and testing are implemented in the LOL [44] dataset. Image Deraining. We use Rain13K [19] dataset for training, while evaluate on Rain100H [50] and Test100 [56] . Evaluation metrics. The common PSNR and SSIM are used to measure performance. Following standard practice, both metrics are computed in the RGB channel, except for the image deraining task, which is evaluated in the Y-channel in YCbCr space, following previous works [23, 27] . 4We compare our method with state-of-the-art lightweight superresolution algorithms, as shown in Table 1 . It can be seen that our Low-light image LLFormer Compacter(Ours) GT URetinexNet Params PSNR SSIM KinD [63] 8,020K 20.86 0.7900 FIDE [47] 8,620K 18.27 0.6550 DRBN [51] 5,270K 20.13 0.8300 KinD++ [60] 8,275K 21.80 0.8338 EnlightenGAN [24] 8,640K 17.48 0.6507 Uformer [42] 5,290K 16.36 0.7110 Restormer [53] 26,130K 22.43 0.8230 URetinex-Net [45] 1,230K 21.32 0.8348 KinD-SKF [46] 8,500K 21.91 0.8350 LLFormer [39] 24,520K 23.65 0.8160Compacter (Ours) 407K 23.76 0.8375Compacter achieves the best performance on almost all five benchmark datasets for all scale factors. In particular, our method achieves an average of 0.05 dB higher PSNR than the latest CAMixerSR (CVPR2024) with 46% fewer parameters. In addition, our method outperforms prominent SRFormer-light (ICCV2023) on average at all scales with less than half of the parameters. Compared to NGswin (CVPR2023), our method achieves up to 0.43 dB PSNR improvement on Manga109 with fewer parameters. These results demonstrate the superiority of our method.In Figure 5 , we also provide a visualization comparison of different lightweight SR methods at ×4 scale. It can be seen that competing algorithms are prone to erroneous textures and patterns. In contrast, the images produced by our method are clearer and visually closer to the ground truth. These visual results also illustrate the advantages of our Compacter in recovering details and textures. We follow [10, 53] to perform blind denoising experiments on the synthetic benchmark datasets generated using additive white Gaussian noise. The lightweight denoising methods used for comparison are derived from [9] . Table 2 shows the performances of different approaches on several benchmark datasets for noise levels 15, 25, and 50. It can be seen that our method achieves the best tradeoff between efficiency and performance. Specifically, our method embraces significantly fewer parameters while obtaining superior performance compared with other methods. For example, for the challenging noise level 50 on the Urban100 dataset, our method outperforms SwinIR by 0.14 dB in PSNR, while the parameters are less than half of it. In addition, the far more perceptually relevant SSIM index shows that our Compacter has a significant advantage over other methods. Figure 6 shows the denoising results of different methods. It can be seen that existing methods fail to recover enough details due to severe noise degradation. In contrast, our method provides better restoration of structure and detail, resulting in a clearer restoration.We conduct low-light enhancement experiments on the LOL dataset. Table 3 shows the quantitative results, from which one can observe that our method achieves better performance while having significantly fewer parameters. Specifically, our Compacter outperforms the acclaimed general-purpose image restoration method Restormer [53] by 1.33dB PSNR while requiring only 1.5% of its parameters. Furthermore, our method also outperforms the much larger LL-Former [39] and KinD-SKF [46] , even though they are specifically tailored for low-light enhancement tasks. We further compared the visual results of the different methods in Figure 7 . Compared to competitors, our method generates images with more natural and vivid colors.We further evaluate our method on the image deraining task. The qualitative results can be viewed in Table 4 . As may be seen, our model improved over the state-of-the-art performances on both datasets. The average PSNR gain of our model over the latest lightweight model RAMiT [10] is 0.11 dB. Figure 8 shows a visual comparison on the more challenging Rain100H. Compared to other algorithms, our model produces images that are clearer and more consistent with the ground truth.For ablation experiments, we train super-resolution models on DIV2K for 200K iterations. The evaluations were performed on the Manga109 dataset [32] for ×2 scale super-resolution. FLOPs are calculated on input image size 3 × 256 × 256. Compact Transformer Layer. We performed ablation experiments to validate their effectiveness by removing CASA or DSGM from the CTL. For a fair comparison, we adjusted the depth of the models so that their parameters and FLOPs were comparable. Table 5 shows that removing CASA leads to severe performance degradation. This illustrates that the ability to model dense interactions is critical for restoration tasks that require per-pixel prediction. In addition, the absence of DSGM also leads to significant performance degradation. DSGM further improves the restoration performance by enhancing the useful information and suppressing the detrimental pixels through dual selective gating. In contrast, the full model achieves the best performance by coupling two complementary components. Compact Adaptive Self-Attention (CASA). We conduct an ablation study to investigate the effectiveness of each component of CASA. As shown in Table 6 , neither a single SWSA nor CWSA can achieve the desired performance, which suggests that unidimensional self-attention cannot realize comprehensive information interaction. Although the combination of the two attention techniques can achieve better performance, the absence of the projection sharing strategy leads to noticeable additional parameters and computational overheads. In addition, adaptive interaction further boosts performance barely costing any parameters and FLOPs. Heterogeneous aggregation also brings a 0.05dB performance improvement at a cost of a limited computational budget. Dual Selective Gated Module (DSGM). We further explored the effectiveness of DSGM design as shown in Table 7 . It can be seen that introducing depth-wise convolution in vanilla FFNs can bring performance improvement with few parameters and FLOPs. This strategy is also widely adopted by advanced image restoration networks [42, 53] . In addition, the gating mechanism helps to further reduce the model parameters and computation but leads to some performance degradation. In contrast, our DSGM utilizing the dual selective gating mechanism can achieve superior performance while barely imposing additional computations. 
conclusion: 
 In this paper, we present a novel compact Transformer for lightweight image restoration. Specifically, we propose compact adaptive self-attention (CASA) for comprehensive information dissemination and interaction within a compact computational unit. CASA simultaneously establishes the global context of spatial and channel as well as their interactions and achieves global-local coupling through heterogeneous aggregation. In addition, we propose a dual selective gated module (DSGM) to achieve dynamic context aggregation through a dual-path structure and gating mechanism. Thanks to two complementary components, our Compacter enables comprehensive pixel-level relational interactions while maintaining desirable model sizes. Extensive experiments demonstrate that Compacter achieves state-of-the-art performance on several restoration tasks with fewer model parameters. 
 
 Appendix:  


Decision: The paper received (2) borderline accept and (1) borderline reject. However, the negative comments are too general, which are not the key limitations of the proposed method. This paper presented a compact and efficient transformer-based model for image restoration. The motivation is clear, and all reviewers recognized the contribution of the reasonable design of the architecture, and the experimental results are sufficient to support the lightweight model and achieve better performance. We congratulate the authors on the acceptance of their paper!

Reviews: review: paper_summary: paper_summary: This paper proposes a novel approach to image restoration with Compacter, a compact and efficient Transformer-based model. The key innovations include the Compact Adaptive Self-Attention (CASA) mechanism and the Dual Selective Gated Network (DSGN). CASA integrates spatial and channel self-attention using shared projections, enhancing global-local feature interactions. DSGN dynamically adjusts and aggregates features through parallel branches and gating mechanisms, improving restoration quality. With significantly fewer parameters, Compacter still achieves competitive or better performance across various image restoration tasks against state-of-the-art methods.
strength_weakness: strength_weakness: - The proposed Compacter achieves impressive results on multiple restoration tasks and the authors provide extensive experimental results to demonstrate the effectiveness. 

- The use of projection sharing and efficient attention mechanisms significantly reduces the number of parameters while outperforming SOTA methods.

- The whole structure is straightforward to follow and easy to modify for different tasks.
limitations: limitations: - Although the experiments show that the model is effective, I think that the authors need to go further and demonstrate why the design of each component is valid. Especially for Tab 6, with only SWSA and 393K parameters, the model can already outperform most of the baseline models. However, in Tab 5, removing CASA leads to severe performance degradation with 749K parameters. I suggest that the authors provide further explanation of this.

- I am confused about some reported results:
    1) For the ablation study, the PSNR of the full model is 39.29 (Tab 5, 6, 7). But in Tab 1 it's 39.40. Why are they different?
    2) For Tab 6, with SWSA or CWSA, the number of parameters does not decline significantly. Why this is the case? If the authors "adjusted the depth of the models so that their parameters and FLOPs were comparable" as in Tab 5, I think it needs to be declared.
    3) For Tab 5, how do authors implement "w/o CASA"? 

- The authors didn't illustrate or cite how the comparison methods (ConvFFN and GDFN) are implemented. Also, the results cannot explain why the design of DSGN is effective, i.e. why choose GELU and Sigmoid?

- In Fig 3, Q, K, and V should come from three different linear layers rather than one linear layer (see Eq. 5).

- In Fig 2 b), there are two "element additions" missing.
paper_summary: paper_summary: This work proposes Compacter for lightweight image restoration by using projection sharing, adaptive interaction, and heterogeneous aggregation. Compacter achieves state-of-the-art performance for a variety of lightweight IR tasks.
strength_weakness: strength_weakness: 1. The paper is easy to follow.
2. The results seems promising.
limitations: limitations: 1. Design of the network is not novel.  I don’t see much new stuff, most of it is just existing work being integrated together.
2. The paper lacks runtime comparison for all experiments. Only the parameters and Flops comparison is not enough. 
3. More results on Real-world denoising, super-resolution and image dehazing is necessary.
paper_summary: paper_summary: In this paper, authors propose a compact Transformer (Compacter) for lightweight image restoration. In Compacter, they design Compact Adaptive Self-Attention (CASA) and dual selective gated network (DSGN). Thanks to these two complementary components, the Compacter achieves state-of-the-art performance on several restoration tasks with fewer model parameters.
strength_weakness: strength_weakness: The paper addresses the challenges of computational cost overhead in image restoration. Extensive experiments demonstrate that the proposed Compacter achieves SOTA results on various lightweight restoration tasks with significantly fewer parameters.
limitations: limitations: 1）The proposed network in this paper builds upon existing CNN-based and Transformer-based approaches for image restoration. The lack of novelty in the paper stems from the fact that the proposed CASA and DSGN do not introduce fundamentally new concepts or techniques.

2）The name of Dual Selective Gated Network should be changed to improve the readability of the article, as it is only a key component of the Compact network. For example, it can be modified to Dual Selective Gated Block or Dual Selective Gated Module.

3）In the experimental section, it is recommended to add some inference speed experiments.

4）In the experimental section, some visual feature maps are suggested to be added to verify the effectiveness of the proposed module.