Flow Matching on General Geometries
Authors: Ricky T. Q. Chen
Yaron Lipman
Venue: ICLR 2024 oral (2024)
Keywords: general manifolds
diffusion models
continuous normalizing flow
Primary area: learning on graphs and other geometries & topologies
TL;DR: We derive sufficient conditions for Conditional Flow Matching on general manifolds, with significant algorithmic improvements to diffusion-based approaches even on simple manifolds, and the first to tackle more general manifolds.

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 real-world non-Euclidean datasets, and we demonstrate tractable training on general geometries, including triangular meshes with highly non-trivial curvature and boundaries.

Published as a conference paper at ICLR 2024
FLOW MATCHING ON GENERAL GEOMETRIES
Ricky Q. Chen
Yaron Lipman
FAIR, Meta
FAIR, Meta and Weizmann Institute of Science y i aep.annmeta.commmmmmmm
rtgiicchenmeta.com com
ABSTRACT
We propose Riemannian Flow Matching (RFM), a simple yet powerful framework for training continuous normalizing flows manifolds. Existing methods for gen- erative modeling on manifolds either require expensive simulation, are inherently unable scale high dimensions, or approximations for limiting quantities that result in biased training objectives. Riemannian Flow Matching bypasses these limitations and offers several advantages over previous approaches: is simulation- free simple geometries, does not require divergence computation, and computes target vector field in closed-form. The key ingredient behind RFM is the construction of relatively simple premetric for definng 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-ar performance on many real-world non- Euclidean datasets, and we demonstrate tractable training general geometries, including triangular meshes with highly non-trivial curvature and boundaries.
INTRODUCTION
While generative models have recently made great advances in fitting data distributions in Euclidean spaces, there still chal- lenges in dealing with data residing in non-Euclideds spaces, specifically on general manifolds. These challenges include scalability to high dimensions (e.g., (Rozen et al., 2021)), the requirement for simulation or iterative sampling during training even for simple geometries like hyperspheres (e.g.. (Mathieu Nickel, 2020; De Bortoli et al., 2022)), and difficulties in constructing simple and scalable training objectives.
C
this work, we introduce Riemannian Flow Matching (RFM), simple yet powerful methodology for learning continuous normalizing flows (CNFs; (Chen et al., 2018)) on general Ric- mannian manifolds M. RFM builds upon the Flow Matching framework (Lipman et al., 2023; A Abergoo & Vanden-Fijnden, 2023; Liu et al., 2023) and learns CNF by regressing an implicitly defined target vector field un (x) that pushes a base distribution p towards target distribution g defined by the train- ing examples. To address the intractability of uy (z)), we employ similar approach Conditional Flow Matching (Lipman et al., 2023), where we regress onto conditional vector fields (x)())) that push p towards individual training examples T1-
Geodesic
Biharmonic
Figure Our approach makes of user-specified premetrics general manifolds to define flows. On select simple manifolds, the geodesic can be computed exactly and leads simulation-free algo- rithm. On general manifolds where the geodesic is only computa- tionally expensive but can lead to degeneracy (e.g., along boundaries), we propose the use of spectral dis- tances (e.g., biharmonic), which can be computed efficiently contin- gent on one-time processing cost.
A key observation underlying our Riemannian generalization is that the conditional vector field necessary for training the CNF can be explicitly expressed in terms of "premetric" d(x, w), which distinguishes pairs of points x and on the manifold. A natural choice for such premetrie is the geodesic distance func- tion, which coincides with the straight trajectories previously used in Euclidean space by prior approaches.

---

