Sharing Key Semantics in Transformer Makes Efficient Image Restoration

Response to Reviewer YUtD:

Q1: Comparison to KiT[1]

A: Please refer to our answer to the 1st question in the shared "Author Rebuttal".

[1] KNN Local Attention for Image Restoration. CVPR'22

Q2: Fig.2 needs to improve regarding the dimension

A: Thanks for this very important suggestion, we have improved Fig.2 (d) and made it much clearer to follow. The revised figure is shown in Fig 5 of the rebuttal pdf file. We also demonstrated the difference in the calculation of the attention between training (Torch-Mask is used to exclude the possible negative connection below the threshold ) and inference (Triton kernel is design and used to greatly reduce the computation cost).

Q3: Top-k selection discussion?

A: Please see the answer to the 2nd question in the shared "Author Rebuttal".

Q4: Comparison to [2] and [3]?

A: We compare the deraining results among AirNet[2] and PromptIR[3] under the single deraining degradation setting. The results are reported in Tab. 7. It shows that:

• Our SemanIR outperforms AirNet by a large margin (i.e., 2.93 dB) on PSNR. The reason why the parameters of AirNet are only 8.93M is that it is built on CNNs.
• Our SemanIR outperforms PromptIR by 0.79dB on PSNR with 27% fewer parameters even though both SemanIR and PromptIR are built with transformers.

Table 7 The single-degradation deraining comparison on the Rain100L dataset.

[2] All-in-One Image Restoration for Unknown Corruption. CVPR'22

[3] PromptIR: Prompting for All-in-One Blind Image Restoration. NIPS'23

Dear Reviewer_YUtD,

We sincerely appreciate your feedback and the positive score of our manuscript, and we are committed to further enhancing the quality of our manuscript based on your suggestions.

In the following, first, we aim to address the concerns regarding the novelty of SemanIR in greater detail. Subsequently, we will provide a more reliable and thorough comparison of SemanIR with the approaches presented in [2] and [3].

Concerns regarding the novelty of SemanIR:

Answer: Although both KiT and SemanIR use KNN search, they are fundamentally different in the following core aspects:

• Design logic: KiT tries to extend the attention field from a local patch to $k$ patches, which follows the local-to-global logic. On the other hand, SemanIR aims to sort out the most similar tokens for a token in the global range efficiently, which in essence is a global method.
• Patch-wise vs. token-wise similarity: Due to the different design logic, KiT computes the similarity between （$r \times r$） patches, then attention is conducted between tokens in $k$ ($r\times r$) similar patches. By contrast, SemanIR directly computes the similarity between tokens, and the attention is done directly between the similar tokens.
• Implementation and efficiency: KiT computes KNN search in each transformer layer and for the sake of efficiency, locality-sensitive hashing is used. On the other hand, SemanIR computes the KNN directly for each token, and a single KNN is shared across all transformer layers in the same stage to improve computational efficiency.

They are all beyond the key-semantic dictionary-sharing strategy. Additionally, we would like further to highlight the following contributions of the proposed SemanIR:

• Attention Calculation: Our approach utilizes the Torch-Mask function during training and the Triton kernel operator during inference. This trade-off ensures accurate backpropagation with the semantically most relevant patches during training and reduces the shape of $K$ and $V$ from $HW \times C$ to $k \times c$ during inference. This optimization significantly reduces inference time, as shown in Table 2, making it practical for deployment.

• Top-k Training Strategy: As detailed in our response to the second question of the "Author Rebuttal" and Appendix D, we decouple the use of $k$ between training and inference. This allows for flexible, randomly sampled $k$ during training and fixed $k$ during inference, enabling effective use of large GPU memory during training while ensuring efficient performance with limited GPU resources during inference. This strategy is applicable to large-scale IR models.

• Extensive Experiments: We have validated our method on 6 diverse IR tasks, achieving state-of-the-art performance on most. The visual comparisons in the appendix further demonstrate the effectiveness of SemanIR across various degradation types.

These aspects, in addition to the dictionary-sharing concept, collectively enhance the originality and novelty of SemanIR. We believe these clarifications strengthen the presentation of SemanIR's contributions. Thank you for the opportunity to address these points, and we hope this clarification meets your expectations.

Response to Reviewer aTgi:

Q1: Top-k selection discussion?

A: See the answer to the 2nd question in the shared "Author Rebuttal".

