Learning Discrete Diffusion on Graphs via Free-Energy Gradient Flows

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, as well as cell type prediction in single-cell trajectory data. Code will be made available at github.com/dariorancati/discrete-gradient-flow.

Many real-world phenomena evolve in order to minimize some kind of energy function: this assumptions has in the past been an effective modelling tool in diverse fields such as neurosciences, developmental biology, and machine learning. The key idea is that, if we find a way to recover this energy function from observed trajectories, we construct a model to perform forecasting that is both cheap to train and run as well as interpretable. However, whenever the quantities we observe can only take only a finite number of values (for example from a finite set of cell types which biologies have identified), the existing mathematical machinery to learn such energy function breaks down. In this work, we present the first implementation of this kind to circumvent this problem: we leverage decade-old mathematical findings in the field of discrete optimal transport to develop our method, providing both theoretical foundation and a practical implementation for this problem. We demonstrate the effectiveness of the implemented model on both synthetic as well as single-cell biological data.