Categorical Flow Matching on Statistical Manifolds

2024-05-26

arXiv.org

AI Summary

Existing discrete data generation methods lack efficient and accurate modeling of categorical distributions. Method: This paper proposes the first flow-matching framework grounded in information geometry. It constructs a Riemannian structure on the categorical statistical manifold via the Fisher–Rao metric and defines optimal transport paths along geodesics, enabling exact likelihood computation without variational lower-bound constraints. Crucially, it integrates information geometry with flow matching for the first time, eliminating reliance on simplistic priors (e.g., uniform or independent distributions). Contribution/Results: We develop a natural-gradient-driven geodesic flow-matching algorithm that supports diffeomorphic mappings and invertible modeling. Empirically, our method substantially outperforms state-of-the-art discrete diffusion and flow models on image, text, and biological sequence generation tasks—simultaneously improving both sample quality and log-likelihood accuracy.

Technology Category

• Machine Learning: Learning with Manifolds
• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Mixed Discrete/Continuous Search

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

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.

Problem

• Probabilistic Graphical Models
• Distribution Matching
• Discrete Generation

Innovation

• Statistical Flow Matching
• Information Geometry
• Optimal Transport