Exploring the Design Space of Diffusion Bridge Models via Stochasticity Control

Limited Discussion on Computational Overhead.

• In our experiments, we observed that the function evaluation step during each sampling cycle is the most computationally intensive, taking 241.598 ms, while the remaining operations require only 9.343 ms (measured on an NVIDIA H100 GPU with a batch size of 200 for the edges2handbags generation task). This demonstrates that function evaluation is the dominant factor in computational cost, which is why we rely on the number of function evaluations (NFE) as a key metric for comparing computational loads across models.
• Compared to DDBM, I2SB, and DBIM, our model employs much simpler transition kernels, characterized by linear linear $\alpha_t$ and $\beta_t$, and quadratic $\gamma_t^2$, which should brings simplification in the calculation. Still, we claim that the improvement is trivial, most of the computing resources are spent on function evaluations.

Lack of Detailed Analysis on the Contribution of Each Stochasticity Control Component

• Ablation studies on sampling algorithms, $\gamma_{\max}$, and $\epsilon_t$. These results are presented in Section 4 (Experiments) and illustrated 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.
• 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 7 in Appendix D.
• Enhanced diversity evaluation. In Table 4, we added more evaluation on different samplers and different choices of $b$.

Lack of Clear Design Guidelines for Stochasticity Control

• Guidelines based on ablation studies. Our ablation study on $\gamma_{max}$ demonstrates that the best choices are around $\gamma_{max}=0.125$ and $\gamma_{max}=0.25$, this results further evaluated through Edges2handbags and DIODE ($256 \times 256$) datasets, as shown in Table 3. Our ablation study on $\epsilon_t$ shows that the best choices are around $0.8$ and $1.0$ for different Linear models, even for different $\gamma_{\max}$ and $b$.

• Additional heuristic guidelines. We included more design guidelines in the Appendix D. We have also added discussions on the selection of $\alpha_t$, $\beta_t$, the shape of $\gamma_t$, $\gamma_{\max}$ and $\epsilon_t$, which potentially helpful for future studies and applications.

• Additional ablation studies. We provided the ablation study on the shape of $\gamma_t$ in the Appendix D, although this analysis extends beyond the original scope of this paper.

Lack of Reliability for the Proposed Evaluation Metric

• 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 thoroughly validated the effectiveness of our proposed metric, AFD, for measuring conditional diversity and demonstrated its complementary role to FID, 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.

Questions on Lack of Theoretical Explanation

We added an illustration of the framework for constructing diffusion bridge models, shown in Fig. 1. The parameters $b$, $\gamma_t$, and $\epsilon_t$ govern the stochasticity introduced at three main stages: preprocessing, training, and sampling. Specifically, $b$ determines the noise added to the base distribution during preprocessing, $\gamma_t$ controls the noise introduced into the transition kernel, impacting both training and sampling, and $\epsilon_t$ regulates the noise added to the sampling SDEs, affecting only the sampling stage.

Questions on the purpose of Figure 1

We updated Figure 1 to illustrate our framework. Due to space constraints, we have moved it to the Appendix D and included additional explanations to provide more clarity.

Thank you for your response. The author’s numerical ablation studies have addressed my primary concerns. Therefore, I raise my score to 6.

Unclear innovations.

• We have moved the background material on stochastic process theory to the Appendix B, focusing the main text on our novel contributions. All theorems and formulas in Sections 3 and 4 are derived by us, ensuring that the originality of our work is clearly highlighted.
• This work is the first to derive sampling SDEs for diffusion bridges conditioned on endpoints, specifically addressing the evolution of $q(X_t \mid x_T)$ for arbitrary Gaussian transition kernels. Unlike Stochastic Interpolants, which model the full data distribution $p_0$, our framework directly addresses sampling from the conditional distribution $p_{0 \vert T}$. This novel focus simplifies both training and inference processes, representing a significant advancement in conditional generation techniques.
• Our framework unifies and simplifies other mainstream diffusion bridge models, such as DBIM and I2SB. This generality demonstrates the versatility and practical impact of our approach.
• Our approach is the first to highlight the issue of lacking conditional diversity in diffusion bridge models and to resolve it by introducing stochasticity into the base distribution.
• Our proposed sampler is novel, achieving much better results than contemporaneous work DBIM.
• Our score reparameterization is extension work, DDBM has the same reparameterization, we extend their work to fit any given translation kernel.

