Compatible Likelihoods for Flow Matching on Manifolds

AI summary

Overview

Research area: Generative modelling — specifically flow matching and continuous normalizing flows extended to Riemannian manifolds, with a focus on how training objectives are derived from likelihood models.

Technical level: Advanced. The paper assumes familiarity with Riemannian geometry (exponential/logarithmic maps, geodesics, tangent spaces, Fréchet/Karcher means), the continuity equation, and variational/likelihood-based formulations of flow matching.

Scope (one sentence): The paper defines a correctness criterion called compatibility for likelihood-based flow matching objectives on manifolds, proves that a recent Riemannian Gaussian method fails it while Riemannian flow matching satisfies it, and proposes a new exponential-family objective that is both compatible and restricted to physically feasible velocities. (Note: the supplied text is truncated partway through Table 3, so several later experimental figures are not available and are flagged as not reported below.)

What This Paper Is About

Flow matching trains a generative model by regressing a neural network onto simple "conditional" velocity fields that carry samples from a base distribution to data. A recent line of work reframes these losses as maximum likelihood under a chosen probability model, which makes it easy to invent new objectives for curved data — but not all such objectives are actually correct. This paper asks what property a likelihood-based flow matching loss must have so that its best possible fit reproduces the true marginal velocity field, shows that a recently proposed manifold method (RG-VFM) violates it with a concrete counterexample on the 1D flat torus, and introduces a replacement objective that satisfies the property while also refusing to place probability mass on velocities that cannot occur on a compact manifold.

Key Contributions

1. The compatibility criterion and a general test for it. The authors define a likelihood-based flow matching objective as compatible if its population minimiser, after decoding to a velocity, equals the conditional expectation of the true conditional velocity under the data posterior. They derive a checkable condition (an affine relationship between the target feature and the model score, required to hold at every parameter value, not just the optimum) and apply it to existing methods.

2. A counterexample showing RG-VFM is incompatible on a flat manifold. On the 1D flat torus — which has no curvature at all — they construct a two-point base and two-point data distribution with independent coupling where the RG-VFM loss is globally minimised by an endpoint prediction that decodes to a velocity pointing away from the data. They also analyse a smoothed version of the construction to rule out the possibility that the failure is an artefact of singular distributions.

3. A quantitative error bound. Theorem 1 bounds the gap between the true marginal velocity and the velocity implied by the fitted likelihood by the product of a distributional misfit term (the chi-squared divergence between the data posterior and the fitted model) and the operator norm of the residual of a least-squares fit. The residual vanishes exactly when the objective is compatible, so a compatible but badly misspecified likelihood can still have zero flow error. A second decomposition adds a "likelihood–decoder mismatch" term for methods that decode a network head into a velocity rather than computing a likelihood moment.

4. A new objective that is compatible and support-aware. They place an exponential family over endpoints with the conditional velocity as sufficient statistic, which is compatible by construction and only puts mass on endpoints that exist on the manifold. Training reduces to a reweighting of the standard Riemannian flow matching loss by the inverse Fisher information (with a stop-gradient), evaluated via pretabulated lookup tables.

Main Findings

• The counterexample is exact and the failure is global, not local. With base distribution ½δ₋α + ½δ₊α and data distribution ½δ₋β + ½δ₊β (independent coupling, α = τβ/(1−τ), β ∈ (π/2, π(1−τ))), the two cross-pair geodesics meet at x = 0 at time t = τ carrying equal and opposite velocities, so the true marginal velocity there is zero. The RG-VFM loss at that point has local minima at μ = 0 and μ = π; because β > π/2 the shorter arc runs through π, making μ = π the unique global optimum, which decodes to a velocity of magnitude π/(1−τ) — the opposite of the correct answer.

• Empirically, RG-VFM splits into two populations. Across 30 seeds on the analytic test case (τ = 1/4, α = π/5, β = 3π/5), the velocity magnitude ‖v‖ at (0, 1/4) is bimodal: two clusters of 15 seeds each, with the upper cluster at 4.123 ± 0.042 and the lower at 0.037 ± 0.019 (overall RG-VFM mean 2.080 ± 2.078). RFM reaches 0.070 ± 0.040 and the proposed method 0.060 ± 0.041, against an ideal value of zero. The upper cluster sits just below the value 4π/3 implied by the analytic calculation because flow speed is bounded on a compact manifold.

