Metric Flow Matching for Smooth Interpolations on the Data Manifold

Authors: kacper  kapuśniak
peter  potaptchik
teodora  reu
leo  zhang
alexander  tong
michael  bronstein
avishek joey bose
francesco di giovanni
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 18:34:51

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.

Main paper: 
 introduction: 
 A central task in many natural and scientific domains entails the inference of system dynamics of an 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 to entire time-trajectories 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 et al., 2015] , where measurements are sparse and static, due to the procedure being expensive and destructive. Consequently, the nature of these tasks demands the design of frameworks capable of reconstructing the temporal dynamics of a system (e.g. cells) from observed time marginals that contain finite samples. This overarching problem specification is referred to as trajectory inference [Hashimoto et al., 2016 , Lavenant et al., 2021] .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 a Schrödinger Bridge [Schrödinger, 1932 , Léonard, 2013] . Specifically, to reconstruct the unknown dynamics t → p * t between observed time marginals p 0 and p 1 , we leverage Conditional Flow Matching (CFM) [Tong et al., 2023b] , a simulation-free framework which constructs a probability path p t through interpolants x t connecting samples of p 0 to samples of p 1 . In general, x t is designed under the assumption of Euclidean geometry, resulting in straight 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 t = 1 2 , using Conditional Flow Matching and METRIC FLOW MATCHING (MFM), respectively. MFM provides a more meaningful reconstruction supported on the data manifold.trajectories. However, in light of the widely accepted "manifold hypothesis" [Tenenbaum et al., 2000, Belkin and Niyogi, 2003] , the target time-evolving density p * t is supported on a curved low-dimensional manifold M ⊂ R d -a condition satisfied by cells in the space of gene expressions [Moon et al., 2018] . As such, straight interpolants stray away from the data manifold M, leading to reconstructions p t that fail to model the nonlinear dynamics generated by the underlying process.Present work. We aim to design interpolants x t that stay on the data manifold M associated with the underlying dynamics. Nonetheless, parameterizing the lower-dimensional 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 et al., 2002 , Hauberg et al., 2012] , where we still work in the ambient space R d , but equip it with a data-dependent 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-free generalization of CFM where interpolants x t are learned by minimizing a geodesic loss that penalizes the velocity measured by the metric g. As a result, x t approximates the geodesics of g and hence tends towards the data, leading to the evaluation of the matching objective in regions of lower uncertainty. Therefore, the resulting probability path p t is supported on the data manifold for all t ∈ [0, 1], giving rise to a more natural reconstruction of the underlying dynamics in the trajectory inference task, as depicted in Figure 1 .Our main contributions are:(1) We prove that given a dataset D ⊂ R d , one can always construct a metric g on R d such that the geodesics connecting x 0 sampled from p 0 to x 1 sampled from p 1 , always lie close to D ( §3.1).(2) We propose METRIC FLOW MATCHING, a novel framework generalizing CFM to the Riemannian manifold associated with any data-dependent metric g. In MFM, before training the matching objective, we learn interpolants that stay close to the data by minimizing a 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 t to stay supported on the data manifold ( §4.1). Through the proposed metric, x t depends on the entire data manifold and not just the endpoints x 0 and x 1 sampled from the marginals. By accounting for the Riemannian geometry induced by the data, MFM generalizes existing approaches that construct x t 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 dynamics ( §5). Additionally, we validate the versatility of the framework through tasks such as 3D navigation with LiDAR point clouds and unpaired translation of images. 
preliminaries and setting: 
 We review Conditional Flow Matching, which forms the basis of METRIC FLOW MATCHING §3. Next, we recall basic notions of Riemannian geometry, with emphasis on constructing geodesics.Notation and convention. We let R d be the ambient space where data points are embedded. A random variable x with a distribution p is denoted as x ∼ p(x). A function φ depending on time t, space x and learnable parameters θ, is denoted by φ t,θ (x), and its time derivative by φt,θ . We also let δ x0 (x) be the Dirac function centered at x 0 and assume that all distributions are absolutely continuous, which allows us to use densities. We denote the space of symmetric, positive definite d × d matrices as SPD(d), and let G(x) ∈ SPD(d) be the coordinate representation of a Riemannian metric at some point x.Conditional Flow Matching. We consider a source distribution p 0 and a target distribution p 1 defined on R d . We are interested in finding a map f that pushes forward p 0 to p 1 , i.e. f # p 0 = p 1 .In line with Continuous Normalizing Flows [Chen et al., 2018] , we look for a map of the form f = ψ 1 , where the time-dependent diffeomorphism ψ t : R d → R d is the flow generated by a vector field u t , i.e. dψ t (x)/dt = u t (ψ t (x)), with ψ 0 (x) = x, for all x ∈ R d . If we define the density path p t = [ψ t ] # p 0 , then p t and u t satisfy the continuity equation and we say that p t is generated by u t .If density path p t and vector field u t are known, we could regress a vector field v t,θ , modeled as a neural network, to u t by minimizing L FM (θ) = E t,pt(x) ∥v t,θ (x) -u t (x)∥ 2 . Since p t and u t are typically intractable though, Conditional Flow Matching (CFM) [Lipman et al., 2023 , Albergo and Vanden-Eijnden, 2022 , Liu et al., 2022] simplifies the problem by assuming that p t is a mixture of conditional paths: p t (x) = p t (x|z)q(z)dz, where z = (x 0 , x 1 ) is sampled from a joint distribution q with marginals p 0 and p 1 , and p t (x|z) satisfy p 0 (x|z) ≈ δ x0 (x) and p 1 (x|z) ≈ δ x1 (x). If ψ t (x|x 0 , x 1 ) denotes the flow generating p t (x|z), then the CFM objective isEQUATIONwhere x t := ψ t (x 0 |x 0 , x 1 ) are the interpolants from x 0 to x 1 . Since L FM and L CFM have same gradients [Lipman et al., 2023 , Tong et al., 2023b] , we can use the tractable conditional paths to learn v θ .As in §3 we design x t to approximate geodesics, we review key notions from Riemannian geometry.Riemannian manifolds. A Riemannian manifold (M, g) is a smooth orientable manifold M equipped with a smooth map g assigning to each point x an inner product ⟨•, •⟩ g defined on the tangent space of M at x. We let G(x) ∈ 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 §A for details). A continuous, positive function p is then a probability density on (M, g), i.e. p ∈ P(M), if p(x)dvol(x) = 1. Naturally, it is possible to define curves γ t , indexed by time t. A geodesic is then the curve γ * t that minimizes the distance with respect to g. Specifically, the geodesic connecting x 0 to x 1 in M, can be found by minimizing the length functional:EQUATION)where γt is velocity. From eq. ( 2), we see that geodesics tend towards regions where ∥G(x)∥ is small. In Euclidean geometry (i.e., M = R d and G(x) = I), γ * t is a straight line with constant speed. 
