Probabilistic flow matching improves modeling of complex data shapes
Probabilistic Geodesic Flow Matching on Location-Scale Families
Machine Learning
Summary
Many AI systems learn to create data by smoothly changing simple noise into real-world-looking examples. Existing methods often assume these changes happen along straight, simple paths, which don’t work well for complicated data with unusual features. The authors extended this approach by allowing more flexible, curved paths better suited for tricky data shapes, like ones with sharp spikes or heavy tails. They showed that this new way can better capture complex data patterns and sometimes outperforms previous methods.
What this means in practice
- •For machine learning engineers: Improve generative models for complex data by using flexible probability paths that better capture irregular distribution shapes.
- •For data scientists in physical sciences: Model complicated scientific data distributions with improved fidelity by applying the method’s geodesic probability paths on location-scale families.
Authors
Zeyuan Yu, Zhi Chang, Shiwei Lan
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.