Q2: Visualization:

A: We set a query region in the input and provided a detailed comparison from the attention-based activation map together with the input query region in Fig. 3 of the rebuttal pdf file. Fig. 3(a) shows the query region input. Fig. 3(b) displays the activation map generated using standard attention mechanisms. Fig. 3(c-f) illustrate activation maps using our key-semantic dictionary with different top-k values ([8, 16, 64, 256]) during inference. To conclude:

• Normal Attention: The activation map from normal attention (Fig. 3(b)) shows connections to regions that may not be semantically related to the query region. This indicates that normal attention can consider irrelevant regions.
• SemanIR with Key-Semantic Dictionary: With the proposed SemanIR, using a smaller top-k value (e.g., top-k = 8 or 16) results in activation maps that connect only a limited number of neighboring patches with the closest semantic similarity. This is consistent with our intended concept, as demonstrated in Fig.1(e) of the original manuscript.
• Effect of Top-K Value: Increasing the top-k value allows for connections to more semantically related regions. However, when the top-k value is set too high (e.g., top-k = 256 as shown in Fig. 3(f)), the activation map may include some semantically unrelated regions. This aligns with the findings and the results depicted in Fig. 3 of our manuscript, where increasing the top-k value beyond a certain point (e.g., from 396 to 512) does not further improve PSNR.

Q3: Non-sharing of key-semantic dictionary?

A: We discuss this issue from the following perspectives:

• Potential Benefits: Constructing a key-semantic dictionary for each layer could indeed enhance performance, as each layer would benefit from a dictionary specifically tailored to its unique context. This approach might allow each layer to utilize a more precisely matched dictionary, potentially improving the semantic relevance and accuracy of the attention process.
• Experimental Constraints: For the sake of training efficiency, the Torch-Mask strategy is used. For a window-size of 32 with key-semantic dictionary sharing, this already leads to an explosion of memory. Creating layer-wise key semantics would further increase the memory footprint significantly.
• Empirical Evidence: As an alternative, we have visualized the attention maps for each layer within the same stage in Fig. 4 of the rebuttal pdf file. It shows: (i) The attention maps exhibit only slight variations across layers, indicating that the semantic focus remains largely consistent. (ii) The activation regions in these maps are very similar, which supports the effectiveness of our approach of sharing the key-semantic dictionary across layers.

Q4: Complexity of key-semantic dictionary:

A: We acknowledge the mistake and appreciate the opportunity to clarify.

• Correction of Complexity: The correct complexity for the similarity calculation process is indeed $\mathcal{O}((HW)^{2}C)$, rather than $\mathcal{O}(HWC)$. We apologize for this error.
• Revised Eq. 5:

$

\begin{aligned} \mathcal{O}(6 \times [4HWC^{2} + 2kHWC] + (HW)^{2}C) \end{aligned}

$

• Revised Eq. 6:

$

\begin{aligned} \mathcal{O}(6 \times [4HWC^{2} + 2(M)^{2}HWC] - (6 \times [4HWC^{2} + 2kHWC] + (HW)^{2}C)) \ = \mathcal{O}((12M^{2} - 12k - HW)HWC) \end{aligned}

$

• Example Calculation: Let’s consider a common setting as we indicate in our Appendix($M = 7$, patch size = 16, $H$ = $W$ = 64, k=512). We have:

$

\begin{aligned} \mathcal{O}((12M^{2} - 12k - HW)HWC) &= \mathcal{O}(12 \times (7 \times 7) \times(16 \times 16) - 12 \times 512 - 64 \times 64) \&= \mathcal{O}(150528 - 6144 - 4096) >> 0 \end{aligned}

$

Despite the errors, the conclusions about the complexity remain valid. We appreciate your understanding and the opportunity to correct these details.

Q5: Why store the similarity, not index?

A: In Eq 2 of our manuscript, we calculate the similarity values to construct the key-semantic dictionary but store only the indices. We have clarified this in the revised manuscript to prevent any misunderstandings.

Q6: Comparison to NAFNet[1]

We have included a performance comparison with [1] in Tab. 6:

Table 6: The deblurring comparison on GoPro.

It shows that SemanIR outperforms the strong baseline [1] in both PSNR and SSIM on the single-image motion deblurring task. We have included this important comparison in our revised manuscript.

