Original: R IEMANNIAN
Revised: R IEMANNIAN

Original: M ETRIC
Revised: M ETRIC

Original: L EARNING
Revised: L EARNING

Original: VIA
Revised: VIA

Original: O PTIMAL
Revised: O PTIMAL

Original: T RANSPORT
Revised: T RANSPORT

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: We introduce an optimal transport-based model for learning a metric tensor from cross-sectional samples of evolving probability measures on a common Riemannian manifold.
Revised: We introduce an optimal transport-based model for learning a metric tensor from cross-sectional samples of evolving probability measures on a common Riemannian manifold.

Original: We neurally parametrize the metric as a spatially-varying matrix field and efficiently optimize our model’s objective using a simple alternating scheme.
Revised: We neurally parametrize the metric as a spatially-varying matrix field and efficiently optimize our model’s objective using a simple alternating scheme.

Original: Using this learned metric, we can nonlinearly interpolate between probability measures and compute geodesics on the manifold.
Revised: Using this learned metric, we can nonlinearly interpolate between probability measures and compute geodesics on the manifold.

Original: We show that metrics learned using our method improve the quality of trajectory inference on scRNA and bird migration data at the cost of little additional cross-sectional data
Revised: We show that metrics learned using our method improve the quality of trajectory inference on scRNA and bird migration data at the cost of little additional cross-sectional data

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: In settings like single-cell RNA sequencing (scRNA-seq) (Tanay & Regev, 2017), we often encounter pooled cross-sectional data : Time-indexed samples (cid:8) x it (cid:9)
Revised: In settings like single-cell RNA sequencing (scRNA-seq) (Tanay & Regev, 2017), we often encounter pooled cross-sectional data : Time-indexed samples (cid:8) x it (cid:9)

Original: N t
Revised: N t

Original: i =1 from an evolving population X t with no correspondence between samples x is and x it at times s ̸ = t .
Revised: i =1 from an evolving population X t with no correspondence between samples x is and x it at times s ̸ = t .

Original: Such data may arise when technical constraints impede the repeated observation of some population member.
Revised: Such data may arise when technical constraints impede the repeated observation of some population member.

Original: For example, as scRNA-seq is a destructive process, any given cell’s gene expression profile can only be measured once before the cell is destroyed.
Revised: For example, as scRNA-seq is a destructive process, any given cell’s gene expression profile can only be measured once before the cell is destroyed.

Original: This data is often sampled sparsely in time, leading to interest in trajectory inference : Inferring the distribution of the population or the positions of individual particles between times { t i } Ti =1 at which samples are drawn.
Revised: This data is often sampled sparsely in time, leading to interest in trajectory inference : Inferring the distribution of the population or the positions of individual particles between times { t i } Ti =1 at which samples are drawn.

Original: A fruitful approach has been to model the evolving population as a time-varying probability measure P t on R D and to infer the distribution of the population between observed times by interpolating between subsequent pairs of measures (cid:0) P t i , P t
Revised: A fruitful approach has been to model the evolving population as a time-varying probability measure P t on R D and to infer the distribution of the population between observed times by interpolating between subsequent pairs of measures (cid:0) P t i , P t

Original: i +1 (cid:1) .
Revised: i +1 (cid:1) .

Original: Some existing approaches to this problem use dynamical optimal transport to interpolate between probability measures, which implicitly encodes a prior that particles travel along straight lines between observations.
Revised: Some existing approaches to this problem use dynamical optimal transport to interpolate between probability measures, which implicitly encodes a prior that particles travel along straight lines between observations.

Original: This prior is often implausible, especially when the evolving population is sampled sparsely in time.
Revised: This prior is often implausible, especially when the evolving population is sampled sparsely in time.

Original: One can straightforwardly extend optimal transport-based methods by allowing the user to specify a spatially-varying metric tensor to bias the inferred trajectories away from straight lines.
Revised: One can straightforwardly extend optimal transport-based methods by allowing the user to specify a spatially-varying metric tensor to bias the inferred trajectories away from straight lines.

Original: This approach is theoretically well-founded and amounts to redefining a straight line by altering the manifold on which trajectories are inferred.
Revised: This approach is theoretically well-founded and amounts to redefining a straight line by altering the manifold on which trajectories are inferred.

Original: Such a metric tensor, however, is typically unavailable in most real-world applications.
Revised: Such a metric tensor, however, is typically unavailable in most real-world applications.

Original: We resolve this problem by introducing an optimal transport-based model in which a metric tensor may be recovered from cross-sectional samples of evolving probability measures on a common manifold.
Revised: We resolve this problem by introducing an optimal transport-based model in which a metric tensor may be recovered from cross-sectional samples of evolving probability measures on a common manifold.

Original: We derive a tractable optimization problem using the theory of optimal transport on Riemannian manifolds, neurally parametrize its variables, and solve it using gradient-based optimization.
Revised: We derive a tractable optimization problem using the theory of optimal transport on Riemannian manifolds, neurally parametrize its variables, and solve it using gradient-based optimization.

Original: We demonstrate our algorithm’s ability to recover a known metric tensor from cross-sectional samples on synthetic examples.
Revised: We demonstrate our algorithm’s ability to recover a known metric tensor from cross-sectional samples on synthetic examples.

Original: We also show that our learned metric tensor improves the quality of trajectory inference on scRNA data and allows us to infer curved trajectories for individual birds from crosssectional samples of a migrating population.
Revised: We also show that our learned metric tensor improves the quality of trajectory inference on scRNA data and allows us to infer curved trajectories for individual birds from crosssectional samples of a migrating population.

Original: Our method is both computationally-efficient, requiring little computational resources relative to the downstream trajectory inference task, and data-efficient, requiring little data per time point to recover a useful metric tensor.
Revised: Our method is both computationally-efficient, requiring little computational resources relative to the downstream trajectory inference task, and data-efficient, requiring little data per time point to recover a useful metric tensor.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Measure interpolation.
Revised: Measure interpolation.

Original: An emerging literature considers the problem of smooth interpolation between probability measures.
Revised: An emerging literature considers the problem of smooth interpolation between probability measures.

Original: Given a finite collection of probability measures P i sampled at times t

Original: i

Original: ( i = 1 , ..., n ), we seek to construct a curve of probability measures ρ t such that ρ t i = P i for all i . This problem generalizes the classical problem of interpolating a set of points x i ∈ R D . Using the theory of optimal transport, we may construct a displacement interpolation (McCann, 1997) between successive measures ( P i , P i +1 ) , yielding a sequence of geodesics between pairs ( P i , P i +1 ) in the space of probability measures equipped with the 2-Wasserstein distance W 2 (Villani, 2008).
Revised: Using the theory of optimal transport, we may construct a displace- ment interpolation (McCann, 1997) between successive measures ( P i , P i +1 ) , yielding a sequence of geodesics between pairs ( P i , P i +1 ) in the space of probability measures equipped with the 2Wasserstein distance W 2 (Villani, 2008).

Original: This generalizes piecewise-linear interpolation to probability measures.
Revised: This generalizes piecewise-linear interpolation to probability measures.

Original: Schiebinger et al.
Revised: Schiebinger et al.

Original: (2019) use this method to infer the developmental trajectories of cells based on static measurements of their gene expression profiles.
Revised: (2019) use this method to infer the developmental trajectories of cells based on static measurements of their gene expression profiles.

Original: Recent works such as (De Bortoli et al., 2021) and (Vargas et al., 2021) provide numerical schemes for computing Schrödinger bridges ; these are an entropically-regularized analog to the displacement interpolation between measures.
Revised: Recent works such as (De Bortoli et al., 2021) and (Vargas et al., 2021) provide numerical schemes for computing Schrödinger bridges ; these are an entropically-regularized analog to the displacement interpolation between measures.

Original: As we seek smooth interpolants, it is natural to ask whether there is a meaningful analogy to cubic splines for interpolating probability measures.

Original: Chen et al.
Revised: Chen et al.

Original: (2018b);
Revised: (2018b);

Original: Benamou et al.
Revised: Benamou et al.

Original: (2018) leverage the variational characterization of cubic splines as minimizers of mean-square acceleration over the set of interpolating curves and develop a generalization in the space of probability measures.
Revised: (2018) leverage the variational characterization of cubic splines as minimizers of mean-square acceleration over the set of interpolating curves and develop a generalization in the space of probability measures.

Original: Chewi et al.
Revised: Chewi et al.

Original: (2021) extend these works by providing computationally efficient algorithms for computing measure-valued splines.
Revised: (2021) extend these works by providing computationally efficient algorithms for computing measure-valued splines.

Original: Hug et al.
Revised: Hug et al.

Original: (2015) modify the usual displacement interpolation between probability measures by introducing anisotropy to the domain on which the measures are defined.
Revised: (2015) modify the usual displacement interpolation between probability measures by introducing anisotropy to the domain on which the measures are defined.

Original: This change corresponds to imposing preferred directions for the local displacement of probability mass.
Revised: This change corresponds to imposing preferred directions for the local displacement of probability mass.

Original: Whereas Hug et al.
Revised: Whereas Hug et al.

Original: hard-code the domain’s anisotropy, our method allows us to learn this anisotropy from snapshots of a probability measure evolving in time.
Revised: hard-code the domain’s anisotropy, our method allows us to learn this anisotropy from snapshots of a probability measure evolving in time.

Original: The authors also demonstrate that physical priors such as incompressibility and rigidity can be introduced at the cost of non-convexity.

Original: Zhang et al.
Revised: Zhang et al.

Original: (2022) apply similar techniques to unstructured animation problems.
Revised: (2022) apply similar techniques to unstructured animation problems.

Original: Ding et al.
Revised: Ding et al.

Original: (2020) propose a non-convex inverse problem for recovering the ground metric and interaction kernel in a class of mean-field games and supply a primal-dual algorithm for solving a grid discretization of the problem.
Revised: (2020) propose a non-convex inverse problem for recovering the ground metric and interaction kernel in a class of mean-field games and supply a primal-dual algorithm for solving a grid discretization of the problem.

Original: Whereas their Eulerian approach inputs a grid discretization of observed densities and velocity fields, our approach is Lagrangian, operating directly on temporal observations of particle positions.
Revised: Whereas their Eulerian approach inputs a grid discretization of observed densities and velocity fields, our approach is Lagrangian, operating directly on temporal observations of particle positions.

Original: Trajectory inference from population-level data.
Revised: Trajectory inference from population-level data.

Original: In domains such as scRNA-seq, we study an evolving population from which it is impossible or prohibitively costly to collect longitudinal data.
Revised: In domains such as scRNA-seq, we study an evolving population from which it is impossible or prohibitively costly to collect longitudinal data.

Original: Instead, we observe distinct cross-sectional samples from the population at a collection of times and wish to infer the dynamics of the latent population from these observations.
Revised: Instead, we observe distinct cross-sectional samples from the population at a collection of times and wish to infer the dynamics of the latent population from these observations.

Original: This problem is called trajectory inference .
Revised: This problem is called trajectory inference .

Original: Hashimoto et al.
Revised: Hashimoto et al.

Original: (2016) study the conditions under which it is possible to recover a potential function from population-level observations of a system evolving according to a Fokker-Planck equation.
Revised: (2016) study the conditions under which it is possible to recover a potential function from population-level observations of a system evolving according to a Fokker-Planck equation.

Original: They provide an RNN-based algorithm for this learning task and investigate their model’s ability to recover differentiation dynamics from scRNA-seq data sampled sparsely in time.
Revised: They provide an RNN-based algorithm for this learning task and investigate their model’s ability to recover differentiation dynamics from scRNA-seq data sampled sparsely in time.

Original: A recent work of Bunne et al.
Revised: A recent work of Bunne et al.

Original: (2022) presents a proximal analog to this approach, modeling population dynamics as a JKO flow with a learnable energy function and describing a numerical scheme for computing this flow using input-convex neural networks.
Revised: (2022) presents a proximal analog to this approach, modeling population dynamics as a JKO flow with a learnable energy function and describing a numerical scheme for computing this flow using input-convex neural networks.

Original: This method addresses initial value problems where one is given an initial measure ρ 0 and seeks to predict future measures ρ t for t > 0 .
Revised: This method addresses initial value problems where one is given an initial measure ρ 0 and seeks to predict future measures ρ t for t > 0 .

Original: In contrast, our method is primarily applicable to boundary value problems, where we are given initial and final measures ρand ρ 1 and seek an interpolant ρ t for 0 < t < 1 .
Revised: In contrast, our method is primarily applicable to boundary value problems, where we are given initial and final measures ρand ρ 1 and seek an interpolant ρ t for 0 < t < 1 .

Original: Schiebinger et al.
Revised: Schiebinger et al.

Original: (2019) use optimal transport to infer future and prior gene expression profiles from a single observation in time of a cell’s gene expression profile.
Revised: (2019) use optimal transport to infer future and prior gene expression profiles from a single observation in time of a cell’s gene expression profile.

Original: Yang & Uhler (2019) propose a
Revised: Yang & Uhler (2019) propose a

