Exploring the Design Space of Diffusion Bridge Models via Stochasticity Control

Authors: Shaorong Zhang, Yuanbin Cheng, Xianghao Kong, Greg Ver Steeg
Primary area: generative models
Keywords: Diffusion bridge models, image translation, diffusion models

Abstract:
Diffusion bridge models effectively facilitate image-to-image (I2I) translation by connecting two distributions. However, existing methods overlook the impact of noise in sampling SDEs, transition kernel, and the base distribution on sampling efficiency, image quality and diversity. To address this gap, we propose the Stochasticity-controlled Diffusion Bridge (SDB), a novel theoretical framework that extends the design space of diffusion bridges, and provides strategies to mitigate singularities during both training and sampling. By controlling stochasticity in the sampling SDEs, our sampler achieves speeds up to $5 \times$ faster than the baseline, while also producing lower FID scores. After training, SDB sets new benchmarks in image quality and sampling efficiency via managing stochasticity within the transition kernel. Furthermore, introducing stochasticity into the base distribution significantly improves image diversity, as quantified by a newly introduced metric.

Decision: reject
Reviewer scores: 6, 5, 6, 6

Reviews and discussion:
<thread>
creation date:
1730442693883

thread_type:
official_review

<thread_object>
creator: reviewer_o3pd

<content>
summary:
This paper proposes the Stochastic-controlled Diffusion Bridge (SDB) model to improve the effect of noise generated by SDEs, transition kernel, and base distribution on sampling when the existing Diffusion Bridge models (DBMs) are sampling image-to-image translation. The authors introduce stochasticity control (SC) in three areas—sampling stochastic differential equations (SDEs), transition kernels, and base distributions—to improve sampling efficiency, image quality, and diversity in output. In this process, the paper presents a metric for measuring the diversity of output. The SC mechanism mitigates singularities during training and sampling, and introduces flexibility in design, allowing SDB to outperform existing models like DDBM and DBIM on benchmark tasks, achieving lower FID scores, faster sampling, and improved image diversity.

strengths:
**1. Existing Model Shortcomings: Identification and Solutions**  
Recognize the limitations of existing diffusion bridge models as stemming from noise introduced during image-to-image translation sampling. The SC mechanism effectively addresses these shortcomings by introducing enhanced flexibility and precise control over the noise at various stages of the diffusion process, thereby mitigating the inherent limitations of conventional diffusion models. 

**2. Experimental Validation**  
The paper provides extensive quantitative comparisons with multiple baseline models, demonstrating superior performance in FID, sampling efficiency, and conditional diversity on I2I tasks.

**3. Expression Reconfiguration**  
The paper introduces score reparameterization and discretization schemes to address stability issues in previous diffusion bridge methods, offering novel insights into stochastic diffusion processes.

weaknesses:
**1. Limited Discussion on Computational Overhead**  
Although the model achieves lower FID scores with fewer function evaluations, the added complexity of SC mechanisms (especially in tuning multiple noise schedules) may increase computational overhead. A direct comparison of computational resources (e.g., time or memory) required by SDB versus other DBMs would provide a clearer picture of the trade-offs.

**2. Lack of Detailed Analysis on the Contribution of Each Stochasticity Control Component**  
While the authors claim that SC enhances image quality, sampling efficiency, and conditional diversity, there is a lack of detailed analysis on how each component (e.g., SC in the transition kernel, SC in the base distribution) individually contributes to these improvements. Although overall performance gains are evident, the authors have divided Stochasticity Control into three distinct components, yet the extent to which each component contributes to the model's improvement is not adequately demonstrated through ablation studies.

**3. Lack of Clear Design Guidelines for Stochasticity Control**  
The paper introduces various strategies for enhancing performance through Stochasticity Control (SC). However, it lacks clear principles for determining optimal parameters in SC design. For instance, the design of the ϵt value used in SC is largely determined empirically, which complicates the systematic selection of SC parameters suitable for specific tasks. Additionally, the proposed method includes several hyperparameters (e.g.,  αt, βt, γt, ϵt), yet there is no concrete guideline on how these should be chosen. While certain values were empirically set during experiments, it remains unclear how these can be applied in more general cases. To address this, it would be helpful if the authors could provide a sensitivity analysis for these key hyperparameters, showing how performance varies across different settings. Additionally, we suggest the authors share any rules of thumb or heuristics they relied on when selecting these parameters, which would provide practical guidance for applying the method in various contexts.

