PaperlayerView exact revision on arXiv

Source-linked AI summary

Compatible Likelihoods for Flow Matching on Manifolds

Lucas Ng, Georgios Batzolis, Mark Girolami

arXiv:2610.31981v1cs.LGstat.ML

TL;DR

Likelihood-based flow matching on manifolds must recover the data- and coupling-induced marginal velocity while respecting which velocities are feasible. This paper derives a compatibility condition and error bound, then proposes an endpoint-space exponential-family objective; it reports improvements across several manifold benchmarks and pose-generation metrics.

  • Problem

    Likelihood-based flow-matching objectives may fail to recover the true marginal velocity, while compatible objectives can still assign likelihood to infeasible velocities on compact manifolds.

    Source: §1 · p. 2

  • Method

    The proposed objective uses an endpoint-space exponential family with conditional velocity as sufficient statistic to be both compatible and support-aware.

    Source: §1 · p. 2, §4 · p. 6

  • Results

    The method recovers the expected velocity in the analytic test, improves NLL on two spherical benchmarks, and improves sample-based metrics on high-dimensional tori and human poses.

    Source: §8 · p. 10

  • Takeaways & Limitations

    Compatibility and support-awareness are independent criteria to consider when choosing a flow-matching likelihood.

    Source: §1 · p. 2

    The method requires manifold-specific mean-map and covariance computations, and quadrature or Monte Carlo estimation on general geometries may be prohibitive.

    Source: §6 · p. 10

Abstract

from arXiv · show

On the simplest non-Euclidean manifold, the 1D flat torus, we exhibit a distribution where a recently proposed extension of Riemannian flow matching is globally optimized by a velocity field that pushes samples away from the data distribution. We trace this failure mode to a property we call compatibility; a likelihood-based flow matching objective is compatible if its population minimizer recovers the marginal velocity field induced by the data and coupling. We give a general condition for compatibility, apply it to existing methods, and bound the worst-case error in terms of a residual component vanishing under compatibility. Next, we note that existing compatible objectives put density on velocities that are infeasible on compact manifolds. We propose an objective built from an exponential family over endpoints, with the conditional velocity as sufficient statistic, which is both compatible and supported only on feasible velocities. This manifests as a reweighting of the flow matching loss which, on constant-curvature manifolds, treats errors in flow speed and direction differently. Empirically, our method recovers the correct velocity field on the test case, outperforms prior methods on high-dimensional product tori, and obtains substantial NLL improvements over prior work on two of four standard geospatial benchmarks on the sphere at moderate runtime overhead. On a large dataset of high-dimensional human poses, our method outperforms prior work according to sample-based metrics.

1 Introduction

Flow matching extends to non-Euclidean data by learning a marginal velocity field from conditional paths, but likelihood-based objectives can differ in whether they recover that field and respect feasible velocities. The paper characterizes these properties, derives a compatibility condition and error bound, and proposes a compatible, support-aware method.

1 Introduction

  • Curved manifolds motivate flow-matching methods that accommodate non-Euclidean data supports; RG-VFM does so by maximizing a Riemannian Gaussian likelihood over conditional endpoints.Source: §1 · pp. 1–2
  • Compatibility means the population optimum recovers the true marginal flow, while support-awareness means the likelihood assigns mass only to achievable velocities.The properties are independent: RFM is compatible but not support-aware on compact manifolds, whereas RG-VFM is support-aware but not compatible.Source: §1 · p. 2
  • The general compatibility condition is satisfied by RFM but not RG-VFM, and applies to flow-matching problems with nonlinear conditional paths.The paper demonstrates the distinction analytically on the flat torus T1.Source: §1 · p. 2
  • Squared local-flow error scales with distributional fit and a residual term; compatibility makes that residual zero even when the likelihood is misspecified.This removes that source of error without requiring a well-fitting likelihood model.Source: §1 · p. 2
  • RFM can assign likelihood to infeasible velocities on compact manifolds because its tangent-space Gaussian omits global constraints such as a maximum flow speed.Source: §1 · p. 2
  • The proposed endpoint-space exponential family uses conditional velocity as its sufficient statistic, yielding a compatible, support-aware objective while leaving the RFM pipeline otherwise unchanged.Experiments include an analytic test on T1, high-dimensional tori, four geospatial datasets, and human poses in SO(3) with nearly 12 million training samples.Source: §1 · p. 2

2 Preliminaries

The preliminaries construct conditional paths from geodesics and describe flow matching as learning their induced marginal velocity. They also frame flow matching as endpoint-likelihood optimization, including a Riemannian Gaussian model for manifolds.

