Adapt and Diffuse: Sample-adaptive Reconstruction via Latent Diffusion Models

We thank the reviewer for the effort reviewing our paper. We are glad that the reviewer thinks that our experiments show that Flash-Diffusion “leads to the improvement of reconstruction quality and the reduction of necessary computational cost”. This is in fact the main focus of our paper. Please find below our response to your concerns.

• Re theoretical backing: We position our work as a thorough empirical study of a concept that is highly underexplored in the modern inverse problem literature: sample-adaptivity. We hope that the reviewer can appreciate the novelty of the severity encoding framework and the clear demonstration of gains provided by Flash-Diffusion both in terms of efficiency and reconstruction performance. That said, it would be interesting to investigate some theoretical properties of the framework in future work and we would be open for suggestions on specific ideas the reviewer thinks can further strengthen our paper.

• Re experimental results: We understand that the reviewer needs more experimental results in order to better gauge the validity of the results. We have added the following new experiments to the manuscript:

  1. A complete set of experiments on a new forward operator: random inpainting.

  2. In-depth ablation studies against shifts in noise level and forward model (Appendix D).

  3. An ablation study on the usefulness of various components of our framework (Appendix C).

  4. A set of experiments on pairing severity encoding with DDIM sampling (Appendix E).

  We hope that these additional experimental results convinced the reviewer on the validity of our method. We are open for further specific suggestions the reviewer would find necessary to incorporate into the paper.

• Re hyperparameters: We detail how we tuned each hyperparameter in Appendix A. We would like to highlight that we have the same number of hyperparameters as other diffusion-based reconstruction methods, such as CCDF-DPS, and only one extra hyperparameter compared with DPS, the correction factor $c$. However, we find that similar values of $c$ work across all forward operators. The hyperparameter $\lambda_{\sigma}$ controls the relative importance of error estimation in the loss, and we tuned it for Gaussian blur over grid search ($[0.1, 1.0, 10]$) and used the same value for all other models with satisfactory results. That is, we find that no extensive hyperparameter tuning is required across tasks. Moreover, we simply scale the reconstruction loss terms $L_{im. rec.}$ and $L_{lat. rec.}$ with their corresponding dimension (pixel space dimension and latent space dimension respectively) without any further hyperparameter tuning across all experiments. Thus, we conclude that severity encoder training is fairly robust against these hyperparameters.

We hope that we have addressed the reviewer’s concerns adequately and are happy to discuss further if there are any remaining concerns.

We appreciate the reviewer’s time reviewing our work. Please find below our response to your concerns.

• Re Weakness 1: First, we would like to point out that the idea of adapting the cost of reconstruction to the severity of degradation has not been explored in detail to the best of our knowledge in the literature. Severity encoding is our novel contribution to fill this gap. We introduce severity encoding as a general framework that can be seamlessly combined with existing diffusion solvers to add adaptivity to such techniques. We don’t believe that this is a limitation of our work, but a property that adds extra flexibility to our reconstruction framework. Moreover, we have indeed discussed the connection of our work to Chung et al. (2022c) in Sections 2 and 3 and demonstrated significant improvements over their technique in terms of reconstruction quality (see CCDF-DPS in Table 1).

• Re Weakness 2: As the reviewer mentions, our severity estimates build upon the assumption that the estimation error in latent space is i.i.d. Gaussian. We agree with the reviewer that this assumption does not necessarily hold in practice. Indeed this is something we discuss in the paper and propose a technique to mitigate it. It would be an exciting direction for future work to investigate more nuanced error models for severity encoder training, such as modeling as Gaussian with general diagonal covariance. We however respectfully disagree about the impracticality of the assumption as our in-depth experiments demonstrate that the learned severity estimates (1) correlate well with the true underlying image degradation level (Figures 3 and 4), (2) can provide useful insight into outlier samples in the dataset in terms of reconstruction difficulty (Figure 4) and (3) can be effectively leveraged by diffusion-based solvers to accelerate sampling and improve reconstruction quality (Table 1 and Figure 6). We are open for specific suggestions on how to further demonstrate the viability of our proposed framework.