• Better training loss did not mean better samples. The 15 RG-VFM seeds in the upper cluster have better training loss than those in the lower cluster, yet perform worse on generative metrics and require more ODE steps. Ranking by Wasserstein-1 distance agrees with ranking by velocity error at (0, 1/4), supporting the claim that recovering the marginal velocity matters for generation; kernel-based metrics favour the proposed method and RFM over both RG-VFM clusters.

• Large inference cost for the incompatible model. Median dopri5 NFE and wall time (tolerance 10⁻⁵) were 497 [426, 578] and 1.92 s [1.65, 2.22] for RFM; 443 [388, 582] and 1.71 s [1.49, 2.24] for the proposed method; and 36,839 [33,796, 60,574] and 141.4 s [130.3, 233.9] for RG-VFM (upper cluster alone: 55,436 [32,228, 70,115] and 212.9 s [124.0, 271.9]). The authors attribute this to rapid spatiotemporal variation in the learned RG-VFM field.

• Consistent MMD advantage on product tori up to dimension 200. On 𝕋^d with a uniform base and wrapped Gaussian data of scale σ = 0.2, Fourier-kernel MMD² (5 seeds, 20,000 samples per seed) shows a consistent ordering below dimension 200: the proposed method best, RFM second, RG-VFM last. Beyond that, the proposed method keeps its advantage but RFM and RG-VFM swap ranks.

• NLL is a broken yardstick in this setting, and the paper says so. Because σ = 0.2 ≪ 2π, the per-dimension NLL lower bound is the differential entropy of the data, ≈ ½log(2πeσ²) ≈ −0.1905, exact to four decimal places. Several NLL values reported by Chen and Lipman (2024) fall below this entropic bound, possibly because they estimated divergence with Rademacher vectors; the authors therefore use exact divergence via jacrev. In their Table 3, RG-VFM records the lowest NLL (for example −0.2171 ± 0.001 at d = 1 and −0.2142 ± 0.003 at d = 2) but these are shaded as below the bound, i.e. uninformative. The proposed method lands at −0.1904 ± 0.001 (d = 1), −0.1900 ± 0.005 (d = 2) and −0.1879 ± 0.004 (d = 5), close to the bound. (Table 3 is truncated at d = 5, so rows for larger d and the remaining baseline values are not available in the supplied text.)

• Gains on real manifold data. The abstract reports substantial NLL improvements over prior work on two of four standard geospatial benchmarks on the sphere at moderate runtime overhead, and better sample-based metrics than prior work on a dataset of human poses in SO(3)¹⁹ with nearly 12 million training samples. Specific dataset names and metric values are not given in the supplied excerpt.

• Design of the carrier mostly collapses. Requiring a centred, isotropic carrier that factorises on product manifolds, localises around x_t as t → 1, and recovers ordinary flow matching in Euclidean space forces a Gaussian-in-distance form; the resulting velocity-space likelihood is Gaussian in Euclidean space, a truncated Gaussian on flat tori, and on spheres upweights errors in speed relative to errors in direction.

Methodology in Plain English

The authors start from the observation that flow matching can be read as fitting a probability model to the unknown "endpoint posterior" — the distribution over where a sample came from, given where it is now. Different choices of that probability model give different training losses. They then ask a simple-sounding question: if you had infinite data and a perfect network, would your loss lead you to the field that actually transports the base distribution to the data? If yes, they call it compatible.

To test compatibility they use a perturbation argument. If you nudge the data distribution in a way that the likelihood cannot detect to first order, but which does change the average velocity, then the objective can be perfectly optimised while getting the flow wrong. This turns into a concrete algebraic check: the quantity you want to predict must be an affine function of the model's score, and this must hold everywhere in parameter space.

They demonstrate the failure concretely by building a small, fully hand-computable example on a circle (the 1D flat torus) where two geodesics cross at equal and opposite speeds, so the correct average velocity is zero — yet the competing method's globally optimal answer points the other way. Because the example uses point masses, they repeat the analysis with smoothed distributions in an appendix.

They then quantify how bad things can get with a bound: the velocity error is at most the product of how badly the probability model fits the data and the size of an unavoidable residual. Compatibility makes that residual zero, which is why a compatible model can be crude as a density model and still get the flow right.