2 Preliminaries

  • RFM follows a minimum-length geodesic at constant speed from a base sample to a data sample, using exponential and logarithmic maps to define conditional velocities.The conditional velocity is the logarithmic map from the intermediate point to the endpoint, scaled by 1/(1-t).Source: §2 · p. 2
  • Conditional velocity fields and their coupling determine conditional probability paths through the continuity equation; learning the marginal velocity field gives a transport map from base to data.Source: §2 · pp. 2–3
  • VFM views flow matching as endpoint-posterior likelihood maximization, with Euclidean squared-error flow matching corresponding to a fixed-covariance Gaussian likelihood.RG-VFM extends the endpoint-likelihood view to manifolds using a Riemannian Gaussian with a learned mean parameter.Source: §2 · p. 3

3 Compatibility of training objectives

The section defines compatibility as recovering the data- and coupling-induced marginal velocity at population optima, then tests likelihood-based objectives against that criterion. It shows RG-VFM can fail on a flat torus, relates flow error to a residual, and identifies infeasible velocities in RFM’s likelihood.

3 Compatibility of training objectives

  • Compatibility requires every population minimizer, after decoding, to recover the true marginal velocity induced by the data and coupling.RFM’s squared-error optimum is the conditional-velocity expectation, whereas RG-VFM’s endpoint optimum is a Fréchet mean that generally agrees only in Euclidean space.Source: §3 · p. 3

3.1 A motivating example: incompatibility on a manifold without curvature

  • At the torus crossing, RG-VFM’s globally optimal endpoint prediction decodes to a nonzero velocity, although the true marginal velocity is zero.The two cross-pair geodesics meet with equal and opposite velocities; the decoded prediction has magnitude π/(1 − τ).Source: §3.1 · p. 4, Figure 1 · p. 4

3.2 The condition for compatibility of likelihood-based objectives

  • For likelihood-based objectives, compatibility requires the velocity feature to be recoverable from the likelihood score through a z-independent linear map; RG-VFM fails this condition on non-Euclidean manifolds.Its score and target feature depend on different logarithmic maps, which have no global affine relationship there.Source: §3.2 · p. 5

3.3 Measuring the flow error

  • A compatible likelihood has zero residual and population flow error even when its fitted density is misspecified.For incompatible families, adding expressive capacity does not ensure the residual vanishes, so flow error can persist.Source: §3.3 · p. 6
  • A decoder can add flow error when its practical velocity prediction differs from the fitted likelihood’s velocity moments.For nonlinear conditional paths, the likelihood parameterization should avoid this decoder mismatch as well as satisfy compatibility.Source: §3.3 · p. 6

3.4 Riemannian flow matching does likelihood maximization implicitly

  • RFM is compatible, but its tangent-space Gaussian likelihood assigns density to infeasible velocities on compact manifolds.The endpoint-to-velocity map is locally injective but not surjective, leaving tangent vectors that correspond to no endpoint.Source: §3.4 · p. 6

4 Method

The method models endpoints with an exponential family whose sufficient statistic is the conditional velocity, making the likelihood compatible while restricting support to feasible manifold endpoints. Its resulting velocity-space likelihood preconditions the flow-matching regression, with geometry-dependent effects and practical lookup-table computation.

4 Method

  • The endpoint exponential family uses conditional velocity as its sufficient statistic, satisfying compatibility while placing likelihood mass only on feasible manifold endpoints.Its score with respect to the natural parameter is the conditional velocity, so the compatibility condition holds.Source: §4 · pp. 6–7
  • The method reweights the RFM loss through a covariance-based preconditioner, without making the preconditioning matrix entries a new design choice.Detaching the covariance in the backward pass gives gradients identical to those of the likelihood objective.Source: §4 · p. 7
  • The carrier is chosen to satisfy centering, isotropy, product factorization, terminal localization, and recovery of ordinary Euclidean flow matching.Among radial carriers, product factorization imposes a Gaussian form, terminal locality sets its decay, and matching Euclidean covariance sets its scale.Source: §4 · p. 7
  • On spheres, the resulting loss upweights flow-speed errors relative to direction errors, reflecting the greater effect of magnitude errors under geodesic convergence.The velocity-space likelihood is Gaussian in Euclidean space and truncated Gaussian on flat tori.Source: §4 · p. 7
  • Symmetry reduces covariance computation to low-dimensional functions that can be pretabulated and queried in O(1) time during training.For flat tori, one 1D lookup table serves all factors; product manifolds use one covariance block per factor.Source: §4 · p. 7

5 Experiments

Experiments evaluate the theoretical predictions and proposed method across analytic, torus, geospatial, and human-pose settings. The method performs best in the reported flat-torus dimension sweep, improves over the reproduced baseline on sphere benchmarks, and outperforms prior work on human-pose sample metrics.