metric flow matching: 
 We introduce METRIC FLOW MATCHING (MFM), a new simulation-free framework that generalizes CFM by constructing probability paths supported on the data manifold. MFM learns interpolants x t in eq. ( 1) whose velocity minimizes a data-dependent Riemannian metric assigning a 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, v θ in 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 t 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 a Riemannian metric in R d , 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] .Assume that p 0 and p 1 are empirical distributions and that we have access to a dataset of samples D = {x i } N i=1 -in 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 t → p * t , with observed time marginals p * 0 = p 0 and p * 1 = p 1 -the 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 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 a low-dimensional manifold M ⊂ R da property satisfied by cells embedded in the space of gene expressions [Moon et al., 2018] :EQUATIONSince any regular time dynamics can be described using the continuity equation generated by some vector field v * t [Ambrosio et al., 2005, Theorem 8.3 .1], [Neklyudov et al., 2023a] , we rely on CFM, and approximate p * t via the probability path p t associated with v t,θ in eq. ( 1). From eq. ( 3), it follows that a valid reconstruction entails supp(p t ) ⊂ M. As the support of p t is determined by the interpolants x t in eq. ( 1), we need x t 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: [Tong et al., 2023b , Shaul et al., 2023] , with x 0 , x 1 sampled from the marginals. In this case, if p t (x|x 0 , x 1 ) ≈ δ xt (x), then the support of p t satisfiesx t = tx 1 + (1 -t)x 0supp(p t ) ⊂ {y ∈ R d : ∃ (x 0 , x 1 ) ∼ q : y = tx 1 + (1 -t)x 0 }.However, the dynamics t → p * t is often nonlinear, as for single-cells [Moon et al., 2018] , meaning that supp(p t ) ̸ ⊆ M. Straight interpolants are hence too restrictive and should instead be supported on M so to replicate actual trajectories from x 0 to x 1 , which are generated by the underlying process.Operating on a lower-dimensional M is challenging though, since it requires different coordinate systems [Schonsheck et al., 2019 , Salmona et al., 2022] and may incur overfitting [Loaiza-Ganem 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 x t should remain close to D. Therefore, instead of changing the dimension, we design x t to minimize a cost in R d that is lower on regions close to D. We achieve this following the "metric learning" approach, [Hauberg et al., 2012] where we equip R d with a suitable Riemannian metric g.Definition 1. A data-dependent metric g on R d is a smooth map g : R d → SPD(d) parameterized by the dataset D = {x i } N i=1 ⊂ R d , i.e. g(x) = G(x; D) ∈ SPD(d), ∀x ∈ R d .We describe a specific metric g in §4, and note that g can also enforce constraints from the task ( §7). Naturally, if G(x; D) = I, we recover the Euclidean metric. We show that we can always construct g so that the geodesics γ * t stay close to the data D. For details, we refer to Theorem B.1 in Appendix §B. Proposition 1 (Informal). Given a dataset D ⊂ R d , let g be any metric such that: (i) The eigenvalues of G(x; D) do not approach zero when x is distant from D; (ii) ∥G(x; D)∥ is sufficiently small if x is close to D. Then for each (x 0 , x 1 ) ∼ q, the geodesic γ * t connecting x 0 to x 1 stays close to D.Given g as in the statement, if x t = γ * t and p t (x|x 0 , x 1 ) ≈ δ xt (x), then the probability path p t generated by v θ in eq. ( 1) has support near D, i.e. supp(p t ) lies close to the data manifold M, which is our goal. Unfortunately, for all but the most trivial metrics g on R d , it is not possible to obtain closedform expressions for the geodesics γ * t . As such, finding the geodesic γ * t necessitates the expensive simulation of second-order nonlinear Euler-Lagrange equations [Hennig and Hauberg, 2014] .Consider a metric g on R d as in Definition 1 whose geodesics γ * t lie close to D as per Proposition 1. We propose a simulation-free approximation to paths γ * t by introducing interpolants of the formEQUATIONwhere η are the parameters of a neural network φ t,η acting as a nonlinear "correction" for straight interpolants. Note that the boundary conditions are met, i.e. that the path x t,η recovers both x 0 and x 1 at times t = 0 and t = 1, respectively. In fact, x t,η reduces to the convex combination between x 0 and x 1 if φ t,η = 0, meaning that x t,η strictly generalize the straight paths used in [Lipman et al., 2023 , Liu et al., 2022] . Towards the goal of learning η so that x t,η approximates the geodesic γ * t , we note that γ * t can be characterized as the path minimizing the convex functional E g below:EQUATIONSince γ * t minimizes E g over all paths connecting x 0 to x 1 , and x t,η in eq. ( 4) satisfies these boundary conditions, we can estimate η by simply minimizing the convex functional E g over x t,η , which leads to the following geodesic objective (the training procedure is reported in Algorithm 1): Sample (x 0 , x 1 ) ∼ q and t ∼ U(0, 1)3:x t,η ← (1 -t)x 0 + tx 1 + t(1 -t)φ t,η (x 0 , x 1 ) ▷ eq. (4) 4: ẋt,η ← x 1 -x 0 + t(1 -t) φt,η (x 0 , x 1 ) + (1 -2t)φ t,η (x 0 , x 1 ) 5:ℓ(η) ← ( ẋt,η ) ⊤ G(x t,η ; D) ẋt,η ▷ Estimate of objective L g (η) from eq. (6) 6:Update η using gradient ∇ η ℓ(η) return (approximate) geodesic interpolants parametrized by φ t,ηEQUATIONGiven interpolants that approximate γ * t and hence stay close to the data manifold, we can then rely on the CFM objective in eq. (1) to regress the vector field v θ . Since g makes the ambient space R d into a Riemannian manifold (R d , g), we need to replace the norm ∥•∥ in eq. ( 1) with the one ∥•∥ g induced by the metric, and rescale the marginals p 0 , p 1 using the volume form induced by g, so to extend p 0 , p 1 to densities in P(R d , g) (see Appendix §A). Similar arguments work for the joint distribution q. We can finally introduce our framework METRIC FLOW MATCHING that generalizes Conditional Flow Matching (1) to leverage any data-dependent metric g, by using interpolants x t,η (4), whose parameters η are obtained from minimizing the geodesic cost E g . The MFM objective can be stated as:EQUATIONA description of METRIC FLOW MATCHING is given in Algorithm 2. As the interpolants x t,η approximate geodesics of g, in MFM the vector field v t,θ is regressed on the data manifold M, where the underlying dynamics p * t is supported, resulting in better reconstructions. Crucially, eq. ( 6) only depends on time derivatives of x t,η . Therefore, MFM avoids simulations and simply requires training an additional (smaller) network φ t,η in eq. ( 4), which can be done prior to training v t,θ .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 in two ways. To begin with, MFM relies on the data or task inducing a 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 is 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 γ * t stay close to the data support, adopting RFM would entail replacing the MFM objective L MFM in (7) withEQUATIONHowever, as argued above, for almost any metric g on R d , geodesics γ * t can only be found via simulations, which in high dimensions inhibits the easy application of RFM. Conversely, MFM designs interpolants that minimize a geodesic cost (6) and hence approximate γ * t 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).We also introduce OT-MFM, a variant of MFM that leverages Optimal Transport to find a coupling q between p 0 and p 1 . 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.We consider a family of metrics g LAND 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 on the dataset D = {x i } N i=1 .Given ε > 0, we let x → g LAND (x) ≡ G ε (x) = (diag(h(x)) + εI) -1 be the "LAND" metric, whereEQUATIONwith σ the kernel size. We emphasize that G ε (x) was introduced by Arvanitidis et al. [2016] -from which we borrow the name-but its algorithmic use in MFM is fundamentally different. In line with Proposition 1, we see that ∥G ε (x)∥ is larger away from D, thus pushing geodesics (2) to stay close to the data support, as desired. While g LAND 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 a variation of g LAND of the form G RBF (x) = (diag( h(x)) + εI) -1 , whereEQUATIONwith K the number of clusters with centers xk and λ α,k the bandwidth of cluster k for channel α (see Appendix §C for details). In particular, h α is realized via a Radial Basis Function (RBF) network [Que and Belkin, 2016] , where ω α,k > 0 are learned to enforce the behavior h α (x i ) ≈ 1 for each data point x i so that the resulting metric g RBF assigns lower cost to regions close to the centers xk . In our experiments, we then rely on g LAND in low dimensions, and instead use g RBF 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 g RBF can be written as:EQUATIONWe see that the loss acts as a geometric regularization, with the velocity ẋt,η penalized more in regions away from the support of the dataset D, i.e. when λ α,k ∥x t,η -xk ∥ is large for all centers xk in eq. ( 10).OT-MFM. In eq. ( 7), samples x 0 , x 1 follow a joint distribution q, with marginals p 0 and p 1 . 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] , we focus on a coupling q that minimizes the distance in probability space between the source and target distributions. Namely, we consider the case where q is the 2-Wasserstein optimal transport plan π * from p 0 to p 1 [Villani et al., 2009] :EQUATIONwhere Π are the probability measures on R d ×R d with marginals p 0 and p 1 and c is any cost. While we could choose c based on g RBF , we instead select c to be the L 2 distance in R d to avoid additional computations and so we can study the role played by x t,η even when q = π * is agnostic of the data manifold.We then propose the OT-MFM objective, where η * minimizes the geodesic loss in eq. ( 11):EQUATION)We note that the case of g LAND is dealt with similarly. Additionally, different choices of the joint distribution q, beyond Optimal Transport, can be adapted from CFM [Tong et al., 2023b] to MFM in eq. ( 7).In MFM we learn interpolants x t that are optimal according to eq. ( 6). Previous works have studied the "optimality" of interpolants, but have ignored the data manifold. Shaul et al. [2023] proposed to learn interpolants x t that minimize the kinetic energy K( ẋt ) = E t,(x0,x1)∼q ∥ ẋt ∥ 2 . However, K 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 asL g (η) = E (x0,x1)∼q [ E g (x t,η )] ≡ E t,(x0,x1)∼q [ ∥ ẋt,η ∥ 2 g(xt,η)]. As a result, the geodesic loss L g (η) ≡ K g ( ẋt,η ) is precisely the kinetic energy of vector fields with respect to g and hence accounts for the cost induced by the data. In fact, choosing G(x) = I, recovers K. We note that the parameterization x t,η in eq. ( 4) is more expressive than x t = a(t)x 1 + b(t)x 0 , which is studied in Albergo and Vanden-Eijnden [2022] , Shaul et al. [2023] . Besides, x t,η not only depends on the endpoints x 0 , x 1 but, implicitly, on all the data points D through metrics such as g RBF .Data-dependent potentials. Energies more general than the kinetic one 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 to CFM, entails designing interpolants x t that minimize U(x t , ẋt ) = K( ẋt ) + V t (x t ), with V t a potential enforcing additional constraints. However, GSBM does not prescribe a general recipe for V t and instead leaves to the modeler the task of constructing V t , based on applications. Conversely, MFM relies on a Riemannian approach to propose an explicit objective, i.e. eq. ( 11), that holds irrespective of the task and can be learned prior to regressing the vector field v t,θ . In fact, eq. ( 11) can be rewritten asEQUATION)where V t,η is parametric function depending on the boundary points x 0 , x 1 (see Appendix §C.1 for an expression). In contrast to GSBM, MFM designs interpolants x t,η that jointly minimize K and a potential V t,η that is not fixed but also updated with the same parameters η to bend paths towards the data. 