[1] Simple Baselines for Image Restoration. ECCV'22

Q7: Positional Embedding (PE)?

A:

• For $\mathcal{D}_{K}$, the PE is not used but it is interesting to explore.
• For attention, we adopted the relative PE. The full PE for all locations is first indexed from the relative positional encoding. During training, the positional encoding and corresponding value in the attention map of dissimilar tokens is masked by being set to infinity.

Q8: Windowed-attention?

A:

• For all experiments, we set window size = 32, which contains 1024 tokens. This are lots of tokens representing non-local connectivity and this size is already larger than the global range for tasks like classification.
• The window-based calculation is chosen for its efficiency and reduced memory usage, as the masking strategy results in high memory consumption during training, making it costly for IR with vanilla ViT.
• Applying the proposed SemanIR on vanilla ViT architecture would be a very interesting direction with proper design, which we would like to try in our future work for a more generalized exploration.

Response to Reviewer hRwb:

Q1: Have the authors considered window-size ablations?

A: The windows indeed contain mixed information from different semantic parts. Yet, it is precisely this semantic distinction that motivates us to develop a selection mechanism for semantic information using KNN. We have conducted ablation studies of the window size (i.e., on both gray and color image denoising with $\sigma=25$). The results are summarized in Tab.4 and Tab.5. With the increase of the window size, the semantic relevant information for each token is increased, thus leading to a PSNR gain for different IR tasks.

Table 4: Gray DN on Set12, BSD68, and Urban100.

• For gray DN (Tab. 4), performance improves with larger window sizes, showing increased PSNR values across datasets.

Table 5: Color DN on Mcmaster, Cbsd68, Kodak24, and Urban100.

• Similarly, for color DN (Tab. 5), larger window sizes yield better results. For instance, in the McMaster dataset, the PSNR increases from 33.20 dB with a window size of 8 to 33.38 dB with a window size of 32.

These findings indicate that larger window sizes enhance performance by capturing more contextual information, which will be discussed in the revised manuscript.

Q2: Discussing how the proposed SemanIR compares to [1][2][3]

A:

• SemanIR vs. [1]: According to Tab. 9 in [1], [1] achieves 32.51 dB, 28.82 dB, and 27.72 dB PSNR on Set5, Set14, and BSD100 datasets, respectively. In comparison, SemanIR reaches 33.08 dB, 29.34 dB, and 27.98 dB PSNR on the same datasets. Both methods are designed for low-parameter settings. A detailed parameter comparison will be included in the revised manuscript. Note that [1] is not specifically optimized for other IR tasks.

• SemanIR vs. [2]: Unlike SemanIR, [2] reuses earlier attention in subsequent layers, suggesting potential complementary benefits. [2] achieves an 11.3% FLOP reduction with performance comparable to the baseline, while SemanIR reduces FLOPs by 12.3% and improves PSNR by 0.39dB. We acknowledge the fairness of these comparisons and will conduct a similar evaluation in our IR settings, expecting similar results.

• SemanIR vs. [3]: Please see the answer to the 1st question in the shared "Author Rebuttal", where we show that our SemanIR is consistently better than DRSformer on both deblurring and deraining.

[1] CAMixerSR. CVPR'24

[2] Skip-Attention. ICLR'24

[3] DRSFormer. CVPR'23

Q3: While maintaining the key-semantic dictionary is interesting, the idea of sharing information across Transformer layers has been explored previously [4].

Answer: Thank you for your feedback. We would like to address the points raised as follows:

• Timing of Related Work: We note that [4] was made available on arXiv on June 1st, 2024, which is after our manuscript submission to NeurIPS. Therefore, our work was not influenced by this recent publication. Nevertheless, we will include the following discussions in the revised manuscript.

• Focus of Our Work: While [4] focuses on reducing computational costs in Vision Transformers (ViTs), our proposed SemanIR method is distinct in its primary application. SemanIR is designed specifically for image restoration, rather than the broader focus of LaViT on computational efficiency within ViTs. Although both works address computational efficiency, the objectives and methodologies are different.

• Difference in Computational Efficiency Approaches: LaViT [4] reduces computational costs by storing attention scores from a few initial layers and reusing them in subsequent layers. However, this approach does not change the computation cost of attention itself; it merely reuses previously computed scores. In contrast, SemanIR introduces a novel approach for reducing computation during both training and inference. During training, SemanIR leverages a pre-constructed semantic dictionary to exclude irrelevant information from other semantically unrelated patches, thus enhancing restoration quality. During inference, our implementation with Triton kernels optimizes attention operations, directly reducing computational costs.