Analytic test case

  • At the analytic T1 test point, RFM and our method approach the ideal zero velocity, while RG-VFM splits into two clusters of 15 seeds each.The lower RG-VFM cluster is a stable population local minimum; the upper cluster is below the predicted 4π/3 magnitude because the velocity magnitude is bounded.Source: §5 · p. 8, Table 2 · p. 8
  • RG-VFM's upper cluster has better training loss but poorer generative performance than its lower cluster, and requires more ODE steps.RG-VFM also has the highest inference NFEs and wall times, consistent with rapid spatiotemporal variation in its learned field.Source: §5 · p. 8

High-dimensional product tori

  • NLL is uninformative for RG-VFM in the flat-torus comparison, and some reported values fall below the entropic lower bound.The authors use exact divergence to mitigate possible issues from Rademacher-vector divergence estimates in prior results.Source: §5 · p. 8, Table 3 · p. 9
  • Our method has the best Fourier-kernel MMD2 ordering below dimension 200 and maintains an advantage beyond 200, while RFM and RG-VFM switch ranks.The dimension sweep reports mean and standard deviation over five seeds, with 20,000 samples per seed.Source: §5 · pp. 8–9, Figure 3 · p. 8

Earth and climate science datasets on the sphere

  • Our method improves over the reproduced baseline on all four sphere datasets and over prior work on the volcano and fire datasets.It still struggles, like RFM, on the earthquake and flood datasets; the reported improvements are on volcano and fire.Source: §5 · p. 9, Table 4 · p. 9

High-dimensional human pose data

  • On the large AMASS human-pose dataset, our method outperforms prior work on both sample-based metrics.The method and RFM-family models were trained on 11,823,971 samples; the HuProSO3 checkpoint had also seen the test samples.Source: §5 · p. 9

6 Limitations

The method requires specialization for each manifold and adds work to compute and invert the mean map and pretabulate covariance components. Its implementation is less extensive on general geometries, where the required numerical work may be prohibitive.

6 Limitations

  • General-geometry use requires devising efficient mean-map computation and inversion, plus pretabulating covariance components.This adds implementation work beyond the manifold-specific objects also used by RFM and RG-VFM.Source: §6 · pp. 9–10
  • The method is less extensively implemented on general geometries, where quadrature or Monte Carlo estimation may be prohibitive.The authors target data with known symmetries and geometric structure rather than arbitrary 3D meshes.Source: §6 · p. 10

7 Related Work

Prior work has framed flow matching as endpoint-likelihood maximization and extended that view to manifolds, while other intrinsic generative approaches incur computational or scaling costs. Recent manifold work also includes geometry-specific methods and accelerated-inference approaches.

7 Related Work

  • Guzmán-Cordero et al. [2025] proposed endpoint exponential families for heterogeneous tabular data, but manifold path nonlinearity means their endpoint construction is not generally compatible.Their approach estimates the marginal field from endpoint predictions using linearity of conditional paths.Source: §7 · p. 10
  • Intrinsic generative models can require random-walk simulation, heat-kernel approximations, expensive divergence computation, or Hutchinson estimation; earlier simulation-free methods scale poorly.The passage lists Riemannian CNFs, Moser Flow, CNF Matching, and score-based and ELBO-based diffusion methods.Source: §7 · p. 10
  • Recent manifold methods often specialize RFM to particular geometries or applications, while another line distills flow maps or consistency models for accelerated inference.The passage also notes alternative methods for single manifold classes and generalized flow maps for manifolds.Source: §7 · p. 10

8 Conclusion

The conclusion characterizes compatibility as necessary for likelihood optimization to recover the right flow and proposes a compatible, support-aware endpoint exponential family. The appendix identifies where it presents theoretical derivations and supporting material.

8 Conclusion

  • At population optima, likelihood maximization generically fails to recover the true marginal velocity; compatibility characterizes when it succeeds.The proposed endpoint exponential family is compatible and support-aware, and experiments report recovery on the analytic case and gains on several data settings.Source: §8 · p. 10

Appendix

  • The appendix assigns proofs and derivations to section A, additional experiments and ablations to B, implementation details to C, manifold specializations to D, protocols to E, and runtime analytics to F.Source: Appendix · p. 15

A.1 Compatibility bound

  • Theorem 1 states assumptions including square-integrability, absolute continuity, an interior optimum, a first-order condition, and regularity for a score identity.The following proof passage defines a residual through a least-squares problem, but the supplied theorem statement does not include its conclusion.Source: §A.1 · p. 15

