PaperlayerView exact revision on arXiv

Source-linked AI summary

Flow Matching on General Geometries

Ricky T. Q. Chen, Yaron Lipman

arXiv:2302.03660v3cs.LGcs.AIstat.ML

TL;DR

Generative modeling on manifolds remains difficult because existing methods can require simulation, scale poorly in high dimensions, or rely on biased approximations. RFM trains continuous normalizing flows using premetric-defined conditional vector fields, achieving state-of-the-art performance across manifold datasets and tractable training on challenging general geometries.

  • Problem

    Existing manifold generative models may require expensive simulation, fail to scale to high dimensions, or use approximations that produce biased training objectives.

  • Method

    RFM trains continuous normalizing flows by regressing conditional vector fields defined through user-specified premetrics, using geodesics or spectral distances on general manifolds.

  • Results

    RFM achieves state-of-the-art performance across manifold datasets and enables tractable training on triangular meshes and manifolds with non-trivial boundaries.

  • Takeaways & Limitations

    RFM provides simulation-free training on simple geometries and closed-form, divergence-free training objectives that extend to general geometries.

  • Takeaways & Limitations

    Compared with approximated score-based methods, RFM’s stated zero-approximation-error guarantee applies to simple geometries with closed-form geodesics; general geometries use spectral-distance approximations.

Abstract

from arXiv · show

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.

1 INTRODUCTION

RFM addresses the difficulty of training continuous normalizing flows on general manifolds by using premetrics to define target vector fields. It is simulation-free on simple geometries and extends to general geometries through efficient spectral distances, achieving strong performance across manifold settings.

  • 1 INTRODUCTION: RFM learns continuous normalizing flows on general Riemannian manifolds by regressing conditional vector fields toward individual training examples.This conditional formulation addresses the intractability of the implicitly defined target field.
  • 1 INTRODUCTION: Premetrics define the conditional vector field, with geodesic distance providing a natural choice that recovers straight Euclidean trajectories.For general manifolds, spectral distances such as biharmonic distances provide an efficient alternative after one-time processing.
  • 1 INTRODUCTION: RFM is completely simulation-free where geodesics have closed forms and avoids solver differentiation, stochastic iterative sampling, and divergence estimation on general geometries.The method retains a relatively simple ODE-based procedure when exact geodesics are unavailable.
  • 1 INTRODUCTION: The paper compares RFM with closely related continuous-time generative models and reports that it is the only considered method addressing general geometries.The comparison is summarized in Table 1.
  • 1 INTRODUCTION: RFM achieves state-of-the-art performance across manifold datasets, remains scalable to higher dimensions, and enables training on triangular meshes and manifolds with non-trivial boundaries.The reported results are on par with or outperform competitive baselines across various settings.

2 PRELIMINARIES

The paper formulates generative modeling on complete connected smooth Riemannian manifolds using tangent-space vector fields, probability paths, and flows. Continuous normalizing flows model these paths through time-dependent deformations, while prior manifold variants often require simulation during training.

  • 2 PRELIMINARIES: The domain is a complete connected smooth Riemannian manifold M with metric g, tangent spaces T_xM, tangent bundle TM, and Riemannian volume element dvol_x.The metric defines inner products on tangent spaces, and div_g denotes Riemannian divergence.
  • 2 PRELIMINARIES: Probability densities on M are non-negative functions integrating to one with respect to the manifold volume element, and probability paths are time-indexed curves in this density space.These paths provide supervision signals for training generative models.
  • 2 PRELIMINARIES: A flow is a diffeomorphism generated by integrating a time-dependent vector field through an ordinary differential equation on M.The final map is obtained from the flow at time t = 1.
  • 2 PRELIMINARIES: A continuous normalizing flow models a probability path through a continuous-time deformation of a base distribution by parameterizing its vector field.The construction uses the Riemannian instantaneous change-of-variables formula.
  • 2 PRELIMINARIES: Existing manifold variants of continuous-time generative models include approaches requiring simulation, alongside some simulation-free variants.This contrast motivates the paper’s flow-matching formulation.

3 METHOD