experiments: 
 We test METRIC FLOW MATCHING on different tasks: artificial dynamic reconstruction and navigation 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, §E and §F.OT-CFM OT-MFM LAND The model. In all the experiments, we test the OT-MFM method detailed in §4.1. As argued in §4.1, for high-dimensional data, we leverage the RBF metric (10) and hence train with the objective (13). We refer to this model as OT-MFM RBF . Conversely, for low-dimensional data we rely on the LAND metric (9), replacing g RBF with g LAND in eq. ( 11) and eq. ( 13). We denote this model as OT-MFM LAND . Both metrics are 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 q, we tested MFM with independent coupling on the Arch task and the single cell datasets, denoted as I-MFM LAND and I-MFM RBF .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 to more meaningful matching than its Euclidean counterpart OT-CFM [Tong et al., 2023b] . For single-cell experiments instead, we also compare with a variety of baselines, including models specific to the single-cell domain [Tong et al., 2020, Koshizuka and Sato, 2023] , Schrödinger Bridge models [De Bortoli et al., 2021 , Shi et al., 2023 , Tong et al., 2023a] , and ODE-flow methods [Tong et al., 2023b , Neklyudov et al., 2023b] .Our first goal is to validate that MFM and specifically the family of interpolants x t,η 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) between the interpolation according to true dynamics and the marginal at time 1/2 obtained using the Euclidean baseline OT-CFM and our Riemannian model OT-MFM LAND . We see that OT-MFM LAND 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 ).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 a 2D sphere embedded in R 3 . We consider source and target distributions that are Gaussians centered at the sphere's poles, with samples lying exactly on 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 lower-dimensional 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-MFM LAND for navigations through surfaces scanned by LiDAR [Legg and Anderson, 2013] . From Figure 2 we find that straight paths result in 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 a task-agnostic metric.To test the advantages of MFM for more meaningful generation, we consider the task of unpaired image translation between dogs and cats in AFHQ [Choi et al., 2020] . 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-MFM RBF . Additionally, a quantitative comparison can be found in Table 2 , where we measure both the quality of images via Fréchet Inception Distance (FID) [Heusel et al., 2017] and the perceptual similarity of generated cats to source dogs via Learned Perceptual Image Patch Similarity (LPIPS) [Zhang et 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, a topic studied extensively §6. We finally test MFM for reconstructing cell dynamics, a central problem in biomedical applications [Lähnemann et al., 2020] , which holds great promise thanks to the advancements of single-cell RNA sequencing (scRNA-seq) [Macosko et al., 2015 , Klein et al., 2015] . Since in 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 v t,θ and the interpolants x t,η -see eq. ( 20) for how to extend x t,η to multiple timepoints. Following the setup of Schiebinger et al. [2019] , Tong et al. [2020 Tong et al. [ , 2023a] ] we 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-results 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 et al., 2018] . In fact, OT-MFM surpasses all baselines, including those that add biases such as stochasticity (SF 2 Tong et al. [2023a] ) or mass teleportation (WLF- UOT Neklyudov et al. [2023b] ). OT-MFM instead relies on metrics such as LAND and RBF to favor interpolations that remain close to the data.Question: Does the paper provide open access to the data and code, with sufficient instructions 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, a 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. Justification: Yes, such information can be found at the beginning of the Experimental Details appendix.• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or 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 as 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 it into the paper). 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 circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
related work: 
 Geometry-aware generative models. The manifold hypothesis [Bengio et al., 2013] has been studied in the context of manifold learning [Tenenbaum 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 Bortoli et al., 2021] 0.953 ± 0.140 0.862 ± 0.023 1.079 ± 0.117 SF 2 M-Sink [Tong et al., 2023a] 1.054 ± 0.087 1.198 ± 0.342 1.098 ± 0.308 SF 2 M-Geo [Tong et al., 2023a] 1.017 ± 0.104 0.879 ± 0.148 1.255 ± 0.179 SF 2 M-Exact [Tong et al., 2023a] 0.920 ± 0.049 0.793 ± 0.066 0.933 ± 0.054 OT-CFM [Tong et al., 2023b] 0.882 ± 0.058 0.790 ± 0.068 0.937 ± 0.054 I-CFM [Tong et al., 2023b] 0.965 ± 0.111 0.872 ± 0.087 1.085 ± 0.099 SB-CFM [Tong et al., 2023b] 1.067 ± 0.107 1.221 ± 0.380 1.129 ± 0.363 WLF-UOT [Neklyudov et 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-MFMLAND 0.916 ± 0.124 0.822 ± 0.042 1.053 ± 0.095 OT-MFMLAND 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 , Loaiza-Ganem 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 a 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 et al., 2020 , Bunne et al., 2022 , 2023 , Huguet et al., 2022 , Koshizuka and Sato, 2023] CNFs in a latent space to enforce that straight paths correspond to interpolations on the original data manifold. Finally, Scarvelis and Solomon [2023] propose a solution to 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 g LAND and g RBF introduced in §4. 
