Exploring the Design Space of Diffusion Bridge Models

We thank the reviewer for the detailed feedback. We are pleased to hear that you believe the paper introduces new ideas and makes useful contributions.

Our claims are mostly combinations or reformulations of existing ideas. Surprisingly, though, there are gaps in the bridge literature that largely stem from the difficulty of combining ideas conceived under different conceptual frameworks (e.g. Stochastic interpolants[1] versus DDBM [2]). We don't believe it is too bold to ask "Have the bridge paths been fully explored?", to which the answer is a simple "no". However, we don't want to give the impression that our paper provides the final grand unified theory of bridge paths, and we will clarify the language accordingly by focusing and narrowing the scope of our statements in several places. Nevertheless, there are significant gaps in the existing literature that our work helps to fill.

The authors merely test a specific gamma schedule and present limited ablations, falling short of a meaningful exploration, even no theoretical validation, why the new gamma schedule works?

We provided additional design guidelines in the Appendix D. The rationale and ablation studies are as follows:

$\alpha_t$ and $\beta_t$. Theoretically, $\alpha_t$ and $\beta_t$ can be freely designed, and future work may explore alternative design choices. However, in this paper, we focus on the simple case where $\alpha_t = 1 - t$ and $\beta_t = t$. The rationale is as follows: consider the scenario where $\alpha_t = 1 - \beta_t$, which represents an interpolation along the line segment between $x_0$ and $x_1$. For the path 

$p_t^{(1)}(x) = \mathcal{N}((1 - \beta_t)x_0 + \beta_t x_1, \gamma_t^2 \mathbb{I})$, where $\beta_t$ is invertible, it is straightforward to construct another path 

$p_t^{(2)}(x) = \mathcal{N}((1 - t)x_0 + t x_1, \gamma_{\beta_t^{-1}}^2 \mathbb{I})$, which achieves the same objective function but uses a different distribution of $t$ during training. Based on this equivalence, setting $\alpha_t = 1 - t$ and $\beta_t = t$ is a reasonable choice. 

The shape of $\gamma_t$. We also conducted an ablation study on $\gamma_t$ with different shapes. Specifically, we assumed $\gamma_t$ has the form $\gamma_t=2 \gamma_{\max} \sqrt{t^k (1 - t^k)}$, the results are shown in Fig. (7). The results indicate that the best performance is achieved when k = 1, which is the exact setting used in this paper.

$\gamma_{\max}$. Our ablation studies on $\gamma_{\max}$ demonstrate that the optimal values of $\gamma_{\max}$ are approximately $0.125$ or $0.25$. Furthermore, the sampling paths corresponding to different choices of $\gamma_t$ are shown in Fig. 8. Adding an appropriate amount of noise to the transition kernel helps in constructing finer details. This setting achieves the best performance across 3 different datasets: Edges 2 handbags, DIODE and ImageNet, demonstrating its generalization.

$\epsilon_t$. We use the setting $\epsilon_t = \eta \left( \gamma_t \dot{\gamma}_t - \frac{\dot{\alpha}_t}{\alpha_t}\gamma_t^2 \right).$ The ablation studies on $\epsilon_t$ demonstrate that the optimal choice of $\eta$ for the DDBM-VP model is approximately $0.3$, while the best choice for the ECSI model with a Linear Path is around $1.0$. Additionally, we present sample paths and generated images under different $\eta$ settings to illustrate heuristic parameter tuning techniques. The results are shown in Figures 10, 11, and 12. Too small a value of $\eta$ results in the loss of high-frequency information, while too large a value of $\eta$ produces over-sharpened and potentially noisy sampled images.

In our experiments, across different tasks, the superior performance is achieved by setting:

$p_{t|0,T}(x_t | x_0, x_T) = \mathcal{N}\left(x_t; (1-t)x_0 + t x_T, \frac{1}{4}\gamma_{\max}^2 t(1-t)\mathbb{I}\right)$

