Probabilistic Geodesic Flow Matching on Location-Scale Families

1 Introduction

• Flow matching predicts a vector field for a noise-to-data probability path, avoiding the expensive ODE simulation used in continuous normalizing-flow training.The paper situates its method within the broader diffusion and flow-matching modeling context.Source: §1 · p. 1
• Gaussian optimal-transport paths may be suboptimal for complex distributions, including non-Gaussian heavy-tailed data.The exponential power family includes the Gaussian at q = 2 and varies regularization and tail behavior through q.Source: §1 · p. 1
• Probabilistic geodesic flow matching extends flow matching to location-scale families and defines geodesics on probability distributions, unlike Riemannian FM’s input-space geodesics.The approach focuses on exponential power distributions and uses the Fisher–Rao metric.Source: §1 · pp. 1–2
• The paper reports improved or comparable performance for probabilistic geodesic flow matching.Its stated contributions include generalizing flow matching beyond Gaussian distributions and interpreting optimal probability paths as energy-minimizing trajectories.Source: §1 · p. 2

2.1 (Conditional) Flow Matching

• Conditional flow matching trains a neural network vector field to match a conditional vector field, then generates samples by solving its ODE to time 1.The conditional paths begin at a noise distribution and end near a data point; the common Gaussian conditional uses a small positive σ_min.Source: §2.1 · pp. 2–3
• The usual Gaussian optimal-transport conditional path linearly interpolates its location and scale between noise and the data endpoint.Its parameters use µ0 = 0, µ1 = x1, σ0 = 1, and σ1 = σ_min > 0.Source: §2.1 · p. 3

2.2 Exponential Power Distribution

• A location-scale family is parameterized by location µ and positive scale σ, with X = µ + σZ for a standard member Z.Its CDF is F0((X − µ)/σ).Source: §2.2 · p. 3
• The exponential power distribution broadens flow matching beyond Gaussian distributions and reduces to a multivariate normal when q = 2.Its parameter q controls central regularization and tail behavior; smaller q gives sharper regularization and heavier tails.Source: §2.2 · p. 3
• The paper focuses on isotropic and diagonal-scale anisotropic exponential power distributions, whose densities are illustrated for selected q values in Figure 1.The distribution is location-scale: scaling and shifting Pq(0, I) yields Pq(µ, σ²I).Source: §2.2 · p. 3, Figure 1 · p. 3

3 Flexible Flow Path in Location-Scale Family

• The paper also studies minimal-energy, nonlinear-interpolant, and probabilistic-geodesic paths by minimizing flow energies under different metric tensors.The metrics include Euclidean and Fisher–Rao.Source: §3 · pp. 4–5
• Any location-scale probability path can be paired with a vector field satisfying the continuity equation, extending the construction beyond Gaussian distributions.The result is distribution-agnostic within the location-scale family.Source: §3 · p. 4
• The framework supports choosing a probability path in a location-scale family before deriving its vector field and sample flow, allowing greater flexibility than the flow-first approach in [20].The path can come from distributions such as Normal, Laplace, Student’s t, or elliptical families.Source: §3 · p. 4

3.1 Minimal Energy Path

• The framework minimizes flow energy using an energy criterion, and the resulting minimal-energy path is the linear OT interpolant in location and scale.For the conditional path, µt = tx1 and σt = 1 − t + tσ_min.Source: §3.1 · p. 5

3.2 Nonlinear Interpolant Path

• A sinusoidal nonlinear interpolant uses µt = x1 sin(πt/2) and σt = cos(πt/2), but does not minimize the stated flow energy.The paper later identifies it as a special case of a Fisher–Rao probabilistic geodesic path.Source: §3.2 · p. 5

4 Probabilistic Geodesic Flow

