Probabilistic Geodesic Flow Matching on Location-Scale Families

Zeyuan Yu; Zhi Chang; Shiwei Lan

stat.MLSep 28, 2026

Probabilistic Geodesic Flow Matching on Location-Scale Families

Authors: Zeyuan Yu, Zhi Chang, Shiwei Lan

Organizations: School of Mathematical & Statistical Sciences Arizona State University Tempe, AZ 85287 USA

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.

Figures & tables

Figure 1 : Probability densities of exponential power distributions ${\mathcal{P}}_{q}({\bf 0},{\bf I})$ for selected $q$ ’s.

Figure 2 : Contrasting trajectories of energy distance (ED) in the Swiss roll, moons, checkerboard, and funnel examples (left four) and class Frechet inception distance (clf-FID) in MNIST (rightmost).

Figure 3 : Sample paths in the 8-Gaussians (36d) example. Only the first two dimensions are plotted.

DIST	Method	$W_{2}$ $\downarrow$	ED $\downarrow$	MMD $\downarrow$	time/epoch
q=1.0	OT	$3.73e{0}\pm 2.49e{-1}$	$1.43e{-1}\pm 1.35e{-2}$	$3.52e{-3}\pm 1.07e{-3}$	$5.90e{-2}\pm 6.85e{-5}$
q=2.0	OT	$3.74e{0}\pm 2.00e{-1}$	$1.41e{-1}\pm 1.18e{-2}$	$3.32e{-3}\pm 7.60e{-4}$	$5.76e{-2}\pm 8.21e{-5}$
q=1.0	RFM	$3.65e{0}\pm 2.01e{-1}$	$1.27e{-1}\pm 1.33e{-2}$	$2.30e{-3}\pm 7.17e{-4}$	$7.68e{-2}\pm 6.59e{-4}$
q=2.0	RFM	$3.69e{0}\pm 2.12e{-1}$	$1.29e{-1}\pm 1.10e{-2}$	$2.48e{-3}\pm 5.29e{-4}$	$7.11e{-2}\pm 5.74e{-4}$
q=1.0	MFM	$4.05e{0}\pm 2.39e{-1}$	$1.26e{-1}\pm 1.16e{-2}$	$2.65e{-3}\pm 7.57e{-4}$	$4.36e{-1}\pm 1.57e{-2}$
q=2.0	MFM	$4.17e{0}\pm 1.83e{-1}$	$1.27e{-1}\pm 1.31e{-2}$	$2.80e{-3}\pm 5.95e{-4}$	$4.26e{-1}\pm 1.62e{-2}$

Table 1 : Neal’s funnel (100d): comparison in

W_{2}

, ED, and MMD (mean

\pm

std over 10 seeds).

Method	Manifold of	Metric	Geodesic	Data type
Riemannian FM [ 6 ]	general (mainly data)	standard Riemannian	general	non-Euclidean
Metric FM [ 17 ]	Data	learned	general	general
Meta FM [ 4 ]	probability densities	$2$ -Wasserstein	general	gmbedded in graph
Categorical FM [ 8 ]	probability measures	Fisher-Rao	circle	discrete
Fisher FM [ 9 ]	probability densities	Fisher-Rao	circle	discrete
PG FM (ours)	probability densities	Fisher-Rao	Poincaré semicircle	continuous, Euclidean

4.Table : Comparison of flow-matching methods.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Figure A.1 : Euclidean vs Fisher-Rao distances, (probabilistic) geodesic probability paths, flow paths for $q=1.0$ and $q=2.0$ (from left to right) on the manifold of exponential power distributions.

Figure B.1 : Sample paths in the Swiss roll (left) and moons (right) examples: VP, OT, sino, PG-2, and PG-1 (from top to bottom).

