Paper: Key-Graph Transformer for Image Restoration (ICLR-2024)

Reviewer comment (Novelty):
The author should discuss with relevant methods. Chen et al. [1] proposed a top-k strategy to selectively choose the most relevant tokens for image restoration (e.g., deraining).

Author response:
**A:** Thanks for the suggestion. DRSformer [a] proposed a topK strategy that chooses the most relevant tokens to model the non-local priors for deraining. However, the KNN selection in DRSformer occurs after every standard attention operation, where the self-attention is already calculated in a standard manner with  $\mathcal{O} (hw  \times  hw)$ complexity. On the contrary, the  KNN selection of the proposed method occurs exactly before the self-attention operation of each Transformer layer guided by the constructed K-Graph. This selection ensures there are only K nodes in Key (Key here indicted the Key of Query,  Key Value im the self-attention operation) instead of original hw elements in Key, which makes the complexity of the following self-attention become $\mathcal{O} (hw \times  K)$.  We also added the corresponding discussion in the related work section (Non-local Priors Modeling in IR) of our updated manuscript. 

[a] Chen Xiang, et al. "Learning a sparse Transformer Network for Effective Image Deraining". CVPR2023

Reviewer comment (Evaluation):
The limitations of the proposed method should be discussed.

Author response:
**A:** Thanks for the kind suggestion. This study faces a task-specific limitation, i.e., each image restoration task requires training a separate network. While efforts have been made to train models for varying degradation levels within specific types, such as image denoising or removal of JPEG compression artifacts, this approach still leads to inefficiencies in model training and reduces the utility of the trained networks. A potential future enhancement involves developing a mechanism enabling a network to handle diverse image degradation types and levels. Another challenge is the substantial parameter requirement of the proposed KGT, which operates within a tens-of-millions parameter budget. Deploying such a large image IR network on handheld devices with limited resources is challenging, if not unfeasible. Therefore, a promising research direction is the creation of more efficient versions of KGT, integrating non-local context more effectively, to overcome these limitations. We also added the limitation part of the proposed KGT in the updated Supp. Mat.

Reviewer comment (Experiments):
The paper mentions that compared with traditional self-attention, the time complexity of this method significantly slows down from quadratic to linear. Therefore, it is more prudent to add a chart comparing the convergence speed of various methods to the paper.

Author response:
**A:** We provide the comparison of different neural network models regarding the parameters, runtime, and performance (measured in PSN) on the Urban100 datasets for image SR in Table 2 (Please refer to Response to Reviewer ruU5 Part2: Q3). We reported the small version (KGT-S) of our method. It shows that 1) Among the methods, HAT [d], and KGT-S achieve first-class PSNR performances, reaching 28.37dB and 28.34dB, while KGT-S is much faster and has 41.7% fewer parameters than HAT. 2) SwinIR [a] runs a bit faster than KGT-S but with a PSNR loss of 0.89 dB. Compared with ART, HAT, and HAT-S, the proposed KGT-S is faster and more accurate.

Additionally, reporting the convergence speed of other comparison methods requires retraining all these methods from scratch. This is not easy and not feasible to do so.  For the comparison, the quantitative results of these methods are directly adopted from their papers, and the visual results of these methods are from their provided checkpoints or directly the provided visual results. As a remedy, we report the training log of the proposed  KGT with two versions (i.e., the small version  KGT-S and the base version KGT-B). It shows that the proposed network converges gradually during the training. We also include this in our Supp. Mat. (Section 3: More Abation Analyses -> Convergence Visualization, Figure 8).

[a] Liang Jingyuyn, et al., "Swinir: Image restoration using swin transformer". ICCVWorkshop2021

[b] Zhang Jiale, et al., "Accurate image restoration with attention retractable transformer". ICLR2023

[c] Chen Zheng, et al., "Cross Aggregation Transformer for Image Restoration". NeurIPS2022

[d] Chen Xiangyu, et al., "Activating More Pixels in Image Super-Resolution Transformer". CVPR2023

Reviewer comment (Writing):
I am uncertain about the usage of "KK". Upon examining the model diagram, I only observe the notation "K" being used, which raises doubts about the origin or meaning of "KK".

Author response:
**A:** We appreciate your diligence in examining our model diagram, and apologize for any confusion caused by the typo "KK". It should be "K". We have rectified this typo in the updated main manuscript for a more accurate representation. Thank you for bringing this to our attention.

Reviewer comment (Reproducibility):
Please provide a detailed description of the model architecture, including the number of blocks in each layer.

Author response:
**A:** We have two base model architectures, the first one (Archi-V1)  is depicted in Fig.1 of our main manuscript which is for image super-resolution while another U-shaped hierarchical architecture  (Archi-V2) is for other IR tasks i.e., image denoising, single image deblurring, removal of JPEG compression artifacts, image demosaicking, and image restoration in adverse weather conditions. The detailed structure of the U-shaped architecture is also shown in our updated Supp. Mat. The detailed number of the KGT stages and the Transformer layers within each KGT stage for both architectures are listed in the following table (Table 3).


Table 3. The details of the KGT stages and KGT layers per stage for both architectures.
|     | Archi-V1 |      | Archi-V2       |        |                |
|----------------------|----------------|--------|----------------|--------|----------------|
|                      | KGT-small      | KGT-base | Down Stages   | Up Stages | Final Stage  |
| Num. of KGT Stages   | 6              | 8        | 4            | 4        | 1            |
| Num. of KGT layer/stage | 6            | 8        | 6            | 6        | 6            |