ShareFormer: Share Attention for Efficient Image Restoration

Thank you sincerely for your thorough review of our paper. We highly value the questions you raised. Below, we have provided detailed responses to each of them and made corresponding improvements to the paper. We hope our responses and revisions sufficiently address your concerns, and we genuinely appreciate your consideration for an improved rating of our paper.

For Q1, Could you please also list the number of parameters of each model in Tables 4 and 6?

Certainly. We have included the parameter number and latency information for each model in the table below. The size of the model used for image denoising (corresponding to Table 4 in the main paper) and JPEG CAR (corresponding to Table 6 in the main paper) is the same. Therefore, we present these in a single table. Additionally, the details from this table have also been added to Appendix C.

PSNR on image denoising was calculated on the CBSD68 dataset, $\sigma=25$. PSNR on image JPEG CAR was calculated on the Classic5 dataset, $q=40$. "N/A" indicates that the corresponding models are too heavy for the NVIDIA RTX 3090 GPU. "-" means that the result is not available.

For Q2, Thus, what if introducing a trainable parameter to scale the residual connection instead of directly adding a residual connection?

Thank you for your insightful suggestion. In fact, the introduction of a trainable residual connection is a design we had considered and experimented with earlier. However, we opted not to pursue this approach. The table below presents the experimental results on the classical $4 \times$ image super-resolution task, the test set is Urban100, where the residual scale was initialized to 1.

It is evident that incorporating diverse types of residuals does not result in a significant performance enhancement for ShareFormer in the classical $4 \times$ super-resolution task. This could be attributed to ShareFormer's intrinsic capacity to regulate value whereby supplementary trainable residual coefficients will not advance the network's performance (or a more detailed network initialization might be necessary).

For W1, I am concerned that the greater the complexity of the system, the higher the likelihood of training instability

Thank you for raising an interesting research question. The correlation between model complexity and training stability is a topic of great importance. We have identified three initial perspectives on the issue:

1. Firstly, it is important to note that trainability is primarily intended to demonstrate that the model can converge rapidly[1], rather than stably[2]. In terms of stability of convergence, it may be more appropriate to use the sharpness as a metric [3].
2. In the three tasks discussed in this article, we did not notice any apparent variation in the training steadiness of models that differ in size. Nonetheless, we suspect that if the number of parameters is exceedingly high, the training instability will intensify as the intricacy of the model increases. Indeed, many studies relevant LLM scaling laws [4] have established that the more extensive the number of parameters, the less stable the training is.
3. It is a consensus within the deep learning field that more complex networks require more training tricks. there is an Figure 1 And Table 3&4 in RCAN-it[5] for complex RCANs within CNN networks, Figure 2 in ConvNeXT[6] also makes use of a number of tricks to give a pump in the performance of convolutional networks, and Table 8. in DeiT[7] proposes a series of tricks to make Transformer more stable to train. All these instances underscore the need for meticulously designed tricks in training to enhance the stability of complex networks.

For W2, In Tables 4, 5, and 6, DRUNet reached similar performance and much lower latency compared with ShareFormer.

On this point of your concern, the choice between traditional CNNs and emerging Transformers is a widespread concern in the whole field. Here are a few of our perspectives:

1. As previously mentioned, DRUNet has approximately double the number of parameters compared to ShareFormer. This has led to a limitation in its practical application due to insufficient lightweight capabilities.
2. Unlike methods based on CNNs, models based on Transformers in the field of image restoration are still in the early and exploratory stages, with a considerable gap from practical application in industrial settings. Current efforts to improve inference speed and trainability, such as ELAN and our ShareFormer, represent steps toward achieving an optimal and practical Transformer network in the future. While it's just a small step, we believe our work holds significant relevance.
3. In reality, our method has made significant advances with progress of 0.38dB on the notoriously difficult Urban100 dataset, compared to DRUNet. We believe this improvement is significant and should be noteworthy in the selection of methods.
4. Unlike tasks like image classification and detection, numerical metrics are not the sole determinant of an image restoration method's effectiveness. Refer to Figure 10 in the main manuscript for a close-up view, which shows that our method excels in restoring the rich and intricate texture of the original image.

References

[1] Jacot A, Gabriel F, Hongler C. Neural tangent kernel: Convergence and generalization in neural networks[J]. Advances in neural information processing systems, 2018, 31.

