Metric Flow Matching for Smooth Interpolations on the Data Manifold
Authors: Kacper Kapusniak
Peter Potaptchik
Teodora Reu
Leo Zhang
Alexander Tong
Michael M. Bronstein
Joey Bose
Francesco Di Giovanni
Venue: NeurIPS 2024 poster (2024)
Keywords: Flow Matching; Riemannian Geometry; single-cell RNA sequencing;
Primary area: generative_models
TL;DR: We generalize Conditional Flow Matching by learning interpolants that stay on the data manifold leading to more meaningful matching

Abstract: Matching objectives underpin the success of modern generative models and rely on constructing conditional paths that transform a source distribution into a target distribution. Despite being a fundamental building block, conditional paths have been designed principally under the assumption of $\textit{Euclidean geometry}$, resulting in straight interpolations. However, this can be particularly restrictive for tasks such as trajectory inference, where straight paths might lie outside the data manifold, thus failing to capture the underlying dynamics giving rise to the observed marginals. In this paper, we propose Metric Flow Matching (MFM), a novel simulation-free framework for conditional flow matching where interpolants are approximate geodesics learned by minimizing the kinetic energy of a data-induced Riemannian metric. This way, the generative model matches vector fields on the data manifold, which corresponds to lower uncertainty and more meaningful interpolations. We prescribe general metrics to instantiate MFM, independent of the task, and test it on a suite of challenging problems including LiDAR navigation, unpaired image translation, and modeling cellular dynamics. We observe that MFM outperforms the Euclidean baselines, particularly achieving SOTA on single-cell trajectory prediction.

Metric Flow Matching for Smooth Interpolations on the Data Manifold
Kacper Kapušniak'; Peter Potaptchik', Teodora Reu', Leo Zhang', Alexander Tong?.,, Michael Bronstein Avishek Joey Bose Francesco Di Giovanni! University Oxford, Mila, Université de Montréal, "AITHYRA
Abstract
Matching objectives underpin the success of modern generative models and rely constructing conditional paths that transform source distribution into a target distribution. Despite being fundamental building block, conditional paths have been designed principally under the assumption of Euclidean geometry, resulting in straight interpolations. However, this can be particularly restrictive for tasks such as trajectory inference, where straight paths might lie outside the data manifold, thus failing to capture the underlying dynamics giving rise to the observed marginals. In this paper, we propose METRIC FLOW MATCHING (MFM), novel simulation- free framework for conditional flow matching where interpolants are approximate geodesies learned by minimizing the kinetic energy of data-induced Riemannian metric. This way. the generative model matches vector fields on the data manifold, which corresponds lower uncertainty and more meaningful interpolations. We prescribe general metrics to instantiate MFM, independent of the task, and test it on suite of challenging problems including LiDAR navigation, unpaired image trans- lation, and modeling cellular dynamics. We observe that MFM outperforms the Eu- clidean baselines, particularly achieving SOTA on single-cell trajectory prediction.
Introduction
central task in many natural and scientific domains entails the inference of system dynamics of underlying (physical) process from noisy measurements. A core challenge, in these application domains such as biomedical ones-e.g. tracking health metrics [Oeppen and Vaupel, 2002] or diseases [Hay et al., 2021] is that one typically lacks access entire trajetraesssosie and can only leverage cross-sectional samples. An even more poignant example is the case of single-cell RNA sequencing [Macosko et al., 2015, Klein al., 2015], where measurements to the procedure being expensive and destructive. Consequently, the nature of these tasks demands the design of frameworks capable of reconstructing the temporal dynamics from observed time marginals that contain finite samples. This overarching problem specification referred as trajectory inference [Hashimoto al., 2016. Lavenant et al., 2021].
sparse and static, due
system (c.g. cells)
To address this challenge, we rely on matching objectives, a powerful generative modeling paradigm encompassing successful approaches including diffusion models [Sohl-Dickstein et al., 2015, Song et al., 2021], flow matching [Lipman et al., 2023, Liu et al., 2022, Albergo and Vanden-Eijnden. 2022], and finding Schrödinger Bridge [Schrödinger, 1932, Léonard, 2013]. Specifically, to reconstruct the unknown dynamics + pt between observed time marginals Po and pn. we leverage Conditional Flow Matching (CFM) [TTng al., 2023b), a simulation-free framework which constructs probability path Pt through interpolants connecting samples of to samples of P1. In general, 2; is designed under the assumption of Euclidean geometry, resulting in straight
Conresponding author: kacper kww.skipuanaa.oa ak&cs uk Code available
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

p MPM
Pu
CFM
p
Figure 1: In orange and violet the source and target distributions. On the left, straight interpolations vs interpolations following a data-dependent Riemannian metric. On the right, densities of reconstructed marginals at time f 2. using Conditional Flow Matching and METRIC FL.ow MATCHING (MFM), respectively. MFM provides more meaningful reconstruction supported on the data manifold.
trajectories. However, in light of the widely accepted "manifold hypothesis" [Tenenbaum al., 2000, Belkin and Niyogi, 2003], the target time-evolving density p is supported on curved low-dimensional manifold M C R condition satisfied by cells the space of gene expressions [Moon et al., 2018]. As such, straight interpolants stray away from the data manifold M, leading to reconstructions Pt. that fail to model the nonlinear dynamics generated by the underlying process.
Present work. We aim to design interpolants Iy that stay on the data manifold M associated with the underlying dynamics. Nonetheless, parameterizing the dive-diiensional manifold M is prone to instabilities [Loaiza-Ganem et al., 2022] and may require multiple coordinate systems [Salmona et al., 2022]. Accordingly, we adopt the "metric learning" approach [Xing al., 2002, Hauberg et al., 2012], where we still work in the ambient space Rd but equip it with data- data-depennent Riemannian metric g whose shortest- paths (geodesics) stay close to the data points, and hence to M [Arvanitidis et al., 2021]. We introduce METRIC FLOW MATCHING (MFM), a simulation-fre free generalization of CFM where interpolants 24 are learned by minimizing geodesic loss that penalizes the velocity measured by the metric g. As a result, It approximates the geodesics of g and hence tends towards the data, lead- ing to the evaluation of the matching objective in regions of lower uncertainty. Therefore, the resulting probability path pt is supported on the data manifold for all t E [0, 1]. giving rise more natural reconstruction of the underlying dynamics in the trajectory inference task, as depicted in Figure
Our main contributions are:
(1) We prove that given a dataset D C Rdd one always construct metric g on Rd such that the geodesics connecting zo sampled from Pu to æ1 sampled from P1. always lie close to D ($3.1).
(2) We propose METRIC FLOW MATCHING, novel framework generalizing CFM to the Riemannian manifold associated with any data-dependent metric g. In MFM, before training the matching objective, learn interpolants that stay close to the data by minimizing cost induced by g ($3). MFM is simulation- free and stays relevant when geodesics lack closed form and hence Riemannian Flow Matching [Chen and Lipman, 2023] is not easily applicable ($3.2).
(3) We prescribe a universal way of designing a data-dependent metric, independent of the specifics of the task, which enforces interpolants x, stay supported on the data manifold ($4.1). Through the proposed metric, 2, depends on the entire data manifold and not just the endpoints and ø1 sampled from the marginals. By accounting for the Riemannian geometry induced by the data, MFM generalizes existing approaches that construct 2, by minimizing energies ($4.2).
(4) We propose OT-MFM, an instance of MFM that relies on Optimal Transport to draw samples from the marginals ($4). Empirically, we show that OT-MFM attains SOTA results for reconstructing single-cell dynamies ($5). Additionally, we validate the versatility of the framework through tasks such as 3D navigation with LiDAR point clouds and unpaired translation of images.
2 Preliminaries and Setting
We review Conditional Flow Matching, which forms the basis METRIC FLOW MATCHING $3. Next, we recall basic notions of Riemannian geometry, with emphasis constructing geodesies.
Notation and convention. We let Rd be the ambient space where data points are embedded. A random variable a with distribution p is denoted as x p(z). A function yp depending time I,, space x and learnable parameters 0, is denoted by o (ar), and its time derivative by ş0..0- We also &zo (x) be the Dirac function centered at 20 and assume that all distributions are absolutely continuous, which allows

---