conclusions and limitations: 
 We have presented METRIC FLOW MATCHING, a simulation-free framework that generalizes Conditional Flow Matching to 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 metrics proposed in §4 is a benefit, we have not investigated how to further encode biases into the metric that are specific to the downstream task-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-dimensional 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.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 x 0 , x 1 , there exists a unique geodesic γ * t starting at x 0 and ending at x 1 . Integration on Riemannian manifolds. To introduce integration on Riemannian manifolds, we use the fact that under the technical assumption that M is orientable, the Riemannian manifold (M, g) has a canonical volume form dvol. This can be used to define a measure on M, where in local coordinates, we have that dvol(x) = |G(x)|dx where dx denotes the Lebesgue measure in R d and | • | denotes the determinant. Hence, for a chart (U i , φ i ) and a continuous function f : M → R which is compactly supported in U i , we define the integralEQUATIONThis definition can be extended to more general functions through the use of partitions of unity and the Riesz-Markov-Kakutani representation theorem.The Riemannian geometry of Metric Flow Matching. Finally, we note that in our work, we take our smooth manifold to be M = R d with the trivial chart U i = R d , φ i = 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 g through G directly in Definition 1 without the need to check consistency across the choice of local coordinates, and we have the trivial identification of T x M to R d . In addition, for a continuous f : M → R, we have the simplified definition of the Riemannian integral (when the integral exists) asEQUATION)We use this to define probability densities on (M, g) to extend CFM to our setting in §B.1. 
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 to biomedical applications. We cannot foresee immediate negative societal impacts associated with our work.Guidelines:• The answer NA means that there is 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 (e.g., deployment of technologies 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 not tied to particular applications, let alone deployments. However, if there is a 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, it 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 as 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 or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.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). 
safeguards: 
 • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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. 
 
 Appendix: To begin with, we provide a formal statement covering Proposition 1 and report a proof below. We introduce some notation. We write γ t ⊂ U when the image set of γ : Assume that for each (x 0 , x 1 ) ∼ q, there exists at least a path γ t connecting x 0 to x 1 whose length is at most Γ and such that γ t ⊂ B δ (D), with δ > 0. Let κ > 0, ρ > δ and g be any data-dependent metric satisfying:[0, 1] → R d is contained in U , i.e.(i) v ⊤ G(x; D)v ≥ κ∥v∥ 2 for each x ∈ R d \ B ρ (D) and v ∈ R d ; (ii) ∥G(x; D)∥ ≤ κ(ρ/Γ) 2 for each x ∈ B δ (D).Then for any (x 0 , x 1 ) ∼ q, the geodesic γ * t of g connecting x 0 to x 1 satisfies γ * t ⊂ B 2ρ (D).Proof of Theorem B.1. We argue by contradiction and assume that there exist x ∈ R d \ B 2ρ (D) and (x 0 , x 1 ) ∼ q such that the geodesic γ * t of g connecting x 0 to x 1 passes through x, i.e. there is a time t 0 ∈ (0, 1) such that γ * t0 = x. By assumption, there exists a path γ t that connects x 0 to x 1 and stays within B δ (D), which means thatx 0 = γ 0 = γ * 0 ∈ B δ (D). Since D is closed, the function x → d E (x, D) := inf y∈D ∥x -y∥, i.e. the Euclidean distance of x from the dataset D, is continuous. In particular, t → d E,D (t) := d E (γ * t , D) is also continuous, due to the geodesic being a smooth function in the interval [0, 1]. From the continuity of d E,D and the fact that d E,D (0) ≤ δ and d E,D (t 0 ) > 2ρ, it follows that there must be a time 0 < t ′ < t 0 such that d E,D (t) ≥ ρ for all t ∈ (t ′ , t 0 ] and d E,D (t ′ ) = ρ.If we unpack the definition of d E,D , we have just shown that γ * t ∈ R d \ B ρ (D) for all t ∈ (t ′ , t 0 ]. Accordingly, we can estimate the length of γ * t with respect to the Riemannian metric g asLen g (γ * t ) = 1 0 ∥ γ * t ∥ g(γ * t ) dt > t0 t ′ ∥ γ * t ∥ g(γ * t ) dt = t0 t ′ ( γ * t ) ⊤ G(γ * t ; D) γ * t dt ≥ √ κ t0 t ′ ∥ γ * t ∥ dt ≥ √ κ∥γ * t ′ -γ * t0 ∥,where in the second-to-last inequality we have used the assumption (i) on the metric having minimal eigenvalue √ κ in R d \ B ρ (D), while the final inequality simply follows from the Euclidean length of any curve between the points γ * t ′ and γ * t0 being larger than their Euclidean distance. We claim that ∥γ * t ′ -γ * t0 ∥ ≥ ρ. To validate the latter point, take ϵ > 0. It follows that there existsx i ∈ D such that ∥x i -γ * t ′ ∥ ≤ ρ + ϵ, because d E,D (t ′ ) = ρ, i.e. γ * t ′ ∈ B ρ (D). If ∥γ * t ′ -γ * t0∥ < ρ, then we could apply the triangle inequality and derive that∥γ * t0 -x i ∥ ≤ ∥γ * t0 -γ * t ′ ∥ + ∥γ * t ′ -x i ∥ < ρ + ρ + ϵ = 2ρ + ϵ.Since ϵ > 0 was arbitrary, it would follow that γ * t0 ≡ x ∈ B 2ρ (D), which is a contradiction to our starting point. Therefore, ∥γ * t ′ -γ * t0 ∥ ≥ ρ, and henceEQUATION)On the other hand, the assumptions guarantee the existence of another path γ t connecting x 0 to x 1 whose image is always contained in B δ (D). It follows thatLen g (γ t ) = 1 0 ∥ γt ∥ g(γt) dt = 1 0 ( γ * t ) ⊤ G(γ * t ; D) γ * t dt ≤ √ κ ρ Γ 1 0 ∥ γt ∥ dt = √ κ ρ Γ Γ = √ κρ,where we have used the assumption (ii) on g and that the length of γ t is at most Γ. We can then combine the last inequality and eq. ( 19), and conclude thatLen g (γ t ) ≤ √ κρ < Len g (γ * t ),which means that we have found a path connecting x 0 to x 1 , whose length with respect to the metric g is shorter than the one of the geodesic γ * t from x 0 to x 1 . This is a contradiction and concludes the proof.In this subsection, we demonstrate how to generalize the CFM objective in eq. ( 1) to the Riemannian manifold (R d , g), defined by a data-dependent metric g as per Definition 1.Densities on the manifold. To begin with, we extend the densities p 0 , p 1 and the joint density q to valid densities on the manifold (R d , g). First, we recall that the Riemannian volume form induced by g can be written in coordinates asdvol(x) = |G(x; D)|dx,where | • | denotes the determinant of the matrix G(x; D), while dx is the standard Lebesgue measure on R d . Accordingly, given a density p ∈ P(R d ), we can derive the associated density p ∈ P(R d , g) by rescaling:p(x) := p(x) |G(x; D)| ,which guarantees that p is continuous, positive, andR d p(x)dvol(x) = R d p(x)dx = 1.This in particular applies to the marginals p 0 , p 1 given by our problem. A similar argument applies to the joint density q whose marginals are p 0 and p 1 , respectively. In fact, we can now define a density q on the product manifold, i.e. q ∈ P(R d × R d , g × g), by simply taking q(x 0 , x 1 ) := q(x 0 , x 1 )|G(x 0 ; D) • G(x 1 ; D)| ,where the denominator is exactly the pointwise volume form induced by the product metric g × g.Therefore, the MFM objective in ( 7) is interpreted asL MFM (θ) = E t,(x0,x1)∼q ∥v t,θ (x t,η * ) -ẋt,η * ∥ 2 g(x t,η * ) = E t R d ×R d ∥v t,θ (x t,η * ) -ẋt,η * ∥ 2 g(x t,η * ) q(x 0 , x 1 ) |G(x 0 ; D) • G(x 1 ; D)|dx 0 dx 1 = E t R d ×R d ∥v t,θ (x t,η * ) -ẋt,η * ∥ 2 g(x t,η * ) q(x 0 , x 1 )dx 0 dx 1 = E t,(x0,x1)∼q ∥v t,θ (x t,η * ) -ẋt,η * ∥ 2 g(x t,η * ), 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 a density on the manifold (R d , g), the same as an expectation over the density q 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 (x; D) to avoid the need for initialization schemes that are specific to the metric. In fact, introducing a high-dimensional, nontrivial norm ∥ • ∥ g in the CFM objective could introduce instabilities in the optimization of CFM. As such, the objective in L MFM effectively reduces to L CFM with interpolants x t,η * minimizing the geodesic loss L g . Equivalently, METRIC FLOW MATCHING consists of a 2-step procedure: we first learn interpolants that approximate the geodesics of a data-dependent metric (Algorithm 1), and then we regress the vector field v θ using the interpolants from the first stage as per the standard CFM objective. A pseudocode for the MFM-pipeline is reported in Algorithm 2.Inference. We finally note that, in principle, the differential equation generated by v θ is defined over the tangent bundle of (R d , g) 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 R d , the Euclidean Euler-step integration provides a firstorder approximation of the Riemannian exponential associated with g (see for example Monera et al. [2014] ). Accordingly, at inference, we can simply rely on the canonical Euler-discrete integration to approximate the exponential map associated with the data-dependent metric g, provided that the step size is small.Require: coupling q, initialized network v t,θ , trained network φ t,η , data-dependent metric g(•)1: while Training do 2:Sample (x 0 , x 1 ) ∼ q and t ∼ U(0, 1)3:x t,η ← (1 -t)x 0 + tx 1 + t(1 -t)φ t,η (x 0 , x 1 ) ▷ eq. (4) 4: ẋt,η ← x 1 -x 0 + t(1 -t) φt,η (x 0 , x 1 ) + (1 -2t)φ t,η (x 0 , x 1 ) 5: ℓ(θ) ← ∥v t,θ (x t,η ) -ẋt,η ∥ 2 g(xt,η)▷ Estimate of objective L MFM (θ) from eq. ( 7)6:Update θ using gradient ∇ θ ℓ(θ) return vector field v t,θWe conclude this section by justifying our claims about the relation of the support of the probability path p t and the interpolants. Consider a family of conditional paths p t (x|x 0 , x 1 ) such that p t (x|x 0 , x 1 ) ≈ δ(x-x t ), where x t are the interpolants connecting x 0 to x 1 , and δ is the Dirac distribution. Assume that x t are straight interpolants, i.e. for any (x 0 , x 1 ) ∼ q, we define x t = tx 1 +(1-t)x 0 . Let us now consider y ∈ R d such that there is no (x 0 , x 1 ) ∼ q : y = tx 1 + (1 -t)x 0 . Accordingly:p t (y) := R d ×R d p t (y|x 0 , x 1 )q(x 0 , x 1 )dx 0 dx 1 = R d ×R d δ(y -x t )q(x 0 , x 1 )dx 0 dx 1 = 0,since there is no (x 0 , x 1 ) in the support of q, such that y = x t . We have then shown that supp(p t ) ⊂ {y ∈ R d : ∃ (x 0 , x 1 ) ∼ q : y = tx 1 + (1 -t)x 0 }, which can be limiting for tasks where the data induce a nonlinear geometry, as validated in Section 5.In this Section, we provide further details on the diagonal metrics introduced in 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 g RBF . 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 λ k associated with cluster C k as:λ k = 1 2 κ |C k | x∈C k ∥x -xk ∥ 2 -2, where κ is a 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 ω α,k are determined by training the loss function,L RBF ({ω α,k }) = xi∈D 1 -hα (x i ) 2 = xi∈D 1 - K k=1 ω α,k exp - λ α,k 2 ∥x i -xk ∥ 2 2 .Two important comments are in order. First, by learning the RBF metric, the framework MFM effectively entails jointly learning a data-dependent Riemannian metric and a suitable matching objective along approximate geodesics x t,η . 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 et 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 is a key advantage in high-dimensions 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 is neither a benefit nor a 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 in eq. ( 6) could always bias to learn interpolants x t,η moving through these regions independent of the starting point-this is a 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 as an issue in practice whenever we adopted the LAND metric for low-dimensional data.We detail how the geometric loss (6) used to learn interpolants x t,η can be recast as a generalization of the GSBM framework in Liu et al. [2024] . In general, for any data-dependent metric g, we can indeed rewrite L g (η) in eq. ( 6) as:L g (η) = E t,(x0,x1)∼q [⟨ ẋt,η , G(x t,η ; D) ẋt,η ⟩] = E t,(x0,x1)∼q ∥ ẋt,η ∥ 2 + ⟨ ẋt,η , (G(x t,η ; D) -I) ẋt,η ⟩ = E t,(x0,x1)∼q ∥ ẋt,η ∥ 2 + V t,η (x t,η , x 0 , x 1 ) ,where the parametric potential V t,η has the formV t,η (x t,η , x 0 , x 1 ) = ⟨ ẋt,η , (G(x t,η ; D) -I) ẋt,η ⟩ ẋt,η = x 1 -x 0 + t(1 -t) φt,η (x 0 , x 1 ) + (1 -2t)φ t,η (x 0 , x 1 ),where we have used the parameterization of x t,η in eq. ( 4). By replacing G(x; D) with the explicit expression given by the RBF metric in eq. ( 10), this provides a concrete formulation for eq. ( 14). We also note that an equivalent perspective amounts to replacing the parametric potential with a fixed function U t (x t,η , ẋt,η ) that depends not only on x t,η as the potentials in Liu et al. [2024] , but also on the velocities ẋt,η . Once again, we emphasize that our framework prescribes explicit parametric potentials V t,η via the diagonal Riemannian metrics in Section 3.1 and hence differs from Liu et al.[2024] which leave to the user the task of designing potentials based on applications.We parameterize the models φ t,η (x 0 , x 1 ) in eq. ( 4) and v t,θ (x t ) in eq. ( 7) as neural networks, and train them separately and sequentially. We start with φ t,η (x 0 , x 1 ), so that interpolants are learned prior to regressing v t,θ (x t ) 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 higher-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 V100 GPUs. Training time for AFHQ varies by GPU, ranging from 12 hours (A100) to 1 day (V100).In the synthetic Arch, Sphere and LiDAR experiments, we parameterized both φ t,η (x 0 , x 1 ) and v t,θ (x t ) 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 a learning rate of 0.0001 for φ t,η (x 0 , x 1 ) and the AdamW optimizer Loshchilov and Hutter [2017] with a learning rate of 10 -3 and a weight decay of 10 -5 for v t,θ (x t ). We used a 90%/10% 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 σ = 0.125 and ϵ = 0.001 in the LAND metric. During inference, we solved for p t using the Euler integrator for 100 steps.To generate the Arch dataset, we follow the experimental setup from [Tong et al., 2020] . We sampled 5000 points from two half Gaussians N (0, 1 2π ) and N (1, 1 2π ). The exact optimal transport interpolant at t = 1 2 was used as the test distribution. We then embedded the points on a 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 (0, 1 2π ) and N (1, 1 2π ), and then scaled them by π. We sampled longitudes uniformly. The exact optimal transport interpolant at t = 1 2 was used as the test distribution. We then embedded the points on a sphere of radius 1 without any noise. The Earth Mover'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, 5] 3 ⊂ R 3 , 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.In the unpaired translation experiments, we utilized the U-Net architecture setup from Dhariwal and Nichol [2021] for both φ t,η (x 0 , x 1 ) and v t,θ (x t ). 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 φ t,η (x 0 , x 1 ) based on training loss. During inference, we solved for p t at t = 1 using the adaptive step-size solver Tsit5 with 100 steps. We trained the RBF metric with κ = 0.5 and ϵ = 0.0001 with k-means clustering, as described in Appendix C. For this experiment, we enforce stronger bending by using ( hα (x)) 8 , in the loss function L g RBF (η) in (11). Note that higher powers of h ensures that regions away from the support of D, where h < 1, are penalized even more in the geodesic objective, which highlights the role played by the biases introduced via the metric.Multiple constraints setting Following a similar approach to Neklyudov et al. [2023b] , we modified our sampling procedure to interpolate between two intermediate dataset marginals, with neural network parameters η shared across timesteps. The interpolation is defined as:x t = t i+1 -t t i+1 -t i x ti + t -t i t i+1 -t i x ti+1 + 1 - t i+1 -t t i+1 -t i 2 - t -t i t i+1 -t i 2 φ t,η (x 0 , x 1 ). (20)Low-Dimensional experiments We parameterized both φ t,η (x 0 , x 1 ) and v t,θ (x t ) as 3-layer MLP networks with a width of 64 and SeLU activation. In the LAND metric, we 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, we used σ = 0.25 and ϵ = 0.001.High-Dimensional experiments Both φ t,η (x 0 , x 1 ) and v t,θ (x t ) were parameterized as 3-layer MLP networks with a width of 1024 and SeLU activation. We set κ = 1.5 and ϵ to 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.We report additional results for single-cell reconstruction using 50 principal components in Table 6 . Again, we note that OT-MFM performs strongly, and it is marginally surpassed only by SF 2 on Multi(50D) (results have not been reproduced). Two important remarks are in order. First, the baseline SF 2 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, we highlight that as we move to more realistic high-dimensional settings (100 principal components, shown in Table 3 ) the advantages of our framework become even more apparent. G Supplementary Sphere Results We visualize the problem setting in Figure 5 . We report the Earth Mover's Distance between the sampled and test sets, as well as the mean distance of the middle points of trajectories to the sphere, to quantitatively compare CFM and MFM on the task. MFM improves significantly over the Euclidean baseline, CFM (Table 8 ). Furthermore, the samples generated by MFM at intermediate times are much closer to the underlying sphere than the Euclidean counterparts (Table 7 ). • The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be 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 a 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 a specific model and empirical evaluation, it may be necessary to either make it 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., 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 submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe 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 or verifying the results.With the technical assumptions that M is second-countable and Hausdorff.Formally, this is a choice of a smooth symmetric covariant 2-tensor field on M which is positive definite at each point. 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: ### Clarity

