Riemannian Metric Learning via Optimal Transport
Authors: Christopher Scarvelis
Justin Solomon
Venue: ICLR 2023 poster (2023)
Keywords: optimal transport
riemannian geometry
manifold learning
time series

Abstract: 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. We neurally parametrize the metric as a spatially-varying matrix field and efficiently optimize our model's objective using a simple alternating scheme. Using this learned metric, we can non-linearly interpolate between probability measures and compute geodesics on the manifold. 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.

Published as a conference paper at ICLR 2023
RIEMANNIAN METRIC LEARNING VIA OPTIMAL TRANSPORT
Christopher Scarvelis MIT CSAIL scarvêmit edu
Justin Solomon MIT CSAIL j jsolomon edu
ABSTRACT
We introduce optimal transport-basee model for learning a metric tensor from cross-sectional samples of evolving probability measures common Riemannian manifold. We neurally parametrize the metric as a spatially efficiently optimize our model's objective using simple alternating scheme. Using this learned metric, we can nonlinearly interpolate between probability measures and compute geodesics on the manifold. We show that metrics learned using our method improve the quality of trajectory inference on scRNA and bird migration data at the cost little additional cross-sectional data.
matrix field and
INTRODUCTION
In settings like single-cell RNA sequencing (scRNA-seq) (Tanay & Regev, 2017), we often encounter pooled cross- poss--eectioal data: Time-indexed samples { Nr :=1 from an evolving population X, with no correspondence between samples xy and a; at times & * t. Such data may arise when technical constraints impede the repeated observation of some population member. 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.
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,] which samples are drawn. A fruitful approach has been to model the evolving population time-varying probability measure P: on TR. and infer the distribution of the population between observed times by interpolating between subsequent pairs of measures (tt, Pt.+1). Some existing approaches to this problem use dynamical optimal transport to interpolate between probability measures, which implicitly encodes prior that particles travel along straight lines between observations. This prior is often implausible, especially when the evolving population is sampled sparsely in time.
One can straightforwardly extend optimal transport- based methods by allowing the user to specify apatially-varying metric tensor to bias the inferred trajectories away from straight lines. This approach theoretically well-founded and amounts to redefining straight line by altering the manifold on which trajectories are inferred. Such a metric tensor, however, is typically unavailable in most real-world applications.
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. 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.
We demonstrate our algorithm's 's ability to recover a known metric tensor from cross-sectional samples on synthetic examples. 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 cross- sectional samples of migrating population. Our method is both computationally-effie 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.
requiring
RELATED WORK
Measure interpolation. An emerging literature considers the problem of smooth interpolation between probability measures. Using the theory of optimal transport, we may construct displace- ment interpolation (McCann, 1997) between successive measures (P;, P.t1), yielding a sequence

---

Published as a conference paper at ICLR 2023
of geodesies between pairs (P,, P.+1) in the space of probability measures equipped with the Wasserstein distance W2 (Villani, 2008). This generalizes piecewise-lincau interpolation to probability measures. Schiebinger et al. (2019) use this method to infer the developmental trajectories of cells based static measurements of their gene expression profiles. 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-regulariied analog to the displacement interpolation between measures.
Chen et al. 8b); Benamou et al. (2018) leverage the variational characterization of cubic splines as minimizers of mean-square acceleration over the set of interpolating curves and develop a general- ization in the space of probability measures. Chewi et al. (2021) extend these works by providing computationally efficient algorithms for computing measure-valued splines.
Hug et al. (2015) modify the usual displacement interpolation between probability measures by introducing anisotropy to the domain on which the measures are defined. This change corresponds to imposing preferred directions for the local displacement of probability mass. Whereas Hug et al. hard-code the domain's 's anisotropy. our method allows us learn this anisotropy from snapshots of probability measure evolving in time. Zhang et al. (2022) apply similar techniques to unstructured animation problems.
Ding et al. (2020) propose a non-convex inverse problem for recovering the ground metric and interaction kernel class mean-field games and supply a primal-dual algorithm for solving a grid discretization of the problem. Whereas their Bulerian approach inputs grid discretization of observed densities and velocity fields, our approach is Lagrangian, operating directly on temporal observations of particle positions.
Trajectory inference from population-level data, In domains such as scRNA-seq, we study an evolving population from which it is impossible or prohibitively costly collect longitudinal data. Instead, we observe distinct cross-sectional samples from the population at a collection of times and wish infer the dynamics of the latent population from these observations. This problem is called trajectory inference.
Hashimoto al. (2016) study the conditions under which it is possible recover potential function from population-level observations of a system evolving according to a Fokker-Planck equation. 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. A recent work of Bunne et al. (2022) presents proximal analog to this approach, modeling population dynamics as a JKO flow with learnable energy function and describing numerical scheme for computing this flow using input-convex neural networks. This method addresses initial value problems where one is given an initial measure Po and seeks to predict future measures Pr for / >> O. In contrast, our method is primarily applicable to boundary value problems, where we are given initial and final measures Po and P1 and seek interpolant Pt. for 0 < t <
Schiebinger et al. (2019) use optimal transport to infer future and prior gene expression profiles from a single observation in time of cell's gene expression profile. Yang & Uhler (2019) propose 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.
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. Tong al. (2020) show that a regularized continuous normalizing flow can provide an efficient approximation to displacement interpolations arising from dynamical optimal transport. These displacement interpolations often do not yield plausible paths through gene expression space, so the authors propose several application- specific regularizers to bias the inferred trajectories toward more realistic paths. Rather than relying on bespoke regularizers, our method supplies a natural approach for learning local directions of particle motion way that is amenable to integration with algorithms like that of Tong et al. (2020).
Riemannian metric learning. Lebanon (2002) develops a parametric method for learning Rieman- nian metric from sparse high-dimensional data and uses this metric for nearest-neighhor classification. Hauberg et al. (2012) construct a metric tensor tensors and show how to compute geodesics and exponential and logarithmic maps on the resulting Riemannian manifold. Arvanitidis et al. (2016) learn a Riemannian metric that encourages geodesics to move towards region of high data density, define a Riemannian normal distribution with respect to
weighted average of a set of learned metric

---

