Title: Sharing Key Semantics in Transformer Makes Efficient Image Restoration
Venue: NeurIPS 2024
Accepted: yes
Keywords: Low-level Vision

Image Restoration

Vision Transformer
Abstract: Image Restoration (IR), a classic low-level vision task, has witnessed significant advancements through deep models that effectively model global information. Notably, the emergence of Vision Transformers (ViTs) has further propelled these advancements. When computing, the self-attention mechanism, a cornerstone of ViTs, tends to encompass all global cues, even those from semantically unrelated objects or regions. This inclusivity introduces computational inefficiencies, particularly noticeable with high input resolution, as it requires processing irrelevant information, thereby impeding efficiency. Additionally, for IR, it is commonly noted that small segments of a degraded image, particularly those closely aligned semantically, provide particularly relevant information to aid in the restoration process, as they contribute essential contextual cues crucial for accurate reconstruction. To address these challenges, we propose boosting IR's performance by sharing the key semantics via Transformer for IR (i.e., SemanIR) in this paper. Specifically, SemanIR initially constructs a sparse yet comprehensive key-semantic dictionary within each transformer stage by establishing essential semantic connections for every degraded patch. Subsequently, this dictionary is shared across all subsequent transformer blocks within the same stage. This strategy optimizes attention calculation within each block by focusing exclusively on semantically related components stored in the key-semantic dictionary. As a result, attention calculation achieves linear computational complexity within each window. Extensive experiments across 6 IR tasks confirm the proposed SemanIR's state-of-the-art performance, quantitatively and qualitatively showcasing advancements. The visual results, code, and trained models are available at: https://github.com/Amazingren/SemanIR.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a dictionary-based image restoration method that leverages the most relevant information to recover images with low computational costs. Specifically, the method constructs a key-semantic dictionary that stores the top-k semantically related regions for each patch and performs attention only among these related regions. Additionally, the dictionary is created only once at the beginning of each transformer stage to reduce the computational burden.\n\nSTRENGTHS:\n1. This paper presents an efficient and effective strategy for utilizing the most relevant patches for image restoration.\n2. Extensive experiments demonstrate the effectiveness of the proposed SemanIR method across several image restoration tasks.\n\nWEAKNESSES:\n1. The novelty of this paper is somewhat limited. Similar to KiT [1], the proposed method SemanIR utilizes the most relevant information for restoration by performing attention only among the related regions. The key difference between KiT and SemanIR is that SemanIR calculates the semantic relation only once and reuses the information in subsequent layers.\n2. Fig. 2 needs further improvement to increase its readability. The selected $\\hat{K},\\hat{V}$ are not shown in the calculation process of (d). Additionally, the dimensions of $\\hat{K},\\hat{V}$ are unclear in the statement. Suppose the dimension of $\\hat{K}$ is $k\\times c$, how could the dimension of $D^{att}_{K}=\\text{Softmax}_K(Q\\hat{K}^T/\\sqrt{d})$ be $hw\\times hw$ ?\n3. The intention behand constructing the random top-k strategy is unclear. Compared to the fixed top-k strategy, what are the benefits of a random $k$? It is apparent that the fixed top-k strategy would have inferior performance since the $k=[64,128,192,256,384]$ is different from the training setting $k=512$. As shown in Fig. 3, $k=256$ seems to be a suitable value for the proposed method. What is the need to increase $k$ to 512? \n4. In the IR in AWC task, some related works [2,3] are not compared.\n\n[1] KNN Local Attention for Image Restoration. CVPR'22\n\n[2] All-in-One Image Restoration for Unknown Corruption. CVPR'22\n\n[3] PromptIR: Prompting for All-in-One Blind Image Restoration. NIPS'23\n\nQUESTIONS:\nSee Weaknesses.\n\nLIMITATIONS:\nThe limitations and broader impact have been discussed in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper propose SemanIR, a novel Transformer architecture for image restoration. The paper propose a new attention mechanism for better efficiency and efficacy, based on that within a degraded image, patches semantically close to the target patch to restore provide major information in the restoration process. They build a key-semantic dictionary, which stores the top-k closest patches, for each patch. This dictionary is shared across different transformer layers. This way, they successfully constrain the cost of attention operation with preservation of its meaningfulness and with global receptive fields. Their approach reach state-of-the-art performance with great efficiency on various image restoration tasks and benchmark datasets.\n\nSTRENGTHS:\n1. The idea of using top-k semantically close patches (tokens) only sounds novel and smart, adaptively saving computational costs depending on the user's choice, with minimal negative effect on performance.\n\n2. The proposed method seems to be generally applicable, with potential to be adopted in tasks other than image restoration.\n\n3. The strong experimental results also confirm its effectiveness.\n\nWEAKNESSES:\n1. Overall, the experiments have been conducted extensively, on a variety of restoration tasks / benchmarks.\nYet, I think there could have been more experiments on the proposed method itself, than comparison to existing works.\nSee the Questions section (#1, #2, #3) below..\n\nAdditional comments:\n\n- I think the notation D is overused. In Eq. 1, it denotes the softmax output, which denotes the dictionary after a few lines (i.e., Eq.2). Although the subscript is slightly different, I found it confusing to read, as D_{i,j} and D(i,j) denotes completely different two things.\n\n- For future studies, open-sourcing the implementation is recommended.\n\nQUESTIONS:\n1. Experiments on top-k selection\n\nThe experiment of fixed top-k  in Figure 3 seems to be conducted with k=512 during training, and k is changed only during training. What happens when using a different value for k during training, matching the k value with the k to use in inference? E.g., k=64 in both training & inference, k=256 in both training / inference. I wonder this since if k is fixed to k=64 both in training / inference without notable performance drop, why would one need to use a random top-k approach, when one could simply use a small k value in training, and opt to use a larger value? In addition, I wonder how the performance would differ when the k value in inference time is set to a value unused in the training process, such as 16 or 32.\n\n2. Visualization of attention maps / key-semantic dictionary?\n\nAs semantically close patches are used only in the attention process, I wonder how the visualized attention maps would look like.\n- Do the attention maps have a smooth attention across the top-k patches, or does it still rely heavily on a few nearest neighbors?\n- Could the authors provide an example visualization of top-k patch matches in the construction process? Does the visualized attention map and top-k match look as intended, similar to the concept visualization at Fig.1(e)?\n\n3. Non-sharing of key-semantic dictionary?\n\nKey-semantic dictionary is shared across multiple layers for efficiency. How does it change when it the dictionary is constructed after every layer? Ignoring the computational cost, does it lead to consistent improvement in performance?\n\n4. Complexity of key-semantic dictionary\n\nTo construct a key-semantic dictionary, the N = H x W tokens initially has to compute similarities against N tokens, Thus, I guess the complexity of the dictionary construction process to be of O(N^2) cost, which would still have quite a burdensome complexity. However, according to the Appendix, the complexity is O(HWC). Can the authors give further explanations on this?\n\n- typo: the K in the second row of Eq.6 seems to be small k, not a capital K.\n\n5. Why store the similarity, instead of index?\n\nIn construction of the key-semantic dictionary, it seems like the dictionary stores the similarity values of the top-k patches (Eq. 2). Is there a reason for storing the similarities values instead of storing the indices of top-k patches only?\n\n6.  NAFNet [11] baseline?\n\nNAFNet is a very strong baseline well-known in image restoration literature. Is there any reason there is no performance comparison to it?\n\n7. How are positional information of tokens (patches) handled?\n\nDoes the proposed attention mechanism consider positional information (either absolute or relative), or is it just the local features only that is used? I wonder how the positional information is handled both in key-semantic dictionary construction and attention layers. Are they just simply neglected?\n\n8. Windowed-attention?\n\nThe proposed method use an efficient attention mechanism while maintaining the global receptive field.\nBut according to the code from the supplementary material, it seems like SemanIR use a window-based approach from Swin-Transformers. Doesn't this limit the global receptive field and make it local, contradictory to the mentioned benefits of the proposed attention mechanism? I believe the proposed method has the potential to generalized application, even on vanilla Transformer architecture. Have the authors tried removal of windowed-attention, and applied it on vanilla Transformer architecture?\n\nLIMITATIONS:\nThe limitations and potential impacts have been discussed appropriately.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes an efficiency-first modification to the self-attention mechanism in ViTs. The main premise of the work is to construct a key-semantic dictionary which relates each key to its semantically-relevant patches, and then share the dictionary across Transformer blocks in the same stage for computational efficiency.\n\nSTRENGTHS:\n**S1.** The proposed method is fairly straightforward, and offers a decent alternative to dense self-attention for image restoration tasks.\n\n**S2.** The proposition to share a dictionary, once constructed, across Transformer blocks allows for benefits in terms of FLOPs (G), and parameters.\n\n**S3.** Extensive experiments are conducted on several image restoration tasks.\n\nWEAKNESSES:\n**W1.** The proposed Key-Semantic dictionary computes the dot product to measure semantic similarity on windowed patches. While diagrammatically it is easier to illustrate, at feature level, due to mixing, etc., these windows might not necessarily contain information specific to one region. Have the authors considered window-size ablations?\n\n**W2.** There are several methods that focus on improving on efficiency of self-attention in context of image restoration. However, there is no comparison with different attention methods discussing how the propose method compares either in distortion metrics or computational performance [1], [2], [3].\n\n**W3.** The main computational efficiency is observed due to sharing the key information across Transformer layers. While maintaining the key-semantic dictionary is interesting, the idea of sharing information across Transformer layers has been explored previously [4].\n\n[1] CAMixerSR: Only Details Need More “Attention”\n\n[2] Skip-Attention: Improving Vision Transformers by Paying Less Attention\n\n[3] Learning A Sparse Transformer Network for Effective Image Deraining\n\n[4] You Only Need Less Attention at Each Stage in Vision Transformers\n\nQUESTIONS:\n**Q1.** For the columnar architecture style, have the authors considered sharing the dictionary across different stages? Specifically in later stages, degraded information has mostly been recovered, so sharing across Transformer stages might be reasonable.\n\nLIMITATIONS:\nThe paper address the limitations, and the societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nUnlike traditional transformers where the multi-head self-attention layer calculates the correlation between one patch and all patches, the method proposed by the authors computes the correlation among the top k semantically similar patches, allowing image restoration with lower computational cost. \n\nAdditionally, by generating the key semantic dictionary only once at the beginning and sharing it across all transformer layers, the computational burden is significantly reduced.\n\nSTRENGTHS:\nThe authors logically explain the proposed method. They provide detailed information on the KNN algorithm and Key-Semantic Dictionary Construction and clearly illustrate how it is utilized through text and figures.\n\nThey also demonstrate the performance of the proposed method through numerous experiments. They conducted various experiments on hyperparameters and quantitatively measured the efficiency of the proposed method by assessing FLOPS, the number of parameters, and runtime.\n\nWEAKNESSES:\nThe explanation of the difference between the KNN matching method used in KiT and DRSformer and the KNN method proposed by the authors is lacking.\n\nThe authors mention that their method differs from previous token merging or pruning methods, but there are no comparative results to show which method reduces computational cost more effectively.\n\n In Figure 2 (d) Key-Semantic Attention should be corrected from \"topk\" to \"top-k\".\n\nQUESTIONS:\nWhat is the specific difference between the KNN matching method used in KiT, DRSformer, and the method proposed by the authors?\n\nHow does the performance of the proposed method compare to previous token merging or pruning methods? If previous methods yield better results, why did the authors choose to use the KNN method?\n\nLIMITATIONS:\nThe authors clearly outline the limitations of their research, informing the readers.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}