Flow Matching on General Geometries

TL;DR

Riemannian Flow Matching (RFM) extends continuous normalizing flows to general manifolds by leveraging premeterics to define target vector fields and by using spectral distances for complex geometries. The framework delivers a simulation-free training paradigm on simple geometries and a scalable approach on general manifolds via RCFlow Matching (RCFM) and conditional flows. It achieves state-of-the-art results across diverse non-Euclidean datasets, including spherical Earth data, toroidal protein representations, high-dimensional tori, and mesh-based geometries with curvature and boundaries. The work broadens CNF applicability to challenging geometric domains while maintaining tractable training and principled handling of manifold structure.

Abstract

We propose Riemannian Flow Matching (RFM), a simple yet powerful framework for training continuous normalizing flows on manifolds. Existing methods for generative modeling on manifolds either require expensive simulation, are inherently unable to scale to high dimensions, or use approximations for limiting quantities that result in biased training objectives. Riemannian Flow Matching bypasses these limitations and offers several advantages over previous approaches: it is simulation-free on simple geometries, does not require divergence computation, and computes its target vector field in closed-form. The key ingredient behind RFM is the construction of a relatively simple premetric for defining target vector fields, which encompasses the existing Euclidean case. To extend to general geometries, we rely on the use of spectral decompositions to efficiently compute premetrics on the fly. Our method achieves state-of-the-art performance on many real-world non-Euclidean datasets, and we demonstrate tractable training on general geometries, including triangular meshes with highly non-trivial curvature and boundaries.

Figures (14)

• Figure 1: Our approach makes use of user-specified premetrics on general manifolds to define flows. On select simple manifolds, the geodesic can be computed exactly and leads to a simulation-free algorithm. On general manifolds where the geodesic is not only computationally expensive but can lead to degeneracy ( e.g., along boundaries), we propose the use of spectral distances ( e.g., biharmonic), which can be computed efficiently contingent on a one-time processing cost.
• Figure 2: The conditional vector field $u_t(x|x_1)$ defined in equation \ref{['e:cond_ut']} transports all points $x\ne x_1$ to $x_1$ at exactly $t=1$.
• Figure 3: Contour plots of geodesic and spectral distances (to a source point) on general manifolds. Geodesics are expensive to compute online and are globally non-smooth. The biharmonic distance behaves smoothly while the diffusion distance requires careful tuning of the hyperparameter $\tau$.
• Figure 4: Visualization of (top) the eigenfunctions that were used to construct target distributions, and (bottom) the learned density & samples from trained models with the Biharmonic distance.
• Figure 5: Riemannian Flow Matching scales incredibly well to higher dimensions as it is simulation-free and all quantities required for training are computed exactly on simple geometries such as tori. Log-likelihoods are in bits.

Theorems & Definitions (4)

• Theorem 3.1
• Proposition 3.2
• proof
• proof