The best results are achieved by setting $\gamma_{\max}=0.125$ and $\gamma_{\max}=0.25$, and the best sampler is achieved by setting $\epsilon_t = \dot{\gamma_t} \gamma_t - \frac{\dot{\alpha_t}}{\alpha_t}\gamma_t^2=\gamma_{\max}^2 / 8$. This setting can be adapted to different kinds of image-to-image translation problems, e.g., edges to handbags, depth image to RGB image, deblurring.

[1] Albergo M S, Boffi N M, Vanden-Eijnden E. Stochastic interpolants: A unifying framework for flows and diffusions[J]. arXiv preprint arXiv:2303.08797, 2023.

[2] Zhou L, Lou A, Khanna S, et al. Denoising diffusion bridge models[J]. arXiv preprint arXiv:2309.16948, 2023.

I still affirm that this work fills some theoretical gaps. However, I am concerned that this writing style is overly exaggerated. Therefore, I maintain my rating unchanged.

Response to Weakness 1

Thank you for your feedback. We agree that noise augmentation in diffusion sampling, as explored in works like EDM and Stochastic Interpolants, shares similarities with our approach. However, a key aspect of ECSI is ensuring the marginal distribution remains consistent by satisfying the Fokker-Planck equation through our derived SDEs, which we believe is a significant contribution.

Additionally, we argue that ECSI’s novelty extends beyond noise scheduling. By modifying the base distribution, we introduce a new strategy to enhance diversity, which is not sufficiently addressed by adding noise alone, as shown in our experiments. This approach is distinct and, to our knowledge, novel, with results demonstrating its effectiveness in improving diversity for diffusion bridge models.

Response to Weakness 2

We appreciate your insights and the opportunity to clarify the distinctions between ECSI and DDBM. While ECSI shares some conceptual similarities with DDBM, its training scheme diverges in key ways, particularly through the integration of noised stochastic sampling and a novel modification of the base distribution. This modification enhances diversity beyond what endpoint conditioning or x0-prediction achieves in the original DDBM framework.

We view ECSI as an advancement akin to an "EDM" version of DDBM, with additional innovations. Specifically: (1) ECSI unifies multiple diffusion bridge models (I2SB, DDBM, DBIM, and EDM), as demonstrated in Table 8; (2) it decouples parameters in the transition kernel, enabling a broader design space; and (3) it supports a more flexible noise schedule during inference. Besides,  we specifically designed a noise scheduling at inference time, i.e., $\epsilon_t= \eta(\gamma_t \dot{\gamma}_t - \frac{\dot{\alpha}_t}{\alpha_t} \gamma_t^2)$, which achieves better results in the experiments. 

Furthermore, We also proposed a novel strategy that modifies the base distribution. Intuitively, this strategy interpolates between standard bridge models and unconditional diffusion models. We found that this strategy diversifies the sampling path in a complementary manner to other stochasticity controls. This strategy significantly improved the sampler diversity.

Response to Weakness 3

In our paper, we provide a rigorous proof for Proposition 1, which establishes the theoretical foundation of ECSI and demonstrates that I2SB [1], DDBM [2], DBIM [3], and EDM can be viewed as special cases within our unified framework, see Appendix B. Due to space constraints, detailed proofs are provided in Appendix B. Please let us know if you had a different type of convergence result in mind that we failed to address in the Appendix.

Response to Question 1

We adopt the same $x_0$ prediction and preconditioning as DDBM [2]. The training details could be found in Appendix F.

Response to Question 2

We acknowledge the importance of minimizing discretization errors in diffusion models. However, like other samplers for continuous diffusion and diffusion bridge models, the ECSI sampler does not fully eliminate sampling discretization error. To our knowledge, no existing sampler for these models can completely mitigate this issue. To reduce discretization error, we recommend increasing the number of sampling steps or adopting higher-order discretization schemes, as these approaches have proven effective in practice.

Response to Question 3

