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

Oct 2, 2026stat.ML

Generalization Bounds for Flow-matching Generative Models for Intrinsically Low-dimensional Data

Despite the remarkable empirical success of flow-matching models, their statistical generalization guarantees remain underdeveloped. Existing analyses often impose restrictive assumptions on the estimated velocity field and yield convergence rates that fail to reflect the intrinsic low-dimensional structure common in real data, such as natural images and molecular geometries. In this work, we study the statistical generalization of flow-matching models for learning an unknown distribution

P_{\mathrm{data}}

from finitely many samples. We derive finite-sample error bounds on the learned generative distribution, measured in the Wasserstein-

p

distance, for all

p\geq 1

. Specifically, given

n

i.i.d. samples from

P_{\mathrm{data}}

, we show that, for every

d>d_p^\ast(P_{\mathrm{data}})

and appropriately chosen network architectures and hyperparameters, the learned distribution

\widehat{P}^{\mathrm{FM}}

satisfies

\mathbb{W}_p(\widehat{P}^{\mathrm{FM}},P_{\mathrm{data}}) \lesssim n^{-1/d}+n^{-1/(2p)}\bigl(\log(1/ξ)\bigr)^{1/(2p)}

with probability at least

1-ξ

, where

d_p^\ast(P_{\mathrm{data}})

denotes the Wasserstein-

p

dimension of the target measure. Our results demonstrate that flow matching naturally adapts to the intrinsic geometry of data and mitigates the curse of dimensionality, as the convergence exponent depends on the intrinsic rather than ambient dimension. These guarantees remain meaningful in high-dimensional regimes and provide a theoretical explanation for the empirical success of flow matching on structured data distributions under substantially more relaxed assumptions than those in existing analyses.

Saptarshi Chakraborty, Quentin Berthet, Peter L. Bartlett

Department of Statistics, University of Michigan · Google DeepMind · Department of Statistics, University of California, Berkeley +1

Riemannian Flow Matching

Ask PDF