Solving Inverse Problems with Latent Diffusion Models via Hard Data Consistency

We would like to thank the reviewer for the insightful comments! Below, we address the reviewer’s concerns in detail:

Comment: the rationale behind the result: why pixel-space diffusion models achieves lower quality than with our method

Response: Thanks for pointing this out! We have the follow hypotheses:

1. LDM has potential to achieve better generation quality than pixel-based diffusion models. This is highlighted in the paper [1] . It shows that LDM is able to outperform DDPM based methods in FID when training on the same dataset. This indicates that LDM may approximate the prior distribution better than pixel-based DDPM models.
2. The gain of our performance may comes from the careful design of resample and hard consistency, which may leads to smaller measurement consistency loss compared to DPS-based methods especially when the forward operator is nonlinear and nonconvex. As shown in Table I in the following, we observe that ReSample is able to achieve better measurement consistency than latentDPS. Also DPS shows significant performance degradation when it comes to nonlinear inverse problems, but the performance of ReSample is potentially more stable as demonstrated in Table 2 and 6 in the paper. This observation further corroborates our hypothesis. We provide insights on this in Appendix section D.
3. Some work in parallel with ours theoretically demonstrates that DPS-based algorithms may encounter performance degradation with higher resolution data, as reported in [2].

Table I: average measurement consistency loss for CT reconstruction for ReSample and Latent-DPS

Comment: Experimental Results on box-inpainting and comparison with PSLD

Response: Thanks for the suggestion! We perform additional experiments for box inpainting task on the CelebA-HQ dataset. We use the same setting as in [3] for generating the box mask. The results are demonstrated in Figure. 7 and Table. 7 in Appendix Section A.1. Both quantitative results and qualitative results demonstrate that we are able to achieve better or comparable results compared to the baselines. We add an additional result section in the Appendix (Section A.1) to discuss about the results of the box-inpainting task. PSLD is able to generate very realistic images, but it sometimes has unstable reconstructions for box inpainting as we observed. We also attach the performance in Table II below.

Table II: Performance of box-inpainting on CelebA-HQ

[1] Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

[2] Rout, Litu, et al. "Solving linear inverse problems provably via posterior sampling with latent diffusion models." arXiv preprint arXiv:2307.00619 (2023).

[3] Chung, Hyungjin, et al. "Diffusion posterior sampling for general noisy inverse problems." arXiv preprint arXiv:2209.14687 (2022).

Thanks again for your review! Your review is crucial for us to improve our manuscript. Please let us know whether our response has addressed your concerns. If there is any remaining question about the manuscript and we will try our best to answer.

Dear reviewer 4YEX,

Thank you very much for the encouraging feedback! We are glad to see that you are satisfied with the response and improved the score of our paper, and willing to answer any further question.

Comment: Possible switch to a plain norm instead of square norm for optimization. Does the theoretical analysis still hold?

Response: Thank you for this interesting observation! Following this comment, we conducted several studies on using a plain norm as our objective function instead of the squared norm. However, we did not observe much of a difference (if any) in the performance of ReSample. We conjecture that for DPS, the algorithm needs to compute the gradient with respect to the score function, whereas we calculate the gradient with respect to the decoder. Therefore, there may be some subtle differences in these two networks that favor a squared norm over a plain norm.

In case of whether our theoretical analysis holds, our analysis mainly focuses on the resampling part of our algorithm, which takes an arbitrarily optimized $\hat{z}_0(y)$ as input, so the variance reduction and unbiasedness results still hold. Since the last theorem is studying unconditional reverse sampling and is not related to the hard consistency, we believe that it still holds for the plain norm.

Comment: Emphasizing strengths other than computational efficiency

Response: Thank you for the great suggestion! We acknowledge that due to the hard data consistency update through the latent space, our algorithm requires slightly more inference time than several pixel-diffusion based approaches, while comes along with better memory efficiency than pixel-diffusion based methods as indicated in Table 4 in our paper. We edit our introduction part to emphasize more on our memory efficiency, the training efficiency of latent diffusion models and the novelty of our algorithm.

Comment: Adding references and discussions about other methods such as CSGM

Response: Thanks for catching this! We have added a related work section in our Appendix section E for a thorough discussion of other methods. We also include citations and a brief discussion of this in our introduction.

Comment: How is our method performing when data fidelity is performed at every time step?