Original: GAN-like solver for unbalanced optimal transport and investigate the effectiveness of their method for the inference of future gene expression states in zebrafish single-cell gene expression data.
Revised: GAN-like solver for unbalanced optimal transport and investigate the effectiveness of their method for the inference of future gene expression states in zebrafish single-cell gene expression data.

Original: As they learn transport maps between probability measures defined at discrete time points, none of the above OT-based methods is suitable for inferring continuous trajectories.
Revised: As they learn transport maps between probability measures defined at discrete time points, none of the above OT-based methods is suitable for inferring continuous trajectories.

Original: Tong et al.
Revised: Tong et al.

Original: (2020) show that a regularized continuous normalizing flow can provide an efficient approximation to displacement interpolations arising from dynamical optimal transport.
Revised: (2020) show that a regularized continuous normalizing flow can provide an efficient approximation to displacement interpolations arising from dynamical optimal transport.

Original: These displacement interpolations often do not yield plausible paths through gene expression space, so the authors propose several applicationspecific regularizers to bias the inferred trajectories toward more realistic paths.
Revised: These displacement interpolations often do not yield plausible paths through gene expression space, so the authors propose several applicationspecific regularizers to bias the inferred trajectories toward more realistic paths.

Original: Rather than relying on bespoke regularizers, our method supplies a natural approach for learning local directions of particle motion in a way that is amenable to integration with algorithms like that of Tong et al.
Revised: Rather than relying on bespoke regularizers, our method supplies a natural approach for learning local directions of particle motion in a way that is amenable to integration with algorithms like that of Tong et al.

Original: P RELIMINARIES

Original: Before describing our method, we review relevant facts from Riemannian geometry and optimal transport.

Original: Riemannian geometry.
Revised: Riemannian metric learning.

Original: A Riemannian manifold M = ( M, g ) is a differentiable manifold M equipped with a Riemannian metric g .

Original: Throughout this paper, we will take M = R D and focus our attention on the metric g .

Original: A Riemannian metric assigns an inner product ⟨· , ·⟩ p to the tangent space

Original: T p M of each p ∈ M in a way that varies smoothly with p .

Original: In the case where M = R D , we may simply identify all tangent spaces T p M with R D ; here a metric g amounts to a spatially-varying inner product on R D .

Original: Since any inner product on R D may be computed as ⟨ x, y ⟩ A

Original: = x T Ay for some A ≻ 0 , a Riemannian metric g on M = R D is specified by a smooth field of positive definite matrices A ( x ) : R D → S D ++ (we use S D
Revised: Lebanon (2002) develops a parametric method for learning a Riemannian metric from sparse high-dimensional data and shows that nearest-neighbor classifiers based on their learned metric outperform classifiers based on TF-IDF cosine similarity.

Original: ++ to denote the set of positive definite D × D matrices).

Original: We will henceforth use ⟨ x, y ⟩ A ( p ) =

Original: x T A ( p ) y to denote the metric at a point p ∈ R D .

Original: An inner product ⟨· , ·⟩ A ( p ) induces a norm ∥ v ∥ A ( p ) :

Original: = (cid:112) ⟨ v, v ⟩ A ( p ) .

Original: Given this norm, we define the length of a continuously differentiable curve γ : [0 , 1] → M to be ℓ ( γ ) =

Original: [-(cid:82) 1 
-]
Revised: [-(cid:82) 1 
-]

Original: ∥ ˙ γ ( t ) ∥

Original: A ( γ ( t )) dt . The A -geodesic distance d A ( x, y ) between two points x, y ∈ M is then the infimum of the length ℓ ( γ ) of continuously differentiable curves γ such that γ (0) = x and γ (1) =
Revised: Hauberg et al.

Original: y .

Original: W 1 on Riemannian manifolds. Feldman & McCann (2002) study Monge’s transport problem on a Riemannian manifold M = ( M, g ) ; we once again consider the case M = R D with g specified by a field ofmatrices A ( x ) ≻ 0 . Given two compactly-supported densities ρ 0 , ρ 1 on M , Monge’s problem seeks a pushforward s of ρ 0 onto ρ 1 solving the following problem:
Revised: (2012) construct a metric tensor as a weighted average of a set of learned metric tensors, show how to compute geodesics and exponential and logarithmic maps on the resulting Riemannian manifold, and use their metric tensor in regression and dimensionality reduction tasks arising from human body shape analysis.

Original: [Equation] Feldman & McCann (2002) show that under mild technical conditions, (1) has a possibly non-unique solution s . Intuitively, this map is a pushforward of ρ 0 onto ρ 1 minimizing the average geodesic distance between matched units of probability mass. The optimal value of (1) is the 1-Wasserstein distance between ρ 0 and ρ 1 .
Revised: Arvanitidis et al.

Original: If ϕ is continuously differentiable on M , then it is 1-Lipschitz on M if and only if ∥∇ ϕ ( x ) ∥

Original: A − 1 ( x ) ≤ 1 for all x ∈ M . Here, we use ∇ ϕ ( x ) to denote the vector of partial derivatives of ϕ at x . Armed with  this local characterization of Lipschitz continuity, we may define the following dual problem to (1): [Equation]
Revised: (2016) learn a Riemannian metric that encourages geodesics to move towards region of high data density, define a Riemannian normal distribution with respect to this metric, and show how to learn the parameters of this model via maximum likelihood estimation.

Original: If the minimum in (2) is attained by some Kantorovich potential ϕ , the optimal values of (1) and (2) coincide and the Lipschitz bound for ϕ is tight at pairs of points ( x, y ) arising from a Monge map s solving (1) (Feldman & McCann, 2002, Lemma 4).

Original: Feldman & McCann (2002, Lemma 10) also show that given a geodesic γ x : [0 , 1] → M between a pair of matched points ( x, s ( x )) and t ∈ (0 , 1) , we have ∇ ϕ ( γ x ( t ))

Original: =

Original: − ˙ γ x ( t ) ∥ ˙ γ x ( t ) ∥ A ( γx ( t )) . Intuitively, ∇ ϕ points along geodesics γ x joining pairs of points in the support of ρ 0 and ρ 1 that are matched by the Monge map s .
Revised: Whereas these methods learn a Riemannian metric from static data, our model learns a metric from snapshots of probability distributions evolving in time on a common manifold.

Original:  M ETHOD
Revised:  M ETHOD

Original: We now describe our method for learning a Riemannian metric from cross-sectional samples of populations evolving in time on a common manifold.
Revised: We now describe our method for learning a Riemannian metric from cross-sectional samples of populations evolving in time on a common manifold.

Original: We learn a metric that minimizes the average 1Wasserstein distance on the manifold between pairs of subsequent time samples from each population.
Revised: We learn a metric that minimizes the average 1Wasserstein distance on the manifold between pairs of subsequent time samples from each population.

Original: We derive a dual formulation of our problem, parametrize its variables by neural networks, and solve for the dual variables and the metric via alternating optimization.
Revised: We derive a dual formulation of our problem, parametrize its variables by neural networks, and solve for the dual variables and the metric via alternating optimization.

Original: 4. M ODEL
Revised: 3. M ODEL

Original: Suppose we have K populations evolving according to latent continuous dynamics over a common
Revised: Suppose we have K populations evolving according to unknown continuous dynamics over a common

Original: Riemannian manifold M = (cid:0) R D , g (cid:1)  .
Revised: Riemannian manifold M = (cid:0) R D , g (cid:1) with unknown metric g .

Original: The metric g is defined at any p ∈ R D by the inner product ⟨ x, y ⟩p = x  T A ( p ) y for A ( p ) ≻ 0 .
Revised: The metric g is defined at any x ∈ R D by the inner product ⟨ u, v ⟩ x = u T A ( x ) v for A ( x ) ≻ 0 .

Original: We model each population as a compactly-supported probability distribution P k with density ρ k on R D being pushed through a latent velocity field v k ( x ) . 
Revised: We model each population as a compactlysupported probability distribution P k with density ρ k on R D being pushed through an unobserved velocity field v k ( x ) . We will learn the metric tensor A ( x ) from temporal snapshots of the populations P k during their evolution on the manifold.

Original: Depending on the nature of our data, we may have sample access to each P k at  times t = 0 and t =
Revised: Depending on the nature of our data, we may have the ability to repeatedly sample from P k at a pair of initial and final times t = 0 and t =

Original: T k or to empirical distributions over S k samples { x k, 0 i }
Revised: T k or have a fixed set of samples { x k, 0 i }

Original: S k i
Revised: S k i

Original: =1 and { x k,T k i } S k i
Revised: =1 and { x k,T k i } S k i

Original: =1 drawn from the  P k at their respective times.
Revised: =1 drawn from the populations P k at their respective times.

Original: For convenience, we denote both the density ρ k and the empirical approximations to P k at times t ∈ { 0 , T k } by ρ k 0 and ρ k 1 respectively.
Revised: For convenience, we denote both the density ρ k and the empirical distributions over samples from P k at times t ∈ { 0 , T k } by ρ kand ρ k 1 respectively.

Original: As we do not observe the continuous dynamics of the populations, we assume that probability mass travels from initial to final positions along A -geodesics.
Revised: As we only observe the initial and final spatial distributions ρ k 0 and ρ k 1 of the populations and do not observe their dynamics, we assume that probability mass travels from initial to final positions along A -geodesics.

Revised: Geodesics are paths that minimize the action (or average kinetic energy) of a particle traveling between points x, y ∈ M ; this least-action interpretation of a geodesic makes it a natural prior on paths in the absence of further information.

Original: We learn a field of positive definite matrices A ( x ) that minimizes the average A -geodesic distance between the initial and final positions of each unit of probability mass.
Revised: We learn a field of positive definite matrices A ( x ) that minimizes the average A -geodesic distance between the initial and final positions of each unit of probability mass from each P k .

Original: Formally, let r k be the map sending a point x ∈ M to its final position r k ( x ) after flowing through the latent velocity field v k from time t = 0 to t =
Revised: Formally, let r k be the map sending a point x ∈ M to its final position r k ( x ) after flowing through the latent velocity field v k from time t = 0 to t =

Original: T k .
Revised: T k .

Original: If we had access to such a solution map, we would ideally solve the following problem:
Revised: If we had access to such a solution map, we would ideally solve the following problem:

Original: [Equation]
Revised: [Equation]

Original: where R ( A ) is a regularizer that excludes the trivial solution A ≡ 0 .
Revised: where R ( A ) is a regularizer that excludes the trivial solution A ≡ 0 .

Revised: Problem (1) optimizes for a non-trivial metric A ( x ) that minimizes the average A -geodesic distance d A (cid:0) x, r k ( x ) (cid:1) traveled by particles x in the population ρ k 0 at time t = 0 to their final positions r k ( x ) at time t =

Original:  Since we do not know the velocity fields v k , however, we also do not know the maps r k in (3).
Revised: T k . Since we do not know the velocity fields v k that encode the particle dynamics, however, we also do not know the maps r k in (1) that specify the correspondence between particle positions x at t =and positions r k ( x ) at final times T k .

Original: Furthermore, as noted in Section 1, we often encounter data for which it is impossible to observe contiguous particle trajectories: A particle that we observe at t = 0 may not be in the sample  t =
Revised: Furthermore, as noted in Section 1, we often encounter data for which it is impossible to observe contiguous particle trajectories: A particle that we observe at t =may not be in the sample at t =

Original: T k .
Revised: T k .

Revised: This issue is unavoidable for destructive measurement processes such as scRNA sequencing; in this setting, a cell whose scRNA profile is observed at t = 0 would be destroyed at this time and hence unobservable at a future time t

Revised: = T k .

Original: To accommodate these limitations, we replace the true matchings of initial and final positions r k with the Monge map , defined in the Section 3 as the solution s kto (1).
Revised: To accommodate these limitations, we replace the true matchings of initial and final positions r k with the Monge map s k that solves the following problem: [Equation] Here s k is a pushforward of ρ k 0 onto ρ k 1 ; we write ρ k 1 = s k

Original:  Substituting (1) into our idealized objective (3), we obtain the following lower bound on (3):
Revised: # ρ k 0 to denote this relationship. This map matches units of mass from the initial and final distributions ρ k 0 , ρ k 1 to minimize their average A -geodesic distance. Substituting solutions to (2) for the maps r k in our idealized objective (1), we obtain the following lower bound on (1):

Original: [Equation]
Revised: [Equation]

Original: Problem (4) is challenging as written: It requires the ability to compute and differentiate geodesic distances d A with respect to an arbitrary metric.
Revised: Problem (3) is challenging as written: It requires the ability to compute and differentiate geodesic distances d A with respect to an arbitrary metric.

Revised: However, the inner optimization problem over maps s k is a collection of decoupled Monge problems (2) whose optimal value is the average 1-Wasserstein distance between ( ρ k 0 , ρ k 1 ) pairs on the manifold with metric A ( x ) .

Original: We mitigate this difficulty by appealing to KantorovichRubinstein duality to rewrite Problem (1) in a  form that is amenable to gradient-based optimization.
Revised: As shown in Appendix B, these problems can be expressed in a dual form (10) which is amenable to gradient-based optimization.

