Flow Matching on General Geometries

Introduction

While generative models have recently made great advances in fitting data distributions in Euclidean spaces, there are still challenges in dealing with data residing in non-Euclidean spaces, specifically Riemannian manifolds. These challenges include scalability to high dimensions (e.g., ), the requirement for simulation or iterative sampling during training even for simple geometries like hyperspheres (e.g., ), and difficulties in constructing simple and scalable training objectives.

In this work, we introduce Riemannian Flow Matching (RFM), a simple yet powerful methodology for learning continuous normalizing flows (CNFs; ) on general Riemannian manifolds $\mathcal{M}$. RFM builds upon the Flow Matching framework and learns a CNF by regressing an implicitly defined target vector field $u_{t}(x)$ that pushes a base distribution $p$ towards a target distribution $q$ defined by the training examples. To address the intractability of $u_{t}(x)$, we employ a similar approach to Conditional Flow Matching , where we define and regress onto conditional vector fields $u_{t}(x|x_{1})$ that push $p$ towards individual training examples $x_{1}$.

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 a “premetric” $\textrm{d}(x,y)$, which distinguishes pairs of points $x$ and $y$ on the manifold. A natural choice for such a premetric is the geodesic distance function, which coincides with the straight trajectories previously used in Euclidean space.

On simple geometries, where geodesics are known in closed form (e.g.,, Euclidean space, hypersphere, hyperbolic space, torus, or any of their product spaces), Riemannian Flow Matching remains completely simulation-free. Even on general geometries, it only requires forward simulation of a relatively simple ordinary differential equation (ODE), without differentiation through the solver, stochastic iterative sampling, or divergence estimation.

On all types of geometries, Riemannian Flow Matching offers several advantages over recently proposed Riemannian diffusion models . These advantages include avoiding iterative simulation of a noising process during training for geometries with analytic geodesic formulas; not relying on approximations of score functions or divergences of the parameteric vector field; and not needing to solve stochastic differential equations (SDE) on manifolds, which is generally more challenging to approximate than ODE solutions . Table 1 summarizes the key differences with relevant prior methods, which we expand on further in Section 4 (Related Work).

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.

Preliminaries

Probability paths and flows on manifolds.

and the final diffeomorphism is defined by setting $\Psi(x)=\psi_{1}(x)$. Given a probability density path $p_{t}$, it is said to be generated by $u_{t}$ from $p$ if $\psi_{t}$ pushes $p_{0}=p$ to $p_{t}$ for all $t\in$. More formally,

where the $\sharp$ symbol denotes the standard push-forward operation and $x_{s}=\psi_{s}(\psi_{t}^{-1}(x))$. This formula can be derived from the Riemannian version of the instantaneous change of variables Formula (see equation 22 in ). Previously, Chen et al., suggested modeling the flow $\psi_{t}$ implicitly by considering parameterizing the vector field $u_{t}$. This results in a deep generative model of the flow $\psi_{t}$, called a Continuous Normalizing Flow (CNF) which models a probability path $p_{t}$ through a continuous-time deformation of a base distribution $p$. A number of works have formulated manifold variants that require simulation in order to enable training, while some simulation-free variants scale poorly to high dimensions and do not readily adapt to general geometries.

Method

We aim to train a generative model that lies on a complete, connected smooth Riemannian manifold ${\mathcal{M}}$ endowed with a metric $g$. Concretely, we are given a set of training samples $x_{1}\in{\mathcal{M}}$ from some unknown data distribution $q(x_{1})$, $q\in\mathcal{P}$. Our goal is to learn a parametric map $\Phi:{\mathcal{M}}\rightarrow{\mathcal{M}}$ that pushes a simple base distribution $p\in\mathcal{P}$ to $q$.

Flow Matching is a method to train Continuous Normalizing Flow (CNF) on Euclidean space that sidesteps likelihood computation during training and scales extremely well, similar to diffusion models , while allowing the design of more general noise processes which we make use of in this work. We provide a 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.

where $t\sim{\mathcal{U}}$, the uniform distribution over $$.

Probability path construction.

Riemannian Flow Matching therefore requires coming up with a probability density path $p_{t}\in\mathcal{P}$, $t\in$ that satisfies the boundary conditions