Thank you for your question regarding the marginal distribution of the ECSI sampler. We confirm that the marginal distribution of the ECSI sampler coincides with the true underlying marginal distribution, as established through a rigorous Fokker-Planck equation analysis. Specifically, Lemma B.2 in Appendix B provides a detailed proof demonstrating how the inclusion of a divergence-free term ensures that the marginal distribution remains consistent with the true distribution. We are happy to provide further clarification or expand on this analysis if needed.

[1] Liu G H, Vahdat A, Huang D A, et al. I $^ 2$ SB: Image-to-Image Schr" odinger Bridge[J]. arXiv preprint arXiv:2302.05872, 2023.

[2] Zhou L, Lou A, Khanna S, et al. Denoising diffusion bridge models[J]. arXiv preprint arXiv:2309.16948, 2023.

[3] Zheng K, He G, Chen J, et al. Diffusion bridge implicit models[J]. arXiv preprint arXiv:2405.15885, 2024.

[4] Karras T, Aittala M, Aila T, et al. Elucidating the design space of diffusion-based generative models[J]. Advances in neural information processing systems, 2022, 35: 26565-26577.

Response to Weakness 1 and Question 2, 3

We sincerely thank the reviewer for their insightful comments and for recognizing the novelty of ECSI compared to Albergo et al. [1]. We acknowledge the strong connection between Proposition 4.1 in our work and the derivations for linear interpolants in Albergo et al. [1, Section 4]. We will revise the text to more explicitly emphasize the similarities as well as the conceptual differences.

However, the strong connection of Proposition 4.1 to the above derivations should be explicitly acknowledged, and a detailed discussion of the conceptual differences between ECSI and SI should be provided. Most importantly, the paper lacks a performance comparison between ECSI and the unconditional SI framework as described in Albergo et al. [1]. This comparison is especially relevant since Albergo et al. [1] already provides, for linear interpolants given by $X_t = \alpha(t) X_0 + \beta(t) X_1 + \gamma(t) Z$. a detailed discussion of the choice of $\alpha(t)$, $\beta(t)$, $\gamma(t)$ in [Section 4, 1].

Could the authors clarify the advantages and conceptual differences of ECSI compared to the spatially linear interpolants presented in Albergo et al. [Section 4, 1]?

In short, both SI and ECSI can be seen as special cases of Conditioned SI. A key advantage of ECSI is its efficiency: while SI need to estimate two terms: $\mathbb{E}[x_0 \mid x_t]$ and $\mathbb{E}[x_1 \mid x_t]$, ECSI only estimates $\mathbb{E}[x_0 \mid x_t, y]$ (where $y=x_1$).

First, SI can be extended to a conditional version. Conditioned Stochastic Interpolants build a marginal probability path $p_{t \mid y}$ using a mixture of interpolating densities: $p_{t \mid y}(x) = \int p_t(x_t \mid x_0, x_1) \pi(x_0, x_1 \mid y) dx_0 dx_1$, where $\pi(x_0,x_1 \mid y)$ is a joint distribution with marginals $\pi_{0 \mid y}(x_0 \mid y)$ and $\pi_{1 \mid y} (x_1 \mid y)$. For linear interpolants given by: $X_t=\alpha_t X_0 + \beta_t X_1 + \gamma_t z$. The conditional kernel $p_t(x_t \mid x_0, x_1)$ is  given by a Gaussian distribution: $p_t(x_t \mid x_0, x_1) = \mathcal{N}(\alpha_t x_0 + \beta_t x_1, \sigma_t^2 I), \forall ~ t \in [0, 1]$. Then we can sample from the conditional distribution $p_{0 \mid y}(x_0 \mid y)$ by running a stochastic process $p_{t \mid y} (x_t \mid y)$ from time $t = 1$ to $t = 0$, which is given by the following SDE:

$dx = b(t, x, y)dt + \sqrt{2 \epsilon_t} dW_t, \quad x_1 \sim p_{1 \mid y},$

where the drift term $b(t, x, y)$ is:

