Hierarchical Information Flow for Generalized Efficient Image Restoration

Table 3: Model scaling-up exploration.

Table 5: Dot production attention vs. cosine similarity attention for model scaling. PSNR reported for SR.

Q1: Comparison with ShuffleFormer and Shuffle Transformer.

Ans: Thanks a lot for the comments. We compare with Random shuffle transformer [1] and Shuffle transformer [2]. The following comparson is added to Appendix D in the revised paper. Both methods use spatial shuffle operations to facilitate non-local information exchange, with one being random and the other deterministic.

1. Random Shuffle Transformer (ShuffleFormer) [1] applies random shuffling on the spatial dimension, which increases the probability of global information existing within a local window. While this operation extends the receptive field globally in a single step, it compromises the relevance of pixels within the window. In contrast, the hierarchical information flow proposed in this paper progressively propagates information from local to global while preserving the relevance of attended pixels. A comparison with ShuffleFormer on image deblurring is presented in the following table. Hi-IR outperforms ShuffleFormer by a significant margin while using 55.5% fewer parameters. This demonstrates the effectiveness of the hierarchical information flow method introduced in this work. The comparson is added to Table 11 of the main paper.

2. Shuffle Transformer employs a spatial shuffle operation to aggregate information from distant pixels or tokens. However, it differs from the proposed Hi-IR in several key aspects. First, Shuffle Transformer does not enable progressive information propagation within a hierarchical tree structure. Second, its shuffle operation is based on a fixed grid size of $g = 8$. The distance between pixels in the shuffled window is $H/g$ and $W/g$ along the two axes, which directly depends on the image size. For large images (e.g., 1024 pixels), this design forces distant pixels to attend to one another, often introducing irrelevant information. Consequently, this operation is unsuitable for image restoration tasks, where image sizes can become extremely large. In contrast, the L2 information flow attention proposed in this paper limits the maximum patch size, thereby constraining the maximum distance between pixels at this stage. This restriction enhances the relevance of pixel interactions, making it more effective for image restoration tasks.

[1] Xiao, Jie, et al. "Random shuffle transformer for image restoration." International Conference on Machine Learning. PMLR, 2023.

[2] Huang, Zilong, et al. "Shuffle transformer: Rethinking spatial shuffle for vision transformer." arXiv preprint arXiv:2106.03650 (2021).

Q2: Typos

Ans: Thanks a lot for pointing out the typos. We have corrected the typos in the revised version.

Q3: Different datasets and scale factors in Tables 3, 4, 5, 6, and 7

Ans: Thank you for highlighting the uniformity of the datasets and scaling factors in the ablation study. We provide the following clarifications.

1. Different test sets in Tables 3, 4, 5, 6, and 7: The ablation study was conducted on super-resolution (SR) tasks, utilizing five test sets: Set5, Set14, BSD100, Urban100, and Manga109. The model was trained once and evaluated on all datasets. However, due to space constraints, we could not display all test results. Notably, the performance gap between different realizations or training methods is more pronounced on Urban100 and Manga109 compared to Set5, Set14, and BSD100. This discrepancy reflects the characteristics of the datasets rather than the methods themselves, as the results are consistent across datasets. To enhance the diversity of reported results, we primarily presented results on the representative dataset Set5 in Tables 4 and 6, Urban100 in Table 7, and across all test sets in Tables 3 and 5.

2. Different scaling factors in Table 3 and Table 5 We indeed conducted experiments for all scaling factors in both ablation studies. However, due to space constraints in the main paper, some results were omitted. The amended results are presented below. We also append the full table in Appendix B.

Thank you for your detailed responses. The novelty of the proposed information flow is relatively limited, which I believe falls short of the acceptance standards of ICLR. Moreover, the proposed information flow does not demonstrate a clear connection to being "Generalized and Efficient," which I believe is an overstatement. I have decided to maintain my rating.

Q1: The details of the permutation operation

Ans: We explained the permutation operation in detail in the response to the common question Q1. The paper is also revised accordingly.

Q2: Technical novelty

Ans: Our response to this question is detailed in the response to the common question Q2. In short, the proposed hierarchical information flow is different from the previous method in the following aspects.

1. It propagates information from local to global with controlled computational complexity. Whereas, window attention is mainly conducted for local patches, e.g. $8\times 8$.
2. L2 information flow attention propagates information at an elevated level while enhancing pixel relevance by constraining the size of the enlarged region. This capability is not achieved by previous methods like ShuffleFormer and Shuffle Transformer.
3. We propose three strategies to systematically guide model scaling-up including removing heavyweight $3\times 3$ convolution, warmup, and using dot product attention. Theoretical analysis of why these strategies work is added to Appendix B.

