ReFIR: Grounding Large Restoration Models with Retrieval Augmentation

[Q1: Variation of the outputs from different LRMs]

In Table 1, why there is a variation of the performance improvement when applying the proposed method to different LRMs? For example, there is a 0.38 dB PSNR improvement in SeeSR but only a 0.03 dB improvement in SUPIR. Additionally, why do the restored images produced by different methods seems still different when using the proposed ReFIR?

In fact, we have also noticed this phenomenon and we think it can be well explained based on Fig.8 in the main paper. As shown in Fig.8, for the latent at the t-th time step, there are two forces in different directions pulling it to produce the latent at the next t−1-th time step. One force is from the internal knowledge of frozen weights in LRMs, and the other is the external knowledge from the retrieved reference image. Therefore, although the external knowledge of different LRMs is the same (the reference image is the same), due to their different internal knowledge (different pre-training parameters), it eventually resulting in different LRMs producing different restoration results.

[Q2: Inference costs from additional reference image]

It seems the proposed method requires additional reference images as input, which may increase inference costs and time.

To ensure efficient implementation, we use batch accerlaration to process LR and Ref paralley, whose cost is roughly similar to increasing batchsize from 1 to 2. Moreover, although using the reference image inevitably incurs computational cost compared to direct inputing LR for inference, the additional reference can mitigate the hulucination of LRMs, resulting in significant performance gains. We provide a comprehensive comparison regarding performance and effectiveness as follows.

[Q3: Efficiency of the adaptive alignment module]

In the spatial adaptive alignment module, considering the size of the similarity matrix sim as HW×HW, which requires computing per-pixel similarity, does this introduce significant inference costs?

We are sorry to cause your confusion. In fact, the computation of the similarity matrix is quite efficient. Since we find the pixel-wise similarity between LR and Ref does not change too much among different layers, we thus compute the all layer shared similarity maxtrix for one UNet block, resulting in only 12 calculation for one timestep. We will make it more clear in the revision.

[Q4: Code open-source]

Will the authors open-source the related details and codes?

Yes! We will release all the code when the paper is accepted.

[Q1: Difference from other works]

Difference form other methods. The proposed method seems relevant to some of works in image editing tasks, such as MasaCrtl and Prompt2prompt, which also use training-free technique to modify the behavior of the diffusion model. A detailed explanation about the differences between the proposed ReFIR method and these methods would be beneficial.

Although both ReFIR and previous image editing work are both training-free, however, we point out that they differ in the following ways.

• The specific techniques are different. Although both modify the attention layer, our ReFIR further introduces the Spatial Adaptive Gating to solve the spatial misalignment between LR and Ref, and proposes strategies to allow the injection of multiple reference images. In addition, we also propose the novel Distribution Alignment to alleviate the domain gap between two diffusion process chains.
• The goals are different. Previous image editing work modifies the attention layer for controllable generation, while our ReFIR aims to inject external knowledge to alleviate the hallucination problem of existing LRMs.
• Finally, in addition to incorporating reference images through training-free attention modification, another contribution of this work is the proposal of a feasible retrieval system to facilitate the acquisition of external knowledge from retrieved databases.

[Q2: Dependency on the Quality of Retrieved Images]

From Table5, it seems the performance of ReFIR heavily relies on the relevance and quality of the retrieved images. If the retrieval system performs sub-optimally, it could adversely affect the restoration outcomes.

To make our ReFIR more robust under conditions when the relevance and quality of the retrieved images is poor, we further explore techniques including fallback strategy and the adaptive filtering policy. Please see the global author rebuttal part for more details.

Using these techniques can further improve the robustness of ReFIR when the retrieval system performs sub-optimally. We will explore more techniques for improvements in the future.

[Q3: Complexity in implementation]

The proposed cross-image injection and spatial adaptive gating mechanisms introduce additional operation. Moreover, the input image also contains extra reference images, which might challenge efficient implementation in practice.

Actually, since the modification of the attention layer only happens in the decoder in the last 20 timesteps, which only accounts for 12% of all the layers, the cross-image injection and spatial adaptive gating rise almost negligible increase in computational cost. We provide the specific cost comparison as follows.

