Prompt-tuning Latent Diffusion Models for Inverse Problems

Thank you for your thoughtful and constructive review.

W1. Method is complicated. There are many hyper-parameters to tune

A. We agree that there are many hyper-parameters involved in P2L. However, we would like to note that the different choices made for specific inverse problems are merely to achieve the best performance, and not because they are necessary. We can still achieve improved performance as opposed to baselines, which are evident in Table 1 and 3, where we use standard GD with 1.0 step size.

W2. Method is expensive in computation. Reports on the inference time are necessary

A. Thanks for pointing this out. We have now included inference times in Appendix C, Table 5, and included a discussion on the main text of Section 4. Please see general comment C2. We would like to 4. while solving inverse problems can lead to a significant improve in quality. We believe that future investigations will yield methods that can build upon P2L to achieve solvers that are faster.

W3. There is no evidence that the proposed method can solve large-scale inverse problems

A. We agree that including full images rather than patches of results in Figure 3 would be beneficial. In Appendix G, we included some example results.

W4. Range space projection is highly related to DDNM [1]

A. The range space discussed in this paper is the range space of the VAE encoder, which is irrelevant from the range space of the forward imaging operator that DDNM considers. We have clarified this in the revised version. In particular, for understanding of whole pipeline, we have derived a new theoretical justification of each step of the algorithm using an alternating minimization perspective. We believe that this derivation provides a theoretical justification of the projection step and help readers to understand the difference from DDNM.

References

[1] Wang et al. “Zero-shot image restoration using denoising diffusion null-space model” ICLR 2023.

Thank you for your thorough constructive review. Please see our response below.

W1. The writing is poor and the submission is hard to follow.

A. The main focus of this work was to propose a first method that leverages text-prompt embedding for solving inverse problems using latent diffusion models. Moreover, we have improved our manuscript also in the theoretical side from an alternating minimization perspective for a better understanding of why our method works better, and on which component our algorithm is improving the previous solvers.

1. Optimizing the text embedding can be thought of as corrections that are made on the posterior PF-ODE trajectory of reverse diffusion that modifies the embedding to meet the measurement conditions. Under this view, we do not need any assumptions on the independence between the noisy latents and the embedding.

2. The Encoder range space projection can be understood as performing MAP optimization under the VAE prior. Per this view, P2L can be seen as the first to leverage both the diffusion and the VAE generative prior into the posterior sampling scheme, further highlighting the strength of the approach.

These modifications offer better understanding on how the method can be derived and understood, and what type of approximation is used, and similarity and differences from the existing LDIS approaches such as PSLD. In this paper, the focus is more on the proof of concept that prompt tuning can yield improved performance, and that we can achieve better results by incorporating the prior learned through the VAE.

We have updated Section 3. We would appreciate it if the reviewer can point us to any specific section that the reviewers found it unclear or hard to follow.

W2.Why do we need iterative optimization similar to EM?

A. As discussed in General Comment C1, in this revision we have provided a unified theoretical perspective using alternating minimization. Moreover, Eq. (6) and (7) justify the need for prompt optimization.

More specifically, if we don’t optimize for the text, we have subpar performance, which is the main point that we want to make in this paper: prior approaches only optimize for the latent variable. If we do not optimize the latent variables, then we cannot find a solution compatible with the measurements. This is why the optimal performance is obtained when simultaneously or alternatively optimize both text and latent representations.

As stated in the answer to weakness 1, we have made our best efforts to clarify the theoretical implications of the proposed work. The optimization of text-prompt embedding is required as it is an unknown variable that is used throughout the reverse diffusion PF-ODE trajectory. The remaining sampling algorithm that synergistically combines optimization in both the pixel- and the latent-space can be regarded as MAP optimization of a specific loss function by splitting the variables.

W3. The proposed projection is similar to DPS [1]. Why the method is named projection is unclear

A. The proposed projection is completely irrelevant to DPS [1]. Note that in [1] there is no VAE or latent diffusion model, nor no exploitation of the text prompt. The two main contributions of this work are: 1. Exploit the text dependency in inverse problems by jointly optimizing the text representation; 2. Exploit the VAE prior to making the framework work on high-dimensions (through better latent diffusion models).

Regarding the specific proposed projection using the VAE prior, we elaborated on why the proposed projection approach yields good regularization in Section 3.2. The proposed method is called projection, as it clearly projects the latents onto the range space of the encoder.

References

[1] Chung et al. “Diffusion posterior sampling for general noisy inverse problems”, ICLR 2023

[2] Rout et al. “Solving linear inverse problems provably with latent diffusion models”, NeurIPS 2023

[3] He, Linchao, et al. "Iterative reconstruction based on latent diffusion model for sparse data reconstruction." arXiv preprint arXiv:2307.12070 (2023).

[4] Song, Bowen, et al. "Solving inverse problems with latent diffusion models via hard data consistency." arXiv preprint arXiv:2307.08123 (2023).