DIST	PATH	$W_{2}$ $\downarrow$	ED $\downarrow$	MMD $\downarrow$	time/epoch
q=1.0	VP	$1.49e{-1}\pm 3.5e{-2}$	$3.69e{-2}\pm 1.5e{-2}$	$3.33e{-4}\pm 2.2e{-4}$	$5.84e{-2}\pm 6.1e{-4}$
q=2.0	VP	$1.42e{-1}\pm 4.5e{-2}$	$2.93e{-2}\pm 2.3e{-2}$	$3.06e{-4}\pm 3.6e{-4}$	$6.24e{-2}\pm 5.3e{-4}$
q=1.0	OT	$1.14e{-1}\pm 2.3e{-2}$	$1.56e{-2}\pm 1.1e{-2}$	$7.50e{-5}\pm 9.1e{-5}$	$5.66e{-2}\pm 5.4e{-4}$
q=2.0	OT	$1.12e{-1}\pm 3.0e{-2}$	$1.69e{-2}\pm 1.6e{-2}$	$9.03e{-5}\pm 1.7e{-4}$	$5.40e{-2}\pm 6.2e{-4}$
q=1.0	sino	$1.33e{-1}\pm 2.5e{-2}$	$2.53e{-2}\pm 1.3e{-2}$	$1.87e{-4}\pm 1.2e{-4}$	$5.74e{-2}\pm 3.7e{-4}$
q=2.0	sino	$1.17e{-1}\pm 1.4e{-2}$	$1.24e{-2}\pm 8.3e{-3}$	$6.05e{-5}\pm 5.0e{-5}$	$5.46e{-2}\pm 4.2e{-4}$

Appendix

Table B.1 : Swiss roll dataset comparison in

W_{2}

, ED, and MMD (mean

\pm

std over 10 seeds).

		Moons			Checkerboard
DIST	PATH	$W_{2}$ $\downarrow$	ED $\downarrow$	MMD $\downarrow$	$W_{2}$ $\downarrow$	ED $\downarrow$	MMD $\downarrow$
q=1.0	VP	$1.33e{-1}\pm 2.4e{-2}$	$2.56e{-2}\pm 1.5e{-2}$	$2.14e{-4}\pm 1.8e{-4}$	$3.58e{-1}\pm 9.6e{-2}$	$4.45e{-2}\pm 2.9e{-2}$	$3.01e{-4}\pm 3.0e{-4}$
q=2.0	VP	$1.06e{-1}\pm 1.7e{-2}$	$1.00e{-2}\pm 1.0e{-2}$	$5.69e{-5}\pm 6.5e{-5}$	$3.53e{-1}\pm 6.2e{-2}$	$4.29e{-2}\pm 1.6e{-2}$	$2.05e{-4}\pm 1.2e{-4}$
q=1.0	OT	$1.17e{-1}\pm 2.1e{-2}$	$1.94e{-2}\pm 1.2e{-2}$	$1.08e{-4}\pm 1.2e{-4}$	$3.15e{-1}\pm 5.3e{-2}$	$3.23e{-2}\pm 1.7e{-2}$	$1.63e{-4}\pm 1.1e{-4}$
q=2.0	OT	$9.83e{-2}\pm 1.3e{-2}$	$9.39e{-3}\pm 8.6e{-3}$	$5.56e{-5}\pm 6.0e{-5}$	$2.96e{-1}\pm 4.1e{-2}$	$2.21e{-2}\pm 1.6e{-2}$	$9.36e{-5}\pm 7.4e{-5}$
q=1.0	sino	$1.25e{-1}\pm 1.8e{-2}$	$2.26e{-2}\pm 1.1e{-2}$	$1.59e{-4}\pm 1.1e{-4}$	$3.59e{-1}\pm 9.1e{-2}$	$4.76e{-2}\pm 2.4e{-2}$	$3.38e{-4}\pm 2.5e{-4}$

Appendix

Table B.2 : Comparison on Moons and Checkerboard datasets in

W_{2}

, ED, and MMD (mean

\pm

std over 10 seeds).

	Swiss Roll		Moons		Checkerboard
DIST	$W_{2}$ $\downarrow$	ED $\downarrow$	$W_{2}$ $\downarrow$	ED $\downarrow$	$W_{2}$ $\downarrow$	ED $\downarrow$
q=0.5	$3.70e{-1}\pm 2.1e{-2}$	$6.76e{-2}\pm 5.8e{-3}$	$3.18e{-1}\pm 1.7e{-2}$	$6.14e{-2}\pm 4.6e{-3}$	$3.92e{-1}\pm 4.2e{-2}$	$3.56e{-2}\pm 1.3e{-2}$
q=0.8	$1.11e{-1}\pm 1.9e{-2}$	$1.02e{-2}\pm 8.8e{-3}$	$1.09e{-1}\pm 1.6e{-2}$	$1.41e{-2}\pm 9.6e{-3}$	$3.19e{-1}\pm 4.1e{-2}$	$2.80e{-2}\pm 1.3e{-2}$
q=1.0	$1.05e{-1}\pm 2.4e{-2}$	$9.56e{-3}\pm 7.8e{-3}$	$1.02e{-1}\pm 1.3e{-2}$	$9.26e{-3}\pm 9.2e{-3}$	$3.17e{-1}\pm 3.8e{-2}$	$2.74e{-2}\pm 1.6e{-2}$
q=1.5	$1.06e{-1}\pm 1.4e{-2}$	$9.13e{-3}\pm 8.6e{-3}$	$9.95e{-2}\pm 1.2e{-2}$	$9.93e{-3}\pm 9.3e{-3}$	$3.19e{-1}\pm 5.4e{-2}$	$3.07e{-2}\pm 1.8e{-2}$
q=2.0	$1.08e{-1}\pm 2.0e{-2}$	$8.66e{-3}\pm 8.0e{-3}$	$1.00e{-1}\pm 1.4e{-2}$	$5.62e{-3}\pm 8.0e{-3}$	$3.05e{-1}\pm 3.2e{-2}$	$2.56e{-2}\pm 6.8e{-3}$