This paper lacks clarity in its exposition, and are not sufficiently careful in stating its claims.

- 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. 

- 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?

### Novelty

Some 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).

- 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. 

- 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.

### Empirical validation

I feel the experimental results are messy and do not provide a coherent analysis.

For 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.

Going through each experiment section, there are concerns regarding the setup or results:

1) The LiDAR scan. 
- 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.
- 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.

2) The AFHQ experiment. 
- 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.
- 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? 
- 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.

3) The single-cell experiment.
- 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?
- 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.
- 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?
- 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.
questions: questions: - How effective is the interpolant when no optimal coupling is used, especially for multimodal data distributions?
- Does learning a full flow matching model outperform just using the L2 optimal coupling and using the interpolant to create x_t samples?
- 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?
- 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?
- 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)?
- Can this framework be adapted to learn a low dimensional manifold (rather than just a metric)?
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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).

Given 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.
questions: questions: 1. Taking into account the weakness mentioned above regarding novelty: could the authors clarify their contributions beyond Eq. (4)?
2. 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).
3. 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.

Minor comments:
- 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.
- [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: limitations: The authors have sufficiently addressed the limitations in their work.
paper_summary: paper_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.
strength_weakness: strength_weakness: I have some minor complaints:
1. 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
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?
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?
4. 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
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
questions: questions: 1. Line 125: "regular" means smooth in some way? Can you be more specific?
2. Why is it preferable to use LAND for lower dimensions?
3. 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: 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)
paper_summary: paper_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.
strength_weakness: strength_weakness: *Major comments*