[5] Song, et al. "Pseudoinverse-guided diffusion models for inverse problems." ICLR 2023

W4. No further analysis on why the proposed projection can target any resolution and only subjective descriptions are proposed.

A. Thanks for pointing this out. We have modified Appendix F to fully elaborate on the analysis. The proposed approach is useful for high-resolution image restoration because we are fully leveraging the power of the pre-trained VAE, and because these VAEs are much more robust to the change in the input resolution. In our method, we are proposing a method that can fully utilize this property, by alternatively optimizing latent and pixel space (through the VAE). In the revised version we have incorporated a theoretical analysis where we leverage the VAE prior within the formulation (now pointed out in Section 3.2).

W5. The practicality of the proposed method is limited as it is slow.

A. Please see general comment C2. We agree that P2L requires many iterations. However, we would like to emphasize that P2L is more a proof of concept that optimizing text-prompt embeddings while solving inverse problems can lead to significantly better results. We believe that future investigations will yield methods that can build upon P2L to achieve solvers that are faster. Also, we would like to mention that there are certain use cases that do not require a fast solver but they require to reach the best possible quality (For example, one could use a "slow" method to generate the best possible result on on real data – where we don't know the ground-truth or reference, and later train a supervised method using the paired data).

W6. Lack of experiments on 256$\times$256 data provides less evidence. Experiments on 512$\times$512 is unreliable. Degradation in performance of DIS could be due to sensitivity of hyperparameters.

A. We respectfully disagree. When running all DIS baselines, we did careful hyper-parameter tuning to yield the best possible result out of all methods. The reason for DIS underperforming on 512$\times$512 resolution is because pixel diffusion models that directly target this resolution are often low quality. This is why latent diffusion models and cascaded models dominate the mainstream when targeting image sizes beyond 256$\times$256. In this work, we specifically focus on solving inverse problems on higher resolutions beyond 256$\times$256, which is why we proceed on the latent space (i.e., Stable Diffusion).

In all prior works that tried to leverage latent diffusion models as priors for solving inverse problems [2,3,4], some experiments involve newly training a VAE+LDM from scratch for 256$\times$256 resolution, specified for a dataset that they are considering. This is not what we aim for, as explicitly pointed out in the last paragraph of Section 1: we are interested in using a foundational generative model that can operate on a fully general domain.

W7. Comparison with $\Pi$GDM [5]

A. Please see modified Table 2, 8. We additionally compare against $\Pi$GDM [5].

References

[1] Chung et al. “Diffusion posterior sampling for general noisy inverse problems”, ICLR 2023

[2] Rout et al. “Solving linear inverse problems provably with latent diffusion models”, NeurIPS 2023

[3] He, Linchao, et al. "Iterative reconstruction based on latent diffusion model for sparse data reconstruction." arXiv preprint arXiv:2307.12070 (2023).

[4] Song, Bowen, et al. "Solving inverse problems with latent diffusion models via hard data consistency." arXiv preprint arXiv:2307.08123 (2023).

[5] Song, et al. "Pseudoinverse-guided diffusion models for inverse problems." ICLR 2023

Dear reviewer fyie,

As the deadline for the reviewer-author discussion phase is fast approaching (there is only a day left), we respectfully ask whether we have addressed your questions and concerns adequately. We would be happy to clear up any additional questions.

Dear reviewer fyie,

We believe that we have addressed the concerns that you have raised. Specifically,

1. We have clarified the theory of P2L by deriving the algorithm through the lens of alternating minimization. This also highlights the need for iterative optimization needed for optimizing each component.

2. We further analyzed on why P2L can solve arbitrary-resolution image restoration problems and included full-resolution results.

3. We made sure that the runtime of P2L is discussed and compared against the baselines. We also highlighted the impact of P2L as the first proof-of-concept and for use in cases where speed is not the primary issue.

4. We included additional comparisons against the baseline that the reviewer suggested and clarified why the experiments were conducted on 512x512 resolution.

We would like to gently remind you that the end of the discussion period is imminent. We would appreciate it if you could let us know whether our comments addressed your concerns.

Best regards, Authors

Thank you for your positive and constructive comments. Please see our response below.

W1. Relation to PSLD [1]

A. Thank you for the thoughtful comment! We would like to clarify the experiment conducted in section 3.2. in the original submission. In this experiment, there is no diffusion component. We only iterate the encoding-decoding using the VAE used in Stable Diffusion v1.4. Hence, the experiment is not showing that PSLD always diverge. In fact, there is no inverse problem here, just the repeated application of the VAE encoder and decoder. Even when additionally using the gluing objective in between the encoding-decoding procedure, we still found that the encoding-decoding does not provide a fixed point solution. That being said, we also agree with the reviewer’s point on potential miscommunication to the readers, so we have removed the section in the revised version.

W2. Gluing projects onto measurements. Trade-off between gluing and proximal optimization?

A. We agree that is an important issue. Table 4 of the modified manuscript discusses that the gluing objective is indeed fast projecting to the measurements. As mentioned in the paper, the proposed proximal optimization step can be done with negligible computation overhead, nearly similar to that of the projection. We newly include the comparison of computation cost in the manuscript. Please refer to the modified final paragraph of Section 4.

W3. Choice of $\tau$? Text-prompt embedding is solved by multiple-step optimization

A. Thank you for pointing this out. Indeed, in prompt embedding optimization, we are solving an optimization problem rather than taking a single gradient step. We modified section 3.1. to give a better explanation about the method. Please consult general comment C1. Also, we include the comparison in the compute time for each algorithm in our manuscript (Table 5).

W4. Imprecise description about PSLD [1]

A. Thanks for pointing this out. We have modified the description to be more precise.

W5. Analysis of the proposed algorithm missing

A. Thanks for the constructive comments. As discussed fully in General Comment C1, rather than using linear model, we have now derived each step of the algorithm from an alternating minimization perspective using VAE prior. This derivation clearly explains why each step of the algorithm is necessary and what type of approximation is required to arrive at the simplified solution. Although we could not provide an exact convergence analysis due to the associated approximation at each step, we believe that the unified framework explains the similarity and difference between our method and existing approaches such as PSLD.

Q1. Assumption on independence

A. In the revised version we have removed this assumption. Please see general comment C1 and modified section 3.1.

Q2. Footnote 3 is not the finding of this paper

A. The footnote was used for emphasizing that using the DPS for pixel- and latent- diffusion models lead to different consequences, and not to propose a new idea. That said, we agree that this may be redundant. The footnote was removed.

Q3. How many $\mathcal{C}$ per $\mathbf{z}_t$ update? Runtime analysis needed

A. Please see general comment 2. We now discuss the compute time that scales with the number of update steps used for $\mathcal{C}$ in the main paper. Please see Table 5.

Q4. $K$ should be inside the optimization block

A. Corrected. We appreciate the careful reading!

References

[1] Rout et al. “Solving linear inverse problems provably with latent diffusion models”, NeurIPS 2023

Thank you for your positive and thoughtful comments. Please see our response below.

W1. The paper lacks a theoretical analysis of, for example, convergence.

A. Please see our general response C1. While we limit the scope of our work and do not consider analysis of convergence, we have significantly extended our explanations on the theoretical motivations of P2L using an alternating minimization framework and clarified the approximation at each step. We also clarified the link to the existing LDIS approach such as PSLD. We believe that this extension adds much value to our manuscript, and we hope the reviewer also finds this interesting.

W2. The results are not very promising.

A. We would like to respectfully disagree that the improvements are subtle. It is widely known that PSNR (distortion) and FID (perception) are two metrics that have trade-off [1,2]. When one opts for better PSNR, it leads to worse FID and vice versa. In this regard, the fact that our method keeps the PSNR level and improves FID by a large margin shows that P2L is actually pushing the pareto frontier forward.

1. The performance gains are mostly seen in the FID score: perceptual quality. The information from the measurement may be sufficient for minimizing the distortion, but insufficient to achieve the best perceptual quality. The text embedding that we optimize for fills in this gap, as can be clearly seen in -2.4, -0.9, -11.7,-4.0 improvements in the FID score.
2. The take-home message from Table 4 (now Table 3) is that the proposed components in the paper are synergistic and provides a better trade-off, and do not hamper each other. Moreover, similar to (1), the perceptual quality measured by FID does significantly increase.
3. The compute time required for latent diffusion model-based inverse problem solvers including P2L and all the baselines are almost perfectly linear on NFE. The difference between the compute time for using projection and proximal optimization is marginal. Moreover, while the difference between gluing and our proximal calibration may be not as pronounced on smaller noise levels, it gets more significant as we increase the noise level.

Q1. Text corresponding to learned prompt?

A. Thank you for the thoughtful comment. Indeed, it would be very interesting and helpful if we could find out the “learned prompts” after the optimization. However, this is not possible as we leverage CLIP encoder, which does not have a corresponding decoding architecture. That said, using a diffusion model that leverages e.g. T-5 encoder which has a corresponding decoder, we believe that the learned prompts would be more explainable, which we leave as a future direction of research.

Q2. Which 'OptimizeEMB' do you use in algorithm 4?

A. The ‘OptimizeEMB’ function is defined in a separate algorithm. In order to clarify this, we include Algorithm 2 in the appendix which corresponds to the definition of the ‘OptimizeEMB’ function with annotations. The same function is used for both Algorithm 1 and 3, which are the vanilla, and the adam versions of P2L.

References

[1] Blau, Yochai, and Tomer Michaeli. "The perception-distortion tradeoff." CVPR 2018.

[2] Delbracio, Mauricio, and Peyman Milanfar. "Inversion by direct iteration: An alternative to denoising diffusion for image restoration." TMLR (2023).