A.2 Motivating example with smooth T1 data

The smooth T1 example contrasts the true marginal field with RG-VFM’s globally optimal decoded velocity, while later appendix sections characterize a Gaussian carrier, its relation to RFM, and the learned field profile.

A.2 Motivating example with smooth T1 data

  • For sufficiently narrow smooth mixture components, the true marginal velocity points toward zero on both sides of the crossing, but RG-VFM’s global optimum is at μ=π.At that optimum, its decoded velocity has magnitude π/(1−t), whereas symmetry gives uτ(0)=0.Source: §A.2 · p. 17

A.3 Characterization of the Gaussian carrier

  • The carrier desiderata—centering, isotropy, product factorization, Euclidean CFM recovery, an RFM limit, and normalizability—narrow the choice to a Gaussian carrier.Terminal locality requires β>0, and matching Euclidean covariance to I selects β=1/2.Source: §A.3 · p. 18

A.4 Limiting behaviour and relationship to Riemannian flow matching

  • On non-compact manifolds with diffeomorphic exponential maps, a special carrier recovers RFM; as t approaches 1, the proposed loss also approaches the RFM sample loss up to a factor of 1/2.On flat manifolds the carrier agrees with the proposed choice up to a constant, while the torus pushforward is truncated.Source: §A.4 · p. 18

B.1 Vector-field profile for the T1 analytic test case

  • RG-VFM’s upper-cluster field nearly jumps at t=1/4 between globally minimizing branches, while RFM and the proposed method remain smooth.The high local Lipschitz constant makes the ODE solver take many small steps; the figure plots vθ_t(x=0) over time.Source: §B.1 · p. 19, Figure 4 · p. 19

B.2 NLL evaluation of RG-VFM

On the S2 checkerboard, RG-VFM’s continuous-flow NLL grows sharply as integration is refined, across fixed-step and adaptive solvers. The carrier ablation favors most tested carriers over c0, while the network uses a barrier outside the invertible mean-map region.

B.2 NLL evaluation of RG-VFM

  • RG-VFM’s NLL rises from 1.7 bits at 50 Euler steps to 723.7 bits at 10,000, while RFM stabilizes near 2.51 bits by 200 steps.Heun and RK4 show qualitatively identical RG-VFM behavior, exceeding 50 bits by 500 steps for both; this rules out a first-order discretization artifact.Source: §B.2 · p. 19, Figure 5 · p. 20
  • With dopri5, RG-VFM’s NLL increases from 16.6 to 533.5 bits as tolerance tightens from 10^-2 to 10^-5, while RFM stays around 2.47 bits.RG-VFM uses about 70× more function evaluations at the same tolerance; runs at tolerance ≤10^-6 exceed the 300-second cap, while RFM completes in under 2 seconds.Source: §B.2 · pp. 19–20, Table 6 · p. 20
  • Across stiff-solver tests, every RG-VFM run hits the 900-second cap, whereas 10 of 12 RFM runs finish with NLLs within 0.04 bits of 2.47.The two RFM exceptions are LSODA and Radau at tolerance 10^-7.Source: §B.2 · p. 20

B.3 Carrier ablation

  • On S2, carriers other than c0 show no clear KL difference except at p = 0, while c0 ranks clearly worst.The ablation uses the wrapped Gaussian target, uniform base, and independent coupling; c0 does not localize the endpoint distribution around x_t as t approaches 1.Source: §B.3 · p. 21

C.1 Network output and domain of the mean map

  • The mean map is invertible in the relative interior of feasible conditional velocities, and the implementation clamps predictions outside that region with an inward-pointing barrier extension.The practical loss agrees with the surrogate loss in the interior; the barrier was active for none of the selected experimental checkpoints.Source: §C.1 · p. 21

C.2 Training algorithms

The appendices describe training and implementation choices, then specialize the surrogate loss to several manifold geometries. They also explain the evaluation metrics used and introduce a Fourier-kernel MMD for high-dimensional tori, with a stated limitation on detecting interactions.

C.3 Implementation on less symmetric manifolds

  • On less symmetric manifolds, the method can compute conditional-velocity means and covariances, numerically invert the mean map, and cache the resulting functions for training.The passage identifies Monte Carlo expectations, offline computation with interpolation, and reuse of learned approximations across datasets as possible implementation routes.Source: §C.3 · p. 22

D Specialization to common geometries

  • The surrogate loss is preconditioned by the conditional-velocity covariance, whose calculation is tractable by exploiting manifold symmetries and pretabulating scalar functions.On the geometries discussed, symmetry reduces covariance calculations to a small number of scalar functions.Source: D Specialization to common geometries · p. 22

