Riemannian Flow Matching on General Geometries

  • 2023-02-07 18:21:24
  • Ricky T. Q. Chen, Yaron Lipman
  • 19

Abstract

We propose Riemannian Flow Matching (RFM), a simple yet powerful frameworkfor training continuous normalizing flows on manifolds. Existing methods forgenerative modeling on manifolds either require expensive simulation,inherently cannot scale to high dimensions, or use approximations to limitingquantities that result in biased objectives. Riemannian Flow Matching bypassesthese inconveniences and exhibits multiple benefits over prior approaches: Itis completely simulation-free on simple geometries, it does not requiredivergence computation, and its target vector field is computed in closed formeven on general geometries. The key ingredient behind RFM is the constructionof a simple kernel function for defining per-sample vector fields, whichsubsumes existing Euclidean cases. Extending to general geometries, we rely onthe use of spectral decompositions to efficiently compute kernel functions. Ourmethod achieves state-of-the-art performance on real-world non-Euclideandatasets, and we showcase, for the first time, tractable training on generalgeometries, including on triangular meshes and maze-like manifolds withboundaries.

 

Quick Read (beta)

loading the full paper ...