Appendix

Table B.3 : PG

q

-sweep comparison across Swiss roll, moons, and checkerboard in

W_{2}

and ED (mean

\pm

std over 10 seeds).

Figure B.2 : Sample paths in the checkerboard (2d, left) and Neal’s funnel (10d, right) examples: VP, OT, sino, PG-2, and PG-1 (from top to bottom).

DIST	Method	$W_{2}$ $\downarrow$	ED $\downarrow$	MMD $\downarrow$	time/epoch
q=1.0	OT	$5.660e{0}\pm 1.43e{-2}$	$5.93e{-2}\pm 1.25e{-2}$	$1.39e{-4}\pm 5.09e{-5}$	$5.92e{-2}\pm 1.27e{-4}$
q=2.0	OT	$5.493e{0}\pm 1.41e{-2}$	$6.62e{-2}\pm 7.93e{-3}$	$2.47e{-4}\pm 4.99e{-5}$	$5.70e{-2}\pm 1.41e{-4}$
q=1.0	Sino	$5.499e{0}\pm 1.03e{-2}$	$4.06e{-2}\pm 6.28e{-3}$	$9.67e{-5}\pm 2.37e{-5}$	$5.82e{-2}\pm 1.11e{-4}$
q=2.0	Sino	$5.329e{0}\pm 2.86e{-3}$	$1.29e{-1}\pm 2.66e{-3}$	$1.17e{-3}\pm 4.85e{-5}$	$5.69e{-2}\pm 8.26e{-5}$
q=1.0	RFM	$5.49e{0}\pm 8.79e{-3}$	$7.28e{-2}\pm 6.73e{-3}$	$2.93e{-4}\pm 4.30e{-5}$	$7.47e{-2}\pm 3.16e{-3}$
q=2.0	RFM	$5.50e{0}\pm 1.61e{-2}$	$6.41e{-2}\pm 7.89e{-3}$	$2.24e{-4}\pm 6.10e{-5}$	$6.89e{-2}\pm 3.31e{-3}$

Appendix

Table B.4 : 8-Gaussians (36d): comparison in

W_{2}

, ED, and MMD (mean

\pm

std over 10 seeds).

Figure B.3 : Energy distance (ED) of OT, sino, PG-2, and PG-1 in the 8-Gaussians (36d, left) and funnel (100d, right) examples.

DIST	PATH	$W_{2}$ $\downarrow$	ED $\downarrow$	MMD $\downarrow$	time/epoch
q=1.0	VP	$2.55e{0}\pm 1.0e{-1}$	$9.62e{-2}\pm 1.5e{-2}$	$3.04e{-3}\pm 1.2e{-3}$	$5.50e{-2}\pm 4.3e{-4}$
q=2.0	VP	$2.45e{0}\pm 4.6e{-2}$	$9.48e{-2}\pm 3.3e{-2}$	$4.10e{-3}\pm 3.0e{-3}$	$5.27e{-2}\pm 1.2e{-4}$
q=1.0	OT	$2.47e{0}\pm 9.2e{-2}$	$3.25e{-2}\pm 1.6e{-2}$	$2.36e{-4}\pm 1.5e{-4}$	$5.59e{-2}\pm 5.6e{-4}$
q=2.0	OT	$2.54e{0}\pm 1.5e{-1}$	$2.69e{-2}\pm 1.3e{-2}$	$1.63e{-4}\pm 8.5e{-5}$	$5.37e{-2}\pm 4.0e{-4}$
q=1.0	sino	$2.51e{0}\pm 8.7e{-2}$	$3.64e{-2}\pm 1.1e{-2}$	$2.76e{-4}\pm 1.7e{-4}$	$5.64e{-2}\pm 2.0e{-4}$
q=2.0	sino	$2.55e{0}\pm 2.5e{-1}$	$3.07e{-2}\pm 8.6e{-3}$	$1.50e{-4}\pm 6.3e{-5}$	$5.30e{-2}\pm 1.3e{-3}$