Questions on why the proposed stochastic control outperforms other methods.

• Ablation studies on sampling algorithms, $\gamma_{\max}$ and $\epsilon_t$. We added ablation studies on different sampling algorithms, $\gamma_{\max}$, and $\epsilon_t$, with results presented in Section 4 (Experiments) and illustrated in Figure 3.

• Analysis of training stability. We added new experiments to analyze the training stability introduced by the proposed reparameterization. These results, which also provide insights into numerical stability, are shown in Figure 6 in Appendix D. Enhanced diversity evaluation. To strengthen diversity evaluation, we included DBIM in Table 4 for a direct comparison with our SDB model and other baselines.

• More design guidelines. We added more design guidelines in the Appendix D. Although these experiments extend beyond the original scope of the paper, we have also added discussions on $\alpha_t$, $\beta_t$, and different shapes of $\gamma_t$ in Appendix D to provide a more comprehensive analysis.

Simplistic experiments, generalization of the model and validatation of AFD.

• Generalization issue. Our experiments are conducted on the Edges2handbags and DIODE datasets, which are widely recognized benchmarks for image-to-image translation tasks. These datasets were chosen because they are commonly used in the literature, making it easier to compare our results with other models and evaluate the effectiveness of our approach. It is important to note that EDM and I2SB—which are special cases of our proposed framework—have already been shown to perform effectively across a wide range of real-world applications. This indirectly demonstrates the generalization capabilities of our framework, as it encompasses these state-of-the-art methods while offering additional flexibility and enhancements. We acknowledge the importance of testing on larger-scale and more complex datasets to better reflect real-world scenarios. While this is beyond the scope of the current paper, in the future we will apply our model to more applications, e.g., brain imaging.

• 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.

• Average Feature Distance. To further validate the robustness of our metrics, we included additional visual results to illustrate that a larger AFD corresponds to greater diversity in colors and textures, aligning with human judgment. We thoroughly validated the effectiveness of our proposed metric, AFD, for measuring conditional diversity and demonstrated its complementary role to FID, 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.

Questions on converting the sampling process into an ODE process.

Yes, it is. In our framework we can simply let $\epsilon_t = 0$ to convert the sampling SDE to ODE.

Dear Reviewer PqyK:

We greatly appreciate your careful reading and thoughtful feedback. Based on your suggestions, we have added clarification on our innovations, ablation studies, design guidelines and new diversity metric test. Please let us know if our responses fully addressed your concerns. We are always open to further discussion and welcome any suggestions to further improve our work!

Theoretic contribution

• Firstly, diffusion bridge matching framework [1], [2], [3] [4] build a pinned process $p(\cdot \mid x_0, x_T)$ using SDEs, then do bridge matching to find $p(X_t)$. However, in the I2I translation tasks, we want to build a link between Gaussian translation kernel $p(x_t \mid x_0, x_T) = \mathcal{N}(x_t; \mu(x_0, x_T), \gamma_t^2 I)$ and the conditional process $p(x_t \mid x_T)$. Our framework achieve this goal: Given transition kernel, we can find related sampling SDEs that related to $p(x_t \mid x_T)$. There are two difference compared to diffusion bridge matching framework: 1. we started from Gaussian transition kernel, 2. our sampling SDEs related to $p(x_t \mid x_T)$ instead of $p(x_t)$.

• Secondly, to evolve sampling SDEs by adding score term, diffusion bridge matching and stochastic interpolants [5] framework need to train at least two models. While our framework only need to train one model to approximate $\mathbb{E}[x_0 \mid x_t, x_T]$, the score can be determined by score reparameterization.

• Thirdly, diffusion Schrödinger bridge is different from I2I translation. We take diffusion Schrödinger bridge matching for example: 1. for image translation tasks, we do not have and care about the reference process, it's unnecessary to match reference process, which leads to alternative optimization for two objectives; 2. Diffusion Schrödinger bridge don't consider divergence free term, while I2I translation should consider. 3. Markovian projection in DSBM is to approximate $\mathbb{Q}$ given $\mathbb{Q}(\cdot \mid x_0, x_T)$, while ours framework is to approximate $\mathbb{Q}(\cdot \mid x_T)$ given $\mathbb{Q} (x_t \mid x_0, x_T)$.

