We introduce Statistical Flow Matching (SFM), a novel and mathematically rigorous flow-matching framework on the manifold of parameterized probability measures inspired by the results from information geometry. We demonstrate the effectiveness of our method on the discrete generation problem by instantiating SFM on the manifold of categorical distributions whose geometric properties remain unexplored in previous discrete generative models. Utilizing the Fisher information metric, we equip the manifold with a Riemannian structure whose intrinsic geometries are effectively leveraged by following the shortest paths of geodesics. We develop an efficient training and sampling algorithm that overcomes numerical stability issues with a diffeomorphism between manifolds. Our distinctive geometric perspective of statistical manifolds allows us to apply optimal transport during training and interpret SFM as following the steepest direction of the natural gradient. Unlike previous models that rely on variational bounds for likelihood estimation, SFM enjoys the exact likelihood calculation for arbitrary probability measures. We manifest that SFM can learn more complex patterns on the statistical manifold where existing models often fail due to strong prior assumptions. Comprehensive experiments on real-world generative tasks ranging from image, text to biological domains further demonstrate that SFM achieves higher sampling quality and likelihood than other discrete diffusion or flow-based models.

翻译：我们提出了统计流匹配（SFM），这是一种新颖且数学严谨的流匹配框架，它建立在参数化概率测度的流形之上，其灵感来源于信息几何学的研究成果。我们通过在分类分布的流形上实例化SFM来证明该方法在离散生成问题上的有效性，该流形的几何性质在以往的离散生成模型中尚未得到充分探索。利用费希尔信息度量，我们为该流形配备了黎曼结构，并通过遵循测地线的最短路径来有效利用其内在几何特性。我们开发了一种高效的训练与采样算法，该算法通过流形间的微分同胚克服了数值稳定性问题。我们独特的统计流形几何视角使得我们能够在训练过程中应用最优传输，并将SFM解释为沿着自然梯度的最速下降方向。与以往依赖变分下界进行似然估计的模型不同，SFM能够对任意概率测度进行精确似然计算。我们证明，SFM能够在统计流形上学习更复杂的模式，而现有模型由于强先验假设往往在此失效。在涵盖图像、文本及生物领域的真实世界生成任务上的全面实验进一步表明，SFM在采样质量和似然度方面均优于其他基于离散扩散或流的模型。