$b(t, x, y) = \dot{\alpha}_t \mathbb{E}[x_0 \mid x, y] + \dot{\beta}_t  \mathbb{E}[x_1 \mid x, y] +( \dot{\gamma}_t + \frac{\epsilon_t}{\gamma_t}) \mathbb{E}[z \mid x, y]$

As $y$ represents null conditioning, we recover the original sampler of Stochastic Interpolants. In the drift term, $\mathbb{E}[x_0 \mid x, y]$, $\mathbb{E}[x_1 \mid x, y]$ and $\mathbb{E}[z \mid x, y]$ are unknown, but we only need to estimate two of them, since $\mathbb{E}[x_t \mid x_t, y]=\alpha_t \mathbb{E}[x_0 \mid x_t, y] + \beta_t \mathbb{E}[x_1 \mid x_t, y] + \gamma_t \mathbb{E}[z\mid x_t, y]=x_t$

We can further reduce the number of unknown term by endpoint-conditioning. Here as we replace condition $y$ to be endpoint $x_1$, the term $\mathbb{E}[x_1 \mid x, x_1]=x_1$. So we have:

$b(t, x, y)=\dot{\alpha}_t \mathbb{E}[x_0 \mid x, x_1] + \dot{\beta}_t x_1 + (\dot{\gamma}_t + \frac{\epsilon_t}{\gamma_t}) \mathbb{E}[z \mid x, x_1]$

This is exactly the sampler for ECSI in Proposition 4.1. Therefore, both SI and ECSI can be seen as special cases of Conditioned SI. A key advantage of ECSI is its efficiency: while SI need to estimate two terms: $\mathbb{E}[x_0 \mid x_t]$ and $\mathbb{E}[x_1 \mid x_t]$, ECSI only estimates $\mathbb{E}[x_0 \mid x_t, x_1]$.

Response to Weakness 2 and Question 4, 5

We thank the reviewer for this insightful feedback on the evaluation of diversity in diffusion models, including some works we were not aware of. Below, we discuss the applicability of each method. We find that Vendi score is a good fit for our evaluation, and we will add this metric to our diversity evaluation results.

(1) $\beta$-Recall, Improved-Recall and Feature Likelihood Divergence are not applicable. In our work, we focus on the conditional diversity of images given a certain condition. For example, given a particular edge image, can we generate diverse images with different colors and textures. Therefore, we want to calculate the diversity of the generated set of images without knowing any prior distribution. Therefore, $\beta$-Recall, Improved-Recall and Feature Likelihood Divergence may not be applicable in our case. While we totally agree that the Vendi Score is well-suited in our case. 

(2) A comparison between AFD and other VS. Both AFD and the VS quantify diversity in the feature space of images, using features extracted from the Inception-V3 model. AFD measures the average pairwise Euclidean distance between feature vectors, making it sensitive to outliers. In contrast, the Vendi Score evaluates diversity by computing the effective number of unique feature patterns, based on the eigenvalues of the similarity matrix, emphasizing the overall structural diversity of the feature set. These metrics are complementary, capturing different aspects of diversity. For example, consider two sets of feature vectors: set $a=[2, 1, 1, 1]$, set $b=[1, 1.2, 1.3, 1.6]$, For AFD, the average pairwise Euclidean distance in set $a$ is larger due to the outlier $2$, resulting in $AFD(a)>AFD(b)$. For the Vendi Score, set $b$ has greater diversity due to its more varied feature values, so $VS(a)<VS(b)$.

(3) Additional experiments. We have included an evaluation on different denoisers and samplers using AFD and VS, as shown in the table below. The results demonstrate that ECSI, particularly with the modified base density, consistently outperforms the DDBM checkpoint with DBIM or ECSI samplers across both AFD and VS metrics at varying NFEs (5, 10, 20). 

Response to Question 1

The authors write in the abstract that “the proliferation of techniques based on incompatible mathematical assumptions have impeded progress.” What are these assumptions, and why do they impede the progress of diffusion bridge models for image-to-image translation?