• The probabilistic geodesic is a semi-circle on the exponential-power manifold that minimizes variational energy.Theorem 4.1 gives this minimizing path.Source: §4 · p. 6
• The isotropic geodesic formula has O(d) numerical complexity, and an anisotropic version decouples it into d independent one-dimensional geodesics.Source: §4 · p. 6
• Along the geodesic, the conditional flow-matching loss decomposes into the flow-matching loss and conditional-vector-field variance, and is bounded from below.Source: §4 · p. 6
• The Fisher–Rao distance between endpoint exponential-power distributions can be smaller than their Euclidean distance for some q values.The section also contrasts q = 1.0 and q = 2.0, noting that smaller q leads to more stringent movement.Source: §4 · p. 6
• As σ1 approaches zero, geodesic length grows without bound and the conditional vector field can make the learned ODE stiff and numerically difficult to solve.The authors use a hybrid path, with tc = 0.85 in numerical experiments, to alleviate this limitation.Source: §4 · p. 7

5 Numerical Results

• Across synthetic, image, and scientific tests, PG and HB are compared with baseline and state-of-the-art flow-matching methods using distribution-discrepancy and generation-quality measures.The evaluation includes 2D benchmarks, higher-dimensional simulations, handwritten digits, scientific datasets, and molecular generation.Source: §5 · p. 7

5.1.1 2d Benchmarks: swiss roll, moons, and checkerboard

• On the 2D benchmarks, PG methods achieve smaller W2, ED, and MMD discrepancies in some cases by an order of magnitude, while their per-epoch time overhead fades as dimension increases.PG paths also reduce energy distance to the target quickly and monotonically, unlike the midpoint energy hump reported for independent-coupling OT paths.Source: §5.1.1 · p. 7

5.1.2 Higher Dimensions: 36d eight-Gaussians, 100d Funnel Distribution

• In the 36-dimensional eight-Gaussians example, PG-1 and HB identify all modes fastest; at q=1.0, HB attains the lowest W2 and ED, while MFM matches its lowest ED at higher per-epoch cost.The reported costs are 4.36e-1 seconds per epoch for MFM and 6.92e-2 for HB; q=1.0 nearly always beats q=2.0 in Table 1, with a few exceptions.Source: §5.1.2 · p. 8, Table 1 · p. 8

5.2 MNIST Handwritten Digits

• On MNIST, OT has the lowest FID, while HB ties for the highest IS and has the lowest clf-FID.The comparison covers OT, PG, and HB at q=1 and q=2; results are reported as mean ± standard deviation over five runs.Source: §5.2 · p. 9, Table 2 · p. 9

5.3 Molecular Generation

• For QM9 molecule generation, PG/HB has slightly lower validity but higher stable-atom and stable-molecule percentages than OT, FFM, and FlowMol.The reported percentages are based on 1,000 generated molecules.Source: §5.3 · p. 9, Table 3 · p. 9

6 Conclusion

• The method extends conditional flow matching beyond Gaussians, focusing on exponential power distributions, and reports numerical performance superior or comparable to state-of-the-art flow-matching models.The exponential power family includes the Gaussian as the special case q = 2.Source: §6 · p. 9
• The probabilistic geodesic path is less conceptually straightforward than the OT interpolant, while its computational overhead diminishes as dimension increases.The authors identify integration into recent one-step methods as future work.Source: §6 · p. 9

AI use statement

• The authors used generative AI to polish presentation and assist with coding and debugging, and report that they reviewed AI-assisted work.They state that co-authors verified and tested LLM-generated code and take responsibility for the final content.Source: AI use statement · p. 10

Reproducibility statement

• The authors state that results are reproducible using Python code to be published after publication and that a demonstrative Jupyter notebook is included in the submitted supplementary materials.Source: Reproducibility statement · p. 10

• The comparison lists PG FM for probability densities with Fisher–Rao metric, Poincaré semicircle geometry, and continuous Euclidean data.It contrasts this with Riemannian FM for general data and general Riemannian geometry, and Meta FM for probability densities embedded in a graph.Source: Method Manifold of Metric Geodesic… · p. 12

A Proofs