[2] Berlyand, Leonid and Jabin, Pierre-Emmanuel and Safsten, C. Alex "Stability for the training of deep neural networks and other classifiers" Mathematical Models and Methods in Applied Sciences 2021, 31, 11, 2345-2390.

[3] Philip M. Long, Peter L. Bartlett. "Sharpness-Aware Minimization and the Edge of Stability". arXiv preprint arXiv:2309.12488.

[4] Qin, Zhen and Li, Dong and Sun, Weigao and Sun, Weixuan and Shen, Xuyang and Han, Xiaodong and Wei, Yunshen and Lv, Baohong and Yuan, Fei and Luo, Xiao and others "Scaling transnormer to 175 billion parameters" arXiv preprint arXiv:2307.14995.

[5] Lin, Zudi and Garg, Prateek and Banerjee, Atmadeep and Magid, Salma Abdel and Sun, Deqing and Zhang, Yulun and Van Gool, Luc and Wei, Donglai and Pfister, Hanspeter "Revisiting RCAN: Improved Training for Image Super-Resolution" arXiv preprint arXiv:2201.11279.

[6] Zhuang Liu and Hanzi Mao and Chao-Yuan Wu and Christoph Feichtenhofer and Trevor Darrell and Saining Xie "A ConvNet for the 2020s" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) 2022.

[7] Touvron, Hugo and Cord, Matthieu and Douze, Matthijs and Massa, Francisco and Sablayrolles, Alexandre and Jegou, Herve "Training data-efficient image transformers & distillation through attention" International Conference on Machine Learning (ICML) 2021.

Thanks for your response. The answer solved my concern about the comparison with DRUnet. It's glad to see you have already done the experiments of learnable residual scale which addressed my concern regarding the excessively high weights of value.

I agree with reviewer nYLY's point W3, which is similar to my point W1: The ablation study is limited and fails to comprehensively assess the contribution of each component. Providing additional relevant ablation studies would make the argument more convincing

For W5, Besides the Image Super-Resolution, the overall improvements appear relatively incremental over strong prior work like SwinIR and Restormer.

Your concern pertains to a wider challenge within the field of image restoration. Let me elaborate on a couple of aspects:

1. Image restoration poses unique challenges compared to other high-level visual tasks in terms of defining metrics and achieving improvements. As stated in the discussion of this survey [4]: "More accurate metrics need to be found for image denoising. PSNR and SSIM are popular metrics for the task of image restoration. PSNR suffers from excessive smoothing, which is very difficult to recognize indistinguishable images. SSIM depends on brightness, contrast and structure, and therefore cannot accurately evaluate image perceptual quality." Metrics such as PSNR are not as absolute as accuracy in classification tasks or IoU in detection tasks.
2. The assessment of image quality remains challenging due to the lack of a flawless objective metric. Conventional measures, such as PSNR and SSIM, oversimplify the multidimensional evaluation of image quality. Therefore, the numerical differences among cutting-edge techniques, especially those founded on Transformers, become less significant. However, Fig. 10 in our paper and detailed perceptual comparisons in Appendix C reveal our approach's ability to recover finer details, including dense dots in real images. This demonstrates the visual superiority of our method over those of the competition.
3. It is crucial to emphasize that we aim to achieve an optimal Pareto balance between performance and latency rather than focusing solely on performance gains. In this respect, our method has a profound impact on the efficient use of algorithms, notably extensive Transformer-based models, in authentic industrial environments.

[4] Chunwei Tian, Lunke Fei, Wenxian Zheng, Yong Xu, Wangmeng Zuo, Chia-Wen Lin, Deep learning on image denoising: An overview, Neural Networks, Volume 131, 2020, Pages 251-275, ISSN 0893-6080.

For W2, The gains in trainability from the residual connections could use more detailed analysis and intuition.

Sharing Attention reduces redundancy while also reducing the efficiency of information flow by restricting the channels for layer-by-layer information flow (you can think of the operation Share as a residual connection without residual branches, i.e., $y = f(x) + x \ \text{and} \ f(x) = 0$). In order to improve the efficiency of information flow while controlling the amount of computation, we introduce residual on V.