**4. Lack of Reliability for the Proposed Evaluation Metric**  
The authors propose Average Feature Distance (AFD) as a metric for assessing diversity. However, there is insufficient analysis to establish the reliability and validity of this metric. To substantiate AFD as a credible measure of diversity, the authors should have demonstrated its application across multiple existing baselines and metrics, thereby enhancing its reliability. In particular, there is a lack of comparative analysis and explanation on how AFD interacts with other evaluation metrics and effectively represents trustworthy conditional diversity. To establish the reliability of AFD, it would be beneficial for the authors to compare AFD scores with human judgments of diversity or apply AFD to datasets with known diversity characteristics. Additionally, case studies illustrating how AFD addresses insights that are missed by similar existing metrics would strengthen its credibility as a measure of diversity.

questions:
**1. Lack of Theoretical Explanation**  
The various concepts and mechanisms introduced in the paper, such as Stochasticity Control and Score Reparameterization, are not sufficiently clear or intuitively explained. While the mathematical descriptions provide a theoretical foundation, the paper would benefit from a figure or diagram that illustrates how Stochasticity Control interacts with other components of the model. Including a main figure to depict the overall framework and interconnections within the proposed model would greatly enhance readers' understanding of the core concepts.

**2. Purpose of Figure 1**  
Figure 1 aims to provide an intuitive illustration of how the transition kernels in image-to-image (I2I) translation can incorporate noise through the configuration of multiple hyperparameters, as proposed by the authors. However, since the ultimate goal of I2I translation is to transform the source image into the target image, evaluating only the final translated result may suffice. Figure 1 merely demonstrates the progressive addition of noise, making it challenging to discern the specific impact of these parameters. A clearer explanation of the intended purpose of Figure 1 would be beneficial.

I would be willing to consider increasing the rate if clear responses are provided to the points raised in the Weakness and Question sections.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730654011820

thread_type:
official_review

<thread_object>
creator: reviewer_zibe

<content>
summary:
The authors introduce the Stochasticity-controlled Diffusion Bridge (SDB), a new theoretical framework that expands diffusion bridge design. SDB mitigates singularities during training and sampling, achieving speeds up to five times faster than the baseline while improving image quality and diversity by managing stochasticity in the transition kernel and base distribution.

Due to my limited familiarity with the field addressed in this paper, it is challenging for me to provide in-depth insights. However, I can offer some feedback from a more application-oriented perspective.

strengths:
The Stochasticity-controlled Diffusion Bridge (SDB) demonstrates significant improvements in image quality, sampling efficiency, and conditional diversity, establishing new benchmarks in these aspects. Furthermore, the authors' introduction of score reparameterization and specially designed discretization effectively addresses challenges related to singularities.

weaknesses:
SDS is a crucial loss function in the text-to-3D domain. Could the authors explore similar applications in this paper? It would be beneficial to demonstrate its extension in a 3D context.

questions:
see weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1730707353202

thread_type:
official_review

<thread_object>
creator: reviewer_pqyk

<content>
summary:
This paper introduces the Stochasticity-Controlled Diffusion Bridge (SDB), a framework that enhances diffusion-based generative models by using a flexible, transition kernel-based approach. Unlike standard DDMs, which start from Gaussian noise, SDB can transition between arbitrary distributions, making it more versatile for tasks like image-to-image translation. Key innovations include a Stochasticity Control (SC) mechanism that regulates noise and drift terms, boosting stability, efficiency, and output diversity. Experiments show SDB operates 5× faster than DDBM and achieves an FID score of 0.89 with just five evaluations, setting a new quality benchmark.

strengths:
1. The paper ensures the sampled distributions adhere to a mathematically sound process. Additionally, the proofs offered for Lemma 1 and Theorem 3 validate that the stochastic control can maintain desired marginal distributions across time

2. The method demonstrates strong performance in preserving alignment and detail within conditional sampling tasks. The model captures and retains critical details between the source and target distributions, resulting in high-fidelity outputs that stay true to the given conditions.

weaknesses:
1.  The presentation of stochastic control in sections 3 and 4 mixes well-known stochastic process theory with the paper’s innovations, making it challenging to identify what is genuinely novel. Clarifying how the proposed SC mechanism specifically advances the field would help distinguish the paper’s contributions from established concepts.

2. The paper doesn’t clearly demonstrate why the proposed stochastic control outperforms other methods. While several design choices are introduced, their practical advantages and rationale are not sufficiently justified, leaving the reader uncertain about the improvements gained from these specific approaches.

3.The experiments appear simplistic and don’t convincingly support the claims, especially regarding generalization to real-world applications. Additionally, the diversity metric (Average Feature Distance) and image quality metrics could be better validated on more extensive benchmarks to strengthen the results.

