Partially Observed Trajectory Inference using Optimal Transport and a Dynamics Prior
Authors: Anming Gu
Edward Chien
Kristjan Greenewald
Venue: ICLR 2025 Poster (2025)
Keywords: optimal transport
trajectory inference
stochastic calculus
optimization
Langevin dynamics
Primary area: optimization
TL;DR: We leverage state space models and optimal transport to solve the latent trajectory inference problem.

Abstract: Trajectory inference seeks to recover the temporal dynamics of a population from snapshots of its (uncoupled) temporal marginals, i.e. where observed particles are \emph{not} tracked over time. Prior works addressed this challenging problem under a stochastic differential equation (SDE) model with a gradient-driven drift in the observed space, introducing a minimum entropy estimator relative to the Wiener measure and a practical grid-free mean-field Langevin (MFL) algorithm using Schr\"odinger bridges. Motivated by the success of observable state space models in the traditional paired trajectory inference problem (e.g. target tracking), we extend the above framework to a class of latent SDEs in the form of \emph{observable state space models}.    
 In this setting, we use partial observations to infer trajectories in the latent space under a specified dynamics model (e.g. the constant velocity/acceleration models from target tracking). We introduce the PO-MFL algorithm to solve this latent trajectory inference problem and provide theoretical guarantees to the partially observed setting. Experiments validate the robustness of our method and the exponential convergence of the MFL dynamics, and demonstrate significant outperformance over the latent-free baseline in key scenarios.

