Learning Discrete Diffusion of Graphs via Free-Energy Gradient Flows
For researchers in generative modeling on discrete spaces, this provides a theoretical framework and efficient method for learning diffusion dynamics on graphs, though validation is limited to synthetic data.
This work introduces the first computational approach to discrete diffusion on graphs using gradient flows, leveraging a novel metric to interpret discrete diffusion paths as free-energy gradient flows. The method recovers the underlying functional via first-order optimality conditions, trains extremely fast without requiring sample trajectories, and is validated on synthetic data across various graph classes.
Diffusion-based models on continuous spaces have seen substantial recent progress through the mathematical framework of gradient flows, leveraging the Wasserstein-2 (${W}_2$) metric via the Jordan-Kinderlehrer-Otto (JKO) scheme. Despite the increasing popularity of diffusion models on discrete spaces using continuous-time Markov chains, a parallel theoretical framework based on gradient flows has remained elusive due to intrinsic challenges in translating the ${W}_2$ distance directly into these settings. In this work, we propose the first computational approach addressing these challenges, leveraging an appropriate metric $W_K$ on the simplex of probability distributions, which enables us to interpret widely used discrete diffusion paths, such as the discrete heat equation, as gradient flows of specific free-energy functionals. Through this theoretical insight, we introduce a novel methodology for learning diffusion dynamics over discrete spaces, which recovers the underlying functional directly by leveraging first-order optimality conditions for the JKO scheme. The resulting method optimizes a simple quadratic loss, trains extremely fast, does not require individual sample trajectories, and only needs a numerical preprocessing computing $W_K$-geodesics. We validate our method through extensive numerical experiments on synthetic data, showing that we can recover the underlying functional for a variety of graph classes.