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.

翻译：连续空间上的扩散模型通过梯度流的数学框架取得了显著进展，特别是借助Jordan-Kinderlehrer-Otto（JKO）方案利用Wasserstein-2（${W}_2$）度量。尽管基于连续时间马尔可夫链的离散空间扩散模型日益流行，但由于将${W}_2$距离直接迁移至此类场景存在固有挑战，基于梯度流的平行理论框架至今仍未建立。在本工作中，我们首次提出解决这些挑战的计算方法，通过利用概率分布单纯形上的适当度量$W_K$，使我们能将广泛使用的离散扩散路径（如离散热方程）解释为特定自由能泛函的梯度流。基于这一理论洞见，我们提出了一种学习离散空间上扩散动力学的新型方法论，该方法直接利用JKO方案的一阶最优条件恢复底层泛函。所得方法优化简单的二次损失函数，训练速度极快，无需个体样本轨迹，仅需计算$W_K$-测地线的数值预处理。通过合成数据的广泛数值实验，我们验证了该方法能够为多种图类恢复底层泛函。