However, due to space constraints, we are only able to explain the trainability of residual on V from an experimental point of view rather than a mathematical one. The NTK spectral analysis[1,2] of the Transformer with shared attention and residual on V is exactly where we will direct our future work.

For W3, The ablation study is limited and does not thoroughly evaluate the contribution of each component.

We add ablation experiments on whether CSAU is more effective, please see the table below:

It is evident that CSAU is more efficient when fewer parameters are used. In fact, CSAU serves as an efficient approach to reducing excessive computation in the Feed Forward Network (FFN) section. It accomplishes this by relocating the gating operation before the FFN computation process, which streamlines the FFN into an identity module, resulting in faster execution.

For W4, Important implementation details like dataset splits for training, validation, and testing are not provided. This makes reproducibility difficult.

The specific training datasets, training and testing details have been provided in Appendix B. Here, we present the hyperparameters for training for the four tasks outlined in the article. This is intended to facilitate the reproduction of this work by others and to address any questions that may arise.

[1] Yang G. Tensor programs ii: Neural tangent kernel for any architecture[J]. arXiv preprint arXiv:2006.14548.

[2] Yang G, Littwin E. Tensor programs iib: Architectural universality of neural tangent kernel training dynamics[C]//International Conference on Machine Learning. PMLR, 2021: 11762-11772.

Thank you for dedicating time and effort to review our paper. Below, we address each of your points and have made corresponding adjustments to the manuscript based on your valuable suggestions. If our responses sufficiently address your queries, would you kindly consider enhancing the rating for our paper?

For W1, ShareFormer seems to be a thoughtful amalgamation of pre-existing techniques.

Thank you for your commendation on the balance achieved between quality and speed in our proposed method. Regarding the innovation of the method, we would like to provide further details on a few points.

1. About Residual on V: I would like to clarify that ELAN and Restormer did not propose this innovation. There may be some misunderstanding. Further, we will demonstrate that residual connections on the feature value facilitate the shared attention operation to counteract its inherent drawbacks, while also ensuring higher performance and no impact on inference speed. Experiments about $4 \times$ classical image super resolution were conducted with different numbers of shared layers, exploring two scenarios of residual connection on Value and other configurations. The results revealed the following findings.

As you can see, if the attention maps are only shared between nearby layers, the model's performance will significantly decrease. This demonstrates why it's important and effective to have a residual connection on value.

2. About the group split strategy: the feature splitting approach in this paper differs from EfficientViT, which we will illustrate with a piece of PyTorch-style pseudo-code.

Following the EfficientViT's official code, they add the previous group's output to the next group's input. It also performs an additional grouping convolution on the query. The group split strategy of EfficientViT is:

for i, qkv in enumerate(self.qkvs):
    if i > 0: # add the previous output to the input
        feat = feat + feats_in[i]
    feat = qkv(feat)
    q, k, v = qkv.chunk(3, dim=1)
    q = self.dws[i](q)
    attn = (q.transpose(-2, -1) @ k) * self.scale
    attn = attn.softmax(dim=-1)
    feat = v @ attn.transpose(-2, -1)
    feats_out.append(feat)
x = self.proj(torch.cat(feats_out, 1))

However, we refrained from adding the output of the former group to the latter, as we believe it is an ineffective (or detrimental) operation for the image restoration task due to the following reason. The feature groups will be inputted into window attentions of different window-sizes. The first group of features will be inputted into a window of size [32, 8] by default, while the second group of features will be inputted into a window of size [8, 32]. The windows will analyse and process texture information in different directions. According to Figure 4 in GRL [3], we could see that "image features in natural images are anisotropic." Therefore, two sets of features will be utilized for extracting heterogeneous information. Overlapping these features through addition will restrict the second set of window attention from learning useful knowledge, as its input includes texture information in the opposite direction. Here it's the group split strategy of ShareFormer:

qkvs = self.qkvs(feat).chunk(split_num, dim=1)
feats_out = []
for qkv in qkvs:
    q, k, v = qkv.chunk(3, dim=1)
    attn = (q.transpose(-2, -1) @ k) * self.scale
    attn = attn.softmax(dim=-1)
    feat = v @ attn.transpose(-2, -1)
    feats_out.append(feat)
x = self.proj(torch.cat(feats_out, 1))

We used EfficientViT's group split strategy to train ShareFormer on the classical $4 \times$ image super-resolution task, and the experimental results proved our point.