Q3: Complexity analysis of the proposed hierarchical information flow

Ans: Thanks a lot for the suggestion. We conducted several analyses on the complexity of the proposed hierarchical information flow in Section 5.1.

1. We explored two approaches to implementing L1 and L2 information flow in Transformer layers: alternating them across consecutive layers or integrating both within a single layer. Our findings indicate that integrating L1 and L2 information flow within the same layer enhances performance. The designed Version 3 unifies L1/L2 information flows conceptually into a single flow with an expanded receptive field.
2. We examined the impact of tree structure depth, investigating depths of 2, 3, and 4. The results show that increasing the depth improves performance but at the cost of higher computational complexity. Balancing efficiency and accuracy, we selected a tree structure with a depth of 3.
3. In addition, we also scale up the model with over 100M parameters. This additional experiment validates the potential to further increase the model size.

Q4: Minor points

Ans: All the minor points mentioned by the reviewer are addressed in the revised version of this paper.

Q5: Further scaling up model to 100M parameters.

Ans: Thanks for the comments. The embedding dimension was increased to 256, and the number of Hi-IR transformer stages was set to 12, with 12 transformer layers in each stage. The model was trained for 200k iterations. The experimental results are presented below. While the network has not yet reached full convergence, the early results indicate that training progresses well as the model is further scaled up. Consistent improvements are also observed for the larger model.

Q6: Colors in Figure 2

Ans: The different colors represent local information. The blending of colors at higher levels in the figure indicates that information gradually propagates beyond the local patch. This information is updated in the caption of this figure.

Q1: Novelty of the paper

Ans: Thanks a lot for the comment. We address the reviewer's concern in the response to the common question Q2. We highlight several points in the following. First, although we used columnar architecture for SR and UNet architecture for the other tasks, we did not claim that the U-Net architecture is a contribution of this paper. Those are standard choices following the literature. Second, the hierarchical information flow proposed in this paper can be regarded as an efficient self-attention mechanism. The new mechanism has the following advantages: 1) propagating information from local to global in a progressive and efficient way; 2) by constraining the practical patch size in each level, the proposed method promotes the relevance of pixels in the same level, which can be a problem for ShuffleFormer and Shuffle Transformer. The detailed comparison with ShufflerFormer and Shuffle Transformer is done as suggested by Reviewer 3HDo. Third, the main contribution of Uformer is validating the UNet architecture for Transformers while the main contribution of Restormer is the self-attention along the channel dimension. Both of them are different from the proposed hierarchical information flow in this paper.

Q2: Claim of generalization and efficiency

Ans: Thank you for the valuable feedback. Our original intention was to convey that generalizing to different IR tasks requires careful consideration of the unique properties of each task. Simply combining computational mechanisms designed for different IR tasks does not necessarily result in an efficient solution. We have revised the statements in the paper to better reflect this perspective.

Q3: Details of the tree-based information flow

Ans: Thanks for the comment. We explain the details of the tree-based information flow and the corresponding permutation operation in the response to the common question Q1. We also updated the paper accordingly.

Q4: Handling images with various degradations

Ans: We thank the reviewer for the suggestions. To validate the generalization capability of the proposed method to different types of degradation, we conducted the following experiments. First, we used the same model for both denoising and JPEG compression artifact removal tasks. Notably, a single model was trained to handle varying levels of degradation. Second, we performed experiments on image restoration under adverse weather conditions, including rain, fog, and snow. Third, we further investigated a one-in-all image restoration setup, encompassing five different tasks with real-world images. The experimental results demonstrate that the proposed method outperforms previous methods by a significant margin. These three sets of experiments collectively highlight that the proposed hierarchical information flow mechanism enables training a single model that generalizes effectively to various types and levels of degradation. We updated the experiments and results in Appendix C.4.

Q4: Model scaling-up strategies

Ans:: Thanks for the comments. We provide deeper analysis as follows. We also updated the analysis in Section B.2 and B.3 of the Appendix.

1. The heavyweight $3 \times 3$ convolutions appears at the end the the Hi-IR stage as shown in Fig. 3c.