Response: We also report additional results of both partial and full data-consistency performance in the task of chest CT reconstruction, as shown in Figure 18 and Table 8 in the Appendix section A.5 (shown as Table III here). Generally, from Table 8, we observe that when the skip step size (a skip step size of $10$ implies that we apply hard data consistency once every $10$ iterations of the reverse sampling process) get below 10, the performance almost saturates, which means we found that the quantitative metrics do not significantly change or only improve marginally with full data-consistency. Nevertheless, visual quality shows slightly sharper images with full data-consistency. We revise the ablation study section in the Appendix (Section A.5) to include both quantitative and qualitative results on this. This observation may imply that by only doing partial hard-consistency, we are able to achieve decent quality reconstructions that is not far from full data-consistency. This also shows the flexibility of our algorithms which can adjust the data consistency steps by balancing the trade-off between reconstruction quality and time cost.

Table III: average performance on chest CT images reconstruction with data consistency performed on each $k$ time step

Comment: A discussion on the PSLD method [6]

Response: Thanks for the suggestion! We have added a discussion a few sentences stressing the similarities and differences of our methods and [6] in the related work section in the Appendix (Section E) with a highlight of blue color.

Comment: Adding references on discussing Tweedie's formula

Response: Thanks for the suggestion! We added references on Tweedie's formula in our background section in the main text.

Comment: Misinformation in Figure 16

Response: Thanks for pointing this out! We remove this plot to remove the misinformation as in the chaotic stage we do not access the measurement (as in our experiments we only start to see patterns of the clean image after the middle of the second stage).

We would like to thank the reviewer for the insightful comments! Below, we address the reviewer’s concerns in detail:

Comment: While the general background on diffusion models is well documented, key references are missing regarding diffusion models in constrained settings as well as latent space optimization methods

Response: Thanks to the reviewer for sharing the reference list! We have added these references in our revised paper with more discussions about these methods in the introduction, methods and a new related work section in Appendix section E. We will add these references to the main paper if extra one page is provided during camera ready.

Comment: The comparison with some other methods seems rather unfair. In particular, ADMM-PnP is ran for only 10 iterations, while it often requires at least 50 to 100 iterations to yield good results; similarly, the FBP-UNet seems to strongly under-perform compared to the (Jin et al.) reference.

Response: Thanks for the question and we would like to clarify the baseline settings used in the paper as follows:

1. We only use 10 iterations because this achieves the best performance. We tune the ADMM-PnP parameters based on images in the validation set, and pick the hyperparameters that give the best performance. We observe that if we apply the pretrained denoiser with too many iterations ( $>10$), it tends to over-smooth the image and cause a drop in performance.

2. For FBP-UNet, we are considering a much more challenging setting than [2]. In [2], they use 50 projections and no measurement noise. We experiment on 25 projections with additional Gaussian noise, which is much sparser. As a result, we can observe FBP gives severe streaking artifacts in this scenario as demonstrated in Figure .4 in our paper, which is not visible in [2]. This is the reason why FBP-UNet does not perform as good as the results reported in [2].

To avoid the confusion, we have revised section C.3 in the Appendix to add more descriptions to clarify the baseline settings and discussions about the baseline results in the revised paper.

Comment: Motivations of the paper not very consistent with results

Response: We appreciate the reviewer’s feedback and the opportunity to clarify our paper’s motivation. We acknowledge that due to the hard data consistency update through the latent space, our algorithm requires slightly more inference time than several pixel-diffusion based approaches, while comes along with better memory efficiency than pixel-diffusion based methods as indicated in Table 4 in our paper.

Because there are many challenges in using latent diffusion models for inverse problems, with one of the main challenges lying in the nonlinearity of the decoder, resulting in numerous local minima that lead to unfavorable results. This motivates us to propose the ReSample method to obtain reconstructions faithful to the measurements with latent diffusion models. To avoid confusion, we have strengthened our contribution of a novel methodology with satisfying performance in the introduction section (highlighted in color in the revised paper). We provide an example on the measurement consistency here in Table I, which demonstrates that our algorithm improves the measurement consistency a lot, which supports the our motivation. In this regard, we believe that our algorithm is an important first step in developing an efficient algorithm that can leverage latent diffusion models to solve inverse problems. We provide insights on this in Appendix section D.

Table I: average measurement consistency loss for CT reconstruction for ReSample and Latent-DPS

On the other hand, one of our primary motivations for this work is the computational demands associated with training diffusion models in pixel space as indicated in our paper's abstract and introduction, where latent diffusion models are able to achieve lower training time and reduced memory usage compared to pixel diffusion models, as observed and reported in prior works [3,4]. Experimental results are consistent with this motivation. We added another subsection in the computational efficiency section of the Appendix (Section A.5) to highlight this point about training efficiency. We also attach Table II below with a comparison of training time and memory of LDM [3] and DDPM [5] here.

