Discrete Diffusion Schrödinger Bridge Matching for Graph Transformation

Authors: Jun Hyeong Kim, Seonghwan Kim, Seokhyun Moon, Hyeongwoo Kim, Jeheon Woo, Woo Youn Kim
Primary area: generative models
Keywords: Schrödinger Bridge, Discrete Diffusion Model, Molecular Optimization

Abstract:
Transporting between arbitrary distributions is a fundamental goal in generative modeling.
Recently proposed diffusion bridge models provide a potential solution, but they rely on a joint distribution that is difficult to obtain in practice.
Furthermore, formulations based on continuous domains limit their applicability to discrete domains such as graphs.
To overcome these limitations, we propose Discrete Diffusion Schrödinger Bridge Matching (DDSBM), a novel framework that utilizes continuous-time Markov chains to solve the SB problem in a high-dimensional discrete state space.
Our approach extends Iterative Markovian Fitting to discrete domains, and we have proved its convergence to the SB.
Furthermore, we adapt our framework for the graph transformation, and show that our design choice of underlying dynamics characterized by independent modifications of nodes and edges can be interpreted as the entropy-regularized version of optimal transport with a cost function described by the graph edit distance.
To demonstrate the effectiveness of our framework, we have applied DDSBM to molecular optimization in the field of chemistry.
Experimental results demonstrate that DDSBM effectively optimizes molecules' property-of-interest with minimal graph transformation, successfully retaining other features. Source code is available [here](https://github.com/junhkim1226/DDSBM).

Decision: accept (poster)
Reviewer scores: 8, 3, 6

Reviews and discussion:
<thread>
creation date:
1730559993491

thread_type:
official_review

<thread_object>
creator: reviewer_b65h

<content>
summary:
Adapt diffusion Schrodinger bridge matching to the discrete domain and instantiate it through molecular optimization

strengths:
1. This method adapts diffusion Schrodinger bridge matching which opens a new avenue for this task.
2. The experimental setting for molecular optimization is a very relevant downstream task for the discrete domain, and it builds up some baselines for that task.
3. The work provides corresponding theoretical guarantees.

weaknesses:
1. Globally, the structure is complete but the writing is not easy to follow. The writing can be enhanced at least in the following way:
* The square brackets around the citation should not appear twice.
* The name of section 2.2 'SOLUTION METHOD' needs to be optimized.
* The supplementary material also contains the body part of the paper.
* The structure section 2/3 is a bit confusing:
  * There should be an introduction to the continuous methods (IPF, IMF) as prior knowledge since the method highly dependent on them, then a transition to the method for the discrete domain. Here, section 2 starts by discussing the discrete domain and introducing the definitions together, which makes it a bit confusing which parts come from previous work and which part comes from this work.
  *  Consequently, the insights about the transition from continuous domain to discrete domain are missing.

2. Under the claim of proposing a discrete model for DSBM, the results only contain 2 molecule datasets with few supportive ablation experiments being given. It will support better the contribution of this work to give toy results for examples except for molecule optimization such as toy datasets with discrete features, or results on other types of graphs, or more ablations, or pure generation results (as in continuous SB papers), or more visualizations of the resulting optimization chain. No need to have them all, but similar supportive evidence would definitely make the experiments more complete.

questions:
How significantly does the choice of graph-matching algorithm impact the overall performance of the method?
If it has a substantial influence, there is another reason that experiments for some simpler modalities without such complex matching may help to clarify the contribution.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730616765272

thread_type:
official_review

<thread_object>
creator: reviewer_7pxr

<content>
summary:
This paper introduces discrete diffusion Schrödinger bridge model using continuous-time Markov chains. The proposed model was validated on molecular optimization problem.

strengths:
Strong performance compared to diffusion bridge model.

weaknesses:
Overall, this paper failed to distinguish its own contributions from existing work and to compared itself with relevant works. First of all, the proposed discrete diffusion Schrödinger bridge matching (DDSBM) highly relies on continuous-time Markov chains (CMTC). However, this paper failed to refer flow matching models [1, 2] which introduce CMTC for discrete data domains. 