2. Warmup is effective for training large models primarily because it mitigates issues related to unstable gradients and helps the optimizer gradually adapt to the model's large parameter space [1,2]. In the early stages of training, the model's parameters are initialized randomly. A high learning rate at this stage can cause large updates, leading to unstable or divergent training due to exploding or vanishing gradients. Warmup starts with a small learning rate and gradually increases it, allowing the optimizer to find a stable path in the loss landscape before applying larger updates. Warmup enables the model to adapt gradually, avoiding overshooting minima and ensuring smoother convergence.

3. Dot product v.s. cosine similarity: We analyze the gradient of dot product and cosine similary as follows. Suppose $\mathbf{q}$ denotes the query and $\mathbf{k}$ denotes the keys. Then dot product and cosine similarity between $\mathbf{q}$ and $\mathbf{k}$ are denoted as \text{dot\\_prod}(\mathbf{q}, \mathbf{k}), \text{cos\\_sim}(\mathbf{q}, \mathbf{k}).

   • The gradient of dot product with respect to $\mathbf{q}$ is denoted as \frac{\partial}{\partial \mathbf{q}} \text{dot\\_prod}(\mathbf{q}, \mathbf{k}) = \mathbf{k}.
   • The gradient of cosine similarity with respect to $\mathbf{q}$ is \frac{\partial}{\partial \mathbf{q}} \text{cos\\_sim}(\mathbf{q}, \mathbf{k}) = \frac{\mathbf{k}}{\|\mathbf{q}\| \|\mathbf{k}\|} - \frac{(\mathbf{q} \cdot \mathbf{k}) \mathbf{q}}{\|\mathbf{q}\|^3 \|\mathbf{k}\|} = \frac{1}{\|\mathbf{q}\|} \left(\mathbf{\hat{k}} - \text{cos\\_sim}(\mathbf{q}, \mathbf{k}) \mathbf{\hat{q}}\right), where $\mathbf{\hat{q}}$ and $\mathbf{\hat{k}}$ are normalized $\mathbf{q}$ and $\mathbf{k}$. The gradients with respect to $\mathbf{k}$ have the similar form. The gradient of cosine similarity involves more terms compared to the gradient of the dot product. This increased complexity in the gradient of cosine similarity makes it more prone to producing large or even unstable gradient values. We conducted a numerical analysis of the gradient values for the two attention methods, with the results presented in Figure 9 of the Appendix. As shown in the figure, the gradient of cosine similarity is indeed more prone to producing large values. This issue becomes more pronounced as the model scales up.

[1] Goyal, P. "Accurate, large minibatch SG D: training imagenet in 1 hour." arXiv preprint arXiv:1706.02677 (2017).

[2] Kalra, Dayal Singh, and Maissam Barkeshli. "Why Warmup the Learning Rate? Underlying Mechanisms and Improvements." arXiv preprint arXiv:2406.09405 (2024).

Q5: Rationale for model design in ablation study

Ans:: Thanks a lot for the question. We align the text better with the aim. This ablation serves two purposes: 1) exploring a better design to place the L1/L2/L3 information flow mechanisms; 2) investigating the influence of the tree depth. When designing the L1/L2 information flow attention mechanism, we need to decide whether to interleave L1/L2 information flow across Transformer layers or to implement them in the same layer. To validate this choice, we developed Version 1 and Version 2. However, Version 2 demonstrated reduced performance despite increased model complexity. To address this issue, we introduced Version 3, where L1 and L2 information flows can be conceptually unified into a single flow with a larger receptive field. We revised the corresponding texts in Section 5.1.

Q1: Permutation operation

Ans: Thanks for the suggestion. The detail of the permutation operation is explained in the response to the common question Q1. The permutation is done in a deterministic way. Removing the permuation causes the L2 information flow to degrade into L1 information flow, which isolates information as discussed in Line 52 and Table 1 of the main paper. As requested, we did an ablation study to show this effect. The following abation study shows that the permuation operation is important to maintain the performance of the network.

Q2: Depth of tree structure and model complexity

Ans: Thanks for the insighful comments. Actually, we have done this ablation study to investigate the complexity and efficiency of the proposed tree structure in Section 5.1. And the result is shown in Table 6. To summarize, we have done the following experiments.