D.1 Flat torus Td

  • On flat tori, the endpoint likelihood factorizes into coordinate-wise truncated Gaussians, yielding a diagonal covariance and one scalar loss coefficient per coordinate.As the feasible range expands to all of R, the loss approaches the RFM loss up to a factor of 1/2.Source: §D.1 · pp. 22–23

D.2 Complete simply connected manifolds of constant curvature

  • On constant-curvature manifolds, rotational symmetry reduces the log-partition function to a function of the natural-parameter norm, enabling one-dimensional quadrature and pretabulation.On S2, errors parallel to the predicted velocity receive greater weight than perpendicular errors; both weights tend to one as t approaches 1.Source: §D.2 · pp. 23–24

D.3 Rotation groups and products SO(3)n

  • On SO(3)^n, the likelihood and volume factorize across rotations, so the covariance separates into 3 × 3 blocks and the same two scalar functions can be reused across factors.The network may use all n rotations to predict each velocity; the factorization applies to the covariance.Source: §D.3 · pp. 25–26

E.1 Evaluation framework

  • For geospatial benchmarks, the authors report NLL because small test sets produced negative unbiased kernel estimates, but this leaves RG-VFM comparisons unavailable because its NLLs are uninformative.They use Wasserstein distances for the analytic two-Dirac test case, and energy distance with the square-root trick and Riemannian Laplace-kernel MMD alongside it.Source: §E.1 · p. 26

E.2 A Fourier-kernel MMD for high-dimensional tori

  • For high-dimensional product tori, the paper introduces a Fourier-kernel MMD pseudometric using the wrapped Gaussian’s closed-form characteristic function to avoid estimation uncertainty.The finite-frequency construction detects coordinate-marginal discrepancies but may miss pairwise interactions at high dimensions.Source: §E.2 · pp. 26–27

E.3 Benchmark protocols

The benchmark protocol follows the RFM data splits and specifies separate solvers for the analytic, torus, and geospatial experiments. The human-pose evaluation uses an in-distribution SO(3)^19 split, and adaptive sampling is compared with an autoregressive baseline.

E.3 Benchmark protocols

  • The analytic, torus, and geospatial experiments use an 80/10/10 train-validation-test split, with solver and evaluation settings varying by dataset.The analytic case uses dopri5 at tolerance 10^-5; the high-dimensional torus uses 1,000 Euler steps with per-step projection; geospatial evaluation uses dopri5 at tolerance 10^-7.Source: §E.3 · p. 27

E.4 Human pose training and evaluation

  • The human-pose experiment trains and evaluates in-distribution on an SO(3)^19 partition of AMASS, using 11,823,971 training samples and a 6-layer, width-1,024 MLP.Evaluation reports NLL with respect to Haar measure and energy distance and Gaussian MMD via chordal distance.Source: §E.4 · p. 28

F.1 Adaptive sampling time on SO(3)n

  • Human-pose sampling results use 1,000-step Euler integration, while the adaptive-sampling timing comparison uses dopri5 for flow models and an autoregressive forward pass for HuProSO3.RFM and the proposed method have the same Euler sampling cost because they share the forward-pass architecture; adaptive sampling time depends on the learned neural approximation.Source: §F.1 · p. 28, Table 8 · p. 28

F.2 End-to-end training throughput in the absence of compilation

The throughput benchmark isolates training-loss runtime under a shared training setup, while profiling reports small allocator-memory overhead and near-negligible per-iteration runtime overhead with compilation. The reported experiments use roughly 250 hours of A100 time, excluding additional untracked project compute.

F.2 End-to-end training throughput in the absence of compilation

  • The training-throughput comparison keeps the architecture and execution setup fixed while changing only the loss, and the proposed method’s throughput penalty is largest at low dimensions and shrinks at high dimensions.The benchmark reports iterations per second over five runs with torch.compile disabled; runtime reflects both pretabulation and lookup/training-loop costs.Source: §F.2 · pp. 28–29, Table 9 · p. 29

F.3 Per-iteration time and memory profiling

  • Profiling finds allocator memory overhead of only a few MiB and near-negligible runtime overhead per iteration with torch.compile.The authors report both allocator and driver memory measures because allocator tracking omits some CUDA allocations and pools.Source: §F.3 · p. 29, Table 10 · p. 30

F.4 Compute accounting

  • The analytic, torus, and geospatial experiments together consumed roughly 250 hours of A100 time, while additional preliminary solver sweeps and prototypes were not fully accounted for.The torus evaluation becomes more time-consuming at higher dimensions, reaching around one hour per run at dimension 400.Source: §F.4 · p. 30

Related papers

Loading 2610.31981v1…