Direct Diffusion Bridge using Data Consistency for Inverse Problems

Summary

This paper focuses on the diffusion model-based inversion problem. The paper first analyzes the current works and unifies them with the Direct Diffusion Bridges. Then, the authors point out that the data consistency is ignored by current works and propose CDDB. Experiments show that the proposed CDDB works well on various inverse tasks like inpainting, super-resolution, and deblurring.

Strengths

1. This paper summarized the current work in a systematic way. 2. The differences between this work and previous works are clear. 3. The paper is supported well by theorems and experiments.

Weaknesses

I am not so familiar with the related works and it is hard for me to judge the novelty and contribution of this paper, but the logic of this paper is reasonable, so I can only give borderline accept and I will increase my score if more experts support this paper. I am also wondering whether the work can be applied to the current popular textual inversion problem, since its current form only supports from one noisy image to a clean image.

Questions

Please refer to the weakness part.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Some limitations are discusses in the Section 5.

Summary

In this paper, the authors propose Consistent Direct Diffusion Bridge (CDDB) a modification of the Direct Diffusion Bridge (DDB) procedure [1,2] that includes a data consistency term similar to [3,4]. The authors propose two ways to introduce the consistency: the first one forgets about the Jacobian term which is hard to compute in the guidance approximation. The second one (called CDDB-deep) compute the gradient through the observation term. It is more costly but can lead to improved results. The authors briefly justify their method in the spirit of [5] and then showcase the efficiency of their method on several inverse problems (inpainting/deblurring) on ImageNet 256x256. [1] Liu et al. (2023) -- SB: Image-to-Image Schrödinger Bridge [2] Delbracio, Milanfar (2023) -- Inversion by Direct Iteration: An Alternative to Denoising Diffusion for Image Restoration [3] Song et al. (2023) -- Pseudoinverse-Guided Diffusion Models for Inverse Problems [4] Chung et al. (2023) -- Improving Diffusion Models for Inverse Problems using Manifold Constraints

Strengths

* The paper is very well-written and the explanations are very clear. The presentation of the method is sound and the authors did a good job at recalling the basics on diffusion models and I2SB. * I really appreciated the discussion on the difference (and ultimately the equivalence) between I2SB [1] and INDI [2]. This discussion is very well presented. * The experiments are strong. The authors have done a very thorough work investigating several metrics on challenging problems in ImageNet 256x256. * While one could complain that the novelty of the method is relatively low (basically combining I2SB with a guidance term), I think the experiments are extensive enough to warrant publication.

Weaknesses

* I think this is overall a good paper and the next point is not the central problem tackled by the paper but I think it should be clarified (or at least the related remarks should be tamed). The authors spend a few lines deriving the variance preserving condition (Equation (12)) but not a lot is said afterwards. The authors basically refer to [1] to justify their method "Furthermore, subsequent noising process using deterministic and stochastic noises can then be used to ensure the transition to the correct noisy manifold". It is not clear to me that the results of [1] apply directly in the context of I2SB. Also, the current work is not self-contained as it relies heavily on the work of [1] and in particular the conjugate gradient method is never explained in the manuscript. I think that in the related work or in the theoretical section the authors should put more effort into giving details about what is [1] and how to adapt it to the current context (even though this is quite clear from reading the paper). In the experiments, the authors claim that that "CDDB generally has higher speed and stability, possibly due to guaranteed convergence". In the paper, no convergence results are provided. I don't see anything that would suggest that CDDB has guaranteed convergence. Overall, I think that the theoretical analysis is the weak point of the paper. * The authors "exclude PiGDM for the baseline comparison in the deblurring problem, as directly leveraging the pseudo-inverse matrix may result in unfair boost in performances". If this is the case then why not include PiGDM and use CDDB-deep which is also preconditioned with the pseudo-inverse? The fact that PiGDM performs (better?) than the proposed method because it makes a smart use of the pseudo-inverse does not seem like a good reason to remove it from the baseline. [1] Chung et al. (2023) -- Fast Diffusion Sampler for Inverse Problems by Geometric Decomposition

Questions

* Another method that performs transfer tasks is Rectified Flow (albeit in a different context) [1] (or other transport based methods like [2]). Since these methods are not optimised to promote the fidelity to the term $Ax=y$ they should have a higher FID score but they might give better LPIPS, especially Rectified Flow which is deterministic. Can the authors comment on that? * I find the comment on the link with Schrodinger Bridge to be quite interesting but also quite confusing. In SB problems the coupling is not known beforehand. Hence, I don't really see how one could leverage a relation of the form $Ax=y$. * is it also possible to use some preconditioning with the pseudo-inverse in CDDB? (since in CDDB-deep "we find that preconditioning with the pseudo inverse as in PiGDM improves performance" it is a natural question. *Line 87: there is no footnote associated with the superscript 1. [1] Liu et al. (2023) -- Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow [2] Shi et al. (2023) -- Diffusion Schrödinger Bridge Matching

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Adressed.

Summary

This paper studies the existing works about Direct Diffusion Bridges (DDB) with a unified scheme and limitations, and proposes a modified inference procedure that imposes data consistency without the need for fine-tuning, called data Consistent DDB (CDDB), as a new diffusion model-based inverse problem solvers.

Strengths

+ The paper is well written and organized. Especially, the writing of the background section clearly describes and sorts out several important relevant works in a unified perspective. + This paper solves an important problem to study the data consistency problem with diffusion models, which can be used to solve inverse problems generally. + The experiments performed on natural images are comprehensive, and results show satisfying reconstruction quality, with comparison methods in both DIS and DDB methods.

Weaknesses

- The paper studies a series of DDB methods which match the clean data/image distribution with measurements distribution, with the key contribution to add the data consistency term into the inference sampling steps. Actually, this kind of data consistency term has been introduced in quite a few previous works about diffusion model-based inverse problem solving, so called DIS/DDS methods in this paper. Although this paper claims the proposed approach is a generalization of the DDS method, after reading the whole manuscript, one essential question is still not very clear why DDB methods should outperform DIS/DDS methods, if they share a similar scheme from some perspective and both consider data consistency constraints during the sampling steps. - As mentioned above, the idea of updating the clean part for data consistency from DDIM formulation has been introduced in some previous works, such as DDS paper [5] and DDNM paper [39]. Although the paper claims that the proposed approach is a generalization of method in [5], the comparison with these two similar papers is missing in the experiments. - Besides, for the proposed DDB methods, it seems to require the paired data for training the distribution matching, different from the series of DIS methods which generally are trained without paired data. Although the paper discuss the relevance with supervised learning frameworks, the proposed method may need to be also compared with other supervised methods or conditional diffusion methods, which are missing in the current experiments and results. Besides, in the Discussion section, the paper claims the CDDB is flexible and does not have to pre-determine the number of forward passes or modify the training algorithm. This flexibility is also carried with all the DIS methods generally which are used to solve inverse problems based on trained diffusion models, but not the unique characteristics from DDB methods, which may need to be clarified clearly in the paper. - In the proposed CDDB (deep) method, it relies on estimating a pseudo-inverse to preserve data consistency in each sampling step. But for solving non-linear inverse problems, it rarely uses a pseudo-inverse formulation, which may not be straightforward to obtain pseudo-inverse for most nonlinear problems. - As mentioned in the Discussion, in order to show the advantage of DDB methods to match two distributions with data consistency, a great example is image translation problem, which is kind of surprising not included in the scope of this paper. But meanwhile, for image translation problem, since the forward matrix A is not explicit, is it still possible to use the proposed CDDB method to preserve the data consistency in the image translation problem?

Questions

Please see the weaknesses for specific questions. Minor comments: - In Table 2, the best results are not bolded correctly in the PSNR of “pool” images as mentioned in the caption?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors discuss the limitations and societal impact of the proposed method at the end of the paper.

Summary

The paper introduces a unified view of several existing diffusion-model-based methods for solving the inverse problem (i.e., DPS and $\Pi$GDM), and proposes a new approach for this problem called Direct Diffusion Bridge (DDB) which is derived from the DDPM ancestral sampling, and to some extent, links to the Image-to-Image Schrodinger bridge (I$^2$SB). The main idea behind DDB is using a trained neural network $G_{\theta^*}$ to compute a denoised image $\hat{x}\_{0|t}$, which is an approximation of $x_0$, from $x_t$. The paper also proposes an improved version of DDB called Consistent DDB (CDDB), which according to the authors, is more beneficial for reconstruction (of $x_0$).

Strengths

- The unified view between DPS and $\Pi$GDM is interesting though I think it’s not hard to figure out as both methods share the same motivation. - Experimental results seem to support the proposed method there are improvements in both perception and signal-to-noise ratio compared to previous works. However, the fairness in settings of different methods should be clarified in the paper.

Weaknesses

1) The presentation in the paper makes the motivation for the proposed method “Direct Diffusion Bridge” (DDB) unclear. From my view, neither the unified view between DPS and $\Pi$GDM (Section 2.2) nor the I$^2$SB model (Section 3.1) motivates the design of the denoised image $\hat{x}_{0|t}$ in the paper. In fact, it seems like in Section 3.1 the authors try to link the posterior of $x_t$ in the I$^2$SB model (Eq. 7) with the posterior of $x_s$ (s < t) in the DDPM (Eq. 10), which, I guess, leads to the name “Direct Diffusion Bridge”. 2) I think the main contribution *in terms of technique* in this paper is a new way to compute $\hat{x}\_0$ (line 2 in Algorithm). Other parts of Algorithm 1 (lines 3 -> 7) were already proposed in previous works (e.g., DPS, DDS). However, the trained network $G\_{\theta^*}$, which plays an important part in the algorithm, is not described in detail in the paper. It is only mentioned briefly at lines 120 -> 121. Besides, no ablation study was conducted to validate the use of $G\_{\theta^*}$ with other methods for computing $\hat{x}\_0$ from $x\_t$ (e.g., Tweedie’s formula). 3) The authors don’t mention how $G\_{\theta^*}$ can handle $x\_t$ with different time steps $t$ in the paper. Does it take both $x_t$ and $t$ as input? In fact, from my view, $G\_{\theta^*}$ is very similar to the network in the paper Consistency Models [1]. However, I don’t see the authors discuss Consistency Models in their paper. Since Consistency Models can be directly used for DIS (a lot of experiments on inverse sampling can be found in [1]), I am keen to see the empirical comparison between this method and Consistency Models. 4) Please correct me if I am wrong but I cannot find the NFE of the proposed method in the main paper or the supplementary material. Therefore, it is hard to fairly compare the performance of the proposed method with other baselines. [1] Consistency Models, Song et al., ICML 2023

Questions

1) What are the architecture and configurations of the residual network $G\_{\theta^*}$? I couldn’t find any detail about $G\_{\theta^*}$ in the main paper or the supplementary material. Does it depend on time $t$ besides $x_t$ 2) What are the technical differences between the method in this paper and the use of Consistency Models for the inverse sampling problem? 3) In Algorithm 1, is $x\_1$ $y$ or Gaussian random noise? Since I$^2$SB can take the corrupted image $y$ as input, do the authors think their method can do so? 4) In Eq 15, $\mathbb{E}[x_t]$ should be $\mathbb{E}[x_0|x_t]$. 5) The main difference between Algorithm 1 and Algorithm 2 is the computation of $g$. In Alg. 1, $g$ is computed as $\nabla_{\hat{x}\_{0|t}} \log p(y|\hat{x}\_{0|t})$ while in Alg. 2, $g$ is computed as $\nabla_{x_t} \log p(y|\hat{x}\_{0|t})$. In my view, Alg. 2 is *more theoretically correct* than Alg. 1 rather than just being more beneficial for reconstruction as mentioned in the paper in lines (181 -> 183). Thus, I hope the authors could make this point clearer in their paper rather than just mentioning CDDB as an alternative for DDB when the inverse problem is non-linear. Even when the inverse problem is linear, we still need the term $\frac{\partial \hat{x}\_{0|t}}{\partial x_t}$ which is the Jacobian in Eqs. 5, 6. 6) At line 181, the authors should be clear about the “U-Net Jacobians”. Is the U-Net here the noise network $\epsilon$ of the diffusion model or the network $G\_{\theta}$? I guess it should be $G\_{\theta}$ since we are considering $\frac{\partial \hat{x}\_{0|t}}{\partial x_{t}}$. The authors give me the impression that they are trying to hide details about $G_{\theta}$, which in fact plays a critical role in the paper.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

Please refer to my questions above

Answer to rebuttal

I would like to thank the author for the rebuttal. I am satisfied with their answer. I have a few other minor comments and questions. > Having said that, indeed, CMs share important aspects with DDBs, especially considering the inference procedure (i.e. Algorithm 1 of [3]). Since the network is distilled to produce a one-step estimation of the endpoint of PF-ODE trajectory, it can directly produce a clean image [...] (In the reply to the Reviewer 7H1C) I fail to see how this is specific to consistency models? If the network is parameterised to give the `x`-predicition and not the score prediction then the same conclusion can be drawn. Am I missing something specific to consistency models? > However, we would like to note that while schrödinger bridge methods are more versatile and can be applied to more general tasks, e.g. image translation, they are generally unstable especially when scaling to high-dimensional data. To the best of our knowledge, there has been not much progress in using general schrödinger bridge methods for solving inverse problems in imaging. I would like to point out that I2SB actually does compute the Schrodinger bridge if the given coupling is the entropic regularized OT one. In that respect I think the authors could tame their comment. However, I agree with the reviewer that so far the SB framework has received less attention. In the paper the authors claim that I2SB is a "Schrodinger bridge with paired data". This is actually not true. I2SB only solves the Schrodinger bridge problem if and only if the paired coupling is the entropic OT coupling which is usually not the case. While I don't blame the authors for this mistake (I2SB being slightly misleading in that respect), I think it would be extremely valuable to correct the narrative. In particular, I2SB does not solve an OT problem (note that this is not necessarily a problem as for paired setting it is not clear why solving a OT problem would be useful). > That said, we are a bit confused with the reviewer’s comment, as we usually see lower FID and better LPIS at the same time regardless of the distribution matching or DDB with paired data. Sorry if my comment was unclear let me clarify. I was simply stating that Schrodinger bridge methods (which truly solve the OT problem like Rectified flow, DSB or DSBM) are trained to minimize the L2 cost between the samples from the blurry and the samples from the clean distribution, see for example the "straight coupling" property discussed in the Rectified Flow paper. As such one would expect that the similarity between the blurry and clean samples would be greater. I was merely asking what are the thoughts of the authors on this point. I am happy to clarify further.

Comment on the authors' rebuttal

I would like to thank the authors for their rebuttal. It has addressed most of my concerns. The authors' answers clarified my misunderstanding of their paper at first sight. I thought their work was a variant of DIS. After the authors' rebuttal, I can see that their work belongs to the class of bridging methods. The reason why the authors didn't give many details about $G\_{\theta}$ and other training settings is now clear. It is because their proposed method CDDB uses exactly the same training procedure and settings of $I^2SB$ and is only different from $I^2SB$ in the integration of the gradient $g$ (Algorithms 1 and 2) during the generation stage. In fact, using $g$ provides more signals to adjust $x\_{t-1}$ but also has a limitation that it assumes $y=Ax$ and may not be applicable to general settings where $y = f(x)$ with $f$ is an arbitrary function. I decide to raise my score to 5.