Appendix

Table B.5 : Funnel example (10d) comparison in

W_{2}

, ED, and MMD (mean

\pm

std over 10 seeds).

Figure B.4 : Sample paths in the alanine dipeptide (30d, left) and QM9 (87d, right): VP, OT, sino, PG-2, and PG-1 (from top to bottom). The last column of each panel represents the real data. High dimensions have been projected to 2d using t-SNE for visualization.

		Alanine dipeptide (30d)			QM9 (87d)
DIST	PATH	$W_{2}$ $\downarrow$	ED $\downarrow$	MMD $\downarrow$	$W_{2}$ $\downarrow$	ED $\downarrow$	MMD $\downarrow$
$1.0$	OT	$1.49e{0}\pm 1.4e{-1}$	$7.75e{-2}\pm 3.1e{-2}$	$4.02e{-4}\pm 2.6e{-4}$	$7.27e{0}\pm 8.8e{-1}$	$1.53e{-1}\pm 1.2e{-2}$	$8.65e{-4}\pm 1.3e{-4}$
$2.0$	OT	$1.41e{0}\pm 1.1e{-1}$	$5.62e{-2}\pm 2.3e{-2}$	$1.71e{-4}\pm 1.2e{-4}$	$6.84e{0}\pm 3.2e{-1}$	$1.31e{-1}\pm 1.6e{-2}$	$6.88e{-4}\pm 1.8e{-4}$
$1.0$	sino	$1.56e{0}\pm 1.0e{-1}$	$8.86e{-2}\pm 3.6e{-2}$	$5.00e{-4}\pm 3.0e{-4}$	$7.05e{0}\pm 2.5e{-1}$	$1.27e{-1}\pm 1.6e{-2}$	$6.29e{-4}\pm 1.7e{-4}$
$2.0$	sino	$1.45e{0}\pm 1.3e{-1}$	$6.59e{-2}\pm 3.3e{-2}$	$2.88e{-4}\pm 2.1e{-4}$	$7.04e{0}\pm 3.4e{-1}$	$1.27e{-1}\pm 3.0e{-2}$	$6.33e{-4}\pm 3.4e{-4}$
$1.0$	PG	$1.62e{0}\pm 1.8e{-1}$	$1.05e{-1}\pm 3.7e{-2}$	$7.40e{-4}\pm 4.9e{-4}$	$6.94e{0}\pm 3.3e{-1}$	$1.63e{-1}\pm 1.9e{-2}$	$9.05e{-4}\pm 2.0e{-4}$

Appendix

Table B.6 : Comparison across Alanine dipeptide (30d) and QM9 (87d) in

W_{2}

, ED, and MMD (mean

\pm

std over 10 seeds).

		scRNA (50d)			Lattice $\phi^{4}$ (64d)
DIST	PATH	$W_{2}$ $\downarrow$	ED $\downarrow$	MMD $\downarrow$	$W_{2}$ $\downarrow$	ED $\downarrow$	MMD $\downarrow$
$1.0$	OT	$5.32e{0}\pm 4.0e{-2}$	$5.17e{-1}\pm 2.0e{-2}$	$1.88e{-2}\pm 1.5e{-3}$	$7.54e{0}\pm 9.0e{-3}$	$2.08e{-1}\pm 4.0e{-3}$	$2.18e{-3}\pm 8.0e{-5}$
$2.0$	OT	$5.33e{0}\pm 5.0e{-2}$	$5.10e{-1}\pm 2.0e{-2}$	$1.82e{-2}\pm 1.3e{-3}$	$7.53e{0}\pm 1.1e{-2}$	$2.06e{-1}\pm 6.0e{-3}$	$2.14e{-3}\pm 1.2e{-4}$
$1.0$	sino	$5.41e{0}\pm 5.0e{-2}$	$5.36e{-1}\pm 2.0e{-2}$	$2.04e{-2}\pm 1.5e{-3}$	$7.50e{0}\pm 1.0e{-2}$	$2.25e{-1}\pm 4.0e{-3}$	$2.64e{-3}\pm 1.1e{-4}$
$2.0$	sino	$5.44e{0}\pm 5.0e{-2}$	$5.46e{-1}\pm 2.0e{-2}$	$2.09e{-2}\pm 1.4e{-3}$	$7.47e{0}\pm 1.0e{-2}$	$2.25e{-1}\pm 4.0e{-3}$	$2.74e{-3}\pm 9.0e{-5}$
$1.0$	PG	$5.71e{0}\pm 4.0e{-2}$	$2.93e{-1}\pm 2.2e{-2}$	$5.07e{-3}\pm 6.0e{-4}$	$7.73e{0}\pm 1.4e{-2}$	$1.55e{-1}\pm 7.0e{-3}$	$1.06e{-3}\pm 8.0e{-5}$