We will re-word this claim in a clearer way. Recent diffusion bridge models excel in image translation but suffer from restricted design flexibility and complicated hyperparameter tuning, whereas Stochastic Interpolants offer greater flexibility but lack essential refinements. We show that these complementary strengths can be unified by interpreting all existing methods within a single SI-based framework.

Response to Question 6

What is the computational cost of calculating the Average Feature Distance (AFD)?

The computational cost of calculating the Average Feature Distance (AFD) is $O(n^2)$, compared to VS $O(n^3)$. AFD is more time efficient, but VS is less sensitive to outliers, as described above.

Response to Minor comments

Is it a typo that the authors write $\bar{f}$, $\bar{g}$ in eq. (1) instead of $f$, $g$? The reviewer would assume with this notation $\bar{f}(t) := f(T-t)$

In the original DDBM paper, they use the same notation for both forward process and reverse process. But we agree that it's better to use different notations for better clarification. We will revise this typo in the revised paper.

There is a reference missing for the derivation of the reverse model in eq. (3) and eq. (4)

There is a reference missing for the derivation of the denoising bridge score matching objective in eq. (5)

We will add appropriate references to the derivations of the reverse model in eqs. (3) and (4), and the denoising bridge score matching objective in eq. (5), citing the relevant source in the revised manuscript.

$\log p_{t \mid 0, T}$ is missing it's input in eq. (5)

line 105, $t x_1$ instead of $t x_0$

(line 113): "by the SDE" instead of "by SDE"

Thanks for the reviewer. We will revise those typos or mistakes in the revised version.

There is no reference in the main paper to the proof of Proposition 4.1 provided in the Appendix.

We will include a reference in the main paper to direct readers to the proof of Proposition 4.1 in the Appendix.

[1] Albergo et al. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797, 2023.

[2] Heuse et al. GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium. NIPS 2017.

[3] Kynkäänniemi et al. Improved Precision and Recall Metric for Assessing Generative Models. NeurIPS 2019.

[4] Alaa et al. How Faithful is your Synthetic Data? Sample-level Metrics for Evaluating and Auditing Generative Models. ICML 2022.

[5] Friedman et al. The Vendi Score: A Diversity Evaluation Metric for Machine Learning. Transactions on Machine Learning Research, 2023.

[6] Jiralerspong et al. Feature Likelihood Divergence: Evaluating the Generalization of Generative Models Using Samples. NeurIPS 2023.

Response to Weakness 1

We provided additional design guidelines in the Appendix D. The rationale and ablation studies are as follows:

$\alpha_t$ and $\beta_t$. Theoretically, $\alpha_t$ and $\beta_t$ can be freely designed, and future work may explore alternative design choices. However, in this paper, we focus on the simple case where $\alpha_t = 1 - t$ and $\beta_t = t$. The rationale is as follows: consider the scenario where $\alpha_t = 1 - \beta_t$, which represents an interpolation along the line segment between $x_0$ and $x_1$. For the path 

$p_t^{(1)}(x) = \mathcal{N}((1 - \beta_t)x_0 + \beta_t x_1, \gamma_t^2 \mathbb{I})$, where $\beta_t$ is invertible, it is straightforward to construct another path 

$p_t^{(2)}(x) = \mathcal{N}((1 - t)x_0 + t x_1, \gamma_{\beta_t^{-1}}^2 \mathbb{I})$, which achieves the same objective function but uses a different distribution of $t$ during training. Based on this equivalence, setting $\alpha_t = 1 - t$ and $\beta_t = t$ is a reasonable choice. 

The shape of $\gamma_t$. We also conducted an ablation study on $\gamma_t$ with different shapes. Specifically, we assumed $\gamma_t$ has the form $\gamma_t=2 \gamma_{\max} \sqrt{t^k (1 - t^k)}$, the results are shown in Fig. (7). The results indicate that the best performance is achieved when k = 1, which is the exact setting used in this paper.