[1] Yang G. Tensor programs ii: Neural tangent kernel for any architecture[J]. arXiv preprint arXiv:2006.14548.

[2] Yang G, Littwin E. Tensor programs iib: Architectural universality of neural tangent kernel training dynamics[C]//International Conference on Machine Learning. PMLR, 2021: 11762-11772.

[3] Y. Li, et al., "Efficient and Explicit Modelling of Image Hierarchies for Image Restoration," in 2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), Vancouver, BC, Canada, 2023 pp. 18278-18289.

Fig.2 omits (a) and (b) mark, while the caption uses (a) and (b).

Thank you sincerely for your careful review of our manuscript. We have addressed the identified issue in the paper accordingly.

Restormer never used window attention.

Thank you for pointing out this issue. It was, in fact, a typo and a description that had been deprecated in the early iterations of the manuscript. It was not relevant to the content discussed in the paper. We have now removed it from the manuscript.

For W1, The presentation of text and figures in this paper should become clearer and more understandable. See Q1, Q2, and Q3.

I hope the above answers on Q1-Q3 have allayed your concerns.

For W2, The paper lacks some ablation studies proving the importance of the proposed components, such as CSAU.

Thanks to your kind reminder, we supplemented the ablation experiment with this module of CSAU on classical image $4 \times$ super-resolution task. The test dataset is Urban100. The results are as follows:

It's clear that CSAU works better and faster with fewer parameters. In fact, CSAU is an efficient way to avoid excessive computation in the Feed Forward Network (FFN) part, because it moves the gating operation from the middle of the FFN computation process to before the FFN computation process and simplifies the FFN into an identity module.

For W3, shared attention seems not novel.

The core of our approach is to implement residual connections on the Value element while completely sharing the attention feature map of the network among adjacent layers. This allows the shared attention operation to offset the shortcomings brought about by it, all the while ensuring that there isn't excessive loss of accuracy and running speed remains unaffected. We conducted experiments on varying numbers of shared layers for the two scenarios of residual connection on Value or otherwise in classical image $4 \times$ super-resolution task, and the findings are as follows (the test set is the Urban100):

As you can see, if the attention maps are only shared between nearby layers, the model's performance will significantly decrease. This demonstrates why it's important and effective to have a residual connection on value.

For W4, A potential unfair issue is observed with respect to SR. See Q4.

I hope the above answer on Q4 has solved your problem.

For W5, Despite the improved inference speed, the performance gains of grayscale denoising, color denoising, and JPEG CAR are not significant.

Your concern addresses a broader challenge inherent to the entire domain of image restoration. Allow me to shed light on a few aspects:

1. Image restoration, distinct from other high-level visual tasks, faces unique challenges in defining metrics and achieving improvements. As stated in the discussion of this review [1]: "More accurate metrics need to be found for image denoising. PSNR and SSIM are popular metrics for the task of image restoration. PSNR suffers from excessive smoothing, which is very difficult to recognize indistinguishable images. SSIM depends on brightness, contrast and structure, and therefore cannot accurately evaluate image perceptual quality." Metrics such as PSNR are not as absolute as accuracy in classification tasks or IoU in detection tasks.
2. The evaluation of image quality remains a complex task, given the absence of a perfect objective metric. Traditional measures like PSNR and SSIM oversimplify the multidimensional assessment of image quality. Consequently, the numeric variations among state-of-the-art methods, particularly those based on Transformers, become less prominent. However, as illustrated in Fig. 10 of our paper and detailed perceptual comparisons in Appendix C, our approach excels in recovering finer details, such as dense dots in real images. This showcases the visual superiority of our method over others.
3. Importantly, our emphasis lies in striking an optimal Pareto balance between performance and latency, rather than pursuing absolute performance gains. In this regard, our approach holds significant implications for the effective application of algorithms, especially large Transformer-based models, in real-life industrial settings.

[1] Chunwei Tian, Lunke Fei, Wenxian Zheng, Yong Xu, Wangmeng Zuo, Chia-Wen Lin, Deep learning on image denoising: An overview, Neural Networks, Volume 131, 2020, Pages 251-275, ISSN 0893-6080.