Original: Using the  dual formulation (2), we may equivalently write Problem (4) as a minimax problem:
Revised: Replacing the inner Monge problems with their dual formulations (10), we may equivalently write Problem (3) as a minimax problem:

Original: [Equation]
Revised: [Equation]

Revised: In Problem (4), we learn a metric A ( x ) that minimizes the average 1-Wasserstein distance on the manifold between pairs of subsequent time samples from each population.

Original: (5) requires neither the computation of geodesic distances on M nor the solution of an assignment problem.
Revised: (4) requires neither the computation of geodesic distances on M nor the solution of an assignment problem.

Original: As such, it is substantially more tractable than the initial objective (4).
Revised: As such, it is substantially more tractable than the initial objective (3).

Original: We provide the details of our implementation of (5) in Section 4.2.
Revised: We provide the details of our implementation of (4) in Section 3.2.

Original: 4. I MPLEMENTATION
Revised: 3. I MPLEMENTATION

Original: Enforcing the Lipschitz constraint.
Revised: Enforcing the Lipschitz constraint.

Original: Problem (5) includes global constraints of form ∥∇ ϕ k
Revised: Problem (4) includes global constraints of form ∥∇ ϕ k

Original: ∥
Revised: ∥

Original: A − 1 ≤ 1 .
Revised: A − 1 ≤ 1 .

Original: These constraints are the Riemannian analog to the Lipschitz constraint in the dual formulation of the 1 -Wasserstein distance on R D .
Revised: These constraints are the Riemannian analog to the Lipschitz constraint in the dual formulation of the 1 -Wasserstein distance on R D .

Original: Constraints of this type are challenging to enforce in gradient-based optimization, and the Wasserstein GAN literature has explored approximations (Arjovsky et al., 2017;
Revised: Constraints of this type are challenging to enforce in gradient-based optimization, and the Wasserstein GAN literature has explored approximations (Arjovsky et al., 2017;

Original: Gulrajani et al., 2017; Miyato et al., 2018).
Revised: Gulrajani et al., 2017; Miyato et al., 2018).

Original: We follow the standard technique introduced by Gulrajani et al.
Revised: We follow the standard technique introduced by Gulrajani et al.

Original: and replace the global constraints ∥∇ ϕ k ∥
Revised: and replace the global constraints ∥∇ ϕ k ∥

Original: A − 1 ≤ 1 with soft penalties of the following form:
Revised: A − 1 ≤ 1 with soft penalties of the following form:

Original: [Equation]
Revised: [Equation]

Original:  where σ x 1 x 0 ( t ) :
Revised: (5) where σ x 1 x 0 ( t ) :

Original: = (1 − t ) x
Revised: = (1 − t ) x

Original: 0 + tx 1 is a line segment between x 0 and x 1 parametrized by t ∈ [0 , 1] .
Revised: 0 + tx 1 is a line segment between x 0 and x 1 parametrized by t ∈ [0 , 1] .

Original: Intuitively, (6) penalizes violation of the Lipschitz constraint ∥∇ ϕ k
Revised: Intuitively, (5) penalizes violation of the Lipschitz constraint ∥∇ ϕ k

Original: ∥
Revised: ∥

Original: A − 1 ≤ 1 along line segments connecting randomly-paired points in X k 0 and X k 1 .
Revised: A − 1 ≤ 1 along line segments connecting randomly-paired points in X k 0 and X k 1 .

Original: Gulrajani et al.
Revised: Gulrajani et al.

Original: (2017, Prop. 1) justify this choice via a standard result in W 1 theory showing that the constraint ∥∇ ϕ k ∥ ≤ 1 binds on line segments connecting pairs of points that are matched by the Monge map s k .
Revised: (2017, Prop. 1) justify this choice via a standard result in W 1 theory showing that the constraint ∥∇ ϕ k ∥ ≤ 1 binds on line segments connecting pairs of points that are matched by the Monge map s k .

Original: Korotin et al.
Revised: Korotin et al.

Original: (2022) show that this method results in accurate approximations to the directions of the gradients ∇ ¯ ϕ k of the true Kantorovich potentials.
Revised: (2022) show that this method results in accurate approximations to the directions of the gradients ∇ ¯ ϕ k of the true Kantorovich potentials.

Original: This observation is sufficient for our purposes; as noted below, our optimization scheme encourages the low-energy eigenvectors of A ( x ) to be well-aligned with the solutions ∇ ϕ k to the inner problem in (5).
Revised: This observation is sufficient for our purposes; as noted below, our optimization scheme encourages the low-energy eigenvectors of A ( x ) to be well-aligned with the solutions ∇ ϕ k to the inner problem in (4).

Original: Choice of regularizer R ( A ) .
Revised: Choice of regularizer R ( A ) .

Original: Without a regularizer R ( A ) , the objective in (5) can be driven to 0 by choosing A ( x ) ≡ αI for arbitrarily small α > 0 , thereby making all pairs of measures arbitrarily close.
Revised: Without a regularizer R ( A ) , the objective in (4) can be driven to 0 by choosing A ( x ) ≡ αI for arbitrarily small α > 0 , thereby making all pairs of measures arbitrarily close.

Revised: This trivial solution would incorporate no useful information from the observed samples from ρ k 0 and ρ k 1 , as it is simply a rescaling of the standard Euclidean metric.

Original: Analogous challenges are widespread in machine learning.
Revised: Analogous challenges are widespread in machine learning.

Original: For example, contrastive losses for representation learning draw “positive pairs” close together and push “negative pairs” apart in a learned embedding space; without the implicit regularization of negative sampling, the contrastive loss leads to a collapsed embedding.
Revised: For example, contrastive losses for representation learning draw “positive pairs” close together and push “negative pairs” apart in a learned embedding space; without the implicit regularization of negative sampling, the contrastive loss leads to a collapsed embedding.

Original: We opt for the following regularizer in our method:
Revised: We opt for the following regularizer in our method:

Original: [Equation]
Revised: [Equation]

Original: where σ x 1 x 0 ( t ) is as in the previous paragraph and ∥ · ∥ 2 F denotes the squared Frobenius norm.
Revised: where σ x 1 x 0 ( t ) is as in the previous paragraph and ∥ · ∥ 2 F denotes the squared Frobenius norm.

Original: Regularizer (7) penalizes ∥
Revised: Regularizer (6) penalizes ∥

Original: A − 1 ∥ 2 F at points drawn using a sampling scheme analogous to that in (6).
Revised: A − 1 ∥ 2 F at points drawn using a sampling scheme analogous to that in (5).

Revised: This is a natural way to exclude trivial solutions of form A ( x ) ≡ αI for arbitrarily small α > 0 , as ∥ A − 1 ∥ 2 F is large for such metrics.

Original: We investigate the impact of the regularization coefficient λ on the learned metric in Appendix C. In Appendix D, we demonstrate the effect of removing the optimal transport term from (5), which is equivalent to setting λ = + ∞ .
Revised: We investigate the impact of the regularization coefficient λ on the learned metric in Appendix E. In Appendix F, we demonstrate the effect of removing the optimal transport term from (4), which is equivalent to setting λ = + ∞ .

Original: Optimization scheme.
Revised: Optimization scheme.

Original: We parametrize the scalar potentials ϕ k  and  the matrix field A − 1 by  neural networks; where we require the evaluation of A ( x ) in downstream applications, we evaluate and then invert
Revised: We parametrize the scalar potentials ϕ k by neural networks to enable the use of gradient-based optimization to solve Problem (4). As A appears in Equations (4), (5), and (6) only via its inverse A − 1 , we directly parametrize the matrix field A − 1 by a neural network; where we require the evaluation of A ( x ) in downstream applications, we evaluate and then invert

Original: A − 1 ( x ) .
Revised: A − 1 ( x ) .

Original: We enforce the positive definiteness of A − 1 (and hence the positive definiteness of A ) by parametrizing it as A − 1 ( x ) =
Revised: We enforce the positive definiteness of A − 1 (and hence the positive definiteness of A ) by parametrizing it as A − 1 ( x ) =

Original: Q ( x ) T Q ( x )+ ηI for a matrix-valued function Q ( x ) :
Revised: Q ( x ) T Q ( x ) + ηI for a matrix-valued function Q ( x ) :

Original: M → R D × D and η > 0 .
Revised: M → R D × D and η > 0 .

Original: After parametrizing our problem variables with neural networks, we optimize the objective in (5) via alternation.
Revised: After parametrizing our problem variables with neural networks, we optimize the objective in (4) via alternation.

Original: In the first phase of our scheme, we solve the inner problem by holding A fixed (initializing it as A ( x ) =
Revised: In the first phase of our scheme, we solve the inner problem by holding A fixed (initializing it as A ( x ) =

Original: A − 1 ( x ) ≡ I ) and solving for the optimal ϕ k .
Revised: A − 1 ( x ) ≡ I ) and solving for the optimal ϕ k .

Original: This step decouples over the potentials ϕ k , and each of the resulting problems is an instance of the dual problem for the 1-Wasserstein distance on M .
Revised: This step decouples over the potentials ϕ k , and each of the resulting problems is an instance of the dual problem for the 1-Wasserstein distance on M .

Revised: We approximate the integrals in (4) as sample means 1 N (cid:80) Ni

Revised: =1 ϕ k

Revised: ( x i ) , where the { x i } are samples from the distributions ρ k 0 and ρ k 1 .

Revised: We likewise approximate the Lipschitz penalty (5) and the regularizer (6) as sample means over draws from ρ k 0 and ρ k 1 and over t drawn uniformly from [0 , 1] .

Original: In the second phase, we solve the outer problem by fixing the  ϕ k from the previous step and solving for the optimal matrix field A ( x ) .
Revised: In the second phase, we solve the outer problem by fixing the optimal ϕ k from the previous step and solving for the optimal matrix field A ( x ) .

Original: We optimize both problems using
Revised: We optimize both problems using

Original: AdamW (Loshchilov & Hutter, 2019).
Revised: AdamW (Loshchilov & Hutter, 2019).

Original: Few alternations are needed in practice to obtain high-quality results.
Revised: Few alternations are needed in practice to obtain high-quality results.

Original: The results in Section 3 show that given a fixed metric defined by A ( x ) , the optimal ∇ ϕ k from the first phase of our scheme point along A -geodesics joining pairs of points ( x, s k ( x )) where x ∼ ρ kand s k solves (1).
Revised: The results in Appendix B show that given a fixed metric defined by A ( x ) , the optimal ∇ ϕ k from the first phase of our scheme point along A -geodesics joining pairs of points ( x, s k ( x )) where x ∼ ρ kand s k solves (2).

Original: At initialization, A ≡ I , and these geodesics are line segments in R D connecting the matched points.
Revised: At initialization, A ≡ I , and these geodesics are line segments in R D connecting the matched points.

Original: Given fixed Kantorovich potentials ϕ k from the first phase, the second phase of our scheme solves for a matrix field
Revised: Given fixed Kantorovich potentials ϕ k from the first phase, the second phase of our scheme solves for a matrix field

Original: A ( x ) that minimizes the regularizer R ( A ) =
Revised: A ( x ) that minimizes the regularizer R ( A ) =

Original: ∥
Revised: ∥

Original: A − 1 ∥ 2 F while also making ∥∇ ϕ k
Revised: A − 1 ∥ 2 F while also making ∥∇ ϕ k

Original: ∥
Revised: ∥

Original: A −large.
Revised: A −large.

Original: Intuitively, this encourages the unit eigenvector u 1 ( x ) corresponding to the minimal eigenvalue λ 1 ( x ) of A ( x ) to be aligned with ∇ ϕ k ( x ) wherever the constraint is enforced.
Revised: Intuitively, this encourages the unit eigenvector u 1 ( x ) corresponding to the minimal eigenvalue λ 1 ( x ) of A ( x ) to be aligned with ∇ ϕ k ( x ) wherever the constraint is enforced.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In this section, we first use synthetic data to demonstrate that our algorithm successfully recovers the correct eigenspaces of a known metric A ( x ) from cross-sectional samples satisfying our model.
Revised: In this section, we first use synthetic data to demonstrate that our algorithm successfully recovers the correct eigenspaces of a known metric A ( x ) from cross-sectional samples satisfying our model.

Original: We then use our method to learn a metric from cross-sectional scRNA data and show that this metric improves the accuracy of trajectory inference for scRNA data that is sampled sparsely in time.
Revised: We then use our method to learn a metric from cross-sectional scRNA data and show that this metric improves the accuracy of trajectory inference for scRNA data that is sampled sparsely in time.

Original: We finally show that by learning a metric from time-stamped bird sightings, we can infer curved migratory trajectories for individual birds given the initial and final points of their trajectories.
Revised: We finally show that by learning a metric from time-stamped bird sightings, we can infer curved migratory trajectories for individual birds given the initial and final points of their trajectories.

Original: Details for all experiments are provided in Appendix B.
Revised: Details for all experiments are provided in Appendix D.

Original: 5. M ETRIC RECOVERY
Revised: 4. M ETRIC RECOVERY