Since the LR chain only receives the corresponding same layer features from the Ref chain, this property thus allow the batch acceleration for efficient implementation. Specifically, the input batch size is increased from 1 (LR only in original NoRef LRM) to 2 (LR+Ref in our ReFIR). Therefore, the batch parallel computation of both LR and Ref make it efficient for practical implementation.

[Q4: Scaling up with lager retrieval database]

How does the retrieval system scale with increasingly large HQ retrieval datasets, and what are the computational implications of scaling?

As suggested, we use ImageNet as the larger database for scaling up the retrieval system. The results are as follows.

It can be seen that increasing dataset size can improve the relevance and thus improves the performance. Furthermore, since the calculation of similarity between LR and retrieved dataset is parallel, we find it brings almost negligible computational cost when scaling up retrieved dataset.

[Q5: Explore differeent retrieval algorithms]

How might different retrieval algorithms (beyond nearest neighbor lookups) impact the performance of ReFIR?

As suggested, we use the TopK retrieval algorithm to explore the impact of different retrieval algorithms. Specifically, we select images with top 3 large similarity as reference inputs, and use the multi-reference injection technique to allow multiple reference images. The experimental results are shown below.

It can be seen that using the TopK retrieval algorithm can bring more relevance to the reference image, thus improving the performance.

[Q1: Reference image presentation]

Thanks for your advice, we will add the retrieved image in Fig.5 and Fig.6 in the revision to improve the presentation quality.