questions:
Is it possible to convert the sampling process into an ODE process by leveraging the Fokker-Planck equation?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730715658013

thread_type:
official_review

<thread_object>
creator: reviewer_1o8f

<content>
summary:
This paper introduces Stochasticity-controlled Diffusion Bridge (SDB), a generalized diffusion bridge-based model framework that enables flexible control of stochasticity. Specifically, SDB considers stochasticity in 1) the transition kernel, 2) sampling SDE, and 3) the base distribution. This is achieved by designing a noise schedule for the transition kernel, regulating the drift term in the sampling SDEs, and adding noise into the base distribution. The paper also proposes a score reparameterization and a discretization sampling scheme to mitigate singularity. Experiments show that the proposed SDB has better image quality, efficiency, and diversity than other baselines in the image-to-image translation task.

strengths:
1. The proposed SDB framework provides a more generalized view of diffusion bridge-based models (DBM). Under the framework, this work highlights the importance of randomness in DBM. It shows that carefully selected randomness in sampling SDEs, transition kernel, and the base distribution can significantly improve sampling efficiency, image quality, and diversity.

2. The proposed score reparameterization and discretization sampling scheme by design is more likely to avoid singularity and instability in training and sampling.

3. In experiments, the proposed method SDB notably outperforms previous DBM methods in quality, efficiency, and diversity, as evaluated by several quantitative metrics.

weaknesses:
1. The main issue of this paper is insufficient evaluation. First, the ablation study is not thorough enough. It is not shown how different choices of $\gamma_{max}$ affect the results. In addition, there's no ablation study to show the effects of the proposed reparameterization and sampling algorithm. Thus, it is unclear how much gain each technique contributes to the final model.  
Regarding the verification of reparameterization, apart from the normal metrics, I also suggest analyzing the numerical stability during training/sampling to show that the proposed technique indeed mitigates singularity.

2. Qualitative comparison is insufficient. In the main paper, only Fig. 4 shows the qualitative comparison between SDB and DDBM. However, a comparison to other methods such as DBIM should also be provided. I also suggest including DBIM in Table 4 so that the SDB is compared to more than one baseline in terms of diversity.

3. The randomness added to the base distribution lacks further explanation. Why convolving $\pi_{cond}$ with a Gaussian rather than adding Gaussian noise to $\pi_{cond}$? What's the kernel size for the Gaussian?

4. DBIM is not discussed in related works.

Typo: Line 212 ODE -> SDE.

questions:
The settings of Table 2 and Table 3 are not matched. For example, Table 2 shows that $\eta = 0.3$ is the best, but Table 3 uses $\eta = 1$. The quantitative scores are also not matched for $\eta = 1$ on edges2handbags. Please explain the inconsistency.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732221277408

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Official Comment by Authors

comment:
We thank the reviewers for their comments. Although reviewers agreed that our more generalized framework highlights the importance of randomness in diffusion bridge models, demonstrating superior performance in FID, sampling efficiency, and conditional diversity on I2I tasks, there were questions about our inovations, how different aspects of our method contribute to the strength of the results, and about the new metric we introduced to measure conditional diversity. 

 We summarize here the main points addressed in responses to individual reviewers.

**1. Clarified inovations**. 