Published as a conference paper at ICLR 2023
this metric, and show how to learn the parameters of this model via maximum likelihood estimation. Whereas these methods learn Riemannian metric from static data, model Iearns metric from snapshots of probability distributions evolving in time on a common manifold.
METHOD
We now describe our method for learning a Riemannian metric from cross-sectional samples of populations evolving in time on a common manifold. We learn metric that minimizes the average 1- Wasserstein distance on the manifold between pairs of subsequent time samples from each population. We derive dual formulation of our problem, parametrize its variables by neural networks, and solve for the dual variables and the metric via alternating optimization.
3.1 MODEL
Suppose we have K populations evolving according unknown continuous dynamies over a common Riemannian manifold .M = (IR g) with unknown metric g. The metric g is defined at any x E R by the inner product (u, v)., u² (( for A(x) > O. We model each population as a compactly- supported probability distribution pe with density p on I2 being pushed through an unobserved velocity field vZ (z). We will learn the metric tensor A(x) from temporal snapshots of the populations pk during their evolution on the manifold.
Depending the nature of our data, we may have the ability to repeatedly sample from pk at a pair of initial and final times t 0 and t Tk or have a fixed set of samples {a {= and (a Ti=1 drawn from the populations ph at their respective times. For convenience, we denote both the density p² and the empirical distributions over samples from p at times t E {0,Th) by p and of respectively. As we only observe the initial and final spatial distributions p and p of the populations and do not observe their dynamics, we assume that probability mass travels from initial to final positions along A-geodesics. Geodesics are paths that minimize the action (or average kinetic energy) of a particle traveling between points 2, y E M: this least-action interpretation of a geodesic makes a natural prior on paths the absence of further information. We learn field of positive definite matrices A(z) that minimizes the average A-geodesic distance between the initial and final positions of each unit of probability mass from each pk
Formally, let p be the map sending point E .M to its final position r (zz) after flowing through the latent velocity field u from time f 0 to t = Tg. If we had access to such a solution map, we would ideally solve the following problem:
K
1
E da (z, dob(() AR(A), k=
inf ARR-SSD
K
where R(A) is regularizer that excludes the trivial solution A = 0. Problem (1) optimizes for ton-trivial metric A(z) that minimizes the average A-geodesic distance d (z, r (x)) traveled by particles æ in the population o at time t 0 to their final positions , r²(z) at time = Té.
Since we do not know the velocity fields wi that encode the particle dynamics, however, we also do not know the maps y in (1) that specify the correspondence between particle positions at l = 0 and positions r* (x) at final times Tg. Furthermore, as noted in Section I, we often encounter data for which it is impossible to observe contiguous particle trajectories: A particle that we observe f 0 may not be the sample at f = Tg. This issue is unavoidable for destructive measurement processes such scRNA sequencing; in this setting, cell whose ScRNA profile is observed & would be destroyed at this time and hence unobservable at a future time t = Te- To accommodate these limitations, we replace the true matchings of initial and final positions b with the Monge map s defined as the solution to the following problem:
Ju da (x, (x))
WI.A P
inf stopp=oly
(x))) (o)
(2)
Here & is a pushforward of p onto p²: we write of = Pi to denote this relationship. This map matches units of mass from the initial and final distributions pā ot to minimize their average A A-geodesic distance. Substituting solutions to (2) for the maps nK in our idealized objective (1), we

---

Published as a conference paper at ICLR 2023
obtain the following lower bound on (I):
inf A:R +82 8 shipp
inf
1 K
Σ dA (z, (x)) dp((a) (A).
=s
EE= Wi.A(P--P))
Problem (3) is challenging written: It requires the ability to compute and differentiate geodesic distances d with respect to an arbitrary metric. However, the inner optimization problem over maps 5 is a collection of decoupled Monge problems (2) whose optimal value is the average 1-Wasserstein distance between (pC: p²) pairs on the manifold with metric A(x). As shown in Appendix B, these problems can be expressed in a dual form (10) which is amenable to gradient-based optimization.
Replacing the inner Monge problems with their dual formulations (10), we may equivalently write Problem (3) as a minimax problem:
(J ('(()dp((()) - -J, (w))do((a)) + (R()..
inf A: S
sup M-R
(4)
IV (1)ll
In Problem (4), learn metric A(x) that minimizes the average I-Wasserstein distance on the manifold between pairs of subsequent time samples from each population. (4) requires neither the computation of geodesic distances on M nor the solution an assignment problem. such, substantially more tractable than the initial objective (3). We provide the details of our implementation of (4) Section 3.2.
3.2 IMPLEMENTATION
Enforcing the Lipschitz constraint. Problem (4) includes global constraints of form V ll S 1. These constraints are the Riemannian analog to the Lipschitz constraint dual formulation 1-Wasserstein distance on RR, Constraints of this type are challenging enforce in gradient-based optimization, and the Wasserstein GAN literature has explored approximations (Arjovsky et al., 2017; Gulrajani et al., 2017; Miyato et al., 2018). We follow the standard technique introduced by Guirajani et al. (2017) and replace the global constraints (/ IlA S with soft penalties of the following form:
SoftPlus (IV (0% (a-10) -
E samPo mpp t~ (0.1)
where o (t) t)xo + tz is a line segment between zo and æ1 parametrized by t [0, 1]. Intuitively, (5) penalizes violation of the Lipschitz constraint ||V || < along line segments connecting randomly-paired points in X and X/. Gulrajani et al. (2017, Prop. 1) justify this choice via a standard result in W1 theory showing that the constraint | < 1 binds on line segments connecting pairs of points that are matched by the Monge map sh. Korotin et al. (2022) show that this method results accurate approximations to the directions of the gradients V of the true Kantorovich potentials. This observation is sufficient for our purposes; noted below, optimization scheme encourages the low-energy eigenvectors of A(x) to be well-aligned with the solutions & to the inner problem in (4).
Choice of regularizer R(A), Without regularizer R(A), the objective in (4) can be driven to 0 by choosing A(x) =\t af for arbitrarily small a 0. thereby making all pairs of measures arbitrarily close. This trivial solution would incorporate no useful information from the observed samples from p and p as it is simply a rescaling the standard Euclidean metric.
We opt for the following regularizer in our method:
K
1
R(A) =
[|A- (z) [2.]
(6)
K
P 2i t~U((011)
where Td (t) is as in the previous paragraph and || || denotes the squared Frobenius norm.

---

Published as a conference paper at ICLR 2023
Regularizer (6) penalizes || at points drawn using a sampling scheme analogous to that in (5). This is natural way exclude trivial solutions of form A(x) al for arbitrarily small a > 0, as ||3 is large for such metrics. We investigate the impact of the regularization coefficient À on the learned metric in Appendix H. In Appendix F, we demonstrate the effect of removing the optimal transport term from (4), which is equivalent to setting À = +00.
Optimization scheme. We parametrize the scalar potentials o² 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 A-1, we directly parametrize the matrix field A- by neural network; where require the evaluation of A( (z) in downstream applications, we evaluate and then invert (z). We enforce the positive definiteness of A (and hence the positive definiteness of A) by parametrizing it as A (x) = (x)) + n/ for a matrix-valued function Q(z) .M- - RR D'D and > 0.
After parametrizing our problem variables with neural networks, optimize the objective in (4) via alternation. In the first phase of our scheme. we solve the inner problem by holding A fixed (initializing as A(x) A (z) = I) and solving for the optimal Ø² This step decouples over the potentials of and each of the resulting problems is an instance of the dual problem for the
1- I-Wasserstein distance M. We approximate the integrals in (4) sample means N Σ o² (z.). where the { samples from the distributions p and p We likewise approximate the Lipschitz penalty (5) and the regularizer (6) as sample means over draws from P and o and over t drawn uniformly from [0, 1]. In the second phase, we solve the outer problem by fixing the optimal Φ² from the previous step and solving for the optimal matrix field A(x). We optimize both problems using AdamW (Loshchilov & Hutter, 2019). Few alternations are needed in practice to obtain high-quality results.
The results in Appendix show that given a fixed metric defined by A(x), the optimal Vask from the first phase of our scheme point along A-geodesics joining pairs of points (ar, (x)) where x P and s solves (2). At initialization, A I, and these geodesics are line segments in RR connecting the matched points.
Given fixed Kantorovich potentials o² from the first phase, the second phase of our scheme solves for matrix field A(x) that minimizes the regularizer R(A) ||? while also making || V large. Intuitively, this encourages the unit eigenvector d1 (z) corresponding to the minimal eigenvalue A (x) of A(z) to be aligned with Vo (z) wherever the constraint is enforced.
4 EXPERIMENTS
In this section, we first use synthetic data to demonstrate that our algorithm successfully recovers the correct cigenspaces of known metric A (z) from cross-sectional samples satisfying our model. We then method 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. We finally show that by learning metric from time-stamped bird sightings, we can infer curved migratory trajectories for individual birds given the initial and final points of their trajectories. Details for all experiments are provided in Appendix D.
4.1 METRIC RECOVERY
We first show that our method recovers the correct eigenspaces of a known metric A(x) from cross-sectional samples XI and X satisfying the model in Section 3.1. We fix initial positions ",,(0 (0) and final positions (1) for a set of N particles v,, E X and also fix spatially-vary metric A(x). We compute A A-geodesics between each (z,(0), xr (1)) pair and define populations T.
X = [x,((t)) ((:::
I, ..., N} for times t; E [0, 1],
We use our method to learn the latent metric A(x:) from the pairs of samples (X*;,X++)) For each example, we plot the eigenvectors of the true metric A(x) and those the learned metric Á(x) on a P x P grid (first row, Figure 1). We also plot the eigenvectors of Á(x) along with the log-condition number log (>a())/A(())) of A(x) (second row, Figure 1); this value is large when the learned metric highly anisotropic. We finally report measure of the alignment of the eigenspaces of the true
metric and its learned counterpart: E(A, A) = D ΣLEX I (ua()), (x))). Here ua(x)) is the unit eigenvector with eigenvalue da of A(x), (x))) is the corresponding eigenvector for A(æ),

---

Published as a conference paper at ICLR 2023
and X' a set of grid points which we plot the eigenvectors of A and A. E(A, A) = when each
eigenvector of A(r) is parallel to the corresponding eigenvector of A(z) at all grid points x: E A.
Figure 1: Row I: Eigenvectors of true metric A (ar) (purple) and learned metric A(x) (orange). Row 2: Log-condition number of learned metric A(x) yellow indicates highly anisotropic A(x). The points are time samples from which A(z) is recovered. Teal points indicate earlier times and yellow points indicate later times, so the pairs of temporal samples (X*: X follow the color gradient. Our method accurately recovers the eigenvectors of the true metric, and the learned metric is highly anisotropic in regions that overlap the observed data.
Our method accurately recovers the cigenspaces of the "Circular" (C(A, = 0.995) and "X Paths" (E(A, A) = 0.916) metrics. It achieves lower accuracy with the "Mass Splitting" metric (E(A, A) 0.839), struggling to capture the discontinuity in its eigenvectors the æ-axis and exhibiting numerical instability to the left of the y-axis, where the training trajectories diverge. 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. 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.
Row 2 of Figure 1 shows that our learned metric is highly anisotropic in regions that overlap with the training data. 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.
4.2 POPULATION-LEVEL TRAJECTORY INFERENCE WITH A LEARNED METRIC
scRNA sequencing (scRNA-seq) allows biologists to study the set of mRNA molecules (the "tran- scriptome") of individual cells at high resolution (Haque et al., 2017). As scRNA-seq a destructive process, any individual cell's RNA can be sequenced only once, impeding the use of this technology to study the change in cell's 's transcriptome over time. This has led interest in methods that use population-level scRNA-sep 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. Schiebinger et (2019) and Tong et al. (2020) both use optimal transport to infer cellular development trajectories but assume a Euclidean metric on gene expression space. In this section, we use our method to relax this strong assumption by learning a metric for the ground space from scRNA data. We incorporate A(x) into a downstream trajectory inference task and show that this strategy yields more accurate trajectories than a baseline without learned metric and bascline using the Euclidean metric.

---

Published as a conference paper at ICLR 2023
ScRNA data. We perform trajectory inference experiments with the scRNA data drawn from Schiebinger et al. (2019). This data consists of force-directed layout embedding coordinates of gene expression data collected over 18 days of reprogramming (39 time points total). We construct populations X for i 1, ....., 39 by drawing 500 samples per time point in the original data; this sampling uses 8.25% of the available data on average. We follow the same procedure as in Section 4.1 to learn a metric tensor A(x) from subsequent pairs of samples (X", X(+1). Learning the tensor takes 16 minutes on a single V100 GPU.
For the downstream trajectory inference task, we keep one out of every k time points in the original data for k 2, .... 19 obtain new collection of time points & with & {nk n E N, < 22 We then perform trajectory inference between subsequent retained time points (te, te42) (using all of the available data for these time points) by optimizing the following objective:
F
Se (xtens, (xtest, (0=1(1)) + A 0
ΣΣ>>
[e.,((((t))) (((OO)))))
(7)
ZEX
Here S. is the Sinkhorn divergence (Feydy et al., 2019) between the target data X and the advected samples dz (1o [X Φz (vD) denotes the solution map resulting from advecting a particle through ceurally-parametrre
vire-ng velocity field Ur,9 for one unit of time; and Ej E [0, 1J. The Sinkhorn divergence used here is Sa 8)
We Wa B) WWa ox) (3, it is a debiased variant of the Wasserstein distance W..
aatropieally-reullrian
This is similar to the method of Tong et al. (2020), where we replace the log-likelihool fitting loss with the Sinkhorn divergence; we found that the Sinkhorn divergence led to stabler training than log-likelihood fitting loss. We use GeomLoss (Feydy et al., 2019) to compute the Sinkhorn divergence between the target and advected samples efficiently.
Figure 2: Mean EMD between left-out samples and corresponding ad- vected samples versus k with our learned metric (orange), the Euclidean metric A(x) = I (green), no regularizer (blue), and interpolants obtained us- ing Schiebinger et al. (2019)'s 's method (red).
We follow Tong et al. (2020) and assess the quality of inferred tra- jectories by measuring the W1 distance (EMD) between left-out time points in the ground truth data and advected samples at correspond- ing times in the inferred trajectories. These distances are of the form Wi (x*: Oizz,, (v=) (X so we compare ground truth samples at each left-out time t; nk + h (h L, ..., k 1) with samples from X advected through ey for time f = '.' average EMD between left out samples and their inferred counterparts and compare our method to a baseline where X = O in (7) ("Without A" in Figure 2) and a Euclidean baseline where A(x)= =/ I A"AA- = Id" Figure 2). Our "Without A" and "A = Id" baselines are comparable to the "Base" and "Base E" models, respectively, from Tong et al. (2020). Both our bascline models and their models learn a velocity field vg.0 that pushes samples at some time & onto samples at future time t241 and use the path followed by these samples as they flow through ve,0 as an inferred trajectory. Our "Without A" baseline and their "Base" model do not regularize vr,0, whereas our "A = Id" baseline and their "Base E" model regularize v4,0 by penalizing its squared norm, which encourages samples to flow along straight paths.
h For each value of k. we record the
We also compare against interpolants obtained via the method of Schiebinger et al. (2019). As their Waddington OT method solves 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. We compute static optimal transport couplings between data at subsequent retained time points (fe, (4+1), linearly interpolate between coupled data points, and measure the W1 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 "Without A" baseline (analogous to the "Base" model of Tong et al. (2020)) and the approach based on Schiebinger et al. (2019) for nearly all values of k. whereas the Euclidean baseline A(x) = (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,
7

---

Published as a conference paper at ICLR 2023
Figure Comparison of inferred trajectories for k 3 (first row) and k 15 (second row). 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.
using our learned metric results in 19.7% mean reduction in EMD between left-out samples and the corresponding advected samples for l 2 10. This observation indicates that including learned metric has a larger impact on the inferred trajectories when the ground truth data is sampled sparsely in time.
Figure 3 compares the trajectories inferred using our method to the non-regularized baseline and ground truth trajectories for A: = 3, where the time sampling is sufficiently dense that the baseline performs well, and & = 15, where our learned metric tensor substantially improves the quality of the inferred trajectories. 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.
These results suggest that in settings where data collection expensive and samples collected at a small subset of times 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.
4.3 INDIVIDUAL-LEVEL TRAJECTORY INFERENCE WITH A LEARNED METRIC
In this section, learn metric A(x) from time-stamped sightings of snow geese during their spring migration. 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.
Migratory paths as geodesics a latent manifold. 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). Given this behavior, a boundary value problem that fixes the endpoints of the birds' migratory trajectories a reasonable model for bird migration. Once the endpoints are fixed, must specify objective that birds plausibly optimize to determine their trajectories. Absent further information, minimizing total kinetic energy over their trajectory is reasonable choice. However, measuring kinetic energy with respect the Euclidean metric leads to straight-line paths, which typically do not agree with observed migration trajectories. We use our method learn metric from untagged snow goose sightings that agrees with the data for this species. 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.
Snow goose data. 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. We treat the sightings at time / as samples from time-indexed spatial distribution of birds p² and train our metric on subsequent pairs of bird distributions (p', pp)) This implies that our method does not have access to any complete goose trajectories when learning the metric. We do not expect untagged goose sightings to contain enough information to predict high-frequee freeuerreyncyy 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. The widespread availability and

---

Published as a conference paper at ICLR 2023
Ground truth
Ours
Baseline
Figure 4: By using a metric A(x) Iearned from time tine-stamped bird sightings, we obtain inferred trajectories (center) that capture the curved structure of the ground truth migratory paths (left). Our method results in a 26.9% reduction mean DTW distance between the inferred and ground truth trajectories relative to the Euclidean baseline (right).
low cost of obtaining untagged bird sighting data motivates the use of our method for bird trajectory inference
The training data is drawn from the eBird basic dataset (Sullivan al., 2009), current of February 2022. We bin the sightings by month of observation and use our algorithm to learn a metric tensor A(x) from populations X consisting of the spatial coordinates of snow goose sightings in month i for 0, .... (i.e. January to June). This training data is depicted in Figure 5 in Appendix D.3.
We use this learned metric tensor compute an Á-geodesic between the initial and final observation of several snow geese along their migratory paths. This data is drawn from the Banks Island Snow Goose study as hosted on Movebank (Kays et al., 2021). We estimate an A-geodesic between initial and final points on each path by learning time-varying velocity field vr optimizing the following objective:
min |2| x11]l2 +
Σ ((e,o((((t))) ((
(8)
As in Section 4.2, the velocity field Ur is neurally parametrized, 0 to S t..t. < .m 1, and we optimize (8) using AdamW. We set the initial condition x00) = zo such that 2o the goose's initial position on its migratory trajectory and use Euclidean norm penalty to force its final position x(1) along the inferred trajectory to match the true final position Z1.
Figure 4 compares the Å-geodesics obtained using the method described above (center plot) to the ground truth goose trajectories (left plot). While population-level data (in the form of untagged snow goose sightings; see Figure 5 in Appendix D.3) 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.
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. We also generate straight-lin -ine 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. Our method results 26,9 26.9% reduction in mean DTW distance between the inferred and ground truth trajectories relative to the Euclidean baseline.
5 DISCUSSION AND CONCLUSION
We have introduced an optimal transport- based method for learning Riemannian metric from cross-sectional samples of populations evolving over time on a latent Riemannian manifold. Our method accurately recovers metrics from cross-sections of populations moving along geodesics on 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.
One key limitation our work is that learns a Riemannian metric on RDD whereas large swaths of data naturally lie on graphs. Future work might consider extending this algorithm learn from data defined on graphs; a recent work of Heitz et al. (2021) presents an optimal transport-basee direction towards solving this problem. Future work may also extend our methods to learn metric from temporal snapshots of Schrödinger bridges.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
The MIT Geometric Data Processing group acknowledges the generous support of Army Research Office grants W91 INF2010168 and W91 INF21 10293, of Air Force Office of Scientific Research award 149550-19-1-033 of National Science Foundation grants IIS- 118-1838071 and CHS-1955697, from the CSAIL Systems that Learn program, from the MIT-IBM Watson Al Laboratory, from the Moyota-CSAIL.A Joint Research Center, from gift from Adobe Systems, and from Google Research Scholar award.
CS acknowledges the support of the Natural Sciences and Engineering Research Council of Canada (NSERC), funding reference number SLDD3-557558-200 and from the 2022 Siebel Scholars award.
REFERENCES
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In Proceedings of the 34th International Conference Machine Learning Volume 70, (CML'17, pp. 214-223. JMLR.org, 2017.
Georgios Arvanitidis, Lars K. Hansen, and Sø ren Hauberg. A locally adaptive nor- mal distribution. In D. Lee, M. Sugiyama, U. Luxburg. I. Guyon, and R. Gamett (eds.), Advances in Neural Information Processing Systems, volume 29, Curran Asso- ciates, Inc., 2016. URL https::/pproceeednggs neurips.co. wwwuripsscccppppr//2220 6/000/6/FF0 file// pdf.
Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution to the Monge- Kantorovich mass transfer problem. Numerische Mathematik, 84:375-393, 01 2000. doi: 10.1007/ 1002110050002
Jean-David Benamou, Thomas Gallouët, and François-Xavier Vialard. Second-order models for optimal transport and cubic splines on the Wasserstein space. Foundations of Computational Mathematics, 19, 01 2018. doi: 10.
Donald J. Berndt and James Clifford. Using dynamic time warping find patterns in time series. In Proceedings of the 3rd International Conference on Knowledge Discovery and Data Mining, AAAIWS' 94, pp. 259-370. AAAI Press, 1994.
Charlotte Bunne, Lactitia Meng-Papaxxanthos Andreas Krause, and Marco Cuturi. JKOnet: Proximal optimal transport modeling of population dynamics. AISTATS, 151, 2022.
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Resa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Asso- ciates, Inc., 2018a. URL ht http:::://ppooceeiiiii newrips. /2tpappr/211/////// file/ 6938 (3 Paper paf.r.pdf.
Yongxin Chen, Giovanni Conforti, and Tryphon Georgiou. Measure-valued spline curves: optimal transport viewpoint. SIAM Journal on Mathematical Analysis, 50, 01 2018b. doi: 10.113//8MM166249
Sinho Chewi, Julien Clancy, Thibaut Le Gouic, Philippe Rigollet, George Stepaniants, and Austin Stromme. Fast and smooth interpolation wasserstein space. In AISTATS, pp. 3061-3069, 2021. URL http: http:///ooeeeeeeeeee eigsediingg.... mll. wttress//v3//chaae
html.
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications score-based generative modeling. In M. Ran- zato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 17695-17709, Curran Asso- ciates, Inc., 2021. URL htttps:/pproceedinn sewrdiinnnnnenuisssseee.e.u cttpapp.ccppppp/eeee// cc/papprr/22///i1ee 94432255r22aaaeeec22111112
-Paper.pd..
Lisang Ding, Wuchen Li, Stanley Osher, and Wotao Yin. mean field game inverse problem, 2020.

---

Published as a conference paper at ICLR 2023
Mikhail Feldman and Robert McCann. Monge's transport problem on Riemannian manifold. Transactions of the American Mathematical Society, 354, 01 2002. doi: 10.2307/2693835.
Jean Feydy, Thibault Séjourné, François-Xavin Vialard, Shun ichi Amari, Alain Trouvé, and Gabriel Peyré. Interpolating between optimal transport and mmd using sinkhorn divergences. AISTATS, pp. 2681-2690, 2019. URL http: proccedooeeeoooooooooo ngs m ...esssss.... vv8 9 feydyl 9a html.
Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, and David Duvenaud. Scalable reversible generative models with free-form continuous dynamics. In International Conference on Learning Representations, 2019. URL https: http//ppeeeevieeeeeeeeeeeeeeeeee net/ forum wtt.:///ooormmmmeooooooooco wwd=:
Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron Courville. Improved training of Wasserstein GANs. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https:/ pttpppeeoceenn..c//ee/ee neurips. ccc/paper/211 /teeeieeeee//eeeeee htw.lle//223b1bc6ee/m wwt.cel1ccccccccceeee. 93652222222722 c0dd0 f f 60wd6-83dpPeee .df.
Ashraful Haque, Jessica Engel, Sarah Teichmann, and Tapio Lönnberg. A practical guide to single- cell ma-sequencing for biomedical research and clinical applications. Genome Medicine, 9, 08 2017. doi: 10.1 3073-/117770000000000
Tatsunori Hashimoto, David Gifford, and Tommi Jaakkola. Learning population-level diffusions with generative rnns. In Maria Florina Balcan and Kilian Q Weinberger (eds.). Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pp. 2417-2426, New York, New York, USA, 20-22 Jun 2016. PMLR. URL https: pww.eedincegggs www.sspresss... / v wtw.shmctooooooooo.oo 6. html.
Sø ren Hauberg, Oren Freifeld, and Michael Black. A geometric take metric learning. In F. Pereira, C.J. Burges, L. Bottou, and K.Q. Weinberger (eds.), Advances in Neural Information Processing Systems, volume 25. Curran Associates, Inc., 2012. URL htwp:: http::tttt/sseoeiitiiiii http://pooceeee..eeeeepo/o neurips. cc/paper/2012/ /ttp:/LLE/5555211111222222211110
Matthieu Heitz, Nicolas Bonneel, David Coeurjolly, Marco Cuturi, and Gabriel Peyré. Ground metric learning on graphs. Journal of Mathematical Imaging and Vision, 63, 01 2021. doi:
10.1007/<10551--000 -220000000.000--009----
Romain Hug, Emmanuel Maitre, and Nicolas Papadakis. Multi-physicss Optimal Transportation and Image Interpolation. ESAIM: Mathematical Modelling and Numerical Analysis, 1966)16771 November 2015, doi: 10.1051/m2an/201 15038. URL httph:t///// /hal wwwhhaa..hhhhhhc archives ouvertes.
St/bal-00998370
Roland Kays, Sarah Davidson, Matthias Berger, Gil Bohrer, Wolfgang Fiedler, Andrea Flack, Julian Hirt, Clemens Hahn, Dominik Gauggel, Benedict Russell, Andrea Koelzsch, Ashley Lohr, Jesko Partecke, Michael Quetting, Kamran Safi, Anne Scharf, Gabriel Schneider, IIoaa Lang, Friedrich Schaeuffelhut, and Martin Wikelski. The movebank system for studying global animal movement and demography. Methods in Ecology and Evolution, 13, 12 2021. doi: 10. 10.1111/04112110..11
Alexander Korotin, Alexander Kolesov, and Evgeny Burnaev. Kantorovich strikes back! wasserstein GANs are not optimal transport? In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2022. URL https: wwwp//eeeeeeww.eeoo//p// net/ forum?
id=VtEEpi -ddlt.....l....
Guy Lebanon. Learning riemannian metrics. In Proceedings of the Nineteenth Conference on Uncertainty in Artificial Intelligence, UAI'03, pp. 362-369, San Francisco, CA, USA, 2002. Morgan Kaufmann Publishers Inc. ISBN 0127056645.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2019. URL https: / net.
forum? hwt.:forttorme id=
Bkg6RiCqY7.
Robert J. McCann. A convexity principle for interacting gases. Advances in Mathematics, 128(1): 153-179, 1997. ISSN 0001-8708. doi: http:://oorooyylll.ooovaaaa////p URL https:
// wwww.wcieeeee.cm rect. rect. com/ / http://c/c....eeeee/////////////

---

Published as a conference paper at ICLR 2023
Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for generative adversarial networks. In International Conference Learning Representations, 2018. URL https:://// wppp/eeeeeiiieeeeeeeeee net forum? i BIQRgziTT
T.B. Mowbray, F. Cooke, and B. Ganter. Snow goose (anser caerulescens), version 1.0. Cornell
Lab of Ornithology: Birds of the World, 2020. doi: 10.2173/bow.snogoo.... URL https:
//birdssofthwwr110 wwtp:/w/g///oooooooooo bow/ species
org/
Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, Lia Lee, Jenny Chen, Justin Brumbaugh, Philippe Rigollet. Konrad Hochedlinger, Rudolf Jaenisch, Aviv Regev, and Eric Lander. Optimal- transport analysis of single-cell gene expression identifies developmental trajectories in reprogram- ming. Cell, 1 17,, 02 2019. doi: 10. 201166..ell.2211..01.
Brian L. Sullivan, Christopher Wood, Marshall J. Iliff, Rick Bonney, Daniel Fink, and Steve Kelling. eBird: A citizen- based bird observation network in the biological sciences, Biological Conservation, 142:2282-2292,, 2009.
Amos Tanay and Aviv Regev. Scaling single-cell genomics from phenomenology to mechanism. Nature, 541:331-338, 01 2017. doi: 10. 10.1038/mature21350.
Alexander Tong, Jessie Huang, Guy Wolf, David van Dijk, and Smita Krishnaswamy, TrajectoryNet: A dynamic optimal transport network for modeling cellular dynamics. In Proceedings of the 37th International Conference on Machine Learning, 2020.
Francisco Vargas, Pierre Thodoroff. Austen Lamacraft, and Neil Lawrence. Solving schrödinger
bridges via maximum likelihood. Entropy, 23(9), 2021. ISSN 1099-4300. doi: 10.3390/e2309 134.
URL https:tpp/// / /www, mwwp. www.c///////... 199p- /0/.443/3//2/2// 9/1134.
Cédric Villani. Optimal transport - Old and new, volume 338, pp. xxii+973. 01 2008. doi:
1111007/778--540-1100
Karren D. Yang and Caroline Uhler. Scalable unbalanced optimal transport using generative adver- sarial networks. ArXiv, abs/1810. 10.11144 1447, 2019.
Paul Zhang, Dmitriy Smirnov, and Justin Solomon. Wassersplines for stylized neural animation. Symposium on Computer Animation, 2022.
A BACKGROUND ON RIEMANNIAN GEOMETRY
A Riemannian manifold M (M,g) is differentiable manifold M equipped with Riemannian metric g. Throughout this paper, we take M = IR and focus our attention on the metric g. A Riemannian metric assigns an inner product ( > to the tangent space T,M of each p E M in way that varies smoothly with p. In the case where M- RD. may simply identify all tangent spaces Tp M with R:D; here a metric g amounts to spatially-varying inner product on R2D. Since any inner product on D may be computed as (u. v) = u² Ae for some > 0, a Riemannian metric on M R²" is specified by smooth field of positive definite matrices A(x) R S (we use SP denote the set of positive definite D matrices). We use (u, o) u² A(xx)u to denote
the metric at a point x E D.
()))) Given this norm, we define the length of a continuously differentiable curve 7 : [0, 1] + .M to be E(>) = fo (t)((e) dt. The A geodesic distance (z. 3) between two points x, y E M is then the infimum of the length E(>) of continuously differentiable curves such that 70)
An inner product (: A)a) induces norm ||0||
and (11 y.
B W, ON RIEMANNIAN MANIFOLDS
Feldman & McCann (2002) study Monge's transport problem on Riemannian manifold M (M, g); we focus on the case M = RD with g specified by a field of matrices A(x) > O. Given two
12

---

Published as a conference paper at ICLR 2023
compactly-suppont densities Po: Pr on M, Monge's problem seeks pushforward & of po onto Pi solving the following problem:
/ da (2))
(Po) = inf s:p smp
Feldman & McCann (2002) show that under mild technical conditions, (9) has possibly non-unique solution s. Intuitively, this map is a pushforward of Po onto #i minimizing the average geodesic distance between matched units of probability mass. The optimal value of (9) is the I-Wasserstein distance between Po and pi. which we denote by Wi,a(po.P1) PI).
If i is continuously differentiable on M, then it is 1-Lipschitz on M if and only if (o(=))la->) S 1 for all M. Here, we use Vo(x)) to denote the vector of partial derivatives of Φ at æ. Armed with this local characterization of Lipschitz continuity, we may define the following dual problem to (9):
(1 ((z)dp((x)- (le)dp((()) I
sup M-
(10)
If the minimum in (10) is attained by some Kantorovich potential o. the optimal values of (9) and (10) coincide and the Lipschitz bound for ø is tight at pairs of points (x, v) arising from Monge map solving (9) (Feldman & McCann, 2002, Lemma 4).
Feldman & McCann (2002, Lemma 10) also show that given a geodesic "/2 [0. 1] - M between a
(o(() pair of matched points (z, s(x)) and t (0, 1), we have Ve (T)(()) T (e)llaaaa)))) Intuitively, V points along geodesics Ye joining pairs of points in the support of po and #1 that are matched by the Monge map
C BACKGROUND ON DYNAMICAL OPTIMAL TRANSPORT
Let and 3 be probability measures defined on Ω C R. The Denamou-Brente formulation of the 2- Wasserstein distance W2(a, (3) defines it as the solution to fluid-dynamical problem (Benamou & Brenier, 2000)
{ [ (((x))le(())11 dr dt
(a, = = min Pe.
(11)
subject to the constraints that Po = a, P1 = 3 and the continuity equation de + V (prve) = 0.
8p.
Solving this problem yields vare-ng velocity field (z(z) that transports a's mass H's along curve of probability measures. The kinetic energy in the integrand of (11) encourages probability mass to travel in straight lines. This assumption is occasionally undesirable, but it is straightforward to modify Eq. (11) to encourage to point in specified directions (Hug et al., 2015):
(' A(z)) dx dr. Petti
min
(12)
Using the language developed Appendix A, A(z) > specifies a metric on the Riemannian manifold M (TR2 D .). and (x:)) A(x))x(x)) (x) Eq. (12) encourages be aligned with the eigenvectors μ, corresponding to the minimal eigenvalues A of the matrices A(x))
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(z) is available priori. This assumption is unrealistic for many problem domains, motivating our model, which learns a metric from cross-sectional samples from populations evolving over time.
13

---

Published as a conference paper at ICLR 2023
D EXPERIMENT DETAILS
D.1 METRIC RECOVERY
Circular data. The ground truth metric tensor for this example is A(z,8) = - -0( () u)?
where v(z, y)
2 +22
To generate the training data, we begin by drawing 100 samples each from 4 isotropic normal distri- butions with standard deviation o = 0.1 whose means are μ4 E (-1, We randomly pair samples from subsequent distributions and compute A A-geodesics between each pair by solving problem (8). We implement (8) in Pytorch using time-invariant vector field vo parametrized by fully connected two-layer neural network with ELU nonlinearities and 64 hidden dimensions. We set X 1 and solve the initial value problem =;()= no (x)()): ==(0) 20 using the explicit Adams solver in torchdiffeq's ode nt with default hyperparameters (Chen et al., 20182) 8a). We optimize the objective using A with learning rate 10 and weight decay 10-3 and train for 100 epochs per pair of samples. We then draw 24 points at equispaced times t, E [0, 1] from each resulting geodesic and aggregate across geodesics to form the observed populations X
(0, 1)}.
We then our method to recover Á (z) from the X We parametrize the scalar potentials in (4) as a single-hidden-laye neural net with 32 hidden dimensions and Softplus activation. We parametrize the matrix field (xr) as (x) = (a=)² (xx) + 10-31, where Q(x) is a two-layer neural network with Softplus activations and 32 hidden dimensions. The strength of the gradient penalty (5) is 10 when training the potentials ø in the first step of alternation and 10 in the next step; it is 1 when training Á-1 The strength of the regularization (6) is 10 We carry out two steps of the alternating scheme by training with Adam with learning rate 10- 2 and weight decay 5 10-1. We train for 300 epochs for the é step and 1,000 epochs for the A step.
We evaluate the alignment score ((A, A) on 100 x 100 grid overlaid on box of radius 1.5.
Mass splitting data. The ground truth metric tensor for this example is A(x, y) I for y 0.
022, y),(,,))))))))))))))) where v(z, 3) (1 v2) for y 2 and v(z, y) (12
v2
To generate the training data, we begin by drawing 100 samples each from a standard normal dis- tribution and from a mixture of two isotropic normal distributions with unit variance and mixture components centered at (10, 10) and (10, 10). We randomly pair samples from subsequent dis- tributions and compute A-geodesics between each pair by solving problem (8) using the method described in the circular data section. We then draw 10 points at equispaced times f; E [o. 1] from each resulting geodesic and aggregate across geodesics to form the observed populations X².
We then use our method to recover Â i=((w) from the X We parametrize the scalar potentials in (4) as a single-hidden-layer neural net with 32 hidden dimensions and Softplus activation. We parametrize the matrix field (=) layer neural network with Softplus activations and 32 hidden dimensions. The strength of the gradient penalty (5) is 2 when training the potentials & and when training Å. The strength of the regularization (6) is 106, We carry out a single step of the alternating scheme by training with AdamW with learning rate 5 * 10-3 and weight decay 10-3, We train for 600 epochs for the 6 step and 20,000 epochs for the A step.
4
Q(=))Q(())
10-31, where Q(x) is a two-
We evaluate the alignment score ((A, A) on 100 x 100 grid overlaid on the rectangular region [-2.5, 15) [-15, 15].
X-Paths data. The ground truth metric tensor for this example is A(x, v) = I v u(x, (2, T. Here we define v(z,8) ()) 3(x, (z),,)) where Us (I = (I v2 and
02(2, 9) ) We then define a(z, v) 1.25 tanh tanL(Re v)) and 3(z, v)
- -1.25 tanh(ReLU ())). Intuitively, should be large in quadrants and 3 and 3 should be large in quadrants 2 and 4.
14

---

Published as a conference paper at ICLR 2023
The training data for this example consists of two trajectories. To generate it, we begin by drawing 100 samples each from isotropic normal distributions with standard deviation = 0.1 centred at (-1,-1), (1, 1) for the first trajectory and (-1, 1), (1, -1) for the second trajectory. 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. We then draw 10 points at equispaced times 4 [0, from each resulting geodesic and aggregate across geodesics to form the observed populations X
We then use our method to recover A(x) from the X We parametrize the scalar potentials in (4) as Lungle-hidden-layer neural net with 32 hidden dimensions and Softplus activation. We parametrize
the matrix field A(x) as A(x) Q=x)" Q(x) + 10 I,, where Q(z)) is a two-layer neural network with Softplus activations and 32 hidden dimensions. The strength of the gradient penalty (5) is 10-3 when training the potentials ø the first alternating step and 10-1 for subsequent steps. It is 1 when training A. The strength of the regularization (6) is 10 We carry out three steps of the alternating scheme by training with AdamW with learning rate 10- and weight decay 5 10 We train for 300 epochs for the @ step and 40,000 epochs for the A step.
We evaluate the alignment score ((A, on a 100 x 100 grid overlaid on a box of radius 1.5.
D.2 SCRNA TRAJECTORY INFERENCE
Data pre-processing. The data for this experiment consists of force-directed layout embedding coordinates of gene expression data from Schiebinger al. (2019) collected over 18 days of re- programming (39 time points total) which we rescale by a factor of 10-3 for increased stability in training. We construct populations X for I, ...: 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.
Learning the metric. We use our method to learn a metric A- (zx) from the X We parametrize the scalar potentials in (4) as a single-hidden-layer neural net with 128 hidden dimensions and Softplus activation. We parametrize the matrix field (ar) as A-'(a)
((x)) T=()) + 10-91, where Q(x) is single-hidden-layer neural network with Softplus activation and 2048 hidden dimensions. We omit the gradient penalty (5) when training the potentials & and set the penalty strength to 10 when training Á. The strength of the regularization (6) is 5 102. We carry out a single step of the alternating scheme by training with AdamW with learning rate 5 * 10 10-3 3 and weight decay 1.5 * 10 2 We train for 100 epochs for the é step and 5,000 epochs for the A step.
Trajectory inference task. We parametrize the time evrring velocity field v1,9 as a fully connected grgee-hidden-beee layer neural network with 64 hidden dimensions and Softplus activations. We follow Grathwohl et al. (2019) and concatenate the time variable to the input each layer of the neural net. We compute particle trajectories st() by solving the initial value problem (t)) vo (x())))); (0) = xo for initial positions 20 X using the midpoint method solver in torchdiffeq's odeint with default hyperparameters (Chen et al., 2018a). The fitting loss for this task is GeomLoss's ' Sinkhorn divergence with p = 2 and blur parameter fixed to 5 10 We fix A = 10-² for both our learned metric Á and for the identity baseline A and choose the intermediate time points (, to be an equispaced sampling of [0, 1] with step size 1 for all experiments.
60
We solve problem (7) for each pair of subsequent time retained points (t;, (x41). 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. We evaluate the inferred trajectories by approximately computing the W, distance (using GeomLoss with p = 1 and blur of 10-6) between left-out time points the ground truth data and advected samples at corresponding time points. For left-out time points of form /. = ki + e for some integer E L.... k- 1, the corresponding advected sample in the equispaced sampling of step to
size I of [0, 1] has index _ 60J.
D.3 SNOW GOOSE TRAJECTORY INFERENCE
Data pre-processing. 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. This data is drawn from the February 2022 version of the eBird basic
15

---

Published as a conference paper at ICLR 2023
dataset (Sullivan al., 2009). We bin the sightings by month of observation, keep 1000 records per month, and convert the sighting locations from fatitude-longitue to zy coordinates using the Mercator projection implemented in Matplotlib's ' Basemap module with rsphere set to and the latinude/longitude of the lower left and upper right corners of the projection set the minimum latitude and longitude and maximum latitude and longitude in the training data, respectively.
The ground truth goose trajectories drawn from the "Banks Island SNGO" study Movebank. This data consists of time-stamped GPS measurements of the locations of 8 snow geese from 2019 2022. We use the same Mercator projection convert the goose location measurements from latitude/longiudee to zy coordinates. We then estimate the initial and final time point of a single migra- tion for each goose. 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} and their respective final time indices are {26000, 9100, 15057, 26500, 9000, 13037, 7201, 10000). The initial location of each goose is the initial condition To in (8), and the final location of each goose is 21.
Learning the metric. We use our method to learn a metric (x) from the X We parametrize the scalar potentials in (4) as a single-hidden-layer neural net with 32 hidden dimensions and a Softplus activation. We parametrize the matrix field (x)) as A A-)(x) = Q(x)² + 10- where Q(x) is a two-layer neural network with Softplus activations and 32 hidden dimensions. The strength of the gradient penalty (5) is 10-6 when training the potentials and when training Å The strength of the regularization (6) 10 We carry out a single a step of the alternating scheme by training with AdamW with learning rate 10-2 and weight decay 10 and 10,000 epochs for the A step.
We train for 2,000 epochs for the e step
Mey
Figure 5: The untagged goose sighting data used to learn the metric for the bird trajectory inference experiments. Each subplot depicts goose sightings in the U.S. and Canada during one month of the spring migration, beginning Jamuary (upper-left) and ending in June (bottom-right), Note that there is no correspondence between points in subsequent time points; any given goose is likely to have been sighted only once.
Trajectory inference task. We parametrize the time-varying velocity field v1,9 as a fully connected three-hidee fire--hddeen-laye neural network with 64 hidden dimensions and ELU nonlinearities. We compute particle trajectories æzt) by solving the initial value problem é(t) (=(((t)); z(0) zo for initial positions Po E X using the dopri5 solver in torchdiffeq's 's ode i nt with default hyperparameters
16

---

Published as a conference paper at ICLR 2023
1.25 10 for geese 82901, 82902, 82906, 82908, and 82909 and use À 2.5 10 for the remaining geese 82905, 82907, 82910. In each case, optimize the objective using A with learning rate 10 3 and no weight decay, set the times łj in (8) be an equispaced sampling of [0, 1] with 32 time points and train for 500 iterations.
(Chen et al., 2018a). We fix
E ABLATION STUDY: IMPACT THE REGULARIZATION COEFFICIENT
In this section, earry out ablation study of the impact of the coefficient À the regularization term R(A) (4). We repeat the metric recovery experiment on the "Circular" example. We follow the procedure in Appendix D.1 but set {5 10a, 10 5 10 10 10 109). (Note that our reported results in Section 4.1 use À = 10
The results of this experiment are recorded in Figure 6. The eigenvectors of the learned metric A(æ) robust to the value of À except at the largest value of X 10a, where the alignment score falls to 0.857. As expected, the log-condition number of A(z) increases somewhat with indicating that larger values of À favor high levels of anisotropy.
F ABLATION STUDY: REMOVING THE OT TERM
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. We repeat the metric recovery experiment on the "Circular" example. We follow the procedure in Appendix D.1 but exclude the OT term
( (()))pp(()) - Imx ot (<)dpp(()))
K Σ
(13)
from our implementation of (4), The results this experiment are presented in Figure 7. Note that 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.
17

---

Published as a conference paper at ICLR 2023
0000
Figure 6: Results of ablation study of the impact of regularization coefficient À in (4). The eigenvec- tors of the learned metric A(x) are robust the value of (rows 1,3). The log-condition number of A(z) increases with À (rows 2,4), indicating that large values of lead increased anisotropy.
18

---

Published as a conference paper at ICLR 2023
10
1.5
35
-10
Figure 7: Learned metric A(x) with the OT term (13) removed from Problem (4). The metric is no longer informed by the data when the OT term is excluded.
19


Reviews:
[{"review_id": "wQ9hYSlRiaz", "reviewer": "Reviewer_v7rg", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes to solve a dynamical optimal transport problem where: given a sequence of probability measures sampled in time, the goal is to recover an accurate trajectory of the measures for all intermediate times. Such a scenario is highlighted over two practical applications - single-cell RNA sequencing and discovering migratory paths of birds based on sparse information about their intermediate geographic locations.\n\nThe paper proposes to model these measures over a general Riemannian manifold instead of a Euclidean space where each point is now endowed with a metric tensor. The eventual goal is then to *learn* this metric from samples and demonstrate improved trajectory prediction. The authors then formulate the trajectory inference problem as an optimal transport problem over the manifold and make use of the dual formulation to find the most convenient way to learn the metric tensor field. Additionally, they also propose a regularizer that is claimed to aid the learning process. \n\nExperiments are shown for 1 synthetic and 2 actual datasets - scRNA and bird migration. The results show improvement over simple baselines like trajectory inference using the Euclidean metric or using no regularization \n", "summary_of_the_review": "Overall, I think this paper is below the threshold for acceptance. Despite rating highly on problem formulation and some clever tricks on using dual problems for convenient estimation of the metric - (1.) There is an occasional lack of background and the discussion lacks structure (2.) More importantly, the results are not super convincing beyond a simple proof of concept. Like specifically - the mass splitting example in Figure 1, and results for bird migration in Figure 4, the lack of baselines in Figure 2, etc. ", "strength_and_weaknesses": "Strengths \n\n- I find the problem setup and the solution formulation to be very interesting and refreshing. Changing the metric and having to learn that from sparse trajectory data is innovative!\n- The results do showcase a proof-of-concept, as most claims made about improvements with a learned metric instead of a trivial one are indeed so in the challenging problems of the experiments (especially Figures 1,2 and 3)\n\nWeaknesses\n\n- Overall the paper lacks global quality in terms of clarity of discussion. The parts corresponding to the scRNA problem are poorly introduced and could do with a better background. \n- Occasionally the discussion feels very hurried in the paper (especially in section 5.9). Also, how do Schiebinger et al. (2019) and Tong et al. (2020) compare with the result in Figures 2 and 3?\n- In addition, the choice of regularizer and the need to parameterize the metric and function \\phi using neural networks feels a bit arbitrary to me. What is the motivation here? in addition to easy use of gradient-based optimization? \n- The result in Figure 4 is not convincing. Ideally, I suppose the predicted trajectories must resemble the ground truth very closely.  but they clearly are a bit too smooth and don't have the high frequencies like the ground truth. Could there be a baseline comparison to show the proposed improvements? \n- Some minor issues - all figures could do with bigger sizes (especially Figure 2) and with much bigger font size for the titles. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "PY0eJwlaMh", "reviewer": "Reviewer_m2sP", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This is a relatively difficult paper to follow. Despite several attempts I have more questions about the paper than answers.  My review is based on my limited understanding of the ideas presented here. The goal here seems to be to analyze data (samples) observed from temporal evolution of a probability distribution. The assumption is that individual samples evolve in time along geodesic paths (with respect to certain nonstandard metric). Thus, the paper seeks to estimate this Riemannian metric from these temporal evolutions. However, one does not have the temporal registration between samples across times, i.e., one does not know which sample at time t matches with which sample at time t+1. Thus, the paper casts it as problem of comparing distributions (approximated by empirical distributions) and hence the use of Monge’s transport problem.  There is some related discussion in the paper on spatial regularization of the metric tensor field. The actual optimization over the metric tensor field is performed using neural network. The paper is motivated by applications in molecular biology, esp. in analyzing scRNA data where there is past literature on using optimal transport for estimating trajectories. ", "summary_of_the_review": "\nThis is a paper about learning Riemannian metric using time series data. The paper is difficult to read and understand. Given my limited understanding, it is using the Monge transport problem for solving for the unknown metric tensor field. The paper presentation and justification of the approach needs improvement. ", "strength_and_weaknesses": "Strengths: \n\nLearning Riemannian metrics from data is an important and commonly studied problem. The use of a time series data for metric learning makes this paper more interesting. \n\nWeakness:\n\nThe paper makes for a confusing read. Perhaps it is intended for an audience where the basic setup is already well known. \n\n(1)\tWhat is the justification behind assuming that the samples follow geodesic paths under some Riemannian metric? \n(2)\tWhy solve for the parallel transport problem rather than any other metric (e.g. Fisher-Rao metric) on the space of distribution space? \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "MoXBsqK__3", "reviewer": "Reviewer_M27J", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper deals with identifying a Riemannian metric for the cross-sectional samples of evolving probability measures on a common manifold. In particular, it makes use of  optimal transport theory to compute the Riemannian metric (that minimizes the 1-Wasserstein distance). The usefulness of the learned metric is shown in different applications. \n", "summary_of_the_review": "I wish to point out that the paper is not in my domain of expertise. I, however, believe that the work is a novel contribution to and should be interesting to the community.", "strength_and_weaknesses": "I believe the paper proposes an interesting use-case of learning the metric. As such it is theoretically sound. I would like to mention that it is not in my expertise to dissect and digest the results completely. The optimization problem is still abstract. It would be great to discuss it a few more details about the precise optimization steps and updates strategies. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "2Kaxp1dMCb0", "reviewer": "Reviewer_QfVV", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a method for learning a Riemannian metric for trajectory data over $\\mathbb{R}^D$ using optimal transport. The key technical novelty is the development of a minimax objective that avoids computing geodesics during training, thereby both improving training speed and stability.", "summary_of_the_review": "A potentially really neat result, but I found the paper too difficult to read for me to provide a useful verdict. I do not recommend acceptance in its current form due to the difficulties I had reading the paper, but I also claim low confidence.", "strength_and_weaknesses": "I found the paper difficult to read and have low confidence on the strengths and weaknesses listed below.\n\n## Strengths\n* The key technical innovation appears to be a minimax objective (5) for Riemannian metric learning that avoids the need for computing geodesics under the current estimate of the metric. This is potentially a big deal as computing geodesics is both expensive and unstable. I did however not understand the derivation of this objective, so I may be missing something here.\n\n## Weaknesses\n* I found the paper difficult to read as I found it non-trivial how to get from one equation to another. I get that sometimes these steps are indeed non-trivial, but a bit of hand-holding would be nice. Specifically, I did not understand how to go from Eq 3 to Eq 4 and from there to Eq 5.\n* As far as I could tell, the regularizer in Eq. 7 is going to be expensive to evaluate as I suppose you have to numerically invert the metric (cubic complexity). Is this a correct reading?\n* I missed citations to some of the many papers on learning Riemannian metrics over $\\mathbb{R}^D$. Even if the applied strategies are different than the current paper, I still think previous work should be acknowledged. See for example \"Learning Riemannian Metrics\", Lebanon (UAI 2003) or \"A Locally Adaptive Normal Distribution\", Arvanitidis et al., NeurIPS 2016.\n\n## Minor comments\n* I would suggest moving the related work section to a later part of the paper as I found that it took too long before I got to the actual contribution of the paper.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster