Consistency Diffusion Bridge Models

Thank you for your supportive comments and valuable feedback. We would like to address the weaknesses as follows:

W1: This paper's development deviates from the fundamental assumptions of the It'o integral-based SDEs.

A: We respectfully disagree with this statement and would like to address that we are not expressing CM in the form of a stochastic differential equation (SDE). The reverse SDE for denoising diffusion and DDBM in Eqn. (2) and Eqn. (7) are only mentioned once in Section 2 using the definition from established literature for completeness (which can be removed without affecting the definition and development of CM if it might cause a presentation issue). The consistency model for denoising diffusion and DDBM is defined with the PF-ODE in Eqn. (3) and Eqn. (8), respectively, which do not depend on the reverse SDE in Eqn. (2) and Eqn. (7).

W2: a few statements need to be addressed. In particular, to get an unbiased estimation of the expectation in Equation 15, one needs to sample $x_t$ and $x_T$ first and then sample $x_0$. However, the authors state that we can achieve this by sampling $x_0$ and $x_T$ first and then $x_t$. Such an estimator is biased except when we use a single sample.

A: We appreciate the reviewer for rigorously pointing this out. We will revise this to: With a single sample $(x, y) \sim q_{\mathrm{data}}(x, y)$ and $x_t \sim q_{t|0T}(x_t|x_0=x, x_T=y)$, the score $\nabla_{x_t}\log q_{t|T}(x_t|x_T=y)$ can be estimated with $\nabla_{x_t}\log q_{t|0T}(x_t|x_0=x, x_T=y)$.

Thank you for acknowledging our contribution and the valuable comments on our work. We would like to provide our responses as follows (in reverse order for better logical flow):

Q1: The approximation of the score as presented in the current paper (that is evaluating the integrand in only one point $x_0$) is relevant when $t \approx 0$, but completely false when $t \approx T$. Do you have an estimation of the error made in the latter case for tractable cases? Does it have an impact on the training for this part of the process? Have you come up with an idea to correct this intrinsic bias?

A: Thank you for the question. We think that the original CM paper [1] has largely addressed the problem of justifying the use of the mentioned data score for consistency training. Specifically, in Theorem 6 of the CM paper, they reached a conclusion that, when the interval between two timesteps for training CM becomes infinitely small, then the gradient of the consistency training loss and the consistency distillation loss with the true score under Euler ODE solver will be identical. Therefore, the mismatch between the data score (posterior) and true score (marginal) can be addressed by shrinking the interval between the two timesteps in training CM. In practice, this becomes a principle for designing a consistency schedule that ensures $t - r(t)$ to be as small as possible, at least at the end of training. Nevertheless, we admit the inaccurate estimation will have an impact and we think one way to address it is by proposing more fine-grained consistency schedules, which is an interesting and important research direction.

Q2: Could you give further details and intuition on the design of the consistency schedules $r(t)$, which seem to be crucial so that consistency works (even if these reasons are heuristic)? In particular, could you give examples of schedules that may be adapted in theory but do not work in practice? It still unclear to me which design should be the best depending on the type of consistency (either distillation or training).

A: Thank you for the question. First, the design space of training consistency model, such as consistency schedules, loss weighting, and distance metric in consistency loss, has become a specialized area of research, and there are several papers (e.g., [2, 3]) target on investigating better designs for these components. In this work, we mainly adopt the design guideline proposed in [3] and verified its effectiveness in a different setting, i.e., the training of a CM for the PF-ODE of DDBMs. In specific, the consistency schedule should gradually shrink $t - r(t)$ to be small to avoid optimization problems and error accumulation due to such small time intervals as much as possible, which is intuitively elaborated in Sec 3.2 in [3]. In our experiments, we found that the gradually shrinking schedule generally yields better performance than a fixed interval and could be applied to both consistency distillation and training. Hence, our conclusions on the consistency schedule are consistent with the intuition and insights of recent works studying this on diffusion models. However, how to rationalize the design of the process shrinking $t - r(t)$ (e.g., in conjunction with the accuracy of the score estimation discussed in Q1) is still an active area of research.

[1] Song, Yang, et al. "Consistency Models." International Conference on Machine Learning. ICML, 2023.

[2] Song, Yang, and Prafulla Dhariwal. "Improved techniques for training consistency models." ICLR 2024.

[3] Geng, Zhengyang, et al. "Consistency Models Made Easy." arXiv preprint 2024.

Thank you for your valuable comments and feedback on our work. We would like to provide our responses as follows:

W1: The significance of DDBM is exaggerated as a new family of generative models

A: Thank you for pointing this out and we agree that DDBM itself can not be called as a family of generative models. We will revise this to "a new formulation of generative modeling".

W2: Given the consistency methods applied to diffusion models transfer trivially to diffusion bridge models, I am of the opinion that the novelty of the methodological contribution is quite limited

A: While the high-level idea and core techniques of utilizing CM to accelerate the sampling of DDBM might appear direct, we would like to emphasize this work still makes some valid technical contributions to the community.

First, we demonstrate the effectiveness and advantages of the presented first-order ODE sampler for DDBM, which could achieve competitive, or even superior performance to strong baselines with a deterministic sampling trajectory. To the best of our knowledge, there is no other better candidate in previous works. For example, DDBM [1] finds that using the ODE sampler in EDM [2] yielded poor performance and proposed a hybrid sampler that alternates between ODE and SDE steps. And I$^2$SB [3] uses a stochastic posterior sampling scheme. Neither of them can naturally fit the CM paradigm.

Second, our presented unified framework of empirical design space (i.e., noise schedule, network parameterization, and preconditioning) can encompass a wide range of diffusion bridges, which is completely decoupled from their different theoretical premises (e.g., either the method belongs to bridge matching or conditioned bridge matching). As a result, people could efficiently reuse the successful empirical designs of previous diffusion bridges in a unified way according to our framework. For example, in our paper, we retrain an image inpainting model with DDBM's mathematical formulation and I$^2$SB's noise schedule, network parameterization, and preconditioning that has comparable performance to the original I$^2$SB. And we can further build a consistency model on top of it in a unified manner.

W3: I would also question the usefulness of data to data generative models which are deterministic. For any ill-posed inverse problem, practitioners would want to be able to sample from many possible uncorrupted sample given a single corrupted sample, including some example provided here such as for inpainting.

A: We would like to argue that the deterministic sampling process of data to data generative models, such as the presented first-order ODE sampler for DDBM, does not necessarily produce a deterministic sample given a fixed input. For the inverse problem, it does produce diverse uncorrupted samples given a single corrupted sample. This is because of the fact that the PF-ODE of DDBM is only well-defined for $0 \le t < T$. Thus, the valid way to simulate the PF-ODE is to start from $q_{T-\gamma|T}(x_{T-\gamma}|x_T)$ for some $\gamma > 0$, which introduces stochasticity and will result in sample diversity given a fixed $x_T$. In practice, once we have a plausible approximation of initial distribution, the ODE sampler will work similarly to its counterpart in noise-to-data generative models. We also provide a concrete illustration of the sample diversity with our first-order ODE sampler in the image inpainting task on the one-page pdf in the common response.

On the other hand, deterministic samplers may also have other advantages such as improved sampling efficiency (while ensuring sample diversity). For example, in our replicated image inpainting model with DDBM's formulation, the first-order ODE sampler is notably better than a first-order SDE sampler when NFE is small (e.g., $<100$). Another piece of evidence is that I$^2$SB's officially released checkpoint also adopts a deterministic sampler that removes the stochasticity of posterior sampling.

W4: For diffusion bridge models with conditioned terminal point, there is no longer any stochasticity and hence the ODE derived cannot match the marginals of the stochastic generative backward process

A: As we explained in W3 and addressed in DDBM [1], the PF-ODE derived by DDBM is well-defined and models the evolution of $q_{t|T}(x_t | x_T)$ when $0 \le t < T$. The singularity caused by the fixed terminal point can (and should) be removed by using a stochastic initial distribution, and doing so will enable valid ODE sampling and consistency modeling. We will add some notes about this point in our revised version. Furthermore, we find that the consistency losses in the current version allow us to sample $t=T$ from $\mathcal U(\epsilon, T)$ and we will revise the timestep distribution to $\mathcal U(\epsilon, T - \gamma)$ with a pre-specified $\gamma > 0$.

W5: There are a number of related works that could be discussed. In particular Augmented bridge matching, which draws an equivalence between denoising diffusion bridge models and conditioned bridge matching.

A: Thank you for your suggestion. Augmented bridge matching is an excellent work that reveals and discusses how conditioning on the terminal point changes the theoretical properties of bridge matching, notably pointing out that conditioned bridge matching (or DDBM) preserves the empirical coupling. We have already mentioned this work in our initial submission and will add more discussion about it and the bridge matching technique in our revised version.

W6: [32,6] consider regular bridge matching without any connection to an optimal coupling and hence although tractable just a bridge matching model and not a schrodinger bridge. This is a common mistake.

A: We thank the reviewer for pointing this out. We will add the discussion about bridge matching and properly classify these works as bridge matching & conditional bridge matching procedures in the related works section.

I was meaning the reverse process is not Markovian but I do not think this is a problem.

Thank you for the clarification, I think you are correct - apologies! I appreciate it and this is very interesting and a good contribution to the paper. I think this proof should be included in the paper, unless I've missed it somewhere. The DDBM paper does not have this initial stochastic step so the second part of Theorem 1 of DDBM is incorrect.

It is not clear to me that the forward SDE and backward ODE proposed have the same marginals. Consider a brownian bridge x_t between 2 points x_0 and x_T. There is a non zero probability that x_t > L, say p(x_t>L|x_0, x_T)>l>0. The heuristic proposed is supposedly valid for any gamma>0, then in the backward ODE process p(x_t>L|x_T) -> 0 as gamma -> 0 hence backward and forward do not coincide in distribution for all gamma > 0 ,so the argument proposed seems incorrect.

The Kolmogorov forward equation typically requires the markov property which does not hold here. It is not clear to me if this is the issue or something else, I am aware of cases where Kolmogorov equations hold for non markov processes but this typically requires more justification than provided

I would be happy to raise my score if this is cleared up but it seems lack of a valid probability flow ODE is a major flaw in the paper.

Thank you for your valuable comments and feedback on our work. We provide our responses as follows:

W1: The proposed method may not be very surprising. This work looks more like a direct combination of CM and DDBM, which have been introduced in previous works.

A: While the high-level idea of utilizing CM to accelerate the sampling of DDBM might appear direct, we would like to emphasize this work still makes some valid technical contributions to the community that can not be accomplished by simply combining the existing techniques from previous works.

First, we demonstrate the effectiveness and advantages of the presented first-order ODE sampler for DDBM, which could achieve competitive, even superior performance to strong baselines with a deterministic sampling trajectory. To the best of our knowledge, there is no other better candidate in previous works. For example, DDBM [1] finds that using the ODE sampler in EDM [2] yielded poor performance and proposed a hybrid sampler that alternates between ODE and SDE steps. And I$^2$SB [3] uses a stochastic posterior sampling scheme. Neither of them can naturally fit the CM paradigm.

Second, our presented unified framework of empirical design space (i.e., noise schedule, network parameterization, and preconditioning) can encompass a wide range of diffusion bridges, which is completely decoupled from their different theoretical premises. As a result, people could efficiently reuse the successful empirical designs of previous diffusion bridges in a unified way according to our framework. For example, in our paper, we retrain an image inpainting model with DDBM's mathematical formulation and I$^2$SB's noise schedule, network parameterization, and preconditioning that has comparable performance to the original I$^2$SB. And we can further build a consistency model on top of it in a unified manner.

W2: The key motivation and necessity to combine DDBM and CM is unclear.

A: We thank the reviewer for this question. For DDBMs, the benefit of utilizing the paired data during training is not only reducing the NFE requirement for sampling but also further improving the saturated performance given numerous NFE steps over unsupervised diffusion models. For example, the DDBM baseline in [1] needs hundreds NFE steps to achieve the saturated performance in some image translation tasks, surpassing the saturated performance of diffusion models. We show that introducing CM to DDBM can achieve these SOTA results $\sim 50 \times$ faster. In our opinion, this is analogous to why we need CM on diffusion models. On the other hand, even though some diffusion bridges (e.g., I$^2$SB) can achieve decent performance in a few NFE steps (e.g, $\ge 10$ steps), we believe that further reducing $2.5 \times$ NFE steps still makes a valid contribution.

W3: Some common metrics like PSNR and SSIM are not reported in experimental results.

A: Thank you for the suggestion. We will add these two metrics in the appendix of our revised paper. We show the results in the comment for a short overview.

W4: The NFE for the compared baselines in Table 2 and 3 may also need to be noted.

A: Thank you for the suggestion, we will add the NFE of the baselines in the revised paper. Most results are directly taken from the DDBM [1] and I$^2$SB [3] and some of them only have a rough number of NFE steps. We put them in the comment below for a simple overview.

W5: For qualitative results, the comparison with previous DDBM method under the same NFE would be appreciated.