- **Streamlined focus**. Background on stochastic processes moved to Appendix B. All theorems and formulas in Sections 3 and 4 are original, emphasizing the uniqueness of our work.
- **Framework clarity**. Added an illustration in Fig. 1 and clarified theoretical contributions in the conclusion.
- **Refined theoretical contributions**. We clafied our theoretical contributions in the Conclusion.  This is the first study to derive sampling SDEs for diffusion bridges conditioned on endpoints, addressing the evolution of  $q(X_t \mid x_T)$  for arbitrary Gaussian transition kernels. Unlike Stochastic Interpolants (Albergo et al., 2023), which model the full data distribution $p_0$, our framework directly targets sampling from the conditional distribution $p_{0 \mid T}$, simplifying both training and inference and marking a significant step forward in conditional generation techniques. 
- **New insight to denoising diffusion models**. Our framework can also bring new insight for denosing diffusion models. We note that [TrigFlow (Lu & Song, 2024)](https://openreview.net/forum?id=LyJi5ugyJx), a main contribution of a contemporaneous submission, employs the same score reparameterization and pre-conditioning techniques. Notably, TrigFlow can be considered a special case of our framework by setting $\alpha_t = \cos(t)$, $\beta_t=0$, $\gamma_t=\sigma_0 \sin(t)$, $t\in[0,\frac{\pi}{2}]$. This further demonstrates the validity, generalization and forward-looking nature of our theoretical framework. We included this discussion in Appendix E to further strengthen our theoretical contribution.



**2. New experiments and design guidelines**. 

- **Ablation studies on sampling algorithms, $\gamma_{\max}$, and $\epsilon_t$**. We added ablation studies on different sampling algorithms, $\gamma_{\max}$, and $\epsilon_t$. These results are now presented in Section 4 (Experiments), as shown in Figure 3. 
- **Additional evaluation on more dataset and $\gamma_{\max}$**. We include additional experiments with $\gamma_{\max}=0.25$ and DIODE ($256 \times 256$) dataset, as shown in Table 3. 
- **Additional baseline**.  While DBIM (Zheng et al., 2024) is a contemporaneous work and its official code has not been released, we note that its sampling algorithm can be incorporated into our framework. To enhance the diversity evaluation, we have added DBIM to Table 4 for direct comparison against our SDB model and other baselines. In Figure 5, we added a qualitative comparison of DDIM, which illustrates the advantages of our approach.
- **Verification of Reparameterization**. We include new experiments analyzing the training stability brought by the proposed reparameterization. These results, including insights on numerical stability, are shown in Figure 6 in Appendix D.
- **Enhanced diversity evaluation**. In Table 4, we added more evaluation on different samplers and different choices of $b$. 
- **Design guidelines on $\alpha_t$, $\beta_t$ and the shape of $\gamma_t$**. Although the design of $\alpha_t$, $\beta_t$ and the shape of $\gamma_t$ beyond the original scope of the paper, we have also added discussions on design guidelines on $\alpha_t$, $\beta_t$, and ablation study on different shapes of $\gamma_t$. See Appendix D for more details.

**3. New diversity metric test**.

- **Visualized images**. We presented more visualized images generated by different models and compared their FID and AFD in Figure 5 and Appendix F.
- **Validation**. We added new experiments to validate that our diversity metric correctly captures diversity in a scenario where we have ground truth, in Appendix A. To this end, we designed two classes of pseudo-generative models capable of controlling the diversity of the generated images. The performance of these models was further evaluated using both FID and AFD on the ImageNet dataset. 



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

[2]. Kaiwen Zheng, Guande He, Jianfei Chen, Fan Bao, and Jun Zhu. Diffusion bridge implicit models.
arXiv preprint arXiv:2405.15885, 2024.

[3]. Cheng Lu and Yang Song. Simplifying, stabilizing and scaling continuous-time consistency models.
arXiv preprint arXiv:2410.11081, 2024.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a novel theoretical framework, Stochasticity-controlled Diffusion Bridge (SDB), to enhance the performance of diffusion bridge models (DBMs) in image-to-image (I2I) translation tasks. The authors claim that by meticulously controlling the stochasticity in sampling SDEs, transition kernels, and base distributions, they can significantly improve sampling efficiency, image quality, and diversity. They further assert that their framework offers strategies to mitigate singularities during both training and sampling.   The authors report achieving up to 5x faster sampling speeds compared to baseline models, along with lower FID scores, indicating improved image quality. They also claim to achieve new benchmarks in image quality and sampling efficiency by managing stochasticity within the transition kernel. Additionally, they introduce a new metric to quantify image diversity and demonstrate that introducing stochasticity into the base distribution significantly improves this metric.  

Strengths:
- The paper presents a novel perspective on controlling stochasticity in diffusion bridge models for I2I translation. The proposed SDB framework introduces innovative mechanisms to manage noise levels in different components of the model.
- The paper demonstrates strong empirical results, showing improvements in terms of both FID scores and sampling efficiency on standard (small scale) datasets. The proposed method reportedly achieves state-of-the-art results on standard I2I translation benchmarks.  
- The paper provides a solid theoretical grounding for the proposed SDB framework.

Weaknesses:
- Limited Practical Impact: Concerns were raised about the limited scope of the experiments and the lack of evidence demonstrating the impact of the proposed method on real-world applications. The focus on theoretical contributions may overshadow the practical implications of the work, but the theoretical contribution per-se don't seem strong enough.
- Clarity and Novelty: Despite revisions, concerns remain about the novelty of the proposed framework, with some reviewers suggesting it might be a technical variation of existing results rather than a genuinely new approach.  
- Experimental Validation: While the paper presents empirical results, questions were raised about the sufficiency of the experimental validation, particularly the choice of datasets and the depth of the ablation studies.  

Despite the potential of the SDB framework and the reported performance improvements, concerns about the novelty, practical impact, and experimental validation of the work persist. The reviewers' doubts about the core contribution's novelty and the limited evidence of its effectiveness in real-world scenarios.