Published as a conference paper at ICLR 2025
PARTIALLY OBSERVED TRAJECTORY INFERENCE US- ING OPTIMAL TRANSPORT AND A DYNAMICS PRIOR
Anming Gu & Edward Chien Boston University
Kristjan Greenewald
MIT-IBM Watson AT Lab; IBM Research Cambridge, MA
Boston, MA {agu2002, edchien .ddu
kr kristjan h. www.hh..weeewwaaaammmmmmmmm wmm...bmm..m com
ABSTRACT
Trajectory inference seeks to recover the temporal dynamics of population from snapshots of its (uncoupled) temporal marginals, i.c. where observed particles are tracked over time. Prior works addressed this challenging problem der a stochastic differential equation (SDE) model with a gradient-driven drift in the observed space, introducing minimum entropy estimator relative the Wiener measure and practical grid-free mean-field Langevin (MFL) algorithm using Schrödinger bridges. Motivated by the success of observable state space models in the traditional paired trajectory inference problem (e.g. target track- ing), we extend the above framework observable state space models. In this setting, we use partial observations to infer trajectories in the latent space under specified dynamics model (e.g. the constant velocity/acceleration models from target tracking). We introduce the PO-MFL al- gorithm to solve this latent trajectory inference problem and provide theoretical guarantees to the partially observed setting. Experiments validate the robustness our method and the exponential convergence of the MFL dynamics, and demon- strate significant outperformance over the latent-free baseline in key scenarios.
class of latent SDEs in the form of
INTRODUCTION
Estimating the temporal dynamics and trajectories of population from collections of unpaired observations' at specific time points is challenging fundamental problem with many potential applications such as single-cell genomic data analysis (Lavenant et al., 2024; Chizat et al., 2022). Existing work has focused trajectory inference in the fully observed setting, where all variables that are important to the underlying dynamics are directly observed with no hidden states. We note that research in signal processing and control theory has overwhelmingly shown the importance of being able to handle latent states in dynamics modeling (Hespanha, 2018; Davis, 2013), since the best physical predictive quantities often unobserved (for instance velocity and acceleration of an object whose positions are observed). Even linear state space models have enjoyed a recent resurgence for modeling text sequence data with large language models (Gu & Dao, 2024).
While in general the problem of recovering a hidden state is not identifiable, (Kalman, 1960) etc. in- troduced systems theory and developed observability conditions on the underlying dynamics model that does allow for such recovery. For instance, in target tracking (Doucet et al., 2001; Bar-Shalom & Li, 1995), oftentimes only a position variable is observed, yet the tracking algorithm uses a state space model that includes hidden velocity state. This hidden state allows for better predicting the future position of the target, improving the trajectory inference by not only better modeling the dynamics, but making it easier to identify which of several targets observed at any given time correspond to the current track. The hidden states themselves are also often of direct interest.
The trajectory inference problem has many similarities to the tracking problem. In particular, at any given time, cloud of points is observed, but these points are not labeled, i.e. there no indication of which points at time l match to the points at time l2. Inferring these "matches" is the task of trajectory inference, and closely parallels the data association problem described by (Kirubarajan
This corresponns distributions being independent, e.g. see (iv) of Theorem C.I.

---

Published as a conference paper at ICLR 2025
& Bar-Shalom, 2004; Rezatofighi al., 2015) in target tracking. As result, motivated to introduce latent state space modeling to the trajectory inference problem.
While itself being fully observed framework without incorporating known dynamics, the opti- mal transport (OT)-based method of (Chizat et al., 2022) is particularly amenable to our aims. proposes to optimize collections of particles at each time step minimize data fit term (a cost between the particle cloud and the observed data points) and a trajectory fit term consisting of the entropic Wasserstein distance between sets of particles at adjacent time points. The entropic OT framework arises naturally from the SDE model (as we will see later), and provides an explicit and robust procedure for obtaining inferred trajectories from unpaired time series data by following the OT plan between time points. Representing the inferred time marginal densities as particles is also particularly amenable to our partially observed framework, as we can have the particles be in the hidden state space and form data fit term to the observations using specified (stochastic) obser- vation model. In many ways this parallels the observation model/hidden state partiele setup used by the particle filter (Arulampalam et al., 2002) and other sequential Monte Carlo methods (Doucet et al., 2001) for the paired-observation trajectory inference setting.
While our focus is theoretical, we envision the latent trajectory inference problem being broadly ap- plicable to a variety of real-world tasks. For instance, the extra smoothing induced by introducing a hidden velocity state could improve trajectory inference directly in many settings, including possibly the genomic data analysis setting mentioned by (Chizat et al., 2022). Another application domain could be survey or medical data. Specifically, when trajectory data is needed, longitudinal studies are often designed where individuals are followed over time and continue to be re-interviewed, but significant logistical challenges are involved in such a procedure (Thomson & Holland, 2003). Tra- jectory inference could allow for different individuals to be sampled at each time point, significantly easing the burden on researchers. Our approach for introducing hidden states (e.g. velocity) would be particularly impactful, as in both social science and medical data, often individual's trajectories (e.g- preferences or health) do carry significant momentum. Finally, our approach could have ad- vantages in private learning of time series models or private synthetic data generation for time series. This is because in (differential) privacy (Dwork et al., 2006; Dwork & Roth, 2014), the goal is to preserve the statistics of an individual, and trajectory inference allows for each individual's 's data to be limited to a single point in time, not requiring a full trajectory record that may be difficult to privatize. Our partial observation framework would allow for even more privacy to be maintained as some variables could remain hidden if an appropriate dynamics model was available.
Note the trajectory inference problem is much more difficult than simply trackinglestimating the quence of state distributions at each time point (e.g. (Kim et al., 2021)). While these time marginals can be extracted from the inferred trajectory distribution, the inferred trajectory distribution is a joint distribution over time, i.c. providing couplings between the distributions at each time point. Cru- cially, this coupling allows for the sampling of trajectories from the learned distribution and implies knowledge of individual dynamics, rather than simply uncoupled group dynamics.
Related works (Saclens et al.. 2019) surveys and compares of single-cell trajectory inference meth- ods. (Schiebinger et al., 2019) introduces the use of OT for trajectory inference; however, the method generates paths that are generally not smooth. (Chewi et al., 2021) uses OT to construct measure- valued splines, which yields smooth paths, (Bunne et al., 2022) models population dynamics as Jordan Saan--nneneleeee--OIe
(JKO) flow, and (Qu et al., 2024) uses OT to analyze gene trajectories.
(Weinreb et al., 2018) consider the limits of trajectory inference from single-cell snapshots in the equilibrium setting. However, as far as we are aware, (Lavenant et al., 2024) was the first work to provide theoretical guarantees of any estimator for trajectory inference. They introduce min- entropy estimator for unknown gradient-drives drift models and prove convergence to the ground truth in the limit the number of observations become dense in the observation period. (Chizat et al., 2022) extends the entropy minimization formulation of (Lavenant et al., 2024) by considering different fitting functional, reducing optimization space, and using MFL dynamics. (Zhang al., 2021) considers the application of these OT frameworks to the steady-state setting with known cell birth and death rates. A recent work (Ventre et al., 2024) builds on (Lavenant al., 2024) and provides theoretical guarantees for trajectory inference in the branching case. Similar in spirit to our work, Shen et al. (2024) consider trajectory inference in the setting where a family of reference measures is specified, and Zhang et al. (2025) considers learning the unbalanced setting via neural networks. (Kim et al., 2021) considers a similar latent space dynamics model as us and aims to
2

---

Published as a conference paper at ICLR 2025
recover the time series of latent population distribution, but noted above their work does not infer a trajectory distribution (specifically, the inferred time marginals are not coupled), and therefore is not designed to generate trajectory samples or characterize the individual dynamics (as opposed to population dynamics). These works, along with ours, do not aim to recover the drift dynamics; however, the works (Bunne et al., 2022; Tong al., 2020; Hashimoto al., 2016) do study this problem.
(Vahdat al., 2021) uses score-based generative modeling in latent space. (Jiao et al., 2024) uses pre-trained encoder and decoder, consider diffusion in latent space, and prove theoretical guarantees that the output distribution is close to the ground truth. (Hamm al., 2024; Zhang et al.. 2023) consider learning latent manifold structures using OT, Song et al. (2023b) consider gradient flow in latent space to study equivariant networks, (Song et al., 2023a) study the latent space of generative models using OT, and (Al-Jarrah et al., 2024) consider the nonlinear filtering problem using partial observations using OT. Others consider applications of Schrödinger bridges with non- -ineee ref- erence measures. For example, (Chen et al., 2016) consider Schrödinger bridges where the prior is any Markov evolution for control theory and (Bunne et al., 2023) show that Schrödinger bridges between Gaussians against reference measures induced by linear SDEs have closed forms.
Notation For probability measures μ. v. the relative entropy (e.g. the KL divergence) is H (Hlv) = (eg(du/da yd if μ & v and too otherwise. For n E N. let [n|] == {l..... n). We use X to denote the latent space and y to denote the observation space. We use the notation (()) to denote the probability distributions over a space. The path space is Ω C( 1]: X)), the set of continuous X-valued paths. In our theoretical discussion, as in (Lavenant et al., 2024), we assume without loss of generality that the end time interval is t If R P(22) is probability measure on the space of paths, its marginal at time & is denoted as R. E P(X)). We generally use the Greek letters μr, P to denote probability distributions on V. respectively. We use 82 to denote a Dirac delta at x. By an abuse of notation, we use 2 = (---) E R for both the squared norm of a vector and the quadratic variation of a stochastic process. For function g : X -> y and measure μa, we use Gylu to denote the pushforward measure, e.g. (:((B) = uوg (B)) for a measurable set B C y.
LATENT TRAJECTORY INFERENCE
We now state the problem of latent trajectory inference, which informally is depicted in Figure 1. At a high level, trajectories are modeled as being generated from a stochastic differential equation model (SDE) which incorporates both known terms (e.g. the latent dynamics model describing known relationships between states") and unknown terms that allow for unknown forces or model misspecification. Note that specifying the dynamics between the latent states crucial to make unobserved states identifialle as discussed later. Specifically, let X, X' be an unobserved state vector evolving according to the following SDE for t E [0,1]:
dX, EEt, X.)dt
X.)dt + VFdB,,
(I)
where {B.} is a Brownian motion, is the known diffusivity parameter, E E C(10, x x) is a known driving vector field (i.c. the dynamics model), and v E C² 1] x X') is an unknown potential function. Let P be the law of the SDE with initial condition Po where P, E P(X)) are the marginals of P at time t E [o. 1]. Our SDE differs from that of (Lavenant et al., 2024; Chizat et al., 2022) by the addition of non-zero known E which allows for idenfiable non-observed latent states.
As this latent space is not observed, we require a model mapping the latent space to the observation space in which our data samples live. Consider smooth function g : X + y transforming X, into the observation space y: Y, = g( (X.)). Suppose we have 7 observation times with 0 < 17" < ..... IT < 1, and we observe N²" i.idd. samples from the marginal distribution of Ygr:
(Y N NT itd. P:P.T Q.T+
(2)
forming empirical distributions / NY プラ Byz. Here, note that our framework allows for varying number of samples at each time point: indexes the time, j indexes the data points, and the T' superscript is used to highlight the dependence on the number of time points.
Forr instance, the fact that velocity state additively impacts future positions states.
3

---

Published as a conference paper at ICLR 2025
Nosttoon
Vebeciy
PD PLL
PO MPL
MFL
30
20
70
1010
10-
50
10
(a) Ground truth.
(b) Reconstructed trajectories.
(c) Velocities.
Figure 1: Constant velocity model, where the variance of the ground truth has been rescaled. We see that our method, PO-MFL, is more robust as the baseline method, FO-MFL, fails to converge, and provides per-particle velocity in contrast to FO-MFL. See Section 4 for the experiment setting.
Our goal is to recover the trajectory distribution P given the marginal snapshots in observation space (T)) In particular, we provide an algorithm to sample trajectories in the latent space (which can then be mapped to trajectories on the observation space). In general, to make this problem well- posed and tractable, we make several assumptions on this very general setup. The first of these is notion of observability that generalizes the ensemble observability introduced in (Zeng et al., 2015):
Definition 1 (Ca-ensembie observability). Suppose X, follows (1). Assume that V is unknown but restricted to a class Cy. We say the tuple (g. E, Cy) is Co-ensemble observable if. given g. E, т, and all marginals Q = P;P, of Y2 for all t E [0, 1]. the marginals P. of Xx are uniquely determined for all t E [0, l..
With this observability assumption, we can infer the latent dynamics solely from the marginals Q:- A discussion of the relationship between this condition and that of classical/ensemble observability is present in Appendix B. There, we also verify the conditions of Definition I for several important setups, e.g., the key velocity-basee dynamics model we use in our experiments.
Theoretical assumptions We use the following assumptions for the theoretical analysis and provide more thorough discussion in Appendix C. Let X. be Polish spaces, where X' is smooth and compact Riemannian manifold or a compact and convex subset of RRd. In the manifold setting, we assume its Ricci curvature K is bounded from below, e.g. K -c.
The path space . = C((o, X) is equipped with the uniform topology and its Borel o-algebra. The probability space on paths P(T) is equipped with the weak topology, e.g. convergence against bounded, continuous functions. Assume our probability space (Ω, F, P) is complete and filtered, where the filtration is with respect to the process {X₂): IP is a probability measure, and if it is not specified, expectations are taken with respect to P. Let Wa,r be the measure induced by the SDE d7, E(t, Z,) dt + VFdB,
Assumption 2.1 (Dynamics and Observation Model), Assume E C( x X X) is known, divergence-free, Lipschitz continuous, and satisfies ||E|| < +00. Assume the observation func- tion g : X + y is smooth, measurable, bounded, and time invariant.
The divergence-free all time if the initial condition is vol,* where vol denotes the uniform measure on X. Lipschitz continuity and |||| < too are technical conditions for our proofs, and there are also some necessary mild technical conditions on the pair (E, v). Finally, by "bounded" g we mean the image of set of finite measure also has finite measure.
assumption is required so that the time marginals of W remain vol for
3 PO-MFL: APPROXIMATE MINIMUM ENTROPY ESTIMATION
We seek to estimate the trajectory distribution by maximizing its ligeliieldoood with respect to the distribution induced by the SDE, subject to matching the observed marginals. Inspired by (Lavenant et al., 2024), we pose the maximum likelihood problem an equivalent mininum eniropy estimation problem connecting temporal snapshots into continuous trajectories, where the entropy is the relative entropy (KL divergence) of the estimated trajectory distribution with respect to the known portion of the SDE. We will show that the optimal point of the minimum entropy objective function converges to the ground truth trajectory distribution. It is not practical, however, directly work with the
Note that if X' has boundary, we need zemo flux condition on
4

---

Published as a conference paper at ICLR 2025
trajectory distribution as it is the minimum entropy objective in trajectory space can be reduced to an OT-based objective, where marginals at adjacent time points are coupled via entropic OT. This reduction allows to perform trajectory inference using only representations of the latent space time marginals, which can be accomplished via sets of particles. These particles can then be optimized via MFL dynamics.
object. In what follows, will show that
3.1 MINIMUM ENTROPY ORJECTIVE FUNCTION
In this section, we specify the minimum entropy objective function on the trajectory distribution. Let (1f) C [0, 1] be our observation times, where At; := t 1." Recall that, in general, we do not have exact measurements of the temporal marginals, we will only have samples from them. As a result, we must introduce a fut function measure the discrepancy between the observation space time marginals of the estimated trajectory distribution R and the observed samples.
Let the observed empirical distribution be P == N Ei 0yz E P(V) for E [2]. Consider the fit function Fit P(O)) R:
T AL (eT PP),
Fil F Qep)
(3)
X i=1
with data- fata-fitting term introduced by (Chizat al., 2022) augmented by observation function g to be
DF DF (o: (.TT / log J exp (- [llo(z) - y11² dR,g(()) dp.(v)
20
(4)
H Rep No) +H(
where N₂, is the Gaussian kernel with variance o², C is a constant, and we use the substitution X' as the optimization will occur the latent
Qe Rg Notice that the inner integral space, while the inner integral is over as the observations over V
This data-fitting term can be interpreted as the negative log-likelihood under the noisy observation model V.J x(X[)) + Zz.j where Y is the observation and Zi., s.id. NO, I). Note DF is jointly convex in (R.y,/i ,P?) and linear in Pf" The main difference compared to (Chizat et al., 2022)
is that our data-fitting term is in observation space, e.g. the addition of the function g.
The minimum entropy estimator introduced in (Lavenant et al., 2024) minimizes F :(() +
F(R) Fit (Qe, Qx7) (
Recall that our key novelties are the fit term in observation space and entropy minimization in path space with respect to divergence-free Markov path measures. Nonetheless, we show that we can still recover the ground truth in the limit as the number of observations becomes dense.
Theorem 3.1 (Consistency (informal, see Thm. C.I)). Suppose P is the SDE (1) with initial con- dition Po E P(X) such that (Pu vol) < +00. Let R TTA,0 E P(2)) be the unique minimizer of (5):
T.A. = arg min F(R).
REP(O)
(4)) heer becomes dense in [0, 1]. then 10 (limT
R..A..) almost surely.
This weak convergence result parallels Theorem 2.3 in (Lavenant et al., 2024), which provides consistency result in the fully observed setting where the entire state vector X, observed and = = 0 identically. Due to these differences, the result in (Lavenant et al., 2024) cannot be directly applied to our setting, and while our proof able follow similar overall structure, dense and nontrivial changes throughout the extensive proof are required. See Appendix C.
3.2 THE REDUCED PROBLEM
As (5) is an infinite-dimensional optimization problem, to apply the MFL dynamics (Hu et al., 2021), need to "reduce" the problem over the space P(X)TT similar (Chizat et al., 2022). As before,
5

---

Published as a conference paper at ICLR 2025
let At; := f/+ ty and At. Consider the following entropic OT cost (see Appendix A), defined for some T; 0, as
/ (1 ()) +t) + min TTH v), (6)
T,, == min TEl((..))
-El((M)))
where pi is the transition probability density of wa.r over the time interval [0. and the cost function is c (ar,)) y) T log(p?, (x, u)). In general, p7, cannot be found explicitly, so we will discuss how to approximate this below. This optimization problem is exactly a Schrödinger bridge problem, e.g. (Chizat al., 2022, App. A); (Léonard, 2012: 2013), which inspires the following reduction. Define the functional G: P(X)T + R for = (M()) of reconstructed temporal marginals. by
(T) that represents family
T-1
G( == Fit (9L)) +
Σ
TT, ((tt))) A4
where O:fl " is an abuse of notation to mean the dlement-wise push-forward of each fz by We consider the reduced objective F :P(2)T + R, defined as
F(U): == G(p) + (H))
(d) ( is the negative differential entropy of the family of measures J.. Similar to (Chizat et al., 2022), we have an equivalence of minimizing F, (5), the objective in
where (m)) = loog(u (9)
path space ((2),, and F, (8), the reduced objective over P(x)" The proof provided in Appendix D.
Theorem 3.2 (Representer theorem). Let Fit P(y)* + R. be any function and let be bounded and divergence-free.
(i) F admits minimizer R then (R. Ri R is minimizer for F.
(i)) If admits a minimizer μ²' E P(X), then a minimizer R for F is built as
R²()= - f (- 1Z (TI....., XT))
where (1 ) the law of WE.T conditioned on passing through Ii,..., PT at times Ef ....,
, f.p respectively and R.q.,., is the composition of the optimal transport plans i/s that minimize (px p (i+1)) for [T 1].
The composition of the transport plans is obtained as:
Re. (dx)..... drr) = (dri, ((2)))(daa)))) (drrlledo....--dddooorlrr--llll
where the OT plans 7a (dr,, dri41) = 7/s (do;+1 |z; (u (da;)) are conditional probabilities (or "dis- integrations"). As in (Chizat et al., 2022), the "reduction" of the optimization space from P(22) to P(X')" is enabled by the Markov property of WE.r which holds for us due to the Lipschitz conti- nuity assumption on and that W remains the uniform measure at all time. Then, Theorem 3.2 allows us to compute a minimizer for F from minimizer for F and its associated OT plans.
3.3 APPROXIMATING THE ENTROPIC OT COST
Although now we have a reduced problem over P(X)",, we are not yet ready to solve the objective function F (8), as the entropic OT cost T v) is defined with the transition function p which generally not available in closed form. We provide an approximation of T(, v) by considering an Fuler-Maruyama discretization (Kloeden & Platen, 1992). Let μ E P(Ω) be stochastic process following the SDE dX,, = -3(1, Xt) dl + /F dB, with an arbitrary initial distribution. Let A!:- ta l and suppose Mt; 14g two time marginals of μ. Recall that if X M4g and
X1, = X1, - /
Xs, ds VF(Br, - B,,)
(10)

---

Published as a conference paper at ICLR 2025
then Xt, Mtg- For small At t1, since 3 and as smooth, the integrand of the second term will be approximately constant over the integration interval. Thus, we can approximate the first two terms of (10) as
(x,,
: X - (11, x At
Finally, the last term of (10) is an isotropic Gaussian with variance 7 Tåt. This suggests approximat- ing the transition kernel pr, as deterministic drift given by the current E, followed by isotropic Gaussian noise.* 4 For = = 0, this would reduce to the kernel used by the baseline method, i.c. the Brownian motion transition kernel.
This provides an intuition for why our approach is more robust than that the baseline method when the true E is non-zero, since (X,,, X1, (O0 7 Af) for small At, while the Xx, -Xa used by the baseline algorithm has non-zero expectation, is non-Gaussian, and often has significantly higher variance. In sense, our Fuler-Maruyama approximation can be considered first order approximation method, while the baseline corresponds to a zeroth order method.
To formalize this intuition, first let (V.,,) -) denote the total variation distance. We have the following bound on the densities, which is special case of (Bras et al., 2022, Thm. 2.1). Note that this implies that the difference in probability between the approximate kernel and true kemel for any event is of order O ((AA))//3) which converges
T oo.
Proposition 3.3. Let X12 be as in (10) and X12 == (X,,(X,,) + VF(B,, -Br,), where Xx, 8g. Then, TV (Xt,, x I Cale (Af)1/3, where the constants C1. C2 > depend only on dim X and the Lipschitz constant of
Applying this approximate transition kernel yields updated, tractable OT terms in the objective func- tion. Specifically, instead of (1 (1(++1) ) in (7), we use Tr, (5, (cftt1tt (t) (((+1)), where
(1, min Ti HH(()l,,, yell(u..)
and Pt (z, y) is the transition probability density of the Brownian motion on X over the time interval [0. tl. This cost is easily computed as Pr, the Gaussian kernel. In particular the cost function is Varadhan's 's approximation
(z, TT log(p,, (5>1(z) v)), and
1 2 llv + At 2(1
C (2.z)
which holds for т; small (Norris, 1997). To wrap up discussion here, it is important to highlight that we require a generalization of (Lavenant et al., 2024, Thm. 2.3) using our path measure W Theorem 3.1, to justify convergence of our estimator when including E in our cost function in the entropic OT problem (6).
3.4 MEAN-FIELD LANGEVIN DYNAMICS AND EXPONENTIAL CONVERGENCE
We provide brief description of the MFL dynamics. See Appendix for more complete dis- cussion. MFL dynamics are designed to minimize functionals of the form F, = G where G P(X) + smooth and H is minus the differential entropy. Using the first-vaiiation Viμl of G given in Proposition E. the MFL dynamics is defined as the solution of the following non-linear
McKean- Vlasov SDE, for 2
V (x) ds + 2T dB!.) + d [14 = Law()!"),
dx. (1)
Law Law(X,)) (X((0) Mo i E (T).
(11)
where d& is the boundary reflection in the sense of the Skorokhod problem, e.g. Tanaka (1979). Note that E affects this equation implicitly via the Schrödinger potentials. For computation, we discretize the MFL dynamics via noisy particle gradient descent, discussed in Appendix E.2.
In (Chizat et al., 2022; Chizat, 2022; Nitanda et al., 2022), it is shown that the MFL dynamics converges exponential rate to the minimizer. We provide the proof for the following similar result for our partially observed setting in Appendix E.5
A lternative discretizations can be considered, c.g. we use a midpoint discretization for the constant- velocity experiments, but for clarity, we introduce the method using standard Fuler discretization. 'The d-torus assumption is required for technical reasons in (Chizat, 2022).

---

Published as a conference paper at ICLR 2025
Tinepcents
Changing Stservations
3.-
30
Figure 3: Average W2 distance between ground truth and PO-MFL. (blue) and FO-MFL (orange) recovered positions "constant velocity" model (App. F). (L) Number time points. (C) Number of observations. (R) Velocity.
Theorem 3.4 (Convergence). Assume X is the d-torus. Let po E P(X)" be such that F(uo) < +oo. Then for t 2 O, there exists unique solution (ps)s20 to the MFL dynamics (11). Let T > 0
and assume that μg has a bounded absolute log-density, holds
(1,s))
)- min F, S CF
min F+),
where C= Be for some a, B O independently of μ and т. Moreover, taking a smooth time- dependent T, that decays asymptotically as â/ log for some ā >> a, it holds Fo(u,) - Fo((²)) S log log s/ logs s > and μ. converges weakly to the min-entropy estimator μ³.
3.5 PO-MFL
one particle Y Velocity
Mean Y Velocity
We summarize proposed latent trajectory inference method in Algorithm 1. We recall that discussion on the cost function (line 4) and MFL dynamics (line 7) can be found in Sections 3.3 and 3.4, respectively. We use the Sinkhorn algorithm for entropic OT, which discuss in Appendix A. Using N iterations of the MFL dynamics, the total runtime for algorithm is (NNTm²), need to solve T entropic OT problems on X m size matrices in each iteration, where assume that the num- ber of iterations in the Sinkhorn algorithm is capped at a constant. We remark that our PO-MFL method has the same runtime as the baseline FO-MFL (fully-observed MFL).
IC
1500
1750
2000
250 Earatin
500
Figure 2: (left) Velocity of one particle end of optimization. (right) Popula- tion velocity at beginning of optimiza- tion, showing exponential convergence.
Algorithm PO-MFL: framework for latent trajectory inference
Require: Collection observations (?ı,..., .er) collection of T time samples (+ (f.'.... tTT) velocity dy- namics E, number of iterations for MFL dynamics N, number of particles m, entropic OT parameter A
1: Initialize m particles for each time: (fi.....,fv) .ur) E x
2: iterations do
3:
for i [T do
4:
A4, (7 C
5:
C. = 1ii + I may + ALE(( mi)))) % Sinkhorn(mm,, mi+1, C,, At,)
6:
7:
end for
8:
m f- MFL(M, 7,
> m = (m.,..., m.), etc.
9:
end for
10:
Output collection of particles ı., trajectories 7/-1 o...
Given the set of observed temporal marginal samples in the observation space V. Algorithm 1 yields set of m particles at each time step i representing the temporal marginal distributions in the latent space X. Simulated trajectories may recovered by sampling from the composition entropic transport plans shown in (9).
8

---

Published as a conference paper at ICLR 2025
4 EXPERIMENTS
In this section, we briefly provide synthetic experiments that demonstrate the advantages of having a dynamies prior. More details on the experimental setup can be found in Appendix F."
"Constant velocity" model We compare the behavior of our method, PO-MFL, that of the baseline, FO-MFL, using the "constant velocity" model popular in tar- get tracking, e.g. (McIntyre & Hintz, 1998). We provide further details of this model and its ensemble observability in Appendix B.I. This model can be inter- preted as introducing velocity state to be inferred, in order to build mo- mentum into the dynamics (an object in motion tends to stay in motion). This is an extremely generic model and makes min- imal assumptions on the underlying data, as evidenced by its use in target tracking.
Posttion
Velocity
2
2.5
0.0
-2
-2.5
-2.5
10
hidden
(a) Ground truth.
PO: MFL
FD MFL
-2
225
0.0
2.5
-2.5
0.0
2.5
In this model, the state space is X (ar y, èt. v) E IR,, with E given the appendix and observations (XX (72,02x2 0x Note that due to non-7ero process noise despite the name, this model does not imply that the velocity is constant in time. The particles are ini- tialized at the origin with velocities set as
(b) Reconstructed trajectories.
t = 5 and = 7, ie. Xg (0,0, 5, 7). The ground truth is shown in Figure la.
(c) Reconstructed velocity from PO-MFL. Note the bimodal velocity estimate.
Our optimization method observes only the positions of the particles, i.e. g(x,, ") (x, y), but it uses = as being a constant velocity prior. Results shown in Figure lb show that PO-MFL is able successfully reconstruct the paths trajectories, while FO-MFL fails to converge. Furthermore, in Figure Ic, we verify that the population average of the particles' velocities matches with the ground truth.
Figure 4: stant velocity" SDE.
Crossing paths experiment under the "con-
Figure 2 (left) displays the velocity of one particle for the last 500 iterations of optimization. Although at each iteration, the velocity stochastic, we can see that the mean at 7. In Figure (right), we plot the average y velocity in the first 400 iterations of optimization, providing empirical evidence of the exponential convergence of our algorithm guaranteed by Theorem 3.4.
Figure 4 shows crossing paths experiment where the population is divided into two groups, one moving right and down, and the other right and up, with their paths crossing in the middle. In this particularly illuminating regime, PO-MFL leverages the "constant velocity" model used in this section to distinguish the downward moving group from the upward moving group. Note that PO- MFL is not told priori which samples belong which group. While FO-MFL here collapses the centroid, we point out that even if its optimization was successful, FO-MFL would prefer V-shaped trajectories, e.g. those that have velocity v for half the time and then - for the remainder, here rather than the correct straight-line trajectories, as it does not retain a hidden velocity state and only seeks match adjacent time points by their relative position via entropic OT.
We provide a variety of additional experiments in Appendix F illustrating how performance changes as the number of observed particles, the spacing of time points, and the underlying ground truth initial velocity affects performance. Figure 3 shows the average W2 distance between the ground truth positions and recovered positions (averaged by time point) across these experiments. Our approach remains significantly more robust as these parameters are varied compared to FO-MFL.
We provide code in http:///tthhh.oooom//mm/miiiiiie

---

Published as a conference paper at ICLR 2025
Ground Tuth
77-
700
292
175
275
320
200
400
500
(b) Samples from the learned trajectory distribution (lines), overlaid with scatter plots the learned par- ticle representations of the time marginals.
(a) Ground truth samples used to learn trajectory distribution.
Figure 5: Wikipedia page traffic data.
Wikipedia traffic data We briefly consider real-world daily traffic data from Wikipedia.? Because some of the webpage traffic has large spikes/outliers, we only consider pages whose daily visits are between 100 and 500. We use the data from 3 pages over the course of 500 days as our ground truth and seek to learn the trajectory distribution.
For PO-MFL, we do not consider partial observations. Instead, we leverage the dynamics prior E to set an autoregressive model Z++1 = D2 + 02 and let g be the identity function. this setting, the state space for FO-MFL is R while the state space for PO-MFL is 2 The first marginal of PO-MFL matches the data for FO-MFL, while the second marginal of PO-MFL is the data lagged by 6 days. For the values of f and 62, take the average of the regression parameters calculated from 30 trajectories drawn from the same distribution.
We try different number of particles port the results in Table 1. To compute the values in the table sample trajectories from the output of the algo- rithm and use this as our empirical measure. We use the Wasserstein distance in 5000 our cost metric. Here, we use just the first dimension for PO-MFL. We see that both the mean and variance from PO-MFL are significantly less than that of FO-MFL. We plot three trajectories in Figure with the m experiment, and note that PO- MFL yields less volatile trajectories more in line with the
3, 6, 15 and re-
PO-MFL
FO-MFL
1.0110.0395
1.61:0.206
1.071000026
1.76:10.303
15
1.05:1000956
1.91
Table 1: Average (normalized) W2 dis- tance between true and sampled trajec- tory distributions for Wikipedia dataset over 100 MCMC trials. Error bars are one standard deviation.
ground truth trajectories.
5 CONCLUSION
We consider the problem of trajectory inference latent space based on indirect observation, ex- tending the theoretical analysis of the min-entropy estimator introduced in (Lavenant et al., 2024) and the MFL dynamics algorithm introduced in (Chizat et al., 2022). Experiments were provided showing that the ability to include simple non-informative latent dynamics models, such the "con- stant velocity" model, and autoregressive models, can dramatically improve the trajectory inference performance the baseline FO-MFL method. For future work, while do here provide some flexibility for model misspecification via the unknown W potential, it would be interesting to further explore the stability of our method when the dynamics model E is misspecified. Further exploration of ensemble observability to stochastic systems would also be highly interesting fundamental di- rection explore. Finally, we will seek explore the various promising empirical use cases outlined in the introduction.
ACKNOWLEDGEMENTS
The authors would like to thank Hugo Lavenant for helpful discussions and the anonymous reviewers for their feedback.
We train 2. from http:///wwwwseooooooooooooeeeeee
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Mohammad Al-Jarrah, Niyizhen Jin, Bamdad Hosseini, and Amirhossein Taghvaei. Nonlinear fil- tering with Brenier optimal transport maps. In Forty-first International Conference on Machine Learning, 2024. URL ht tss //pp////eee/eeeeeeeeeeeeeeeeeeee
Bertt//ormmmdddbbe blzDxD 6bKt.
Jason Altschuler, Jonathan Niles- Weed, and Philippe Rigollet. Near-linear time approximation algo- rithms for optimal transport via Sinkhorn iteration. Advances in Neural Information Processing Systems, volume 30, 2017. URL https::/pproeeeedss....ett
cc/ hwpp/pppr... fars__fils//___ Pawp:e/ddd.d.e-aae.........
filel 4 4ta.11/1111111efecccaceeeeeeeea0
Mare Amaudon, Ana Bela Cruzeiro, Christian Léonard, and Jean-Claude Zambrini. An entropic interpolation problem for incompressible viscous fluids. Annales de l'Institul Henri Poincaré", Probabiliiéés et Statistiques, 56(3):2211 2235, 2020. doi: 10. 001111/99- URL https: ///oi. ott/100.gg1...0 1tapr///..111//// 1214/19-A1HPP11 036.
M Sanjeev Arulampalam, Simon Maskell, Neil Gordon, and Tim Clapp. A tutorial particle filters for online nonlinear/non-Gaussie Bayesian tracking. IEEE Transactions on signal processing. 10(2):174-188, 2002.
Dominique Bakry, Ivan Gentil, and Michel Ledoux. Analysis and geometry of Markov diffusion operators. 2013. URL httpt::////// / /ww.. http:///ooooooo..o//11////1/ otg/11..100/997-- 3-1---------1 0219-0027-222
Yaakov Sar-Shalom and Xiao-Rong Li. Millitarge-mulliieesoor tracking: principles and tech- niques, volume 19. YBS publishing Storrs, CT, 1995.
Jean- David Benamou, Guillaume Carlier, Simone Di Marino, and Luca Nenna. An entropy min- imization approach to second-order variational mean-field games. Mathematical Models and Methods in Applied Sciences, 29500):1553-1583 2019.
Pierre Bras, Gilles Pagès, and Fabien Panloup. Total variation distance between two diffusions in small time with unbounded drift: application to the Euler-Maruyama scheme. Electronic Journal of Probability, 27, January 2022. ISSN 1083-6489.
doi: 10. 10.1214/22-ejp881. URL http:
/ //xxx doi org/r 100 1214/22--1.J888
Charlotte Bunne, Laetitia Papaxanthos, Andreas Krause, and Marco Cuturi. Proximal optimal trans- port modeling of population dynamics. In International Conference on Artificial Intelligence and Statistics, pp. 6511-6.28. PMLR, 2022.
Charlotte Bunne, Ya-Ping Hsieh, Marco Cuturi, and Andreas Krause. The Schrödinger bridge be- tween Gaussian measures has closed form. In International Conference on Artificial Intelligence and Statistics, pp. 5802-5833 PMLR. 2023.
Fan Chen, Zhenjie Ren, and Songbo Wang. Uniform-in-time propagation of chaos for mean field Langevin dynamics, 2023. URL https: httpttt/////iiii wrrpiii.oorg/aaaee abs///////// 2212.
Yongxin Chen, Tryphon T. Georgiou, and Michele Pavon. On the relation between optimal transport and Schrödinger bridges: A stochastic control viewpoint. Journal of Optimization Theory and Applicationss, 169(2):671 691, 2016. doi: 10. 10.1007/410957--15-000333 URL. https: /oi. org/10 11111100000001100000-0-----000
Sinho Chewi, Julien Clancy, Thibaut Le Gouic, Philippe Rigollet, George Stepaniants, and Austin Stromme. Fast and smooth interpolation on Wasserstein space. In Proceedings of The 24th International Conference Artificial Intelligence and Statistics, volume 130 of Proceedings of Machine Learning Research, pp. 3061-3069, PMLR, 13-15 Apr 2021. URL https:// proceedings mlr. www.m....eess....mmm.mmmmmmmmmee www.pppssssslleeee chtwwcewwwaaa/ea .1t.mmll............
Lénaïc Chizat. Mean-field Langevin dynamics Exponential convergence and annealing. Transac- tions on Machine Learning Research, 2022. ISSN 2835-8856. URL /ttpps://ppeeeeview
net. Eorum? wtt.ummmm==3Dqz..... gEm.
Lénaïc Chizat. Stephen Zhang, Matthieu Heitz, and Geoffrey Schiebinger. Trajectory inference via mean-field Langevin in path space. In Advances in Neural Information Processing Systems, 2022. URL https::/p... http//pppeeeeeeeeeeeeeeeeeeeeeee net./ forumr 1d=fft 4

---

Published as a conference paper at ICLR 2025
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Ad- vances in Neural Information Processing Systems, volume 26. Curran Associates, Inc., 2013. URL https:/ htttpp::/pproccediiggs...
htcpippp..cpppppp/e wc/ppper.p.p.../.. .pff.
Mark Davis. Stochastic modelling and control. Springer Science & Business Media, 2013.
Arnaud Doucet, Nando De Freitas, Neil James Gordon, et al. Sequential Monte Carlo methods in practice, volume Springer, 2001.
C. Dwork and A. Roth. The Algorithmic Foundations of Differential Privacy. Foundations and trends in theoretical computer science. Now, 2014. URL https://bocks: googl e com/ books?
Cynthia Dwork, Frank McSherry, Kobbi Nissim, and Adam Smith. Calibrating noise sensitivity in private data analysis. In Shai Halevi and Tal Rabin (eds.). Theory of Cryptography, pp. 265-284, Berlin, Heidelherg, 2006. Springer Berlin Heidelberg. ISBN 978-3-540-32732-55
Z. Gajic, Z. Gajić, and M. Lelić. Modern Control Systems Engineering. Prentice-Han international series in systems and control engincering. Prentice Hall, 1996. ISBN 9780131341166. URL https: //bokks. google com/books ? d 48 89SAAAAMAAJ.
Albert Gu and Tri Dao. Mamba: Lincar-time sequence modeling with selective state spaces. In First
Conference on Language Modeling, 2024. URL https: httpss:///pppeeieeeeeee net hwtp:/.eteefraaaaaae forum?id-
tEYskwlVY2.
Keaton Hamm, Caroline Moosmüller, Bernhard Schmitzer, and Matthew Thorpe. Manifold learning in Wasserstein space, 2024. URL https: org/ wwarrg//b////// ab/////s////1/// 08549.
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 petps::/prooceddiee
ml preess.p////..... v48 8/ hwshimotoll 6. html.
Joao P Hespanha. Linear systems theory. Princeton university press, 2018.
E.P. Hsu. Stochastic Analysis on Manifolds. Graduate studies in mathematics. American Mathemat- ical Society, 2002. ISBN 9780821808022 URL https: /btoksoooss/////// beoooks.gogglo gwogle..oooooooooooooooooooo. com/books"
id=GDEPCgARQBRU.
E.P. Hsu. A brief introduction to Brownian motion on a Riemannian manifold, 2008.
URL pettss:/ppeoplee math. harvard, htt..aaaa.e.deddeettteaa//aaaaaa aecb///ctt//omme//ee/ttlt/t/ pdf.
Kaitong Hu, Zhenjie Ren, David Šižka, and Lukasz Szpruch. Mean-field Langevin dynamics and energy landscape of neural networks. Annales de l'Institut Henri Poincaré, Probabilités et Statis- tiques, 57(4):2043 2065, 2021. doi: 111211/0-AAHH1111 URL https: doi. or. 1214/20 11144/20AAHH1111
Adel Javanmard, Marco Mondelli, and Andrea Montanari. Analysis of two-layer neural network via displacement convexity. The Annals of Statistics, 48(6):3619 3642, 2020. doi: 10.1214/ 20-AOS 1945. URL httph:tt///te... :tttss//ddi.. .rgg 10. / 945.
Yuling Jiao, Lican Kang, Huazhen Lin, Jin Liu, and Heng Zuo. Latent Schrödinger bridge diffusion model for generative learning, 2024. URL https: /arxivi v. org/ abs 2404 13309.
Rudolf Kalman. On the general theory of control systems. In Proceedings First International Conference Automatic Control, Moscow, USSR, pp. 481-492, 1960.
Jin Won Kim, Amirhossein Taghvaei, Yongxin Chen, and Prashant G. Mehta. Feedback particle filter for collective inference, 2021. URL https: v www.rrwww.oro//b...o.o/..///bare /2010 0010 0 06655.
12

---

Published as a conference paper at ICLR 2025
T. Kirubarajan and Y. Bar-Shalom. Probabilistic data association techniques for target tracking in clutter. Proceedings of the IEEE, 92(3):536-551) 2004. doi: 10.1109/PROC..2033883.
P.E. Kloeden and E. Platen. Numerical Solution of Stochastic Differential Equations. Applica- tions of mathematics stochastic modelling and applied probability. Springer, 1992. ISBN 9783540540625. URL attps://boote google wwtpogee.co/bbooooooooo www/books 7id=7bxZAOOAAAA
Hugo Lavenant, Stephen Zhang, Young-Heon Kim, and Geoffrey Schiebinger. Toward a mathemat- ical theory of trajectory inference. The Annals of Applied Probability, 34(1A):428 500, 2024.
doi:
URL https: doi org 12144/231.211111311 196199
Peter Li and Shing Tung Yau. On the parabolic kernel of the Schrödinger operator. Acta Mathemat- ica, 156(none):153 201, 1986. doi: 10.1007/BBO22399203. URL httph:://doooe http:/doo..oog//// JDC7/BF0239920)
Christian Léonard. From the Schrödinger problem the Kantorovich problem. Jour- nal of Functional Analysis, 262(4):1879-1920, 2012. ISSN 0022-1236. doi: https://doi.
URL
articcle/ss1/50022233110000
Christian Léonard. A survey of the Schrödinger problem and some of its connections with optimal transport, 2013. URL https:://///.. /tttt::s://axxi...e/ www/abgg/.....ooooo...ee...e..ee 11008 0215.
Simone Di Marino and Augusto Gerolin. An optimal transport approach for the Schrödinger bridge problem and convergence of Sinkhorn algorithm. Journal of Scientific Computing, 85 (2):27, 2020. doi: 10.1007/s10015-0000002222 000-111-020-013300
URL https: / /ttp:///ci.iee htt////cccooor//////////// attps://ccc..orr/10110....
Gregory A MeIntyre and Kenneth J Hintz. Comparison of several maneuvering target tracking models. In Signal processing, sensor fusion, and target recognition VII, volume 3374, pp. 48-63. SPIE, 1998.
Atsushi Nitanda, Denny Wu, and Taiji Suzuki. Convex analysis of the mean field Langevin dynam- ics. In International Conference on Artificial Intelligence and Statistics, pp. 9741-9757. PMLR, 2022.
James R. Norris. Heat kernel asymptotics and the distance function in Lipschitz Riemannian manifolds. Acta Mathematica, 179(1):79 103, 1997. doi: 10.1007/BFF02392720. URL https: //t:://do.... atoii.org/1... 1007 1007//0022/000220000
B. Øksendal. Stochastic Differential Equations: An Introduction with Applications. Universitext. Springer Berlin Heidelberg, 2010. ISBN 9783642143946
URL
google.
com/books btt.oooksiidggg7AAAAa......
Sabriel and Marco Cuturi. Computational optimal transport, 2020. URL / org/abs/ / 803. 00567.
Rihao Qu, Xiuyuan Cheng, Esen Sefik, Jay S. Stanley III. Boris Landa, Francesco Strino, Sarah Platt, James Garritano, Ian D. Odell, Ronald Coifiman, Richard A. Flavell, Peggy Myung, and Yuval Kluger. Gene trajectory inference for single-cell data by optimal transport metries. Nature Biotechnology, Apr 2024. ISSN 1546-1696. doi: 10. 15810084411887--22-02228686 URL https: //doi. org..org/11 10. 1038 s 41587- 024 02186-3.
Seyed Hamid Rezatofighi, Anton Milan, Zhen Zhang, Qinfeng Shi, Anthony Dick, and Ian Reid. Joint probabilistic data association revisited. Proceedings of the IEEE international conference computer vision, pp. 3047-3055, 2015.
Wouter Saelens, Robrecht Cannoodt, Helena Todorov, and Yvan Saeys. A comparison of single-cell trajectory inference methods. Nature Biotechnology, 37(5):547-554, May 2019. ISSN 1546-1696. doi: 10. 10.10088441377019900711999 URL https: : wedii.or////e 1038////1////
051587-019-0077
13

---

Published as a conference paper at ICLR 2025
Filippo Santambrogio. Optimal Transport for Applied Mathematicianss: Calculus of Variations, PDEs, and Modeling. Progress Nonlinear Differential Equations and Their Applications. Birkhääser Basel, 2015. ISBN 978-3-319 20827-5. doi: https:: www. spr aggrnngrr.... nger. com gp http//bo//////oooo/oooo/ /00k//78333 778331 9208275.
URL
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 S. Lander. Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, org/10. htg/100..066f..ee111119.....0..0 URL https: /wwww//wwww.//////////////.//... stttps:/www.sciieneeeeeee www.sccienceiirect..... ctt.coo////a/////. cromscciecee/ articcle/ssi/50022377113333
2019. ISSN 0092-8674. doi: https://doi.
Yunyi Shen, Renato Berlinghieri, and Tamara Broderick. Multi-marginal Schrödinger bridges with iterative reference refinement, 2024. URL
org/ abss 2408.06277.
Yue Song, Andy Keller, Nicu Sebe, and Max Welling, Latent traversals in generative models as potential flows. In Proceedings of the 40th International Conference on Machine Learning, ICML'23, 2023a.
Yue Song, T. Anderson Keller, Nicu Sebe, and Max Welling. Flow factorized representation learn- ing. In sevrthy Conference on Neural Information Processing Systems, 2023b. URL https: bttts::/epenzeee cpenreview.ewwece///ee/eeroeee.m net wwt/tnnumiiddnnnnnaa... htt/://mmmmnnyySnSAAA------
Taiji Suzuki, Denny Wu, and Atsushi Nitanda. Mean-field Langevin dynamics: Time-space dis- cretization, stochastic gradient, and variance reduction. In Thirth-seventh Conference on Neu- ral Information Processing Systems, 2023. URL https: WKKLttts////peeeeee/eee wpwprevieeeeeeeeeeo//ee//e/oeewe net wee//forumme forum?id= 9STYRIVx6u.
Hiroshi Tanaka. Stochastic differential equations with reflecting boundary condition in convex gions. Hiroshima Mathematical Journal, 9(I):163 - 177, 1979. doi: URL httpt: http://......11.....2... http:///...3311//////22223333333
Rachel Thomson and Janet Holland. Hindsight, foresight and insight: The challenges of longitudinal qualitative research. International Journal of Social Research Methodology, 6(3):233-244, 2003.
Alexander Tong, Jessie Huang, Guy Wolf, David van Dijk, and Smita Krishnaswamy. TrajectoryNett A dynamic optimal transport network for modeling cellular dynamics. In Proceedings of the 37th International Conference Machine Learning, 2020.
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. In Advances in Neural Information Processing Systems, 2021. URL https:://///.. httthtts////peeeeeiiiiee /pppnrevie....
net/forum? 1d==9TYG0j-wtG.
Vatiwutipong and N. Phewchean. Alternative way derive the distribution of the multivariate Omstein-Uhlenbeek process. Advances in Difference Equations, 2019.
Elias Ventre, Aden Forrow, Nitya Gadhiwala, Parijat Chakraborty, Omer Angel, and Geoffrey Schiebinger. Trajectory inference for branching SDE model of cell differentiation, 2024. URİ.
https://arxiv org/ abcp//ab//22/ 2307 07687.
Caleb Weinreb, Samuel Wolock, Betsabeh K. Tusi, Merav Socolovsky, and Allon M. Klein. Fun- damental limits on dynamic inference from single-cell snapshots. Proceedings of the National Academy of Science, 115(10):2467-E2277 March 2018. doi: 10.1073/pnas. 17147233
Shen Zeng, Steffen Waldherr, Christian Ebenbauer. and Frank Allgöwer. Ensemble observability of linear systems. IEEE Transactions on Automatic Control, 61, 07 2015, doi: 10.9109/T...2015. 2463631.
Stephen Zhang, Anton Afanassiev, Laura Greenstreet, Tetsuya Matsumoto, and Geoffrey Schiebinger. Optimal transport analysis reveals trajectories in steady-state systems. PLOS Computational Biology, 17(12):1-29, 12 2021. doi: 10.137//joumal..pbii pom.0 1009466. URL
https://doi.. http///i.oogg/.... 1371 /journal ppbii 1009466.
14

---

Published as a conference paper at ICLR 2025
Stephen Zhang, Gilles Mordant, Tetsuya Matsumoto, and Geoffrey Schiebinger. Manifold learning with sparse regularised optimal transport, 2023. URL https://arx....or//aaaat http://.rrrraa/////////// absr/077 09816.
Zhenyi Zhang, Tiejun Li, and Peijie Zhou. Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport. In The Thirteenth International Conference on Learn- ing Representations, 2025. URL https: / http///pppeeeeenneeee/e net forum 20
15

---

Published as a conference paper at ICLR 2025
A ENTROPIC OPTIMAL TRANSPORT
We provide a brief exposition to entropic OT. We refer the reader to (Cuturi, 2013; Peyré & Cuturi, 2020) for a more thorough introduction.
Let X, X be arbitrary spaces, c X x y + R be a cost function, and μ4, v be probability measures on X.Y, respectively. The entropic OT problem is
/
(p, = inf
xell(p..) Xxy (Lz.)) dy) + cH(TT @v),
where [1,,,)) v) is the set of all probability measures A' Y with marginals μ X and V H the relative entropy, and e is the regularization parameter. By standard duality theory, this is equivalent the following problem
I pdu++ f (idu+e +e E0111111 (1- (--fetetotows-deee (ooo)dd((=dddd du du(z)du(w))
SEL PELAM) (E Te v))
max
which admits a unique solution up translation (sp + 1G, 1)) for E R. Furthermore, the functions (p, 1) satisfy the following conditions:
[v(x)) -c log exp( - c(., c(r,y)))duu) vhyy = - log
- c(2r, b)) ()(r).
In the discrete (empirical measure) setting, these potentials give motivation for the Sinkhorn al- gorithm, which we describe Algorithm 2. Here, @ taken to be element-wise multiplication. (Altschuler et al., 2017) shows that the entropic OT problem can be solved in approximately linear time.
Algorithm 2 Sinkhorn
Require: Probability measures μl, v, matrix C, regularization parameter c, number of iterations NV
1: 6
2: K t exp(-C/e)
3: for 1ri=1..... 1...., do
4:
(b0)) v @ KJ (i-1)
5:
(p6)) t μ (i----)
6: end for
7: Output transport plan (X)) K diag (1)))
B ENSEMBLE OBSERVABILITY FOR LINEAR SYSTEMS
Recall that in classical observability (Gajic et al., 1996), the goal is to recover the dynamics of a single particle, while here we want to recover the dynamics of a probability distribution. The notion of ensemble observability introduced in (Zeng al., 2015) tackles this problem. Consider the non-stochastic model with linear E(X) = AzX + Bs:
dx, = - (A=X, X, + Bz) dt
(12)
with initial condition Po and linear observations Y, = g(X.) = Cg.,, , + Dy. For shorthand, we denote this system (Az, Be, Dg).
The following the definition of ensemble observability as introduced (Zeng et al., 2015): it does not consider stochasticity.
Definition (Ensemble observability (Zeng et al., 2015, Def. 1)). The linear system (12) is ensem- ble observable if given marginals 9rP, Y for all E [0, 1), the marginals P of X₁ uniquely determined for all t E [0, 1].
We can consider Definition in Definition 1, we exactly recover ensemble observability. (Zeng et al., 2015) showed that classical observability is necessary condition for ensemble observability, and provided several sufficient conditions as well. For a random variable X, we denote sx to be its characteristic function. We assume the following the initial distribution X
an extension of ensemble observability. In particular, if we let
16

---

Published as a conference paper at ICLR 2025
Assumption B.1. Let Xg be such that (oxo(s)) is real real-analytic for all non-zero R"
This assumption not very strong and most "nice" distributions satisfy it, e.g. if they admit have a density with respect to the Lebesgue measure. Recall that by (Gajic al., 1996, Thm. 5.2). classical observability holds if and only the observability matrix O = [Cz, A= C= has rank n. (Zeng et al., 2015) provides two useful sufficient conditions for ensemble observability of such systems.
Proposition B.2 (Zeng et al., 2015, Thm. 8)). Under Assumption B.I. if [As, Dg) is observable and rank Cg
1, then (Ag, Bz, Cg. Dg) is ensemble observable.
Corollary B.3 ((Zeng et al., 2015, Cor. 8)). Under Assumption B.1, if X R and
(2, Cg. Dg) is observable, then (As, Be, Cg. Dg) is ensemble observable.
We now show that these two conditions can be carried over to stochastic systems, i.e., those with T > O. Consider adding stochasticity to (12) with the model
d.,, = (A= + Bz) dt + vTdB,
(13)
with initial condition Po. {B,} } is an -aauue Brownian motion, and observations Y3 = C. X +
Dy- We have the following result:
Corollary B.4. Suppose (Az, Bz, Cg. Dg) is ensemble observable (with T = 0). Then the system (13) with known v >> O is ensemble observable.
Proof. It is easy to see via direct calculation, e.g. using an integrating factor, that the solution to (13)
-o-o (f ds Bs V&+vV] e dB,,
x; -
Art Ao
(14)
where we use matrix exponentials. Using arguments similar to those in (Vatiwutipong & Phewchean, 2019), can characterize the covariance:
Σ, = Cov [Xi] = S'" -11
We also know that
-(foo-dttt-) ds ds)
Pe B[X,] e Ast A
Bz.
Then as the first two terms on the right-hand side of (14) have zero variance and the Itô integral of deterministic integrand is normally distributed with mean zero, we know that (14) is distributed as
X₂ ~ ((e. M.).
As we know т, and the corresponding observability matrix to the system has full-rank, see that the pushforward (to observation space) of every term in (14) is also fully recoverable as well. Note that it is possible to deconvolve, e.g., using the fact that deconvolution is equivalent to division in the Fourier domain, the known Gaussian noise from X,, and hence the system is ensemble observable. This concludes the proof.
Next, we extend Corollary B.3 to independent processes and apply Corollary B.4.
Proposition B.5. Let X' Xi x XXy, where each X; R22 Suppose (Ag, Bz, Cy. Dg.) is ob- servable for each i E [n). Further, suppose the initial condition joint distribution for the Xo,; satis- fies Assumption B.I with [No,]]:: [Xo,3]2 conditionally independent conditioned on [No,]2, (Xos)1, for each + i. If noise parameter > 0 is known and {Be} is an X-valued Brownian motion, then the system
A
Bz
g1
D
B=
C
D
Thhese are not the only concrete conditions provided therein. Furthermore, more general sufficient condi- tion provided which is possible check numerically.
17

---

Published as a conference paper at ICLR 2025
is ensemble observable. Furthermore, the system remains ensemble observable under (known) per- mutations.
Proof. This follows from simple application of Corollaries B.3 and B.4.
B. EXAMPLE: "CONSTANT VELOCITY" MODEL
The two-dimensional "constant velocity" model, so named because the velocity would be constant there were no process noise (T 0), uses a state vector
X (z, a. 1) E R
where here (a, y) are two-dimensional positional coordinates, and (æ, () is the current two- dimensional velocity. The "constant velocity" dynamics model uses [(X) Agx where
0
0
A=
0
0
which is very simple matrix simply implying that the rate of change of X1 x is given by the current state vector X3 = and similarly for y.
Having defined the dynamics, in this model, only the positions are observed. In other words, the observations g(X) = //2 X2, i.e. C= = [/2, 02x2] and D= =
Note that this experimental setting satisfies Proposition B.5 as the x and dynamics are independent.
Hence, is sufficient to check if the system observable. Here = [1,0] 0] and = [0, 1;0, 0J. so the observability matrix for each of these subsystems becomes
-
which is the identity and thus full rank. By observability theory, the system is classically observable, and by the results above, ensemble observable as well.
Note that ensemble observability can be extended non-zero v this setting. For instance, in the "constant velocity" model of the main text, V [0; Φ for constant but unknown z will serve simply as drift term on the mean of the hidden velocity state. Since without this drift the mean of the velocity is constant, this drift will be identifiable and the system will be ensemble observable.
Finally, as a brief remark, note that in Definition 1, we require V to be restricted to class of functions Cup as otherwise the SDE (1) may fail to satisfy classical observability. Further exploration of Csp classes is left to future work.
C PROOFS FOR CONSISTENCY
We make this section as self-contained as possible, although we suppress some of the longer details when they are very similar to certain corresponding results in (Lavenant et al., 2024). Whenever do, we point the fuller arguments in (Lavenant et al., 2024). We will use 4 to denote a-wide heat kernel, which is equivalent to the Gaussian convolution Ng. The main result in this section is the following:
Theorem C.1 (Consistency, (formal version of Thm. 3.1)). Let P be the law of the SDE given in (1), restated below:
dX, -E(( di (V X,) + VFdB,
with initial condition Po E P(X) such that (Poovoll) < +oo. Assume we have the following:
X + y is a smooth, measurable, bounded, time invariant function, and (g. E,Cx) is (i) Co Co-ensemble observable.
(ii For every T 1, we have sequence of ordered observation times (t berween and 1, and {1]
1T Ji=1 becomes dense in [0. 1] T +00.
18

---

Published as a conference paper at ICLR 2025
(ii) For each T and each E [7), we have Nf 2 random variables (V2) N. and distributed according to 9:Pet
which are i...d
(iv) The variables Y.J and Y are sampled independently from their respective distributions except when (T, = , (T',''J'
Consider the functional (5), restated below:
F(R): == Fit " . .: (
and let RTA E P(2) be its unique minimizer:
R.... arg min F(R).
REP(I)
Then, we have the weak convergence
lim
Jim RTAa)) =
o -- 770000
almost surely.
Proof. We use Theorem C.5 to take the limit + By the law of large numbers and the weak convergence assumption, have Pt Pt, almost surely. Define R be the limit of R... as - +00. By Theorem C.5, is the unique minimizer of
R + FA (R) = 1 / DF Pi) dt +
X J
By definition of the data- tata-fitting term, the functional Gx,a in Theorem C.17 differs from Fi, only by constant. We see that
Gx (R) (R. H(P,)
so R must also be the unique minimizer for Gx.-- Finally, we use Theorem C.17 take the limit of RA A. as a and A 0. This concludes the proof.
C.1 VARIATIONAL. CHARACTERIZZTTON THE SDE
We recall some previously introduced preliminaries and notation. Ω = C( 1] X) is the set of X'-valued paths, (Xe)c_(o.1] our canonical process, and F is the Borel d-algebra generated by the random variables X, for S such that (Fohejooa) is a filtration. We use the notation | (..) E for the quadratic variation of process (similarly use this notation for cross- variation).
For the Girsanov transforms to be martingales, we have the following mild technical assumption. Assumption C.2 (Novikov conditions on E). Assume that the following Novikov conditions hold:
1 (/' | ds
E
exp
+00
2
and
1
E
cxp
E+ v ds
+00.
2
Also, assume that there exists C to0 such that f' 1 ds S C and So | VV ds SC
This last condition is so that we can apply Girsanov's on manifolds, e.g. (Hsu, 2002, Thm. 8.1.2).
Proposition C.3 (analogous to (Lavenant et al., 2024, Prop. 2.11)). Let P be the law of the SDE in (1). Then the Cadon-Nikodhe derivative of P with respect to WE.r given W -aee. by
dP
dP
T (0,
X1)
(X)
(Xo) cxp
dWET
dvol
(15)
( 0,4
1 2 ( vv) + + _AV)
(s. X.) do do)
exp
2
19

---

Published as a conference paper at ICLR 2025
To prove this proposition, do not martingale characterization as in (Lavenant et al., 2024), but directly use the Girsanov theorem (which by our assumption on E, can be applied on manifolds) and the Itō formula.
Proof. By the chain rule, we have
dP
dPo
dP
dW
dWE. (X)
(Xo)
dvol
dW
d'W..r'
The first term follows from an averaging argument identical to that of (Lavenant al., 2024, Prop.
2. 2.1 For the second term, recall that P is the measure induced by the process dXx, = -(3 +
V dt + v dB, and W is the measure induced by the process dY, dBd We have
dP
1 VF Jo
/ ( VV) dB,
27 / E+ (VP((s, ds
1
exp
dWr
/'
1
/
exp
VW X3) dB,
| (s, Xz) ds
2T
V= / [ss Xa) dB.
1
exp
(EL + 2E, (8, ds
2T 0
(o) (11, + {
0,4 2
v4
2 -Δッ
(s, ds)
0
V / E(s, (ss.x))23
exp
(E3 + 2(E, VV) X,) ds
2T Jo
E(s, X.) dB )
exp
(1, + VT 1 exp ( - /
(0
Xo) (o)
(16)
T Av 2
(8, X.) ds
where the first line follows from the Girsanov theorem, (Øksendal, 2010, Thm. 8.6.6) and the third line follows from Itô's formula, (Øksendal, 2010, Thm. 4.2.1). Letting W be the measure induced by the process dZ,-- = -E,, Zi) dt + /T dB, we have
dW
1 2r Jo
[ EdB,
dW,T exp
E ds
(17)
VT
by Girsanov. Combining (16) and (17) yields (15). The claim follows.
The next result is the variational characteristic of the SDE.
Theorem C.4 (analogous to (Lavenant et al., 2024, Thm. 2.1)). Suppose (g, E, Cv) is Co-ensemble observable. Let E [0, x X X be a smooth function and [0,1] X R be smooth potential. Let P be the law of the SDE
dx, Et,, X2) di (V, Xx) di + /r dB,
with initial condition Po P(X) such that H(Polvol) toc. If P(L) is such that R. g:P. for all t E [0, 1), we have
H(P S H(R(WE))
with equality and only R
The argument follows that of (Lavenant al., 2024) with ensemble observable assumption and reference measure. Here, the proof is the same, but now we use the fact that our reference measure "cancels out" the stochastic integral, e.g. see Proposition C.3.
Proof. Let P be the law of the solution of (1) and suppose R E P(2)) is another path measure such that H (R(WE.r)) WE.r) +00. L,tp,r p, E (22, WE.T) denote the Radon-Nikodym derivative P. R with respect to WE.r, respectively. By strict convexity of x + x log we have
llgg p log (1 log p).
20

---

Published as a conference paper at ICLR 2025
5a--aammst everywhere, with equality if and only if p. Integrating with respect to W
see
-((WSS))-(((WS-)) 2 ER + log Ep logp).
(18)
Using Proposition C.3, we have
dPo dvol
(o, - (1)
En 1 log
ER
log
(Xo)
7
V (E, vw) + T A X,) ds 2
1
0
2
By definition of ensemble observabilityy, if g. P. for all E [0, 1]. this implies R. P, for all & E [0, 1]. Because this expression only depends on the temporal marginals of R as the Radon- Nikodym derivative of P with respect to Wa.r does contain a stochastic integral, the right-hand side of (18) vanishes if gy 9;P, for all t E [0, 1]. This concludes the proof.
C.2 THE MAIN TECHNICAL RESULT: THEOREM C.5
Theorem C.5 (analogous to (Lavenant et al., 2024, Thm. 2.7)). Fix and assume we have the following:
(i) For every T E N, we have sequence of ordered observation times {t. YT a sequence of data o (a collection of probability measures on Y); and a sequence of non-negative weights
{w] (TTV Ti=1
(ii) There exists a -v(Y)-valued continuous curve p E C((0, 1] P())) such that the following weak convergence holds: for all continuous functions : [0. 1] x R.
lim Σ w. / all? , (dx) 1'' / alt, alt,x)p(dz)) dt.
T
1=1
For each T, let R E P(2) be the unique minimizer of
R PR)) = - + w ( R.y
(19)
Then as > +00, the sequence {RT} converges weakly on P((2) to the unique minimizer of
= 7 x //// D....)..P..
Before proving this theorem, we state the following result that is immediate from the non non-negativity of our data-fitting term.
Fact C.6 (Non-negativity) With the assumptions of Theorem C.5, the functionals Fr and F are bounded from below by
Proof of Theorem C.5. The argument follows that of (Lavenant et al., 2024). Let R be minimizer of and R" be the minimizer of Fr. By optimality of the minimizers, we have R, and Go = RT. Using Proposition C.15, we can find a sequence R that converges weakly to R as T +00 such that
F(R) 2 sup Fr(RR) 2 lim sup min Fr. lim (T(R"). T-++00
T-4400
T-4400 P(D)
In particular, the sequence is bounded, which by Fact C.6, implics the sequence H H (RTWWE.)) is bounded. Then from the compactness of the sublevel sets of the entropy, we have limit point R of sequence {R²). Using the optimality of R and Proposition C. we see
F(R)= < F(GoR) < inf Fr (RT).
T-++o
21

---

Published as a conference paper at ICLR 2025
Thus, have equalities everywhere,
F(R) sup Fr(RR)
lim Fr(R"). T-++00
T-++000
Then we see
Fr FT(R)- FT FT(
min Fr P(R)
converges to 0 as 7 +00. By (Lavenante et al., 2024, Lem. B.3), relative entropy is 1-convex with respect to the total variation, i.c. if p, q. are three probability measures,
+ q 2
1 2
(((r)) - 221le alliv.
H
Since the data-fitting term is also convex, the full objective (-(()) is 1-convex with respect to the total variation. By a classic strong convexity argument, since (r converges to the minimum value (n) Fr achieved at R |R R R must also converge to 0 as T + +oc. Recall that TV convergence stronger than weak convergence. Then, using the weak convergence R to R. we see that R converges weakly to R as T > +00. This concludes the proof.
D
The remainder of Appendix C.2 dedicated towards proving Theorem C.5.
C.2.1 HEAT FLOW AND REGULARIIZTION THE MARGINALS
Recall that we use i3 to denote the heat flow with width We use the heat flow regularize the marginals. First, we have the following result showing that the density of T R. is continuous jointly in and æ.
Proposition C.7 (analogous to (Lavenant al., 2024, Prop. 2.12)). Let 0. There exists constant C depending only on and E for which the following hold:
(i) For each R E P(L)). its heat flow regularization &.R, has density p(5) (t, (with respect to the volume measure) that satisfies for all t E 1),1) x E X, we have
( (1, C.
(ii) For all 11 [0, and æT, 22 X, we have
1o (1, - / (12, c (Vrv ( C 12| dx(
Proof. The first estimate is directly from (Lavenant et al., 2024). The second estimate follows from (Lavenant et al., 2024) and the following proposition.
Proposition C.8 (analogous to (Lavenant al., 2024, Lem. 2.13)). There exists constant C depending only on X and E\\ such that for each R E P(?),
ER (X,,, Xx,) S C (H(R) + + +C+Co")
0 - 12/-
Proof. The argument follows that of (Lavenant et al., 2024). For any 17 > 0, using the dual repre- sentation of entropy with the function X + 1dxx(x,,) Xt,). we have
EER|dx )1 & (R| + log Ewa.. (exp(ndx (X,,, )1,)
Using an upper bound the heat kernel from (Li & Yau, 1986, Cor. 3.1) and that || t00, we have the following bound on the transition probability for WE.r;
Cel (rt)d/2 exp
Crt dx(x,v)²
(~ v,
Crt
Then the remainder the argument of the proof of (Lavenant et al., 2024, Prop. 2.13) yields the desired result.
n
22

---

Published as a conference paper at ICLR 2025
C.2.2 HEAT FLOW AND ENTROPY ON THE SPACE OF PATHS
We introduce an auxiliary variational problem in which all the temporal marginals are fixed.
Definition 3. Let C([0, P(X)) be P(X))-valedd continuous curve with respect to the weak topology. Define the problem Ap(p) to be
(p) = inf (TH(R)WE") I 10, = 9tPt)-
REP(S)
We use the convention that (p) = +oo if the above problem has no admissible competitor.
Using a dual representation of A. we can use PDE theory to solve this problem. First, we give martingale characterization of class of stochastic processes:
Proposition C.9. Suppose WE. is the law of the SDE dX, -Edt V dB, with arbitrary initial distribution. Let yp [0, x X 7 R be smooth function. Then, the process whose value at t [0, 1] is given by
(o) exp (p(t.))) -((,x)
L/" Oop [pp2= -E, Vg) + 2 As
(s, (s, (s) ds) ds))
(20)
an Fi-martingale under Wa.r
Proof. By Assumption C.2 on E, the process
(E, VΦ) + I Mf Xi) (0, - [ (E,
2 _ (8, Xo ds
(21)
an Fi-martinglee under Wa.r by (Øksendal, 2010, Thm 8.3.1). We calculate the quadratic varia- tion (M+), similar to (Hsu, 2008, Prop. 1.3.1). First, we have
x,) +00 Xo)² + M
[-
[s, ds
J
I'
=o) + M
(p,, Vp) AAp"l(s, ds.
J
Using Ité's formula, we have
X,) (o) +2 L ((s, Xodp (+)++((9)) (M²)
L' p(s, p(s,X.)dMM + I x(s. ve) AApl(s, X.)dds + (M*).
= (o) +2
Equating the bounded variation parts, we see
1 (M
//
2-p Vp) + Ap² + 2 Vip) rومl (s, ds
1
/" [Ap² pApl (s, Xs
2
// Vy(s, X2) ds.
Then, (20) is the exponential martingale of MY.
Here, recall that Assumption C.2 ensures that (21) is martingale. Otherwise, it is only local mar- tingale and we would need to check the L convergence of the stopped process with an increasing sequence of stopping times that goes +oo. This is a standard argument in stochastic calculus, e.g. see (Øksendal, 2010).
Now give the dual representation mentioned above.
23

---

Published as a conference paper at ICLR 2025
Proposition C.10 (analogous to (Lavenant et al., 2024, Prop. 2.15)). Let C(O, P(X)) be P(Z'--valud continuous curve. We have
(p) (polvol)))
+ sup { {-/ ((0, x00)))))())) - S. Jx (0:pp + /Vpp - (E, (5) + T ( pe(dtr di )
where the supremum is taken over all ip C² 1] x X) such that (p1, = 0.
Proof. The argument follows that of (Lavenant et al., 2024). From the duality result of (Arnaudon al., 2020, Prop. 2.3), we have
Ar (p) - (Poollllll)) + sup (L (ppe do - /x [og log. exp (L. v()) po po(dann]
where WE, is the measure such that W = 6g and the supremum is taken over ab E C( X)). Here, the characterization holds because the argument in (Arnaudon et al., 2020) only requires the reference measure to be uniform at all marginals.
Let - = ) (0 / (E, + App) for some smooth y satisfying the terminal condition (11 x) = 0. By Proposition C.9, we see
((0,XX))
=fo ((l, 1.
Ewar exp
T
As Xo 0 under W.. have
lx log log exp (F(((,
wbl,, (.(1, xx)
d }]
1 [Toge (p(0.2) po(do)
1 / x (0(0, xx) Paldx).
The remaining argument of (Lavenant et al., 2024) follows through.
The main idea is that there is contraction of Ar under the heat flow, which we can think of space-time counterpart of the contraction of entropy under the heat flow.
Proposition C.11 (analogous to (Lavenant et al., 2024, Prop. 2.16)). Let E C((0, 1] P(X)) be -vvll- (X)) valuedd continuous curve and for 2 O, define the new curve plg) : t + DoPr. Let K be
lower hound on the Ricci curvature of the manifold X. Then, for any 2 0, we have
(p(o)) S -2K
2K - Ap)..
Proof. Consider the dual formulation in Proposition C.10. y with boundary condition -)1,-) O, then by the self-adjointness property of the heat semigroup, we have
: (0, 1] x - a C2 function
f ((0, )p0 (4) + 1. (aip + + - vp) + - Ap) p. dt
- - -10000000)000 )po + L.' (8.0.0 + - V 2,, Vp)
+ T A 1
p dt,
where 0,0.p = 0 by Schwarz's theorem and $.Ap Ab, by simple calculation. By properties of the carré du champ operator (Bakry et al., 2013, Cor. 3.3.19) and expanding out the inner product, we see that (#.E, Vb.4)) S 2k (E, V.p). Thus, letting é e2K* x and
= e2Ks we have
-f (Oupt 1 ( v (E, Vip) [
p(0.))pb
2 _ pi dt
T
S - -2Ks 1/ ((00))pm -f" [.] 1 ( } V - (3. VG)- + _ _
pe at
2
S <e- (A+(p) (Pu||vol))
24

---

Published as a conference paper at ICLR 2025
where the last inequality is due to Proposition C.10. Taking supremum over ye, that (p Se 2K* Ar(() + [H( (P.ooolool - e 2K (po /vol]]
By (Lavenant et al., 2024, Eq. B.3), the second term in the right-hand side is always non-positive, so the claim follows.
Next, we define the regularizing operator 9, that acts at the level of laws the space of paths. Definition 4, For each R E P(()) with H (R(WE.)) < +oo and for each $ 2 0. define
G.(R) arg (in(H(R/WS")) Vt E [o, (0.1). g: 9: g.
REP(O)
That is, among all probability distributions on the space of paths whose marginals in hidden space coincide with t ++ 94 9, the measure G, (R) P(2)) is the one with the smallest entropy.
Note that (R) is well-defined because thanks to Proposition C.11, A, A,(( I e-2K* Ar Ar(,))
(R < +oc, so the minimization problem has admissible solu- tions. Since sublevel sets of entropy are compact, there exists a minimizer, and from strict convexity of the entropy functional, it unique. Now note that
Ar (R.):) = HQ..(R (R(T))
This gives us the following result. Proposition C.12 (analogous (Lavenant et al., 2024, Prop. 2.18)). For each E P(22) such that (R|WE.7) < +oc, we have the following:
(i) For any s 2 O, (R) S e 2K* -22 ((Go(R)) <e
(ii) G. (R) converges to (R) weakly as s + 0+.
Proof. The argument follows that of (Lavenant et al., 2024). The first property is rewriting of Proposition C.11 together with the definition of G. and Ar. The second property follows from our observability assumption and analogous argument that of the proof of (Lavenant al., 2024, Prop. 2.18). Consider the following sequential characterization. Let {s JeE be sequence with 8n n + +0c. By the contraction estimate (i) and that the Ricci curvature is bounded from below, we know that (G,,, R/WE.)) is uniformly bounded n. Let R. be any limit point of G.,, R. Notice that this limit point exists due to the compactness of the sublevel sets of H-| (-E)..
We show that R. R. by standard analytic argument. We consider subsequence (which we do not relabel) G. R. that converges to R. - +oo. The marginals of R. agree with those of R as we easily see that the marginals of G.,, R are { Re ReJuuj. and $ f > f in L vol) as
8n - 0. Then, using the lower semi continuity of entropy, the definition of G.,, and the contraction estimate for A, we have
< lim inf H9G., W. +=+000
lim inf A, (4. R()c)l
++++00
liminff
-2K A, (RR))) A, ((R)).)
This shows that R R, which concludes the proof.
C.2.3 THE FITAING TERM
We recall the definition of the data-fittin fata-fitting term here:
DF (gp R. a. = S, log Sr exp (- [19(x)) - y/13
(-looo)))
=112)) d (x) dpf(())
20
HoTT 19g & No) + (((T) + +C,
where N, is the Gaussian kernel. First, we have the following result, which is immediate from properties of entropy.
25

---

Published as a conference paper at ICLR 2025
Proposition C.13. The function - (F p) is convex and lower semi continuous P(X).
We will require a quantitative control on the effect of the heat flow on the data- data-fitting term.
Proposition C.14 (analogous to (Lavenant et al., 2024, Prop. 2.22)). Assume X + y is measure preserving. Let p,r P(X). There exists a constamt C depending only X, 9, and such that for every > 0,
(g Dr,, 9:p) - (F + S-C..
Above, for simplification of the notation, we pushforward both parameters of the data-fitting term by This makes the argument below much cleaner.
Proof. The argument follows that of (Lavenant et al., 2024), but is much simpler due to our different data-fitting term. In particular, we do not need bound on the Fisher information. By abuse of notation, denote r E (X,, vol) the density of r with respect to vol. Denote r(s, :) to be the density of o or with respect It satisfies the heat equation
Ör Ōs
= Ar.
Then, we have
ds J - log // exp
[1((x) - g(1)112)
DF DF
r(s, x))oo(
ply)vol(dy)
ds
202
/ log
[1g(ar)
{(5)112
ly
exp
20²
Os
rs, 2s)lo))
p(y)vol(dy)
I log exp
[19(x)) -g(u))13
Ar(s, v)vol(de)
ply)vol(dy)
20²
C / plu)vol(dy) SC
where the inequality follows from properties of the Gaussian integral and the fact that f TArss x))ol(dx) = 1. Integrating yields the desired result.
C.2.4 TWO RESULTS ON LIMITS OF FUNCTIONALS
We require two results of the functional Fr defined in (19). We use these for the r"convergence theory required in the proof of Theorem C.5.
Proposition C.15 (analogous (Lavenant al., 2024, Prop. 2.24)). Use the notation and assump- tions of Theorem C.5. Suppose R P(2)) with F(R) < too and G. = R. Then there exists a sequence R.T which converges weakly to R as T +00 and
limsup Fr [ < F(R).
T-+
Proof. The argument follows that of (Lavenant al., 2024). Let 0. Combining Proposition C. 14 for the data-fitting term and Proposition C.12 for the relative entropy on the space of paths, we that
1 L' DF DF pT) dt
F(G.(R)) H
X J
1 L' DF
S Te 2K -2K H W
so we have
lim F(G.R) S F(R).
8-40
Suppose that (2, ()), Ay) measure spaces with I Lehesgue measure. g is measure preserving if for every Borel B X, (A) A).
26

---

Published as a conference paper at ICLR 2025
Now - log[9,Rl, R/e is continuous function of t and x by Proposition C.7, the weak convergence of p" to p to write, for & > O,
T
DF 9. 17 P. - -1(' DF(o dt.
Σ DF
lim T T-+00
This implies for all O, we have 1im (T, F)).R), so it is sufficient to let R Go R for sequence {ST)T21 that decays to 0 sufficiently slowly T too. This concludes the proof.
Proposition C.16 (analogous (Lavenant al., 2024, Prop. 2.25)). Use the notation and assump- tions of Theorem C.5. For each T 1, let R² P(2)) and assume that it converges weakly to some R E (()) as T 7 oc. Then
R) < inf Fr (R").
T-+000
Proof. The argument follows that of (Lavenant et al., 2024). Assume that lirm inf (r("") too otherwise we are done. Then, up to subsequence (that we do not relabel), we have supy H (R" ((T" < too. Combining Proposition C.14 for the data-fitting term and Proposi-
tion C.12 for the relative entropy on the space of paths, we have
Fr(G.R") TH (W
Σ
w DF Ro
TE 2K H(R. ( + X
DF R
sc
Now we rewrite the above as
FT FT(RR)) FT(,,R") C(s),
where
T|e s/p H (R [ +
sc
A
upper bounded by quantity independent of I' and lim, C(s) = 0. For the data-fitting term, define the sequence of functions a. (t, x) to be
a (t, == log [f eex
(=1(())
[1((2) - x012
d R (2)
202
which is parametrized by T. Notice that from the definition of the data-fitting term, we have
T
T
of DF o Σ
Σ
T
a
x) (dx).
For fixed O. the family of functions a (f, indexed by T' is uniformly equicontinuous due to being continuous and Proposition C.7. Then there exists a subsequence (that we do not relabel) that converges uniformly on [0, y as oo to the function
d (t, x) = log f
11((z) x
d d&, R(())
cxp
20
/ exp
(1((z) 11
= - log
d(G.Rl((2)
20
Using this uniform convergence with the weak convergence of P Pt. we
lim
T
w DF (I
lim T T-++2
Σ wh
[ aT x)) (dx)
T +-+
L / (,(t, dt
L DF P.) dt.
27

---

Published as a conference paper at ICLR 2025
Using lower semi continuity of entropy, we have F(G,R) I lim iiminf--oo Fr,g, Thus, for each $ >> 0, we have
inf Fr F(G.R) C(s).
Finally, we use Proposition C.12 to take 8 - 0+ using the lower semi continuity of and the convergence of G.R to GoR when + +
C.3 T-CONVERGENCE, TAKING - 0, 0
Theorem C.17 (analogous to (Lavenant et al., 2024, Thm. 2.91). Let E P(Ω) with H(P/WSr) +oc. For each À > 0 and 0, les R be the minimizer of the functional
2 H(P:RR, * NNo) dt.
R + GA = -
Then, as h + -, O, the measure R^o converges the minimizer of R + H(R among all measures such that Rt g.P. for all & E [0, 1]. Furthermore, if P is the law of the SDE in (1), then R converges to P.
Proof. The argument follows that of (Lavenant et al., 2024). First, consider := G.P P(l)) as competitor in Gx Using the contraction estimate given by Proposition C.12, we have
min GA,a -Gxo(RA")) S TH S re Ko H(G.P)WE.)) P(O)
As > - DO and H (PIWE.)) by assumption, we see that Gx,o (RAa) uniformly bounded in and a. Thus, HR A. |WEr) is uniformly bounded as well. Due to (Lavenant et al., 2024, Prop. B.2), this implies that the family R^ belongs to a compact set in the weak topology. Let R be any limit point in the limit as A + 0, o - O. We only need to show that R = GoP. Note that
H HRR (I S Gx. G S T 2k H(G.P/W
By taking and using the lower semi continuity of entropy, H(((EEr) s H(G.P/WS.))
Now using Fatou's 's lemma and joint lower semi continuity of the entropy, we have
[ H(P.RR, + No) lim inf
['
H(P.P.|RA di
x-0,0000
lim inf x-
AsupGa.
0.
X
Thus, it follows that 9 R. = Pp for almost every 1. Therefore, by definition of Go. we have R P. This concludes the proof.
D REDUCED FORMULATION
D.1 PROOF OF THEOREM 3.2
We use the following result to prove Theorem 3.2. Here, the statement is identical to that of (Chizat et al., 2022), but we consider different reference measure.
Lemma D.1 (analogous to (Chizat et al., 2022, Prop. B.2)). There exists constant such
that, for any R P(2) and tf. ...., I collection of time instants, it holds
HR
H(R ......
F of (.....
T-1
RST R) HR, Je | (R,y &
HR W$R
28

---

Published as a conference paper at ICLR 2025
The first inequality (1) becomes equality if and only if
R(-) = / X
Wa, (-11,............. .T) dR 17 (T1,.... IT),
where
is the law of W conditioned on passing through 21,.... Fr times 6 ....../; ty respectively. In addition, the second inequality (*) becomes an equality if and only if R Markovian.
Proof. Using the fact that E divergence-free and that W has the Markov property, the proof from (Chizat et al., 2022) holds. We provide the full proof for completeness. The first inequality (t) and the equality case follows from the behavior of entropy with respect a Markov measure under conditioning, e.g. (Léonard, 2012, Eq. 11). In particular, we have
((R(WE)) HR ......4 of W of (F....eT)
+ (R(- 2T W (1/11 TT)) dR, (z.. (21....... IT),
where the second term vanishes if and only if the conditional distributions R(- FT) follow the law of W=, for R.T....f? almost every (#1...., ØT). The second inequality (*) follows from (Benamou et al., 2019, Lem. 3.4), which states
T-1 H
T-1 H
H W
W G
Σ
( E,
with equality and only if R. ..... of is Markovian. As in (Chizat et al., 2022), we reorganize the terms
Without loss of generality, assume that R,y are absolutely continuous with density dR,y' (x) == (x()) and let Vx be the Lebesgue volume of X. Since W is the uniform measure on X' for every ff we have
"f"
H(R
H(R (T) + log Vx.
Letting 7. == 7( 1 we also have
1 w E QTT (dx, dy) = - Pr. (z,y) dx dy.
Vx
Thus, we see that for any u E P(X) with finite differential entropy and y E I(u, v), we have
Vx dx @ dy pT,
dn
H(> W
log
dy(l,
(7
d7 du du
log Vx + log
dr
1 did v) dx dy
log Vx H(r (M v)) H(H) + H()).
where the last line follows from (Marino Gerolin, 2020, Lem. 1.6). Now using the fact that R E 1l(RTr R,y we have
E= log Vx + Σ HRRggg P (R. R +) Σ H(Rg))
which proves the formula.
Theorem D.2 (Thm. 3.2, restated). Let Fit P(Y)T - R be any function and let E be bounded and divergence-free.
(i) If F admits a minimizer R then (Rig.... ..... R is a minimizer for F.
29

---

Published as a conference paper at ICLR 2025
(ii) F admits minimizer μ P(XY". then minimizer R for F is built
R²(-) = / 1xz W (-I1,......,.
(f....... (II.....,
where WE, (1.,,....... AT) is the law of W conditioned on passing through ..,.... FT at times f ...., by respectively and RuT.... ....G. is the composition of the optimal transport plans V/s that minimize T,,=(u 11 (x(++))), for [T-1] 1].
Proof. The proof from (Chizat et al., 2022) holds using our transition probability densities and OT plans. We provide it for completeness. First, note that minimizer Rs E P(()) of F(R) =
Fit(Q Qie 9ey ( W2.r) of the form in Lemma D.1. Let f( == R be its marginals
and 7000 = R o..
, ; which clearly satisfies 700) E 11( 11+1)). Using C' == log Vx, we see
that
F(R*) = Fit(ouu)
9:11 (T) +
Pr. (u @ () +TH((H) +C H(7 (1
Fit(gent () ................................ 9:p4 (T)
T T (H ,litt)) + (H( +C,
where the inequality becomes an equality if and only if R (T, = = is optimal in the definition of T (1)) p ((+1). The claim follows.
E MEAN-FIELD LANGEVIN DYNAMICS
Recall that the MFL. dynamics is defined as the solution of (11), which we restate below:
dx (8) = VV (xx! ds + 2 d + d
Law (x,()) = 140
Law (X.").
i [T].
where d (1) is the boundary reflection the sense of the Skorokhod problem. The family of laws {2 of this stochastic process are characterized by the following system PDEs:
0.1 = (VV() (m)) + ! (0)
(22)
which are coupled via the quantity V ( [p.]. The link between (11) and (22) follows from the itô-Tanaka formula, see e.g. (Javanmard et al., 2020, Lem. C.3). This a multi-species PDE where
each of the species (9) attempts to minimize At, Fit> """"""""""""")"))"">" + via drift-diffusion dynamics, and it connected to (1-1) and j (++1) via Schrödinger bridges.
E. PROPERTIES OF G AND F
We describe some properties of functions G (7) and F (8).
Recall that the first-variation of G P(X)T + R at is the unique (up to an additive constant) function Vin E C(X)" such that for all v E P(X)T.
T
1 +)) - = Σ vo (p]()) d(v (x).
lim c+0
Proposition E.1 (analogous to (Chizat et al., 2022, Prop. 3.2)). in each of its inputs (but not jointly), weakly continuous and its first- is given for P(X)" and E [T] by
The function G is convex separately
V [1]
5Fit 8p
[m4]
P1.111
(T
t?
1
17
30

---

Published as a conference paper at ICLR 2025
and
Fiit (p(()
At /
No(g(z) 1) (No * (y)
dō(y),
where (pt.) v1.s) E C (X) are the Schrödinger potentials for T = (6))). with the con- vention that the corresponding term vanishes when involves v1.0 or PT,T+1- The function F jointly convex and admits unique minimizer μt, which has an absolutely continuous density (again
denoted by μ² characterized by
(n* xe- (1 for E [T].
Here, the Schrödinger potentials are classically L by standard (entropic) OT theory, but we can extend them to C functions, as discussed in Chizat et al. (2022).
Proof. The argument similar to that of (Chizat al., 2022). The properties G and its first- variation are clear. In particular, the first-variation of T follows from the fact that g is smooth and (Santambrogio. 2015, Prop. 7.17), and the first-variation of Fit follows by direct calculation. The convexity of G' follows from the convexity of Tr,= and the fact that the pushforward of g is linear. The joint convexity of F. its unique minimizer, and the characterization of the minimizer follow directly from the argument in the proof of (Chizat et al., 2022, Prop. 3.2).
E.2 NOISY PARTICLE GRADIENT DESCENT
Let m N be the number of particles used in the discretization for each of the time marginals 14()) For computation, we approximate the MFL dynamics by running noisy gradient descent on the function Gm ( + R defined as Gm (X) = G(Øz), where
1
Σ
250
x
j=1
From (Chizat, 2022, Prop. 2.4), we see that m the discretization of (11):
Gm (X) VV (x
Thus, this yields
x
X (1) + 1| -V p [ プー 8 14
2 ((1)
(8) Mo
(23)
where > 0 is a step-size, the 2 are i.i.d. standard Gaussian variables, and all the particles should be projected onto X' at each step if X' has boundaries. The MFL. dynamics are recovered the limit and 7 0 (Suzuki al., 2023; Nitanda al., 2022; Chizat, 2022).
Recently, (Suzuki et al., 2023; Chen et al., 2023) have shown a uniform-in-time propagation of chaos for the MFL dynamics: the "distance" between the m-particle distribution and the -ffttile-prtticle limit is order O (+) for all / > 0.
E.3 EXPONENTIAL CONVERGENCE
Theorem E.2 (Thm. 3.4, restated). Assume the d-torus. Let ob E P(X)" such that F(Po) < +x. Then for 2 O, there exists unique solution (ps)a20 to the MFL dynamics (11),
Let r > 0 and assume that po has a bounded absolute log- deg-denityy, it holds
(r()s) min F- Ce min F+),
where BBee-a/T for some 8, O independently of μa and T. Moreover, taking a smooth time- dependent that decays asymptotically as â log s for some ā a, holds Fo(1u,) Fo(u²) S log log s/ log s + and converges weakty to the min- min-entropy estimator μt.
Proof. As in (Chizat et al., 2022), we simply need to verify the assumptions in (Chizat, 2022, Thm. 3.2). Recall that the objective function of the form F G HH. The stability and regularity
31

---

Published as a conference paper at ICLR 2025
of the first-variatin V, (Chizat, 2022, Assumption 1), is immediate from (Chizat et al., 2022, Prop. C.2) and that is bounded. The convexity of Fo and existence of a minimizer for F,, (Chizat, 2022, Assumption 2), follows from Proposition E.1.
For the uniform Sog-Soolev inequality (LSI), (Chizat, 2022, Assumption 3), first note that the ith component the first-variaton of Fo given by [p] + T log log Define D := diam X and E : diamngte yx where D < +o: by assumption and E < too as is bounded. Note that V [14] +000 the gradient formula for 8Fit 55 is non-negative and is bounded by Ee Ee 1/////2 /220)) and by (Chizat et al., 2022, App. A, Eq. 17), the Schrödinger potential (pi,4 has an oscillation bounded by
[m]
sup GCT NT, (2. x,EEX''
inf 2,yEX
T4 (I,3) D2/2.
Following the argument in the proof of (Chizat, 2022, Thm. 3.3), the probability measure propor- tional to - log // satisfies LSI with constant ae for some a, independent
of 8,6, M4-
Then, (Chizat, 2022, Thm. 3.2) guarantees the exponential convergence with rate e with C = 2TP. Furthermore, the convergence result with simulated annealing follows from (Chizat, 2022, Thm. 4.1).
-
F ADDITIONAL EXPERIMENTS
F.1 SETTINGS FOR EXPERIMENTS IN MAIN TEXT
All experiments were run on MI Macbook Air with 16 GB of RAM. Synthetic experiments take a few minutes to run, and Wikipedia experiments take few hours to run.
"Constant velocity" model this experiment, the diffusivity parameter set at T 0.05. Par- ticles are initialized from Xg (0, 0.1 7) and simulated over the time interval E [0. 5] with marginals sampled at evenly spaced intervals. Both PO-MFL and FO-MFL are applied using 100 particles, observe 32 particles at each time point, and we use kernel width of a = 1.0 for the data-fitting term. The optimization procedure initialized with 7j = 0.5 and continues for 2,000 iterations. The number of Sinkhorn iterations for entropic OT capped 500 iterations.
For the crossing paths experiment, the diffusivity parameter is set at r = 0.0005, and the time inter- val is (0, 2.25), and marginals are sampled 10 evenly spaced intervals we use 50 particles.
Wikipedia data In this experiment, the diffusivity parameter is set at r = 0.001. Particles are initialized uniformly over the interval [100, 300). We fitting term. The optimization is initialized with = 0.5 and continues for 2,000 iterations. The number of Sinkhorn iterations for entropic OT is capped at 250 iterations. We scale the data by 1/50 for stability during optimization.
kernel width of a 1.0 for the data-
The three websites we use are:
F.2 CIRCULAR MOTION MODEL
In the circular motion experiment, the diffusivity parameter is set at T 0..002. Particles are initialized from Xo ~ (N00, 0.1 and simulated over the time interval E 3.,3.14 with marginals sampled at 15 evenly spaced intervals. Both PO-MFL and FO-MFL are applied using m = 100 particles, we observe 32 particles at each time point, and we use kemel width of a 1.0 for the data-fitting term. The optimization procedure is initialized with n = 0.5 and continues for 4,000 iterations. The number of Sinkhorn iterations for entropic OT capped 500 iterations.
32

---

Published as a conference paper at ICLR 2025
Velocity: (0.5, 1.0)
Velocity: (3.0, 6.0)
PO-MFLL
MFL
PO-MFLL
MFL
1.0
1.0
Velocity: (1.0, 2.0)
Velocity: (3.5, 7.0)
PO.MFL
MFL
PO-FLL
MFL
15
10
10
2
2
Velocity: 31.5
Velocity: (4.0, 8.0)
PO
MFL
PO-MFL
MFL
15
15
10
10
01
0
0.0
2.5
5.0
7.5
0.0
2.5
5.0
7.5
Velocity: (2.0, 4.0)
Velocity: (4.5, 9.0)
PO-MFL
7.5
3.0
5.0
10
10
2.5
2.5
3.0
0.0
0
0
0
5
0
5
Velocity: (2.5, 5.0)
Velocity: (5.0, 10.0)
10
10
20
20
10
10
0
Figure 6: Varying velocity. We see that as the ground truth initial velocity increases, MFL breaks while PO-MFL remains robust.
In our second model, the particles (0, 0, 0) E S x R.2 represent a constant acceleration model on the unit circle, starting from the initial condition (0, 0.5, 1). Here, acceleration. In this experiment, we only observe the position, e.g. g(0,.,-) 0. In Figure 7a, show the ground truth with position the left and angular velocity on the right. In Figure 7b, we show that PO-MFL successfully reconstruct the positions while although FO-MFL converges, it does recover the ground truth. In Figure 7c, we show that the reconstructed velocity matches that of the ground truth in Figure 7a.
angular velocity and angular
In the following sections, if parameter is not stated, we assume the same setting of parameters as in the main text.
F.3 VARYING VELOCITY
Experiments varying the mean of the ground truth initial velocity distribution are shown in Figure 6. At the endpoints, observe 32 particles, and in the intermediate stages, observe just 2 particles per time point. Note that the small velocity regime, although MFL converges, it converges to the wrong distribution.
33

---

Published as a conference paper at ICLR 2025
PO-MFL
MTL
1.0
1.0
0.5
-10
-1.0
-0.5
0.0
(a) Ground truth.
(b) Reconstructed position marginals.
2
10
12
Time step
(c) Reconstructed angular velocity marginals from PO-MFL.
Figure 7: Cireular motion model.
Observation
PO-MFL
MFL
Observation
8
PO-MFL
MFL
10
10
10
10
L4
0.0
2.5
5.0
0.0
2.5
5.0
0.0
2.5
5.0
0.0
2.5
5.0
0.0
2.5
5.0
0.0
5.0
Observation
PO-MFL
Observation
PO-MFL
10
10
10-
10-
10-
10-
0-
0.0
2.5
5.0
0.0
2.5
5.0
0.0
2.5
5.0
0.0
2.5
5.0
0.0
2.5
5.0
0.0
2.5
5.0
Observation
PO-MFL
MFL
Observetion 32
PO-MFL
MfL
10-
10-
5.0
5.0
5.0
Figure 8: Varying number of observations the intermediate times. Increasing number of observa- tions improves the optimization.
F.4 VARYING NUMBER OF OBSERVED PARTICLES
Figures 8 and 9 show results when the number of observed samples at the intermediate time points varied (the number of observations the endpoints is held constant 32). Here, try the same settings as above, but now we consider velocity (4, ú) (2, 4). We try the number of observations 1, 2, 4, 8, 16, 32, 128, 256. Even in a large number of observation regime, FO-MFL algorithm is not capable of reconstructing the full trajectory, instead clustering around the center of the overall trajectory.
F.5 VARYING TEMPORAL SAMPLING DENSITY At
In Figure 10, we show results for increasing the density of temporal sampling. At the endpoints, we observe 32 particles, and in the intermediate stages, we observe 2 particles. FO-MFL was sensitive
34

---

Published as a conference paper at ICLR 2025
Observation 128
PO-MFL
10
10
10-
0.0
2.5
5.0
0.0
2.5
5.0
0.0
2.5
5.0
Observation
256
PO-MFL
10
10
10-
0
0.0
2.5
Figure 9: A high observation regime at the intermediate time points.
to hyperparameter values as we needed to try different parameters to get semi-reasonable results for the figure. We used σ 0.1, 25.
Timepoints
PO-MFL
Timepoints =
PO-MFL
7.5
7.5
7.5
5.0
5.0
5.0
2-
2.5
2.5
2.5-
0.0
0.0
0.0
Timepoints
PO-MFL
MFL
Timepoints 10
PO-WIL
HTL
7.5
7.5
7.5
5.0
2.5
2.5
2.5
0.
Timepoints
PO-MFL
Timepoints
7.5
7.5
7.5
7.5
7.5
7.5
5.0
5.0
5.0
5.0
5.0
5.0
2.5
2.5
2.5
0.0
0.0
0.0
0.0
0.0
0.0
0
2
2
2
Timepoints
PO-MFL
MFL
Timepoints
PO-MFL
5.0
5.0
5.0
2.5
2.5
2.5
2.5
2.5
2.5
0.0
0.0
0.0
0.0
0.0
0.0
Timepoints
PO-MFL
MFL
Timepoints 15
PO-MFL
MFL
7.5
7.5
7.5
7.5
7.5-
7.5
5.0
5,0
5,0
5.0
5.0
5.0
2.5
2.5
2.5
2
Figure 10: Varying the number of observed time points for fixed time window, i.e. varying At. Note that PO-MFL is always robust. FO-MFL does better with more observations, but the method still tends to collapse inwards because its model suggests that, in expectation, particles should not be moving (as the Brownian motion reference measure has 0 expectation).
35


Reviews:
[{"review_id": "GfR5I3mrWR", "reviewer": "Reviewer_mqcs", "summary": "This paper proposes a latent variable model for trajectory inference, i.e., inferring a time-indexed collection of particles. We assume a latent state vector $(X_t)$ follows some trajectory driven by an SDE, and project the state vector into the observation space via a function $g$, and the objective is to recover the latent trajectories that give rise to the observations. \n\nTo infer the latent trajectories, we formulate a minimum entropy functional (an approximation of the KL-divergence with Gaussian kernel smoothing) that can be optimized by an entropy-regularized gradient flow that can be simulated via a McKean-Vlasov process, i.e., the PO-MFL algorithm. \n\nThe paper evaluates PO-MFL empirically via toy examples and Wikipedia traffic data, showing improvements compared to the original MFL.", "questions": "- I don't understand precisely how optimal transport fits into the framework -- eq. 11 is a Wasserstein gradient flow w.r.t. a convexified functional in eq. 8. Is it possible to simplify the framework via Wasserstein gradient flow?\n- The paper mentions an unknown potential function $\\Psi$ in eq. 1, but does not mention outside the appendix afterwards. Could you explain if inferring $\\Psi$ is left for future work, or if it is omitted somewhere in the later notations?\n- I don't understand the explanation at the beginning of Section 3.4. Why do we need to increase the entropy by $\\epsilon$ when $G + \\tau H$ is already convex?", "limitations": "", "rating": 3, "confidence": 2, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "While I can't seem to understand part of the paper, I think the overall approach of applying latent variable models to trajectory inference is is a sound and novel approach. I think trajectory inference for latent variables are also a key open problem in this field, and this paper proposes an approach to solving this problem. I also think that the paper's main idea is technically sound for the parts I am able to understand. \n\nThe empirical illustration looks convincing and showcases the validity of the algorithm. While I have some reservations about whether or not it is a fair comparison with MFL, I think the paper's approach of inferring latent trajectories + a push-forward function g is overall more preferable than trajectory inference in the observation space.", "weaknesses": "While I think the paper presents an interesting approach to trajectory inference, I find the paper's presentation overall unsatisfactory. The explanations are hard to follow, and the notations are perplexing. I think the lack of readability makes it unsuitable to grant an acceptance for the current version, and it makes it difficult to assess the contribution.\n\n- The paper is hard to follow overall. While there is a \"notations\" section, new notations appear throughout the rest of the paper. Many notation can be inferred from context, but that seems to defeat the purpose of a notation section.  \n- Some of the confusing notation include 2 different ways to denote convolving variables with a Gaussian density in eq. 4, e.g., both $\\mathcal{N}_\\sigma$ and $\\hat{\\rho}_i^{T, h}$. $\\gamma$ is both a multiplicative constant in eq. 3 and part of the OT formulation in eq. 6. \n- Some of the theoretical explanations in the paper, such as $C_\\Psi$-ensemble observability and the general assumptions on $\\mathcal{X}$ and $\\mathcal{Y}$ are mentioned in the paper, but they are not really included in the discussion.\n- There are a number of \"filler\" type of sentences about \"recall that\", \"note that\" which often are re-statements or clarifications. I think it is helpful to reformulate the explanation in a more cogent way. Perhaps it is a good idea to dwell on the most important concepts of the paper, and relegate some less relevant but still beneficial explanation to the appendix. \n- Some notations seem unnecessarily complex. For example, I don't understand why the time variables need to be formulated as $t_T^T$. \n- The paper makes many references to previous works in the abstract, which makes it look more like a background overview. I think the abstract needs to be restructured in a more self-contained fashion. \n- PO-MFL seems to empirically outperform other trajectory inference methods, and latent variable models are ostensibly harder to infer. But this experimental setting (e.g., Figure 4) seems more beneficial, as the latent velocity variables are easier to infer.", "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": "", "extra_scores": {}}, {"review_id": "Jv879kusJv", "reviewer": "Reviewer_tXu2", "summary": "This work proposes methodology for sampling trajectories of a diffusion which is only partially (but repeatedly) observed at a finite number of time points. At each iteration, the proposed methodology uses optimal-transport ideas to construct a joint distribution over the states of the diffusion at observation times from their marginal distributions (the latter are approximated via samples/\"particles\"). At the end of each iteration, a mean-field Langevin algorithm is then used to optimise the particle locations.", "questions": "1. What can you say about how this method scales with the dimension of $X_t$?\n\n2. What do you do if $\\Xi$ and $\\tau$ are unknown?\n\n3. Having a the same diffusion parameter $\\tau$ for all components of $X_t$ (e.g., for both the evolution of the positions and of the velocities) seems very limiting. Is there any hope of relaxing this assumption?\n\n4. I didn't understand the \"crossing paths\" experiment discussed in Lines 477--500 and Figure 4. Can you explain this more clearly?\n\n5. What is the real-world significance of the \"Wikipedia traffic data\" experiment. In this case, all the traffic data has already been observed. What is the real-world utility of the estimated trajectories?\n\n6. In the caption of Table 1, what does \"MCMC\" refer to here? I didn't see MCMC methods being mentioned anywhere else.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. **Clarity.** I believe that this work is overall fairly well written (though see some of the comments/questions below, especially with regards to the explanation of some of the experiments).\n\n2. **Novelty/Originality.** The proposed methodology seems new. The main innovation is the extension of a previous approach to allow for a non-zero drift term $\\Xi$. This seems fairly moderate, but IMHO sufficient given that it requires non-trivial extensions to the theoretical guarantees which the authors provide.\n\n3. **Quality/Correctness.** This paper comes with a number of rigorously derived theoretical guarantees (though I did not have time to verify all details nor check the derivations in the appendix).", "weaknesses": "**Main issues:**\n\n\n1. **Lack of comparison with simple Monte Carlo alternatives.** I believe the method needs to be compared against a simple Monte Carlo alternative like a particle filter (with backward sampling to obtain trajectories) or a simple MCMC method. The method is currently only compared against an alternative, denoted \"MFL\", which seems to be so heavily biased that it makes almost anything look good by comparison.\n\n### EDIT: 2024-11-22: As explained in my comment below, Monte Carlo alternatives may likely need to be modified in a not-completely-trivial way to be applicable here. Therefore, I agree that it is fine not to compare against these. \n\n\n2. **Seemingly high computational cost.** The actual computation time and total number of iterations used in the numerical experiments should be mentioned in the main text. The only mention I found is in Appendix F.1 where it says \"[s]ynthetic experiments take a few minutes to run, and Wikipedia experiments take a few hours to run.\" However, the phrase \"a few minutes\" is too vague. Further, the synthetic experiments are fairly small (the dimension of $X_t$ seems to be between 2 and 4). I would expect that a simple Monte Carlo alternative (like those mentioned above) would run in mere seconds rather than minutes.\n\n\n3. **Variance of target distribution seems poorly captured.** As shown in Figure 1c, and commented on in Lines 467--469, the *average* of the sampled velocities seems to match the ground truth. However, the variance seems to be heavily overestimated (or is this just a visual artifact of using scatter plots with different numbers of points?). Additionally, Figure 1a seems to indicate that the variance of the ground truth increases over time but this does not seem to be reflected in the trajectories produced by the proposed method shown in Figure 1b.\n\n\n4. **Lack of clarity about the real-data experiment.** I find the \"Wikipedia traffic data\" experiment impossible to understand. In Line 511, the dynamics(?) $x_{t+1} = \\theta_1 x_t + \\theta_2 x_{t-6}$ are stated. But it is not clear to me how this discrete-time autoregression relates to the SDE fro (1)? That is, what are $\\Xi$ and $\\tau$ here? Furthermore, are the three websites modelled/estimated separately or somehow jointly (e.g., do the trajectories shown in Figure 5b represent a single website or all three)?\n\n\n**Minor issues:**\n\n* In Line 32, it should be explained what \"unpaired observations\" means.\n\n* The term \"MFL\" appears multiple times throughout the manuscript but sometimes refers to the minimisation algorithm (also used within the proposed methods) and sometimes to the algorithm used as a benchmark in the numerical experiments. I found this confusing, especially because it makes it difficult to find the reference for the benchmark algorithm.\n\n* Figure 2: The tick labels on the 1st axis in the left panel are confusing. Later in the text it is explained that these the final 500 iterations rather than Iterations 1 to 500.\n\n* All axes need to be labelled in all figures. And the 1st-axis labels in Figure 5a are all overlapping and thus unreadable.\n\n* Figures should not appear in the text before they are referenced. This is particularly confusing with Figure 1 which appears on Page 2 but is only referenced/used on Page 9. But there are similar problems with Figures 2 and 3.\n\n* The number of footnotes seems excessive. Many of these, e.g., Footnote 2 an 3, should be included in the text.\n\n* Bibliography: Most entries in the bibliography are missing the journal/conference name. There are also many typos (e.g., missing capitalisation such as \"wasserstein\", \"schroedinger\", \"brownian\", etc).\n\n\n\n**Typos:**\n\n* L135: $H(\\mu|\\nu)$ is defined for probability measures $\\mu$ and $\\nu$ here but is later used with non-probability measures.\n\n* L144: The last line in the \"notation\" paragraph seems incomplete since neither $g$ nor $\\xi$ have been previously defined and it should probably be stated explicitly that $g_{\\sharp\\mu}$ denotes the pushforward measure.\n\n* L231: \"measurement\" -> \"measurements\"\n\n* Equation 4: Define the \"$*$\" operator.\n\n* L251: What is $\\hat{\\mu}_i^{T, h}$?\n\n* L253: \"computationally more effective\" -> \"computationally more efficient\"\n\n* L280: Equation 7 is referenced before it appears in the text.\n\n* L289: Shouldn't the first term in the cost function be $-\\tau_i$ rather than $-\\Delta t_i$?\n\n* L335: Shouldn't \"approximate (10)\" be replaced by \"approximate the first two terms in (10)\"?\n\n* L338: There seems to be a space missing at the end of the sentence.\n\n* L362: What is \"$\\xi$\" (i.e., without a superscript) here?\n\n* L367: Use natbib's \\citet rather than \\citep\n\n* Algorithm 1, Line 8: Is there a \"hat\" accent missing on $\\mathbf{m}$ in the comment?\n\n* Algorithm 1, Lines 8 and 10 (and in the \"Require\" line): Shouldn't this be \"$T$\" rather than \"$t$\"?", "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": "", "extra_scores": {}}, {"review_id": "XrYDrh9IA3", "reviewer": "Reviewer_VsTM", "summary": "The paper proposes an extension of existing methods for the trajectory inference problem, focusing on where states are partially observed only through \"snapshots\" at a given timepoint. These snapshots, for example, capture position but not velocity. They are also unlabelled, for example, you cannot know whether a particle at timepoint t_1 corresponds to a particle at timepoint t_2. The objective is to infer the trajectory distribution of the particles over time. \n\nThe authors introduce the PO-MFL (Partially observed mean field Langevin) algorithm to address this challenging problem, developing a deep theoretical contribution that robustly explores their algorithm, including a weak consistency result (section 3.1) showing that their proposed objective function minimises to the true law of the SDE, under some assumptions and in the limit, in ideal conditions. As well as a convergence result (section 3.4), indicating that the PO-MFL algorithm’s approximation of the true trajectory distribution improves over time and eventually stabilises, showing the robustness of PO-MFL.\n\nFinally, the paper grounds their theory in robust experiments - for example, showing empirically that the convergence from Theorem 3.4 holds in Figure 2. Additional experiments show the usefulness of the method, testing robustness across various settings on synthetic and real-world data.\n\nI congratulate the authors on a well constructed paper.", "questions": "1. How does the computational complexity scale against other methods such as regular MFL?\n\n2. It would improve the strength of the experiments to compare to other methods. Would comparing to something like rejection sampling based methods work? Even with a high computational expense, it could be worth it to compare to as a baseline?\n\n3. Are there any scenarios (within the Partially Observed Trajectory Inference setting) in which PO-MFL will break down? What restrictions does the method have, if you know of any? Are there any limiting factors?\n\n4. Are there any limitations or concerns related to relying on Assumption 2.1? Does the restriction for $\\Xi$ to be known, divergence-free and Lipschitz continuous restrict the applicability to some real world systems that have more complex dynamics?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. The paper introduces a well developed framework, extending mean-field Langevin dynamics with latent dynamics modelling for partially observed trajectory inference, making substantial advancements in the field. \n\n2. The paper is theoretically robust - providing a wealth of theory that seeks to capture the applicability and performance of the method, such as the consistency and convergence results.\n\n3. The experiments are varied on synthetic and real-world data, provide insights into the method, and empirically back up theoretical claims.\n\n4. The paper is extremely well written, providing a good background into the subject and gradually introducing concepts to build up to a final goal.\n\n5. The work is framed well within the global context of the problem, highlighting relevant works and providing comparisons and gaps in the literature when needed. Applications of the problem are presented in many cases, showcasing the general applicability of the need for solving the Partially Observed Trajectory Inference problem.", "weaknesses": "1. The computational complexity of the algorithm is $O(NTm^2)$, which is not explored further than it is stated, nor is it compared against other methods in the experiments section. Some details on real-world runtime could be useful, due to its (presumably) large scaling cost.\n\n2. Experiments show comparisons to regular MFL, which is not suited for the unobserved case, and therefore there are no comparisons to methods that are suited for this domain. However, I am unfamiliar with the literature, so not sure if there are other methods that can be used in this case.\n\n3. Whilst there is some discussion on the limitation towards model misspecification in the summary section, limitations of the methodology are not discussed in great detail.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)