Bures-Wasserstein Flow Matching for Graph Generation

TL;DR

This paper introduces BWFlow, a novel graph generation method that models joint node and edge evolution using optimal transport, resulting in smoother probability paths and improved training and sampling performance.

Contribution

It presents a theoretically grounded framework for constructing smooth probability paths in graph generation by modeling graphs as Markov random fields and leveraging optimal transport.

Findings

BWFlow achieves competitive graph generation performance.

It demonstrates better training convergence compared to existing methods.

It enables more efficient and reliable sampling of graphs.

Abstract

Graph generation has emerged as a critical task in fields ranging from drug discovery to circuit design. Contemporary approaches, notably diffusion and flow-based models, have achieved solid graph generative performance through constructing a probability path that interpolates between reference and data distributions. However, these methods typically model the evolution of individual nodes and edges independently and use linear interpolations in the disjoint space of nodes/edges to build the path. This disentangled interpolation breaks the interconnected patterns of graphs, making the constructed probability path irregular and non-smooth, which causes poor training dynamics and faulty sampling convergence. To address the limitation, this paper first presents a theoretically grounded framework for probability path construction in graph generative models. Specifically, we model the joint evolution of the nodes and edges by representing graphs as connected systems parameterized by Markov random fields (MRF). We then leverage the optimal transport displacement between MRF objects to design a smooth probability path that ensures the co-evolution of graph components. Based on this, we introduce BWFlow, a flow-matching framework for graph generation that utilizes the derived optimal probability path to benefit the training and sampling algorithm design. Experimental evaluations in plain graph generation and molecule generation validate the effectiveness of BWFlow with competitive performance, better training convergence, and efficient sampling.