Is "specifically addressing the evolution of for arbitrary Gaussian transition kernels" a vague and over claim.

• Our framework gives sampling SDEs for any Gaussian translation kernel with the form $\mathcal{N}(x_t; \alpha_t x_0 + \beta_t x_T, \gamma_t^2I)$, it's better to say "specifically addressing the evolution of $q(X_t \mid x_t)$ for arbitrary Gaussian transition kernels with the form $\mathcal{N}(x_t; \alpha_t x_0 + \beta_t x_T, \gamma_t^2I)$" , but it should be good.

Dataset and visualization

For datasets and visualization, we also visualized depth image to real images with size $256 \times 256$ in the Appendix.

[1] Liu, Xingchao, et al. "Let us build bridges: Understanding and extending diffusion generative models." arXiv preprint arXiv:2208.14699 (2022). [2] Peluchetti S. Non-denoising forward-time diffusions[J]. arXiv preprint arXiv:2312.14589, 2023. [3] Shi Y, De Bortoli V, Campbell A, et al. Diffusion Schrödinger bridge matching[J]. Advances in Neural Information Processing Systems, 2024, 36.

[4] Ye, Mao, Lemeng Wu, and Qiang Liu. "First hitting diffusion models for generating manifold, graph and categorical data." Advances in Neural Information Processing Systems 35 (2022): 27280-27292.

[5] 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.

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.

Thank you for highlighting the significance of Score Distillation Sampling (SDS) in the text-to-3D domain and suggesting its exploration within the scope of our work. While our proposed sampling framework, which focuses on diffusion bridges and conditional diversity, is inherently flexible, exploring 3D applications is beyond the scope of this paper, which could be explored in future research.

Thank you for your response ~

Insufficient Evaluation

• 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.
• 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.
• Verification of Reparameterization. In Table 4, we demonstrate that reparameterizations used in DDBM and EDM can be incorporated into our framework. Their experimental results and discussions further support our arguments. Additionally, we include new experiments analyzing the training stability brought by the proposed reparameterization. These results, including insights on numerical stability, are shown in Figure 7 in Appendix D. This further substantiates our claim that the proposed technique mitigates singularity issues effectively.
• Validation of AFD. We thoroughly validated the effectiveness of our proposed metric, AFD, for measuring conditional diversity and demonstrated its complementary role to FID, 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.

Insufficient qualitative comparison.

• 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.
• Adding baseline DBIM in Fig. 5 and Table 4. While DBIM 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 address this, we implemented DBIM within our framework and conducted comparisons. These results, including experiments with different samplers, $b$, and NFEs, provide further evidence supporting our conclusions. 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.

Question on randomness added to the base distribution.

The convolution of $\pi_{\text{cond}}$ with a Gaussian is equivalent to adding Gaussian noise to the samples. This operation ensures that the randomness is introduced in a way consistent with the probabilistic framework.

DBIM is not discussed in related works.

Thank you very much for your reminder, we added DBIM in our related works.

Typo: Line 212 ODE -> SDE.

Thank you very much for your reminder, the typo has been corrected.

Questions on the setting of $\eta$ in Table 2 and 3.

Thank you very much for your question. Results in Table 2 are based on the DDBM pre-trained model using our proposed sampling algorithm. The optimal $\eta$ for DDBM under our evaluation was found to be 0.3. Results based on our SDB model are based on our SDB model, which achieved the best performance at $\eta=1.0$. To address the reviewer's concern, we have conducted an ablation study for $\eta$ in the context of the SDB model, shown in Fig. 3. Additionally, we have added explanations in the main text to clarify the differences.

Dear Reviewer 1o8f;

We greatly appreciate your careful reading and thoughtful feedback. Based on your suggestions, we have added additional ablation studies and qualitative comparison, we also added more discussion about inovations, design guidance and metric test. Please let us know if our responses fully addressed your concerns. We are always open to further discussion and welcome any suggestions to further improve our work!