Original: We first show that our method recovers the correct eigenspaces of a known metric A ( x ) from cross-sectional samples X k 0 and X k 1 satisfying the model in Section 4.1.
Revised: We first show that our method recovers the correct eigenspaces of a known metric A ( x ) from cross-sectional samples X k 0 and X k 1 satisfying the model in Section 3.1.

Original: We fix initial positions x n (0) and final positions x n (1) for a set of N particles x n ∈ X and also fix a spatially-varying metric A ( x ) .
Revised: We fix initial positions x n (0) and final positions x n (1) for a set of N particles x n ∈ X and also fix a spatially-varying metric A ( x ) .

Original: We compute A -geodesics between each ( x n (0) , x n (1)) pair and define populations
Revised: We compute A -geodesics between each ( x n (0) , x n (1)) pair and define populations

Original: X t i = { x n ( t i ) : n = 1 , ..., N } for times t i ∈ [0 , 1] , i = 0 , ..., T .
Revised: X t i = { x n ( t i ) : n = 1 , ..., N } for times t i ∈ [0 , 1] , i = 0 , ..., T .

Original: We use our method to learn the latent metric A ( x ) from the pairs of samples ( X t i , X t i +1 ) .
Revised: We use our method to learn the latent metric A ( x ) from the pairs of samples ( X t i , X t i +1 ) .

Original: For each example, we plot the eigenvectors of the true metric A ( x ) and those of the learned metric ˆ A ( x ) on a
Revised: For each example, we plot the eigenvectors of the true metric A ( x ) and those of the learned metric ˆ A ( x ) on a

Original: P × P grid (first row, Figure 1).
Revised: P × P grid (first row, Figure 1).

Original: We also plot the eigenvectors of ˆ A ( x ) along with the log-condition number log ( λ 2 ( x ) / λ 1 ( x ) ) of ˆ A ( x ) (second row, Figure 1); this value is large when the learned metric is highly anisotropic.
Revised: We also plot the eigenvectors of ˆ A ( x ) along with the log-condition number log ( λ 2 ( x ) / λ 1 ( x ) ) of ˆ A ( x ) (second row, Figure 1); this value is large when the learned metric is highly anisotropic.

Original: We finally report a measure of the alignment of the eigenspaces of the true metric A and its learned counterpart: ℓ ( A, ˆ A ) = 1 D |X| (cid:80) x ∈X (cid:80)
Revised: We finally report a measure of the alignment of the eigenspaces of the true metric A and its learned counterpart: ℓ ( A, ˆ A ) = 1 D |X| (cid:80) x ∈X (cid:80)

Original: Dd =1 |⟨ u d ( x ) , ˆ u d ( x ) ⟩| .
Revised: Dd =1 |⟨ u d ( x ) , ˆ u d ( x ) ⟩| .

Original: Here u d ( x ) is the unit eigenvector with eigenvalue
Revised: Here u d ( x ) is the unit eigenvector with eigenvalue

Original: λ
Revised: λ

Original: d of A ( x ) , ˆ u d ( x ) is the corresponding eigenvector for ˆ A ( x ) , and X is a set of grid points at which we plot the eigenvectors of A and ˆ A .
Revised: d of A ( x ) , ˆ u d ( x ) is the corresponding eigenvector for ˆ A ( x ) , and X is a set of grid points at which we plot the eigenvectors of A and ˆ A .

Original: ℓ ( A, ˆ A ) = 1 when each eigenvector of A ( x ) is parallel to the corresponding eigenvector of ˆ A ( x ) at all grid points x ∈ X .
Revised: ℓ ( A, ˆ A ) = 1 when each eigenvector of A ( x ) is parallel to the corresponding eigenvector of ˆ A ( x ) at all grid points x ∈ X .

Original: Our method accurately recovers the eigenvectors of the true metric, andthe learned metric is highly anisotropic inregions that overlap theobserved data.
Revised: Our method accurately recovers the eigenvectors of the true metric, and metric is highly anisotropic in that overlap the

Original: Our method accurately recovers the eigenspaces of the “Circular” ( ℓ ( A, ˆ A ) = 0 . 995 ) and “X Paths” ( ℓ ( A, ˆ A ) = 0 . 916 ) metrics.
Revised: Our method accurately recovers the eigenspaces of the “Circular” ( ℓ ( A, ˆ A ) = 0 . 995 ) and “X Paths” ( ℓ ( A, ˆ A ) = 0 . 916 ) metrics.

Original: It achieves lower accuracy with the “Mass Splitting” metric ( ℓ ( A, ˆ A ) = 0 . 839 ), struggling to capture the discontinuity in its eigenvectors at the x -axis and exhibiting numerical instability to the left of the y -axis, where the training trajectories diverge.
Revised: It achieves lower accuracy with the “Mass Splitting” metric ( ℓ ( A, ˆ A ) = 0 . 839 ), struggling to capture the discontinuity in its eigenvectors at the x -axis and exhibiting numerical instability to the left of the y -axis, where the training trajectories diverge.

Original: Note, however, that the alignment score ℓ ( A, ˆ A ) is in part measured at points that do not lie on the trajectory of the training data.
Revised: Note, however, that the alignment score ℓ ( A, ˆ A ) is in part measured at points that do not lie on the trajectory of the training data.

Original: We would not expect our method to accurately recover the metric in these regions; the fact that it largely does so for the “Circular” and “X Paths” examples reflects desirable smoothness properties of the neural parametrization of the Kantorovich potentials and the metric tensor.
Revised: We would not expect our method to accurately recover the metric in these regions; the fact that it largely does so for the “Circular” and “X Paths” examples reflects desirable smoothness properties of the neural parametrization of the Kantorovich potentials and the metric tensor.

Original: Row 2 of Figure 1 shows that our learned metric is highly anisotropic in regions that overlap with the training data.
Revised: Row 2 of Figure 1 shows that our learned metric is highly anisotropic in regions that overlap with the training data.

Original: In the “Mass Splitting” and “X Paths” examples, ˆ A ( x ) also has small condition number near the origin where the two paths diverge and cross, respectively; this behavior reflects expected uncertainty in low-energy directions of motion in this region.
Revised: In the “Mass Splitting” and “X Paths” examples, ˆ A ( x ) also has small condition number near the origin where the two paths diverge and cross, respectively; this behavior reflects expected uncertainty in low-energy directions of motion in this region.

Original: 5. P OPULATION - LEVEL TRAJECTORY INFERENCE WITH A LEARNED METRIC
Revised: 4. P OPULATION - LEVEL TRAJECTORY INFERENCE WITH A LEARNED METRIC

Original: The use of optimal transport for inferring cellular development trajectories from scRNA data is well-established.
Revised: scRNA sequencing (scRNA-seq) is a recently-developed technique that allows biologists to study the set of mRNA molecules (the “transcriptome”) of individual cells at high resolution (Haque et al., 2017). As scRNA-seq is a destructive process, any individual cell’s RNA can be sequenced only once, impeding the use of this technology to study the change in a cell’s transcriptome over time. This has led to interest in methods that use population-level scRNA-seq data to infer the temporal evolution of an individual cell’s scRNA-seq profile. Optimal transport-based techniques are well-established tools for solving these trajectory inference problems.

Original: Schiebinger et al.
Revised: Schiebinger et al.

Original: and Tong et al.
Revised: and Tong et al.

Original: both use optimal transport to infer cellular development trajectories but assume a Euclidean metric on gene expression space.
Revised: both use optimal transport to infer cellular development trajectories but assume a Euclidean metric on gene expression space.

Original: In this section, we use our method to relax this strong assumption by learning a metric for the ground space from scRNA data.
Revised: In this section, we use our method to relax this strong assumption by learning a metric for the ground space from scRNA data.

Original: We incorporate ˆ A ( x ) into a downstream trajectory inference task and show that this strategy yields more accurate trajectories than a baseline without a learned metric and a baseline using the Euclidean metric.
Revised: We incorporate ˆ A ( x ) into a downstream trajectory inference task and show that this strategy yields more accurate trajectories than a baseline without a learned metric and a baseline using the Euclidean metric.

Original: scRNA data.
Revised: scRNA data.

Original: We perform trajectory inference experiments with the scRNA data drawn from
Revised: We perform trajectory inference experiments with the scRNA data drawn from

Original: Schiebinger et al.
Revised: Schiebinger et al.

Original: This data consists of force-directed layout embedding coordinates of gene expression data collected over 18 days of reprogramming (39 time points total).
Revised: This data consists of force-directed layout embedding coordinates of gene expression data collected over 18 days of reprogramming (39 time points total).

Original: We construct populations X t i for i = 1 , . . .
Revised: We construct populations X t i for i = 1 , . . .

Original: , 39 by drawing 500 samples per time point in the original data; this sampling uses 8.25% of the available data on average.
Revised: , 39 by drawing 500 samples per time point in the original data; this sampling uses 8.25% of the available data on average.

Original: We follow the same procedure as in Section 5.1 to learn a metric tensor ˆ A ( x ) from subsequent pairs of samples ( X t i , X t
Revised: We follow the same procedure as in Section 4.1 to learn a metric tensor ˆ A ( x ) from subsequent pairs of samples ( X t i , X t

Original: i +1 ) .
Revised: i +1 ) .

Original: Learning the tensor takes 16 minutes on a single V100 GPU.
Revised: Learning the tensor takes 16 minutes on a single V100 GPU.

Original: For the downstream trajectory inference task, we keep one out of every k time points in the original data for k = 2 , ..., 19 to obtain a new collection of time points ¯ t ℓ with ℓ ∈ { nk : n ∈ N , n < 39 k } .
Revised: For the downstream trajectory inference task, we keep one out of every k time points in the original data for k = 2 , ..., 19 to obtain a new collection of time points ¯ t ℓ with ℓ ∈ { nk : n ∈ N , n < 39 k } .

Original: We then perform trajectory inference between subsequent retained time points (¯ t ℓ , ¯ t ℓ +1 )
Revised: We then perform trajectory inference between subsequent retained time points (¯ t ℓ , ¯ t ℓ +1 )

Original: (using all of the available data for these time points) by optimizing the following objective:
Revised: (using all of the available data for these time points) by optimizing the following objective:

Original: [Equation]
Revised: [Equation]