I have two main questions and comments, which I expect to be addressed in the revised version.

* 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.

* 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. 
In 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.

*Minor comments*

I appreciate comprehensive reviews in section 6. There are a few missing references, and I sugges to include them. 
* Geometry-aware generative models:
   - 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), 
   - 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)
   - geometry-based regularization for learning accurate manifolds: Neighborhood Reconstructing Autoencoders (NeurIPS 2021) and On explicit curvature regularization in deep generative models (TAG-ML 2023).
questions: 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: limitations: Yes

Comments: title: Thank you for additional comments (part 2)
Thank you for the additional comments
Thank you
Kindly awaiting more feedback
Response to official comment by Reviewer mgNS
Re: Response
Responses
Rebuttal (3/3)
Rebuttal (2/3)
comment: > Cost and Complexity: Here is a simple question: Given the LAND or RBF metric definition, do you have access to all geodesics in closed form? If so, then I agree the method would be simulation-free and very computationally efficient; 


We politely disagree with the reviewer. Simulation-free training means that we are able to jump to any time $t$ of the interpolants vector field $u\_t (x|z)$ without simulating the path up to $<t$. **An interpolant does not need to be a geodesic for simulation free training to still hold**. Our work never sought to claim that MFM is able to extract exact geodesics in a simulation free manner but rather the learned interpolants for estimating the metric-induced velocity does not require solving a differential equation (and/or backpropagating through it) and learning requires only pointwise evaluations. 

> The proposed workaround is to actually fit a parametric approximation to the geodesics …

Certainly, if our network successfully finds the global optimum by minimising the Dirichlet energy then we do recover a geodesic of the data manifold. The step of approximating geodesics instead of simulating the Euler Lagrange equations is a key contribution of our work that shows one can leverage benefits from metrics via simulation-free training to lead to measurable benefits for trajectory inference tasks. 
We disagree with the reviewer’s assessment that this cost is expensive as it can be thought of as a 1 time pre-processing step which notably is not more expensive than the original trajectory inference problem itself. We argue that the performance gains of MFM in domains such as single-cell justify this pre-processing overhead.

### Regarding the manifold hypothesis 


> The proposed method works in the original data space.

In general, the reviewer is correct in identifying that if we care about *explicitly* parameterizing a lower-dimensional space, then one needs coordinate maps.For the sphere, the fact that the points do not lie *exactly* on the surface  is expected, considering that we are assuming we do not know the underlying manifold. In practice, for the application of trajectory inference one does not need the explicit manifold, but simply cares about reconstructing trajectories whose intermediate points are closer to the underlying manifold. Empirically, we have shown the impact of our metric with the single-cell experiments, where it is even hard to define what the ground-truth manifold is and  adopting the metric approach is favorable to finding an explicit parameterization of the lower-dimensional manifold since it requires the notion of a norm in the ambient space.

We emphasize that “just a metric in the ambient space” can actually induce a non-trivial Riemann tensor and hence a “proper” Riemannian manifold. More generally, the idea of learning a distance in the ambient space rather than parameterising a lower dimensional space **is not a new approach and we are certainly not the first ones arguing for it**, but plenty of previous works (referenced in our manuscript) have adopted this angle (referred to as **metric learning**). As such, we believe our work is using this perspective in the context of generative modeling and has been substantially motivated by previous literature.
We thank the reviewer for taking the time to continue engaging with us. We believe certain aspects of our response may have added confusion rather than clarity which we apologize for. We now attempt to answer the points raised by the reviewer’s latest response.

> "to the best of our knowledge, existing works do not use parametric metrics" …  But I believe here the authors are using parametric as a synonym for learned; 


Yes, “parametric” here refers to a metric that learned via gradient based optimization. In contrast, we use non-parametric to refer to a metric that does not require learning via gradient based optimization.

>however, even then, this claim is certainly false. There has been many works on learning parametric metrics (...) and many works on using latent spaces for transport problems where the metric is induced by the autoencoder  

We also agree with the reviewer that our work is not the first to learn a data-dependent metric. Our principal claim, that we clarify, is that MFM is the first method that employs a data-dependent (parametric or non-parametric) Riemannian metric within the context of flow matching. While it is certainly true that DSBM and GSBM have experiments that use the latent space of an auto-encoder, they do not invoke the Riemannian formalism which would have required the use of the metric to define distances, vector fields, and interpolants. As a result, we view the methodology of GSBM in its techniques to bias paths complementary to the approach presented in this paper. Indeed, there are similarities between the goals of our work and GSBM—which will be included in a larger discussion in the main paper—but crucially the key difference lies in the objective and mechanism used to bias paths. We understand how these techniques may not have been initially clear but we will use this response as a guide to further increase clarity and transparency of our novelty claims and the difference between methods.

> the authors did not address my concern on reproducing FID: DSBM reports 14.16 and GSBM reports 12.39 FID.

Please note we did attempt to answer this in our original rebuttal but we now include additional detail that we believe could strengthen the original response. We first highlight the 3 main reasons why our experimental setup differs from GSBM which lead to discrepancies between FID numbers between GSBM and MFM: 
- The chief reason is that in our experiment we use the latent space of a pre-trained VAE autoencoder in StableDiffusionV1 as opposed to a more shallow VAE used in GSBM. As a result these GSBM and MFM in our current draft do generative modeling using two very different latent spaces. 
- The GSBM codebase resizes each input image to 64 x 64 while we resized it to 128 x 128 to use our pretrained autoencoder.
- GSBM and DSBM operate directly on the ambient space and do not actually generate via the decoder of the VAE. In MFM as we heavily use the Riemannian metric for latent flow matching we generate using the decoder of the VAE.

We also took this time to investigate the reviewers’ great suggestion and attempted to compute an FID using the *saved checkpoints of the UNET and VAE from GSBM*  with their provided sampling notebook. Unfortunately, we were unable to reproduce the FID of 14.16 and achieved an FID of 29.5 which is lower but closer to the one we report.

> Instead, it seems the authors have deliberately chosen not to compare against the different choices of potentials in these works.

The purpose of the image translation experiments was to compare CFM with MFM. In particular, we wanted to show how the **same** metric, i.e. RBF, can be fit to different applications and does not require adjustments specifics to the task. We hope the reviewer can see the merit in that. Regarding potential comparisons with GSBM for example, the setups are actually entirely different so simply comparing reported FID would not be accurate (see our previous point for more details). 


> "in favor of potentials (that are generally not defined)" --> I am not sure what is meant by not defined here

We apologize for the vague terminology. By “not defined”, we mean as in the case of GSBM that the potential has to be defined by the modeler. In our setting instead, we wanted to propose a choice of metrics e.g. LAND and RBF that could be readily applied to different downstream tasks, without explicit encoding of the task outside of using data samples. 

> Actually, it is worth highlighting that my concern regarding the LiDAR experiment has also not been addressed (...)

We appreciate the reviewers' comments. We note that LAND in itself is not a contribution of this work but rather its use in flow matching. The main purpose of LiDAR is to visually demonstrate the impact of using a data dependent metric (in contrast to CFM) on the trajectory in flow matching. We argue that the current experiment adequately serves this goal and does not gain additional benefit from comparing to GSBM.
Unfortunately, the authors did not seem to understand my concerns about clarity and the paper's claims. I want to clarify that the only concern that was addressed by the rebuttal was that the choice of metric has advantages regardless of the choice of coupling used during training. 

**Metric:** 