Thank you for your meticulous review and valuable insights into our paper. Below, we will respond to each of the questions you raised and make corresponding modifications to the manuscript. If our responses address your concerns and the revised manuscript meets your satisfaction, we kindly ask for your consideration in potentially enhancing the rating of our paper. Your feedback is immensely appreciated.

For Q1: Where is the exact part that the shared attention is operated?

We have made appropriate changes to Figures 2 and Figure 3 and their captions in the article to make them clearer.

In fact, both Attention Map and Feature Value obtained in SPSA are fed into CSAU. For easier understanding, we summarise the whole procedure as Share Attention Block (SAB).In SAB, there exist two layers, the first one contains two modules: SPSA, GDFN, and the second one contains one module, which is CSAU. Among them, SPSA is used to compute the Attention Map $Attn\_{l-1}$ and Feature $V\_{l-1}$, while GDFN is classically used as Feed Forward Network (FFN); for the $l$th layer, which is the CSAU module, it receives the output of GDFN $X\_{l}$, SPSA's attention map $Attn\_{l-1}$ and feature $V\_{l-1}$, and obtains the output according to Eq. 5 and Eq. 7 in the article.

Specifically, we can use the following formulas for further explanation:

$$Q_{l-1}, K_{l-1} = W_q(X_{l-1}), W_k(X_{l-1});$$ $$Attn_{l-1} = Softmax((Q_{l-1}K_{l-1}^T)/s);$$ $$V_{l-1} = W_v(X_l{l-1});$$ $$X_{l} = W_{l-1}(Attn_{l-1} V_{l-1}) + X_{l-1};$$

The above four lines of equations represent the specific work of the SPSA;

$$X_{l} = GDFN(X_{l});$$

The formula above represents the specific work of the GDFN;

$$G_{l}, V_{l} = W_g X_{l}, W_v{X_{l}} + V_{l-1};$$ $$Attn_l = Attn_{l-1};$$ $$X_{l+1} = W_{l}(G_{l} \odot (Attn_l V_l)) + X_l;$$

The above three lines of equations represent the specific work of the CSAU.

We hope this explanation gives you a clearer understanding of our network structure, and you can also view the updated Figure 2 and Figure 3 for a more intuitive understanding of the model's design options.

For Q2: From ConViT, how can you draw the conclusion that the concentrated ERF implies an amplified locality bias of the network?

We suspect there might be a misunderstanding in our presentation. In Section 4.2 of our paper, the statement is as follows: "As elucidated by Barzilai et al. (2023), the ensemble behavior of the network results in its ERF carrying more weight towards the central region of the receptive field." Therefore, we did not draw this conclusion from ConViT. This particular conclusion was obtained from Section 5.2 of "A kernel perspective of skip connections in convolutional networks" (Barzilai et al. 2023) [1], a result that has been rigorously proven mathematically.

For Q3, In Appendix D, why did you compare the attention maps of layers 1 and 3?

Thank you for your inquiry, which brings up a crucial point.

1. In Appendix D, we compared layers 1 and 3 because these two layers happened to exhibit an apparent redundancy, making it more beneficial to illustrate for readers' comprehension. However, we did not anticipate that this would lead to misunderstandings, and we sincerely apologize for that. To enhance clarity and aid understanding, we have now provided a complete set of attention maps visualization for all 24 layers of SwinIR-L in Appendix D. In reality, redundancy in attention computation might occur between adjacent layers (as shown in layer 4 and layer 5 of block 3 in the updated Appendix D) or even across distant layers (such as the previously compared layer 1 and layer 3 of block 1, and the layer 0 and layer 5 of block 4 in the updated Appendix D). The reasons for redundancy appearing across distant layers and how to determine which two layers exhibit significant redundancy are valuable questions that we acknowledge and plan to address in our future works.
2. The use of SwinIR for comparison was not intended to convey any specific meaning. We simply aimed to illustrate that redundancy in attention computation is prevalent in various Transformers, and SwinIR, being one of the most famous and representative Transformer networks in the field of image restoration, served as a natural example to demonstrate this phenomenon. However, we also think your suggestion is meaningful. We are currently training ShareFormer without shared attention maps and observing the computational redundancy it introduces. We plan to update these results in the Appendices within the next two days.

[1] Daniel Barzilai, Amnon Geifman, Meirav Galun, Ronen Basri. "A Kernel Perspective of Skip Connections in Convolutional Networks" ICLR 2023.