A: We thank the reviewer for the suggestion. Regarding the performance under a few NFE of the I$^2$SB baseline, we want to address that the DDBM (ODE-1) baseline in Table 3 is our replicated version of I$^2$SB on the center image inpainting task. As demonstrated in the answer of W1, we train an image inpainting model with DDBM's mathematical formulation and I$^2$SB's noise schedule, network parameterization, and preconditioning that has comparable quantitative performance to the original I$^2$SB. We put the results of both models under the same NFE in the comment for a simple overview.

As shown in the table, our replicated model can achieve the reported FID of 4.9 in I$^2$SB paper with 10 NFE steps. On the other hand, I$^2$SB only provides checkpoints with a fully deterministic forward process (i.e., the "OT-ODE" in their paper), which does not align with the DDBM formulation. Hence, we choose our version as the main baseline. We also add a qualitative comparison between I$^2$SB and CDBM in the one-page pdf.

Besides, it seems like the mentioned CDDB [4] paper still needs $>50$, even thousands of NFE steps to achieve the saturated performance, which, again, justifies our motivation of introducing CM to achieve comparable performance with much fewer NFE steps.

W6: What is the computational complexity and parameter generalization robustness to retrain the model on a new task?

A: The cost of the paired data training stage is mostly concentrated on the base DDBM model. For reference, training a DDBM from scratch on the DIODE (256 $\times$ 256) image translation task would need $\sim500$k iterations. In our experiment that replicates I$^2$SB, it takes us about 200k iterations to fine-tune a DDBM for center inpainting from an unsupervised diffusion model on ImageNet 256 $\times$ 256, which is far easier than learning the diffusion model on this dataset. Once we have the base DDBM model, the consistency distillation/fine-tuning typically takes $<100$k iterations using the data for training the DDBM, which is more computationally efficient compared to training the base DDBM model itself.

For generalization robustness and efficiency on a new task, it would be an interesting research direction to figure out whether there is, or how to build a foundation model (such as a diffusion model) that can be efficiently transferred to a diffusion bridge on various new tasks by leveraging paired data.

Thank you for your feedback and we would like to provide the response to your follow-up question as follows:

Q1: Note that the PSNR and SSIM results in Table for W3 of DDBM (ODE-1, NFE=2) are largely higher than the proposed method CBD / CBT (NFE=2). Could the author explain more about the reason?

Thank you for the question. We would like to argue that, when evaluating generative models such as diffusion models that hold a certain level of sample diversity, the distortion metrics laid on the pixel space such as PSNR and SSIM may not be the proper major evaluator that aligns with sample quality. For example, DDBM (ODE-1, NFE=2) not only surpasses CBD / CBT (NFE=2) but also surpasses DDBM (ODE-1, NFE=100) in terms of PSNR & SSIM. One should agree that using hundreds of NFEs yields better sample quality than only two NFEs on the image translation task, though NFE=2 has better PSNR & SSIM than NFE=100, which can be justified by the blurry samples under NFE=2. Thus, following the community of diffusion models, we focus on FID as the major metric and claim the accelerated ratio according to this metric, where CBD / CBT (NFE=2) clearly surpasses DDBM (ODE-1, NFE=2). On the other hand, the results in Table for W3 actually further support our claim that CBD / CBT (NFE=2) has comparable sample quality to DDBM (ODE-1, NFE=100) measured by various metrics.

Q2: It is kind of confusing about the argument with CDDB though. Although the CDDB paper also reports that more NFEs can help, the key is to compare in the same setting?

We apologize for the confusion caused by our argument with CDDB. In this paper, the target of the proposed consistency distillation/tuning is to achieve comparable performance in a few NFEs to their corresponding base model under numerous NFEs. Hence, we mainly compare the performance between the distilled/fine-tuned model (i.e., CBD/CBT) and the corresponding base model (DDBM ODE-1) under different NFEs. We agree that comparing different methods in the same setting is important but we think this should be conducted among different distillation/fine-tuning methods under the same base models targeted for accelerated sampling, which, to the best of our knowledge, is still a vacant area for the diffusion bridge community.

Q3: The CDDB paper reports that "20 NFE CDDB outperforms 1000 NFE I2SB in PSNR by > 2 db". Can it suppose that CDDB would be a stronger baseline than I2SB to compare with?

We would like to address that CDDB is actually orthogonal to our method since it is a pure inference technique designed for the sampling fashion of diffusion bridges. The sampling process for CBT/CBD can fit into the inference process where CDDB is applicable and thus these two methods can be incorporated together. Moreover, the CDDB requires access to the Gaussian linear measurement for inverse problems, which might be a bit tricky to act as a baseline for fair comparison.