"to the best of our knowledge, existing works do not use parametric metrics" --> To be clear, by parametric (in contrast to non-parametric) I simply meant that the metrics do not depend on a large data set. But I believe here the authors are using parametric as a synonym for learned; however, even then, **this claim is certainly false**. There has been many works on learning parametric metrics over data sets / point clouds [1,2], and many works on using latent spaces for transport problems where the metric is induced by the autoencoder, such as the DSBM and GSBM methods [3,4]. **Relatedly, it may be worth noting that the authors did not address my concern on reproducing FID: DSBM reports 14.16 and GSBM reports 12.39 FID**. Both are significantly better than the values reported by this paper (37-41), and this is especially concerning because existing works have open source implementations available. Instead, it seems the authors have deliberately chosen not to compare against *the different choices of potentials* in these works.

"in favor of potentials (that are generally not defined)" --> I am not sure what is meant by not defined here, as I understand these prior works work for any potential. **Since the potentials can be arbitrarily defined, it fully includes the Lagrangian point of view of Riemannian geodesics.** As mentioned in Appendix C.1, one can always convert from the cost defined using a Riemannian metric to the regular Euclidean metric plus a potential.

"do not adopt a Riemannian formalism to assess the velocity of the path" --> I believe the authors are describing how to extract a Riemannian manifold from a data set. There is significant work in this area (see the literature on constructing Riemannian metrics from point clouds e.g. [1, 2]) and is often used for LiDAR. Actually, it is worth highlighting that **my concern regarding the LiDAR experiment has also not been addressed**: compared to the GSBM experimental result which show two modes of trajectories [4; Figure 4], the results from this paper only show one mode: why is this? For me, I am not convinced that this LAND metric is as good as the choices of potential compared to existing works.

**Cost and Complexity:** Here is a simple question: Given the LAND or RBF metric definition, do you have access to all geodesics in closed form? If so, then I agree the method would be simulation-free and very computationally efficient; however, from my understanding, **the answer is no. And the proposed workaround is to actually fit a parametric approximation to the geodesics.** This, as the authors also agree, have significant computational overhead and can increase costs to double the amount for CFM due to this extra step of needing to approximate geodesics with neural networks.

**Manifold hypothesis:** Taking wikipedia as a proper "definition", the manifold hypothesis, if it were to be true, would imply at least two things:

1) "Machine learning models only have to fit relatively simple, low-dimensional, highly structured subspaces within their potential input space (latent manifolds)."

This, as also reflected by the rebuttal, does not hold for the proposed method. The proposed method works in the original data space. 

2) "Within one of these manifolds, it’s always possible to interpolate between two inputs, that is to say, morph one into another via a continuous path along which all points fall on the manifold."

As noted in the review, and also agreed upon by the rebuttal response, the intermediate samples (see Figure 3 for images, and the rebuttal pdf for sphere and LiDAR) certainly do not seem to lie on the data manifold. *This is a subtle point that I raised which I believe the authors have not understood.* **If one were to actually define a manifold (along with a tangent plane and a corresponding metric on that plane), then interpolations will never leave the manifold, no matter how poorly the model/distribution is fit.** This is the key idea behind the literature on constructing Riemannian manifolds as described above, where we can fully construct a proper manifold and not just a metric in the ambient space. However, the proposed approach does not learn a subspace and due to this, we see that the points are some distance away from the manifold (shown for both sphere and LiDAR visualizations in the rebuttal pdf). 

[1] "Neural FIM for learning Fisher information metrics from point cloud data" Fasina et al. 2023.

[2] "Approximating the Riemannian Metric from Point Clouds via Manifold Moving Least Squares" Sober et al 2020. 

[3] "Diffusion Schrödinger Bridge Matching" Shi et al 2023.

[4] "Generalized Schrödinger Bridge Matching" Liu et al 2023.
Thank you for taking the time to engage with us further during this rebuttal. We also appreciate the increase in score. However, we would like to take this last chance to clarify some misconceptions about our work you have raised.  

-   **Metric:** We would like to stress that, to the best of our knowledge, existing works do not use parametric metrics. GSBM or WLF argue in favor of potentials (that are generally not defined), but do not adopt a Riemannian formalism to assess the velocity of the path, which we believe to be what sets our work apart, independent of the choice of the metric. We note that we consider both parametric and non-parametric metrics in our work, as outlined in section 4.1 and Appendix C of the original manuscript.
-   **Cost and Complexity:**  We would like to politely note that MFM does not require numerical solving of trajectories at any stage of the training since it is a simulation free algorithm. Further, the complexity of the training is at most twice the standard CFM, as we always use the same network architecture for both training phases. At inference, MFM does not require metric meaning it is exactly as efficient as the standard CFM. All these points, make MFM efficiency comparable with GSBM and WLF.
-   **Manifold Hypothesis:**  The results from the arch experiment lifted onto a sphere (general rebuttal), suggest that MFM successfully preserves the underlying geometry without explicitly inducing a lower-dimensional manifold --- our approach maintains proximity to the manifold without the need for explicit parameterization.

We hope this message helps to clarify the lingering concerns. We will add key points discussed in the rebuttal, including extended comparisons with GSBM and WLF, in the revised version of the manuscript.
Thanks for the additional ablations. I have increased my score, and would further encourage providing additional clarity and honesty in highlighting the contributions. I agree with the points raised by the authors in the rebuttal: (1) the use of a predefined non-parameteric metric is the main concept that sets this apart from existing works which only use parametric metrics, (2) however this comes at a higher computational cost compared to existing works due to (i) numerically solving individual trajectories as well as (ii) the evaluation of a non-parametric metric also involving neural network evaluations, and (3) the proposed method is not fully justified as a solution to the typical manifold hypothesis discussed in machine learning (i.e., computational cost is not dependent on the data manifold but still on the original space) and is focused purely on learning a metric in the original data space.
We thank the reviewer for engaging with us, for the detailed and valuable feedback, and finally for endorsing our submission.
We thank the reviewer again for your time and feedback that allowed us to strengthen the paper with new experiments and clarifications during this important rebuttal period. As the end of the rebuttal period is fast approaching we were wondering if our answers in the rebuttal were sufficient enough to address the important concerns raised regarding 1.) clarity of our claims, 2.) the technical novelty that distinguishes our proposed approach MFM and Wasserstein Lagrangian Flows/GSBM, and 3.) the experimental evaluation. We highlight that our global response includes new ablations on the use of couplings following the great suggestions by the reviewer.

We would be happy to engage in any further discussion that the reviewer finds pertinent, please let us know! Finally, we are very appreciative of your time and effort in this rebuttal period and hope our answers are detailed enough for the reviewer to consider a fresher evaluation of our work with a potential score upgrade if it's merited.
Thank you for your detailed rebuttal. I believe you've answered all my questions. Since most of my comments were minor details, my overall judgement of the paper has not changed, so I will leave my score as is.
We appreciate the reviewer taking the time to engage with us further during this rebuttal.
Our goal with the new ARCH on a sphere experiment was to test whether trajectories with MFM with the LAND metric lie on the sphere *without actually parametrizing the sphere*. Given this idealized setting, we see that MFM matches our intuitions and generates samples that lie on the sphere.
If the spherical inductive bias ---i.e. the exact parametrization of the sphere---was known to us apriori then RFM and MFM may coincide if the sampling of data demands it. However, note that if data is irregularly sampled on the sphere then trajectories can themselves bend on the sphere, and not obey "shortest paths" on the sphere which is a setting RFM cannot model but is handled by MFM.


In general, we agree with the reviewer's point that in many practical settings data may not be sampled homogeneously on an underlying manifold, and in these cases, the LAND metric may not be optimal. We would like to note however that our MFM framework can be instantiated with any choice of metric, including a parametrized/learned one such as RBF which we found works well for higher dimensions. The key point we highlight is that any metric used in MFM is informed by the sampling of data (training set) which alleviates the need for an exact parametrization of the manifold.

We thank the reviewer again for allowing us to clarify these technical aspects of MFM. We would be happy to answer any further questions the reviewer has, otherwise we politely encourage the reviewer to consider increasing their score as they originally suggested they are inclined to.
Dear authors,

I am in general positive with the rebuttal. Here are a few additional comments.

> We emphasize that we manage to attain this without explicitly parameterizing the lower-dimensional space but simply relying on the LAND metric.