Table II: Efficiency and Performance of training 100k iterations of LDMs and DDPMs on CT images on a V100 GPU

Comment: Is Stage 0 necessary

Response: Thanks for pointing this out. We agree that existing works have achieved great progress on finding a good initial point/noise for reconstruction problems [7, 8]. However, it is still challenging to find the right time step to insert the noise, as some degradation is mild and some degradation is severe, as mentioned in [9]. Hence we resort to the method of starting from a random noise in our paper. Furthermore, the approximation of $z_0$ via Tweedie’s formula may be inaccurate during the early sampling stages, potentially leading us to a poor local minimum. Thus, we believe that stage 0 is important for improving the performance of the current algorithm of ReSample. We agree with the reviewer that this is indeed a very interesting research question that could be a potential follow-up of our work in the future work.

Comment: Other editing

Response: We added the PSNR metrics on images for the additional results for box-inpainting and the ablation study on data consistency time steps in the Appendix A.2 and A.5. We modify Figure.6 to be more clear about the x-axis (i.e. ReSample frequency)

[1] Ryu, Ernest, et al. "Plug-and-play methods provably converge with properly trained denoisers." International Conference on Machine Learning. PMLR, 2019.

[2] Jin, Kyong Hwan, et al. "Deep convolutional neural network for inverse problems in imaging." IEEE transactions on image processing 26.9 (2017): 4509-4522.

[3] Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

[4] Luo, Simian, et al. "Latent Consistency Models: Synthesizing High-Resolution Images with Few-Step Inference." arXiv preprint arXiv:2310.04378 (2023).

[5] Nichol, Alexander Quinn, and Prafulla Dhariwal. "Improved denoising diffusion probabilistic models." International Conference on Machine Learning. PMLR, 2021.

[6] Rout, Litu, et al. "Solving linear inverse problems provably via posterior sampling with latent diffusion models." arXiv preprint arXiv:2307.00619 (2023).

[7] Liu, Guan-Horng, et al. "I $^ 2$ SB: Image-to-Image Schr" odinger Bridge." arXiv preprint arXiv:2302.05872 (2023).

[8] Chung, Hyungjin, Jeongsol Kim, and Jong Chul Ye. "Direct Diffusion Bridge using Data Consistency for Inverse Problems." arXiv preprint arXiv:2305.19809 (2023).

[9] Fabian, Zalan, Berk Tinaz, and Mahdi Soltanolkotabi. "Adapt and Diffuse: Sample-adaptive Reconstruction via Latent Diffusion Models." arXiv preprint arXiv:2309.06642 (2023).

Dear reviewer azNZ,

Thank you very much for the encouraging feedback! We are glad to see that you are satisfied with the response and willing to improve the rating of our paper.

Thanks a lot for sharing the paper with us which looks very interesting! We will use the DPIR algorithm in our future work.

Thanks again for your review and the response! Your review is crucial for us to improve our manuscript. If there is any remaining question about the manuscript and we will try our best to answer.

We would like to thank the reviewer for the insightful comments! Below, we address the reviewer’s concerns in detail:

Comment: Scaling up to high-resolution images

Response: To scale up the algorithm to high-resolution images, one possible way would be using an arbitrary-resolution (as long as divisible by 4 for instance) random noise as the initial noise of the LDM. Due to the convolutional nature of UNet of the LDM and the autoencoder, one can output arbitrary resolution generated images corresponding to the resolution of the latent vector, as demonstrated in [1]. In the revision, we have added this discussion in the Appendix section A.3 and provide an example of reconstructed high-resolution images in the task of random inpainting.

Comment: Closed-form expression for data consistency

Response: Thanks for pointing this out! We observe that in a noiseless setting with linear inverse problems, there is indeed a formula that would work, which is given by $\hat{z}_{0}(y) = E(D(\hat{z}_0) - (A^{+}AD(\hat{z}_0) - A^{+}y))$, where $E$ is the encoder. We revise the paper to provide an additional discussion on this formula in the Appendix section B. We also already applied a similar formula in the CT reconstruction task (where the pseudo-inverse is implemented by Conjugate Gradient) as demonstrated in Appendix B.1, and we revise this subsection for a clear description on how we use these formulas for our proposed optimization problem.

Comment How to apply ReSample to an unknown linear operator