Riemannian Flow Matching adapts Flow Matching to Riemannian manifolds by learning tangent-space vector fields through conditional flows defined by premetrics. Geodesic premetrics yield simulation-free training on simple geometries, while general manifolds require forward ODE simulation without differentiating through the solver.

  • 3.1 FLOW MATCHING ON MANIFOLDS: Riemannian Flow Matching trains a continuous normalizing flow by regressing a target vector field that transports base distribution p toward data distribution q on a manifold.The Flow Matching objective compares predicted and target tangent vectors using the manifold’s Riemannian metric.
  • 3.1 FLOW MATCHING ON MANIFOLDS: Conditional probability paths start from p and concentrate at each training example x1 at t = 1, then marginalize over q(x1) to form the population path.This construction satisfies the required boundary conditions p0 = p and p1 = q.
  • 3.2 CONSTRUCTING CONDITIONAL FLOWS THROUGH PREMETRICS: A premetric d(x, y) defines conditional flows by measuring closeness, requiring non-negativity, positivity, and non-degeneracy.The scheduler κ(t) controls the prescribed monotonic decrease of the premetric toward zero.
  • 3.2 CONSTRUCTING CONDITIONAL FLOWS THROUGH PREMETRICS: The resulting conditional vector field is the minimal-norm solution that decreases the premetric and is defined everywhere except at x = x1.The theorem establishes that this field generates the desired conditional flow and reaches x1 at t = 1.
  • 3.2 CONSTRUCTING CONDITIONAL FLOWS THROUGH PREMETRICS: On general manifolds and premetrics, training requires forward ODE simulation to obtain xt, but it does not require differentiating through the solver.This limitation contrasts with the simulation-free case available for simple geometries.
  • 3.2 CONSTRUCTING CONDITIONAL FLOWS THROUGH PREMETRICS: Using geodesic distance makes conditional trajectories geodesics, enabling closed-form computation of xt and simulation-free training on manifolds with closed-form geodesics.Examples include Euclidean space, hyperspheres, hyperbolic space, high-dimensional tori, and some matrix Lie groups.

4 RELATED WORK

Prior manifold generative models face topology, simulation, approximation, and divergence-estimation challenges. RFM addresses these issues with exact conditional vector fields and simulation-free training on simple geometries.

  • Manifold-to-Euclidean flows can create theoretical and numerical problems when the manifold is not homeomorphic to Euclidean space.
  • Diffusion-based manifold methods require simulated random walks because the manifold Ornstein–Uhlenbeck SDE lacks closed-form solutions.
  • Approximating conditional score functions can bias denoising-score-matching gradients, whereas RFM uses exact conditional vector fields.
  • Implicit score matching avoids conditional-score approximation but requires divergence computation whose Hutchinson-estimator variance worsens with dimension and non-Euclidean geometry.
  • RFM is simulation-free on simple geometries, avoids divergence computation, and extends to general manifolds using premetrics that need not satisfy every metric axiom.
  • RFM builds on Euclidean Flow Matching by regressing directly onto generating vector fields with ODEs rather than score functions or SDEs.

5 EXPERIMENTS

Experiments evaluate RFM on diverse manifold datasets, including spheres, tori, high-dimensional tori, triangular meshes, and manifolds with boundaries. The method maintains performance at higher dimensions and supports general mesh geometries.

  • The experiments cover earth and climate data, protein structures, high-dimensional tori, closed manifolds, and maze-shaped manifolds with non-trivial boundaries.
  • On high-dimensional tori, RFM shows no significant performance drop as dimension increases, unlike methods relying on approximations or stochastic divergence estimation.
  • For triangular meshes, exact geodesics are too expensive for frequent training evaluations, motivating spectral distances as practical alternatives.
  • Diffusion distance required careful tuning of τ, while biharmonic distance worked out-of-the-box and provided better smoothness properties.
  • Figure 5 reports log-likelihoods in bits while illustrating RFM's scaling to higher-dimensional simple geometries such as tori.
  • On maze-shaped manifolds, RFM trains with biharmonic distance using k=30 eigenfunctions for distributions constrained by non-trivial boundaries.

6 CONCLUSION

The paper concludes that RFM provides scalable continuous normalizing-flow training on manifolds, with simulation-free exactness on simple geometries and tractable training on general geometries and boundaries.

  • RFM is presented as a highly scalable approach for training continuous normalizing flows on manifolds.
  • On geometries with closed-form geodesics, the method is simulation-free and introduces zero approximation errors.
  • The paper demonstrates tractable training on general geometries, including closed manifolds and manifolds with boundaries.
  • The theoretical development adapts Conditional Flow Matching derivations and proofs to the Riemannian setting.

B PROOF OF THEOREM 3.1

Theorem 3.1 establishes that the premetric-defined conditional flow follows the intended distance-scaling path and has minimal norm among flows satisfying the same constraint.

  • Theorem 3.1 states that the flow generated by the proposed conditional vector field satisfies the target path equation.
  • Along the flow, differentiating the premetric reduces the trajectory condition to an ODE for the time-dependent distance.
  • Solving the ODE yields d(x_t, x_1)=d(x, x_1)κ(t), and uniqueness establishes the desired path relation.
  • The proposed vector field is parallel to the premetric gradient and therefore gives the minimal-norm solution of the associated linear system.

C PROOF OF PROPOSITION 3.2