Published as a conference paper at ICLR 2024
On simple geometries, where geodesics are known in closed form (e.g., Euclidean space, hyper- sphere, hyperbolic space, torus, or any of their product spaces), Riemannian Flow Matching remains completely iomulation-frre. Even general geometries, it only requires forward simulation of relatively simple ordinary differential equation (ODE), without differentiation through the solver, stochastic iterative sampling, or divergence estimation.
On all types of geometries, Rieman- nian Flow Matching offers several ad- vantages over recently proposed Rie- mannian diffusion models (De Bor- toli et al., 2022; Huang et al.. 2022). These advantages include avoiding it- erative simulation of noising process during training for geometries with an- alytic geodesic formulas; not relying on approximations of score functions or divergences of the parameteric vec- tor field; and not needing to solve stochastic differential equations (SDE) on manifolds, which is generally more challenging to approximate than ODE solutions (Kloeden al., 2002; Hairer et al., 2006; Hairer, 201 I). Table summarizes the key differences with relevant prior methods, which we expand on further in Section 4 (Related Work).
Simulation free cn simple goo.
Clesed-fonn farget vector field
Does not require divergence
Bon- Homo al. (2022)
De Bectoll er al. (2022) (DSM)
x
De Boctoli et al (2022) (ISM)
X
X
Huang al. (3022)
X
X
Riemannion FM (Ouns)
Table I: Comparison of closely related methods for training continuous- time generative models on Riemannian mani- folds. Additionally, we are the only among these works to consider and tackle general geometries.
Empirically, we find that Riemannian Flow Matching achieves state-of-the-art performance on manifold datasets across various settings, being on par or outperforming competitive baselines. We also demonstrate that our approach scales to higher dimensions without sacrificing performance, thanks to our scalable closed-form training objective. Moreover, we present the first successful training of continuous-time deep generative models on non-trivial geometries, including those imposed by discrete triangular meshes and manifolds with non-trivial boundaries that represent challenging constraints on maze-shaped manifolds.
2 PRELIMINARIES
Riemannian manifolds. This paper considers complete connected, smooth Riemannian manifolds M with metric 9 as basic domain over which the generative model is learned. Tangent space to M at E M is denoted M,,, and g defines an inner product over T M denoted (u, b).- u, v E T.M. TM = UzeM } x M is the tangent bundle that collects all the tangent planes of the manifold. U = {we} denotes the space of time dependent smooth vector fields (VFs) u, [0, 1] x M + TM, where (x) E T.M for all x E M: div,(u,, is the Riemannian divergence w.r.t. the spatial (x) argument. We will denote by dvol. the volume element over M, and integration of function f: . + R over M is denoted J fx))dvol.. For readers who are looking for a more comprehensive background on Riemannian manifolds, we recommend Gallot et al. (1990).
Probability paths and flows on manifolds. Probability densities over M are continuous non- negative functions : : R4, such that p(=)dvol, 1. The space of probability densities over M is marked P. probability path p: a curve in probability space P. : [0, 1] + P; such paths will be used as supervision signal for training our generative models. A flow is diffeommophism : M + . defined by integrating instantaneous deformations represented by a time-dependent vector field E U. Specifically, a time-dependent flow, the : .M- - MM, is defined by solving the following ordinary differential equation (ODE) on M over E [0, 1].
d [x(x)) (v)())),
vo(x)) x,,
and the final diffeomorphism is defined by setting V(x) = w/ (s). Given a probability density path Pt, it is said to be generated by un from p if th pushes Po = p to Pe for all i E [0, 1j. More formally,
f. (u.)
(x) log((
where the # symbol denotes the standard push-forwane operation and x, = ei. ( This formula can be derived from the Riemannian version of the instantaneous change of variables Formula (see
2

---

Published as a conference paper at ICLR 2024
equation 22 in Ben-Hamu al., 2022)). Previously, Chen et al. (2018) suggested modeling the flow 2h implicitly by considering parameterizing the vector field ug. This results in a deep generative model of the flow tt,, called a Continuous Normalizing Flow (CNF) which models probability path pr through a continuous-time deformation
base distribution p. A number of works have formulated manifold variants (Mathieu & Nickel, 2020; Lou al., 2020; Falorsi, 2020) that require simulation in order to enable training, while some simulation- free variants (Rozen et al., 2021; Ben-Hame et al., 2022) scale poorly high dimensions and do not readily adapt to general geometries.
3
METHOD
We aim to train generative model that lies on complete, connected smooth Riemannian manifold M endowed with metric g. Concretely, we are given set of training samples 21 E M from some unknown data distribution ((11). E P. Our goal is to learn parametric map t : .M- + MM that pushes a simple base distribution E P to q.
3.1 FLOW MATCHING ON MANIFOLDS
Flow Matching Lipman et al. (2023) is method to train Continuous Normalizing Flow (CNF) on Euclidean space that sidesteps likelihood computation during training and scales extremely well, similar diffusion models (Ho et al., 2020; Song et al., 2020b), while allowing the design of more general noise processes which enables this work. We provide brief summary and make the necessary adaptation to formulate Flow Matching on Riemannian manifolds. Derivations of the manifold case with full technical details are in Appendix A.
Riemannian Flow Matching, Flow Matching (FM) trains a CNF by fitting a vector field v E U. i.e., (x( E M, with parameters 0 IR," to an a priori defined target vector field U that is known to generate a probability density path E P over M satisfying Po and p1 = On manifold endowed with Riemannian metric g. the Flow Matching objective compares the tangent vectors (2),, (z) E TrM using the Riemannian metric g at that tangent space:
Lus px)) ) 1| (x) (x))²
where ([0,, 1], the uniform distribution over [0, 1].
Probability path construction. Riemannian Flow Matching therefore requires coming up with probability density path p. P. e [0. 1] that satisfies the boundary conditions
Pu =
P1 =
and corresponding vector field (VF) (x)) which generates pa(x)) from p in the sense of equation 2. One way to construct such a pair is to create per-sample conditional probability paths Pe (x(a)) satisfying
po(x|<1) = p(x), pi (x/ Re 8z, (x),
(5)
where 62, (I) is the Dirac distribution over M centered at ø1. One can then define pe(x) as the marginalization of these conditional probability paths over q(x1).
(z( - po(2(=1))(() ddvol..)
(6)
which satisfies equation 4 by construction. It was then proposed by Lipman et al. (2023)-whie we verify for the manifold setting- to define (x) as the "marginalization" of conditional vector fields un (x(xx1) that generates pr (xt (2) (in the sense detailed in Section 2).
ux(x)) = =// "*(3)/1)) -(mn((rr())) Pr ((((1)9((1)))) dvolays
pe(
which provably generates (ar).. However, directly plugging (x)) into equation is intractable as computing (z) is intractable.
Riemmanian Conditional Flow Matching. A key insight from Lipman al. (2023) is that when the targets Pt and u4 are defined as in equations and 7, the FM objective is equivalent to the following Conditional Flow Matching objective.
(amx(()) (r.,((z))) [lx((x) un (<|<x))1
(8)
as long as uy (z( (x[ (11) is a vector field that generates Pe (x(21) from p.
3

---

Published as a conference paper at ICLR 2024
To simplify this loss, consider the conditional flow, which we denote via the shorthand,
Algorithm Riemannian CFM
Require: base p. target scheduler k Initialize parameters 0 of
Zy (2
(9)
while not converged do
defined the solution to the ODE in equation with the VF ( and the initial condition (2o (m)) arg- Furthermore, since sampling from Pr (x| can be done with 4 (zol where wo ~- p(zo), we can reparametrize equation 8 as
sample time (0 1)
sample training example 21 q sample noise Xip ~ p
if simple geometry then
2, = expx, (M(t)log,, (ro)
else if general geometry then
CRCM (x,g(()))p((r))) (10r(2)) -iel? (10)
2₁ = solve_ Doolve_oDE(O,, t). To, (x(2=11)
end if
where i; 4/dt 2₁ ue (x1).
(() = -) lxll
0 = opt wwwiiimmitte(((ee
Riemannian Conditional Flow Matching (RCFM) has three requirements: parametric vector field 2t that outputs vectors on the tangent planes, the of the appropriate Riemannian metric l1-1ly- and the design of (computationally tractable) conditional flow ( 21) whose probability path satisfies the boundaries conditions in equation 5. We discuss this last point in the next section. Generally, compared to existing methods for training generative models on manifolds, RCFM both simple and highly scalable; the training procedure summarized in Algorithm and a detailed comparison can be found in Appendix D.
end while
3.2 CONSTRUCTING CONDITIONAL FLOWS THROOGII PREMETRICS
We discuss the construction of conditional flows the (2(21) on M that concentrate all mass at zı at time 1: Figure provides an illustration. This ensures that equation 5 will hold (regardless of the choice of p) since all points are mapped to 2 at time / = 1, namely
( = æ1, for all æ E M.
(11)
On general manifolds, directly constructing the that satisfies equation 11 can be overly cumbersome. Alternatively, we propose an approach based on designing a premetric instead, which has simple properties that, when satisfied, characterize conditional flows which satisfy equation 1. Specifically, we define premetric as d M x M - R satisfying:
1. Non-negative: d(x, v) 2 for all a, E M.
2. Positive: d(z, y) r = y.
3. Non-degenerate: dd / off
Figure 2: The conditional vector
field uy (x(2r) defined in equa- tion 13 transports all points x # æ1 to æ1 at exactly 1 =
We convention Vd(x. v_d(x, y). Such premetric denotes the closeness of a point x to x1, and we aim to design conditional flow th (2|2²) that monotonically decreases this premetric. That is, given a monotonically decreasing differentiable function x(t) satisfying x(0) = 1 and x(1 O, we want to find 1ús that decreases d(-, x1) according
d(th (2olaz)) .1) ((t)d(zo, æ1),
(12)
Here k(t) acts as scheduler that determines the rate at which d(-[a1) decreases. Note that at = 1. we necessarily satisfy equation 11 since .i is the unique solution to d(- (x1)) due the "positive" property of the premetric. Our next theorem shows that th (x(xx) satisfying equation 12 results in the following vector field,
d log k(t)
Vd(z, x1)
ue (x(xx)
d(x,))))
(13)
Vd(x, 2 e1?
2
The "non-degenerate" property guarantees this conditional vector field is defined everywhere * 21-
4

---

Published as a conference paper at ICLR 2024
Theorem 3.1. The flow (2((2) defined by the vector field (x(2/11) in equation 13 satisfies equation 12, and therefore also equation 11. Conversely, out of all conditional vector fields that satisfy equation 12, this u,(x)az (1) is the minimal norm solution.
A more concise statement and full proof of this result can be found in Appendix B. Here we provide proof for the first part: Consider the scalar function a(t) d(r,, æ1), where It th ( (x(xx) the flow defined with the VF in equation 13. Differentiation w.r.t time gives
_ a() = (Vd(ar, 21) it) = (Vd(z,, xi) u(), - dloog k(t) a((),
dt
The solution of this ODE is a(l) = =(1)d(xr, xi). which can be verified through substitution, and hence proves d x1) x(t)d(zz, 21). Intuitively, u (x(xa)1) is the minimal norm solution since does not contain orthogonal directions that do not decrease the premetric.
A simple choice make in this paper for the scheduler is x(t) 1 t, resulting in conditional flow that linearly decreases the premetric between 24 and 21. Using this, we arrive at a more explicit form of the RCFM objective,
vd(a,, 21)
Lucraa (0) = 1e
(e((ee) + d(ro: 21)
(14)
Vd )
For general manifolds M and premetrics d, training with Riemmanian CFM will require simulation in order to solve for It, though it does not need to differentiate through dt. However, on simple geometries RCFM can become completely simulation-free by choosing the premetric to be the geodesic distance, as we discuss next.
Geodesic distance. A natural choice for the premetric d(x, y) over a Riemannian manifold M is geodesic distance d, (ar, y). Firstly, we note that when using geodesic distance as our choice of premetric, the flow ( -inceince it is the minimal norm solution equivalent to the geodesic path, i.e., shortest path, connecting zμ and ø1.
Proposition 3.2. Consider a complete, connected smooth Riemannian manifold (M, g) with geodesic distance dg (2, y). In case d(x,, y) = dg (x, y) then 2, thr (zo (x1) defined by the conditional VF in equation 13 with the scheduler x(())
is a constant speed geodesic connecting Zo to TT.
This makes it easy compute geodesics, e.g., Euclidean space, the hypersphere, the hyperbolic space, the high- nigh-diimnsional torus, and some matrix Lie Groups. In particular, the geodesic connecting zo and æ1 can be expressed in terms of the exponential and logarithm maps,
simple manifolds, which define manifolds with closed-form
3e exp21 (ns(1) log,, (mm)), E [0, 1].
(15)
This formula can simply be plugged into equation 10, resulting in a highly scalable training objective. list of simple manifolds that we consider can be found in Table
Euclidean geometry. With Euclidean geometry M = R", and with standard Euclidean norm d(x, y/ = ||a 2| the conditional VF (equation 13) with scheduler s(()) = reduces to the VF used by Lipman et (2023), ue (a (x1) = 1-1 and the RCFM objective takes the form
-
Cacru (0) g (xxe) + 20 x1 11
which coincides with the Euclidean case of Flow Matching presented in prior works (Lipman et al., 2023: Liu et al., 2023).
3.3 SPECTRAL DISTANCES ON GENERAL GEOMETRIES
Geodesies can be difficult to compute efficiently for general geometries, especially since it needs to be computed for any possible pair of points. Hence, we propose using premetrics that can be computed quickly for any pair of points M contingent Riemannian manifolds, we consider the use of approximate spectral distances as an alternative to the geodesic distance. Spectral distances actually offer some benefits the geodesic distance such as
one-time upfront cost. In particular, for general
5

---

Published as a conference paper at ICLR 2024
Geodesic
Biharmonic
Diftusion
Diffusion==
Diffusion 7 TO
Figure 3: Contour plots of geodesic and spectral distances (to source point) on general manifolds. Geodesics expensive compute online and globally non-smooth. The biharmonic distance behaves smoothly while the diffusion distance requires careful tuning of the hyperparameter т.
robustness to topological noise, smoothness, and are globally geometry-aware (Lipman et al., 2010). Note however, that spectral distances do not define minimizing (geodesic) paths, and will require simulation of un (2(21) in order to compute conditional flows Iz.
Let Pp: M R be the eigenfunctions of the Laplace-Beltrami sponding cigenvalues A, i.e., they satisfy Agtp Aipi, for = 1,2,..., then spectral distances are of the form
Beltrami operator Ag over M with corre-
(=) = Σ
p.(x)))
(16)
where R - R4 is some monotonically decreasing weighting function. Popular instances of spectral distances include:
Diffusion Distance Coifman & Lafon (2006): w(A) exp(-27A), with parameter т.
2. Biharmonic Distance Lipman al. (2010): (AA) A-
In practice, truncate the infinite series in equation 16 to the smallest A: eigenvalues. These eigenfunctions can be numerically solved as a one-time preprocessing cost prior to training. Furthermore, we note that using an approximation of the spectral distance with finite k is sufficient for satisfying the properties of the premetric, leading to no bias in the training procedure. Lastly, consider manifolds with boundaries and show that solving eigenfunctions using the natural, or Neumann, boundary conditions ensures that the resulting (z(2(z1) does not leave the interior of the manifold. Detailed discussions on these points above can be found in Appendix G. Figure visualizes contour plots of these spectral distances for manifolds with non-trivial curvatures.
RELATED WORK
Deep generative models on Riemannian manifolds. Some initial works suggested constructing normalizing flows that map between manifolds and Euclidean spaces of the same intrinsic dimension (Gemici et al., 2016; Rezende et al., 2020; Bose et al., 2020), often relying on the tangent space some pre-specified origin. However, this approach is problematic when the manifold is not homeomorphic to Fuclidean space, resulting in both theoretical and numerical issues. On the other hand, continuous- models such as continuous normalizing flows bypass such topogical constraints and flow directly the manifold itself. To this end, number of works have formulated continuous normalizing flows on simple manifolds (Mathieu & Nickel, 2020; Lou et al., 2020; Falorsi, 2020), but procedure. More recently,
these rely on maximum likelihood for training, costly simulation-bast simulation- free training methods for continuous normalizing flows on manifolds have been proposed (Rozen al., 2021; Ben-Hame al., 2022); however, these scale poorly to high dimensions and do not adapt to general geometries.

---

Published as a conference paper at ICLR 2024
Table 2: Test NLL on Earth and climate science datasets. Standard deviation estimated over 5 runs.
Volcano 827
Earthquake 6120
Flood 4875
Fire 12809
Dataset size (train val + test)
CNF-bused
Riemannian CNF (Mathios & Nickel, 2020)
-6.05
0.
1.11=0.19
00.80
Moser Flow (Kozen et 2021)
#0.17
-0.16:40.06
=.57
-1.28
CNP Matching (Ben-Homa et sl. 2022)
=0.38
+0.08
0.25
=1,40
Riemannian Score-Based (De Bontll et sl, 2022)
-4.92=00.5
0.119
0.48=0.17
-1.33
ELBOO-based
Riemannian Diffusion Model (Huung al., 2022)
-0.61
+0.40
0.43-a0
-1.382005
Ours
Riemannian Flow Matching "/ Gendesic
-7.93
-0.28
0.42=005
-1.8620
Riemannian diffusion models. With the influx of diffusion models that allow efficient simulation- free training Euclidean space (Ho et al., 2020; Song et al., 2020b), multiple works have attempted to adopt diffusion models to manifolds (Mathieu & Nickel, 2020; Huang et al., 2022). However, due to the reliance on stochastic differential equations (SDE) and denoising score matching (Vincent, 2011), I), these approaches necessitate in-training simulation and approximations when applied to non-Euclidean manifolds.
First and foremost, they lose the simulation-free sampling of It ~ (=(x((1) that is offered in the Euclidean regime; this is because the manifold analog of the Ornstein-Uhlenbed SDE does not have closed-form solutions. Hence, diffusion-basee methods have to resort to simulated random walks as a noising process even on simple manifolds (De Bortoli et al., 2022: Huang al., 2022).
Furthermore, even simple manifolds, the conditional score function is not known analytically, so De Bortoli et al. (2022) proposed approximating the conditional score function with either an eigen- function expansion or Varhadan's 's heat-kernel approximation. These approximations lead to biased gradients in the denoising score matching framework. We find that the heat-kernel approximations can potentially be extremely biased, even with hundreds with eigenfunctions (see Figure 10). In contrast, we show in Figure 11 that our framework can satisfy all premetric requirements even with a small number of eigenfunctions for the spectral distance approximation hence guaranteeing that the optimal model distribution is the data distribution. See detailed discussions in Appendix G.I.
A way bypass the conditional score function is to use implicit score matching (Hyvärinen & Dayan, 2005), which Huang et al. (2022) adopts for the manifold case, but this instead requires divergence computation of the large neural nets during training. Using the Hutchinson estimator (Hutchinson, 1989; Skilling, 1989; Grathwohl et al., 2019; Song et al., 2020a) for divergence estimation results in more scalable algorithm, but the variance of the Hutchinson estimator scales poorly with dimension (Hutchinson, 1989) and is further exacerbated on non -un-Euclidean manifolds (Mathieu & Nickel, 2020).
Finally, the use of SDEs as noising process requires carefully constructing suitable reverse-timg time processes that approximate either just the probability path (Anderson, 1982) or the actual sample trajectories (Li al., 2020), whereas ODE solutions are generally well-defined in both forward and reverse directions (Murryy& & Miller, 2013).
In contrast these methods, Riemannian Flow Matching is simulation-free on simple geometries, has exact conditional vector fields, and does not require divergence computation during training. These properties are summarized Table 1, and a detailed comparison of algorithmic differences to iiffussionbbaed approaches is presented in Appendix D. Lastly, for general Riemannian manifolds we show that the design of a relatively simple premetric is sufficient, allowing the use of general distance approximate spectral distances with finite
functions that don't 't satisfy all axioms of a truncation- concation-going beyond what is currently possible with existing Riemannian diffusion methods.
Euclidean Flow Matching. Riemannian Flow Matching built on top of recent simulation-fit methods that work with ODEs instead of SDEs, regressing directly onto generating vector fields instead of score functions (Lipman et al., 2023; Albergo & Vanden- Eijjden, 2023; Liu et al., 2023; Neklyudov et al., 2022), resulting in arguably simpler approach to continuous-ting modeling without the intricacies of dealing with stochastic differential equations. In particular, Lipman et al. (2023) shows that this approach encompasses and broadens the probability paths used
generative
7

---

Published as a conference paper at ICLR 2024
Table Test NLL on protein datasets. Standard deviation estimated over runs.
General (2D)
Glycine (2D)
Proline (2D)
Pre-Pro (2D)
RNA (7D)
Dataset (train + + test)
138808
15333
1834
0910
9488
Mixture of Power Spherical (iwing 20E2)
1.15
2.08
0.27
1.34
4.081038
Ricmannian Diffusion Model (Hum at 20220
1.04+0012
1.970011
0.12=a011
1.24=0004
-3.70=0.90
Riemannian Flow Matching " Geodesic
1.01
1.90
0.15
1.18
-5.20=0.07
Eigenfunction
Model
Bunny (k=10)
Bunny (k=50)
Bunny (k=100)
Spot (k=10)
Spot (k=50)
Spot (k=100)
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.
by diffusion models while remaining simulation-free; Albergo & Handen-Eijnded (2023) discusses an interpretation based on the use of we also make explicit the construction of the marginal probability path pt (xx) and vector field un (x);; Liu et al. (2023) shows that repeatedly fitting to model's own samples leads to straighter trajectories; and Neklyudov et al. (2022) formulates an implicit objective when de (z) is gradient field.
to our conditional flows 4/e (x( except
5 EXPERIMENTS
We consider data from earth and climate science, protein structures, high- nigh-dimensional tori, compli- cated synthetic distributions on general closed manifolds, and distributions on maze-shaped manifolds that require navigation across non-trivial boundaries. Details regarding training setup is discussed in Appendix H. Due space constraints, additional experiments on hyperbolic manifold and manifold over matrices can be found in Appendix J, which are endowed with nontrivial Riemannian metrics. Table 5 provides all details of the simple manifolds and their geometries. Details regarding the more complex mesh manifolds can be found in the open source code, which we release for reproducibility!!
Earth and climate science datasets the sphere. We make use of the publicly sourced datasets (NOAA, 2020a;b; Brakenridge, 2017; EOSDIS, 2020) compiled by Mathieu & Nickel (2020). These data points lie on the 2-D sphere, simple manifold with closed form exponential and logarithm maps. We therefore stick to the geodesic distance and compute geodesics in closed form as in equation 15. Table 2 shows the results alongside prior methods. We achieve sizable improvement over prior works on the volcano and fire datasets which have highly concentrated regions that require a high fidelity. Figure 8 shows the density of our trained models.
Protein datasets on the torus. We make use of the preprocessed protein (Lovell et al., 2003) and RNA (Murray et al., 2003) datasets compiled by Huang al. (2022). These datasets represent torsion angles and can be represented on the 2D and 7D torus. We represent the data on flat torus, which is isometric to the product of 1-D spheres used by prior works (Huang al., 2022; De Bortoli et al., 2022) and result in densities that are directly comparable due to this isometry. Results are displayed in Table 3, and we show learned densities of the protein datasets in Figure Compared Huang et al. (2022), we see a significant gain in performance particularly on the higher dimensional 7D torus, due to the higher complexity of the dataset.
Scaling to high dimensions. We next consider the scalability of our method in the case of high-dimensional tori, following the exact setup in De Bottoli al. (2022). We com-
httpss: :/tps:////ii / http://cco//cccccoeeeeeeeeeeeeee htt//riemannnn an---a--
8

---

Published as a conference paper at ICLR 2024
pare to Moser Flow (Rozen et al., 2021), which does not scale well into high dimensions, and Riemannian Score-based (De Bortoli et al., 2022) using implicit score matching (ISM). As shown in Table I, this objective gets around the need to approximate conditional score functions, but it requires stochastic divergence estimation, introducing larger amounts of variance at higher
tog-lleellhood
dimensions. In Figure 5 we plot log-likelihood values, across these two baselines and our method with the geodesic construction. We see that our
Rlemannian Rlw Matching
Riemanniaa Score-based
method performs steadily, with no significant drop in performance at higher dimensions since we do not have any reliance on approximations.
Moser Flow
101
(Dimension)
Manifolds with non-trivial curvature. We next experiment with general closed manifolds using spectral distances described in Section 3.3. Specifically, we experiment on manifolds described by triangular meshes. For meshes, computing geodesic distances on-the-fly is too expensive for our use case, which requires hundreds of evalu- ations per training iteration. Fast approximations to the geodesic distance between two points are (n lo (Kimmel & Sethian, 1998), while exact geodesic distances require O( log n) (Surazh- sky et al., 2005), where n is the number of edges. On the other hand, computing spectral distances (k)), where & < i.e., it does not scale with the complexity of the manifold after the one-time preprocessing step. As our manifolds, we use the Standard Bunny (Turk & Levoy, 1994) and Spot the Cow (Crane al., 2013). Similar Rozen al. (2021), construct distributions by computing the k-th eigenfunction, thresholding, and then sampling proportionally to the eigenfunction. This done on high-resolution mesh that the distribution non-trivial on each triangle. Figure contains visualizations of the eigenfunctions, the learned density, and samples from a trained model that transports from uniform base distribution. We used k=200 eigenfunctions, which is sufficient for our method to produce high fidelity samples.
Figure 5: Riemannian Flow Matching scales incredibly well to higher dimensions it is simulation-free and all quantities required for training are computed exactly on simple geome- tries such as tori. log-likelihoods are in bits.
In Table 4, we report the test NLL of models trained using either the diffusion distance or the biharmonic distance. We had to carefully tune the diffusion distance hyperparameter 7 while the biharmonic distance was straightforward to use out-of-the-box and it has better smoothness properties (see Figure 3).
Stanford Bunny
Spot the Cow
K=10
k=50
h=100
h-10
I-50
Aw-300
Risminnia CFM
Diffusion (y-
1.48:am
1.534co:
0.87asan
0.55isae
L.M
Bibarmonic
1.55
1.49.um0 uo0
1.02
1.08
1.29. cas
Table 4: Test NLL on mesh datasets.
Manifolds with boundaries. Lastly, we ex- periment with manifolds that have boundaries. Specifically, we consider randomly generated mazes visualized in Figure 6. We set the base distribution to be Gaussian in the middle of the maze, and set the target distribution be a mixture of densities corners of the maze. These mazes are represented using triangular meshes, and we use the biharmonic distance using k= eigenfunctions. Once trained, the model represents a single vector field that transports all mass from the source distribution to the target distribution with no crossing paths. We plot sample trajectories in Figure 6 (b) and (d), where it can be seen that the learned vector field avoids boundaries of the manifold and successfully navigates to different modes in the target distribution.
SS
(a)
(b)
(c)
(d)
Figure 6: (a, c) Source (cyan) and target (yel- low) distributions on a manifold with non- trivial boundaries. (b, d) Sample trajectories from a CNF model trained through RCFM with the Biharmonic distance.
6 CONCLUSION
We propose Riemannian Flow Matching a highly-scalable approach for training continuous normalizing flows on manifolds. Our method is completely simulation-fre and introduces zero approximation errors on simple geometries that have closed-form geodesics. We also introduce benchmark problems for general manifolds and showcase for the first time, tractable training on general geometries including both closed manifolds and manifolds with boundaries.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
Ricky T. Q. Chen would like to thank Chin-Wei Huang for helpful discussions. Additionally, we acknowledge the Python community (Van Rossum Drake Jr. 1995; Oliphant, 2007) for developing the core set of tools that enabled this work, including PyTorch (Paszke et al., 2019), Py PyTorch Lightning (Falcon & team, 2019), Hydra (Yadan, 2019), Jupyter (Kluyver et al., 2016), Matplotlib (Hunter, 2007), seaborn (Waskom et al., 2018), numpy (Oliphant, 2006; Van Der Walt et al., 201 1), SciPy (Jones et al., 2014) pandas (McKinney, 2012), geopandas (Jordahl et al., 2020), torchdiffeq (Chen, 2018), libigl (Panozzo & Jacobson, 2014), and PyEVTK (Herrera, 2019).
REFERENCES
Michael Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. International Conference Learning Representations, 2023.
Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their
Applications, 12(3):313-326 1982.
Alexandre Barachant, Stéphane Bonnet, Marco Congedo, and Christian Jutten. Classification of covariance matrices using riemannian-base kernel for bci applications. Neurocomputiing, 1 112: 172-178, 2013.
Mikhail Belkin and Partha Niyogi. Laplacian eigenmaps for dimensionality reduction and data representation. Neural computation, 1376)113331136, 2003.
Heli Ben- Ben-Hamu, Samuel Cohen, Joey Bose, Brandon Amos, Aditya Grover, Maximilian Nickel, Ricky T. Q. Chen, and Yaron Lipman. Matching normalizing flows and probability paths on manifolds. International Conference on Machine Learning, 2022.
Benjamin Blankertz, Guido Dornhege, Matthias Krauledat, Robur- Müller, and Gabriel Curio. The non-invasive berlin brain computer interface: fast acquisition of effective performance in untrained subjects. Neurolmage, 37(2):539-550) 2007.
Joey Bose, Ariella Smofsky, Renjie Liao, Prakash Panangaden, and Will Hamilton. Latent variable modelling with hyperbolic normalizing flows. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume I 119 of Proceedings of Machine Learning Research, pp. 1045-1055. PMLR, 13- 18 Jul 2020.
G.R. Brakenridge. Global active archive of large flood events. http: htt////looooee/e///ooo rvatory. colorado. www.eeeww...eeeeeeeeeeeee index..///nnee/eeeeeeeeee html, 2017. Dartmouth Flood Observatory, University of Colorado..
Clemens Brunner, Robert Leeb, Gernot Müller-Putz, Alois Schlögl, and Gert Pfurtscheller. Bei competition 2008-graz data set a. Institute for Knowledge Discovery (Laboratory of Brain- Computer Interfaces). Graz University of Technology, 16:1-6, 2008.
Ricky T. Q. Chen. torchdiffeg, 2018. torchdiffeq.
URL https: httpttuuuu..cmmtttigii..iecceeee
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances neural information processing systems, 31, 2018.
Ronald R Coifman and Stéphane Lafon. Diffusion maps. Applied and computational harmonic analysis, 21(1):5-30, 2006.
Keenan Crane, Ulrich Pinkall, and Peter Schröder. Robust fairing via conformal curvature flow. ACM Transactions on Graphics (TOG), 3204):1-10, 2013.
Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thornoon, Yee Whye Teh, and Amaud Doucet. Riemannian score-based generative modeling. Advances in Neural Information Processing Systems, 2022.
10

---

Published as a conference paper at ICLR 2024
Zhijie Deng, Jiaxin Shi, Hao Zhang, Peng Cui, Cewu Lu, and Jun Zhu. Neural eigenfunctions structured representation learners. arXiv preprint arXiv:2210.12633 2022.
EOSDIS.
Active fire data.
httpss//earthddteenaa...gg
worth-observatior-eo wewaattaeearrel-ttieeeeee het-reeelttmmeeiiessseeeee-- wttpe//eiiiaaaatttteeeeeeeeeeeee
2020. Land, Atmosphere Near real-time Capability for EOS (LANCE) system operated by NASA's Earth Science Data and Information System (ESDIS).
William Falcon and The PyTorch Lightning team.
Pytorch lightning, 2019. URL https://
github. www.oo////oooo. Lighhnigghtiiggnnnnnnngi A/Iiighhiingg
Luca Falorsi. Continuous normalizing flows on manifolds. PhD thesis, University of Amsterdam, 2020.
Sylvestre Gallot, Dominique Hulin, and Jacques Lafontaine. Riemannian geometry, volume 2. Springer, 1990.
Mevlana C Gemici, Danilo Rezende, and Shakir Mohamed. Normalizing flows on riemannian manifolds. arXiv preprint arXiv:1611.02304, 2016.
Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. FFJORD:
Free-form continuous dynamics for scalable reversible generative models. International Conference on Learning Representations, 2019.
Ernst Hairer. Solving differential equations on manifolds. Lecture Notes, Université de Geneve, 2011.
Emst Hairer, Marlis Hochbruck, Arieh Iserles, and Christian Lubich. Geometric numerical integration. Oberwolfach Reports, 3(1):805-882, 2006.
Paulo Herrera. Pyevtk, 2019. URL https:ttt/////
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in
Neural Information Processing Systems, 33:6840-6851, 2020.
Chin Wei Huang, Milad Aghajohari, Avishek Joey Bose, Prakash Panangaden, and Aaron Courville. Riemannian diffusion models. Advances in Neural Information Processing Systems, 2022.
John Hunter. Matplotlib: A 2d graphies environment. Computing in science engineering, 9(3): 90, 2007.
M.F. Hutchinson. A stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines. 18: 11059-1776,,,0 01 1989.
Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
Eric Jones, Travis Oliphant, and Pearu Peterson. {SciPy): Open source scientific tools for (Python). 2014.
Peter Jones, Mauro Maggioni, and Raanan Schul. Manifold parametrizations by cigenfunctions of the laplacian and heat kernels. Proceedings of the National Academy of Sciences, 105(6): 1803-1808, 2008
Kelsey Jordahl, Joris Van den Bossche, Martin Fleischmann. Jacob Wasserman, James McBride, Jeffrey Gerard, Jeff Tratner, Matthew Perry, Adrian Garcia Badaracco, Carson Farmer, Geir Ame Hjelle, Alan D. Snow, Micah Cochran, Sean Gillies, Lucas Culbertson, Matt Bartos, Nick Fu- bank, maxalbert, Aleksey Bilogur, Sergio Rey, Christopher Ren, Dani Arribas-Bal Bel, Leah Wasser, Levi John Wolf. Martin Journois, Joshua Wilson, Adam Greenhall, Chris Holdgraf, Filipe, and François Leblanc. geopandas/geoopandas: v0.8.1, 2020.
Ron Kimmel and James A Sethian. Computing geodesic paths on manifolds. Proceedings of the national academy of Sciences, 95(15):8431-8435 1998.
Peter Eris Kloeden, Eckhard Platen, and Henri Schurz. Springer Science & Business Media, 2002.
11

---

Published as a conference paper at ICLR 2024
Thomas Kluyver, Benjamin Ragan-Kelley, Fernando Pérez, Brian E Granger, Matthias Bussonnier, Jonathan Frederic, Kyle Kelley, Jessica B Hamrick, Jason Grout, Sylvain Corlay, et al. Jupyter notebooks-a publishing format for reproducible computational workflows. In ELPUB, pp. 87-90, 2016.
R Leeb, C Brunner, G Müller-Putzz A Schlögl, and GJGUOT Pfurtscheller. Bei competition 2008-2000 -gaa data set b. Graz University of Technology, Austria, pp. 1-6, 2008.
Xuechen Li, Ting-Kan Kam Leonard Wong, Ricky Q. Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. In International Conference on Artificial Intelligence and Statistics, pp. 2870-3884 PMLR, 2020.
Yaron Lipman, Raif M Rustamov, and Thomas A Funkhouser. Biharmonic distance. ACM Transac- tions on Graphics (TOG), 29(3)1-111, 2010. 2010.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. International Conference on Learning Representations, 2023.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning generate and
transfer data with rectified flow. International Conference on Learning Representations, 2023.
Aaron Lou, Derek Lim, Isay Katsman, Leo Huang, Qingxuan Jiang, Ser Nam Lim, and Christopher M De Sa. Neural manifold ordinary differential equations. Advances Neural Information Processing Systems, 15:17548-17555 2020.
Simon C Ian W Davis, W Bryan Arendall III, Paul IW De Bakker, J Michael Word, Michael G Prisant, Jane S Richardson, and David C Richardson. Structure validation by ca geometry: é. v and cB deviation. Proteins: Structure, Function, and Bioinformatics, 50(3):437-450 2003.
Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. Advances in Neural Information Processing Systems, 33:2503-2515, 2020.
Robert J McCann. Polar factorization of maps on riemannian manifolds. Geometric & Functional Analysis GAFA, 11(3):589-600) 2001.
Wes McKinney. Python for data analysis: Data wrangling with Pandas, NumPy, and IPython. O' O'Reilly Media, Inc.", 2012.
Maher Moakher and Philipp G Batchelor. Symmetric positive-defintte matrices: From geometry applications and visualization. Visualization and processing of tensor fields, pp. 285-298, 2006.
Francis J Murray and Kenneth S Miller. Existence theorems for ordinary differential equations. Courier Corporation, 2013.
Laura JW Murray, W Bryan Arendall III, David Richardson, and Jane S Richardson. Rna backbone is rotameric. Proceedings of the National Academy of Sciences, 10012)):39044133000 2003.
Kirill Neklyudov, Daniel Severo, and Alireza Makhzani. Action matching: A variational method for learning stochastic dynamics from samples. arXiv preprint arXiv:2210.066600 2022.
NOAA. Global significant earthquake database. htwpt:tt//t/tcttt //att wwd.ooo....... noaa..noaa... www.eeonaaao....... www.dd..ooa...ooo...oooo..o
iso?id=gov woaa..a....mmooooooo ngde. mgg. hazards :tdddd:::ll111: 2020a. National Geophysical Data Center / World Data Service (NGDC/WDS): NCEI/WDS Global Significant Earthquake Database.
NOAA National Centers for Environmental Information.
NOAA. Global significant volcanic eruptions database. https:
nodc. nowa. gov/
egi-bin/iie issossso0 iioooiid-aae id-gov. noaa .nddc... gogdc.mgg..o hazards: G1014777 2020b. National Geo- physical Data Center / World Data Service (NGDC/WDS): NCEI/WDS Global Significant Volcanic Eruptions Database. NOAA National Centers for Environmental Information.
Travis Oliphant. A guide to NumPy, volume 1. Trelgol Publishing USA, 2006.
Travis E Oliphant. Python for scientific computing. Computing in Science & Engineering, 9(3): 10-20, 2007.
12

---

Published as a conference paper at ICLR 2024
Daniele Panozzo and Alec Jacobson. Libigl: A A++ library for geometry processing without mesh data structure. 2014.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-perfornance deep learning library. In Advances in neural information processing systems, pp. 8026-8037, 2019.
David Pfau, Stig Petersen, Ashish Agarwal, David GT Barrett, and Kimberly Stachenfeld. Spectral inference networks: Unifying deep and spectral learning, arXiv preprint arXiv: 02001, 2018.
Boris T. Polyak and Anatoli Juditsky. Acceleration of stochastic approximation by averaging. 1992.
Prajit Ramachandran, Barret Zoph, and Quoe V Le. Searching for activation functions. arXiv preprint arXiv:1710.05941, 2017.
Danilo Jimenez Rezende, George Papamakarios, Sébastien Racaniere, Michael Albergo, Gurtej Kan- war, Phiala Shanahan, and Kyle Cranmer. Normalizing flows on tori and spheres. International Conference on Machine Learning, pp. 8083-8092. PMLR, 2020.
Noam Rozen, Aditya Grover, Maximilian Nickel, and Yaron Lipman. Moser flow: Divergence- based generative modeling on manifolds. Advances in Neural Information Processing Systems, 34: 17669-17680, 2021.
John Skilling. The eigenvalues of mega-dimensional matrices. In Maximun Entropy and Bayesian Methods, pp. 455-466. Springer, 1989.
Yang Song. Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: A scalable approach to density and score estimation. Uncertainty in Artificial Intelligence, pp. 574-584. PMLR, 2020a.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:201 1.3456, 2020b.
Vitaly Surazhsky, Tatiana Surazhsky, Danil Kirsanov, Steven J Gortler, and Hugues Hoppe. Fast exact and approximate geodesics on meshes. ACM transactions on graphics (TOG), 24(3):553-560, 2005.
Greg Turk and Mare Levoy. Zippered polygon meshes from range images. In Proceedings of the 2 annual conference Computer graphics and interactive techniques, pp. 311-318, 1994.
Stefan Van Der Walt, S Chris Colbert, and Gael Varoquaux. The numpy array: structure for efficient numerical computation. Computing in Science & Engineering, 13(2):22, 2011.
Guido Van Rossum and Fred L Drake Jr. Python reference manual. Centrum voor Wiskunde en Informatica Amsterdam, 1995.
Cédric Villani. Optimal transport: old and new, volume 338. Springer, 2009.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computa- fion, 23(7):1661-1674, 2011.
Michael Waskom, Olga Botvinnik, Drew O'Kane, Paul Hobson, Joel Ostblom, Saulius Lukauskas, David C Gemperline, Tom Augspurger, Yaroslav Halchenko, John B. Cole, Jordi Warmenhoven, Julian de Ruiter, Cameron Pye, Stephan Hoyer, Jake Vanderplas, Santi Villalba, Gero Kunter, Eric Quintero, Pete Bachant, Marcel Martin, Kyle Meyer, Alistair Miles, Yoav Ram. Thomas Brunner, Tal Yarkoni, Mike Lee Williams, Constantine Evans, Clark Fitzgerald, Brian, and Adel Qalieh. mwaskom/seabom: v0.9.0 (july 2018), July 2018. URL attpps///don http///oo.oo///////////
1313201.
Omry Yadan. Hydra a framework for elegantly configuring complex applications. Github, 2019.
URL https:tp://// / thub.
13

---

Published as a conference paper at ICLR 2024
A CONDITIONAL FLOW MATCHING ON MANIFOLDS
We provide the necessary derivations and proofs for the Conditional Flow Matching over a Riemannian manifolds; the proofs and derivations from Lipman et al. (2023) are followed "as-is", with the necessary adaptation to the Riemannian setting.
Assumptions. We will notations and setup from Section 2. Let p(-|x1) [0, P be a (conditional) probability path sufficiently smooth with integrable derivatives, strictly positive Pe(z(z1) and po(z(z1) p. where p is our source density. Let u(-[æı) be (conditional) time-dependent vector field, sufficiently smooth with integrable derivatives and such that
J. (1/10r Pa(x) )dvol, < co.
Further assume ur (x(22) generates pr (x((x1)) from p in the sense equation 2, i.e., we denote by ()((1(11) the solution to the ODE (equation l):
vy(x m1) ([(W((r)))))ln)) dt
(17)
vo(x (((x)a))= =x
(18)
then
pe(-(an) = [whe-]aan (+)1)))))
(19)
Proof of the marginal VF formula, equation 7. First, the Mass Conservation Formula Theorem (see, e.g., Villani (2009)) implies that po(r))r) and ux(a)ar))))))))))))))))(())))( satisfy
dt Pr (2|21) + (=))))
(20)
where div, is the Riemannian divergence with metric g.
Next, we differentiate the marginal (x( w.n.t.t t: d SmL (((()))))(() di
d di D(())
-div, [/
ddvol,,
=diivg (p((o) / ue(at (1)())))((()) pr(x)
dvol.
= -divv (x)(
where in the first and second equalities changed the order of differentiation and integration, and in the second equality we used the mass conservation formula for u, (x(zz) In the previous to last equality multiplied and divided by pe(x). In the last equality defined the marginal vector field W, as in equation 7.
RCFM loss equivalent to RFM loss. We will now show the equivalence of the RCFM loss (equation 8) and the RFM loss (equation 3). First note the losses expand as follows:
Lara(() Ex.p)()) [[ve(x) ux(x))1 Ex,p,()) (x)( 2 (cx(x), un (x)) + [1)x()))];
Cacra (8) E 1.(())) [[ve(x) u,(0|11))) MI3 (a(a)a)
E (g((2) [|0(())]1 2 (v.(x), (((===))) [u(()/21)1) pe(0((1)
Second, note that
ge (x((ar) ||vel1, L'' //m (et dt
= f'' / (1
= -E |loell
14

---

Published as a conference paper at ICLR 2024
Lastly,
Le J. (0a(1), u((k))))) (1(221 (dvol)(())))dddooooo
(o (or(x), (z(2221))
ddvol, dt
!'= I. (mnla)s / we(alan (e (0( Javola
dvol, dl
1'" Im (mnlo)) Iox un(xlx)) pr (r dvol,,
p p((
Pe(x) dt
I'' I. (vo(2), (x))),, pe (c dt (a)) (ve(2)), un (x)))
We got that CRCTM( and (0) differ by constant,
(e)))v) dt - L.¹ S. [ww. dt const S.'
that does not depend on 0.
B PROOF OF THEOREM 3.1
Theorem 3.1. The flow ()(())) defined by the vector field ux(x(z1) in equation 13 satisfies equation 12, and therefore also equation 11. Conversely, out of all conditional flows thx (x(21) defined by a vector fields ū(x()1) that satisfy equation 12, this ue (ar (r)) is of minimal norm.
Proof. Let 1 be the flow defined by equation 13 in the sense of equation 1. Differentiat- ing the time-dependent function d( x1) w.r.t time gives
dt d( 21) (Vd(z,, (V) #1) (Vd(z,, 21). us (x)(aa))) dloog x(() a()),
(21)
dt
This shows that the function a(t) = d(z;, .1) satisfies the ODE
d ((1) d log x(t) a(t):
dt
dt
with the initial condition a(0) d(x, ((,x) 21). General solutions to this ODE are of the form a(t) cx(t), where e > is a constant set by the initial conditions. This can be verified by substitution. The constant is set by the initial condition,
d(r, x1) = (0)) = cx((0 = c.
This gives att) = d(æ, x1 (x)) as the solution. Due to uniqueness of ODE solutions we get that equation 12 holds.
Conversely, consider (x((rr) satisfying equation 12. Differentiating both sides of this equation w.r.t. i and then using equation again we get
drs(t) d(x,,1)) dx(t) d( #1) di x(()
dllog k<()) d d(x,,
(1),
di
dt
let ue (x((1) denote the VF defining the diffeomorphism (x(())1) in the sense of equation then the last equation takes the form
(Vd(z,, 21) (=(2)(11)) dlog k(t)
d(z,, x1).
(22)
This equation provides under-determined linear system for u, (v(æ1) ) M at every point y with zon-zero probability, which in our case is all y E M as we assume Pe (y((x) > for all E M. As can be seen in equation 21, (x(x))1) defined in equation 13 is also satisfying this equation. Further note, that since u, (x (2)) defined in equation 13 satisfies u, (xxxx) || Vd(x, 21) (proportional) is the minimal norm solution to the linear system in equation 22.
15

---

Published as a conference paper at ICLR 2024
C PROOF OF PROPOSITION 3.2
Proposition 3.2. Consider a complete, connected smooth Riemannian manifold (M,g) with geodesic distance (2, y). case d(x, y) (z y) then It (zoo(z))) defined by the conditional VF in equation 13 with the scheduler x<(t) - t is geodesic connecting xo to æ1.
Proof. First, note that by definition (1 x1)) = z0. and (1 = ø1. Second, from Proposition 6 in McCann (2001) we have that
2 d - log
where log is the Riemannian logarithm map. From the chain rule we have
1 V # V (x, y)² = dg (z, v) (z,
Since the logarithm map satisfies |log, (w)l dg (2r, y) we have that
(z, u)1l, =1
(23)
Now, computing the length of the curve we get
L.
|bbell dt f [x<( dt
L.¹
do (2, x1) (xr,
dt
1-t
Vdg 2: )1
= -1(( 1'
Vd vd x11
dt
| (Tts 21 11
{( dt
=d 21)
dg(ro. 21)
where in the second equality we used the definition of the conditional VF (equation 13) with k(t) = 1 - t. in the third equality we used Theorem 3.1 and equation 12, and in the fourth equality we used equation 23, Since I, realizes minimum of the length function, a geodesic.
16

---

Published as a conference paper at ICLR 2024
D ALGORITHMIC COMPARISON TO RIEMANNIAN DIFFUSION MODELS
Algorithm 2 Riemannian Diffusion Models
Algorithm 3 Riemannian Flow Matching
Require:
base distribution p(ær), target ((ro)
Require:
base distribution p(zo). target ((x1)
Initialize parameters of st
Initialize parameters of vb
while not converged do sample time / 20, 7)
while not converged do sample time f ~ (4(0, 1)
sample training example xo g(30)
sample training example 21 (
sample noise To ~ p(zo)
if simple geometry then
% simulate Geometric Random Walk
It exp2, (#(t)1oq,, (x1))
= solve Solve_sDE((O) (DE( ro)
else if general geometry then
St sslle_ODE((0) ( zo, ur (x)
end if
if denoising score matching then
% approximate conditional score
eig-expansion
log (x) x
Varhadan
{()) = (xx,; 0) V log pe ( 2r))
% closed-fone regression target us ((relas)
else if implicit score matching then
{(0) [ 0) 2e (x)
% estimate Riemmanian divergence
sample E ~ [0,
div,s: R= T 8et + s llog det g(m+)
dz,
das
{(0) = 2 [0x; 0)1| + div
end if
= 1==ppimizer_step((/ end while
= end while
Figure 7: Algorithmic comparison between Riemannian Diffusion Models (De Bortoli al., 2022; Huang et al., 2022) and our Riemannian Flow Matching. Note time is reversed between these formulations. In red, we denote expensive computational aspects (sequential simulation during training), biased approximations (for the score function), and stochastic estimation (for divergence) that may not scale well. Also note that Geometric Random Walk does not converge to the stationary prior distribution unless simulated for an infinite amount of time, in practice requiring tuning T as a hyperparameter depending on the manifold. On simple manifolds, Riemannian Flow Matching bypasses all computational inconveniences and in particular is completely simulation-free
17

---

Published as a conference paper at ICLR 2024
E LIMITATIONS
As can be seen from Figure 7, our method still requires simulation of 2t on general manifolds. This sequential process can be time consuming, and more parallel or simulation-free approach to constructing would be more favorable. Furthermore, the spectral distances require eigenfunction solvers which may be computationally expensive on complex manifolds. Using approximate methods such neural eigenfunctions (Pfau et al., 2018; Deng et al., 2022) may be a possibility. One major advantage of our premetrie formulation that these eigenfunctions need not be perfectly solved in order to satisfy the relatively simple properties of our premetric.
F ADDITIONAL FIGURES
Volcano
Earthquake
Lattbude
Latibude
-56
Longitude
Longitude
Flood
Fire
Lastuse
Lastudo
Longitude
Longitude
Figure 8: Data samples and Ramachandran plots depicting log likelihood for protein datasets.
Soodi
General
Glycine
Proline
Pre-Pro
Figure 9: Data samples and Ramachandran plots depicting log likelihood for protein datasets.
G ADDITIONAL DISCUSSION
G.1 ON THE USE APPROXIMATE SPECTRAL DISTANCES THE PREMETRIC
Computation cost. The smallest & eigenvalues and their eigenfunctions need only be computed once as a pre-pree pre-processing step. On manifolds represented as discrete triangular meshes, this step can be done in matter of seconds. Afterwards, spectral distances can be computed very efficiently for all pairs of points. Note however, that training with RCFM does still require simulating for Pt, but as the vector fields (equation 13) do not contain neural networks, the flows can be solved efficiently in practice. This results in a similar cost to diffusion-base methods (Huang et al., 2022; De Bortoli et al., 2022) that require simulation even for simple manifolds.
18

---

Published as a conference paper at ICLR 2024
4.734
3.887
4.251
2.717
1.938
0.095
0.01
2.436
2.243
1.875
0.935
0.05
1.683
1.253
0.729
0.120
0,000
0.000
k: = 8
k 15
k k= 35
k 120
= 440
Figure 10: A visualization of the approximated heat kernel in equation 24 on 2D sphere for different values. Above each visualization we show the relative error of the conditional score function: (x,(2,)z)) [|] V log pr (z)) Vz, log (:(2)/20)) //VV.. log (zx where is located at the top of the sphere. We see that especially at small t values (near data distribution), there is significant error in the score function approximation even when using hundreds of eigenfunctions.
k=3
k=
k 15
35
k 120
440
Figure 11: Spectral distances (biharmonic), i.e., d(x,,) :) where is at the top of the sphere, visualized using isolines. We see that is very easy to satisfy the properties of a premetric with very few eigenfunctions, importantly the Positive property: d(x, v) 0 iff y. which allows to concentrate probability perfectly at every point of the manifold, and the Non-degenerate property Vd(x, y) f 0 iff + y (satisfied here everywhere except the antipodal point) that allows to properly define the flow in equation 13 almost everywhere.
Sufficiency with finite k. One may wonder we pay any approximation costs when using finite the answer is no. In fact, we only need as many eigenfunctions as it takes to be able to distinguish (almost) every pair of points on the manifold. Put differently, we don't need compute the spectral distances perfectly, only sufficiently enough that the conditions of our premetric are satisfied. Regarding the question what is enough? This is only understood partially: for local neighborhoods the number of required eigenfunctions is the manifold dimension (but not necessarily the first ones), a property proven in (Jones et al., 2008, Theorem 2). Nevertheless, the use of spectral distances computed with k smallest eigenvalues is equivalent to computing Euclidean distances in k-dimensional Euclidean embedding using the same eigenfunctions; this embedding is known to preserve neighborhoods optimally (Belkin & Niyogi, 2003).
19

---

Published as a conference paper at ICLR 2024
As comparison, Riemannian score-based generative models (De Bortoli et al., 2022) suggests heat kernel approximation
pe(2, = Σ Aut (z))))(((z))
(24)
1=0
resulting in the approximated conditional score function
Vx, log (x) Vz, log e Au (((o)))(a)))
(25)
However, equation 24 is only correct as k + oo. This is manifested in practice as large amounts of error even when using hundreds of eigenfunctions. In Figure 10, visualize the heat kernel approximation in equation 24 as well as relative error in the score function, taken in expectation w.r.t. po(r))))) i.e., the relative error is weighted higher in regions where we will evaluate the conditional score function during training. We see that at small time values close to the data distribution, the conditional score function may not even have the first significant digit correct (in expectation).
In contrast, we visualize the spectral distances using the biharmonic formulation in Figure 11, where we see that we already have an extremely accurate premetric with very few eigenfunctions. In particular, the required properties of premetric are satisfied even for just h: 3, roughly corresponding to the manifold dimension. Higher values of it simply refine the spectral distance but are not necessary.
G.2 MANIFOLDS WITH BOUNDARY
In considering general geometries, we also consider the case where M has a boundary, denoted Q.. In this case, need to add another condition to our premetric to make sure ur (x|21) will not flow particles outside the manifold. Let (z) E T&M denote the interior cittioo-poining normal direction at a boundary point
aM. We add the following condition to our premetric:
4. Boundary: (Vd(z, v), m(x)), S 0, Vy E M, E aM.
If the premetric satisfies this condition, then the conditional VF in equation 13 satisfies
(me(<(<1), n(x)), 2 0
implying that the conditional vector field does not point outwards on the boundary of the manifold.
Spectral distances at boundary points. In case A has boundary we want to make sure the spectral distances in equation 16 satisfy the boundary condition. To ensure this, we can simply solve eigenfunctions yos using the natural, or Neumann, boundary conditions, i.e., their normal derivative at boundary points vanish, and we have (V)))()),,n(x)) 0 for all æ E θ.M. This property implies that V ddw (z, y)², n(x))) O, satisfying the boundary condition of the premetric.
H EXPERIMENT DETAILS
Training setup. All experiments are run a single NVIDIA V100 GPU with 32GB memory. We tried our best to keep to the same training setup as prior works (Mathieu & Nickel, 2020; De Bortoli al., 2022; Huang et al., 2022); however, as their exact data splits were not available, we used our own random splits. We followed their procedure and split the data according to 80% train, 10% val, and 10% test. We used seeds values of 0-4 for our five runs. We used the validation stopping based on the validation NLL., and then only computed the test NLL using the checkpoint that achieved the best validation NI.L. We used standard multilayer perceptron for parameterizing vector fields where time is concatenated as input to the neural network. We generally used 512 hidden units and tuned the number of layers for each type of experiment, ranging from 6 to 12 layers. We used the Swish activation function (Ramachandran et al., 2017) with learnable parameter. We used Adam with a learning rate of le-4 and an exponential moving averaging on the weights (Polyak & Juditsky, 1992) with decay of 0.999 for all of our experiments.
for early
High-dimensional tori. We use the same setup as De Bortoli et al. (2022). The data distribution is a wrapped Gaussian the high-dimensional tori with uniformly sampled mean and scale of 0.2.
20

---

Published as a conference paper at ICLR 2024
Table Riemannian manifolds with known geodesics that we use in our experiments. The operator @ denotes Möbius addition. Exponential maps, logarithm maps, and inner products are used during training with Riemannian Conditional Flow Matching. Note x, M, u, v E M, and |1-11; {t: g) The last column log [x)) denotes the logarithm of the absolute value of the determinant of the metric tensor at x E M: this is used during log- ongllikelhood computation (see equations 37 and 33).
Mandted M
(spp(()
(6 s/s
Inglo(e)]
soow([o))s + rip di (1(1)))))))
aroncw (ll plottiiicc
(a.e)
N-D 2.2x)
(2x)
arcean (siula alveoddy
(a.e)
N-D Hypebolic (x 8M Iril <
ze (1outh (1-ii))) rit)
(1- [4f)) panks (I "1-x i oll) prr
T- (4.5)
Ning -
N N SPD matrices
X xx0(X-5000-000
xt xcllexX--xx--- xxx-xxxxxxxxxxxxxx
o (X-1xX-)1))
2i (og + (N + 1y X
We use a MLP with hidden layers of size 512 to parameterize Uy, train for 50000 iterations with a batch size of 512. We then report the log- likelihood per dimension (in bits) on 20000 newly sampled data points.
Triangular meshes. We use the exact open source mesh for Spot the Cow. For the Stanford Bunny, we downsample the mesh to 5000 triangles and work with this downsampled mesh. For constructing target distributions, we compute the eigenfunctions on a 3-times upsampled version of the mesh threshold the k-th eigenfunction at zero, then normalized to construct the target distribution. This target distribution uniform on each upsampled triangle, and further weighted by the area of each triangle. On the actual mesh we work with, this creates complex non-uniform distributions on each triangle. We also normalize the mesh so that points always lie in the range of (-1, 1).
Maze manifolds. We represent each maze manifold using triangular meshes in 2D. Each cell is represented using a mesh of 8x8 squares, with each square represented as two triangles. two neighboring cells connected, then we connect them using either 2x8 or 8x2 squares; if two neighboring cells are not connected (i.e., there is a wall), then we simply do not connect their meshes, resulting in boundaries where all the triangles are the same size, We randomly create maze structures based on breadth- first-search algorithm, represent these using meshes, and we then normalize the mesh that points lie in the range of (0, 1).
the manifold. This produces manifolds represented by triangular meshes
Vector field parameterization. We parameterize vector fields as neural networks in the ambient space and project onto the tangent space at every x. That is, similarly to (Rozen et al., 2021) we model
(z) = g(z) 1 Pr. um) Tx)))
(26)
where π is the projection operator onto the manifold, i.e.,
T(z) = arg min æ ylly,
(27)
yEM
and Py is the orthogonal projection onto the tangent space at y.
We also normalize the vector field using g(x) , which cancels out the effect of g on the Riemannian norm and makes standard neural network parameterization more robust changes in the metric, iee.
1 = (g 10) و(6
N -1v0 = vTv --
= =1u112-
(28)
We found this bypasses the need to construct manifold-specific initialization schemes for our neural networks and leads to more stable training.
ligeliiel computation. We solve for the log-density log Pı (x), for an arbitrary test sample x E M, by using the instantancous change of variables (Chen et al., 2018), namely solve the ODE
d
( 21 f((z))
ddivv (- 2e (xx) (vi) (xx))
(29)
di
where div, is the Riemannian divergence. We solve backwards in time from / time = 0 with the initial conditions
( 21 (1(1))
=
(30)
21

---

Published as a conference paper at ICLR 2024
and compute the desired log-density at x via
log p1 (x) = Pog - fo(2).
(31)
For manifolds that are embedded in an ambient Euclidean space (e.g., hypersphee,, flat tori, triangular meshes), the parameterization in equation 26 allows to compute the Riemannian divergence directly in the ambient space (Rozen et al., 2021, Lemma 2). That is,
div, (we) divæ (vx) = Σ (o)(
(32)
For general manifolds with metric tensor g (e.g., Poincaré ball model of hyperbolic manifold, the manifold of symmetric positive definite matrices), we compute the Riemannian divergence as
1 t VE log det g
div, (u)) divg (v)) + 2
(33)
where divg is the standard Euclidean divergence and V Dan gradient.
Ø )ar is the Buclidean
Wrapped distributions. effective way to define a simple distribution p E over a manifold M of dimension d is pushing some simple prior ÿ defined on some euclidean space Rd via chart 6: _ + Mr for example ø = exp, M + .M the Riemannian exponential map. Generating sample p(x) is done by drawing sample ~ p(2), E Rd and computing ((z). To compute the probability density p(x) at some point x E M,, we integrate over some arbitrary domain Ω M.
[ p(z) =
p(o(z)) g(2)dz.
(34)
(0-)(0))
where (z) = (0,6(2), (z))) i,j E [d]. is the Riemannian metric tensor in local coordinates,
and (,6(2) = 0d(z) 8z, From this integral we get that
p(2) p(p())) /dt g(z)
(35)
and therefore
plo p()-
p(x)
(36)
Vdet (x)-
and in log space
log p(x) = log p( (x)) log det g (x)).
(37)
Numerical accuracy. On the hypersphere, NLL values were computed using an adaptive step size ODE solver (dopri5) with tolerances of 1e-7. On the high dimensional flat torus and SPD manifolds, we use the same solver but with tolerances of 1e-5. 5. We always check that the solution does not leave the manifold by ensuring the difference between the solution and its projection onto the manifold is numerically negligible.
On general geometries represented using discrete triangular meshes, used 1000 Euler steps with projection after every step for evaluation (NLL computation and sampling after training). During training on general geometries, we solve for the path 2; using 300 Euler steps with projection after every step. In order to avoid division by zero during the computation of the conditional vector field in equation 13, we solve H, from = = 1- E, where e is taken to be le-5; this effectively flows the base distribution to non-degenerate distribution around xı that approximates the Dirac distribution, similar to the role of aii of Lipman et al. (2023).
See Hairer et al. (2006) and Hairer (2011) for overviews ODE solving on manifolds.
22

---

Published as a conference paper at ICLR 2024
(a) Geodesic paths
(b) Learned flow
(c) Model samples
Figure 12: (a) Geodesic paths on a hyperbolic manifold (represented using the Poincaré disk model) originating from single point (blue star). (b) Learned CNF where blue samples are from p(x0) and orange samples are from ((x1). The CNF respects the geometry of the hyperbolic manifold, and learns to transport along geodesic paths.
I EMPIRICAL RUNTIME ESTIMATES
During Riemannian Flow Matching training, there
main computational considerations:
(i) solving for Fr, and (ii) computing the training objective. In comparison to diffusion models, Conditional Flow Matching has the clear advantage in (ii), since don't need to estimate conditional score function through an infinite series (as in DSM), nor do we require divergence estimation (as in ISM). In the following, focus on runtime for different ways (i) solving for xęi
Simulation of ODE/SDE (200 steps) on a flat torus:
6.36 iterations/ second
Simulation-free a flat torus:
104.04 iterations second
Simulation of ODE/SDE (200 steps) on the bunny mesh:
0.422 iterations / second
These numbers were benchmarked on Tesla V100 GPU, with batch size 64. The runtime is for the full training loop, but the main difference between the three lines is how xa is solved while all others (i.e. architecture) are fixed.
Generally, the bunny and other mesh manifolds are more expensive due the projection operator applied after every step, which we implemented rather naively for meshes. However, comparing to iteratively solving an ODE/SDE even on simple manifolds, we see a significant speedup of roughly 17x even when taking into account the full training loop (including gradient descent etc). This shows the efficiency gains from using simulation-Fi training over simulation- -based
ADDITIONAL EXPERIMENTS
Here we consider manifolds with constrained domains and non-trivial metric tensors, specifically, hyperbolic space and manifold of symmetric positive definite matrices, equipped with their standard Riemannian metrics. See Table for a summary of the geometries of these manifolds.
J.1 HYPERBOLIC MANIFOLD
We use the Poincaré disk model for representing hyperbolic space 2-D. Figure 12 visualizes geodesic paths originating from single point the manifold, a learned CNF using Riemannian Conditional Flow Matching, and samples from the learned CNE. Our learned CNF respects the geometry of the manifold and transports samples along geodesic paths, recovering near-optimal transport map in line with the Riemannian metric. Similarly, due to the use of this metric, the CNF never transports outside of the manifold.
23

---

Published as a conference paper at ICLR 2024
Table 6: Test set evaluation on EEG datasets. Stdev. estimated over runs.
BCI-IV-2b 6x6 (2ID)
BCI-IV-2a x5 (325D)
BCI-IV-1 59x59 (1770D)
Geodesic
Norm
NLL
Valid
NLL
Valid
NLL
Valid
Euclidean
Euclidean
-1158=026
100=om
-776.07
81.23
N/A
0=:00
Riemannian
Euclidean
-61.64=022
100-om0
-277.06=087
91.47=5.91
0=a00
Euclidean
Riemannian
-52.22=967
100=0.00
-267 4277447
100
-1167.63+
10000.00
Riemannian
Riemannian
-61.76=0224
100
-271.5421.17
100
-1209.88:53.5
100:0000
J.2 MANIFOLD OF SYMMETRIC POSITIVE MATRICES
We use the space of symmetric positive definite (SPD) matrices with the Riemannian metric (Moakher & Batchelor, 2006). We construct datasets using electmenncepallgrrpply (EEG) data collected by Blankertz et al. (2007); Brunner et al. (2008); Leeb et al. (2008) for a Brain- ComppeCe Interface (BCI) competition. We then computed the covariance matrices of these signals, following standard preprocessing procedure for analyzing EEG signals (Barachant et al., 2013).
In Table 6, report estimates of negative log- likelihood (NLL) and the percentage of simulated samples that are valid SPD matrices (i.e., samples which lie on the manifold). We ablate and note the importance of the Riemannian geodesic and the Riemannian norm during training.
Riemannian geodesic. We compare using Riemannian geodesics (i.e., setting the premetric be the geodesic distance) and Euclidean geodesics (i.e,, setting the premetric be the L2 distance). In comparing between different paths, we find that the Riemannian one generally performs better as respects the geometry of the underlying manifold. Figure 13 visualizes the space of 2 SPD matrices a convex cone. It displays how the Riemannian geodesic behaves- its flow always becomes perpendicular to the boundary when it gets close to boundary and therefore does leave the manifold-whereas the Euclidean geodesic ignores this geometry.
Figure 13: Comparison of extrapolated geodesic paths on SPD manifold, visu- alized (black) Riemannian geodesic. (violet) Euclidean geodesic.
Riemannian norm. We also compare between using the Riemannian norm (i.e, 11-11" for training and using the Eu- clidean norm (i.e., 11-112). Theoretically, the choice norm does not affect the optimal u, (z), which will equal to u, (æ); however, when v, is modeled with limited capacity neural network, the choice of norm can very important it affects which regions the optimization focuses (in particular, regions with large metric tensor). particular, on the SPD manifold, similar to hyperbolic, the metric tensor increases to infinity in regions where the matrix is close being singular (i.e., ill-conditioned). We find that especially for larger SPD matrices, using the Riemannian norm important to ensure U; does not leave the manifold during simulation (that is, the simulated result is still SPD matrix).
cone.
24


Reviews:
[{"review_id": "m9Bz3U2aqm", "reviewer": "Reviewer_FTs2", "summary": "This paper proposes Riemannian flow matching (RFM), a generalization of Euclidean flow matching techniques to arbitrary (complete and connected) Riemannian manifolds. The basic idea is similar to the Euclidean setting, where the authors seek to learn a vector field pushing a tractable base distribution to the data distribution, which is achieved by conditional flow matching on a specified path of probability measures. However, the Riemannian setting introduces several challenges, chiefly in the construction of said probability paths and corresponding flows. The authors demonstrate how to construct such flows using general premetrics, a special case of which is the geodesic distance. Practically, the method enables is simple, scalable approach for learning continuous-time normalizing flows on general manifolds.", "questions": "- The base measure used throughout the paper is the uniform distribution. Is there any benefit to considering non-uniform base distributions? While the theory easily admits arbitrary base distributions, are there practical challenges in using non-uniform distributions?\n- Some more exposition regarding practical choices surrounding premetric choices would improve the paper. For instance, can we say anything about what makes a particular premetric \"good\"? Are there principled reasons someone would choose the diffusion distance over the biharmonic distance? Are spectral distances chosen in Section 3.2 primarily for practical reasons, or are there other aspects that make the corresponding premetrics desirable? Is there a compelling reason spectral distances use the spectrum of Laplace-Beltrami operator, rather than some other operator?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The problem of learning models with geometric constraints is both difficult and important, and this paper solves several key practical challenges present in prior work in this area (namely, simulation-free and scalable training).\n- The paper is a joy to read -- the paper is atypically clear in its aims, methods, and results.\n- The empirical support for the method is clear and convincing across a diverse set of settings, and the evidence provided by the experiments generally supports the conclusions drawn by the authors.\n- All theoretical claims made in the paper are, to the best of my knowledge, complete and correct.", "weaknesses": "- It would have been interesting to have an experiment exploring the effect of $k$ on the learned model (i.e. the number of terms taken in the spectral distance, Equation 16). \n- In the same vein, additional experiments on the choice of scheduler $\\kappa(t)$ could provide valuable practical insights regarding the design of these models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YngXdTTkRs", "reviewer": "Reviewer_tLmW", "summary": "This work presents a Flow matching framework for generative modeling on Riemannian manifolds, where the paper proposes a novel and simple conditional flow defined on general manifolds through the premetric. Experimental results show that the proposed approach can effectively model data on diverse manifolds and can scale to higher dimensions which previous diffusion models failed.", "questions": "- Please address the questions in the Weakness.\n\n- Why is the result of the proposed method for Flood dataset in Table 2 highlighted in bold even though CNF Matching shows better result?\n\n- Can the other diffusion models applied to general manifolds such as the triangular mesh? For example, RSGM seems to be applicable when using the denoised score matching with the Varadhan approximation.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written and easy to follow with sufficient background and related works. \n\n- The proposed framework provides important advantages over previous methods: (1) does not require divergence computation for training and is thereby scalable to higher dimensions and (2) readily adaptable to general manifolds.\n\n- Especially, the construction of the conditional flows using the premetrics is novel which makes the proposed method applicable to general geometries.\n\n- This work presents the application to data on general manifolds such as manifolds with non-trivial curvature (triangular mesh) or manifolds with boundaries (maze), whereas previous diffusion models were limited to simple manifolds (e.g., sphere, torus) and their product manifolds.", "weaknesses": "- The advantage of simulation-free training on simple geometries is not clear. How fast is the training of the proposed approach compared to the training of previous diffusion models? As RSGM (Riemannian Score-based Generative Model) observes that they achieve similar test NLL with only 5 in-training simulation steps (Section O.1), I think that the simulation-free training does not provide significant time efficiency.\n\n- The reason for the outperformance of the proposed framework is not clear. Why does the proposed method outperform the Riemannian Diffusion model? Further, why does CNF Matching show superior performance on Earthquake and Flood datasets over the proposed method?\n\nI would like to raise my score if the above concerns are sufficiently addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DnmUW0z4QA", "reviewer": "Reviewer_q47k", "summary": "The authors propose a method Riemannian Flow Matching for finding continuous normalizing flows on manifolds. The main idea is to construct flows from noise samples towards samples from the training distribution in such a way that the entire flow concentrates around the specific training example. The flows can be constructed by flowing along gradients of the Riemannian distance, or, as the authors show, gradients of a premetric. Such premetrics can be constructed to give flows that are efficient to evaluate.", "questions": "No questions.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- well-written and clearly presented paper\n- I believe the main idea with the conditional flows arising from premetrics is both interesting and efficient\n- the method avoids some of limitations of other approaches including being simulation-free in some cases and avoiding divergence computations", "weaknesses": "I have not identified substantial weaknesses.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (oral)