and a corresponding vector field (VF) $u_{t}(x)$ which generates $p_{t}(x)$ from $p$ in the sense of equation 2. One way to construct such a pair is to create per-sample conditional probability paths $p_{t}(x|x_{1})$ satisfying

where $\delta_{x_{1}}(x)$ is the Dirac distribution over ${\mathcal{M}}$ centered at $x_{1}$. One can then define $p_{t}(x)$ as the marginalization of these conditional probability paths over $q(x_{1})$.

which satisfies equation 4 by construction. It was then proposed by Lipman et al., —which we verify for the manifold setting—to define $u_{t}(x)$ as the “marginalization” of conditional vector fields $u_{t}(x|x_{1})$ that generates $p_{t}(x|x_{1})$ (in the sense detailed in Section 2),

which provably generates $p_{t}(x)$. However, directly plugging $u_{t}(x)$ into equation 3 is intractable as computing $u_{t}(x)$ is intractable.

Riemmanian Conditional Flow Matching.

A key insight from Lipman et al., is that when the targets $p_{t}$ and $u_{t}$ are defined as in equations 6 and 7, the FM objective is equivalent to the following Conditional Flow Matching objective,

as long as $u_{t}(x|x_{1})$ is a vector field that generates $p_{t}(x|x_{1})$ from $p$.

To simplify this loss, consider the conditional flow, which we denote via the shorthand,

defined as the solution to the ODE in equation 1 with the VF $u_{t}(x|x_{1})$ and the initial condition $\psi_{0}(x_{0}|x_{1})=x_{0}$. Furthermore, since sampling from $p_{t}(x|x_{1})$ can be done with $\psi_{t}(x_{0}|x_{1})$, where $x_{0}\sim p(x_{0})$, we can reparametrize equation 8 as

where $\dot{x}_{t}=\nicefrac{{d}}{{dt}}\;x_{t}=u_{t}(x_{t}|x_{1})$.

Riemannian Conditional Flow Matching (RCFM) has three requirements: a parametric vector field $v_{t}$ that outputs vectors on the tangent planes, the use of the appropriate Riemannian metric $\left\|\cdot\right\|_{g}$, and the design of a (computationally tractable) conditional flow $\psi_{t}(x|x_{1})$ 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 can be simple and highly scalable; the training procedure is summarized in Algorithm 1 and a detailed comparison can be found in Appendix D.

2 Constructing Conditional Flows through Premetrics

We discuss the construction of conditional flows $\psi_{t}(x|x_{1})$ on ${\mathcal{M}}$ that concentrate all mass at $x_{1}$ at time $t=1$; Figure 2 provides an illustration. This ensures that equation 5 will hold (regardless of the choice of $p$) since all points are mapped to $x_{1}$ at time $t=1$, namely

Non-negative: $\textrm{d}(x,y)\geq 0$ for all $x,y\in{\mathcal{M}}$.

Non-degenerate: $\nabla\textrm{d}(x,y)\neq 0$ iff $x\neq y$.

We use as convention $\nabla\textrm{d}(x,y)=\nabla_{x}\textrm{d}(x,y)$. Such a premetric denotes the closeness of a point $x$ to $x_{1}$, and we aim to design a conditional flow $\psi_{t}(x|x_{1})$ that monotonically decreases this premetric. That is, given a monotonically decreasing differentiable function $\kappa(t)$ satisfying $\kappa(0)=1$ and $\kappa(1)=0$, we want to find a $\psi_{t}$ that decreases $\textrm{d}(\cdot,x_{1})$ according to

Here $\kappa(t)$ acts as a scheduler that determines the rate at which $\textrm{d}(\cdot|x_{1})$ decreases. Note that at $t=1$, we necessarily satisfy equation 11 since $x_{1}$ is the unique solution to $\textrm{d}(\cdot|x_{1})=0$ due to the “positive” property of the premetric. Our next theorem shows that $\psi_{t}(x|x_{1})$ satisfying equation 12 results in the following vector field,

$u_{t}(x|x_{1})=\frac{d\log\kappa(t)}{dt}\textrm{d}(x,x_{1})\frac{\nabla\textrm{d}(x,x_{1})}{\left\|\nabla\textrm{d}(x,x_{1})\right\|^{2}_{g}},$ (13)