Response: This is a very interesting question! To our best knowledge, we get to know there are some recent works [1,2,3,4] studying on the blind inverse problems with diffusion models. One interesting venue of these works is to train an image diffusion model and a parametrized operator diffusion model. At inference time, sampling is performed simultaneously from both diffusion models. Our algorithm can be incorporated with this kind of method by replacing the image diffusion model with the latent diffusion models. We leave this as an interesting follow-up research direction that we would like to explore more in the future.

[1] Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

[2] Chung, Hyungjin, et al. "Parallel diffusion models of operator and image for blind inverse problems." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

[3] Murata, Naoki, et al. "Gibbsddrm: A partially collapsed gibbs sampler for solving blind inverse problems with denoising diffusion restoration." arXiv preprint arXiv:2301.12686 (2023).

[4] Levac, Brett, Ajil Jalal, and Jonathan I. Tamir. "Accelerated motion correction for MRI using score-based generative models." 2023 IEEE 20th International Symposium on Biomedical Imaging (ISBI). IEEE, 2023.

Thanks again for your review! Your review is crucial for us to improve our manuscript. Please let us know whether our response has addressed your concerns. If there is any remaining question about the manuscript and we will try our best to answer.

We would like to thank the reviewer for the insightful comments! Below, we address the reviewer’s concerns in detail:

Comment: data consistency update is done with $z_{t-1}$ different from prior works

Response: Sorry for the confusion. We have a slightly different notation from prior works. Prior works write the reverse sampling from $z_{t}$ to $z_{t-1}$ in their algorithm pseudo-code, whereas our algorithm denotes the reverse sampling from $z_{t+1}$ to $z_t$, and so indeed we only feed-forward the diffusion model once for each time step. Therefore, our work is actually doing data consistency updates on $z_t$, which is the same as in prior works. Figure 2 may illustrate this point more clearly with the blue arrows in the figure indicating the data consistency operations. Please let us know if this clarifies your comment, and we would be happy to follow up.

Comment: Performance of full data-consistency on every time-step

Response: We revised the ablation study section in the Appendix (Section A.5) to report the inference time v.s. data consistency with different time steps of $k$ for chest CT reconstruction (we perform data consistency on every $k$ step, i.e. if $k=1$, we perform data consistency on every time step)

We also report additional results of both partial and full data-consistency performance in the Table I below. We also revise the ablation study section in the Appendix (Section A.5) to include both quantitative and qualitative results on this, as shown in Figure 18 and Table 8. Generally, from Table I, we observe that when the skip step size (a skip step size of $10$ implies that we apply hard data consistency once every $10$ iterations of the reverse sampling process) get below 10, the performance almost saturates, which means we found that the quantitative metrics do not significantly change or only improve marginally with full data-consistency. Nevertheless, visual quality shows slightly sharper images with full data-consistency. We revise the ablation study section in the Appendix (Section A.5) to include both quantitative and qualitative results on this. This observation may imply that by only doing partial hard-consistency, we are able to achieve decent quality reconstructions that is not far from full data-consistency. This also shows the flexibility of our algorithms which can adjust the data consistency steps by balancing the trade-off between reconstruction quality and time cost.

Table I: Inference times in mm:ss and performance for performing data consistency on each k step for chest CT reconstruction

Comment: There is no description of $\gamma$ in figure 15

Response: Thanks for pointing this out! We have added the description of $\gamma$ in Figure 15 in the revised manuscript.

Dear reviewer 6mow,

Thanks for the prompt response!

the proposed method seems to compute the score twice: for unconditional sampling $s_\theta (z_{t+1}, t+1)$ and for the hard data consistency $s_\theta (z'_t, t)$. Is that just a typo?

Thank you for catching this! We use the $s_\theta (z_{t+1}, t+1)$ for both unconditional sampling and hard data consistency. We just uploaded a revised version of the paper to make it clear that we only compute the score once per each time step (we also highlight our resample part of the algorithm in blue) at the top of page 5.

Do the authors believe that this applies similarly to other images such as CelebA?

This is a great question! Indeed we use $k=10$ for every experiment we report in this work including CelebA-HQ, FFHQ, and CT images. We observe that $k=10$ works well for every experiment we reported in this work while striking a good balance between inference efficiency and reconstruction quality. Intuitively, the next couple steps of samples after the data consistency step can retain similar semantic information to some extent, so we guess that a mild partial consistency like $k=10$ should be applicable to a variety of situations.

Thanks again for your review and your follow-up discussion! Your review is crucial for us to improve our manuscript. Feel free to let us know if there is any remaining question about the manuscript and we will try our best to answer.