[4] You Only Need Less Attention at Each Stage in Vision Transformers. CVPR'24

Q4: Sharing the dictionary across different stages?

Answer: Thank you for your valuable suggestion regarding dictionary sharing across different stages of our columnar architecture. We appreciate the insight and acknowledge that this is a compelling idea worth exploring further. To address your suggestion, we conducted an analysis to visualize the self-similarity of features at the beginning of each stage in our model, which consists of a total of six stages. The visualization results are shown in Fig.1 of the rebuttal pdf file. Based on our observations:

• Stage-wise Semantic Similarity: We noted that there are still significant differences in the semantic similarity maps across various stages. This suggests that sharing the dictionary from the early stages to the later stages could potentially lead to a performance drop due to the divergence in feature representations.

• Adjacency-based Sharing: Despite the variability across stages, we observed that adjacent stages exhibit similar semantic similarities. This indicates that it might be feasible to share the dictionary every two or three stages. Such an approach could reduce computational costs while maintaining performance.

Given these findings, we recognize the potential benefits of exploring dictionary sharing. We plan to investigate this perspective in more detail in our future work to assess its impact on performance and efficiency.

I have gone through the authors' response, both to my questions, and to other reviewers'. I thank the authors for responding to the comments, and for providing dictionary-sharing visualizations in the attached pdf. Further, authors have provided comparisons with other methods focusing on computational efficiency with respect to the attention mechanism, including the run analysis with DRSFormer. The proposed method, SemanIR, either scores higher in tasks, or is faster, or both. I am satisfied with the response, and do not have any follow up questions. Therefore, I would like to raise my score to accept.

Response to Reviewer xhpu:

Q1: What's the difference between SemanIR and other methods like KiT or DRSformer?

A: Please refer to our 1st answer in the shared "Author Rebuttal".

Q2: What is the difference between SemanIR and token merging and pruning methods?

A: The key focus of our method significantly diverges from previous token merging or pruning techniques, aiming not only for efficiency but, more crucially, for effective performance in regression-based dense prediction image restoration tasks. While improving efficiency aligns with the goals of token merging methods, the proposed SemanIR is specifically tailored for image restoration within a regression-based paradigm, as mentioned in Line 92 of our original manuscript. For image restoration, every token encodes information about a local patch. Merging or pruning tokens will lead to a loss of information in corresponding patches, which is not preferred in image restoration [4]. Our method’s suitability for image restoration can be attributed to the following points:

• Utilization of Semantic Information: Unlike token merging or pruning methods [1, 2, 3], which reduce the number of tokens by combining or removing redundant ones with higher semantic similarity, SemanIR leverages these semantic-close tokens to enhance restoration. By ensuring that each token can benefit from others, our approach improves the overall quality of restoration.

• Preservation of Detail: Token merging and pruning methods may reduce computational costs but can lead to the loss of crucial detailed information, such as texture, color, and local structures. In contrast, SemanIR maintains detailed information by focusing on semantic relevance rather than merely reducing token counts.

• Effective KNN Strategy: In SemanIR, for each degraded pixel, we select its top-k most semantically close neighbors to contribute to the restoration of that pixel. This KNN strategy not only excludes the negative contributions from semantically unrelated neighbors but also enhances efficiency, distinguishing it from traditional approaches.

To conclude, although token merging and pruning techniques, such as those in [1, 2, 3], are effective for reducing computational complexity and are often suited for classification tasks [3], they are less appropriate for dense prediction tasks like IR. SemanIR’s approach addresses the specific needs of image restoration tasks more effectively by balancing efficiency with the preservation of detailed image information.

[1] Token merging: Your ViT but faster. ICLR'23

[2] DynamicViT: Efficient vision transformers with dynamic token sparsification. NeurIPS'21

[3] A-ViT: Adaptive tokens for efficient vision transformer. CVPR'22

[4] Skip-Attention: Improving Vision Transformers by Paying Less Attention. ICLR'24

Q3: Notation correction:

A: Thank you for pointing this out. We have revised Fig. (d) in the PDF file (Fig. 5) to include additional details, making the figure easier to follow and understand.