The “non-degenerate” property guarantees this conditional vector field is defined everywhere $x\neq x_{1}$.

The flow $\psi_{t}(x|x_{1})$ defined by the vector field $u_{t}(x|x_{1})$ in equation 13 satisfies equation 12, and therefore also equation 11. Conversely, out of all conditional vector fields that satisfy equation 12, this $u_{t}(x|x_{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)=\textrm{d}(x_{t},x_{1})$, where $x_{t}=\psi_{t}(x|x_{1})$ is the flow defined with the VF in equation 13. Differentiation w.r.t. time gives

The solution of this ODE is $a(t)=\kappa(t)\textrm{d}(x,x_{1})$, which can be verified through substitution, and hence proves $\textrm{d}(x_{t},x_{1})=\kappa(t)\textrm{d}(x,x_{1})$. Intuitively, $u_{t}(x|x_{1})$ is the minimal norm solution since it does not contain orthogonal directions that do not decrease the premetric.

A simple choice we make in this paper for the scheduler is $\kappa(t)=1-t$, resulting in a conditional flow that linearly decreases the premetric between $x_{t}$ and $x_{1}$. Using this, we arrive at a more explicit form of the RCFM objective,

For general manifolds ${\mathcal{M}}$ and premetrics d, training with Riemmanian CFM will require simulation in order to solve for $x_{t}$, though it does not need to differentiate through $x_{t}$. However, on simple geometries RCFM can become completely simulation-free by choosing the premetric to be the geodesic distance, as we discuss next.

A natural choice for the premetric $\textrm{d}(x,y)$ over a Riemannian manifold ${\mathcal{M}}$ is the geodesic distance $\textrm{d}_{g}(x,y)$. Firstly, we note that when using geodesic distance as our choice of premetric, the flow $\psi_{t}(x_{0}|x_{1})$—since it is the minimal norm solution—is equivalent to the geodesic path, i.e., shortest path, connecting $x_{0}$ and $x_{1}$.

Consider a complete, connected smooth Riemannian manifold $({\mathcal{M}},g)$ with geodesic distance $\textrm{d}_{g}(x,y)$. In case $\textrm{d}(x,y)=\textrm{d}_{g}(x,y)$ then $x_{t}=\psi_{t}(x_{0}|x_{1})$ defined by the conditional VF in equation 13 with the scheduler $\kappa(t)=1-t$ is a geodesic connecting $x_{0}$ to $x_{1}$.

This makes it easy to compute $x_{t}$ on simple manifolds, which we define as manifolds with closed-form geodesics, e.g., Euclidean space, the hypersphere, the hyperbolic space, the high-dimensional torus, and some matrix Lie Groups. In particular, the geodesic connecting $x_{0}$ and $x_{1}$ can be expressed in terms of the exponential and logarithm maps,

This formula can simply be plugged into equation 10, resulting in a highly scalable training objective. Figure 1 (Right) shows an illustration for a sphere where $x_{t}$ is chosen as the geodesic path and is computed exactly in closed-form. A list of simple manifolds that we consider can be found in Table 5.

Euclidean geometry.

which coincides with the Euclidean case of Flow Matching presented in prior works , demonstrating that these are special cases of Riemannian Flow Matching when using geodesic paths.

3 Spectral Distances on General Geometries

Geodesics 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 on ${\mathcal{M}}$ contingent on a one-time upfront cost. In particular, for general Riemannian manifolds, we consider the use of approximate spectral distances as an alternative to the geodesic distance. Spectral distances actually offer some benefits over the geodesic distance such as robustness to topological noise, smoothness, and are globally geometry-aware . Note however, that spectral distances do not define minimizing (geodesic) paths, and will require simulation of $u_{t}(x|x_{1})$ in order to compute conditional flows $x_{t}$.

Diffusion Distance : $w(\lambda)=\exp(-2\tau\lambda)$, with a hyperparameter $\tau$.

Biharmonic Distance : $w(\lambda)=\lambda^{-2}$.

In practice, we truncate the infinite series in equation 16, and use only $k$ eigenfunctions corresponding to the smallest $k$ eigenvalues (i.e., the dominant terms in the sum, excluding the first eigenvalue which is zero). These $k$ eigenfunctions can be numerically solved as a one-time preprocessing step 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, we consider manifolds with boundaries and show that solving eigenfunctions using the natural, or Neumann, boundary conditions ensures that the resulting $u_{t}(x|x_{1})$ does not leave the interior of the manifold. More detailed discussions on these points can be found in Appendix G. Figure 3 visualizes contour plots of these spectral distances for manifolds with non-trivial curvatures.

Related Work

Some initial works suggested constructing normalizing flows that map between simple manifolds and an Euclidean space of the same intrinsic dimension , often relying on the tangent space at some pre-specified origin. However, this approach is problematic when the manifold is not homeomorphic to Euclidean space, resulting in both theoretical and numerical issues. On the other hand, continuous-time models such as continuous normalizing flows bypass such topogical constraints and flow directly on the manifold itself. To this end, a number of works have formulated continuous normalizing flows on simple manifolds , but these rely on maximum likelihood for training, a costly simulation-based procedure. More recently, simulation-free training methods for continuous normalizing flows on manifolds have been proposed ; however, these scale poorly to high dimensions and do not adapt to general geometries. Another direction is when the manifold is unknown but assumed to be embedded in an ambient Euclidean space, a number of works have discussed methods for learning the manifold directly .

Manifold diffusion models.

With the influx of diffusion models that allow efficient simulation-free training on Euclidean space , multiple works have attempted to adopt diffusion models to manifolds . However, due to the reliance on stochastic differential equations (SDE) and denoising score matching , these approaches are bound to introduce in-training simulation and approximations when generalized to manifolds.

First and foremost, they lose the simulation-free sampling of $x_{t}\sim p_{t}(x|x_{1})$ that is offered in the Euclidean regime; this is because the manifold analog of the Ornstein–Uhlenbeck SDE does not have closed-form solutions. Hence, diffusion-based methods have to resort to simulated random walks as a noising process even on simple manifolds .

Furthermore, again even on simple manifolds, the conditional score function is not known analytically, so De Bortoli et al., proposed approximating the conditional score function with either eigenfunction-expansion or Varhadan’s heat-kernel approximations. These approximations lead to biased gradients in the denoising score matching framework.

Finally, the use of SDEs as a noising process requires carefully constructing suitable reverse-time processes that approximate either just the probability path or the actual sample trajectories , whereas ODE solutions are generally well-defined in both forward and reverse directions .

In contrast to 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 in Table 1, and a detailed comparison of algorithmic differences to diffusion-based 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 functions that don’t satisfy all axioms of a metric—such as approximate spectral distances with finite truncation—going beyond what is currently possible with existing Riemannian diffusion methods.

Euclidean Flow Matching.

Riemannian Flow Matching is built on top of recent simulation-free methods that work with ODEs instead of SDEs, regressing directly onto generating vector fields instead of score functions , resulting in an arguably simpler approach to continuous-time generative modeling without the intricacies of dealing with stochastic differential equations. In particular, Lipman et al., shows that this approach encompasses and broadens the probability paths used by diffusion models while remaining simulation-free; Albergo and Vanden-Eijnden, discusses an interpretation based on the use of interpolants—equivalent to our conditional flows $\psi_{t}(x|x_{1})$, except our flows are dependent only upon the target data point $x_{1}$ and we also make explicit the construction of the marginal probability path $p_{t}(x)$ and vector field $u_{t}(x)$; Liu et al., shows that repeatedly fitting to a model’s own samples leads to straighter trajectories; and Neklyudov et al., formulates a variational objective when $u_{t}(x)$ is taken to be a gradient field.

Experiments

We empirically validate our approach on both real-world data and synthetic distributions that lie on a Riemannian manifold. We consider data from earth and climate science, protein structures, high-dimensional tori, complicated 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. Ablation experiments on hyperbolic manifold and symmetric positive definite matrices can be found in Appendix I. Table 5 provides a summary of the geometries.

We make use of the publicly sourced datasets compiled by Mathieu and Nickel, . These data points lie on the 2-D sphere, a 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 mean and standard deviation of negative log likelihood (NLL) on the test set computed over 5 random runs with different splits of the dataset, as was done in prior works. We achieve a 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 and RNA datasets compiled by Huang et al., . These datasets represent torsion angles and can be represented on the 2D and 7D torus. We represent the data on a flat torus, which is isometric to the product of 1-D spheres used by prior works and result in densities that are directly comparable due to this isometry. Table 3 contains results aggregated over 5 random splits of the dataset. We match the performance of Huang et al., on the 2D tori datasets as these models are likely close to optimal already. However, we see a significant gain in performance on the higher dimensional 7D torus, due to the higher complexity of the dataset. We show learned densities of the protein datasets in Figure 9.

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 Bortoli et al., , we construct a wrapped Gaussian distribution on the N-D torus with a uniformly sampled mean and fixed variance of $0.2$. We compare to Moser Flow , which does not scale well into high dimensions, and Riemannian Score-based using implicit score matching (ISM).

As shown in Table 1, while this objective gets around the need to approximate conditional score functions, it requires stochastic divergence computation during training, introducing larger amounts of variance at higher dimensions. In Figure 5 we plot the test log-likelihood values, normalized by dimension, across these two baselines and our method with the geodesic construction. We see that our method performs steadily, with no significant drop in performance at higher dimensions since we do not have any stochastic estimation or reliance on approximations.

Manifolds with non-trivial curvature.

We next experiment with general closed manifolds using spectral distances as described in Section 3.3. Specifically, we experiment on manifolds described by triangular meshes, where computing eigenfunctions of the Laplace-Beltrami operator amounts to eigendecomposition of a sparse positive definite matrix, which can be done on the CPU rather efficiently. As our manifolds, we use the Standard Bunny and Spot the Cow . We construct synthetic distributions by computing the $k$-th eigenfunction, thresholding, and then sampling proportionally to the eigenfunction. This is done on an upsampled version of the mesh so that the distribution is non-trivial on each triangle. Figure 4 contains visualizations of the eigenfunctions, the learned density, and samples from a trained model that transports from a uniform base distribution. We used $k$=200 eigenfunctions to compute spectral distances. Our method is able to produce high fidelity samples on these discrete mesh manifolds.

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 $\tau$ while the biharmonic distance was straightforward to use out-of-the-box and it has better smoothness properties (see Figure 3).

Manifolds with boundaries.

Lastly, we experiment with manifolds that have boundaries. Specifically, we consider randomly generated mazes visualized in Figure 6. We set the base distribution to be a Gaussian in the middle of the maze, and set the target distribution to be a mixture of densities at corners of the maze. These mazes are represented using triangular meshes, and we use the biharmonic distance using $k$=30 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.

Conclusion

We propose Riemannian Flow Matching as a highly-scalable approach for training continuous normalizing flows on manifolds. Our method bypasses many of the inherent inconveniences in prior methodologies, is completely simulation-free on simple geometries that have closed-form geodesics, and applies readily to general geometries contingent on a one-time preprocessing cost. We find that our approach achieves state-of-the-art performance on real-world manifold datasets, with no drop in performance when scaling to high dimensions, and we showcase for the first time, tractable training on general geometries including both closed manifolds and manifolds with boundaries.

Ricky T. Q. Chen would like to thank Chin-Wei Huang for his help in clarifying his prior works. Additionally, we acknowledge the Python community for developing the core set of tools that enabled this work, including PyTorch , PyTorch Lightning , Hydra , Jupyter , Matplotlib , seaborn , numpy , SciPy pandas , geopandas , torchdiffeq , libigl , and PyEVTK .

References

Appendix A Flow Matching Loss on manifolds

We provide the necessary derivations and proofs for the Conditional Flow Matching over a Riemannian manifolds; the proofs and derivations from are followed "as-is", with the necessary adaptation to the Riemannian setting.

Assumptions. We will use notations and setup from Section 2. Let $p(\cdot|x_{1}):\rightarrow\mathcal{P}$ be a (conditional) probability path sufficiently smooth with integrable derivatives, strictly positive $p_{t}(x|x_{1})>0$, and $p_{0}(x|x_{1})=p$, where $p\in\mathcal{P}$ is our source density. Let $u(\cdot|x_{1})\in\mathfrak{U}$ be a (conditional) time-dependent vector field, sufficiently smooth with integrable derivatives and such that

Further assume $u_{t}(x|x_{1})$ generates $p_{t}(x|x_{1})$ from $p$ in the sense of equation 2, i.e., if we denote by $\psi_{t}(x|x_{1})$ the solution to the ODE (equation 1):

Proof of the marginal VF formula, equation 7. First, the Mass Conservation Formula Theorem (see, e.g., ) implies that $p_{t}(x|x_{1})$ and $u_{t}(x|x_{1})$ satisfy

Next, we differentiate the marginal $p_{t}(x)$ w.r.t. $t$:

where in the first and second equalities we changed the order of differentiation and integration, and in the second equality we used the mass conservation formula for $u_{t}(x|x_{1})$. In the previous to last equality we multiplied and divided by $p_{t}(x)$. In the last equality we defined the marginal vector field $u_{t}$ 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:

We got that ${\mathcal{L}}_{\scriptscriptstyle\text{RCFM}}(\theta)$ and ${\mathcal{L}}_{\scriptscriptstyle\text{RFM}}(\theta)$ differ by a constant,

Appendix B Proof of Theorem 3.1

Let $x_{t}=\psi_{t}(x|x_{1})$ be the flow defined by equation 13 in the sense of equation 1. Differentiating the time-dependent function $\textrm{d}(x_{t},x_{1})$ w.r.t. time gives

This shows that the function $a(t)=\textrm{d}(x_{t},x_{1})$ satisfies the ODE

with the initial condition $a(0)=\textrm{d}(x,x_{1})$. General solutions to this ODE are of the form $a(t)=c\kappa(t)$, where $c>0$ is a constant set by the initial conditions. This can be verified by substitution. The constant $c$ is set by the initial condition,

This gives $a(t)=\textrm{d}(x,x_{1})\kappa(t)$ as the solution. Due to uniqueness of ODE solutions we get that equation 12 holds.

Conversely, consider $x_{t}=\psi_{t}(x|x_{1})$ satisfying equation 12. Differentiating both sides of this equation w.r.t. $t$ and then using equation 12 again we get

If we let $u_{t}(x|x_{1})$ denote the VF defining the diffeomorphism $\psi_{t}(x|x_{1})$ in the sense of equation 1 then the last equation takes the form

This equation provides an under-determined linear system for $u_{t}(y|x_{1})\in T_{y}{\mathcal{M}}$ at every point $y=x_{t}$ with non-zero probability, which in our case is all $y\in{\mathcal{M}}$ as we assume $p_{t}(y|x_{1})>0$ for all $y\in{\mathcal{M}}$. As can be seen in equation 21, $u_{t}(x|x_{1})$ defined in equation 13 is also satisfying this equation. Further note, that since $u_{t}(x|x_{1})$ defined in equation 13 satisfies $u_{t}(x|x_{1})\parallel\nabla\textrm{d}(x,x_{1})$ (proportional) it is the minimal norm solution to the linear system in equation 22. ∎

Appendix C Proof of Proposition 3.2

Consider a complete, connected smooth Riemannian manifold $({\mathcal{M}},g)$ with geodesic distance $\textrm{d}_{g}(x,y)$. In case $\textrm{d}(x,y)=\textrm{d}_{g}(x,y)$ then $x_{t}=\psi_{t}(x_{0}|x_{1})$ defined by the conditional VF in equation 13 with the scheduler $\kappa(t)=1-t$ is a geodesic connecting $x_{0}$ to $x_{1}$.

First, note that by definition $\psi_{0}(x_{0}|x_{1})=x_{0}$, and $\psi_{1}(x_{0}|x_{1})=x_{1}$. Second, from Proposition 6 in McCann, we have that

where $\log$ is the Riemannian logarithm map. From the chain rule we have

Since the logarithm map satisfies $\left\|\log_{x}(y)\right\|_{g}=\textrm{d}_{g}(x,y)$ we have that

Now, computing the length of the curve $x_{t}$ we get

where in the second equality we used the definition of the conditional VF (equation 13) with $\kappa(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 $x_{t}$ realizes a minimum of the length function, it is a geodesic. ∎

Appendix D Algorithmic comparison to Riemannian diffusion models

Appendix E Limitations

As can be seen from Figure 7, our method still requires simulation of $x_{t}$ on general manifolds. This sequential process can be time consuming, and a more parallel or simulation-free approach to constructing $x_{t}$ would be more favorable. Furthermore, the spectral distances require eigenfunction solvers which may be computationally expensive on complex manifolds. Using approximate methods such as neural eigenfunctions may be a possibility. One major advantage of our premetric formulation is that these eigenfunctions need not be perfectly solved in order to satisfy the relatively simple properties of our premetric.

Appendix F Additional figures

Appendix G Additional Discussion

Computation cost. The smallest $k$ eigenvalues and their eigenfunctions need only be computed once as a pre-processing step. On manifolds represented as discrete triangular meshes, this step can be done in a 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 $x_{t}$, 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-based methods that require simulation even for simple manifolds.

Sufficiency with finite $k$. One may wonder if we pay any approximation costs when using finite $k$; 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 to compute the spectral distances perfectly, only sufficiently enough that the conditions of our premetric are satisfied. Regarding the question of what $k$ 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 [33, Theorem 2]. Nevertheless, the use of spectral distances computed with $k$ smallest eigenvalues is equivalent to computing Euclidean distances in a $k$-dimensional Euclidean embedding using the same eigenfunctions; this embedding is known to preserve neighborhoods optimally .

G.2 Manifolds with Boundary

In considering general geometries, we also consider the case where ${\mathcal{M}}$ has a boundary, denoted $\partial{\mathcal{M}}$. In this case, we need to add another condition to our premetric to make sure $u_{t}(x|x_{1})$ will not flow particles outside the manifold. Let $n(x)\in T_{x}{\mathcal{M}}$ denote the interior-pointing normal direction at a boundary point $x\in\partial{\mathcal{M}}$. We add the following condition to our premetric:

Boundary: $\left\langle\nabla\textrm{d}(x,y),n(x)\right\rangle_{g}\leq 0$, $\forall y\in{\mathcal{M}},x\in\partial{\mathcal{M}}$.

If the premetric satisfies this condition, then the conditional VF in equation 13 satisfies

implying that the conditional vector field does not point outwards on the boundary of the manifold.

Spectral distances at boundary points. In case ${\mathcal{M}}$ 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 $\varphi_{i}$ using the natural, or Neumann, boundary conditions, i.e., their normal derivative at boundary points vanish, and we have $\left\langle\nabla_{g}\varphi_{i}(x),n(x)\right\rangle_{g}=0$ for all $x\in\partial{\mathcal{M}}$. This property implies that $\left\langle\nabla_{x}\textrm{d}_{w}(x,y)^{2},n(x)\right\rangle_{g}=0$, satisfying the boundary condition of the premetric.

Appendix H Experiment Details

Training setup. All experiments are run on a single NVIDIA V100 GPU with 32GB memory. We tried our best to keep to the same training setup as prior works ; 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 set for early stopping based on the validation NLL, and then only computed the test NLL using the checkpoint that achieved the best validation NLL. We used standard multilayer perceptron for parameterizing vector fields where time is concatenated as an 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 with a learnable parameter. We used Adam with a learning rate of 1e-4 and an exponential moving averaging on the weights with a decay of 0.999 for all of our experiments.

High-dimensional tori. We use the same setup as De Bortoli et al., . The data distribution is a wrapped Gaussian on the high-dimensional tori with a uniformly sampled mean and a scale of 0.2. We use a MLP with 3 hidden layers of size 512 to parameterize $v_{t}$, 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 is 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. If two neighboring cells are 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 on the manifold. This produces manifolds represented by triangular meshes where all the triangles are the same size. We randomly create maze structures based on a breadth-first-search algorithm, represent these using meshes, and we then normalize the mesh so that points lie in the range of (0, 1).

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 we model

where $\pi$ is the projection operator onto the manifold, i.e.,

and $P_{y}$ is the orthogonal projection onto the tangent space at $y$.

We also normalize the vector field using $g(x)^{-\frac{1}{2}}$, which cancels out the effect of $g$ on the Riemannian norm and makes standard neural network parameterization more robust to changes in the metric, i.e.,

We found this bypasses the need to construct manifold-specific initialization schemes for our neural networks and leads to more stable training.

Log-likelihood computation. We solve for the log-density $\log p_{1}(x)$, for an arbitrary test sample $x\in{\mathcal{M}}$, by using the instantaneous change of variables , namely solve the ODE

where $\text{div}_{g}$ is the Riemannian divergence. We solve backwards in time from $t=1$ to time $t=0$ with the initial conditions

and compute the desired log-density at $x$ via

For manifolds that are embedded in an ambient Euclidean space (e.g., hypersphere, flat tori, triangular meshes), the parameterization in equation 24 allows us to compute the Riemannian divergence directly in the ambient space [63, Lemma 2]. That is,

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

where $\text{div}_{E}$ is the standard Euclidean divergence and $\nabla_{E}=(\frac{\partial}{\partial x_{1}},\ldots,\frac{\partial}{\partial x_{d}})^{T}$ is the Euclidean gradient.

where $g_{ij}(z)=\left\langle\partial_{i}\phi(z),\partial_{j}\phi(z)\right\rangle_{g}$, $i,j\in[d]$, is the Riemannian metric tensor in local coordinates, and $\partial_{i}\phi(z)=\frac{\partial\phi(z)}{\partial z_{i}}$. From this integral we get that

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. 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, we used 1000 Euler steps with a projection after every step for evaluation (NLL computation and sampling after training). During training on general geometries, we solve for the path $x_{t}$ 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 $x_{t}$ from $t=0$ to $t=1-\varepsilon$, where $\varepsilon$ is taken to be 1e-5; this effectively flows the base distribution to a non-degenerate distribution around $x_{1}$ that approximates the Dirac distribution, similar to the role of $\sigma_{\text{min}}$ of Lipman et al., .

See Hairer et al., and Hairer, for overviews on ODE solving on manifolds.

Appendix I Additional Experiments

Here we consider manifolds with constrained domains and non-trivial metric tensors, specifically, a hyperbolic space and a manifold of symmetric positive definite matrices, equipped with their stadard Riemannian metrics. See Table 5 for a summary of the geometries of these manifolds.

We use the Poincaré disk model for representing a hyperbolic space in 2-D. Figure 10 visualizes geodesic paths originating from a single point on the manifold, a learned CNF using Riemannian Conditional Flow Matching, and samples from the learned CNF. Our learned CNF respects the geometry of the manifold and transports samples along geodesic paths, recovering a 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.

I.2 Symmetric Positive Matrices

We use the space of symmetric positive definite (SPD) matrices with the Riemannian metric . We construct datasets using electroencephalography (EEG) data collected by Blankertz et al., , Brunner et al., , Leeb et al., for a Brain-Computer Interface (BCI) competition. We then computed the covariance matrices of these signals, following standard preprocessing procedure for analyzing EEG signals .

In Table 6, we 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.

We compare using Riemannian geodesics (i.e., setting the premetric to be the geodesic distance) and Euclidean geodesics (i.e., setting the premetric to be the $L_{2}$ distance). In comparing between different paths, we find that the Riemannian one generally performs better as it respects the geometry of the underlying manifold. Figure 11 visualizes the space of $2\times 2$ SPD matrices as a convex cone. It displays how the Riemannian geodesic behaves—its flow always aligns with the boundary and does not leave the manifold—whereas the Euclidean geodesic ignores this geometry.

Riemannian norm.

We also compare between using the Riemannian norm (i.e., $\smash{\left\|\cdot\right\|_{g}^{2}}$ for training and using the Euclidean norm (i.e., $\smash{\left\|\cdot\right\|_{2}^{2}}$). Theoretically, the choice of norm does not affect the optimal $v_{t}(x)$, which will equal to $u_{t}(x)$; however, when $v_{t}$ is modeled with a limited capacity neural network, the choice of norm can be very important as it affects which regions the optimization focuses on (in particular, regions with a large metric tensor). In particular, on the SPD manifold, similar to hyperbolic, the metric tensor increases to infinity in regions where the matrix is close to being singular (i.e., ill-conditioned). We find that especially for larger SPD matrices, using the Riemannian norm is important to ensure $v_{t}$ does not leave the manifold during simulation (that is, the simulated result is still a SPD matrix).