[1] Campbell, Andrew, et al. "Generative Flows on Discrete State-Spaces: Enabling Multimodal Flows with Applications to Protein Co-Design." Forty-first International Conference on Machine Learning.

[2] Gat, Itai, et al. "Discrete flow matching." arXiv preprint arXiv:2407.15595 (2024).

And, this is not the first paper that solving Schrödinger bridge (SB) problem in discrete state spaces. There are works for discrete SB [3, 4]. However, the relevance and difference between the proposed method were not discussed.

[3] Harchaoui, Zaid, Lang Liu, and Soumik Pal. "Asymptotics of discrete Schrödinger bridges via chaos decomposition." Bernoulli 30.3 (2024): 1945-1970.

[4] Chow, Shui-Nee, et al. "A discrete Schrodinger bridge problem via optimal transport on graphs." calculus of variations 20.33 (2021): 34.

questions:
1. What is the relevance and difference between the proposed method and existing discrete SB models in Weakness?

2. How does the performance of molecular optimization compare with discrete flow matching models?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730664406569

thread_type:
official_review

<thread_object>
creator: reviewer_az9a

<content>
summary:
This work presents Discrete Diffusion Schrödinger Bridge Matching (DDSBM), a framework that adapts continuous-time Markov chains (CTMCs) to address the Schrödinger Bridge (SB) problem within high-dimensional discrete spaces. Using the Iterative Markovian Fitting (IMF) technique, DDSBM enables optimal graph modifications by minimizing structural changes, which is particularly valuable for molecular optimization tasks in drug and material discovery. In this context, the framework aligns with graph edit distance (GED), allowing efficient property-driven modifications while preserving the molecule's structure. Experiments show DDSBM achieves minimal structural shifts and successfully maintains desirable molecular properties, outperforming traditional graph translation methods.

strengths:
- The Schrödinger bridge problem is well studied in the continuous-state diffusion literature, but has, to the best of my knowledge, not yet been applied to the discrete setting. The paper nicely bridges this gap, hence proposing a valuable contribution to the discrete diffusion literature.
- The paper looks sound and technically strong.
- Empirical results seem to indicate that this approach outperforms previous baselines by a large margin.

weaknesses:
- The choice of the FCD metric the evaluate the graph structure distributions is surprising, since it is typically used to assess distribution learning from a chemical perspective. Why not use NSPDK for example ?
- I’d like to see some error bars on your results.
- The method description is quite technical, and might be hard to understand for the average graph machine learning practioner. It could be worth it to include a short tutorial on the Schrödinder Bridge problem to ease the reader understanding.

questions:
- What is the Hungarian algorithm ? You mention it in section 3.3. I'd appreciate if you could drop a reference or give some details in appendix.
- The baselines outperform your approach in terms of Validity. Even though the results on FCD and NLL indicate that your method is stronger, I'd appreciate if you could comment on this metric.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732542768434

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
First of all, we thank all reviewers for their thoughtful and constructive reviews.

We are pleased to share the revised version of our manuscript, which has been updated based on the valuable feedback provided. The changes are highlighted in **blue**. We combined the appendix with to the main manuscript. Below, we briefly summarize the key changes:



**[About Readability]**

1. **Revised the flow from Section 2 to Section 4:**
    
    We reconstructed our manuscript to clearly distinguish previous works and our contributions. The manuscript now first explains the related works (**Section 2**) and theoretical backgrounds about Schrödinger Bridge (SB) problem in the continuous space (**Section 3.1**). Following that, we outline about how one can extend SBP to discrete state space (**Section 3.2**). Building on these, we finally present our main contributions, which are proving the convergence of IMF iteration in discrete setting (**Section 3.3**), and its practical application methods in discrete space (**Section 4**).  
    
2. **Added a tutorial on the Schrödinger Bridge problem and Iterative Markovian Fitting (IMF):**
    
    For general readers, we added the tutorial in **Appendix A** with explanations of the theoretical backgrounds and linked each part of the tutorial directly to the relevant sections in the manuscript. We also included some insights that SB in discrete space can be intuitively thought of as finding coupling that best preserves the initial states, which aligns with the main motivation behind developing DDSBM.
    