• Re Weakness 3:

  • Re efficiency: We highlight the efficiency of our method by demonstrating that any non-adaptive/fixed starting diffusion time would result in worse performance than our sample-adaptive method. This key experiment is depicted in Figure 6, which we understand the reviewer found confusing. We updated the figure and its caption to more clearly reflect our efficiency results. To highlight our findings based on Figure 6 and Table 1: we achieve better reconstruction quality using approx. 100 diffusion steps on average than DPS with 1000 diffusion steps. Moreover, using any fixed/non-adaptive starting time in the latent diffusion framework achieves worse performance than sample-adaptive Flash-Diffusion.

  • Re using other sampling techniques: Thank you for suggesting the combination with DDIM, this is a great idea. Even though we are familiar with the efficiency of DDIM in image generation, we are unaware of general (including nonlinear) inverse problem solvers that leverage latent space DDIM and achieve performance comparable to state-of-the-art in reconstruction. We would appreciate it if the reviewer could point us to such techniques, we would gladly compare them with Flash-Diffusion. That said, we added two new sets of experiments to address the reviewer’s question on pairing severity encoding with DDIM.

    1. We replace DDPM sampling (N=1000) with DDIM sampling (N=20 and N=100) in Flash-Diffusion in order to investigate the viability of such accelerated sampling schemes for latent space reconstruction. We observe that Flash-Diffusion obtains high quality reconstructions with as low as 10 NFEs, however there is a clear degradation in performance in terms of image quality metrics (see Table 5). We hypothesize that this can be due to the inaccuracy of LDPS in the accelerated sampling setting, which is an interesting direction for future work. Furthermore, we find that Flash-Diffusion reconstruction quality is robust in terms of the DDIM sampling stochasticity parameter. See a detailed discussion and experimental results in Appendix E.

    2. We demonstrate that naively running DDIM sampling in latent domain with LDPS, even though greatly reduces sampling cost, is not competitive with other diffusion-based solvers. In particular, we experiment with LDPS with DDIM sampling of N=100, which closely matches the average number of steps Flash-Diffusion uses in the DDPM setting. We observe that reconstruction quality has extreme variance: the technique produces reasonable reconstructions on some samples, but completely breaks down on others resulting in poor performance on average (see Table 6). A comprehensive discussion is added to Appendix E.

Thank you for the insightful review. We are glad that the reviewer agrees that our work on sample-adaptive inference “addresses a major limitation of existing methods” and that our experiments are “well-designed” and the analysis is “comprehensive”. Please find below our response to your concerns.

• Re discussion on potential limitations: Thank you for pointing this out along with Reviewer Cd1z. We agree that understanding the limitations of severity encoding is crucial. Therefore, we added an in-depth robustness study investigating the effects of noise level perturbations and mismatch in forward model on severity encoding. To summarize our findings:

1. Noise perturbations in test time: Interestingly, we find that severity encoding maintains its performance at noise levels lower than in the training setting (Figures 8 and 9), which results in steady reconstruction performance of Flash-Diffusion. At higher unseen noise levels severity encoding accuracy decreases, however the overall degradation in Flash-Diffusion performance is more graceful than in case of other supervised baselines (Figure 9). We hypothesize that this effect is due to adaptivity: the severity estimates are still useful in finding an appropriate starting time, however not as accurate as at familiar noise levels. Moreover, we observe that the severity encoder manages to scale the compute effort at noise levels unseen during training: for measurement noise variances outside the training setting, the predicted number of necessary diffusion steps is proportional to the noise level (Figure 10). However, at extreme noise perturbations, the severity encoder completely breaks down, assigning essentially the same severity estimate to any input.

2. Forward model perturbations in test time: we observe a consistent degradation in performance when the forward model is drastically changed compared to the training setup (e.g. training on Gaussian blur and evaluation on nonlinear blur, Table 4). However, we find that as long as the operators are somewhat similar, the drop in performance is more modest. For instance, we show that a severity encoder trained on non-linear blur has high accuracy on estimating the degradation severity of images corrupted by Gaussian blur, as long as the blur amount is high (Table 3 and Figure 11). This observation may hint at the potential for generalizing to unseen operators when the severity encoder is trained on a mix of known operators, which is an exciting direction for future work. We have also highlighted these limitations in the main text in Section 4 and we are planning to incorporate the above experiments into the main manuscript in the camera-ready version (given one extra page). With respect to strategies to mitigate severity encoder failure: this is a very interesting question. As we observe in the robustness experiments, the severity encoder remains accurate at lower unseen noise levels. Therefore, a technique to improve noise robustness would be to train the encoder on slightly higher noise levels than expected in test time in order to “proof” it against potential failure when the noise level is perturbed.

