Paper: Solving Inverse Problems with Latent Diffusion Models via Hard Data Consistency (ICLR-2024)

Reviewer comment (Evaluation):
It is not clear why training the model in pixel space (Table 1 and 2) achieves lower quality results compared with the proposed ReSample algorithm. The rationale behinds this result remains unclear.

Author response:
***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].


|Anatomical site|ReSample (Ours)|Latent-DPS|
| -------- | ------- |------- | 
|Chest	|***5.36e-5***	|2.99e-3|
|Abdominal|	***6.28e-5***|	3.92e-3|
|Head	|***1.73e-5***	|2.03e-3|

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

Reviewer comment (Experiments):
The experimental results primarily involve random inpainting, a task considered relatively straightforward. It would be valuable to explore how the proposed algorithm performs in more complex inpainting scenarios, such as the removal of large pixel regions, as demonstrated in the approach outlined by Rout et al. in 2023.

Author response:
***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. 

| Algorithm | PSNR | SSIM    | LPIPS |
| -------- | ------- |------- |------- |
| DPS |  22.85|0.861 | 0.127 |
| DDRM | 24.33 | 0.860 | 0.120 |
| DMPS | 22.01 |0.803 |0.233 |
| PSLD | 23.99 | 0.787 |0.201 |
| Latent-DPS | 23.14 | 0.784 |0.199 |
| ReSample (Ours) | ***24.67*** | ***0.892*** |***0.093*** |


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

Reviewer comment (Evaluation):
The primary incentive for training the model in a lower-dimensional latent space appears to be a reduction in computational complexity. Nevertheless, the experimental outcomes presented in Table 1 and 2 demonstrate that the proposed algorithm surpasses conventional techniques in terms of reconstruction quality. Could you provide insights into the potential factors contributing to this outcome?

Author response:
***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].


|Anatomical site|ReSample (Ours)|Latent-DPS|
| -------- | ------- |------- | 
|Chest	|***5.36e-5***	|2.99e-3|
|Abdominal|	***6.28e-5***|	3.92e-3|
|Head	|***1.73e-5***	|2.03e-3|

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

Reviewer comment (Novelty):
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 (see references below).

Author response:
***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.

Reviewer comment (Evaluation):
The motivations of the paper are not very consistent with the results (e.g. regarding the advantage of latent sampling on inference speed).

Author response:
***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.

|Anatomical site|ReSample (Ours)|Latent-DPS|
| -------- | ------- |------- | 
|Chest	|***5.36e-5***	|2.99e-3|
|Abdominal|	***6.28e-5***|	3.92e-3|
|Head	|***1.73e-5***	|2.03e-3|

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. 


| Model type | Training Time | Training Memory | FID |
| -------- | ------- |------- | ------- |
| LDM | ***133 minutes***  |***7772MB*** |  ***100.48***   |
| DDPM | 229 minutes | 10417MB | 119.50  |


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

Reviewer comment (Evaluation):
It does not appear clearly in the paper that working in the latent domain improves the computational efficiency despite being a main motivation of this work by the authors. Table 7 is disappointing in that respect. Maybe the authors could emphasize other strengths of their method instead of this one (for instance, the successful results despite the difficulty to sample in constrained settings [1, 2, 3]).

Author response:
***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.

Reviewer comment (Experiments):
While the choice of solving (10) makes sense in practice, checking that the proposed sampling scheme works (or not) when (10) is solved at each iteration would be very insightful. So far, it is only backed by "Because of the continuity of the sampling process, after each data-consistency optimization step, the samples in the following steps can still remain similar semantic or structural information to some extent." How does your method perform if a data fidelity step was performed at each iteration?

Author response:
***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.

| k | PSNR | SSIM    |
| -------- | ------- |------- |
| 1 | 31.74| 0.922|
| 4 | 31.74 | 0.923 |
| 10 | 31.72 | 0.922 |
| 20 | 31.23 | 0.918 |
| 50 | 30.86 | 0.915 |
| 100 | 30.09 | 0.906 |
| 200 | 28.74 | 0.869 |

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

Reviewer comment (Novelty):
The authors mention that (Rout et al., 2023) proposes a similar method; adding few sentences stressing the similarities and differences with the proposed method may be insightful.

Author response:
***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.

Reviewer comment (Novelty):
The theoretical derivations in Appendix appear in a more general form in [6]. Adding the reference when introducing Tweedie's formula may be of interest.

Author response:
***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.

Reviewer comment (Evaluation):
Adding PSNR metrics when displaying images would be informative (e.g. in Fig. 3 and 4, but maybe more importantly in the appendix where there is no page limit).

Author response:
***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)

Reviewer comment (Presentation):
Figure 6 is not very clear since the plot on the right does not correspond to the problem displayed on the left. Furthermore, the legend states "ReSample frequency" while the plot states "# of iterations", maybe the graph should either choose frequency in x-axis or the label should be updated.

Author response:
***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)

Reviewer comment (Presentation):
If I understand Algorithm 1 correctly, the "if t \in C" condition is not activated before Stage 2 of Fig. 16. In this context, I do not understand how information about the measurement can be visible in Stage 1. Isn't Fig. 16 misleading with that regards?

Author response:
***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).

Reviewer comment (Experiments):
More generally, is Stage 0 really necessary? I would naively that for solving eq. (10) and sampling from that may provide a good initial step, allowing to reduce the number of iterations.

Author response:
***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.

Reviewer comment (Experiments):
How do you scale up the algorithm to high-resolution images?

Author response:
***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.

Reviewer comment (Theory):
To solve equation 10 - the L2 least square, I assume you are using gradient descent or conjugate gradient descent, one possible way to accelerate it is to use close form expression on D(z) if A is a linear operator, Do you have any thoughts on this path?

Author response:
***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.

Reviewer comment (Experiments):
Now, I assume the linear operator A is known, but for deblurring task, the system operator can be unknown, can you elaborate on how you could adopt your algorithms to this scenario?

Author response:
***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.

Reviewer comment (Experiments):
In the algorithm, at time point $t$, the data consistency update is done with $z_{t-1}$, rather than $z_t$ which is differ from prior works of diffusion-based inverse problem solver including DPS, DDNM and so on. In fact, to compute as in the algorithm, we need to feed-forward the diffusion model twice for each time step, which lead to two times longer sampling time. I think this is the reason that the paper emphasize the *partial* hard data-consistency update works well in the section 3.1. What if we use $z_t$ from the previous step without additional unconditional DDIM step? It seems reasonable because the mean of previous $z_t$ is also follows the unconditional sampling path. If I'm overlooking anything, please pointing it out.

Author response:
***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.

Reviewer comment (Experiments):
How long it takes if we apply the hard data-consistency for every time-steps? Also, can authors provide the quantitative result on it? I wonder how marginal the performance gap between full hard data-consistency and partial hard data-consistency (which is a proposed method) is, as the performance would be further improved if we use ReSample more frequently according to the figure 6 and figure 15.

Author response:
***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.


|Algorithm | k | PSNR | SSIM    | Inference Time|
|---------| -----| ------- |------- |------- |
| | 1 | 31.74| 0.922| 31:20 |
| | 4 | 31.74 | 0.923 | 8:24 |
|ReSample (Ours)| 10 | 31.72 | 0.922 | 3:52 |
| | 20 | 31.23 | 0.918 | 2:16 |
| | 50 | 30.86 | 0.915 | 1:20 |
| | 100 | 30.09 | 0.906 | 1:08 |
| | 200 | 28.74 | 0.869 | 0:54 |


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

Reviewer comment (Presentation):
There is no description of $\gamma$ in figure 15.

Author response:
***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.