1. We investigated three versions of the information flow mechanism with different complexities in Figure 4. In Version 1, we use L1 information flow and L2 information flow alternatively in the consecutive Transformer layers. In Version 2 and 3, both L2 and L3 information flow are implemented in a single Transformer Layer. Version 3 differs from Version 2 in that the projection layer between the L1 and L2 attention is removed and the embedding dimension of $Q$ and $K$ is reduced by half. This leads to models with different complexities including Version 1 (~14.5M), Version 2 (~19.37M), and Version 3 (~15.84M). Compared with Version 1, two consecutive Transformer Layers in Version 2 and Version 3 have two more information flow attention operations, which can be regarded as a deeper tree. Yet, the performance of Version 2 becomes worse with increased model complexity, which is due to the design. By contrast, Version 3 works better than Version 1 with deepened tree structure and increased model complexity. Considering the tradeoff between PSNR gains and model complexity, we use Version 3 throughout this paper. In short, well-designed deeper tree structures lead to improved model performance but with increased model complexity.
2. In addition to the investigation of L1 and L2 information flow, we also ablate the effects of L3 information flow. Thus, we remove L3 information flow from the network for all the three versions mentioned above, which leads to tree structures with reduced depth. In particular, the depth of the tree is reduces to 2 for the Version 1 model. By comparing the 'with L3' and 'w/o L3' columns of Table 3, we conclude the importance of L3 information flow.
3. Moreover, during the rebuttal phase, we implemented a deeper tree structure similar to Version 1. We added another information flow attention. Thus, three information flow attention operations alternate in the network, and propagate information in a $8\times 8$, $64 \times 64$, and $256 \times 256$ patch, respectively. The experimental results for the tree structure with depth 4 and another two tree structures with depth 2 and 3 discussed above are shown. The performance of the tree structure with depth of 4 is improved but the model size is also increased.

Q3: Addtional feedback

Ans:: All the addtional feebacks are addressed in the revised version.

Dear Reviewer 99w8,

Thank you for the comments on our paper.

We have provided a response and a revised paper on Openreview based on the comments. Since the discussion phase ends on Nov. 26, we would like to know whether we have addressed all the issues. Please consider raising the scores after this discussion phase.

Thank you

Q2: Novelty of the paper. (Reviewer wAuz Q1, Reviewer VDnP Q2)

Ans: The novelty of the paper comes with the follwing aspects.

1. We propose a hierachical information flow mechanism to efficiently and progressively propagate information from local to the global field. This is different from previous works such as SwinIR[1], Uformer[2], Restormer[3], ShuffleFormer[4], Shuffle Transformer[5], etc.
   • The proposed L2 information flow attention facilitates the propagation of information between dispersed pixels within an enlarged region and at an elevated level. Specifically, we avoid expanding to the entire image at this stage for two reasons: 1) the computational complexity of attention across the global image is prohibitively high; and 2) not all global image information is relevant to the reconstruction of a specific pixel. Compared with SwinIR[1], the receptive field is larger. Compared with ShuffleFormer[4] and Shuffle Transformer[5], this mechanism promotes pixel relevance.
   • The main contribution of Uformer[2] is validating the UNet architecture for Transformers while the main contribution of Restormer[3] is the self-attention along the channel dimension. Both of them are different from the proposed hierarchical information flow in this paper.
   • Although we use Columnar architecture for SR and UNet architecture for the other tasks, we do not claim the general architecture as the contribution of this paper. Those are standard choices following the literature.
2. We conduct thorough experiments and analysis to study model scaling-up for image restoration. We propose three strategies for IR mdoel scaling-up including removing heavyweight $3\times 3$ convolution, warmup, and using dot production for self-attention. Both experimental results and theoretical analysis (Appendix B) is done. In particular:
   • Removing heavyweight $3\times 3$ convolution from the network avoids intialized weight parameters with small values which leads to vanishing gradients and slow convergence.
   • Warmup helps because it mitigates issues related to unstable gradients in the early phase of training and helps the optimizer gradually adapt to the model’s large parameter space.
   • Dot product attention works better than cosine similarity attention becase the gradient of cosine similarity is more prone to producing large or even unstable values.
3. We conducted experiments on various image restoration problems including image super-resolution, denoising, motion deblurring, defocus debluring, removing rain, fog, haze, and snow from the image. The thorough analysis validated the generalizability of the proposed method.

[1] Liang, Jingyun, et al. "Swinir: Image restoration using swin transformer." Proceedings of the IEEE/CVF international conference on computer vision. 2021.

[2] Wang, Zhendong, et al. "Uformer: A general u-shaped transformer for image restoration." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

[3] Zamir, Syed Waqas, et al. "Restormer: Efficient transformer for high-resolution image restoration." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

[4] Xiao, Jie, et al. "Random shuffle transformer for image restoration." International Conference on Machine Learning. PMLR, 2023.

[5] Huang, Zilong, et al. "Shuffle transformer: Rethinking spatial shuffle for vision transformer." arXiv preprint arXiv:2106.03650 (2021).