3. **Clarified our contribution:**
    
    Our major contribution lies in introducing DDSBM to solve SBP in discrete state spaces and applying it to the graph transformation problem. However, the previous manuscript lack of sufficient explanation. To address this, we have additional explanation for our contribution in the beginning of **Section 3.3** and **Section 4**. We also explicitly highlighted that our key contribution lies in introducing a diffusion model-based generative framework capable of tackling high-dimensional graph transformation tasks, which has not been addressed in prior works. (**Section 3.3, line 193-197**)

We hope that these revisions adequately address the reviewers' concerns and enhance the quality and clarity of our paper. We are grateful for the time and effort the reviewers have invested in providing their insightful feedback.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732543606782

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Additional Experiments and Results

comment:
For the experiments, we would like to express our sincere gratitude to the reviewers for their valuable feedback and constructive suggestions, which have significantly improved our work on DDSBM. We have carefully addressed all the comments and updated the manuscript accordingly. Below, we outline the major changes made in the revised version:

1. **Error Bars on Experiments:** 

    We have reran our main experiments on the ZINC and Polymer datasets three times to calculate the average performance and standard deviations. This enhancement provides statistically more reliable results. Please see **Table 1**, **Table 2**, **Appendix Table 5**, and **Appendix Table 6** for the updated results.

2. **Additional Metrics in Experiments:** 

    We have included the NSPDK (Neighborhood Subgraph Pairwise Distance Kernel) metric in all experiments, as suggested. Previously, we described FCD (Fréchet ChemNet Distance) as a graph structural metric. With the addition of NSPDK, we have reclassified FCD as a chemical property metric for a clarity. We appreciate the suggestion of a more appropriate metric.
3. **Ablation Studies:** 

    We have added extensive ablation study results and analyses in **Appendix F.2**. This includes an ablation study on different types of graph matching algorithms. From this experiment, we empirically demonstrated that graph matching can accelerate the convergence of IMF iterations in practice. Additionally, we moved influence of initial data pairing methods on IMF iteration convergence, which was previously discussed in **Section 5.4**, as an another type of ablation for practical acceleration of DDSBM training. 

4. **Unconditional Generation Experiments:** 

    We conducted experiments on the QM9 dataset, commonly used as a benchmark for molecular generative models, and the synthetic Community-20 and Planar datasets frequently used for general graph benchmarks. DDSBM achieved the best or second-best performance in most metrics on both QM9 and Community-20. Furthermore, we consistently showed that DDSBM maintains the overall graph structure effectively, as evidenced by the low NLL (Negative Log-Likelihood) for both initial and generated data. Please see **Appendix G** for detailed results.

5. **Training Hyperparameters:** 

    We have added detailed information on the hyperparameters used during training, including the number of training epochs. **Table 4** provides this information for the main experiments on the ZINC250K and Polymer datasets, and the new **Appendix Table 9** includes the hyperparameters for the unconditional generation tasks.

6. **Trajectory Images:** 

    We have included trajectory images of the generation process for both the main results and the new experiments conducted during the revision. These can be found in **Appendix Figures 16-20**. We believe that visualizing the generation results that were analyzed only with NLL and other metrics would provide an intuitive understanding of the generation process.

7. **Impact of the Number of Function Evaluations (NFEs):** 

    We compared the generation performance (Validity, FCD) of DDSBM and DBM on the QM9 task across different NFEs (5, 10, 20, 25, 50, 100). The results show that DDSBM maintains performance reasonably well even at low NFEs, whereas DBM shows a more pronounced performance degradation under the same conditions. We attribute this to that DDSBM finds optimal generation paths that can be predicted with less NFEs.

We believe that these additional results and discussions have significantly improved the quality of our paper. Once again, we thank to all the reviewers for their time and effort in reviewing our paper.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The authors introduce Discrete Diffusion Schrödinger Bridge Matching (DDSBM), a new method for transforming data in discrete spaces like graphs using continuous-time Markov chains. DDSBM improves key properties of molecules with minimal changes to their structure, preserving other important features.

Strengths: see comments by reviewers.

Weaknesses: The main major weakness after rebuttal was raised by reviewer 7pXr, who requested comparisons with the latest discrete flow matching methods. 

Overall the paper was well received the rest of the reviewers.