For their own method, they choose a probability model over endpoints whose "sufficient statistic" is exactly the conditional velocity. This makes the score automatically affine in the velocity, so compatibility holds by construction, and it never assigns probability to endpoints off the manifold. Differentiating the negative log-likelihood gives a weighted regression: instead of weighting every velocity error equally, they weight it by the inverse covariance (Fisher information) of the model at the predicted velocity. To avoid computing that covariance online at every step, they exploit symmetry — on flat tori the covariance is diagonal with each entry depending on a single scalar, so one 1D lookup table serves all dimensions; on constant-curvature spaces isotropy reduces it to two scalar functions of magnitude; on product manifolds it is block diagonal — giving O(1) table lookups during training.

Experiments cover four settings: the hand-computable circle example (30 seeds), a dimension sweep on flat tori, four geospatial benchmarks on the sphere, and human pose data in SO(3)¹⁹.

Why This Matters

The paper's central contribution is a correctness checklist for a fast-growing family of methods. Reframing flow matching as likelihood maximisation has made it easy to propose new objectives for graphs, manifolds, and other exotic data supports; the authors show that one such published objective can be globally optimal and still learn a field that pushes samples away from the data, on the simplest possible curved space. That is a cautionary result for anyone designing losses by choosing a "geometrically natural" density.

It also reframes a debate in the literature. Zaghen et al. (2026) argued that RG-VFM captures full geometric structure while RFM lacks curvature information; the counterexample here is on a manifold with no curvature, showing the issue is the objective rather than the geometry. Conversely, the authors grant that RFM has a genuine defect — it assigns probability to velocity vectors that no endpoint on a compact manifold could produce — and their method fixes both problems at once.

Real-world applications:

• Geospatial and atmospheric modelling, where data live on the sphere (the paper evaluates on four standard geospatial benchmarks) and where well-calibrated likelihoods matter for downstream scientific use.
• Molecular and protein structure generation, where conformations live on products of rotation groups and tori and where unphysical intermediate states are costly.
• Robotics and motion capture, where pose data on SO(3)¹⁹ is the paper's large-scale test case (nearly 12 million training samples) and where generated motion must respect joint constraints.
• Directional and periodic data in fields such as winds, ocean currents, or crystallography, which are naturally described by tori and spheres rather than Euclidean space.

Industry relevance: the method is a drop-in reweighting of the existing RFM pipeline, requires no new network architecture, and uses pretabulated lookup tables that different datasets on the same manifold can share. The reported "moderate runtime overhead" and the dramatic reduction in ODE solver steps versus RG-VFM (roughly two orders of magnitude in the circle experiment) matter for deployment, where sampling cost often dominates.

Future Directions

• Auditing other likelihood-based objectives. The authors explicitly anticipate future work proposing "natural-looking" likelihood-based flow matching losses that have never been checked against the compatibility condition, and offer their affine-score test as the tool for doing so.

• Choosing and stress-testing the carrier. The carrier function is described as "apparently open" but mostly pinned down by structural constraints; an appendix studies sensitivity to it. A principled theory of when the product-factorisation, terminal-locality, and Euclidean-recovery requirements can be relaxed remains open.

• Extending the tabulation strategy to manifolds with fewer symmetries. The O(1) lookup-table trick relies on symmetry to keep the number of scalar functions small. On manifolds without large isometry groups — or with inhomogeneous geometry — the table dimension could grow, and how gracefully the method degrades is not established in the supplied text.

• Reconciling NLL-based evaluation with sample quality. The paper documents that RG-VFM achieves the best (indeed, impossibly low) NLL while generating worse samples and requiring far more solver steps, and that some published RFM NLL values sit below the entropic bound. Establishing evaluation protocols that are not fooled by this — exact divergences, entropic sanity checks, or sample-based metrics — is a live problem the results point to but do not solve.

Target Audience

Researchers and practitioners working on generative models for non-Euclidean data — particularly those building or extending Riemannian flow matching, Riemannian diffusion, or variational flow matching. It will be most valuable to readers comfortable with differential geometry and likelihood theory who are choosing or designing a training objective for manifold-valued data, and to reviewers and methodologists who need a criterion for judging newly proposed likelihood-based losses. Readers interested only in applications can extract the practical message (prefer compatible, support-aware objectives; verify before adopting a new likelihood) without following the derivations.

Authors’ abstract

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.