• Re ablation studies: This is a great point. We added a new set of experiments in Appendix C, where we ablate the effect of the various components of Flash-Diffusion as the reviewer suggested. As seen in Table 2, adaptivity via severity encoding has the largest contribution to performance. We obtain comparable results with and without latent space diffusion, however latent diffusion is preferable since (1) it has greatly reduced computational cost and (2) degradation severity is better captured in latent space. We hope the reviewer finds this new ablation study useful.

We are happy to discuss any further questions or concerns and would like to ask the reviewer to consider raising their score in case we addressed all major issues.

We are grateful for the helpful feedback. We are glad that the reviewer thinks that the “main idea of the paper is explained well and supported by experiments” and that the proposed sample-adaptive computation is natural and novel. Please find below our response to the issues the reviewer has raised.

• Re acknowledgment of limitations and robustness studies (bullets 1 and 3): Thank you for raising this point. We agree that it is crucial to understand the limitations of the severity encoder, especially in the face of realistic test-time perturbations, including noise level mismatch and forward model shifts. To address your point, we have performed a new set of robustness studies that dive deeper into this question. To summarize our findings:

1. Noise perturbations in test time: Interestingly, we find that severity encoding maintains its performance at noise levels lower than in the training setting (Figures 8 and 9), which results in steady reconstruction performance of Flash-Diffusion. At higher unseen noise levels severity encoding accuracy decreases, however the overall degradation in Flash-Diffusion performance is more graceful than in case of other supervised baselines (Figure 9). We hypothesize that this effect is due to adaptivity: the severity estimates are still useful in finding an appropriate starting time, however not as accurate as at familiar noise levels. Moreover, we observe that the severity encoder manages to scale the compute effort at noise levels unseen during training: for measurement noise variances outside the training setting, the predicted number of necessary diffusion steps is proportional to the noise level (Figure 10). However, at extreme noise perturbations, the severity encoder completely breaks down, assigning essentially the same severity estimate to any input.

2. Forward model perturbations in test time: we observe a consistent degradation in performance when the forward model is drastically changed compared to the training setup (e.g. training on Gaussian blur and evaluation on nonlinear blur, Table 4). However, we find that as long as the operators are somewhat similar, the drop in performance is more modest. For instance, we show that a severity encoder trained on nonlinear blur has high accuracy on estimating the degradation severity of images corrupted by Gaussian blur, as long as the blur amount is high (Table 3 and Figure 11). This observation may hint at the potential for generalizing to unseen operators when the severity encoder is trained on a mix of known operators, which is an exciting direction for future work. We also agree that these limitations should be better highlighted in the main manuscript. Therefore, we added a discussion to Section 4 on the above limitations and we plan to move the robustness studies to the main paper given an extra page in the camera-ready version.

• Re other heuristics for degradation severity (bullet 2): We have a discussion on the challenges of quantifying degradation severity in pixel space in Section 3. That said, it would be interesting to compare the estimates from severity encoding with a heuristic baseline. We are open for suggestions in this direction for a simple and natural baseline.

• Re bullet 3 on robustness to measurement variations: See discussion above where we simultaneously address bullets 1 & 3.

• Re consistency models (bullet 4): Thank you for the references. Combining severity encoding with consistency models is indeed an interesting direction that can potentially result in few-shot high-fidelity reconstructions. Our adaptive shortcut idea works as long as SNR can be analytically evaluated for each time step. The crucial challenge would be that of maintaining data consistency, which has not been thoroughly investigated in the context of consistency models to our knowledge. This is a great suggestion however for future investigation.

Please let us know if you have any further questions or suggestions. We hope that we managed to address all major concerns, in which case we would really appreciate it if the reviewer could update their rating given the generally favorable review.