Discrete Diffusion Schrödinger Bridge Matching for Graph Transformation

2024-10-02

arXiv.org

AI Summary

Existing diffusion bridge models struggle with distribution shift in high-dimensional discrete state spaces—particularly graph-structured data—and rely on joint distribution assumptions, limiting them to continuous domains. Method: We propose the first continuous-time Markov chain Schrödinger bridge framework tailored for discrete graph data. It avoids joint distribution assumptions and models discrete diffusion via iterative Markovian fitting. Theoretically, we show that node- and edge-wise independent update dynamics correspond to entropy-regularized optimal transport driven by graph edit distance. Results: In molecular optimization, our method significantly improves target properties while minimizing structural modifications, thereby preserving essential chemical features of the original molecule. Empirical evaluation demonstrates its effectiveness and superiority on realistic discrete structure generation tasks.

Technology Category

• Search and Optimization: Mixed Discrete/Continuous Search
• Machine Learning: Graph-based Machine Learning
• Constraint Satisfaction and Optimization: Mixed Discrete/Continuous Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Scalable techniques for the creation, curation, publication, maintenance, and consumption of large, Web-based, structured, reusable, knowledge graphs and ontologies

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""odinger 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 $href{https://github.com/junhkim1226/DDSBM}{here}$.

Problem

• Develops a framework for graph transformation using discrete diffusion
• Solves the Schru00f6dinger Bridge problem in high-dimensional discrete spaces
• Applies the method to molecular optimization in chemistry

Innovation

• Utilizes continuous-time Markov chains
• Extends Iterative Markovian Fitting to discrete domains
• Adapts framework for graph transformation with entropy-regularized optimal transport