Appendix

Table B.7 : Comparison across scRNA (50d) and Lattice

\phi^{4}

(64d) in

W_{2}

, ED, and MMD (mean

\pm

std over 10 seeds).

Explore similar work

May 8, 2026cs.LG

Generalized Wasserstein Flow Matching: Transport Plans, Everywhere, All at Once

Flow matching has recently emerged as a flexible and efficient framework for generative modelling by learning deterministic transport dynamics between probability measures. In this work, we extend flow matching to the space of probability measures over probability measures, introducing a Wasserstein-on-Wasserstein (WoW) formulation. Leveraging the nested Wasserstein geometry, we show that measures over transport plans naturally induce velocity fields that realize metameasure flows. This yields a principled generalization of Wasserstein flow matching via coupled outer and inner transport plans. To address the substantial computational cost of WoW transport, we propose scalable approximations based on sliced and linear Wasserstein distances, enabling efficient training while promoting numerically stable, near-straight trajectories. Our framework unifies and extends existing approaches to point cloud and set generation, providing a practical and theoretically grounded method for generative modelling in WoW spaces.

Moritz Piening, Richard Duong, Gabriele Steidl

Institut für Mathematik, Technische Universität Berlin, Germany

Velocity Field Generative Models Point Clouds Transport

Ask PDF

May 7, 2026cs.LG

Flow Matching with Arbitrary Auxiliary Paths

We introduce a new generative modeling framework, \textbf{Flow Matching with Arbitrary Auxiliary Paths (AuxPath-FM)}, which generalizes conditional flow matching by incorporating an auxiliary variable drawn from an arbitrary distribution into the probability path. Unlike prior methods that restrict auxiliary components to Gaussian noise, AuxPath-FM allows the variable

η

to follow any distribution, producing trajectories of the form

X_t = a(t)X_1 + b(t)X_0 + c(t)η

. We theoretically demonstrate that this construction preserves the continuity equation and maintains a training objective consistent with the marginal formulation. This flexibility enables the design of diverse probability paths using various priors, including Gaussian, Uniform, Laplace, and discrete Rademacher distributions, each offering unique geometric properties for generative flows. Furthermore, our framework allows for specialized tasks such as label-guided generation by encoding structured semantic information into the auxiliary distribution. Overall, AuxPath-FM provides a principled and general foundation for probability path design, offering both theoretical generality and practical flexibility for diverse generative modeling tasks.

Xin Peng, Ang Gao

School of Physical Science and Technology, Beijing University of Posts and Telecommunications, Beijing, China

Conditional Flow Matching Pathways

Ask PDF

Jun 23, 2026cs.AI

The Geometry Behind Diffusion and Flow Matching: Gradient Flows and Geodesics in Wasserstein Space

The space

\mathcal{P}_2(\mathbb{R}^d

) of probability measures with finite second moment carries a natural geometry: the quadratic Wasserstein distance W_2 makes it a complete metric space and, following Otto, a (formal) Riemannian manifold whose geodesics are the optimal-transport interpolations. On this manifold, the gradient flow of the free energy F(rho) = KL(rho || π) is exactly the Fokker-Planck equation, and its implicit-Euler discretization is the JKO scheme. This is the geometry underlying diffusion models: the forward process descends the free energy, and each denoising step realizes one JKO step, which recovers DDPM, DDIM, NCSN/SMLD, and Energy Matching; this is one scheme, not separate theories. The same manifold supports a second variational principle. Its geodesics - the minimum-action curves of the Benamou-Brenier formula - are precisely the optimal-transport paths that Flow Matching learns. Fixing both endpoints and following the geodesic, generation becomes a deterministic ODE along a straight line, hence far fewer sampling steps. Placing both families of models on one manifold makes their relationship exact: diffusion follows a free-energy gradient flow, an initial-value problem; optimal-transport Flow Matching follows a Wasserstein geodesic, a boundary-value problem. The two reach the same endpoints along different paths.

Yian Yao, Weiwei Zhang

from the mortal world · from heaven

Riemannian Flow Matching Wasserstein Gradient Flows Wasserstein Distance Geodesic Fokker--Planck Equation Variational+1

Ask PDF