Probabilistic Geodesic Flow Matching on Location-Scale Families

2026-09-28

AI Summary

This study addresses the limitation of existing flow matching methods that rely on Gaussian optimal transport, which struggle to handle complex distributions such as heavy-tailed ones. To overcome this, we propose Probabilistic Geodesic Flow Matching, extending the framework to location-scale families by defining geodesic paths on the manifold of probability distributions. By integrating neural ordinary differential equations (neural ODEs), our method enables optimal path planning within the probability space, effectively transcending the constraints of Euclidean geometry. This approach significantly enhances the capacity of generative models to capture non-uniform structures. Extensive experiments demonstrate that our method accurately models complex distributions across both synthetic and scientific datasets, achieving performance that surpasses or matches state-of-the-art models.

Abstract

Flow matching (FM) has recently emerged as a promising framework for generative modeling due to its conceptual simplicity and strong empirical performance. In FM, samples are transported along a vector field parameterized by a neural network, inducing a probability path that evolves from a simple noise distribution to the target data distribution, governed by an ordinary differential equation (ODE). However, existing FM approaches predominantly rely on probability paths derived from optimal transport (OT) between Gaussian distributions, which may be suboptimal for capturing complex data with inhomogeneous structures such as heavy tail or sharp contrast. In this work, we generalize FM to the broader class of location-scale families for handling data inhomogeneity and introduce a novel class of probability paths defined as geodesics on the manifold of probability distributions. We name this approach probabilistic geodesic flow matching to distinguish it from prior geodesic (Riemannian) FM methods defined in input space. We argue that Euclidean OT-based paths are not necessarily optimal in probability space and may limit modeling flexibility. Through synthetic benchmarks and scientific datasets at different scales, we demonstrate that the proposed method more effectively captures complex distributions, leading to improved or comparable performance compared with SOTA geometry-motivated generative models.

Problem

• Flow Matching
• Generative Modeling
• Location-Scale Families
• Optimal Transport
• Probability Paths

Innovation

• Flow Matching
• Location-Scale Families
• Probabilistic Geodesic
• Probability Manifold
• Generative Modeling