us to use densities. We denote the space of symmetric, positive definite d x d matrices as (d) and let G(z) E (PD(d) be the coordinate representation of a Riemannian metric at some point æ.
Conditional Flow Matching. We consider a source distribution Po and a target distribution pl defined on Rd. We are interested in finding map f that pushes forward Po p1, i.e. f = Pi- In line with Continuous Normalizing Flows [Chen et al., 2018], we look for a map of the form f= ti, where the time-dependent diffeomorphism d/x : IR + R is the flow generated by a vector field uz, i.e. dxb,(r)/dt = ue (1%x (z)). with vo (2) = I, for all æ E R If we define the density path Pt [er] then 4s and u, satisfy the continuity equation and we say that pr is generated by ut.
If density path Pr and vector field Un are known, we could regress a vector field U1 modeled as a neural network, to u, by minimizing (FM(0) = p.:) || - uy (x)) Since Pt and un are typically intractable though, Conditional Flow Matching (CFM) [Lipman al., 2023, Albergo and Vanden- Vanden-Eijnden, 2022, Liu et al., 2022] simplifies the problem by assuming that pt is a mixture of conditional paths: pt(x)) = f Pt (x((z()) (ce)2)(()))) where
(zo, is sampled from a joint distribution
q with marginals Po and Pr. and pt (x](z satisfy Pu(z)|z) 820 (x) and pr (x]2) & (x). If T/e 2o, 21) denotes the flow generating (x) then the CFM objective is
(CFM (0) = (mo,m))) (l - 2
where 2; := ( (o, 21) are the interpolants from zo to TT. Since LFM and LCrM have same gra- dients [Lipman et al., 2023, Tong et al., 2023b], we can use the tractable conditional paths to learn vg. As in $3 we design 2t to approximate geodesics, we review key notions from Riemannian geometry.
Riemannian manifolds. A Riemannian manifold (M, .) is smooth orientable manifold M equipped with a smooth map assigning to each point x an inner product (: defined on the tangent space of . at x. We let Gxx) E SPD(d) be the matrix representing g in coordinates, at any point x, with d the dimension of M. Integration is taken with respect to the volume form dvol (see Appendix SA for details). A continuous, positive function p is then a probability density on g).g.. i.c. p E P(MM), if p(z))vol((r) = 1. Naturally, it is possible to define curves y/ indexed by time t. A geodesic is then the curve Yf that minimizes the distance with respect to g. Specifically, the geodesic connecting To to 21 in M, can be found by minimizing the length functional:
L lotel dt,
arg min
[H (Te, G(Md)))
71 = Tel To 70
where i is velocity. From eq. (2), we see that geodesics tend towards regions where [[G(z)] is small. In Euclidean geometry (i.e., M = Rd and G(z) = I), 72 is a straight line with constant speed.
3 Metric Flow Matching
We introduce METRIC FLOW MATCHING (MFM), new simulation-fire free framework that generalizes CFM by constructing probability paths supported on the data manifold. MFM learns interpolants xt eq. (1) whose velocity minimizes a data-dependent Riemannian metric assigning lower cost to regions with high data concentration. Consequently, the CFM objective is evaluated in areas of low data uncertainty, and the corresponding vector field, Vi eq. (1), learns to pass through these regions.
We structure this section as follows. In $3.1 we discuss the trajectory inference problem and how straight interpolants x; in CFM result in undesirable probability paths whose support is not defined on the data manifold. We remedy this problem by choosing to represent the data manifold via Riemannian metric in IRd, such that geodesics avoid straying away from the samples in the training set. In $3.2 we introduce MFM and compare it with Riemannian Flow Matching [Chen and Lipman, 2023].
3.1 Metric learning
Assume that Po and P1 are empirical distributions and that we have access to a dataset of samples D = {x; | inn practice D is constructed concatenating samples from both the source and target distributions. We are interested in the problem of trajectory inference [Lavenant et al., 2021], where we need to reconstruct an unknown dynamics / ~ pe. with observed time marginals p Po and pi = -he extension to multiple timepoints is easy. In many realistic settings, including single-cell RNA sequencing [Macosko et al., 2015], time measurements are sparse, and leveraging biases from
3

---

the data is hence key to achieving a faithful reconstruction. For this reason, we invoke the "manifold hypothesis" [Bengio et al., 2013], where the data arises from owwdimensional manifold M C R= a property satisfied by cells embedded in the space of gene expressions [Moon et al., 2018]:
supp(p?) = {ar E Rd PE (xx) > 0} c M,, 2 0.
Since any regular time dynamics can be described using the continuity equation generated by some vector field ot [Ambrosio et al., 2005, Theorem 8.3.1]. [Ncklyudov et al., 2023a], we rely on CFM, and approximate pi via the probability path Pt associated with vt,0 in eq. (1). From eq. (3), it follows that a valid reconstruction entails supp(p) C M. As the support of pt is determined by the interpolants 2; in eq. (1), need 24 to be constrained to stay on M. However, in the classical CFM setup, this condition is violated since interpolants are often straight lines, agnostic of the data's support: 2t (æ (1 ()zo [Tong et al.. 2023b, Shaul et al., 2023], with æn, xi sampled from the marginals. In this case, if pa (2r xn ) R 6n, (2), then the support of pt satisfies
supp(pt) C {y R 3 (20. q y tz + (1 - t)moo.
However, the dynamics + pē is often nonlinear, as for single-cells [Moon et al., 2018], meaning that supp(p:) Σ M. Straight interpolants are hence too restrictive and should instead be supported on M so to replicate actual trajectories from Zo to 11, which are generated by the underlying process.
Operating on lower-dimensional M is challenging though, since it requires different coordinate sys- tems [Schonsheck et al., 2019, Salmona et al., 2022] and may incur overfitting Sraiza-Ganemm et al., 2022, 2024]. Nonetheless, a key property posited by the "manifold hypothesis" is that M concentrates around the data points D [Arvanitidis et al., 2022, Chadebec and Allassonnière, 2022]. As such, interpolants 2y should remain close to D. Therefore, instead of changing the dimension, we design r, to minimize a cost in IRd that is lower on regions close to D. We achieve this following the "metric learning" approach, [Hauberg et al., 2012] where we equip Rdd with a suitable Riemannian metric g-
Definition 1. A data-dependent metric g on Rdd is a smooth map g Rd + SPD(d) parameterized by the dataset D = {2;} ) c Rd i.e. ((x) = G D) E SPD(d), Vz E R-
We describe specific metric g in $4, and note that g can also enforce constraints from the task ($7). Naturally, if Gz:; D) = I, we recover the Euclidean metric. We show that we can always construct g so that the geodesics 7 stay close to the data D. For details, we refer to Theorem B.1 in Appendix $B.
Proposition (Informal). Given dataset D C Rd, let g be any metric such that: (i) The eigenvalues of Gz D) do not approach zero when x distant from D; (ii) [[G[;; D)Il sufficiently small if is close to D. Then for each (Zo: 21) ~ q. the geodesic 7tt connecting
to x, stays close to D.
Given g in the statement, if x; 7t and (o(r/au. x1) = 82, (x), then the probability path pe generated by ei in eq. (1) has support near D, i.e. supp(p:) lies close to the data manifold M, which is our goal. Unfortunately, for all but the most trivial metrics on R2d, it is not possible to obtain closed- form expressions for the geodesies %t. As such, finding the geodesic of necessitates the expensive simulation of second-order nonlinear Euler- Lagrang equations [Hennig and Hauberg, 2014].
3.2 Parameterization and optimization of interpolants
Consider a metric g on R4 as in Definition whose geodesics 7: lie close to D as per Proposition 1. We propose a simulation--ree approximation to paths 70 by introducing interpolants of the form
It.y (1 t) + tæ + +(1 - (
(4)
where 1) are the parameters of a neural network sr,,, acting as a nonlinear "correction" for straight interpolants. Note that the boundary conditions are met, i.c. that the path zt,, recovers both 2o and Ii at times t 0 and t 1, respectively. In fact, Ft,, reduces to the convex combination between To and 11 if (pt., = 0, meaning that Zt.,, strictly generalize the straight paths used in [Lipman et al., 2023, Liu et al., 2022]. Towards the goal of learning so that Zt.,, approximates the geodesic 78. we note that 7," can be characterized as the path minimizing the convex functional &g below:
Th
arg min 7o =oo =
Eg(T)), Ey(N) = E, [ G(n;
(5)
Since Ye minimizes Eg over all paths connecting zo to x1, and zt.,, in eq. (4) satisfies these boundary conditions, we can estimate by simply minimizing the convex functional E, over It.n. which leads to the following geodesic objective (the training procedure is reported in Algorithm 1):

---

Algorithm Pseudocode for training of geodesic interpolants
Require: coupling ş. initialized network Pt.,, data-dependent metric G(: (D)
1: while Training
2: Sample (No, 11) ~ q and ~ 2(0, 1)
3: It,n f (1 (1-18 t) + txj + t(1 t)),,, m (1)
> eq. (4)
4:
Tt. 21 20 + +(1 t) Pt (2o, 21) (1 2t) (n,
5:
((r) < (i G D D)
D Estimate of objective (g) from eq. (6)
6:
Update y using gradient v,,((1)
return (approximate) geodesic interpolants parametrized by Pt,
Cg (1) == E (1 [Eg (x
[(tr,)) G D)æz.
(6)
Given interpolants that approximate 7e and hence stay close to the data manifold, we can then rely on the CFM objective in eq. (1) to regress the vector field up. Since y makes the ambient space Rd into
Riemannian manifold (R2 .g), we need replace the norm ||-II in eq. (1) with the one II-II, induced by the metric, and rescale the marginals Po- P1 using the volume form induced by g. so to extend Po. to densities (IR .9) (see Appendix $A). Similar arguments work for the joint distribution q. We can finally introduce our framework METRIC FLOW MATCHING that generalizes Conditional Flow Match- ing (1) to leverage any data-dependent metric g, by using interpolants Tt.,, (4), whose parameters y) are obtained from minimizing the geodesic cost Eg. The MFM objective can be stated as:
(0) = (E,,(ro [1,,,((Fe.))) -e,n (
Tj" = arg min Cn(n).
(7)
As the interpolants It.n approximate geodesics of g, MFM the vector field vr.0 regressed on the data manifold M, where the underlying dynamics P is supported, resulting in better reconstructions. Crucially, eq. (6)
A description of METRIC FLOW MATCHING is given in Algorithm 2.
only depends on time derivatives æt..- Therefore, MFM avoids simulations and simply requires training an additional (smaller) network fP.,, in eq. (4), which can be done prior to training vt,0-
MFM versus Riemannian Flow Matching. While CFM has already been extended to Riemannian manifolds in the Riemannian Flow Matching (RFM) framework of Chen and Lipman [2023], MFM crucially differs from RFM two ways. To begin with, MFM relies on the data or task inducing Riemannian metric on the ambient space which is then accounted for in the matching objective. This is in sharp contrast to RFM, which instead assumes that the metric of the ambient space given and is independent of the data points. Secondly, RFM does not incorporate conditional paths that are learned. In fact, in the scenario above where g is a metric whose geodesics 7 stay close to the data support, adopting RFM would entail replacing the MFM objective LMPM in (7) with
CRMM (0) = (w0.aar))) | 46 (l1)))
(8)
However, as argued above, for almost any metric on R4,, geodesics of can only be found via simu- lations, which in high dimensions inhibits the easy application of RFM. Conversely, MFM designs in- terpolants that minimize a geodesic cost (6) and hence approximate 7/ without incurring simulations.
Learning Riemannian Metrics in Ambient Space
In this section, we focus on a concrete choice of g. which can easily be used within MFM ($4.1) I). We also introduce OT-MFM, variant of MFM that leverages Optimal Transport to find coupling q between Po and Pl. Next, in $4.2 we discuss how MFM generalizes recent works that find interpolants that minimize energies by accounting for the Riemannian geometry induced by the data.
4.1 A family of diagonal metrics: LAND and RBF
We consider a family of metrics GLAND as in Definition 1, independent of specifics of the data type or task. For the ease of exposition, we omit to write the explicit dependence the dataset D N
5

---

Given E > O, we let x + (x) = (z) = (diag(h(x)) + +II)
-1 be the "LAND" metric, where
ha(x)) = Σ (xtt x²)² exp ( (- || x
N
1
(9)
20
with a the kernel size. We emphasize that Ga(r) was introduced by Arvanitidis et al. [2016] from which we borrow the name bbut its algorithmic use MFM is fundamentally different. In line with Proposition we see that [[G{(z)] is larger away from D. thus pushing geodesics (2) to stay close to the data support, as desired. While gLAND is flexible and directly accounts for all the samples in D, in high-dimension selecting in eq. (9) can be challenging. For these reasons, we follow Arvanitidis et al. [2021] and introduce variation of GLAND of the form GRBF (z) (diag ((())) +=I) where
K
ha(x))
Σ k=1
Work (x) exp (- Aa.k -- || 2
(-
=x1l²), 1
15 a
(10)
with K' the number of clusters with centers dy and X the bandwidth of cluster &s for channel (see Appendix $C for details). In particular, ho is realized via a Radial Basis Function (RBF) network [Que and Belkin, 2016], where Wa,k > 0 are learned enforce the behavior ha (2:4) as 1 for each data point 24 so that the resulting metric Jeur assigns lower cost to regions close to the centers &g. In our experiments, we then rely on GLAND in low dimensions, and instead use Janr in high dimensions. We also note that all metrics considered are diagonal, which makes MFM more efficient. Explicitly, given the interpolants in eq. (4), the geometric loss in eq. (6) with respect to Jnn can be written as:
d
(itm) h (T)..)
yna (1) = (a Egeer (#t.n)] = (x0,21)))~
(11)
+E
We see that the loss acts as a geometric regularization, with the velocity dt.., penalized more in regions away from the support of the dataset D, i.e. when Aa || -ix || is large for all centers ît. in eq. (10),
OT-MFM. In eq. (7), samples 2o. 2 follow joint distribution q, with marginals po and p1. Since we are interested in the problem of trajectory inference, with emphasis on single-cell applications where the principle of least action holds [Schiebinger, 2021], focus on coupling q that minimizes the distance in probability space between the source and target distributions. Namely, we consider the case where q the 2- Wasserstein optimal transport plan x² from po to P1 [Villani al., 2009]:
/ JR"
arg min xell
2 (x, (I..))d((x.)))
(12)
where 1l are the probability measures on Ed x1 with marginals po and p1 and cis any cost. While we could choose based on Jess. we instead select e to be the L2 distance in Rdd to avoid additional compu- tations and so we can study the role played by Fi. even when q = x is agnostic of the data manifold. We then propose the OT-MFM objective, where 17" minimizes the geodesic loss in eq. (11):
LOT-MFMane (0) ((.(Fr))))
[1 iz.." 11 (
(13)
We note that the case of fLAND dealt with similarly. Additionally, different choices of the joint distri- bution g. beyond Optimal Transport, can be adapted from CFM [TTng et al., 2023b] to MFM in eq. (7).
4.2 Understanding MFM through energies
In MFM we learn interpolants x; that are optimal according eq. (6). Previous works have studied the "oottimaity"" of interpolants, but have ignored the data manifold. Shaul et al. [2023] proposed to learn interpolants x; that minimize the kinetic energy K(++)) = E2, (20.21)) || 1l2 However, IC assigns each point in space the same cost, independent of the data, and leads to straight interpolants that may stray away from D. Conversely, our objective in eq. (6) can equivalently be written as
(n) = E Jug [ (It.n)) = Ee, )>gg || 112 (
As result, the geodesic loss (T) = Kg ()) is precisely the kinetic energy of vector fields with respect g and hence accounts for the cost induced by the data. In fact, choosing G(x) I, recovers

---

X. We note that the parameterization 2t.n in eq. (4) is more expressive than 2; = a(t)=1 + b(t) which is studied Albergo and Vanden- [2022], Shaul al. [2023]. Besides, 21,, not only depends on the endpoints 20 but, implicitly, on all the data points 'D through metrics such as yame-
K. have been considered in Neklyudov et al. [2023a], Liu et al. [2024], Neklyudov et al. [2023b]. In particular, adapting GSBM [Liu et al., 2024] from the stochastic Schrödinger bridge setting CFM, entails designing interpolants Ie that minimize (( =z) = ((,;) + Vz() with V potential enforcing additional constraints. However, GSBM does not preseribe a general recipe for Vi and instead leaves to the modeler the task of constructing Vg. based on applications. Conversely, MFM relies Riemannian approach propose an explicit objective, i.e. eq. (11), that holds irrespective of the task and can be learned prior to regressing the vector field U0 In fact, eq. (11) can be rewritten as
Data-dependent potentials. Energies more general than the kinetic
y (n) = (x.,.z)))~~ (, + Vi,n(Fe.m 2o, x1)]
(14)
where Vy,,, is parametric function depending on the boundary points To. I (see Appendix $C.1 for an expression). In contrast GSBM, MFM designs interpolants 2y,, that jointly minimize K. and poten- tial Vi,,, that is not fixed but also updated with the same parameters to bend paths towards the data.
5 Experiments
We test METRIC FLOW MATCHING different tasks: artificial dynamic reconstruction and naviga- tion through LiDAR surfaces $5.1; unpaired translation between classes in images $5.2; reconstruction of cell dynamics. Further results and experimental details can be found in Appendices $D, SE and $F.
The model. all the experiments, we test the OT-MFM method detailed in $4.1. As argued in $4.1, for high-dimensional data, we lever- age the RBF metric (10) and hence train with objective (13). We refer to this model as OT-MFM Conversely, for low-dimensional data we rely on the LAND metric (9), replacing Parr with gLasp in eq. (11) and eq. (13). We de- note this model as OT-MFM Both metrics task- independent and do not require further manipulation from the modeler. To assess the impact of metric learning in MFM even without using Optimal Transport for the coupling we tested MFM with independent coupling the Arch task and the single cell datasets, denoted I-MFM and I-MFM
OT-CFM
OT-MFMM LAND
Baselines. MFM generalizes CFM by learning interpolants that account for the geometry of the data. As such, in $5.1 and $5.2, we focus on validating that OT-MFM leads more mean- ingful matching than its Euclidean counterpart OT-CFM [Tong et al., 2023b]. For single-cell periments instead, we also compare with a vari- ety of baselines, including models specific the single-cell domain [TTnn et al., 2020, Koshizuka and Sato, 2023], Schrödinger Bridge mod- els [De Bortoli et al., 2021, Shi et al., 2023, Tong al., 2023a], and ODE-flow methods [Tong et al., 2023b, Neklyudov et al., 2023b].
R1H
Figure 2: Interpolants for the Arch dataset (top row). Sphere datasct (middle row) and over LiDAR scans of Rainier (bottom row). all cases, learning in- terpolants that minimize the LAND metric (9) lcads to more meaningful matchings.
5.1 Synthetic experiments and LiDAR
Our first goal is to validate that MFM and specifically the family of interpolants It.n in eq. (4), are expressive enough to model matching of distributions when the data induce a nonlinear geometry.

---

To begin with, we consider the Arch dataset in Tong et al. [2020]. where we have access to the underlying paths, In Table 1 we report the Wasserstein distance (EMD) be- tween the interpolation according to true dynamics and the marginal at time 1/2 obtained using the Euclidean base- line OT-CFM and our Riemannian model OT-MFM We see that OT-MFM is able to faithfully reconstruct the dynamics by learning interpolants that minimize the geodesic cost induced by the metric in eq. (9). Conversely, straight interpolants fail to stay on the manifold generated by the dynamics (see Figure 2).
Table 1: Wasserstein distance be- tween reconstructed marginal at time 1/2 and ground-truth.
Method
EMD (4)
I-CFM
0.6120 + 0.014
OT-CFM
0,6081 + 0.023
I-MFM
0.1210 + 0.020
OT-MFMMAND
0.0813 + 0.009
To further showcase the ability of MFM to learn trajectories that stay close to an unknown underlying manifold without any added noise, we conduct an experiment on 2D sphere embedded in R.. We consider source and target distributions that are Gaussians centered at the sphere's poles, with samples lying exactly the sphere. MFM significantly improves over the Euclidean baseline, CFM, with samples at intermediate times being much closer to the underlying sphere than those from the Euclidean counterpart. This improvement is achieved without explicitly parameterizing the diversdimensoonal space, relying solely on the LAND metric. Quantitative results, including EMD and the distance of the interpolating paths to the sphere, are reported in Appendix G.
Finally, to highlight the ability of MFM to generate meaningful matching, we also compare the interpolants of OT-CFM and OT-M for navigations through surfaces scanned by LiDAR [Legg and Anderson, 2013]. From Figure 2 we find that straight paths result unnatural trajectories, whereas OT-MFM manages to construct meaningful interpolations that better navigate the complex surface. While OT-MFM can be further enhanced using potentials specific to the task, similar to [Liu et al., 2024], here we focus on its ability to provide meaningful matching even with task-agnostic metric.
5.2 Unpaired translation in latent space
To test the advantages of MFM for more meaningful generation, consider the task of unpaired image translation between dogs and cats in AFHQ [Choi et al., 2020J. Specifically, we perform unpaired translation in the latent space of the Stable Diffusion v1 VAE [Rombach et al., 2022]. Figure 3 reports a qualitative comparison between OT-CFM and OT-MFMas-- Additionally, a quantitative comparison can be found in Table 2, where measure both the quality of images via Fréchet Inception Distance (FID) [Heusel et al., 2017] and the perceptual similarity of generated cats source dogs via Learned Perceptual Image Patch Similarity (LPIPS) [Zhang al., 2018]. We see that OT-MFM improves upon the Euclidean baseline. Our results using MFM further highlight the role played by the nonlinear geometry associated with latent representations, topic studied extensively $6.
OT-CFM
OT-MFMREP 15000
0,000
t=].0
t=000
1=1000
Table FID (L) and LPIPS (L)
values on AFHQ for OT-CFM and OT-M
Method
FID
LPIPS
OT-CFM
41.42
0.512
OT-MFM
37.87
0.502
Figure 3: Qualitative comparison for image translation. By designing inter- polants on the data manifold, OT-MFM better preserves input features.
8

---

5.3 Trajectory inference for single-cell data
We finally test MFM for reconstructing cell dynamics, a central problem in biomedical ap- plications [Lähnemann et al., 2020], which holds great promise thanks the advancements of single-cell RNA sequencing (scRNA-seq) [Macosko et al., 2015, Klein et al., 2015].
Since scRNA-seq trajectories cannot be tracked, we only assume access to K unpaired distributions describing cell populations at K' time points. We then apply the matching objective in eq. (7) between every consecutive time points, sharing parameters for both the vector field v1,0 and the interpolants Tt,0- -se eq. (20) for how to extend Ft.y to multiple time- points. Following the setup of Schiebinger et al. [2019], Tong et al. [2020, 2023a] we
Table 3: Wasserstein-1 distance averaged over left- out marginals for 100-dim PCA single-cell data for corresponding datasets. Results averaged over 5 runs.
Method
Cite (100D)
Multi (100D)
SF M-Gen
44.498 + 0,416
52.203 1.957
SF² M-Fxact
46.530 + 0.426
52.888 + 1.986
OT-CFM
45,393 +: 0.416
54.814 +: 5,858
I-CFM
48.276 3.281
S7.262 1: 3.855
WLF-OT
44.821 0.126
55.416 6.097
perform leave-one-out interpolation, where we measure the Wasserstein-1 distance between the k-th left-out density and the one reconstructed after training on the remaining timepoints. We compare OT-MFM and baselines over Embryoid body (EB) [Moon et al., 2019], and CITE-seq (Cite) and Multiome (Multi) data from [Lance et al., 2022]. In Table 4 and Table 3 we consider the first 5 and 100 principal components of the data, respectively-resut with 50 principal components can be found in Appendix D. We observe that OT-MFM significantly improves upon its Euclidean counterpart OT-CFM, which resonates with the manifold hypothesis for single-cell data [Moon al., 2018]. In fact, OT-MFM surpasses all baselines, including those that add biases such as stochasticity (SF2 Tong et al. [2023a]) or mass teleportation (WLF-UOT Neklyudov al. [2023b]). OT-MFM instead relies on metrics such as LAND and RBF to favor interpolations that remain close to the data.
WLF-UOT
43.731 + 1.375
54.222 5.827
WLF-SB
46.131 + 0.083
55.065 + 5.499
I-MFM
45.987 1 4.014
54.197 +: 1.408
OT-MFMaur
41.784 + 1.020
50,906 + 4.627
Related Work
Geometry-aware generative models. The manifold hypothesis [Bengio et al., 2013] has been studied in the context of manifold learning [TTnennbumm et al., 2000, Belkin and Niyogi, 2003] and metric learning [Xing et al., 2002, Weinberger and Saul, 2009, Hauberg et al., 2012]. Recently, this has also
Table 4: Wasserstein-1 distance (L) averaged over left-out marginals for 5-dim PCA representation of single-cell data for corresponding datasets. Results are averaged over 5 independent
Method
Cite
EB
Multi
Reg. CNF [Finlay et al., 2020|
0.825 0.429
TrajectoryNet [Tong al., 2020]
0.848
NLSB (Koshizuka and Sato, 2023]
0.970
DSBM [Shi et al., 2023]
1.705 + 0.160
1.775 0.429
1.873
0.631
DSB [De Bortoli al., 2021]
0.953 0.140
0.862 0.023
1.079 0.117
SF* M-Sink [Tong et al., 2023a]
1.054 0.087
1.198 + 0.342
1.098 0.308
SE² M-Geo [Tong et al., 2023a]
1.017 + 0.104
0.879 + 0.148
1,255 0.179
SE2 M-Exact ['Tong al., 2023a]
0.920 + 0.049
0.793 0.066
0.933 0.054
OT-CHM ['long et al., 2023h]
0.882 + 0.058
0.790 0.068
0.937
0.054
I-CFM [Tong al., 2023b]
0.965 + 0.111
0.872 0.087
1.085 0.099
SB-CFM [Tong al., 2023b]
1.067 + 0.107
1.221 0.380
1.129 0.363
WLF-UUT [Neklyudov al., 2023b]
0.733 0.063
0.738 0.014
0.911
0.147
WLF-SB [Neklyudov et al., 2023b]
0.797 0.022
0.746 0.016
0.950
0.205
WLF-OT [Neklyudov et al., 2023b]
0.802 0.029
0.742 0.012
0.949 0.211
I-MHM MF-
0.916 0.124
0.822
0.042
1.053
0.095
0.724 0.070
0.713 0.039
0.890 0.123

---

been analyzed in relation to generative models for obtaining meaningful interpolations [Arvanitidis et al., 2017, 2021, Chadebec and Allassonniére, 2022], diagnosing model instability [Cornish et al., 2020, LLoaiz--Ganen et al., 2022, 2024], assessing the ability to perform dimensionality reduction [Stanczuk et al., 2022, Pidstrigach, 2022] and for the improved learning and representation of curved, low-dimensional data manifolds [Dupont et al., 2019, Schonsheck et al., 2019, Horvat and Pfister, 2021, Yonghyeon et al., 2021, De Bortoli, 2022, Jang et al., 2022, Lee et al., 2022, Lee and Park, 2023, Nazari et al., 2023]. Closely related is the extension of generative models to settings where the ambient space itself is a Riemannian manifold [Mathieu and Nickel, 2020, Lou et al., 2020, Falorsi, 2021, De Bortoli et al., 2022, Huang et al., 2022, Rozen et al., 2021, Ben-Hamu et al., 2022, Chen and Lipman, 2023, Jo and Hwang, 2023]. Particularly, Maoutsa [2023] utilized data-dependent geodesic solver [Arvanitidis et al., 2019] to refine drift estimation in stochastic differential equations. Several studies extended beyond the standard linear matching process to meet specific task requirements, but have not accounted for the geometry that the data naturally forms [Liu et al., 2024, Bartosh et al., 2024, Neklyudov et al., 2023b].
Trajectory inference. Reconstructing dynamics from cross-sectional distributions [Hashimoto et al., 2016, Lavenant et al., 2021] is an important problem within the natural sciences, especially in the context of single-cell analysis [Macosko et al., 2015, Moon et al., 2018, Schiebinger et al., 2019]. Recently, diffusion and Continuous Normalizing Flow (CNFs) based methods have been proposed [Tong al., 2020, Bunne al., 2022, 2023, Huguet et al., 2022, Koshizuka and Sato, 2023] but require simulations, whereas Tong et al. [2023a], Neklyudov et al. [2023a], Palma et al. [2023] allow for simulation-frs -free training. In particular, Huguet et al. [2022], Palma et al. [2023] regularize CNFs in a latent space to enforce that straight paths correspond to interpolations on the original data manifold. Finally, Scarvelis and Solomon [2023] propose solution the trajectory inference problem for cellular data, which depends on a regularization of vector fields with respect to a learned metric similar to eq. (6). Crucially though, their framework requires simulations in training and is not immediately extended to more general matching objectives and applications as for MFM. In this regard, we observe that one can also adopt the metric-learning scheme of Scarvelis and Solomon [2023] in MFM, replacing ILAND and Faur introduced in $4.
7 Conclusions and Limitations
We have presented METRIC FLOW MATCHING, a simulation-free framework that generalizes Conditional Flow Matching design probability paths whose support lies on the data manifold. In MFM, this is achieved via interpolants that minimize the geodesic cost of a data-dependent Riemannian metric. We have empirically shown that instances of MFM using prescribed task-agnostic metrics, surpass Euclidean baselines, with emphasis on single-cell dynamics reconstruction. While the universality of the metries proposed in $4 is benefit, we have not investigated how to further encode biases into the metric that are specific to the downstream task-a a topic reserved for future work. Additionally, the principle of learning interpolants that minimize a geodesic cost can also be adapted to score-based generative models such as diffusion models, beyond CFM. When relying on the OT coupling, standard limitations of using OT for high- -igeesinenii problems with large datasets may arise. Lastly, our approach requires the data to be embedded in Euclidean space for the interpolants to be defined; it is an interesting direction to explore how one can learn interpolants that minimize a data-dependent metric even when the ambient space itself is not Euclidean.
Acknowledgements
KK is supported by the EPSRC CDT in Health Data Science (EP/S02428X/1). PP and LZ are supported by the EPSRC CDT in Modern Statistics and Statistical Machine Learning (EP/S023151/1) and acknowledge helpful discussions with Yee Whye Teh. AJB is partially supported by an NSERC Post-doc fellowship and an EPSRC Turing AI World-Leading Research Fellowship. FDG is supported by EPSRC Turing AI World-Leading Research Fellowship. MB is partially supported by EPSRC Turing AI World-L Leading Research Fellowship No. HP/X040062/1 and EPSRC AI Hub on Mathematical Foundations of Intelligence: An "Erlangen Programme" for AI No. EP/Y028872/1
10

---

References
M. S. Albergo and E. Handen-Eijnded Building normalizing flows with stochastic interpolants. arXiv preprint arXiv:2209 /5571, 2022. (Cited on pages 1, 3, and 7)
L. Ambrosio, N. Gigli, and Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005, (Cited on page 4)
G. Arvanitidis, L. K. Hansen, and S. Hauberg. A locally adaptive normal distribution. Advances in Neural Information Processing Systems, 29, 2016. (Cited on pages 6 and 20)
G. Arvanitidis, L. K. Hansen, and S. Hauberg. Latent space oddity: on the curvature of deep generative models. arXiv preprint arXiv: 1379, 2017. (Cited on page 10)
G. Arvanitidis, Hauberg, Hennig, and M. Schober. Fast and robust shortest paths on manifolds learned from data. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 1505-1515. PMLR, 2019, (Cited on page 10)
G. Arvanitidis, S. Hauberg, and B. Schölkopf. Geometrically enriched latent spaces. In International Conference on Artificial Intelligence and Statistics, pages 631-639. PMLR, 2021. (Cited on pages 6, 10, and 20)
G. Arvanitidis, M. González-Duque, A. Pouplin, D. Kalatzis, and S. Hauberg. Pulling back infor- mation geometry. 25th International Conference on Artificial Intelligence and Statistics, 2022. (Cited on page
G. Bartosh, D. Vetrov, and C. A. Naesseth. Neural flow diffusion models: Learnable forward process for improved diffusion modelling. arXiv preprint arXiv:2404. /2940, 2024. (Cited on page 10)
M. Belkin and P. Niyogi. Laplacian eigenmaps for dimensionality reduction and data representation. Neural computation, 15(6):1373-1333 2003. (Cited on pages and 9)
H. Ben- Ben-Hamu, S. Cohen, J. Bose, B. Amos, A. Grover, M. Nickel, R. T. Chen, and Lipman. Matching normalizing flows and probability paths on manifolds, arXiv preprint XXix:2207.04711 2022. (Cited on page 10)
Y. Bengio, A. Courville, and P. Vincent. Representation learning: A review and new perspectives. IEEE transactions on pattern analysis and machine intelligence, 35(8):1798-1888 1828, 2013. (Cited on pages and 9)
C. Bunne, L. Papaxanthos, A. Krause, and M. Cuturi. Proximal optimal transport modeling of population dynamics. In International Conference on Artificial Intelligence and Statistics, pages PMLR, 2022. (Cited on page 10)
C. Bunne, Y.-P. Hsieh, M. Cuturi, and Krause. The schrödinger bridge between gaussian measures has a closed form. In International Conference on Artificial Intelligence and Statistics, pages 5802 5833. PMLR, 2023. (Cited on page 10)
C. Chadebec and Allassonnière. geometric perspective on variational autoencoders. Advances in Neural Information Processing Systems, 35:19618-1966 2022. (Cited on pages 4 and 10)
R. T. Chen and Y. Lipman. Riemannian flow matching on general geometries. arXiv preprint arXiv:2302.03660, 2023. (Cited on pages 2, 3, 5, 10, and 19)
R. Т. Chen, Y. Rubanova, J. Bettencourt, and D. K. Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018 (Cited on page 3)
Y. Choi, Y. Uh J. Yoo, and J.-W. Ha. Stargan v2: Diverse image synthesis for multiple domains. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pages (188-8197, 2020. (Cited on pages 8 and 22)
R. Cornish, A. Caterini, G. Deligiannidis, and A. Doucet. Relaxing bijectivity constraints with continuously indexed normalising flows. In International conference on machine learning, pages 2133-2143 PMLR, 2020. (Cited on page 10)

---

V. De Bortoli. Convergence of denoising diffusion models under the manifold hypothesis. arXiv preprint arXiv:2208.0531/ 2022. (Cited on page 10)
V. De Bortoli, J. Thornton, Heng, and A. Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 34: 17695- 11709, 2021. (Cited on pages and 9)
V. De Bortoli, E. Mathieu, M. Hutchinson, J. Thornton, W. Teh, and A. Doucet. Riemannian score- based generative modelling, Advances in Neural Information Processing Systems, 35:2406-2422, 2022. (Cited on page 10)
Dhariwal and A. Nichol. Diffusion models beat gans on image synthesis. Advances neural information processing systems, 34:8780-877 8794, 2021. (Cited on page 21)
E. Dupont, A. Doucet, and Y. W. Teh. Augmented neural odes. Advances in neural information processing systems, 32, 2019. (Cited on page 10)
L. Falorsi. Continuous normalizing flows on manifolds. arXiv preprint arXiv:2104. 14959, 2021. (Cited on page 10)
C. Finlay, J. Jacobsen, I.. Nurbekyan, and A. Oberman. How to train your neural ode: the world of jacobian and kinetic regularization. International conference on machine learning, pages 1164-3164. PMLR, 2020. (Cited on page
T. Hashimoto, D. Gifford, and T. Jaakkola. Learning population- diffusions with generative mnns. International Conference on Machine Learning, pages 2417-2426. PMLR, 2016. (Cited on pages and 10)
S. Hauberg, O. Freifeld, and M. Black. A geometric take on metric learning. Advances in Neural Information Processing Systems, 25, 2012. (Cited on pages 2, 4, and 9)
A. Hay, L. Kennedy-Shaffer, Shaffer, S. Kanjilal, N. Lennon, S. B. Gabriel, M. Lipsitch, and M. Mina. Estimating epidemiologic dynamics from cross-sectional viral load distributions. Science, 373 r6b52)ee 2021. (Cited on page 1)
Hennig and S. Hauberg. Probabilistic solutions differential equations and their application to riemannian statistics. In Artificial Intelligence and Statistics, pages 347-355. PMLR, 2014. (Cited on page 4)
M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter. Gans trained by two time-scale update rule converge to local nash equilibrium. Advances in neural information processing systems, 30, 2017. (Cited on pages and 22)
C. Horvat and J.PP. Pfister. Denoising normalizing flow. Advances in Neural Information Processing Systems, 34:9099-9111, 2021. (Cited page 10)
C.WW. Huang, M. Aghajohari, J. Bose, P. Panangaden, and A. C. Courville. Riemannian diffusion models. Advances Neural Information Processing Systems, 35:2750-2761, 2022. (Cited on page 10)
G. Huguet, D. S. Magruder, A. Tong, O. Fasina, M. Kuchroo, G. Wolf, and S. Krishnaswamy. Manifold interpolating optimal-transport flows for trajectory inference. Advances in neural information processing systems, 35:29705-29718. 2022. (Cited on page 10)
C. Jang, Y. Lee, Y.K. Noh, and F. C. Park. Geometrically regularized autoencoders for non-euclidean data. In The Eleventh International Conference on Learning Representations, 2022. (Cited on page 10)
J. Jo and J. Hwang. Generative modeling manifolds through mixture of riemannian diffusion processes. arXiv preprint arXiv:2310.07216, 2023. (Cited on page 10)
D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint 2014. (Cited pages 21 and 22)
D. P. Kingma and M. Welling, Auto-encoding variational bayes. ICLR, 2014. (Cited page 22)
12

---

A. M. Klein, L. Mazutis, I. Akartuna, N. Tallapragada, A. Veres, V. Li, L. Peshkin, D. A. Weitz, and M. W. Kirschner. Droplet barcoding for single-cell transcriptomics applied to embryonic stem cells. Cell, 161(5)11177-2211 2015. (Cited on pages and 9)
T. Koshizuka and I. Sato. Neural lagrangian schrödinger bridge: Diffusion modeling for population dynamics. In International Conference on Learning Representations. 2023. (Cited on pages 7, 9, and 10)
D. Lähnemann, J. Köster, E. Szczurek, D. McCarthy, S. Hicks, M. D. Robinson, C. A. Vallejos, K. R. Campbell, N. Beerenwinkel, A. Mahfouz, et al. Eleven grand challenges in single-cell data science. Genome biology, 21:1-35, 2020. (Cited on page
C. Lance, M. Luecken, D. B. Burkhardt, R. Cannoodt, Rautenstrauch, A. Laddach, A. Ub- ingazhibov, Z.-J. Cao, K. Deng. S. Khan, et al. Multimodal single cell data integration challenge: results and lessons learned. BioRxiv, pages 2022 04, 2022. (Cited on pages and 22)
H. Lavenant, S. Zhang, Y.-H. Kim, and G. Schiebinger. Towards a mathematical theory of trajectory inference. arXiv preprint wrXiv:2102.0990 2021. (Cited pages 1, 3, and 10)
J. M. Lee. Smooth manifolds. Springer, 2012. (Cited on page 16)
Y. Lee and F. C. Park. On explicit curvature regularization in deep generative models. In Topological, Algebraic and Geometric Learning Workshops 2023, pages 505-518. PMLR. 2023. (Cited on page 10)
Y. Lee, S. Kim, Choi, and F. Park. A statistical manifold framework for point cloud data. In International Conference Machine Learning, pages 12378-12402, PMLR, 2022. (Cited on page 10)
N. Legg and S. Anderson. Southwest flank of Mt.Rainier wa, 2013. URL https:://oi.org/// 5069/G9PZ56R1 Accessed: 2024-05-21. (Cited on pages and 21)
C. Léonard. A survey of the schr\ " odinger problem and some of its connections with optimal transport. arXiv preprint arXiv:1308. 0215, 2013. (Cited on page 1)
Y. Lipman, R. T. Q. Chen, H. Ben-Hamu. M. Nickel, and M. Le. Flow matching for generative modeling. International Conference on Learning Representations (ICLR), 2023. (Cited on pages 1, 3, and 4)
G.-H. Liu, Y. Lipman, M. Nickel, B. Karrer, E. Theodorou, and R. T. Chen. Generalized schrödinger bridge matching. In ICLR, 2024. (Cited on pages 7, 8, 10, 20, and 21)
X. Liu, C. Gong, and Q. Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprini arXiv:2209.03000 2022. (Cited pages 3, and 4)
G. Loaiza-Ganem, B. L. Ross, C. Cresswell, and L. Caterini. Diagnosing and fixing manifold overfitting in deep generative models. arXiv preprint arXiv:2204.07112. 2022. (Cited on pages 2, 4, and 10)
G. Doaiza-Ganem, B. L. Ross, R. Hosseinzadeh, A. L. Caterini, and J. C. Cresswell. Deep generative models through the lens of the manifold hypothesis: A survey and new connections. arXiv preprint arXiv:2404.02954, 2024. (Cited on pages and 10)
I. Loshchilov and F. Hutter. Decoupled weight decay regularization. arXiv preprint arXiv: 2017. (Cited pages 21 and 22)
A. Lou, D. Lim, I. Katsman, L. Huang, Q. Jiang, S. N. Lim, and C. M. De Sa. Neural manifold ordinary differential equations. Advances in Neural Information Processing Systems, 33:17548- 17558, 2020. (Cited on page 10)
E, Macosko, A. Basu, R. Satija, Nemesh, K. Shekhar, M. Goldman, Tirosh, A. R. Bialas, N. Kamitaki, E. M. Martersteck, et al. Highly parallel genome-wide expression profiling of individual cells using nanoliter droplets. Cell, 16165) 2015. (Cited on pages 1, 3, 9, and 10)
13

---

D. Maoutsa. Geometric constraints improve inference of sparsely observed stochastic dynamics. arXiv preprint 2rXiv:2304.00423, 2023. (Cited on page 10)
E. Mathieu and M. Nickel. Riemannian continuous normalizing flows. Advances in Neural Informa- fion Processing Systems, 33:2503-225 2510,, 2020. (Cited page 10)
M. G. Monera, A. Montesinos-Amillbe and E. Sanabria-Codeste The taylor expansion of the exponential map and geometric applications. Revista de la Real Academia de Ciencias Exactas, Fisicas y Naturales. Serie A. Matematicas, 108:881-906, 2014, (Cited on page 19)
K. R. Moon, Stanley II, Burkhardt, D. van Dijk, G. Wolf, and S. Krishnaswamy. Manifold baanid methods for analyzing single-cell ma-sequencing data. Current Opinion in Systems Biology, 7:36-46, 2018, (Cited on pages 2, 4, 9, and 10)
K. R. Moon, D. Van Dijk, Z. Wang. S. Gigante, D. B. Burkhardt, W. S. Chen, K. Yim, A. v. d. Elzen, M. J. Hir, R. R. Coifman, et al. Visualizing structure and transitions in high-dimensional biological data. Nature biotechmology. 37(12):1482-1492 2019. (Cited on pages and 22)
P. Nazari, S. Damrich, and F. A. Hamprecht. Geometric autoencoders you see is what you decode. arXiv preprint arXiv:2306. 76338, 2023. (Cited on page 10)
K. Neklyudov, R. Brekelmans, D. Severo, and A. Makhzani, Action matching: Learning stochastic dynamics from samples. In International Conference Machine Learning, pages 25858-25889, PMLR, 2023a. (Cited on pages 4, 7, and 10)
K. Neklyudov, R. Brekelmans, A. Tong, L. Atanackovic, Q. Liu, and A. Makhzani. A computational framework for solving wasserstein lagrangian flows. arXiv preprint arXiv:2310.10649, 2023b. (Cited on pages 7, 9. 10, and 23)
J. Oeppen and J. W. Vaupel. Broken limits to life expectancy, 2002. (Cited on page 1)
A. Palma, S. Rybakov, L. Hetzel, and F. J. Theis. Modelling single-cell ma-seq trajectories flat statistical manifold. In NeurIPS 2023 Al for Science Workshop, 2023. (Cited on page 10)
J. Pidstrigach. Score-based generative models detect manifolds. Advances in Neural Information Processing Systems, 35:35852-3555 35865, 2022. (Cited on page 10)
Q. Que and M. Belkin. Back to the future: Radial basis function networks revisited. In Artificial intelligence and statistics, pages 1375-1383. PMLR, 2016. (Cited page 6)
R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10695, 2022. (Cited on pages and 22)
N. Rozen, A. Grover, M. Nickel, and Y. Lipman. Moser flow: Divergence-based generative generaaa modeling on manifolds. Advances in Neural Information Processing Systems, 34:17669-17680, 2021. (Cited on page 10)
A. Salmona, V. De Bortoli, J. Delon, and A. Desolneux. Can fush-forward generative models fit multimodal distributions? Advances in Neural Information Processing Systems, 35:10766-10779, 2022. (Cited on pages and 4)
C. Scarvelis and J. Solomon. Riemannian metric learning via optimal transport. In ICLR, 2023. (Cited on page 10)
G. Schicbinger. Reconstructing developmental landscapes and trajectories from single-cell data. Current Opinion in Systems Biology, 27:1 100511 2021. (Cited on page
G. Schiebinger, J. Shu, M. Tabaka, B. Cleary, V. Subramanian, A. Solomon, J. Gould, S. Liu, S. Lin, P. Berube, et al. Opimmllttrasspott analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, 76(44)928-943 2019. (Cited on pages and 10)
Schonsheck, Chen, and R. Lai, Chart auto-encoders for manifold structured data. arXiv preprint 10099 2019, (Cited pages 4 and 10)
14

---

E. Schrödinger. Sur la théorie relativiste de l'électron et I' l'interprétation de la mécanique quantique. In Annales de I' l'institut Henri Poincaré, volume 2, pages 269-310, 1932. (Cited on page
N. Shaul, R. T. Chen, M. Nickel, M. Le, and Y. Lipman. On kinetic optimal probability paths for generative models. In International Conference Machine Learning, pages 30883-30907. PMLR, 2023. (Cited on pages 4, 6, and 7)
Y. Shi, V. De Bortoli, A. Campbell, and A. Doucet. Diffusion schrödinger bridge matching. In seirth- Conference on Neural Information Processing Systems, 2023. (Cited on pages 7 and
Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. International Conference on Machine Learning (ICML), 2015. (Cited on page 1)
Y. Song. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score-based gener- ative modeling through stochastic differential equations. International Conference on Learning Representations (ICLRR, 2021. (Cited on page I)
Stanczuk, G. Batzolis, T. Deveney, and C.-B. Schönlieb. Your diffusion model secretly knows the dimension of the data manifold. arXiv preprint arXiv:2212.1266 2022. (Cited on page 10)
J. B. Tenenbaum, V. d. Silva, and C. Langford. A global geometric framework for nonlinear dimensionality reduction. science, 290(5500):3319-2223 2000. (Cited on pages 2 and 9)
A. Tong, J. Huang, G. Wolf, D. Van Dijk, and S. Krishnaswamy. Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. International conference on machine learning. pages
PMLR, 2020. (Cited on pages 7, 8, 9, 10, 21, and 22)
A. Tong, N. Malkin, K. Fatras, L. Atanackovic, Y. Zhang, G. Huguet, G. Wolf, and Y. Bengio. Simulation-Free schrödinger bridges via score and flow matching, 2023a. (Cited on pages 7, 9, 10, and 22)
A. Tong, N. Malkin, G. Huguet, Y. Zhang, J. Brotor-Brooos, K. Fatras, G. Wolf, and Y. Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. arXiv preprint ar 0tt:23022.0448 2023b. (Cited pages I, 3, 4, 7,9, and 22)
C. Villani et al. Optimal transport: old and new, volume 338. Springer, 2009. (Cited page 6)
K. Q. Weinberger and L. K. Saul. Distance metric learning for large margin nearest neighbor classification. Journal of machine learning research, 10(2), 2009. (Cited on page 9)
E. Xing, M. Jordan, J. Russell, and A. Ng. Distance metric learning with application to clustering with side-information. Advances in neural information processing systems, 15, 2002. (Cited on pages and 9)
L. Yonghyeon, S. Yoon, M. Son, and F. C. Park. Regularized autoencoders for isometric representation learning. In International Conference on Learning Representations, 2021. (Cited on page 10)
R. Zhang, P. Isola, A. A. Efros, F. Shechtman, and O. Wang. The unreasonable effectiveness of deep features as a perceptual metric. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 586-595, 2018 (Cited on pages 8 and 22)
15

---

Outline of Appendix
In Appendix A we give brief overview of relevant notions from differential and Riemannian geometry. Appendix B provides more details for Section 3 including the formal statement and proof of Proposition 1. We also derive how to rigorously extend flow matching to the Riemannian manifold induced by data-dependent metric. We report pseudocode for MFM in Algorithm 2. Appendix C provides more details regarding the data-dependent Riemannian metrics we use, and relevant training procedures. In Appendix D we supply more details regarding the various experiments. Appendix E, Appendix F, and Appendix H contain supplementary figures and tables for single-cell reconstruction, unpaired latslation, and LiDAR tasks, respectively. Finally, in Appendix G, outline the quantitative evaluation results for MFM on the sphere experiment.
All exact hyperparameters used and reproducible code can be found at https: //gp:::////i /ttp::/ggg/tt.////eeeee thubthu..cc com/ tkaniakk/etric- wlww-watchhie .iit.
A Primer on Riemannian Geometry
Riemannian geometry is the study of smooth manifolds M equipped with a (Riemannian) metric g- Intuitively, this corresponds to spaces that can be considered locally Euclidean, and which allow for a consistent notion of measuring distances, angles, curvature, shortest paths, etc. on abstract geometric objects. We provide primer on the relevant concepts of Riemannian geometry for work, covering smooth manifolds and their tangent spaces, Riemannian metrics and geodesics, and integration on Riemannian manifolds. For a comprehensive introduction, see Lee [2012].
Smooth manifolds, Formally, we say a topological space" . is d-dimensional smooth manifold if we have collection of charts (U,, PP)) where Us open subsets of M and U,U; = M, pp: U, > R are homeomorphisms onto their image, and when U2 n UJ A 0, the transition maps Py :p (u)(n oo (U, nUj) diffcomorphisms. These charts allow to represent quantities on . through the local coordinates obtained by mapping back to Euclidean space via PP well allowing us define smooth maps between manifolds.
In particular, we define smooth path in M passing through æ E M as a function y (-c,c) + M, for some 0 and 7(0) x, such that the local coordinate representation (o; 7 of ry is smooth in the standard sense (for any suitable chart (U,, pp))). The derivatives +(0) of smooth paths y passing through x E M form a d-dimensional vector space called the tangent space T,M at x. We can represent T&M in local coordinates with yi by the identification of +(0) E T.M to Rd where (a) ..., aa(())
pp; op: r()).
Riemannian structure. To introduce geometric notions of length and distances to smooth manifolds, define Riemannian metric g M as map providing smooth assignment of points M on the manifold to a positive definite inner product ((())))))) defined on the corresponding tangent space T.M. In local coordinates about E M given by a suitable chart (U,, spt). we have the local representation (v, w))(a; e' G(x) w where Gxx) E Raxd is a positive definite matrix (which implicitly depends on the choice of chart). We call the pair (M, .) a Riemannian manifold.
Through the Riemannian metric, we can now define the norm of tangent vectors by ||e||a()) =
(v, v) 1/2 for u E Ty,,, as well as the length of smooth path y [0. 1] M by g(x)
Len(>) I'' (1((t)ll(sso))d
(15)
We can then define a geodesic 7 between zo and 21 in M as the shortest possible path between the two points i.e.
7" = arg min [' (11(0)) di, 7(0) zo, 7(1) x1-
(16)
Witt the technical assumptions that M second-countable Hausdorff. Formalll, this is choice of a smooth symmetric covariant 2-tensor field on M which is positive definite at each point.
16

---

We assume that all Riemannian manifolds being considered are (geodesically) complete. meaning that geodesics can be extended indefinitely. In particular, for any pair of points 20: 21, there exists unique geodesic y starting at To and ending at 211-
Integration Riemannian manifolds. To introduce integration Riemannian manifolds, we the fact that under the technical assumption that M is orientable, the Riemannian manifold (M, g) has canonical volume form dvol. This can be used define measure on M, where in local coordinates, we have that dvvl(z) //G(a) where de denotes the Lebesgue measure in Rd and denotes the determinant. Hence, for chart (U,, (P:) and continuous function f .M- R which is compactly supported in U,, we define the integral
J fdvol= = [ Sexcus) fo (x)
(17)
This definition can be extended to more general functions through the use of partitions of unity and the Riesz-Markor-Kakutaa representation theorem.
The Riemannian geometry of Metric Flow Matching. Finally, we note that in our work, we take our smooth manifold to be M = ER with the trivial chart U; = R : = id. This means we can work in the usual Euclidean coordinates instead of requiring charts and local coordinates, simplifying our framework. For example, we can define through G directly in Definition I without the need to check consistency across the choice of local coordinates, and we have the trivial identification of M to Rd. In addition, for continuous f' M - IR, we have the simplified definition of the Riemannian integral (when the integral exists) as
[., fdvol= [ f(x) (I(
(18)
We use this to define probability densities on (M,g) to extend CFM to our setting in $B.1.
B Additional Details for Section
To begin with, we provide a formal statement covering Proposition 1 and report proof below. We introduce some notation. We write " C U when the image set of y [o, 1] + Idd is contained in U, iee. 7/ E U for each & E [0, 1]. Moreover, we let
BD( == fy E Rd 3x, E D.t. lly =
denote the set of points in Rd whose distance from the dataset DD is most r > Theorem B.1 (Formal statement of Proposition 1). Consider closed dataset DD C TRdd (e.g. finite). Assume that for each (Iu, ~ g. there exists at least path T connecting Zo to æ1 whose length is at most I and such that 7/t C (D). with 5 0. Let K > 0.p > and g be any data-dependers metric sattisfying: vi)
TT)TTG(( G(x; D)v 2 ri|lu||? for each x E Id (D) and v E Rd: (ii) (x; D)Il S <(p/T)² for each x E BD))). Then for any (zo, 21) q. the geodesic of of g connecting To to Di satisfies 72 C (D).
Proof of Theorem B.I. We argue by contradiction and assume that there exist x E H (wo, x1) g such that the geodesic 7t of g connecting xo to Ij passes through z, i.c. there is time to E (0, 1) such that vi = x.
B2o(DD and
By assumption, there exists a path 7 that connects zo 24 and stays within (D),), which means that To = 70 E Bs (D). Since D is closed, the function x + dæ(æ, D) == inf | y/1, i.e. the Euclidean distance of x from the dataset D, is continuous. In particular, is also continuous, due to the geodesic being a smooth function in the interval (0, 1]. From the
de. (l) == dm( D)
continuity of d and the fact that dg, (0) S d and de, > 2p, it follows that there must be time 0 < t' < to such that (t) 2 for all t E (t', tol and dt, =
we unpack the definition of dg,p, we have just shown that of E Rd \ B2( for all t E (V' (o]. Accordingly, we can estimate the length of yf' with respect to the Riemannian metric g as
Leng(ri) = 1.' (nVlla(t))) dt >
1.a0 1 (002)) dt
1/ (( dt 2 Vk 1 11 2 Valne -7%
17

---

where in the second-to-lase inequality we have used the assumption (i) on the metric having minimal eigenvalue / in I (D),) while the final inequality simply follows from the Euclidean length of any curve between the points y and "to being larger than their Euclidean distance.
We claim that | - %i. l| 2 p. To validate the latter point, take 0. follows that there exists I2 D such that ||.; - > I| S p e, because dg. = p. i.e. 77 E Bp(D). If the - %i < p.
then we could apply the triangle inequality and derive that
Inio =,1l S I -f + ll zill < p p +
Since > 0 was arbitrary. it would follow that yo B2p which is contradiction to starting point. Therefore, lh % p. and hence
Leng(?)) Vnp.
(19)
On the other hand, the assumptions guarantee the existence of another path "/e connecting z0 to ø1 whose image is always contained in B3 (D). It follows that
Leng - =fo (Hxellatr) di = S/ V(GE)T (H+)TG((;:D))
G(T;; di
S vi f J. [Htll = V
where we have used the assumption (ii) on g and that the length of is at most I. We can then combine the last inequality and eq. (19), and conclude that
(nn,(n)) S Vkp < Leng(())
which means that have found path connecting 2o 11. whose length with respect the metric is shorter than the one of the geodesic ? from To to æ1. This is a contradiction and concludes the proof.
B.1 Extending CFM to the manifold (RR,
In this subsection, we demonstrate how to generalize the CFM objective in eq. (1) to the Riemannian manifold (R .g). defined by a data-dependent metric g as per Definition 1.
Densities on the manifold. To begin with, we extend the densities Po: P1 and the joint density g to valid densities on the manifold (R .g). First, recall that the Riemannian volume form induced by g can be written in coordinates as
dvol(x) = ( D)ldr,
where | denotes the determinant of the matrix Gæ; D), while dr is the standard Lebesgue measure on I Accordingly, given density p E P(Rd), we can derive the associated density á E P(R by rescaling:
p(x)
p(x)
VGGr, DJ
which guarantees that ş is continuous, positive, and
f ((r)dvol(z)) / p(x)dz 1. J
This in particular applies to the marginals Po, Pi given by our problem. A similar argument applies the joint density whose marginals are Po and pr, respectively. In fact, we can now define a density á on the product manifold, i.e. 9 E P(R x Rd .g x g), by simply taking
( x1
((00,11) = / /G D) . G(t;;
18

---

where the denominator is exactly the pointwise volume form induced by the product metric g. Therefore, the MFM objective in (7) is interpreted as
(0) = (ao [I12,(( - =2,0r [
E.
(1.0(1e - N1, I (0 x1) ) (mo;; D). S(ni:D)ddrdda
Jaxxa
S (ur.((ee))))) ,,,07 11 ) 2t )dkode
E,,
( [Io,,((es))) - (1
which is exactly eq. (7), We highlight that for this reason, we slightly abuse notation, and write an expectation with respect to q. regarded as density the manifold (It ,g), the same an expectation over the density g with respect to Lebesgue measure.
The matching objective. We also emphasize that, similarly to Riemannian Flow Matching [Chen and Lipman, 2023]. we normalize the regression objective by G=1/2(z; 'D to avoid the need for initialization schemes that are specific the metric. In fact, introducing high-dimensional, non- trivial norm || ll, in the CFM objective could introduce instabilities in the optimization of CFM. As such, the objective LMFM effectively reduces to LCFM with interpolants ",,,"" minimizing the geodesic loss Lg Equivalently, METRIC FLOW MATCHING consists of a 2-step procedure: we first learn interpolants that approximate the geodesics of data-dependent metric (Algorithm 1), and then we regress the vector field using the interpolants from the first stage as per the standard CFM objective. A pseudocode for the MFM -ii--iieline is reported in Algorithm 2.
Inference. We finally note that, in principle, the differential equation generated by va is defined the tangent bundle of (RR" ..9) and hence solving it through Euler discretization, would require adopting the exponential map associated with the metric g to project the tangent vector to the manifold. However, since the underlying space is still Rd the Euclidean Euler-step integration provides first- order approximation of the Riemannian exponential associated with (see for example Monera et al. [2014]). Accordingly, inference, we can simply rely on the canonical Euler-discrete integration to approximate the exponential map associated with the delendeneddt metric g. provided that the step size is small.
Algorithm 2 Pseudocode for METRIC FLOw MATCHING
Require: coupling q. initialized network v1,,, trained network Pt., data-dependent metric g(-)
1: while Training do
2: Sample (Io, .1) ~ q and ~ 4(001)
3: It. (1 - t) + tx, + +[1 t) (I),
eq. (4)
4:
it.n + 21 - 2o + 1(1 - t) 01) + (1 (t)
5:
E(O) < I|v..0 (It,,) - d (()
& Estimate of objective LMPMM(0) from eq. (7)
6:
Update using gradient (oe
return vector field 02,0
B.1.1 Support of Pt
We conclude this section by justifying our claims about the relation of the support of the proba- bility path Pt and the interpolants. Consider family of conditional paths pe Pa To, x11) such that Pe Pe(z/00 ø1) Re 6(x-xz), where x, are the interpolants connecting Zo to I1, and ō the Dirac distribu- tion. Assume that ur are straight interpolants, i.e. for any (Ir,, .1) g. we define T, lan (1-1)aro Let us now consider y Id such that there is no (zo, 21) g y tz + t) Accordingly:
(v(y) [
Spexed Pr Pr ( x1 (1)g((0 æ (1foo.z))ddo dddddr1 {
8(v x )q( x )xzodx = 0,
JR
since there is no x1) the support of y, such that y = wt. We have then shown that supp(p)) C (v E IRd (2o. g: y tz + (1 f)
which can be limiting for tasks where the data induce nonlinear geometry, validated in Section 5.
19

---

C Additional Details on the Riemannian Metrics used
In this Section, we provide further details on the diagonal metrics introduced Section 4, which we adopt in METRIC FLOw MATCHING. In particular, we comment on important differences between the LAND metric and the RBF metric. We finally derive an explicit connection between MFM and recent methods that learn interpolant minimizing generalized energies.
Learning the RBF metric Dou. For the RBF metric in (10), we follow the metric design of Arvanitidis et al. [2021] and find the centroids by performing k-means clustering. Similarly, we define the bandwidth d associated with cluster C₁ as:
-2
A
K CE
1
Sx 1x
2
rEC₁
where K is tunable hyperparameter. We note that the bandwidth is chosen to assign smaller decay in (10) to the centroids of clusters that have high spread to better enable attraction of trajectories. Finally, the weights Wa,k & determined by training the loss function,
K _1--om
CRBF ((BF)( = Σ (1- ho (ail) = Σ IEDD
-
Wa exp (- Aa.k 10 - -mx1(3)
(
2
IGDD
Two important comments are in order. First, by learning the RBF metric, the framework MFM effectively entails jointly learning data-dependent Riemannian metric and suitable matching objective along approximate geodesics zt,,- Second, we note that more general training objectives can be adopted for h, to enforce specific properties for the metric that are task-aware. We reserve the exploration of this topic for future work.
LAND vs RBF metrics. While similar in spirit, since they both assign a lower cost to regions of space close to the support of the data points D, the LAND metric [Arvanitidis al., 2016] and the RBF metric [Arvanitidis et al., 2021] differ in two fundamental aspects. Crucially, RBF is learned based on the data points, requiring less tuning than LAND, which key advantage high- -ithemmiessinn and the main motivation for why we resort to RBF for experiments on images and single-cell data with more than 50 principal components. Additionally, the RBF metric assigns similar cost to regions with data and is, in principle, more robust than the LAND metric to variations in the concentration of samples in D. While this neither benefit nor downside in general, we observe that if a metric such as LAND consistently assigns much lower cost to regions of space with higher data concentration, then the geodesic objective eq. (6) could always bias to learn interpolants Ft.n moving through these regions independent of the starting point-this is consequence of the fact that we never compute the actual length of the paths to avoid simulations. Note though that this has not been observed an issuc in practice whenever adopted the LAND metric for dineeieniional data.
C.1 Connection between Riemmanian approach and data potentials
We detail how the geometric loss (6) used to leam interpolants Tt.,, be recast as generalization of the GSBM framework in Liu et al. [2024]. In general. for any data-dependent metric g, we can indeed rewrite (n) in eq. (6) as:
(n) = (x0, (ir G( D)
Ee,(20,21)ng ||. 12 + (Jim (G D) I)
E2,(20.21>ng [|.,,]11 1/2 12
x1)]
where the parametric potential Vi,, has the form
Ven To, rz) (it... G (G D) I)
òt,, = 21 - To + t(1 - t) 21) + (1- 2t) 21);
where we have used the parameterization of It.,, eq. (4). By replacing G(r;; D) with the explicit expression given by the RBF metric in eq. (10), this provides concrete formulation for eq. (14). We
20

---

also note that an equivalent perspective amounts to replacing the parametric potential with fixed function Ug (..)) that depends not only on 2t,, as the potentials in Liu al. [2024], but also on the velocities Tt.g-- Once again, we emphasize that our framework prescribes explicit parametric potentials Ve., via the diagonal Riemannian metries Section 3.1 and hence differs from Liu et al.
[2024] which leave to the user the task of designing potentials based on applications.
D Experimental Details
We parameterize the models (t.,, x1) in eq. (4) and (1,0(1r)) in eq. (7) neural networks, and train them separately and sequentially. We start with Pt.,, (Zo, 21), so that interpolants are learned prior regressing v1,0(2,) in the matching objective. Synthetic Arch, Sphere, LiDAR, and single-cell experiments were run on a single CPU. A single run of the LiDAR architecture and 5-dimensional single-cell experiments trains in under 10 minutes, while sigher-dimensional single-cell experiments typically train in under an hour. The unpaired translation experiment on AFHQ was trained on a GPU cluster with NVIDIA A100 and V 100 GPUs. Training time for AFHQ varies by GPU, ranging from 12 hours (A (110)) to day (V100).
D.1 Synthetic Arch, Sphere and LiDAR experiments
In the synthetic Arch, Sphere and LiDAR experiments, we parameterized both (t,, 21) and (1,((ze) as 3-layer MLP networks with a width of 64 and SeLU activation. The networks were trained for up to 1000 epochs with early stopping (patience of 3 based on validation loss). We employed the Adam optimizer Kingma and Ba [2014] with learning rate of 0.0001 for (P..n (Tg. 21) and the AdamW optimizer Loshchilov and Hutter [2017] with a learning rate of 10-3 and weight decay of 10- for v (2.). We used a 90% train/validation split. Training samples served as source and target distributions and for calculating the LAND metric, while validation samples were used for early stopping. We used o = 0.125 and e = 0.001 in the LAND metric. During inference, we solved for p. using the Euler integrator for 100 steps.
Synthetic Arch To generate the Arch dataset, we follow the experimental setup from [Tong et al., 2020]. We sampled 5000 points from two half Gaussians N(, 2x and N N(1, 2 The exact optimal transport interpolant at / 1 was used as the test distribution. We then embedded the points half circle of radius 1 with noise N (0, 0.1) added to the radius. The Earth Mover's Distance between the sampled and test sets was averaged across five independent runs.
Synthetic Sphere To generate the Sphere dataset, we sampled 5,000 points. We sampled the latitudes from two half Gaussians N ((, 2 and N(1, 21, and then scaled them by æ. We sampled longitudes uniformly. The exact optimal transport interpolant at 1/2 was used as the test distribution. We then embedded the points on a sphere of radius without any noise. The Earth Mover's 's Distance between the sampled and test sets, and the mean distance of the middle points of trajectories to the sphere (both reported in Appendix G) were averaged across five independent runs.
LiDAR For LiDAR, the data consists of point clouds within [-5, C R representing scans of Mt. Rainier [Legg and Anderson, 2013]. The source and target distributions are generated as Gaussian Mixture Models and projected onto the LiDAR manifold as in Liu et al. [2024]. We then standardized all the LiDAR, source, and target points.
D.2 Unpaired translation in latent space
In the unpaired translation experiments, we utilized the U-Net architecture setup from Dhariwal and Nichol [2021] for both (To, 21) and (z)) The exact hyperparameters are reported in Table 5. We used the Adam optimizer Kingma and Ba [2014] for both networks and applied early stopping only for (1 31) based on training loss. During inference, we solved for pt / using the adaptive step-size solver Tsit5 with 100 steps. We trained the RBF metric with x = 0.5 and = 0.0001 with k-means clustering, as described in Appendix C. For this experiment, we enforce stronger bending by using (ha (z))², in the loss function Lynag (1) in (11). Note that higher powers of h ensures that regions away from the support of D. where < 1, are penalized even more in the geodesic objective, which highlights the role played by the biases introduced via the metric.
21

---

Our method operates in the latent space of the Stable Diffusion v1 VAE Kingma and Welling [2014], Rombach al. [2022], except for the k-means clustering step in RBF metric pretraining, which was performed in the ambient space.
Table 5: U-Net architecture hyperparameters for unpaired image translation on AFHQ.
(4,,
(r,0(1:))
Channels
128
128
ResNet blocks
Channels multiple
1,
2, 2, 2
Heads
Heads channels
64
64
Attention resolution
16
Dropout
0.1
Batch size
256
256
Epochs
100
8k
Learning rate
le-4
le-4
EMMA-decap
0.9999
0.9999
Dataset We used the Animal Face dataset from Choi et al. [2020], adhering to the splitting predefined by dataset authors for train and validation sets, with validation treated the test set. Standard preprocessing was applied: upsizing to 313x256, center cropping to 256x256, resizing to 128x 128, and using VAE encoders for preprocessing. Finally, we computed all embeddings using pretrained Stable Diffusion v1 VAE Rombach et al. [2022]. FID Heusel et al. [2017] and LPIPS Zhang al. [2018] were computed using the validation sets. FID measured with respect the cat validation set, while LPIPS Zhang et al. [2018] between pairs of source dogs and generated cats.
D.3 Trajectory inference for single-cell data
We performed both low-dimensional and high-dimensiona single-cell experiments following the setups in Tong et al. [2023a.b]. For each experiment, the single-cell datasets were partitioned by excluding an intermediary timepoint, resulting in multiple subsets. Independent models were then trained on each subset. Test metrics were calculated on the left-out marginals and averaged across all model predictions.
We employed the Adam optimizer Kingma and Ba [2014] with a learning rate of 10 for (r.n XI) and the AdamW optimizer Loshchilov and Hutter [2017] with a learning rate of 10-3 and a weight decay of 10-5 for (1z)) During inference, solved for p. t being left-out marginal using the Euler integrator for 100 steps.
We used 90%/10% train/validation split, excluding left-out marginals from both sets. Training samples served as source and target distributions and for calculating the metries, while validation samples were used for early stopping. We note that these settings slightly more restrictive (and realistic) than those reported for SFaM-Gee Tong et al. [2023a], where the left-out timepoint was also included in the validation set.
Embryoid Body dataset We used the Embryoid Body (EB) data Moon al. [2019] preprocessed by Tong et al. [2020], focusing the first five whitened dimensions. The dataset, consisting of five time points over 30 days, was used to train separate models across the full-time scale, each time leaving out one of the time points 1, 2, or 3.
Cite and Multi datasets We utilized the Cite and Multi datasets from the Multimodal Single-cell Integration Challenge at NeurIPS 2022 Lance et al. [2022], preprocessed by Tong et al. [2023a]. These datasets include single-cell measurements from CD4+ hematopoietic stem and progenitor cells for 1000 highly variable genes across four time points (days 2, 3, 4, and 7). We trained separate models each time leaving out one of the time points or The data was whitened only for 5-dimensional experiments.
22

---

Multiple constraints setting Following similar approach to Neklyudov et al. [2023b], we modified our sampling procedure interpolate between two internediate dataset marginals, with neural network parameters shared across timesteps. The interpolation is defined as:
1i+1 - +
1 1 1141 li
li+1 (1-(
(1-
1,+1 1++1 (
1141 (
1 -- -1)
)) (2,0 (20)
Jesss
1i41 4i
Low-Dimensiona experiments We parameterized both fP.n (Zo, F1) and (2,) as 3-layer MLP networks with width of 64 and SeLU activation. In the LAND metric, used 0.125 and = 0.001 for all EB sets and the first leave-out time point in the Cite and Multi datasets. For the second leave-out time point in Cite and Multi, used 0.25 and 0.001.
High-Dimennsonal experiments Both (2,, 21) and (g,0(x)) were parameterized as 3-layer MLP networks with width of 1024 and SeLU activation. We set K 1.5 and be the complement of the final metric pretraining loss to maintain consistent regularization across datasets and leave-out timesteps.
Baselines We reproduced the results reported by Neklyudov et al. [2023b] to ensure consistent reporting of standard deviations and the same versions of EB dataset used across all experiments. The standard deviations were calculated across all leave-out timesteps and seeds for each dataset and dimension. We used the provided code and hyperparameters for training, averaging the results across 5 seeds.
E Supplementary Single-Cell Reconstruction
report additional results for single-cell reconstruction using 50 principal components Table Again, we note that OT-MFM performs strongly, and it is marginally surpassed only by SF² on Multi(50D) (results have not been reproduced). Two important remarks are order. First, the baseline SF2 M-Geo leverages a geodesic cost in the formulation of the Optimal Transport coupling, which enforces similar biases to OT-MFM. In fact, as argued in Section 4, OT-MFM can similarly consider data-aware costs in the formulation of the optimal transport coupling. In this work though, we wanted to focus on the ability of interpolants to lead to meaningful matching in settings where the optimal transport coupling is agnostic of the data support. Additionally, highlight that move to more realistic nigh-dimensional settings (100 principal components, shown in Table 3) the advantages of our framework become even more apparent.
Table Wasserstein-1 distance averaged over left-out marginals (L better) for 50-dim PCA represen- tation of single-cell data for corresponding datasets. Results are averaged over 5 independent runs.
Method
Cite (50D)
Multi (50D)
SF2 M-Geo
38.524 + 0.293
44.795 + 1.911
SI M-Exact
40.009 + 0.783
45.337 + 2.833
OT-CFM
38.756 + 0.398
47.576 = 6.622
I-CFM
41.834 + 3.284
49.779 + 4.430
WLF-UOT
37.007 1.200
46.286 + 5.841
WLF-SB
39.695 1.935
47.828 6.382
WLF-OT
38.352 + 0.203
47.890 = 6.492
OT-MFFM Rar
36.394 + 1.886
45.16 + 4.96
23

---

Supplementary Unpaired Translation Results
OT-CFM
OT-MFMMBBI
OZ
1-0.25
LOTO
L-1
L.000
L-
1-0,50
T-A77
Figure 4: Additional qualitative comparison for the task of unpaired transla- tion between OT-CFM and OT-MFM
24

---

G Supplementary Sphere Results
We visualize the problem setting in Figure We report the Earth Mover's Distance between the sampled and test sets, well as the mean distance of the middle points of trajectories to sphere, to quantitatively compare CFM and MFM on the task. MFM improves significantly over the Euclidean baseline, CFM (Table 8). Fur- thermore, the samples generated by MFM at intermediate times are much closer to the un- derlying sphere than the Euclidean counterparts (Table 7).
Sourc toe
Figure 5: Problem Setup
Table 7: Mean Distance of reconstructed tra- jectories time 1/2 from the sphere. Results averaged over 5 runs.
Table 8: Wasserstein-1 distance between recon- structed marginal at time 1/2 and ground-truth. Results averaged over 5 runs.
Method
Distance from Sphere (L)
Method
EMD (4)
OT-CFM
0.519 + 0.002
OT-CFM
0.525 0.003
OT-MFM LAND
0.085 + 0.005
OT-MFM LAND
0.340 + 0.074
H Supplementary LiDAR Visualizations
Figure 6: Supplementary Visualizations of MFM Interpolants on LiDAR
25

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and precede the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No] or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the
relevant information is Not Available.
Please provide short (1-2 sentence) justificatin right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]" ", is perfectly acceptable to answer "[No] "provided proper justification is given (c.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). ). general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only the provided macros for your answers.
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: The abstract and introduction clearly state our claims, assumptions, and outline contributions (see the end of the Introduction). The claims supported by theoretical and experimental results found in the main body and the provided appendix.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
26

---

Justification:
The limitations of this work can be found in the final section of the main body.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "I "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model sele--ppeiiiainnn. asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: All assumptions in theoretical results are clearly stated (both main body and appendix), and proofs are complete and correct.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: The main section of experiments alongside the Experimental Details section in the Appendix clearly outlines all the technical information, hyperparameters, and hardware needed to reproduce the results.
27

---

Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-soneee open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: The anonymous repository provided contains the code, python environment requirements, full set of hyperparameters for each result, and links to the datasets. The link to the repository can be found in the Experimental Details appendix.
Guidelines:
The answer NA means that paper does not include experiments requiring code. Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
wtttssawwsssnooooooooooooooooooe conooicyy) for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ //nips ppppuulccuuuuds..... /wtoeeeobeissiiiiiceaaaaae sesiiisssooooooo. seiionnnicc) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
28

---

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The experimental Details section in the Appendix clearly outlines all the data splits, hyperparameters, and optimizers used.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: All single-cell and synthetic experiments have been independently run across 5 seeds, and standard deviations have been reported. For single-cell experiments, the standard deviation is measured both across seeds and leave-out timepoints.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the eтror bar is the standard deviation or the standard error of the mean.
OK to report I-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Yes, such information can be found the beginning of the Experimental Details appendix.
29

---

Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: / /neuriis....// cwtrrpp..cpppppee
Answer: [Yes]
Justification: The authors reviewed the NeurIPS Code of Ethics, and the research conducted in this paper conforms with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Introduction outlines the context of the trajectory inference problem, and its potential positive impact biomedical applications. We cannot foresee immediate negative societal impacts associated with our work.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (c.g., deploymentt of technologics that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
30

---

Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Datasets used were already publicly available. We do not provide pretrained models with high risk for misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes)
We cite and acknowledge any assets used, both in the paper and the repository provided.
Justification:
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets apart form the provided anonymous code with hyperparameters for reproducibility.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
31

---

The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjeets
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
32


Reviews:
[{"review_id": "SHDbBMYCnh", "reviewer": "Reviewer_hBAX", "summary": "This paper proposes an instantiation of flow matching where the interpolants are learned by minimizing the kinetic energy defined by a nonparametric metric defined over a set of data points (a weighted L2 norm of the velocity field). This metric is defined through another weighted normal distribution, with learnable weights that ensure the metric is at the right scale.  The interpolants are defined using another neural network and optimized to reduce the kinetic energy.", "questions": "- How effective is the interpolant when no optimal coupling is used, especially for multimodal data distributions?\n- Does learning a full flow matching model outperform just using the L2 optimal coupling and using the interpolant to create x_t samples?\n- How does the framework work when you have multiple time snapshots? Is the metric time-dependent (by depending on different snapshots given time) or do you fit it to a single fixed-time snapshot?\n- This part confuses me: to justify the metric, it is described that p_t \"stays close\" to a reference dataset D, however, the evaluation settings use K separate data distributions (at K different time values). How do you determine what D is?\n- Given two disjoint distributions separated apart, does the metric still learn some non-trivial p_t between them (or does the metric degenerate to the Euclidean one when far away from data)?\n- Can this framework be adapted to learn a low dimensional manifold (rather than just a metric)?", "limitations": "There is a computational limitation when using larger training set sizes, as the majority of machine learning datasets are now in the order of millions and billions. It does not seem feasible to use a nonparametric method such as the metric proposed here.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 1, "strengths": "Learning a metric over a low-dimensional manifold could potentially be very interesting for high-dimensional machine learning applications.", "weaknesses": "### Clarity\n\nThis paper lacks clarity in its exposition, and are not sufficiently careful in stating its claims.\n\n- The paper introduces their framework by discussing the \"manifold hypothesis\": that high-dim data lies in a low-dim manifold. However, the choice of metric does not induce a lower dimensional manifold, as it is induced by a simple gaussian mixture model. For some of the cellular experiments, the manifold seems to actually be found by PCA and taking the first few components to define the space of the manifold. \n\n- The paper uses very vague terminology. For instance, there is repeated mention of a \"more meaningful matching\" being learned. But what is a \"matching\" and how do you compare between two? I believe the authors are referring to the learned time-dependent probability density at intermediate times?\n\n### Novelty\n\nSome of the claims about novelty are overly strong and could be more clear in precisely describing this paper's contribution of using a data-dependent metric (which is an interesting direction, I just feel the paper unnecessarily sugar coats its contribution instead of stating it objectively).\n\n- The proposed algorithm is essentially the same as Wasserstein Lagrangian Flows and GSBM, except these two existing works actually take a further step and learn the optimal transport coupling between x0 and x1 induced by the choice of metric. In contrast, I believe the main contribution of this work lies in its choice of metric defined over a finite data set. The paper repeatedly states that it is a generalization of the existing works but in fact I think of this work as an instantiation of a subset of the existing frameworks with a particular choice of metric. \n\n- One claim is that \"MFM is simulation-free and stays relevant when geodesics lack closed form\" when compared to riemannian flow matching. I think this is contrasting to the case of riemannian flow matching where the interpolants are solved with an ODE solver. However, here the geodesics are also not known in closed form and instead are solved by a neural network. I think of the proposed approach just as a time-parallel method for approximating the geodesic, and feel that the above claim is too strong. There is in fact a non-trivial optimization problem happening because MFM does not have closed-form geodesics, and the reliance on a neural network further suggests it is not as simple a procedure as this claim makes it sound like.\n\n### Empirical validation\n\nI feel the experimental results are messy and do not provide a coherent analysis.\n\nFor instance, part of the motivation of a data-dependent metric is to impose a more useful probability path p_t. However, the probability path p_t is determined by both the coupling q(x0, x1) and the interpolant. The use of an optimal coupling should have a very strong influence on the resulting p_t, and the choice of interpolant is not independently studied with an independent coupling.\n\nGoing through each experiment section, there are concerns regarding the setup or results:\n\n1) The LiDAR scan. \n- It is unclear what the aim of this experiment is. In the paper it is mentioned that this uses a different V_t than GSBM. [Suggestion:] However, one direct comparison that should be made is to take the same V_t as GSBM but simply replace their kinetic energy with your metric-based kinetic energy. This can help answer whether the data-dependent metric is useful for this setting.\n- It is hard to tell from the visualization, but do the samples from p_t stay close to the LiDAR data points or are they \"floating\"? [Suggestion:] i.e., is the condition || x_t - nn(x_t) || < max_i,j || x_i - x_j || satisfied? Here x_t is sampled from the learned p_t, nn returns the nearest neighbor of x_t in the LiDAR dataset, and the right-hand-side here is the maximum distance between points in the LiDAR dataset.\n\n2) The AFHQ experiment. \n- Here the FID values are extremely high, which indicates that the model is poorly fit. Looking at values reported by GSBM and the original StarGAN v2 where this dataset was introduced, it seems reasonable to expect FID values in the range of 10-20? [Suggestion:] use the same code as existing works to reproduce their results and perform a direct comparison.\n- Qualitatively looking at the samples, it doesn't seem clear which method has the better samples at t=1/2. I feel this setting doesn't showcase a non-trivial p_t since it is done in the latent space of a VAE, so we know that the samples follow a normal distribution. Given this, I'm not particularly convinced that a different interpolant provides something more interpretable. [Suggestion:] Perhaps it'd be good to visualize a 2D PCA of the trajectories? \n- Again, here I feel the role of the interpolant is heavily diminished when compared to the role of the optimal coupling. I believe LPIPS is computed based on pairs of (x0, x1), and that the values are extremely close seems to suggest the following: the choice of interpolant does not (or has very little) influence on the learned coupling. Since the goal of unpaired translation is to learn how to translate, it seems one should not be using the optimal coupling for training. [Suggestion:] test different interpolants while using the independent coupling.\n\n3) The single-cell experiment.\n- Here the setup is given K time points, while one time point is missing. However, we know that by construction, the probability density at this missing time point only depends on the two closest time points. [Suggestion:] Does training on all K-1 time points actually offer an improvement over training just over the nearest time points?\n- Furthermore, there is the issue with the choice of coupling influencing the results. [Suggestion:] Testing different interpolants while using the independent coupling would be good to have here.\n- When the optimal coupling is given, is there even a need to learn the flow? It seems this experiments can be solved by sampling pairs (x0, x1) from the coupling q and using the interpolant x_t. [Suggestion:] Does learning a flow actually provide any improvements on top of this baseline?\n- I am a bit confused regarding the different dimensions (number of components used from PCA), since the initial motivation of this work is to learn the low-dimensional manifold. I feel both this and the high-dimensional image setting could be set up where the metric directly influences the dimension of the flow model, which would significantly help support the initial justification.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their detailed feedback and constructive comments, which gave us an opportunity to improve our work significantly. We are glad to hear that the main thrust of our work could be very “interesting” for higher dimensional ML applications. We kindly point the reviewer to our **general rebuttal** and 1 page pdf, which includes several new ablations and additional results. In this rebuttal, we address the main concerns of the reviewer: (i) Clarity on metric learning vs manifold learning; (ii) Novelty; (iii) Experiments using independent coupling. \n\nBelow the rebuttal, i.e. in the following comments, we tackle any standing point from the review.\n\n### Metric learning vs Manifold learning\n\n> _The paper introduces their framework by discussing the \"manifold hypothesis\": that high-dim data lies in a low-dim manifold. However, the choice of metric does not induce a lower dimensional manifold._ \n \nThe reviewer makes a pertinent observation that the choice of metric in MFM does not induce a lower dimensional manifold. To clarify why this is **intended**, given our data-points, we are interested in defining a notion of path between samples $x\\_0$ and $x\\_1$ that preserve the underlying geometry. There are now two equivalent ways to proceed: (i) One can find explicit coordinate representations of the underlying manifold and build the path through these coordinates (*manifold learning*); (ii) Define the path in the ambient space but change the geometry of the ambient space so that regions away from the samples (and hence the manifold) are highly penalized (*metric learning*). In practice we achieve the same result, which is having paths bending according to where the samples are. By leveraging the metric learning approach (ii), we do not need to prescribe a specific lower dimension or find a coordinate representation. \n\nTo support our point further, we have followed the suggestion of reviewer mgNS and run MFM on the arch task when the samples belong to the 2D sphere in $\\mathbb{R}^3$. We have found that MFM not only improves significantly over the Euclidean baseline CFM, but crucially that the samples generated by MFM at intermediate times are much closer to the underlying sphere than the Euclidean counterpart (see Table 4 in the 1 page pdf and our **general rebuttal**). We emphasize that we manage to attain this **without** explicitly parameterizing the lower-dimensional space but simply relying on the LAND metric. We hope this provides you with strong empirical evidence as to why the metric approach can help us design flows that stay close to a lower-dimensional manifold even when we do not know it.\n\n### Novelty\n\n> _The proposed algorithm is essentially the same as Wasserstein Lagrangian Flows and GSBM (…)_\n\nWe value the reviewers opinion about the similarities between our proposed MFM and WLF and GSBM.  We would like to gently push back against the reviewers assertion that these methods take a further step than MFM. In particular, we disagree with the assertion that WLF and GSBM learn an OT coupling **induced by the choice of the metric**, since no Riemannian metric (of any kind) is proposed or studied. As such, we argue WLF and GSBM cannot be considered generalizations of MFM despite their similarities. \n\nWe kindly refer to the **general rebuttal** where we have provided details on technical differences between MFM and WLF and GSBM.\n\n\n> _One claim is that \"MFM is simulation-free and stays relevant when geodesics lack closed form\" when compared to riemannian flow matching. (...)._\n\nWe appreciate the reviewers comment regarding the simulation free nature of the approximate geodesics learned in MFM. At present, there is no computational method that can find exact geodesics without simulations and we certainly do not hope to claim that MFM exactly solves these. In contrast, we believe that the strength of MFM lies in the fact that the problem of trajectory inference benefits from approximate geodesics that can easily be plugged in a conditional flow matching framework in a computationally efficient manner—i.e. they are simulation free. We further highlight that such approximate geodesics can be found thanks to our proposed objective Eq. 6, i.e. \n\n$$\n\\mathcal{L}\\_g(\\eta) = \\mathbb{E}\\_{(x\\_0,x\\_1)\\sim q,t} \\left[(\\dot{x}\\_{t,\\eta})^\\top \\mathbf{G}(x\\_{t,\\eta};\\mathcal{D})\\dot{x}\\_{t,\\eta}\\right]\n$$\n \n\nwhere we minimize the Dirichlet energy of the path—whose optimum is in fact a geodesic. We point to the general rebuttal and to the 1 page pdf, where we showed how MFM manages to generate a flow whose intermediate samples are close to the lower-dimensional manifold by learning approximate geodesics of LAND through Eq. (6).  We hope the reviewer may now agree that our empirical evidence enables us to claim both technical and empirical novelty of using learned (simulation-free) geodesics for trajectory inference.\n\n### Experiments using independent coupling\n\n> _(...) The choice of interpolant is not independently studied with an independent coupling._\n\nWe thank the reviewer for this suggestion. In the **general rebuttal** we have provided additional experiments comparing CFM with independent coupling, i.e. I-CFM, and MFM with independent coupling, i.e. I-MFM, on the arch task and the single cell RNA datasets (in both 5D and 100D). We briefly summarize here the main takeaways: \n- I-MFM generally surpasses I-CFM.\n- We further stress that, differently from CFM, MFM also uses the coupling q in the first stage where it optimizes paths based on the metric, which justifies why using Optimal Transport for the coupling is even more beneficial for MFM than CFM.\n\nWe hope that in light of the new experiments, the reviewer can see that the benefits of MFM are not just due to the choice of the coupling, since I-MFM and OT-MFM both surpass their respective Euclidean counterparts I-CFM and OT-CFM.", "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": "MgyGAHQtpz", "reviewer": "Reviewer_mgNS", "summary": "The authors introduce a method for trajectory inference based on conditional flow matching (CFM) that takes a formulation of the trajectories using Riemannian geometry. The Riemannian metric is built following the manifold hypothesis and prior work, resulting in a flow matching with a data-dependent metric encouraging trajectories to stay close to the so-called data manifold. Furthermore, the sampling distribution used in the optimisation of the objective function is improved with optimal transport concepts, properly matching the joint distribution of interpolant endpoints to the main task of interest: single-cell dynamics.", "questions": "1. Taking into account the weakness mentioned above regarding novelty: could the authors clarify their contributions beyond Eq. (4)?\n2. I understand that in other domains, e.g. single-cell data, meaningfullness of trajectories might be more obvious to define, but I do not understand the claim about meaningfullness of trajectories between pictures of cats and dogs. Especially after looking at additional results in the appendix, I believe the setting is fundamentally flawed to evaluate the proposed model, with both models being similarly bad in many cases. What is the expected \"good\" case? Could the authors clarify their motivations for using this specific example? (By the way, the 6th row of Fig. 4 in the Appendix does not seem to be the same interpolants in both cases).\n3. Can a comparison with RFM be made in a setting where the underlying \"true\" appropriate metric is known? The Arch data set or some variant in hyperbolic geometry seem to be interesting cases for me.\n\nMinor comments:\n- Why are there different citations for the manifold hypothesis? At first (introduction), a set of papers are cited, but then later in Section 3.1, another paper is cited.\n- [l. 160] \"boundary conditions\": at this point I was defensively looking for some boundary conditions I might have missed, but later noticed the authors likely meant the endpoints $x_0$ and $x_1$. Is that the case? Could this be clarified at this at this point when first mentioning \"boundary conditions\"?", "limitations": "The authors have sufficiently addressed the limitations in their work.", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper is very well written and easy to follow. Its main originality lies in the use of a data-dependent Riemannian metric with conditional flow matching to learn trajectories between distributions which are discouraged to stray from the data manifold. The elegant incorporation of a data-dependent metric means the method works well without excessively requiring expert knowledge about the problem domain and an appropriate metric to be chosen.", "weaknesses": "The technical contributions allow CFMs to be employed with a Riemannian metric and without deep understanding of an appropriate metric to be used for the task of interest. However, most of the work to get there was already proposed before. In particular, beyond Eq. (4) followed by then adapting all distances to be Riemannian (as required by the problem setting), I see little \"on top\" of the work by Arvanitidis et al. (2016, 2021) and Tong et al. (2023b).\n\nGiven the technical contribution builds little on top of existing work, and taking into account question 2 below regarding the image translation experiment, I believe the paper showed limited applicability beyond single-cell dynamics. I expected a stronger experimental setting showing applicability in more problems where adopting a Riemannian metric is the advantageous choice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed feedback and constructive comments, which gave us an opportunity to improve our work significantly. We are pleased to see that the reviewer found our work “very well written and easy to follow” and that our idea is “elegant” and that “the method works well without excessively requiring expert knowledge”.  We now address key clarification questions raised by the reviewer.\n\n### Technical novelty\n\n> _Beyond Eq. (4)(...) I see little \"on top\" of the work by Arvanitidis et al. (2016, 2021) and Tong et al. (2023b)._\n\nWe value the reviewer’s comments regarding the technical contribution of MFM over prior works. We kindly refer to the **general rebuttal** for a detailed discussion on the main technical contributions of MFM, while here we provide a short summary.\n\nFirst, we hope that drawing connections from metric learning—such as the works of Arvanitidis et al—and flow matching, is in itself an important contribution that can foster new works in this space bridging different communities further. We believe that showing the role played by the data geometry (not the ambient geometry) for generative modeling is an important research direction.\n\nMore crucially, in terms of technical contributions, we believe that the novelty is not just in the parameterization adopted in (4), but also, crucially, in the optimization in Eq. (6). In fact the objective in Eq. (6) is based on evaluating a *data-dependent metric Dirichlet energy* over the paths, once we sample the boundary points $x\\_0,x\\_1$ according to the coupling q. A key contribution in our work is that we can learn approximate geodesics using a simulation free learning objective which we find is quite useful for downstream generative modeling applications like trajectory inference. Moreover, we believe that optimizing using a velocity-induced regularization is an essential part of our framework, along with aligning the geodesic sampling with the joint distribution q. Crucially though, our framework is not bound to use a specific metric, and we simply adopted existing ones to showcase flexibility and ease of use. We hope this addresses the reviewer’s concerns, and clarifies the technical novelty of our paper.\n\n\n### Questions\n\n> _Taking into account the weakness mentioned above regarding novelty: could the authors clarify their contributions beyond Eq. (4)?_\n\nWe have addressed this point above when it comes to technical contributions.\n\n> _I understand that in other domains, e.g. single-cell data, meaningfullness of trajectories might be more obvious to define, but I do not understand the claim about meaningfullness of trajectories between pictures of cats and dogs. (...)_\n\nWe agree that meaningfulness of trajectories can be better evaluated over actual dynamical systems such as for single cell data. The experiments on unpaired image translation were meant to showcase that MFM can be on par, or in fact better, than the Euclidean baseline even on settings it was not mainly designed for. In particular, we sought to attach a visual correspondence to the learned interpolation from MFM. In contrast, one can consider another exotic interpolant which may traverse other classes—or entirely off the image manifold—in trying to go from cats->dogs. Note that to further try assessing the meaningfulness of interpolations we adopted the LPIPS metric.\n\n> _Can a comparison with RFM be made in a setting where the underlying \"true\" appropriate metric is known? The Arch data set or some variant in hyperbolic geometry seem to be interesting cases for me._\n\nThis is a great suggestion! We address this with new experiments in our 1 pg PDF. We have extended the arch task on the 2D sphere embedded in $\\mathbb{R}^3$. We have found that MFM not only improves significantly over the Euclidean baseline CFM (see Table 5 in the 1 page pdf), but crucially that the samples generated by MFM at intermediate times are much closer to the underlying sphere than the Euclidean counterpart (see Table 4 and Figures 1,2,3 in the 1 page pdf and comments in our general rebuttal). We emphasize that we manage to attain this **without** explicitly parameterizing the lower-dimensional space but simply relying on the LAND metric. Finally, we note that RFM uses the ground-truth geodesics of the *standard* metric on the sphere, meaning that in general how well RFM is able to solve trajectory inference problems on the manifold depends on how well the geodesics provided by the data-agnostic metric, resemble trajectories of the underlying system. Conversely, MFM directly learns from data and can approximately recover a curved, lower-dimensional manifold from the samples.\n\n\n**Minor comments**\n\n\n\n> _Why are there different citations for the manifold hypothesis?_\n\nThis is just because we wanted to cite key papers in the main body of the paper (beyond the related work section) and the distinction (i.e. not citing all of them in both cases) is due to space constraints and formatting.\n\n> _[l. 160] \"boundary conditions\": at this point I was defensively looking for some boundary conditions I might have missed, but later noticed the authors likely meant the endpoints  and . Is that the case? Could this be clarified at this at this point when first mentioning \"boundary conditions\"?_\n\nYes, by boundary conditions we mean that the paths recover $x\\_0$ and $x\\_1$ at times 0 and 1, respectively. We will add an extra sentence here to improve clarity.\n\n### Final comment\n\nWe thank the reviewer again for their review and detailed comments that helped strengthen the paper—particularly through the addition of MFM experiments on the sphere. We believe we have answered to the best of our ability all the questions here and in the general rebuttal. We hope this  allows the reviewer to consider upgrading their score if they see fit. We are also more than happy to answer any further questions.", "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": "JNMjtsQZGo", "reviewer": "Reviewer_q2sv", "summary": "This paper proposes metric flow matching, a variant of conditional flow matching where the interpolated distributions lie on the data manifold. They first learn trajectories between $p_0$ and $p_1$ which minimize a data-dependent kinetic energy, then use these trajectories instead of straight lines for conditional flow matching. The model is applied to single-cell trajectory prediction, where performance is better than the non-metric flow matching baseline.", "questions": "1. Line 125: \"regular\" means smooth in some way? Can you be more specific?\n2. Why is it preferable to use LAND for lower dimensions?\n3. Am I correct that the LIDAR data is just an example of a 2d manifold embedded in 3d? I found it confusing at first since from fig 2 I thought it was encouraged to follow low-altitude trajectories and that maybe the height was being used to inform the metric. Perhaps you can chose a different visualization/update the description to avoid this confusion", "limitations": "The only limitation given is that the data must be embedded in Euclidean space. I am sure the authors can think of others, such as using a diagonal metric, and whatever trade-offs are involved in the OT scheme they use (see Weaknesses above)", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper is very well written, with great attention to detail. Everything is easy to follow, including the more technical details. The authors address an interesting problem (how to get more interpretable and useful flow matching trajectories) which has received surprisingly little attention given the current popularity of flow matching. They find a suitable application in single-cell trajectory prediction and it may improve the quality of image translation models (it's hard to be confident since the image experiments are limited). To my knowledge the work is an original contribution and a useful addition to the generative modelling community.", "weaknesses": "I have some minor complaints:\n1. It seems that the dataset $\\mathcal{D}$ is the concatenation of samples from both $p_0$ and $p_1$ but this is not clearly stated\n2. There is no justification of the use of diagonal metrics except that it reduces computation cost. Surely there must be some downside? Is it possible to compare the performance to non-diagonal metrics in the 2d example?\n3. It's not clear how the OT is implemented. Is it a batchwise scheme? Does this introduce a bias? If not batchwise, is it scalable?\n4. In tables 3 and 4, the MFM results are bolded. While the mean value is indeed the lowest, we cannot be confident that the true value is lower than some of the other values (e.g., WLF-UOT) with the stated error bounds. While I understand that this is common practice, I personally find it unscientific to report MFM values in bold and not others which could plausibly be lower than the MFM values\n5. Line 289: I don't agree that the interpolations are better semantically. If there is a difference, it is too small to be able to make such a subjective claim", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and positive appraisal of our work. We are thrilled that the reviewer viewed our work to be “well-written, with great attention to detail” and that we “address an interesting problem” and is an “original contribution” that is a “useful addition to the generative modelling community”. We now provide responses to the main questions raised by the reviewer.\n\n ### Weaknesses\n\n> _1. It seems that the dataset $\\mathcal{D}$ is the concatenation of samples from both $𝑝_0$ and $p_1$  but this is not clearly stated_\n\nYes that is correct—we will make this more explicit, since we work in a “generalized” trajectory inference setup, we think of samples from $p\\_0$ and $p\\_1$ as samples from the same dynamical system evaluated at different times so that the manifold we refer to is indeed the one spanned by the trajectories.\n> _2. There is no justification of the use of diagonal metrics except that it reduces computation cost. Surely there must be some downside? Is it possible to compare the performance to non-diagonal metrics in the 2d example?_\n\nThat's a great question! Working with diagonal matrices, may potentially have some expressive power limitations in high dimensions, specifically for settings where capturing a precise underlying metric is crucial—note that in the applications considered in our submission it is hard to identify a ground-truth metric. In general though, there are theoretical reasons as to why our choice is not particularly limiting; for example, in the 2D case any metric is actually locally conformally flat (follows from the existence of isothermal coordinates). Accordingly, for 2D manifolds, one would not lose expressive power by actually taking the matrix $\\mathbf{G}$ to be a multiple of the identity pointwise. Nonetheless, our framework is not constrained to diagonal metrics and more general cases are indeed possible. To further support this point, we refer to the additional experiments shared in the 1 page pdf and **general rebuttal**, where we showed how intermediate samples generated using MFM are much closer to the underlying lower dimensional manifold (a sphere) than the ones generated using CFM (see Table 4 and Figures 1,2,3). Crucially, MFM only leverages the diagonal LAND metric defined in the ambient space based on samples and is never given information about the sphere. \n> _3. It's not clear how the OT is implemented. Is it a batchwise scheme? Does this introduce a bias? If not batchwise, is it scalable?_\n\nOur method relies on using mini-batch OT which has been a standard tool in the generative modeling / Machine Learning literature (see for example Tong et al., 2023 b,c). In this case we can easily trade off the cost of OT with the batch size and in practice this does not create a severe overhead in training our models. We find that OT helps improve training stability and leads to better empirical performance. We note however that MFM with independent coupling still surpasses CFM with independent coupling, as reported in our global 1 pg response where we compare CFM and MFM without OT (Table 1,2,3).\n\n\n> _4. In tables 3 and 4, the MFM results are bolded. (..) I personally find it unscientific to report MFM values in bold and not others which could plausibly be lower than the MFM values_\n\nThank you for your feedback. We agree with your observation and will remove bold numbers from Tables 3 and 4. \n\n> _5. Line 289: I don't agree that the interpolations are better semantically. If there is a difference, it is too small to be able to make such a subjective claim_\n\nWe again agree that this may be entering subjective territory a little, so we refine our statement and claims.\n\n\n### Questions\n\n> _1. Line 125: \"regular\" means smooth in some way? Can you be more specific?_\n\n\nYes, it does mean smooth in some way. To improve clarity, when citing the paper of Ambrosio et al we refer to their explicit statement, where assumptions are stated clearly.\n\n> _2. Why is it preferable to use LAND for lower dimensions?_\n\nThat's an excellent question. We empirically found LAND to be marginally better in lower dimensions, which we believe to be reasonable given that in this setting we are building the metric using all samples directly without any clustering being involved beforehand.\n> _3. Am I correct that the LIDAR data is just an example of a 2d manifold embedded in 3d? (...)_\n\nYes that is correct, as in we do not provide height information for the metric since we wanted to test MFM in a setting where the metric was built agnostic of the downstream application. To address your concerns, in the additional 1 page pdf we have provided more views which we will add in the appendix of the revised version to improve clarity.\n\n> _The only limitation given is that the data must be embedded in Euclidean space. I am sure the authors can think of others, such as using a diagonal metric, and whatever trade-offs are involved in the OT scheme they use (see Weaknesses above)_\n\n\n\nThank you for the feedback. We have now also added a line mentioning the trade-offs associated with the OT scheme. In light of our previous comment, we do not see any serious weakness with using diagonal metrics and in general the framework does not require one.\n\n### Final comment\n\nWe hope we have addressed the main concerns of the reviewer, and that the addition of experiments of MFM without OT and MFM on a sphere have further improved the strength of our submission. We hope that our answers enable the reviewer to continue endorsing our paper and potentially even upgrading the score if the reviewer deems it. We are also more than happy to engage in any further discussion.", "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": "Ywzl68T9Aq", "reviewer": "Reviewer_CNRE", "summary": "This work proposes a metric flow matching algorithm, where interpolants are approximate geodesics learned by minimizing the kinetic energy of a data-induced Riemannian metric. This targets the trajectory inference problem, such as single-cell trajectory prediction.", "questions": "* Both $g_{RBF}$ and $g_{LAND}$ capture data density and ensure that the interpolants lie within the data support. In principle, this algorithm appears to be designed to ensure that inferred trajectories lie within the data support. Are there other trajectory inference approaches where algorithms essentially do the same thing (i.e., identify trajectories that mostly lie within the data support and connect two data distributions)? If so, is the advantage of your method that it can provide directional information when using a general Riemannian metric? Can a 'good' Riemannian metric be defined for trajectory inference that includes direction? How does this compare to ``Riemannian Metric Learning via Optimal Transport (ICLR 2023)''?", "limitations": "Yes", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* This paper clearly addresses the motivation for solving the trajectory inference problem by proposing a solution that naturally connects the recently introduced flow matching method with a data-induced Riemannian metric. The paper is well-written, with preliminaries explained compactly and effectively. \n* Additionally, it thoroughly discusses the differences with Riemannian flow matching. The experiments are well compared with recent studies.", "weaknesses": "*Major comments*\n\nI have two main questions and comments, which I expect to be addressed in the revised version.\n\n* I believe the trajectory inference problem with unpaired cross-sectional data has been well addressed in \"Riemannian Metric Learning via Optimal Transport (ICLR 2023).\" In this work, the authors learn the Riemannian metric without introducing an auxiliary objective, but by alternatively minimizing the optimal transport objective. Your approach and this work seem to have a fundamental difference, as the objectives used for learning the Riemannian metric are different. This difference should be clearly noted, emphasizing the advantages of your approach. I hope this is included in the Introduction section with some emphasis.\n\n* You have used L2 distance for \\( c(x,y) \\) in (12) due to computational issues. While this is understandable, I believe it could significantly impact the results. The two main factors affecting the resulting trajectory via flow matching are the Riemannian metric and the choice of coupling \\( q \\). In the original flow matching algorithm, the quality of the resulting generation was more important than the intermediate trajectory. However, in this study, the trajectory itself is the objective, making it crucial to carefully choose these two factors. \nIn the current version of the paper, this issue seems to be mentioned very briefly and almost as if it were insignificant. It appears necessary to include a more detailed discussion and address this point more thoroughly.\n\n*Minor comments*\n\nI appreciate comprehensive reviews in section 6. There are a few missing references, and I sugges to include them. \n* Geometry-aware generative models:\n   - Learning geometry-preserving latent spaces: Regularized Autoencoders for Isometric Representation Learning (ICLR 2022) and Geometric Autoencoders--What You See is What You Decode (ICML 2023), \n   - Learning data manifolds and latent spaces by exploiting ambient space metrics: A statistical manifold framework for point cloud data (ICML 2022)'' and Geometrically regularized autoencoders for non-Euclidean data (ICLR 2023)\n   - geometry-based regularization for learning accurate manifolds: Neighborhood Reconstructing Autoencoders (NeurIPS 2021) and On explicit curvature regularization in deep generative models (TAG-ML 2023).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their time and thoughtful review of our work, which gave us an opportunity to improve our work significantly. We are glad to hear that the reviewer found our work “well-written” and that it “proposes a solution that naturally connects the recently introduced flow matching method with a data-induced Riemannian metric” with “experiments that are well compared with recent studies”. We first refer to the **general rebuttal**, where we have provided additional experiments and visualizations (see the 1 page pdf). Below, we address the key clarification points and questions raised in the review one by one. \n\n\n### Weaknesses\n\n> _I believe the trajectory inference problem with unpaired cross-sectional data has been well addressed in \"Riemannian Metric Learning via Optimal Transport (ICLR 2023).\" (..) Your approach and this work seem to have a fundamental difference, as the objectives used for learning the Riemannian metric are different. This difference should be clearly noted, emphasizing the advantages of your approach._\n\nWe appreciate the reviewer for pointing out this highly relevant work. We have already added a citation to Scarvelis & Solomon 2023 work in a revised version of the paper. We will further comment extensively on key differences: (i) the objectives being learned and (ii) the fact that we first optimize the paths and then match the vector field rather than regularizing directly the vector field (as in their Eq. (7)). Crucially, since MFM relies on flow matching, our framework does not require simulations during training when tackling the trajectory inference problem, which is an advantage compared to the work of Scarvelis and Solomon. \n> _You have used L2 distance for ( c(x,y) ) in (12) due to computational issues. While this is understandable, I believe it could significantly impact the results. The two main factors affecting the resulting trajectory via flow matching are the Riemannian metric and the choice of coupling ( q ). (...)  It appears necessary to include a more detailed discussion and address this point more thoroughly._\n\nWe appreciate the reviewer's concern. It is correct that the quality of the matching is given by both the choice of a coupling $q$—that is used to sample the boundary points $x_0$ and $x_1$—and the choice of the interpolants connecting such boundary points. First, we note that choosing a non-Euclidean cost for the OT objective is an interesting and non-trivial problem on its own. For arbitrary Riemannian metrics it is often unlikely to hope for a closed form expression to the metric-induced distance and thus OT cost. As a result, this requires us to simulate all pairs of matching to compute $c(x,y)$ which is computationally prohibitive from a practical standpoint. Furthermore, we argue an advantage of our framework is showing that even when we fix the cost to be Euclidean for the OT coupling, by learning data-dependent paths we can improve upon the Euclidean baseline. \nTo further decouple the impact of the interpolants from the choice of the coupling, we conducted further experiments—found in our 1 pg global response PDF—comparing both CFM and MFM on the arch dataset and the single-cell ones, where we chose the coupling q to simply be the independent one. We agree with the reviewer that aspects of this discussion need further clarifications and we will include a larger section to explicitly outline our motivation and claims, along with the addition of the new experiments using independent coupling.\n\n\n### Minor comments\nThanks for sharing these works, we will add relevant citations in the revised version of our paper.\n\n### Questions\n> _Both $g\\_{\\rm RBF}$ and $g\\_{\\rm LAND}$ capture data density and ensure that the interpolants lie within the data support. In principle, this algorithm appears to be designed to ensure that inferred trajectories lie within the data support. Are there other trajectory inference approaches where algorithms essentially do the same thing (i.e., identify trajectories that mostly lie within the data support and connect two data distributions)?_\n\n\nRegarding the choice of the metrics, we adopted those that seem to be the easiest/most efficient ones to show versatility of our framework and importantly reduce any computational overhead compared to the Euclidean FM baseline. Crucially though, we also wanted to highlight how our framework can, in principle, be applied beyond trajectory inference (as for the image interpolation problem), which is why we did not go overly-specific on the choice of the metric based on trajectories. However we believe this to be an interesting direction, i.e. how to choose task-dependent metrics to further improve the MFM framework. We hope this can be addressed in future work by us and the community more generally.\n\n### Final comment\n\nWe hope that our responses here in conjunction with the general rebuttal and the additional experiments help answer the great questions raised by the reviewer. We politely encourage the reviewer to continue asking more questions or if possible consider a fresher evaluation of our paper with a potential score upgrade.", "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:
We thank the reviewers for their thoughtful feedback and constructive questions that have helped us improve the submission significantly. We are glad to see that reviewers found our work “well-written, with great attention to detail” (R q2sv), that our proposed framework is “elegant” (R mgNS), “naturally connects the recently introduced flow matching method with a data-induced Riemannian metric” (R CNRE), and “could be very interesting for higher dimensional ML applications” (R hBAX). In this general response we address two points: (1) an overview of the new experiments we ran to address shared questions raised in the reviews; (2) the technical novelty and contributions of this work.

We discuss the additional experiments we conducted during the rebuttal period. *We will refer to the attached 1 page document*. 

### Additional experiments & Ablations

**Independent coupling on arch and single cell (R hBAX, R CNRE)**

To assess the impact of metric learning in MFM even without using Optimal Transport for the coupling $q$, we tested MFM with independent coupling on the Arch task and the single cell datasets (on both 5D and 100D). The results are reported in **Tables 1,2,3**  in the attached pdf, and highlight two key takeaway messages:
- I-MFM generally surpasses I-CFM

- Differently from CFM, MFM also uses the coupling q in the first stage where it optimizes paths based on the metric, which justifies why using Optimal Transport for the coupling is even more beneficial for MFM than CFM.

We hope that in light of these experiments, reviewers can see that the benefits of MFM are not just due to the choice of the coupling, since I-MFM and OT-MFM both surpass their respective Euclidean counterparts I-CFM and OT-CFM.

**Arch experiment on a explicit manifold (R hBAX, R mgNS, R q2sv)**

To assess the ability of MFM to learn trajectories that stay close to an unknown, underlying manifold, we followed the great suggestion of R mgNS and ran MFM on the arch task defined on a specific lower dimensional space, i.e. a 2D sphere embedded in $\mathbb{R}^3$. We see that MFM not only improves significantly over the Euclidean baseline CFM (see Table 5 in the attached pdf), but crucially that the samples generated by MFM at intermediate times are much closer to the underlying sphere than the Euclidean counterpart (see Table 4 and Figures 1,2,3 in the 1 page pdf). We emphasize that we managed to attain this **without** explicitly parameterizing the lower-dimensional space but simply relying on the LAND metric. We hope this provides the reviewers stronger empirical evidence as to why the metric approach can help us design flows that stay close to a lower-dimensional manifold even when we do not know it.


**LIDAR visualizations (R hBAX, R q2sv)**

To address the reviewers' concerns about the faithfulness of $p_t$ to the LiDAR data, in our additional 1 page pdf, we have provided more views of the learned paths from OT-CFM and OT-MFM (see Figure 4). The new visualizations clearly indicate that MFM paths 1.) do not intersect the manifold and 2.) bend closely around the manifold which highlights the geometric inductive bias. 

### Technical novelty and contributions (R hBAX, R mgNS)

In light of some concerns about technical contributions of our submissions, we reiterate and clarify what we believe to be the key novelty of our work.

**Connecting metric learning to flow matching** 

To the best of our knowledge, this is the first work that links metric learning to recent state-of-the-art generative frameworks such as flow matching and emphasizes the role played by the data geometry (*not* the ambient geometry) for generative modeling. 

**Learning interpolants that approximate geodesics**

The key technical contribution of our work consists in learning stochastic interpolants that approximate geodesics of a data-dependent metric in a simulation-free manner. At present, there is no computational method that can find exact geodesics without simulations. Instead, MFM proposes to find approximate geodesics by minimizing the **metric Dirichlet energy of the path** in Eq. 6, whose minimizer is the geodesic, i.e.

$$
\mathcal{L}\_g(\eta) = \mathbb{E}\_{(x\_0,x\_1)\sim q,t} \left[(\dot{x}\_{t,\eta})^\top \mathbf{G}(x\_{t,\eta};\mathcal{D})\dot{x}\_{t,\eta}\right].
$$

To our knowledge, this is a novel objective and, importantly, can be fully decoupled from the matching objective used to learn the vector field generating the flow. Finally, our framework is general and is not bound to use a specific metric.

**Comparisons with GSBM/WLF**

In light of reviewer hBAX’s comments about similarities between MFM and WLF/GSBM—which we discuss in Section 4.2 of our paper—we comment further on **technical differences** between these frameworks. Since GSBM is an explicit matching counterpart of WLF, we focus on comparing MFM to WLF.
- (i) In contrast to WLF or GSBM, our method, MFM, separates the optimization of the paths from the matching objective. We argue this is beneficial since it avoids introducing additional challenges when learning the vector field $v\_\theta$. 
- (ii) None of the Lagrangians considered in WLF account for the data distribution and/or the geometry induced by it. As such, the choice of Lagrangian always needs to be given by the _user based on overall considerations, e.g. unbalanced OT, or one needs to define a potential V_—that only depends on positions and not velocities, as discussed in 4.2. In MFM instead, we explicitly account for the empirical samples to learn the paths. If we adopt the same procedure for optimization (i.e. 2-stage learning, as clarified in (i)), then our approach can be interpreted as learning paths by minimizing a Lagrangian whose potential depends on both positions and velocities and is further dependent on the whole data distribution through the metric.

Meta-review: The paper proposes to endow flow matching methods with the non-parametric LAND metric to thereby better adapt to the data distribution. Several significant critiques were raised, especially regarding the validity of the conducted experiments. The reviewers were split and did not reach a consensus. The authors are expected to take the reviewer's feedback into account when finalizing the manuscript.
Decision: Accept (poster)