Yes, that is indeed interesting. We must however be a bit skeptical and say that in general it is hard to assume that the data is sampled so nicely along the underlying manifold. I wonder which amount of jitter or lack of proper sampling would make the LAND metric unreliable with real data. As such, one would then put the RFM requirement of having the exact metric versus the LAND approach of relying on the data as both **desirable**, but not guaranteed to be available/satisfied. I originally meant the experiment to compare that exact thing to see how much we "sacrifice" compared to ground truth by relying purely on data, but the example provided is already showing something positive in my opinion.

After reading the other reviews and their rebuttals, I am inclined to increase my score to Accept.
We thank the reviewer for their time and engaging with us in the rebuttal. We are glad that the reviewer has found our rebuttal to increase the clarity of the paper and we hope the reviewer can also upgrade their score as they mention in their rebuttal response. We are also more than happy to answer any lingering questions the reviewer has, please let us know!
Thank you for addressing all the points that I raised. I believe the clarity of the paper will increase, and I am therefore inclined to raise the score to accept.
> _When the optimal coupling is given, is there even a need to learn the flow?_

That's an interesting suggestion! Unfortunately, single cell RNA is known to have a destructive generative process which means that different time marginals can contain varying numbers of *unpaired* observations. Operationally, this means there is no 1-1 correspondence between populations at different times. Consequently, evolving the particles along a path $x\_t$ may be an ill-posed problem and instead, we must assess the probability path $p\_t$. We achieve that by taking the pushforward of $p_0$ using the flow generated by the vector field $v_\theta$ learned using flow matching. We understand how this technical point may not have been clear given the complex nature of single cell data and now hope the reviewer agrees that their suggestion, while interesting, cannot be employed in this particular experiment.

> _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. (...)_

We appreciate the reviewer's concern. We politely point out that trajectory inference using PCA of single-cell data is not an innovation of our work but is a large subfield in its own right (see [2],[3]). Practitioners often resort to PCA as actual single cell data is both very high dimensional and noisy which provides a computationally intractable domain for learning trajectories prior to simulation free generative modeling. In this context we argue our reported benchmarks are standard in the literature. We also point again to our new experiments on the sphere (see **general rebuttal**), highlighting the ability of MFM to learn trajectories that stay closer to an unknown, underlying manifold.


### Questions

> _How effective is the interpolant when no optimal coupling is used, especially for multimodal data distributions?_ 

Please see our response/additional experiments on using independent coupling.


> _Does learning a full flow matching model outperform just using the L2 optimal coupling and using the interpolant to create x_t samples?_

Please see our response above on why using just the interpolants for single cell RNA can be reductive due to the lack of a 1-1 correspondence between samples and the need to evaluate the time-evolution of the **population density**.

> _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?_
	
Please see the reply to the first point of single cell experiment setup.

> _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?_

Please see the reply to the first point of single cell experiment setup.


> _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)?_

This is an interesting question and we believe that given the setting you described, if the supports of the marginal distribution is fully separated, then there is no meaningful signal for the nonlinear correction $\varphi\_{t,\eta}$ in Eq. (4) and hence the path should remain linear, perhaps after adding a penalization over the norm of the weights of the MLP $\varphi\_{t,\eta}$. 

> _Can this framework be adapted to learn a low dimensional manifold (rather than just a metric)?_

Please see our response above on why MFM is not designed to learn a low dimensional manifold and why this is not an issue in practice (we point again to the new experiments of MFM on the sphere).

### Final comment

We thank the reviewer again for their valuable feedback and great questions that enabled us to include new results that have strengthened our paper. We hope that our rebuttal addresses their questions and concerns—particularly in regard to how MFM does not need to prescribe a lower-dimensional space, as shown by the new experiments on the sphere, and the comparison with CFM using independent couplings. We kindly ask the reviewer to consider upgrading their score if the reviewer is satisfied with our responses. We are also more than happy to answer any further questions that arise.

### References
 
[1]:  “A geometric take on metric learning”, Hauberg et al., Advances in Neural Information Processing Systems, 2012.

[2]: “Best practices for single-cell analysis across modalities.”, Heumos, L., Schaar, A.C., Lance, C. et al. Nat Rev Genet (2023). https://doi.org/10.1038/s41576-023-00586-w

[3]:~ https://www.sc-best-practices.org/preprocessing_visualization/dimensionality_reduction.html

[4]: “Contrastive Learning for Unpaired Image-to-Image Translation” Park et al., Computer Vision-ECCV (2020)
In these comments we now address any standing question/concern.
> _For some of the cellular experiments, the manifold seems to actually be found by PCA_

In single-cell data, current best practices suggest performing non-linear dimensionality reduction on the top 10-100 principle components. We follow this standard approach, learning a non-linear metric on the top 5/100 PCs, please see [2] and [3] below. 
> _The paper uses very vague terminology. For instance, there is repeated mention of a "more meaningful matching"_

We clarify that the term matching (introduced in Score Matching and Flow Matching [Song et. al 2019, Lipman et. al 2023]) refers approximating a target distribution $p_1$ with a learned one $p_{\theta}$ via a regression objective over vector fields—i.e. $|| v\_{\theta}(x\_t, t) - u\_t(x\_t | x\_0, x\_1)||\_2$. For trajectory inference tasks “more meaningful matching” refers to learning a matching that better respects the task description, meaning that the reconstructed trajectories replicate those generated by the underlying physical system. 

We hope that our answer here addresses the reviewer's comment and we will update the main text to include these clarifications. 

**1. Lidar scan**
> _It is unclear what the aim of this experiment is (...)_

We thank the reviewer for their thoughtful suggestions. The primary purpose of our Lidar experiments is to simply contrast MFM to the Euclidean baseline (OT-CFM) and visually illustrate how the metric can affect the learned paths. While using more task-specific potentials—e.g. that account for the height information—as done in GSBM, is possible within our framework, we believe that this goes beyond the goal of this experiment which serves to highlight the benefit of incorporating data geometry **without** providing additional details from the downstream task. 

To address the reviewer's concern about the faithfulness of $p\_t$ to the LiDAR data, in our **general rebuttal** and additional 1 page pdf, we also provide more views of the learned paths from OT-CFM and OT-MFM . 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. 

**2. The AFHQ experiment**
> _Here the FID values are extremely high (...)_

We acknowledge the reviewers' healthy skepticism regarding our reported FID values. We wish to highlight that other papers e.g. [4] have empirically reported FID values for interpolation of cats→dogs to be around 70. As such we believe that our values are acceptable. 
> _Qualitatively looking at the samples, it doesn't seem clear which method has the better samples at t=1/2.(...)_

We agree with the reviewer’s observation. Note however that in this application the goal is not much about having better samples at the intermediate times, but having samples at $t=1$ that are more similar to those at $t=0$ as measured by the LPIPS metric, which we have reported in Table 2. In general though, the main goal of this experiment was to showcase how MFM can be used along with or instead of CFM, even for settings outside the trajectory inference task which is what was mainly designed for.  
> _Again, here I feel the role of the interpolant is heavily diminished when compared to the role of the optimal coupling.(...)_

We thank the reviewer for their suggestions. As noted above and detailed in the **general rebuttal**, we have conducted experiments comparing CFM and MFM using independent coupling for both the Arch task and single cell RNA sequencing to validate that I-MFM surpasses I-CFM.

**3. The single-cell experiment**
> _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(...)_

This is an interesting question. In our framework, the metric $\mathbf{G}$ is constructed using two consecutive dataset marginals—based on a dataset $D\_{i,j}$, which is a concatenation of samples from these marginals (excluding $p\_{\text{out}}$). For instance, considering the scRNA dataset with densities $p\_0$, $p\_1$, $p\_2$, and $p\_3$, and excluding $p\_1$ as the left-out density, we use separate metrics $\mathbf{G}_{0,2}$ and $\mathbf{G}\_{2,3}$ for the pairs $\{p\_0, p\_2\}$ and $\{p\_2, p\_3\}$. In this regard, the definition of the metric already follows the procedure you suggested.

In contrast, a *single* time-dependent interpolant network $\varphi_{t,\eta}$ and a *single* vector field networks $v\_{t, \theta}$ are used for all time-steps ($t\_0$, $t\_2$ and $t\_3$) to ensure continuity of trajectories across different times. This also ensures that we follow standard procedures adopted by the baselines reported in Tables 3 and 4. 
> _Issue with the choice of coupling (...)_

We have addressed this point above and in the **general rebuttal**. Once again, we highlight that I-MFM is better than I-CFM.