Proposition 3.2 shows that choosing geodesic distance as the premetric makes the conditional flow trace a geodesic from x0 to x1. The proof establishes the endpoints and verifies that the resulting curve has minimum length.

  • With d(x, y) equal to the geodesic distance, the conditional vector field with scheduler κ(t) = 1 − t defines a geodesic connecting x0 to x1.The curve starts at x0 when t = 0 and ends at x1 when t = 1.
  • The proof uses the endpoint identities ψ0(x0|x1) = x0 and ψ1(x0|x1) = x1.
  • Applying the Riemannian logarithm map and the chain rule yields the velocity expression used to analyze the curve’s length.
  • The curve realizes a minimum of the length function, and therefore is a geodesic.

E LIMITATIONS

The method remains computationally constrained on general manifolds: constructing paths requires simulation, while spectral distances may require expensive eigenfunction solvers on complex geometries.

  • On general manifolds, Riemannian Flow Matching still requires simulating xt, making the sequential process potentially time consuming.
  • Computing spectral distances can also be computationally expensive on complex manifolds because it requires eigenfunction solvers.
  • Approximate neural eigenfunctions are suggested as a possible remedy, and exact eigenfunction solutions are not required for the premetric properties.

G.1 ON THE USE OF APPROXIMATE SPECTRAL DISTANCES AS THE PREMETRIC

The paper argues that finite spectral decompositions can provide useful premetrics without reproducing heat-kernel scores exactly. Very few eigenfunctions can satisfy the required geometric properties, including on manifolds with boundaries.

  • Finite spectral distances need only distinguish almost every pair of manifold points rather than approximate the distances perfectly.
  • For local neighborhoods, the number of required eigenfunctions is the manifold dimension, although they need not be the first eigenfunctions.
  • Using the k smallest eigenvalues is equivalent to Euclidean distances in a k-dimensional eigenfunction embedding that preserves neighborhoods optimally.
  • The heat-kernel approximation is exact only as k →∞, and at small time values its conditional score may lack even the first significant digit in expectation.
  • For the biharmonic spectral distance, the premetric properties are already satisfied at k = 3, while larger k values only refine the distance.
  • On manifolds with boundaries, Neumann boundary conditions make eigenfunction normal derivatives vanish, satisfying the premetric boundary condition and preventing outward-pointing conditional fields.

H EXPERIMENT DETAILS

The experiments evaluate RFM across wrapped and constrained manifold distributions, triangular meshes, and hyperbolic geometry. They also measure implementation choices, likelihood computation, numerical integration, and training runtime.

  • Training setup: Experiments use five runs with seeds 0–4, an 80%/10%/10% train-validation-test split, and a single NVIDIA V100 GPU with 32GB memory.
  • High-dimensional tori: High-dimensional torus experiments use wrapped Gaussians with uniformly sampled means and scale 0.2.
  • Training setup: The model uses a three-hidden-layer MLP with layers of size 512, trains for 50000 iterations with batch size 512, and reports bits-per-dimension log-likelihood.
  • Triangular meshes: Triangular-mesh experiments include Spot the Cow and a Stanford Bunny downsampled to 5000 triangles, with target distributions constructed from thresholded eigenfunctions.
  • Maze manifolds: Maze manifolds are built from triangular meshes whose disconnected neighboring cells create boundaries, with structures sampled using breadth-first search.
  • Vector field parameterization: Vector fields are parameterized in the ambient space and projected onto the manifold and its tangent space at each point.
  • Runtime: Simulation-free training reaches 104.04 iterations per second on a flat torus, compared with 6.36 iterations per second for 200-step simulation.
  • Runtime: The reported runtime advantage is roughly 17x over iterative ODE/SDE solving, including the full training loop.

J.2 MANIFOLD OF SYMMETRIC POSITIVE MATRICES

The SPD-manifold experiments use EEG-derived covariance matrices to evaluate likelihood, manifold validity, and the effects of geometric choices during training. Riemannian geodesics generally outperform Euclidean paths, while the Riemannian norm helps preserve validity during simulation, especially for larger matrices.

  • Dataset: EEG signals are converted into covariance matrices that form the SPD-manifold datasets.The experiments use standard preprocessing for EEG covariance analysis.
  • Evaluation: Table 6 evaluates models using negative log-likelihood and the percentage of simulated samples that remain valid SPD matrices.The experiments also ablate the Riemannian geodesic and Riemannian norm.
  • Riemannian geodesic: Riemannian geodesics generally perform better than Euclidean geodesics because they respect the underlying manifold geometry.Figure 13 visualizes the contrast between the two paths on the 2 × 2 SPD cone.
  • Riemannian norm: The choice of norm does not change the theoretically optimal vector field, but it changes which regions limited-capacity networks emphasize during optimization.The effect is especially relevant where the metric tensor becomes large near singular, ill-conditioned matrices.
  • Riemannian norm: For larger SPD matrices, using the Riemannian norm is important for preventing simulated vector fields from leaving the SPD manifold.This preserves the validity of simulated outputs when matrices approach singularity.
Loading 2302.03660v3…