Discrete Diffusion Schrödinger Bridge Matching for Graph Transformation

TL;DR

This work extends the Schrödinger Bridge framework to high-dimensional discrete spaces by formulating Discrete Diffusion Schrödinger Bridge Matching (DDSBM) using continuous-time Markov chains on finite state spaces.DDSBM adapts Iterative Markovian Fitting (IMF) to discrete-time, introducing reciprocal and Markov projections with a convergence guarantee to the SB solution, and couples this with a graph-transformation setting where the OT cost corresponds to the graph edit distance (GED).The authors demonstrate molecular optimization tasks—transforming graphs with minimal structural edits while achieving target properties—and provide evidence that DDSBM can outperform diffusion-bridge baselines and latent-graph methods on both small-molecule and polymer datasets, including when using pseudo-optimal initial couplings.A key contribution is the combination of discrete SB theory with graph permutation matching and GED-based costs, enabling principled, minimal-edit transformations in graph-structured data with practical applications in chemistry.

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

Figures (22)

• Figure 1: A schematic illustration of DDSBM transforming (a) bridge process to (b) Schrödinger Bridge in discrete spaces.
• Figure 2: A schematic illustration of the lazy gas experiment
• Figure 3: A schematic comparison between lazy gas and busy gas
• Figure 4: A schematic comparison between Schrödinger bridge and aribtrary bridge
• Figure 5: A schematic comparison between Iterative Proportional Fitting (IPF) and Iterative Markovian Fitting (IMF) method. The image are reproduced based on peluchetti2023diffusion. In the notation, $P(\Gamma, \cdot)$ represents the collection of path measures where only the initial distribution is fixed at $\Gamma$, $P(\cdot, \Xi)$ represents the collection of path measures where only the terminal distribution is fixed at $\Xi$, and $P(\Gamma, \Xi)$ represents the collection of path measures where both the initial and terminal distributions are fixed at $\Gamma$ and $\Xi$, respectively.

Theorems & Definitions (20)

• Definition 3.1
• Definition 3.2
• Theorem 3.3
• Theorem B.1
• proof
• Proposition B.2
• Lemma B.3
• proof
• Lemma B.4
• proof