[Q2: Improving ReFIR's ability in real-world scenarios]

Tables 1 and 2 highlight a significant performance ... detailed discussion is necessary.

To further improve the ability of ReFIR in real-world applications, we further explore several techniques to improve its robustness. The related analysis and experiments are shown in the global author rebuttal.

Equipped with the fallback mechanisms (not using reference or generating Ref like CoSeR) as well as the adaptive filtering strategies (relevance/quality/task-based filtering), our ReFIR can alleviate the dilemma when high-correlated and high-quality images are scarce or unavailable, enhancing its ability in real-world scenarios.

[Q3: Moving the retrieval stage details to main body]

Due to the 9-page limit for submission this year (1 page less than usual), we are sorry for placing the image retrival details in the appendix. We will add more details about the image retrieval stage in the main body in the revision.

[Q4: Computaion and time cost from Source Reference Chain]

The paper should discuss the computational cost and time consumption associated with the "Source Reference Chain" process. While the performance improvement is valuable, quantifying the trade-off in terms of computational resources is essential.

Since the "Target Restoration Chain" only receives the corresponding same time-step and layer features from the "Source Reference Chain", this property thus allow the batch acceleration for efficient implementation. Specifically, these two chains are implemented sharing one LRM weight, and simply increasing the batch size from 1 (LR only in original NoRef LRM) to 2 (LR+Ref in our ReFIR).

Therefore, the computational cost and time consumption of ReFIR is roughly similar to the original LRM with batchsize set to 2, with a slight overhead increase from the additional cross-image injection. The comparison of model complexity before and after incorporating our ReFIR are given in Table7 of Addpendix, and here we also give this reuslt as follows:

TableA. Comparison on computation and time cost. The input resolution is $2048\times2048$, and is evaluated on one 80G A100 GPU.

[Q5: Dicussion on patchfied images ]

Consider discussing whether the model retrieves features from the corresponding chunk of the reference image or employs a different strategy to handle chunked inputs.

The patchfy is usually used when the input resolution is large. Although it is intuitive to use the corresponding patch from the large Reference Image as the reference branch's input, it can inevitably incur significant performance degradation when the LR and Ref possess spatial misalignment.

Here, we introduce another solution which simply uses existing techniques to deal with this problem. Specifically, we employ the Multi-reference injection, shown in Appendix-A, in which we use all the patches from the reference image to guide the restoration of a certain LR input patch. In this way, the LR image can maintain the global receptive filed from the Ref. In experiment, we choose the number of patches as 4 for a 2048x2048 input from RealPhoto60 dataset, resulting in 4 small patches with size 512x512. The results are shown as bellow.

To some extent, introducing the multi-reference injection can alleviate the perofrmace drop while maintaining modest memory footprint. However, similar to exsting SD-based restoration works, the inherent tradeoff between whole-image receptive filed and the low GPU cost still exist. Our framework may potentially benefit from future acceleration works.

[Q6: Computation mode for retrieval]

It would be helpful to clarify whether the features used for matching in the reference image retrieval process are pre-computed or computed online.

Since the retrival database is agnostic to the specific LR inputs, we thus pre-compute the retrival vectors for fast inference. We will make it more clear in the revision.

[Q7: Converting refernece image as internal knowledge]

Due to the backpropagation of LRM needs huge GPU costs, e.g. 64 Nvidia A6000 GPUs for training SUPIR, we are not able to practically convert reference images as internal knowledge through fully fine-tuning.

However, we point out that the idea of injecting downstream knowdge through fine-tuning has been explored in the image genration community, such as DreamBooth. During the subject-driven adaptation of dreambooth, the authors of dreambooth also encountered diffusion model's overfitting for specific subjects, and they propose the prior preservation loss to solve this problem.

As for the image restoration, it is difficult to get calibration datasets for prior preservation since the content of LR images can be various, unlike the similar topic in image generation. Moreover, in real-world applications, it is also difficult to obtain scene-specific data for training in advance, and even if we get it, it requires additional training time. In contrast, our ReFIR framework does not need the scene-specific data in advance and can adapt to specific scenes in a training-free manner.

Thank you very much for your positive feedback. We are delighted that our responses have addressed your concerns. We will further improve our work in revision based on the reviewers' comments and discussions.

[Q1: Reliance on the image quality and relevance]

The effectiveness of ReFIR is significantly dependent on the quality and relevance of the images retrieved from external databases. This reliance could limit the framework's effectiveness in scenarios where highly relevant and high-quality reference images are scarce or unavailable. Addressing this, the authors could explore mechanisms to assess and ensure the relevance and quality of retrieved images dynamically during the restoration process or develop fallback strategies when suitable images are not found.

We understand your concerns, and we have thus conducted a thorough experiment presented in the global author rebuttal part.

With these newly developed fallback and filtering strategies, our ReFIR is further improved under the real-world retrieval settings, relieving the dependence on retrieval system.

[Q2: Computational overhead from retrieval and attention modification]

While the paper mentions that the framework does not require additional training, it does not thoroughly address the potential increases in computational overhead and latency introduced by the retrieval process and the modification of self-attention layers. This could be particularly challenging when deploying in real-time or resource-constrained environments. Future iterations of the framework could benefit from a detailed analysis of computational costs and potential optimizations to streamline the retrieval and integration processes.

In order to reduce the computational overhead of the retrieval process, we pre-calculated the feature vectors of all images in the retrieval database before inference. Furthermore, the cosine similarity between the LR image vectors and all retrieval vectors is computed in parallel. These strategy results in an almost negligible (less than 3% inference time) cost of computational overhead.

The modification of self-attention layers only happens in the last 20 timestep in the decoder layers, i.e., only 12% attention layers are modified while the left is kept intact.

These analysis is also supported by practice, in which we find these two process only take up <5% inference time, with most computational cost coming from the original LRM. Future LRM acceleration (e.g. pruning, quantization) will benefit our ReFIR, and we will explore more efficient implementation in the future.

[Q3: Testing on broader degradations]

The experimental setup primarily demonstrates the effectiveness of ReFIR under controlled conditions with specific types of image degradations. To fully validate the robustness and generalizability of the framework, additional testing across a broader spectrum of real-world scenarios and more diverse degradation types would be beneficial. This would help in understanding the limitations and operational range of ReFIR in practical applications, ensuring it can effectively handle unexpected or uncommon degradation patterns.

As suggested, we test the robustness of our ReFIR framework on another two challenging real-world mixed degradations, i.e., low-resolution$\times$4+noise$\sigma$=50, as well as low-resolution$\times4$+JPEG$q=30$. The results are as follows.

TableA low-resolution$\times$4+noise$\sigma$=30 on the WR-SR dataset

TableB low-resolution$\times4$+JPEG$q=30$ on the WR-SR dataset

It can be seen that our ReFIR framework maintains its effectiveness on real-world hybrid degradations, demonstrating its robustness and generalizability.

[Q4: Possible bias from retrival database]

We agree with you. Currently, we only use publicly available, high-quality academic data as the retrieval database. In the future, we will pay attention to content diversity and potential bias in constructing larger scale retrieved data, e.g., we could use VLMs or LLMs for automatic filtering followed by human review.