• The OT interpolant is separately characterized as minimizing flow energy, while the appendix suggests a Fisher–Rao-norm loss would yield a tighter bound.Source: A Proofs · pp. 13, 17
• The location-scale path uses affine interpolation of location and scale parameters, and the resulting flow has an affine-transformation solution.The appendix also states that the corresponding geodesic path is affine in time.Source: A Proofs · pp. 12–13
• Under the exponential-power Fisher metric, geodesics are hyperbolic and can be expressed through a Poincaré semicircle parametrization.The appendix gives equivalent semicircle and hyperbolic parametrizations; the Gaussian case recovers the well-known Gaussian Fisher–Rao distance.Source: A Proofs · pp. 14, 16
• The CFM loss decomposes into the FM loss and conditional-vector-field variance, with that variance providing a lower bound independent of the neural network.Equality holds when the learned field matches the marginal vector field almost everywhere in space and time.Source: A Proofs · pp. 16–17

B.1 Setup for Numerical Experiments

• Experiments use a five-layer dense network with 512 hidden dimensions for simulations and a classic UNet for MNIST, with midpoint ODE sampling.The simulation network uses sinusoidal time encoding and Adam at 1e-4; MNIST uses AdamW at 2e-4 and 100 sampling steps.Source: §B.1 · p. 17
• Swiss-roll and moons sample paths compare VP, OT, sino, PG-2, and PG-1 in that top-to-bottom order.Source: Figure B.1 · p. 18
• The experiments measure synthetic-to-real distribution discrepancy with energy distance and RBF-kernel MMD, and report training time per epoch.FID scores for image applications are evaluated on 50,000 samples.Source: §B.1 · pp. 17–18

B.2 2d Benchmarks

• The two-dimensional benchmark tables compare methods using W2, ED, and MMD on Swiss roll, moons, and checkerboard.Source: Table B.1 · p. 18, Table B.2 · p. 18

B.3 Higher Dimensional Simulations

• Neal’s funnel has a cusp at one end and a fat tail at the other, and the appendix tests it at d = 10/100 using 10,000 training samples.Source: §B.3 · pp. 18, 20
• PG methods achieve the lowest or comparable discrepancy metrics in the higher-dimensional simulations.The tables compare W2, ED, and MMD for the 36-dimensional eight-Gaussians and 10-dimensional funnel examples.Source: Table B.4 · p. 19, Table B.5 · p. 19, §B.3 · p. 20
• The higher-dimensional appendix includes sample-path visualizations for checkerboard and Neal’s funnel, and energy-distance plots for eight-Gaussians and funnel examples.Source: Figure B.2 · p. 20, Figure B.3 · p. 20

B.4 Scientific Datasets

• The benchmarks include Alanine dipeptide (30d), QM9 (87d), scRNA (50d), and Lattice ϕ4 (64d).Source: §B.4 · p. 20
• Across scientific datasets, PG methods demonstrate better or comparable performance.The appendix also visualizes Alanine dipeptide and QM9 sample paths projected to two dimensions with t-SNE.Source: §B.4 · p. 20, Table B.6 · p. 21, Table B.7 · p. 21

B.5 De novo Molecule Generation

• De novo generation requires one sample to jointly specify all atoms’ 3D coordinates, atom types, formal charges, and pairwise bond orders.Source: §B.5 · p. 20
• QM9 molecule samples span 54–2,433 dimensions, averaging roughly 103, with the bond block accounting for 83% at n = 29.Source: §B.5 · p. 22

• The anisotropic extension replaces the scalar scale with a per-coordinate scale vector.Source: C Anisotropic Probabilistic Geodesic… · p. 23
• The anisotropic Fisher–Rao metric is a product of rescaled Poincaré half-plane metrics, making joint squared distance and geodesic energy additive.Source: C Anisotropic Probabilistic Geodesic… · pp. 23–24
• The joint manifold is a Riemannian product, so its geodesics are Cartesian products of coordinate-wise geodesics.The geodesic equations are decoupled across coordinates.Source: C Anisotropic Probabilistic Geodesic… · p. 23
• For a coordinate with nonzero location displacement, the geodesic is a semicircle in the (µ_i, λσ_i) upper half-plane; with zero displacement, location stays fixed while scale interpolates geometrically.Source: C Anisotropic Probabilistic Geodesic… · pp. 23–24
• The anisotropic path’s training-time cost matches the isotropic counterpart because its per-sample quantities are precomputed in closed form without an ODE solve.Source: C Anisotropic Probabilistic Geodesic… · p. 25