Original: and corresponding advected samples versus k , with our learned metric (orange), the Euclidean metric A ( x ) ≡ I (green), and no regularizer (blue).
Revised: with our learned metric (orange), the Euclidean metric A ( x ) ≡ I ( regularizer ( interpolants obtained using Schiebinger et al.

Original:  Here S ϵ is the Sinkhorn divergence (Feydy et al., 2019) between the target data X ¯ t ℓ +1 and the advected samples ϕ t =1 ( v θ )
Revised: (2019)’s method Here S ϵ is the Sinkhorn divergence (Feydy et al., 2019) between the target data X ¯ t ℓ +1 and the advected samples ϕ t =1 ( v θ )

Original: [ X ¯ t ℓ ] ; ϕ t =1 ( v θ ) denotes the solution map resulting from advecting a particle through a neurally-parametrized time-varying velocity field v t,θ for one unit of time; and t j ∈ [0 , 1] .
Revised: [ X ¯ t ℓ ] ; ϕ t =1 ( v θ ) denotes the solution map resulting from advecting a particle through a neurally-parametrized time-varying velocity field v t,θ for one unit of time; and t j ∈ [0 , 1] .

Original: The Sinkhorn divergence used here is S ϵ
Revised: The Sinkhorn divergence used here is S ϵ

Original: ( α, β ) =
Revised: ( α, β ) =

Original: W ϵ ( α, β ) − 12 W ϵ ( α, α ) − 12 W ϵ ( β, β ) ; it is a debiased variant of the entropically-regularized 2-Wasserstein distance W ϵ .
Revised: W ϵ ( α, β ) − 12 W ϵ ( α, α ) − 12 W ϵ ( β, β ) ; it is a debiased variant of the entropically-regularized 2-Wasserstein distance W ϵ .

Original: This is similar to the method of Tong et al.
Revised: This is similar to the method of Tong et al.

Original: (2020), where we replace the log-likelihood fitting loss with the Sinkhorn divergence.
Revised: (2020), where we replace the log-likelihood fitting loss with the Sinkhorn divergence.

Original: We use
Revised: We use

Original: GeomLoss (Feydy et al., 2019) to compute the Sinkhorn divergence between the target and advected samples efficiently.
Revised: GeomLoss (Feydy et al., 2019) to compute the Sinkhorn divergence between the target and advected samples efficiently.

Original: We follow Tong et al.
Revised: We follow Tong et al.

Original: and assess the quality of inferred trajectories by measuring the W 1 distance (EMD) between left-out time points in the ground truth data and advected samples at corresponding times in the inferred trajectories.
Revised: and assess the quality of inferred trajectories by measuring the W 1 distance (EMD) between left-out time points in the ground truth data and advected samples at corresponding times in the inferred trajectories.

Original: These distances are of the form
Revised: These distances are of the form

Original: W 1 (cid:0) X t i , ϕ t =
Revised: W 1 (cid:0) X t i , ϕ t =

Original: ˆ t i ( v θ )
Revised: ˆ t i ( v θ )

Original: [ X ¯ t ℓ ] (cid:1) , so we compare ground truth samples at each left-out time t
Revised: [ X ¯ t ℓ ] (cid:1) , so we compare ground truth samples at each left-out time t

Original: i = nk
Revised: i = nk

Original: + h ( h = 1 , ..., k − 1 ) with samples from X ¯ t ℓ advected through v θ for time
Revised: + h ( h = 1 , ..., k − 1 ) with samples from X ¯ t ℓ advected through v θ for time

Original: ˆ t
Revised: ˆ t

Original: i = hk .
Revised: i = hk .

Original: For each value of k , we record the average EMD between left out samples and their inferred counterparts and compare our method to a baseline where λ = 0 in (8) and a Euclidean baseline where A ( x ) ≡ I .
Revised: For each value of k , we record the average EMD between left out samples and their inferred counterparts and compare our method to a baseline where λ = 0 in (7) and a Euclidean baseline where A ( x ) ≡ I .

Original: Figure 2 shows that our learned metric improves the quality of trajectory inference for nearly all values of k , whereas the Euclidean baseline A ( x ) ≡ I generally results in less accurate trajectories than the non-regularized baseline.
Revised: We also compare against interpolants obtained via the method of Schiebinger et al. As their Waddington OT method solves a static optimal transport problem and pushes cells through transport maps between fixed time steps, a direct comparison against our dynamical method is not possible. However, we attempt to replicate their validation by geodesic interpolation as closely as possible.

Original: As expected, the gap between our results and both baselines increases for large values of k ; for example, using our learned metric results in a 19.7% mean reduction in Trajectories inferred using our learned metric tensor (second column) more closely follow the manifold structure of the ground truth data than the non-regularized baseline trajectories (first column), where particles follow nearly straight-line paths between observed time points.
Revised: We compute static optimal transport couplings between data at subsequent retained time points (¯ t ℓ , ¯ t ℓ +1 ) , linearly interpolate between coupled data points, and measure the W 1 distance (EMD) between left-out time points in the ground truth data and advected samples at corresponding times in the interpolated trajectories. Figure 2 shows that our learned metric improves on the no-metric baseline (analogous to the “Base” model of Tong et al. and the approach based on Schiebinger et al.

Original:  EMD between left-out samples and the corresponding advected samples for k ≥ 10 .
Revised: for nearly all values of k , whereas the Euclidean baseline A ( x ) ≡ I (analogous to the “Base+E” model of Tong et al. (2020)) generally results in less accurate trajectories than the non-regularized baseline. As expected, the gap between our results and both baselines increases for large values of k ; for example, using our learned metric results in a 19.7% mean reduction in EMD between left-out samples and the corresponding advected samples for k ≥ 10 .

Original: This observation indicates that including a learned metric has a larger impact on the inferred trajectories when the ground truth data is sampled sparsely in time.
Revised: This observation indicates that including a learned metric has a larger impact on the inferred trajectories when the ground truth data is sampled sparsely in time.

Revised: Trajectories inferred using our learned metric tensor (second column) more closely follow the manifold structure of the ground truth data than the non-regularized baseline trajectories (first column), where particles follow nearly straight-line paths between observed time points.

Original: Figure 3 compares the trajectories inferred using our method to the non-regularized baseline and ground truth trajectories for k = 3 , where the time sampling is sufficiently dense that the baseline performs well, and k = 15 , where our learned metric tensor substantially improves the quality of the inferred trajectories.
Revised: Figure 3 compares the trajectories inferred using our method to the non-regularized baseline and ground truth trajectories for k = 3 , where the time sampling is sufficiently dense that the baseline performs well, and k = 15 , where our learned metric tensor substantially improves the quality of the inferred trajectories.

Original: Whereas the non-regularized baseline simply advects particles from the base measures to the corresponding targets along nearly-straight paths, the learned metric biases the trajectory to follow the highly-curved paths taken by the ground truth data.
Revised: Whereas the non-regularized baseline simply advects particles from the base measures to the corresponding targets along nearly-straight paths, the learned metric biases the trajectory to follow the highly-curved paths taken by the ground truth data.

Original: These results suggest that in settings where data collection is expensive and samples collected at a small subset of times are of primary interest, our method enables the plausible inference of particle positions at intermediate time points at the cost of little additional data and computation.
Revised: These results suggest that in settings where data collection is expensive and samples collected at a small subset of times are of primary interest, our method enables the plausible inference of particle positions at intermediate time points at the cost of little additional data and computation.

Original: 5. I NDIVIDUAL - LEVEL TRAJECTORY INFERENCE WITH A LEARNED METRIC
Revised: 4. I NDIVIDUAL - LEVEL TRAJECTORY INFERENCE WITH A LEARNED METRIC

Original: In this section, we learn a metric ˆ A ( x ) from time-stamped sightings of snow geese during their spring migration.
Revised: In this section, we learn a metric ˆ A ( x ) from time-stamped sightings of snow geese during their spring migration.

Original: We then compute an ˆ A -geodesic between the initial and final point of GPS-tagged snow geese and show that this provides a reasonable coarse approximation to the geese’s ground truth trajectories.
Revised: We then compute an ˆ A -geodesic between the initial and final point of GPS-tagged snow geese and show that this provides a reasonable coarse approximation to the geese’s ground truth trajectories.

Original: Migratory paths as geodesics on a latent manifold.
Revised: Migratory paths as geodesics on a latent manifold.

Original: Many migratory bird species return to their summer breeding sites and wintering grounds with high spatial fidelity, sometimes returning to within 500 meters of their usual breeding location (Mowbray et al., 2020).
Revised: Many migratory bird species return to their summer breeding sites and wintering grounds with high spatial fidelity, sometimes returning to within 500 meters of their usual breeding location (Mowbray et al., 2020).

Original: Given this behavior, a boundary value problem that fixes the endpoints of the birds’ migratory trajectories is a reasonable model for bird migration.
Revised: Given this behavior, a boundary value problem that fixes the endpoints of the birds’ migratory trajectories is a reasonable model for bird migration.

Original: Once the endpoints are fixed, we must specify an objective that birds plausibly optimize to determine their trajectories.
Revised: Once the endpoints are fixed, we must specify an objective that birds plausibly optimize to determine their trajectories.

Original: Absent further information, minimizing total kinetic energy over their trajectory is a reasonable choice.
Revised: Absent further information, minimizing total kinetic energy over their trajectory is a reasonable choice.

Original: However, measuring kinetic energy with respect to the
Revised: However, measuring kinetic energy with respect to the

Original: Euclidean metric leads to straight-line paths, which typically do not agree with observed migration trajectories.
Revised: Euclidean metric leads to straight-line paths, which typically do not agree with observed migration trajectories.

Original: We use our method to learn a metric from untagged snow goose sightings that agrees with the data for this species.
Revised: We use our method to learn a metric from untagged snow goose sightings that agrees with the data for this species.

Original: This metric summarizes the factors that birds may use to modify their migratory paths locally, such as local weather conditions, food availability, and predatory pressures.
Revised: This metric summarizes the factors that birds may use to modify their migratory paths locally, such as local weather conditions, food availability, and predatory pressures.

Original: Snow goose data.
Revised: Snow goose data.

Original: The training data for this experiment consists of time-stamped sightings of untagged snow geese ( Anser caerulescens ) across the U.S. and Canada during their spring migration.
Revised: The training data for this experiment consists of time-stamped sightings of untagged snow geese ( Anser caerulescens ) across the U.S. and Canada during their spring migration.

Original: Birdwatchers report the latitude-longitude coordinates of their goose sightings to eBird, which aggregates this data and makes it available to researchers.
Revised: Birdwatchers report the latitude-longitude coordinates of their goose sightings to eBird, which aggregates this data and makes it available to researchers.

Original: We treat the sightings at time t as samples from a time-indexed spatial distribution of birds ρ t and train our metric on subsequent pairs of bird distributions ( ρ t , ρ t +1 ) .
Revised: We treat the sightings at time t as samples from a time-indexed spatial distribution of birds ρ t and train our metric on subsequent pairs of bird distributions ( ρ t , ρ t +1 ) .

Original: This implies that our method does not have access to any complete goose trajectories when learning the metric .
Revised: This implies that our method does not have access to any complete goose trajectories when learning the metric .

Original: Furthermore, any given goose’s position was likely recorded at most once in the training data.
Revised: Furthermore, any given goose’s position was likely recorded at most once in the training data.

Revised: We do not expect untagged goose sightings to contain enough information to predict high-frequency detail in migratory trajectories, but this data is far cheaper to obtain than complete migratory trajectories, which are typically recorded via GPS trackers attached to individual geese.

Revised: The widespread availability and low cost of obtaining untagged bird sighting data motivates the use of our method for bird trajectory inference

Revised: The training data is drawn from the eBird basic dataset (Sullivan et al., 2009), current as of February 2022.

Revised: We bin the sightings by month of observation and use our algorithm to learn a metric tensor ˆ

Original:  We use this learned metric tensor to compute an ˆ A -geodesic between the initial and final observation of several snow geese along their migratory paths.
Revised: A ( x ) from populations X i consisting of the spatial coordinates of snow goose sightings in month i for i = 0 , ..., 5 (i.e. January to June). This training data is depicted in Figure 5 in Appendix D.3. We use this learned metric tensor to compute an ˆ A -geodesic between the initial and final observation of several snow geese along their migratory paths.

Original: This data is drawn from the Banks Island Snow
Revised: This data is drawn from the Banks Island Snow

Original: Goose study as hosted on Movebank (Kays et al., 2021).
Revised: Goose study as hosted on Movebank (Kays et al., 2021).

Original: We estimate an ˆ A -geodesic between initial and final points on each path by learning a time-varying velocity field v t,θ optimizing the following objective:
Revised: We estimate an ˆ A -geodesic between initial and final points on each path by learning a time-varying velocity field v t,θ optimizing the following objective:

Original: [Equation]
Revised: [Equation]

Original: As in Section 5.2, the velocity field v t,θ is neurally parametrized, 0{+
+} = t 0 ≤ · · · t j ≤ · · · t m = 1 , and we optimize (9) using AdamW.
Revised: As in Section 4.2, the velocity field v t,θ is neurally parametrized, {+
+} = t 0 ≤ · · · t j ≤ · · · t m = 1 , and we optimize (8) using AdamW.

Original: We set the initial condition x (0) =
Revised: We set the initial condition x (0) =

Original: x 0 such that x 0 is the goose’s initial position on its migratory trajectory and use a Euclidean norm penalty to force its final position x (1) along the inferred trajectory to match the true final position x 1 .
Revised: x 0 such that x 0 is the goose’s initial position on its migratory trajectory and use a Euclidean norm penalty to force its final position x (1) along the inferred trajectory to match the true final position x 1 .

Original: Figure 4 compares the ˆ A -geodesics obtained using the method described above (right plot) to the ground truth goose trajectories (left plot).
Revised: Figure 4 compares the ˆ A -geodesics obtained using the method described above (center plot) to the ground truth goose trajectories (left plot).

Original: While population-level data (in the form of untagged snow goose sightings) does not provide sufficient information to reconstruct the migratory paths of individual geese perfectly, the inferred trajectories accurately capture the curved structure of the ground truth trajectories.
Revised: While population-level data (in the form of untagged snow goose sightings) does not provide sufficient information to reconstruct the migratory paths of individual geese perfectly, the inferred trajectories accurately capture the curved structure of the ground truth trajectories.

Original: We evaluate our method’s performance by computing the dynamic time warping (DTW) distance (Berndt & Clifford, 1994) between each inferred goose trajectory and the corresponding ground truth trajectory.
Revised: We evaluate our method’s performance by computing the dynamic time warping (DTW) distance (Berndt & Clifford, 1994) between each inferred goose trajectory and the corresponding ground truth trajectory.

Original: We also generate straight-line paths (i.e. Euclidean geodesics) between the initial and final points of each ground truth goose trajectory and compute DTW distances between these paths and the ground truth.
Revised: We also generate straight-line paths (i.e. Euclidean geodesics) between the initial and final points of each ground truth goose trajectory and compute DTW distances between these paths and the ground truth.

Original: Our method results in a 26.9% reduction in mean DTW distance between the inferred and ground truth trajectories relative to the Euclidean baseline.
Revised: Our method results in a 26.9% reduction in mean DTW distance between the inferred and ground truth trajectories relative to the Euclidean baseline.

Original:  D ISCUSSION AND C ONCLUSION
Revised:  D ISCUSSION AND C ONCLUSION

Original: We have introduced an optimal transport-based method for learning a Riemannian metric from cross-sectional samples of populations evolving over time on a latent Riemannian manifold.
Revised: We have introduced an optimal transport-based method for learning a Riemannian metric from cross-sectional samples of populations evolving over time on a latent Riemannian manifold.

Original: Our method accurately recovers metrics from cross-sections of populations moving along geodesics on a manifold, improves the quality of trajectory inference on sparsely-sampled scRNA data at low data and compute cost, and allows us to approximate individual trajectories of migrating birds using information from untagged sightings.
Revised: Our method accurately recovers metrics from cross-sections of populations moving along geodesics on a manifold, improves the quality of trajectory inference on sparsely-sampled scRNA data at low data and compute cost, and allows us to approximate individual trajectories of migrating birds using information from untagged sightings.

Original: One key limitation of our work is that it learns a Riemannian metric on R D , whereas large swaths of data naturally lie on graphs.
Revised: One key limitation of our work is that it learns a Riemannian metric on R D , whereas large swaths of data naturally lie on graphs.

Original: Future work might consider extending this algorithm learn from data defined on graphs.
Revised: Future work might consider extending this algorithm learn from data defined on graphs.

Original: In this setting, our optimal transport-based algorithm would correspond to learning per-edge costs from observations of flows on a graph; a recent work of Heitz et al.
Revised: In this setting, our optimal transport-based algorithm would correspond to learning per-edge costs from observations of flows on a graph; a recent work of Heitz et al.

Original: (2021) presents another optimal transport-based technique for solving this problem.
Revised: (2021) presents another optimal transport-based technique for solving this problem.

Original: Future work may also extend our methods to learn a metric from temporal snapshots of Schrödinger bridges.
Revised: Future work may also extend our methods to learn a metric from temporal snapshots of Schrödinger bridges.

Original:  R EPRODUCIBILITY S TATEMENT
Revised: A B ACKGROUND ON R IEMANNIAN GEOMETRY

Revised: A Riemannian manifold M = ( M, g ) is a differentiable manifold M equipped with a Riemannian metric g . Throughout this paper, we take M = R D and focus our attention on the metric g . A

Revised: Riemannian metric assigns an inner product ⟨· , ·⟩ p to the tangent space T p M of each p ∈ M in a way that varies smoothly with p .

Revised: In the case where M = R D , we may simply identify all tangent spaces

Revised: T p M with R D ; here a metric g amounts to a spatially-varying inner product on R D .

Revised: Since any inner product on R D may be computed as ⟨ u, v ⟩ A =

Revised: u T Av for some A ≻ 0 , a Riemannian metric g on

Revised: M = R D is specified by a smooth field of positive definite matrices

Revised: A ( x ) : R D → S D

Revised: ++ (we use

Revised: S D

Revised: ++ to denote the set of positive definite D × D matrices).

Original: To enable the reproduction our results we have attached our anonymized source code to this submission as a zip file.
Revised: We use ⟨ u, v ⟩ A ( x ) = u T A ( x ) v to denote the metric at a point x ∈ R D . An inner product ⟨· , ·⟩ A ( x ) induces a norm ∥ v ∥ A ( x ) : = (cid:112) ⟨ v, v ⟩ A ( x ) .

Revised: Given this norm, we define the length of a continuously differentiable curve γ : [0 , 1] → M to be ℓ ( γ ) =

Original: {+(cid:82) 1 
+}
Revised: {+(cid:82) 1 
+}

Revised: ∥ ˙ γ ( t ) ∥

Revised: A ( γ ( t )) dt .

Revised: The

Revised: A -geodesic distance d A ( x, y ) between two points x, y ∈ M is then the infimum of the length ℓ ( γ ) of continuously differentiable curves γ such that γ (0) =

Revised: x and γ (1) =

Revised: y .

Original: Appendix B includes a thorough description of the procedures used to generate synthetic data (Appendix B.1), pre-process the scRNA data (Appendix B.2), and pre-process the snow goose data (Appendix B.3) used in our experiments.
Revised: Feldman & McCann (2002) study Monge’s transport problem on a Riemannian manifold M = ( M, g ) ; we focus on the case M = R D with g specified by a field of matrices A ( x ) ≻ 0 . Given two compactly-supported densities ρ 0 , ρ 1 on M , Monge’s problem seeks a pushforward s of ρ 0 onto ρsolving the following problem:

Revised: [Equation]

Revised: Feldman & McCann (2002) show that under mild technical conditions, (9) has a possibly non-unique solution s .

Revised: Intuitively, this map is a pushforward of ρ 0 onto ρ 1 minimizing the average geodesic distance between matched units of probability mass.

Revised: The optimal value of (9) is the 1-Wasserstein distance between ρ 0 and ρ 1 , which we denote by W 1 ,

Revised: A

Original: It also includes exhaustive descriptions of our model parametrizations and hyperparameters used in each experiment.
Revised: ( ρ 0 , ρ 1 ) . If ϕ is continuously differentiable on M , then it is 1-Lipschitz on M if and only if ∥∇ ϕ ( x ) ∥ A − 1 ( x ) ≤ 1 for all x ∈ M . Here, we use ∇ ϕ ( x ) to denote the vector of partial derivatives of ϕ at x .

Revised: Armed with this local characterization of Lipschitz continuity, we may define the following dual problem to (9):

Revised: [Equation]

Revised: If the minimum in (10) is attained by some Kantorovich potential ϕ , the optimal values of (9) and (10) coincide and the Lipschitz bound for ϕ is tight at pairs of points ( x, y ) arising from a Monge map s solving (9) (Feldman & McCann, 2002, Lemma 4).

Original: Our source code includes all data required for our experiments, either directly or in the form of code that automatically downloads the relevant data.
Revised: Feldman & McCann (2002, Lemma 10) also show that given a geodesic γ x : [0 , 1] → M between a pair of matched points ( x, s ( x )) and t ∈ (0 , 1) , we have ∇ ϕ ( γ x ( t )) = − [Equation] Intuitively, ∇ ϕ points along geodesics γ x joining pairs of points in the support of ρ 0 and ρ 1 that are matched by the Monge map s .

Original: A B ACKGROUND ON DYNAMICAL OPTIMAL TRANSPORT
Revised: C B ACKGROUND ON DYNAMICAL OPTIMAL TRANSPORT

Original: Let α and β be probability measures defined on Ω ⊆ R D .
Revised: Let α and β be probability measures defined on Ω ⊆ R D .

Original: The Benamou-Brenier formulation of the 2-Wasserstein distance W 2 ( α, β ) defines it as the solution to a fluid-dynamical problem (Benamou &
Revised: The Benamou-Brenier formulation of the 2-Wasserstein distance W 2 ( α, β ) defines it as the solution to a fluid-dynamical problem (Benamou &

Original: Brenier, 2000)
Revised: Brenier, 2000)

Original: [Equation]
Revised: [Equation]

Original: subject to the constraints that ρ 0 =
Revised: subject to the constraints that ρ 0 =

Original: α , ρ 1 = β and the continuity equation ∂ρ t ∂t + ∇ · ( ρ t v t ) = 0 .
Revised: α , ρ 1 = β and the continuity equation ∂ρ t ∂t + ∇ · ( ρ t v t ) = 0 .

Original: Solving this problem yields a time-varying velocity field v t ( x ) that transports α ’s mass to β ’s along a curve ρ t of probability measures.
Revised: Solving this problem yields a time-varying velocity field v t ( x ) that transports α ’s mass to β ’s along a curve ρ t of probability measures.

Original: The kinetic energy in the integrand of (10) encourages probability mass to travel in straight lines.
Revised: The kinetic energy in the integrand of (11) encourages probability mass to travel in straight lines.

Original: This assumption is occasionally undesirable, but it is straightforward to modify Eq.
Revised: This assumption is occasionally undesirable, but it is straightforward to modify Eq.

Original: (10) to encourage v to point in specified directions (Hug et al., 2015):
Revised: (11) to encourage v to point in specified directions (Hug et al., 2015):

Original: [Equation]
Revised: [Equation]

Original: Using the language developed in Section 3, A ( x ) ≻ 0 specifies a metric g on the Riemannian manifold M = (cid:0) R D , g (cid:1) , and v t ( x ) T A ( x ) v t ( x ) =
Revised: Using the language developed in Appendix A, A ( x ) ≻ 0 specifies a metric g on the Riemannian manifold M = (cid:0) R D , g (cid:1) , and v t ( x ) T A ( x ) v t ( x ) =

Original: ∥
Revised: ∥

Original: v t ( x ) ∥ 

Revised: v t ( x ) ∥

Original: A ( x ) .
Revised: A ( x ) .

Original: Eq.
Revised: Eq.

Original: (11) encourages v to be aligned with the eigenvectors u 1 corresponding to the minimal eigenvalues λ 1 of the matrices
Revised: (12) encourages v to be aligned with the eigenvectors u 1 corresponding to the minimal eigenvalues λ 1 of the matrices

Original: A ( x ) .
Revised: A ( x ) .

Original: While works like (Hug et al., 2015; Zhang et al., 2022) investigate the modeling applications of anisotropic optimal transport, they only consider the case where the Riemannian metric A ( x ) is available a priori.
Revised: While works like (Hug et al., 2015; Zhang et al., 2022) investigate the modeling applications of anisotropic optimal transport, they only consider the case where the Riemannian metric A ( x ) is available a priori.

Original: This assumption is unrealistic for many problem domains, motivating our model, which learns a metric from cross-sectional samples from populations evolving over time.
Revised: This assumption is unrealistic for many problem domains, motivating our model, which learns a metric from cross-sectional samples from populations evolving over time.

Original: B E XPERIMENT DETAILS B. M ETRIC RECOVERY
Revised: D E XPERIMENT DETAILS D. M ETRIC RECOVERY

Original: Circular data.
Revised: Circular data.

Original: The ground truth metric tensor for this example is A ( x, y ) =
Revised: The ground truth metric tensor for this example is A ( x, y ) =

Original: I − v ( x, y ) v ( x, y ) T , where v ( x, y ) =
Revised: I − v ( x, y ) v ( x, y ) T , where v ( x, y ) =

Original: (cid:18) − y √ x 2 + y 2 , x √ x 2 + y(cid:19) .
Revised: (cid:18) − y √ x 2 + y 2 , x √ x 2 + y(cid:19) .

Original: To generate the training data, we begin by drawing 100 samples each from 4 isotropic normal distributions with standard deviation σ = 0 .
Revised: To generate the training data, we begin by drawing 100 samples each from 4 isotropic normal distributions with standard deviation σ = 0 .

Original: 1 whose means are µ i ∈ { (1 , 0) , ( − 1 , 0) , ( − 1 , − 1) , (0 , 1) } .
Revised: 1 whose means are µ i ∈ { (1 , 0) , ( − 1 , 0) , ( − 1 , − 1) , (0 , 1) } .

Original: We randomly pair samples from subsequent distributions and compute A -geodesics between each pair by solving problem (9).
Revised: We randomly pair samples from subsequent distributions and compute A -geodesics between each pair by solving problem (8).

Original: We implement (9) in Pytorch using a time-invariant vector field v θ parametrized by a fully connected two-layer neural network with ELU nonlinearities and 64 hidden dimensions.
Revised: We implement (8) in Pytorch using a time-invariant vector field v θ parametrized by a fully connected two-layer neural network with ELU nonlinearities and 64 hidden dimensions.

Original: We set λ = 1 and solve the initial value problem ˙ x ( t ) =
Revised: We set λ = 1 and solve the initial value problem ˙ x ( t ) =

Original: v θ ( x ( t ))
Revised: v θ ( x ( t ))

Original: ; x (0) = x 0 using the explicit Adams solver in torchdiffeq’s odeint with default hyperparameters (Chen et al., 2018a).
Revised: ; x (0) = x 0 using the explicit Adams solver in torchdiffeq’s odeint with default hyperparameters (Chen et al., 2018a).

Original: We optimize the objective using AdamW with learning rate 10 − 3 and weight decay 10 − 3 and train for 100 epochs per pair of samples.
Revised: We optimize the objective using AdamW with learning rate 10 − 3 and weight decay 10 − 3 and train for 100 epochs per pair of samples.

Original: We then draw 24 points at equispaced times t i ∈ [0 , 1] from each resulting geodesic and aggregate across geodesics to form the observed populations
Revised: We then draw 24 points at equispaced times t i ∈ [0 , 1] from each resulting geodesic and aggregate across geodesics to form the observed populations

Original: X t i .
Revised: X t i .

Original: We then use our method to recover ˆ A − 1 ( x ) from the X t i .
Revised: We then use our method to recover ˆ A − 1 ( x ) from the X t i .

Original: We parametrize the scalar potentials in (5) as a single-hidden-layer neural net with 32 hidden dimensions and Softplus activation.
Revised: We parametrize the scalar potentials in (4) as a single-hidden-layer neural net with 32 hidden dimensions and Softplus activation.

Original: We parametrize the matrix field ˆ A − 1 ( x ) as A − 1 ( x ) =
Revised: We parametrize the matrix field ˆ A − 1 ( x ) as A − 1 ( x ) =

Original: Q ( x ) T Q ( x ) + 10 − 3 I , where Q ( x ) is a two-layer neural network with Softplus activations and 32 hidden dimensions.
Revised: Q ( x ) T Q ( x ) + 10 − 3 I , where Q ( x ) is a two-layer neural network with Softplus activations and 32 hidden dimensions.

Original: The strength of the gradient penalty (6) is 10 − 3 when training the potentials ϕ
Revised: The strength of the gradient penalty (5) is 10 − 3 when training the potentials ϕ

Original: in the first step of alternation and 10 − 4 in the next step; it is 1 when training ˆ A − 1 .
Revised: in the first step of alternation and 10 − 4 in the next step; it is 1 when training ˆ A − 1 .

Original: The strength of the regularization (7) is 10 3 .
Revised: The strength of the regularization (6) is 10 3 .

Original: We carry out a two steps of the alternating scheme by training with AdamW with learning rate 10 − 2 and weight decay 5 ∗ 10 − 1 .
Revised: We carry out a two steps of the alternating scheme by training with AdamW with learning rate 10 − 2 and weight decay 5 ∗ 10 − 1 .

Original: We train for 300 epochs for the ϕ step and 1,000 epochs for the A step.
Revised: We train for 300 epochs for the ϕ step and 1,000 epochs for the A step.

Original: We evaluate the alignment score ℓ ( A, ˆ A ) on a 100 × 100 grid overlaid on a box of radius 1.5.
Revised: We evaluate the alignment score ℓ ( A, ˆ A ) on a 100 × 100 grid overlaid on a box of radius 1.5.

Original: Mass splitting data.
Revised: Mass splitting data.

Original: The ground truth metric tensor for this example is A ( x, y ) =
Revised: The ground truth metric tensor for this example is A ( x, y ) =

Original: I − v ( x, y ) v ( x, y ) T , where v ( x, y ) =
Revised: I − v ( x, y ) v ( x, y ) T , where v ( x, y ) =

Original: (cid:16) 1 √ 2 , 1 √ 2 (cid:17) for y ≥ 0 and v ( x, y ) = (cid:16) 1 √ 2 , − 1 √ 2 (cid:17) for y < 0 .
Revised: (cid:16) 1 √ 2 , 1 √ 2 (cid:17) for y ≥ 0 and v ( x, y ) = (cid:16) 1 √ 2 , − 1 √ 2 (cid:17) for y < 0 .

Original: To generate the training data, we begin by drawing 100 samples each from a standard normal distribution and from a mixture of two isotropic normal distributions with unit variance and mixture components centered at (10 , 10) and (10 , − 10) .
Revised: To generate the training data, we begin by drawing 100 samples each from a standard normal distribution and from a mixture of two isotropic normal distributions with unit variance and mixture components centered at (10 , 10) and (10 , − 10) .

Original: We randomly pair samples from subsequent distributions and compute A -geodesics between each pair by solving problem (9) using the method described in the circular data section.
Revised: We randomly pair samples from subsequent distributions and compute A -geodesics between each pair by solving problem (8) using the method described in the circular data section.

Original: We then draw 10 points at equispaced times t i ∈ [0 , 1] from each resulting geodesic and aggregate across geodesics to form the observed populations
Revised: We then draw 10 points at equispaced times t i ∈ [0 , 1] from each resulting geodesic and aggregate across geodesics to form the observed populations

Original: X t i .
Revised: X t i .

Original: We then use our method to recover ˆ A − 1 ( x ) from the X t i .
Revised: We then use our method to recover ˆ A − 1 ( x ) from the X t i .

Original: We parametrize the scalar potentials in (5) as a single-hidden-layer neural net with 32 hidden dimensions and Softplus activation.
Revised: We parametrize the scalar potentials in (4) as a single-hidden-layer neural net with 32 hidden dimensions and Softplus activation.

Original: We parametrize the matrix field ˆ A − 1 ( x ) as A − 1 ( x ) =
Revised: We parametrize the matrix field ˆ A − 1 ( x ) as A − 1 ( x ) =

Original: Q ( x ) T Q ( x ) + 10 − 3 I , where Q ( x ) is a twolayer neural network with Softplus activations and 32 hidden dimensions.
Revised: Q ( x ) T Q ( x ) + 10 − 3 I , where Q ( x ) is a twolayer neural network with Softplus activations and 32 hidden dimensions.

Original: The strength of the gradient penalty (6) is 2 when training the potentials ϕ and 1 when training ˆ A .
Revised: The strength of the gradient penalty (5) is 2 when training the potentials ϕ and 1 when training ˆ A .

Original: The strength of the regularization (7) is 10 6 .
Revised: The strength of the regularization (6) is 10 6 .

Original: We carry out a single step of the alternating scheme by training with
Revised: We carry out a single step of the alternating scheme by training with

Original: AdamW with learning rate 5 ∗ 10 − 3 and weight decay 10 − 3 .
Revised: AdamW with learning rate 5 ∗ 10 − 3 and weight decay 10 − 3 .

Original: We train for 600 epochs for the ϕ step and 20,000 epochs for the A step.
Revised: We train for 600 epochs for the ϕ step and 20,000 epochs for the A step.

Original: We evaluate the alignment score ℓ ( A, ˆ A ) on a 100 × 100 grid overlaid on the rectangular region [ − 2 . 5 , 15] × [ − 15 , 15] .
Revised: We evaluate the alignment score ℓ ( A, ˆ A ) on a 100 × 100 grid overlaid on the rectangular region [ − 2 . 5 , 15] × [ − 15 , 15] .

Original: X-Paths data.
Revised: X-Paths data.

Original: The ground truth metric tensor for this example is A ( x, y ) =
Revised: The ground truth metric tensor for this example is A ( x, y ) =

Original: I − v ( x, y ) v ( x, y ) T .
Revised: I − v ( x, y ) v ( x, y ) T .

Original: Here we define v ( x, y )
Revised: Here we define v ( x, y )

Original: =
Revised: =

Original: α ( x, y ) v 1 ( x, y ) + β ( x, y ) v 2 ( x, y ) , where v 1 ( x, y ) =
Revised: α ( x, y ) v 1 ( x, y ) + β ( x, y ) v 2 ( x, y ) , where v 1 ( x, y ) =

Original: (cid:16) 1 √ 2 , 1 √ 2 (cid:17) and v 2 ( x, y ) = (cid:16) 1 √ 2 , − 1 √ 2 (cid:17) .
Revised: (cid:16) 1 √ 2 , 1 √ 2 (cid:17) and v 2 ( x, y ) = (cid:16) 1 √ 2 , − 1 √ 2 (cid:17) .

Original: We then define α ( x, y ) = 1 .
Revised: We then define α ( x, y ) = 1 .

Original: 25 tanh( ReLU ( x · y )) and β ( x, y )
Revised: 25 tanh( ReLU ( x · y )) and β ( x, y )

Original: = − 1 .
Revised: = − 1 .

Original: 25 tanh( ReLU ( − x · y )) .
Revised: 25 tanh( ReLU ( − x · y )) .

Original: Intuitively, α should be large in quadrants 1 and 3 and β should be large in quadrants 2 and 4.
Revised: Intuitively, α should be large in quadrants 1 and 3 and β should be large in quadrants 2 and 4.

Original: The training data for this example consists of two trajectories.
Revised: The training data for this example consists of two trajectories.

Original: To generate it, we begin by drawing 100 samples each from isotropic normal distributions with standard deviation σ = 0 .
Revised: To generate it, we begin by drawing 100 samples each from isotropic normal distributions with standard deviation σ = 0 .

Original: 1 centred at ( − 1 , − 1) , (1 , 1) for the first trajectory and ( − 1 , 1) , (1 , − 1) for the second trajectory.
Revised: 1 centred at ( − 1 , − 1) , (1 , 1) for the first trajectory and ( − 1 , 1) , (1 , − 1) for the second trajectory.

Original: We randomly pair samples from subsequent distributions along each trajectory and compute A -geodesics between each pair by solving problem (9) using the method described in the circular data section.
Revised: We randomly pair samples from subsequent distributions along each trajectory and compute A -geodesics between each pair by solving problem (8) using the method described in the circular data section.

Original: We then draw 10 points at equispaced times t i ∈ [0 , 1] from each resulting geodesic and aggregate across geodesics to form the observed populations
Revised: We then draw 10 points at equispaced times t i ∈ [0 , 1] from each resulting geodesic and aggregate across geodesics to form the observed populations

Original: X t i .
Revised: X t i .

Original: We then use our method to recover ˆ A ( x ) from the X t i .
Revised: We then use our method to recover ˆ A ( x ) from the X t i .

Original: We parametrize the scalar potentials in (5) as a single-hidden-layer neural net with 32 hidden dimensions and Softplus activation.
Revised: We parametrize the scalar potentials in (4) as a single-hidden-layer neural net with 32 hidden dimensions and Softplus activation.

Original: We parametrize the matrix field
Revised: We parametrize the matrix field

Original: ˆ A ( x ) as A ( x ) =
Revised: ˆ A ( x ) as A ( x ) =

Original: Q ( x ) T Q ( x ) + 10 − 3 I , where Q ( x ) is a two-layer neural network with Softplus activations and 32 hidden dimensions.
Revised: Q ( x ) T Q ( x ) + 10 − 3 I , where Q ( x ) is a two-layer neural network with Softplus activations and 32 hidden dimensions.

Original: The strength of the gradient penalty (6) is 10 −when training the potentials ϕ
Revised: The strength of the gradient penalty (5) is 10 −when training the potentials ϕ

Original: in the first alternating step and 10 − 4 for subsequent steps.
Revised: in the first alternating step and 10 − 4 for subsequent steps.

Original: It is 1 when training ˆ A .
Revised: It is 1 when training ˆ A .

Original: The strength of the regularization (7) is 10 3 .
Revised: The strength of the regularization (6) is 10 3 .

Original: We carry out three steps of the alternating scheme by training with AdamW with learning rate 10 − 2 and weight decay 5 ∗ 10 − 3 .
Revised: We carry out three steps of the alternating scheme by training with AdamW with learning rate 10 − 2 and weight decay 5 ∗ 10 − 3 .

Original: We train for 300 epochs for the ϕ step and 40,000 epochs for the A step.
Revised: We train for 300 epochs for the ϕ step and 40,000 epochs for the A step.

Original: We evaluate the alignment score ℓ ( A, ˆ A ) on a 100 × 100 grid overlaid on a box of radius 1.5. 
Revised: We evaluate the alignment score ℓ ( A, ˆ A ) on a 100 × 100 grid overlaid on a box of radius 1.5. D

Original: B. SC RNA TRAJECTORY INFERENCE
Revised: SC RNA TRAJECTORY INFERENCE

Original: Data pre-processing.
Revised: Data pre-processing.

Original: The data for this experiment consists of force-directed layout embedding coordinates of gene expression data from Schiebinger et al.
Revised: The data for this experiment consists of force-directed layout embedding coordinates of gene expression data from Schiebinger et al.

Original: (2019) collected over 18 days of reprogramming (39 time points total) which we rescale by a factor of 10 − 3 for increased stability in training.
Revised: (2019) collected over 18 days of reprogramming (39 time points total) which we rescale by a factor of 10 − 3 for increased stability in training.

Original: We construct populations X t i for i = 1 , ..., 39 by randomly drawing 500 samples per time point in the original data; this corresponds to using 8.25% of the available data on average.
Revised: We construct populations X t i for i = 1 , ..., 39 by randomly drawing 500 samples per time point in the original data; this corresponds to using 8.25% of the available data on average.

Original: Learning the metric.
Revised: Learning the metric.

Original: We use our method to learn a metric ˆ
Revised: We use our method to learn a metric ˆ

Original: A − 1 ( x ) from the X t i .
Revised: A − 1 ( x ) from the X t i .

Original: We parametrize the scalar potentials in (5) as a single-hidden-layer neural net with 128 hidden dimensions and Softplus activation.
Revised: We parametrize the scalar potentials in (4) as a single-hidden-layer neural net with 128 hidden dimensions and Softplus activation.

Original: We parametrize the matrix field ˆ A − 1 ( x ) as A − 1 ( x ) =
Revised: We parametrize the matrix field ˆ A − 1 ( x ) as A − 1 ( x ) =

Original: Q ( x ) T Q ( x ) + 10 − 9 I , where
Revised: Q ( x ) T Q ( x ) + 10 − 9 I , where

Original: Q ( x ) is a single-hidden-layer neural network with Softplus activation and 2048 hidden dimensions.
Revised: Q ( x ) is a single-hidden-layer neural network with Softplus activation and 2048 hidden dimensions.

Original: We omit the gradient penalty (6) when training the potentials ϕ and set the penalty strength towhen training ˆ A .
Revised: We omit the gradient penalty (5) when training the potentials ϕ and set the penalty strength towhen training ˆ A .

Original: The strength of the regularization (7) is 5 ∗ 10 2 .
Revised: The strength of the regularization (6) is 5 ∗ 10 2 .

Original: We carry out a single step of the alternating scheme by training with AdamW with learning rate 5 ∗ 10 − 3 and weight decay 1 .
Revised: We carry out a single step of the alternating scheme by training with AdamW with learning rate 5 ∗ 10 − 3 and weight decay 1 .

Original: 5 ∗ 10 − 2 .
Revised: 5 ∗ 10 − 2 .

Original: We train for 100 epochs for the ϕ step and 5,000 epochs for the A step.
Revised: We train for 100 epochs for the ϕ step and 5,000 epochs for the A step.

Original: Trajectory inference task.
Revised: Trajectory inference task.

Original: We parametrize the time-varying velocity field v t,θ as a fully connected three-hidden-layer neural network with 64 hidden dimensions and Softplus activations.
Revised: We parametrize the time-varying velocity field v t,θ as a fully connected three-hidden-layer neural network with 64 hidden dimensions and Softplus activations.

Original: We follow
Revised: We follow

Original: Grathwohl et al.
Revised: Grathwohl et al.

Original: and concatenate the time variable to the input to each layer of the neural net.
Revised: and concatenate the time variable to the input to each layer of the neural net.

Original: We compute particle trajectories x ( t ) by solving the initial value problem ˙ x ( t ) = v θ ( x ( t )); x (0)
Revised: We compute particle trajectories x ( t ) by solving the initial value problem ˙ x ( t ) = v θ ( x ( t )); x (0)

Original: =
Revised: =

Original: xfor initial positions x
Revised: xfor initial positions x

Original: 0 ∈ X t i using the midpoint method solver in torchdiffeq’s odeint with default hyperparameters (Chen et al., 2018a).
Revised: 0 ∈ X t i using the midpoint method solver in torchdiffeq’s odeint with default hyperparameters (Chen et al., 2018a).

Original: The fitting loss for this task is GeomLoss’s Sinkhorn divergence with p = 2 and blur parameter fixed to 5 ∗ 10 − 2 .
Revised: The fitting loss for this task is GeomLoss’s Sinkhorn divergence with p = 2 and blur parameter fixed to 5 ∗ 10 − 2 .

Original: We fix λ = 10 − 1 for both our learned metric ˆ A and for the identity baseline
Revised: We fix λ = 10 − 1 for both our learned metric ˆ A and for the identity baseline

Original: A
Revised: A

Original: =
Revised: =

Original: I and choose the intermediate time points t j to be an equispaced sampling of [0 , 1] with step size 160 for all experiments.
Revised: I and choose the intermediate time points t j to be an equispaced sampling of [0 , 1] with step size 160 for all experiments.

Original: We solve problem (8) for each pair of subsequent time retained points ( t i , t i +1 ) .
Revised: We solve problem (7) for each pair of subsequent time retained points ( t i , t i +1 ) .

Original: In each case, we optimize the objective using AdamW with learning rate 10 − 3 and weight decay 10 − 3 and train for 10,000 iterations.
Revised: In each case, we optimize the objective using AdamW with learning rate 10 − 3 and weight decay 10 − 3 and train for 10,000 iterations.

Original: We evaluate the inferred trajectories by approximately computing the W 1 distance (using GeomLoss with p = 1 and blur of 10 − 6 ) between left-out time points in the ground truth data and advected samples at corresponding time points.
Revised: We evaluate the inferred trajectories by approximately computing the W 1 distance (using GeomLoss with p = 1 and blur of 10 − 6 ) between left-out time points in the ground truth data and advected samples at corresponding time points.

Original: For left-out time points of form t = ki + ℓ for some integer ℓ ∈ 1 , ..., k − 1 , the corresponding advected sample in the equispaced sampling of step size 160 of [0 , 1] has index ⌊ jk ∗ 60 ⌋ .
Revised: For left-out time points of form t = ki + ℓ for some integer ℓ ∈ 1 , ..., k − 1 , the corresponding advected sample in the equispaced sampling of step size 160 of [0 , 1] has index ⌊ jk ∗ 60 ⌋ .

Original: B. S NOW GOOSE TRAJECTORY INFERENCE
Revised: D. S NOW GOOSE TRAJECTORY INFERENCE

Original: Data pre-processing.
Revised: Data pre-processing.

Original: The training data for learning the metric in this experiment consists of time-stamped sightings of untagged snow geese ( Anser caerulescens ) across the U.S. and Canada during their spring migration.
Revised: The training data for learning the metric in this experiment consists of time-stamped sightings of untagged snow geese ( Anser caerulescens ) across the U.S. and Canada during their spring migration.

Original: This data is drawn from the February 2022 version of the eBird basic dataset (Sullivan et al., 2009).
Revised: This data is drawn from the February 2022 version of the eBird basic dataset (Sullivan et al., 2009).

Original: We bin the sightings by month of observation, keep 1000 records per month, and convert the sighting locations from latitude-longitude to xy coordinates using the
Revised: We bin the sightings by month of observation, keep 1000 records per month, and convert the sighting locations from latitude-longitude to xy coordinates using the

Original: Mercator projection implemented in Matplotlib’s Basemap module with rsphere set to 5 and the latitude/longitude of the lower left and upper right corners of the projection set to the minimum latitude and longitude and maximum latitude and longitude in the training data, respectively.
Revised: Mercator projection implemented in Matplotlib’s Basemap module with rsphere set to 5 and the latitude/longitude of the lower left and upper right corners of the projection set to the minimum latitude and longitude and maximum latitude and longitude in the training data, respectively.

Original: The ground truth goose trajectories are drawn from the “Banks Island SNGO” study on Movebank.
Revised: The ground truth goose trajectories are drawn from the “Banks Island SNGO” study on Movebank.

Original: This data consists of time-stamped GPS measurements of the locations of 8 snow geese fromto 2022.
Revised: This data consists of time-stamped GPS measurements of the locations of 8 snow geese fromto 2022.

Original: We use the same Mercator projection to convert the goose location measurements from latitude/longitude to xy coordinates.
Revised: We use the same Mercator projection to convert the goose location measurements from latitude/longitude to xy coordinates.

Original: We then estimate the initial and final time point of a single migration for each goose.
Revised: We then estimate the initial and final time point of a single migration for each goose.

Original: For geese with ID { 82901 , 82902 , 82905 , 82906 , 82907 , 82908 , 82909 , 82910 } , their respective initial time indices are { 20000 , 0 , 0 , 20000 , 0 , 0 , 0 , 0 } and their respective final time indices are { 26000 , 9100 , 15057 , 26500 , 9000 , 13037 , 7201 , 10000 } .
Revised: For geese with ID { 82901 , 82902 , 82905 , 82906 , 82907 , 82908 , 82909 , 82910 } , their respective initial time indices are { 20000 , 0 , 0 , 20000 , 0 , 0 , 0 , 0 } and their respective final time indices are { 26000 , 9100 , 15057 , 26500 , 9000 , 13037 , 7201 , 10000 } .

Original: The initial location of each goose is the initial condition x 0 in (9), and the final location of each goose is x 1 .
Revised: The initial location of each goose is the initial condition x 0 in (8), and the final location of each goose is x 1 .

Original: Learning the metric.
Revised: Learning the metric.

Original: We use our method to learn a metric ˆ
Revised: We use our method to learn a metric ˆ

Original: A − 1 ( x ) from the X t i .
Revised: A − 1 ( x ) from the X t i .

Original: We parametrize the scalar potentials in (5) as a single-hidden-layer neural net with 32 hidden dimensions and a Softplus activation.
Revised: We parametrize the scalar potentials in (4) as a single-hidden-layer neural net with 32 hidden dimensions and a Softplus activation.

Original: We parametrize the matrix field ˆ A − 1 ( x ) as A − 1 ( x ) =
Revised: We parametrize the matrix field ˆ A − 1 ( x ) as A − 1 ( x ) =

Original: Q ( x ) T Q ( x ) + 10 − 3 I , where Q ( x ) is a two-layer neural network with Softplus activations and 32 hidden dimensions.
Revised: Q ( x ) T Q ( x ) + 10 − 3 I , where Q ( x ) is a two-layer neural network with Softplus activations and 32 hidden dimensions.

Original: The strength of the gradient penalty (6) is 10 − 6 when training the potentials ϕ and 1 when training ˆ A − 1 .
Revised: The strength of the gradient penalty (5) is 10 − 6 when training the potentials ϕ and 1 when training ˆ A − 1 .

Original: The strength of the regularization (7) is 10 9 .
Revised: The strength of the regularization (6) is 10 9 .

Original: We carry out a single a step of the alternating scheme by training with
Revised: We carry out a single a step of the alternating scheme by training with

Original: AdamW with learning rate 10 − 2 and weight decay 10 − 3 .
Revised: AdamW with learning rate 10 − 2 and weight decay 10 − 3 .

Original: We train for 2,000 epochs for the ϕ step and 10,000 epochs for the A step.
Revised: We train for 2,000 epochs for the ϕ step and 10,000 epochs for the A step.

Original: Trajectory inference task.
Revised: Trajectory inference task.

Original: We parametrize the time-varying velocity field v t,θ as a fully connected three-hidden-layer neural network with 64 hidden dimensions and ELU nonlinearities.
Revised: We parametrize the time-varying velocity field v t,θ as a fully connected three-hidden-layer neural network with 64 hidden dimensions and ELU nonlinearities.

Original: We compute particle trajectories x ( t ) by solving the initial value problem ˙ x ( t ) = v θ ( x ( t )); x (0) =
Revised: We compute particle trajectories x ( t ) by solving the initial value problem ˙ x ( t ) = v θ ( x ( t )); x (0) =

Original: x 0 for initial positions x 0 ∈ X t i using the dopri5 solver in torchdiffeq’s odeint with default hyperparameters (Chen et al., 2018a).
Revised: x 0 for initial positions x 0 ∈ X t i using the dopri5 solver in torchdiffeq’s odeint with default hyperparameters (Chen et al., 2018a).

Original: We fix λ = 1 .
Revised: We fix λ = 1 .

Original: 25 ∗ 10 2 for geese 82901, 82902, 82906, 82908, and 82909 and use λ = 2 .
Revised: 25 ∗ 10 2 for geese 82901, 82902, 82906, 82908, and 82909 and use λ = 2 .

Original: 5 ∗ 10 2 for the remaining geese 82905, 82907, 82910.
Revised: 5 ∗ 10 2 for the remaining geese 82905, 82907, 82910.

Original: In each case, we optimize the objective using AdamW with learning rate 10 − 3 and no weight decay, set the times t j in (9) be an equispaced sampling of [0 , 1] with 32 time points and train for 500 iterations.
Revised: In each case, we optimize the objective using AdamW with learning rate 10 − 3 and no weight decay, set the times t j in (8) be an equispaced sampling of [0 , 1] with 32 time points and train for 500 iterations.

Original: C A BLATION STUDY : I MPACT OF THE REGULARIZATION COEFFICIENT
Revised: E A BLATION STUDY : I MPACT OF THE REGULARIZATION COEFFICIENT

Original: In this section, we carry out an ablation study of the impact of the coefficient λ on the regularization term R ( A ) in (5).
Revised: In this section, we carry out an ablation study of the impact of the coefficient λ on the regularization term R ( A ) in (4).

Original: We repeat the metric recovery experiment on the “Circular” example.
Revised: We repeat the metric recovery experiment on the “Circular” example.

Original: We follow the procedure in Appendix B.1 but set λ ∈ { 5 ∗ 10 2 , 10 3 , 5 ∗ 10 3 , 10 4 , 5 ∗ 10 4 , 10 5 } .
Revised: We follow the procedure in Appendix D.1 but set λ ∈ { 5 ∗ 10 2 , 10 3 , 5 ∗ 10 3 , 10 4 , 5 ∗ 10 4 , 10 5 } .

Original: (Note that our reported results in Section 5.1 use λ = 10 3 .)
Revised: (Note that our reported results in Section 4.1 use λ = 10 3 .)

Original: The results of this experiment are recorded in Figure 5.
Revised: The results of this experiment are recorded in Figure 6.

Original: The eigenvectors of the learned metric ˆ A ( x ) are robust to the value of λ except at the largest value of λ = 10 5 , where the alignment score falls to0.857.
Revised: The eigenvectors of the learned metric ˆ A ( x ) are robust to the value of λ except at the largest value of λ = 10 5 , where the alignment score falls to

Original: As expected, the log-condition number of ˆ A ( x ) increases somewhat with λ , indicating that larger values of λ favor high levels of anisotropy.
Revised: As expected, the log-condition number of ˆ A ( x ) increases somewhat with λ , indicating that larger values of λ favor high levels of anisotropy.

Original: D A BLATION STUDY : R EMOVING THE OT TERM
Revised: F A BLATION STUDY : R EMOVING THE OT TERM

Original: In this section, we briefly demonstrate that if we remove the optimal transport term from (5), the resulting metric is not informed by the training data.
Revised: In this section, we briefly demonstrate that if we remove the optimal transport term from (4), the resulting metric is not informed by the training data.

Original: We repeat the metric recovery experiment on the “Circular” example.
Revised: We repeat the metric recovery experiment on the “Circular” example.

Original: We follow the procedure in Appendix B.1 but exclude the OT term
Revised: We follow the procedure in Appendix D.1 but exclude the OT term

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: from our implementation of (5).
Revised: from our implementation of (4).

Original: The results of this experiment are presented in Figure 6.
Revised: The results of this experiment are presented in Figure 7.

Original: Note that the learned metric is no longer informed by the data; the low-energy eigenvectors point in the direction of the vertical axis everywhere and the log-condition number is simply an increasing function of distance from the origin.
Revised: Note that the learned metric is no longer informed by the data; the low-energy eigenvectors point in the direction of the vertical axis everywhere and the log-condition number is simply an increasing function of distance from the origin.