$\gamma_{\max}$. Our ablation studies on $\gamma_{\max}$ demonstrate that the optimal values of $\gamma_{\max}$ are approximately $0.125$ or $0.25$. Furthermore, the sampling paths corresponding to different choices of $\gamma_t$ are shown in Fig. 8. Adding an appropriate amount of noise to the transition kernel helps in constructing finer details. This setting achieves the best performance across 3 different datasets: Edges 2 handbags, DIODE and ImageNet, demonstrating its generalization.

$\epsilon_t$. We use the setting $\epsilon_t = \eta \left( \gamma_t \dot{\gamma}_t - \frac{\dot{\alpha}_t}{\alpha_t}\gamma_t^2 \right).$ The ablation studies on $\epsilon_t$ demonstrate that the optimal choice of $\eta$ for the DDBM-VP model is approximately $0.3$, while the best choice for the ECSI model with a Linear Path is around $1.0$. Additionally, we present sample paths and generated images under different $\eta$ settings to illustrate heuristic parameter tuning techniques. The results are shown in Figures 10, 11, and 12. Too small a value of $\eta$ results in the loss of high-frequency information, while too large a value of $\eta$ produces over-sharpened and potentially noisy sampled images.

In our experiments, across different tasks, the superior performance is achieved by setting:

$p_{t|0,T}(x_t | x_0, x_T) = \mathcal{N}\left(x_t; (1-t)x_0 + t x_T, \frac{1}{4}\gamma_{\max}^2 t(1-t)\mathbb{I}\right)$

The best results are achieved by setting $\gamma_{\max}=0.125$ and $\gamma_{\max}=0.25$, and the best sampler is achieved by setting $\epsilon_t = \dot{\gamma_t} \gamma_t - \frac{\dot{\alpha_t}}{\alpha_t}\gamma_t^2=\gamma_{\max}^2 / 8$. This setting can be adapted to different kinds of image-to-image translation problems, e.g., edges to handbags, depth image to RGB image, deblurring.

Response to Weakness 2

We appreciate the reviewer’s feedback regarding the scalability of ECSI. Our experiments were conducted on three datasets—64x64 Handbags, 256x256 DIODE, and 256x256 ImageNet—which are standard benchmarks widely used in prior works, including I2SB [1], DDBM [2], and DBIM [3]. These datasets were chosen to ensure fair comparisons with existing methods. While we acknowledge the importance of evaluating scalability on higher-resolution or more complex industrial applications, the consistent performance improvements observed across these datasets provide strong evidence of ECSI’s efficiency and quality. To further address scalability concerns, we are open to extending our evaluation to larger-scale, higher-resolution datasets in future work and can include a discussion of this in the paper to highlight potential applicability in such settings.

Response to Weakness 3

We thank the reviewer for highlighting the need for deeper analysis of the interaction between path stochasticity (controlled by $\gamma_t$) and sampler stochasticity (controlled by $\epsilon_t$). In our experiments, we explicitly address this interplay by selecting the sampler stochasticity based on the path stochasticity, as defined by the relation $\epsilon_t= \eta(\gamma_t \dot{\gamma}_t - \frac{\dot{\alpha}_t}{\alpha_t} \gamma_t^2)$, where $\eta$ controls the stochasticity strength at inference time. This is also an innovation of ECSI and this selection has not been demonstrated in previous work. We observed that for different noise levels of path stochasticity (controlled by $\gamma_t$), the best $\eta$ is achived at $\eta=1$. We will include these details in the main text and emphasize our contribution on specifically designing of $\epsilon_t$.

[1] Liu G H, Vahdat A, Huang D A, et al. I $^ 2$ SB: Image-to-Image Schr" odinger Bridge[J]. arXiv preprint arXiv:2302.05872, 2023.

[2] Zhou L, Lou A, Khanna S, et al. Denoising diffusion bridge models[J]. arXiv preprint arXiv:2309.16948, 2023.

[3] Zheng K, He G, Chen J, et al. Diffusion bridge implicit models[J]. arXiv preprint arXiv:2405.15885, 2024.