Solving Schrödinger Bridges via Maximum Likelihood

Introduction

Analysis of cross-sectional data is ubiquitous in machine learning and science. Temporal data are typically sampled at discrete intervals due to technological or physical constraints. This means information between time points is lost. This motivates the need to model the stochastic evolution of a process between sampled time points. The classical Schrödinger bridge problem (?; ?) finds the most likely stochastic process that evolves a distribution $\pi_{0}({\bm{x}})$ to another distribution $\pi_{1}({\bm{y}})$ consistently with a pre-specified Brownian motion. We consider a more general dynamical Schrödinger bridge problem for any pre-specified diffusion prior. Practically, this generalization allows us to exploit domain knowledge, e.g. oceanic and atmospheric flows might be interpolated from empirical measurements using previously established dynamics as priors.

In the classical set up numerical approaches to solve the Schrödinger bridge are mainly based on the Sinkhorn-Knopp algorithm (?). However, extending those algorithms to more general diffusion priors and marginals requires complex adaptations. We introduce an iterative proportional maximum likelihood (IPML) algorithm to solve the general Schrödinger bridge problem. The IPML algorithm also obtains a good approximation for the dynamics of the underlying physical process that solves the SBP. This contrasts to previous approaches (?; ?; ?) that estimate the value that extremises the SBP objective. Practically, this means we obtain physically interpretable solutions that we can leverage for downstream tasks.

IPML’s inspiration is from probabilistic numerics (PN) (?). We combine a PN styled formulation with the iterative proportional fitting procedure (IPFP) (?; ?). The algorithm iteratively simulates trajectories that converge to the SBP. We prove that IPML converges in probability at each iteration. Our numerical experiments show the algorithm can be implemented with Gaussian process (GP) models of the drifts. GPs allow us to incorporate functional prior information. We demonstrate the practical use of our algorithm on real-world embryoid cells data (see Figure 5) by quantitatively and qualitatively comparing the performance of our algorithm to state-of-the-art deep learning methods and optimal transport techniques. To summarise the main contributions of our work are:

We recast the iterations of the dynamic IPFP algorithm as a regression based maximum likelihood objective. This is formalised in Theorem 1 and Observation 1. This differs to prior approaches such as (?) where their maximum likelihood formulation solves a density estimation problem. This allows the application of many regression methods from machine learning that scale well to high dimensional problems. Note that this formulation can be parametrised with any method we chose GPs however neural networks would also be well suited.

We solve the aforementioned regression objectives using GPs (?) motivated by the connection between the drift of stochastic differential equations and GPs (?).

We provide a conceptual and empirical comparison with the approach by (?) and detail why the density estimation formulation in (?) scales poorly with dimension.

Technical Background

Our solution has three components. (1) We reformulate the SBP as a dynamical system giving a stochastic differential equation (SDE) with initial value (IV) and final value (FV) constraints (?). (2) We reverse the system, reformulating the FV constraint as an IV constraint. Both IV problems are solved through a stochastic control formulation (Section 2.1). (3) We iterate the IV and FV constrained problems to converge to the full boundary value SDE.

The dynamic version of the Schrödinger bridge is written in terms of measures over the space of trajectories, that describe the stochastic dynamics defined over the unit interval.

(Dynamic Schrödinger problem) The dynamic Schrödinger problem is given by

When sampling solutions, IV constraints are trivially solved by initialisation of samples, but FV constraints are more problematic. However, if we can reformulate the SBP as a time reversed diffusion the FV constraint becomes an IV constraint (Section 3.2). Here we review how the diffusion is reversed. We reverse time (see Figure 2) in the random variable ${\mathbf{x}}(t)$ described by the diffusion in Equation 2 such that ${\mathbf{x}}^{-}(t)={\mathbf{x}}(1-t)$. The reverse time diffusion ${\mathbf{x}}^{-}(t)$ is also an Itô process. For a modern application of time reversal in Machine Learning see (?).

(?) If ${\mathbf{x}}^{+}(t)$ obeys the SDE

then ${\mathbf{x}}^{-}(t)={\mathbf{x}}^{+}(1-t)$ obeys

where ${\bm{\beta}}^{-}(t)$ is a Brownian motion adapted to the reverse filtration $({\mathcal{F}}^{-}_{i})_{i\in T}$, that is ${\mathcal{F}}^{-}_{t}\!\!\subseteq{\mathcal{F}}^{-}_{s},\;s\!\leq\!t$

Furthermore, the dual drift ${\bm{u}}^{-}({\mathbf{x}},t)$ satisfies Nelson’s duality relation:

where $p({\mathbf{x}},t)$ solves the associated Fokker–Planck equation.

A variety of proofs can be found for this result (?; ?; ?; ?; ?). Note the formulation of Equation 3 varies slightly across studies we use the one from (?) and (?).See page 87, definition of osmotic velocity in (?). ∎

Stochastic Control Formulation

Now that the FV constraint has been converted to an IV constraint we cast the problem into a stochastic control formulation to estimate the drift of each diffusion process. Following from the dynamic formulation, the control formulation casts the problem explicitly in terms of stochastic differential equations. The control formulation is used to enforce constraints as initial value problems. Furthermore, the drift based formulations of the SBP admit a reverse time formulation which starts the chain at the end of the interval and progresses the dynamics backwards in time to the start.

This theorem follows by a direct application of the disintegration theorem (Appendix C) followed by Girsanov’s theorem. A detailed proof can be found in (?). ∎

Using the above decompositions, we can solve the SBP by minimising either decomposition in equation 4 over the space of random processes ${\mathbf{b}}^{\pm}(t)$ that satisfy a valid Itô SDE drift. The backwards and forwards objectives are respectively:

While we do not directly use the stochastic control formulations, the drift based formulation serves as inspiration for our iterative scheme. Specifically, the existence and parametrisation of an optimal drift as shown in (?; ?).

(?) The objectives in Eq equation 6 have optimal drifts:

where the potentials $\phi,\hat{\phi}$ solve the Schrödinger system (?).

This Lemma is key in formulating our ML approach to IPFP since it justifies our parametrisation of the drift in terms of a deterministic function ${\bm{u}}^{\pm*}$ i.e. ${\mathbf{b}}^{\pm*}(t)={\bm{u}}^{\pm*}({\mathbf{x}}^{\pm}(t),t)$. For a brief introduction to the Schrödinger system and potentials see Appendix B.

2 Iterative Proportional Fitting Procedure

We now have two boundary value constrained diffusion processes solved through a stochastic control formulation. We use the iterative proportional fitting procedure (IPFP) to alternate between the forward and backward formulation with only one initial value constraint enforced at a time (see Figure 3), such that convergence to the full boundary value problem is guaranteed.

The measure theoretic version of IPFP we introduce is an extension of the continuous IPFP, initially proposed by (?) to a more general setting over general probability measures (See (?; ?)). The convergence of IPFP has been shown in (?) and further extensions and results have been presented in (?; ?). The idea behind this family of approaches is to alternate minimising KL between the two marginal distribution constraints,

Methodology

In this section we introduce IPML by providing the theoretical foundations of our algorithm and convergence guarantees. Then, we propose a practical implementation of IPML based on a Bayesian non-parametric model (GP). These components allow us to solve the general SBP problem.

We present a novel approach to approximately solve the empirical Schrödinger bridge problem by exploiting the closed form expressions of the half bridge problem. Rather than parametrising the measures in the half bridge and solving the optimisation numerically, we seek to directly approximate the measure that extremises the half bridge objective. We do this by using Gaussian processes (?) to estimate the drift of the trajectories sampled from the optimal half bridge measure.

We start from the following observation, which tells us how to sample from the optimal half bridge distribution (see Appendix D for proof).

Then, we can sample from the solution to the half bridges,

via simulating trajectories (e.g. using the Euler-Maruyama (EM) method) following the SDEs

Intuitively, we are performing a cut-and-paste-styled operation by cutting the dynamics of the shortened (unconstrained) time interval and pasting the constraint to it at the corresponding boundary.

Then carrying out maximum likelihood estimation of the time reversed drift ${{\bm{u}}_{0}}^{-}$ on time reversed samples:

and ${\bm{u}}^{-}$ is the drift of our estimator SDE:

Note that w.l.o.g. the above result also holds for estimating the forward drift from backward samples. The combination of Observation 1 and Theorem 1 constitute one the main contributions of this work as they allow us to solve half bridges with a simple regression objective. It is important to highlight that in practice for most interesting SDEs we can only sample approximately using consistent schemes, for these settings our proof of Theorem 1 does not hold, however we believe it is possible to adapt the result to this case.

2 On the Need for Time Reversal

which is a forward SDE with a terminal hitting condition. Unlike a forward SDE with an initial value problem the above SDE is not trivial to sample from, one approach would be to solve the backwards Kolmogorov equation subject to the terminal condition. However, this approach requires (1) solving a parabolic PDE that typically involves mesh based methods that do not scale well in high dimensions and (2) carrying out density estimation on the samples from $\pi_{1}$, also problematic in high dimensions. However, if we consider the time reversed process, the terminal condition becomes an initial value problem, $d{\mathbf{x}}^{-}(t)={\mathbf{b}}^{-}(t)+\sqrt{\gamma}d{\bm{\beta}}^{-}(t),\quad{\mathbf{x}}(0)\sim\pi_{1},$ for which it is easy to sample consistent trajectories via the Euler-Maruyama (EM) discretisation without mesh based methods or additional density estimations.

3 Iterative Proportional Maximum Likelihood (IPML)

Combining the KL minimisation routines in the original IPFP algorithm with our MLE based drift estimation gives Algorithm 2. The routine DriftFit fits a dual drift on the sampled trajectories, and can be parametrised by any function estimation procedure with consistency guarantees. The SDESolve routine generates $M$ trajectories using the EM method.

Drift Estimation with Gaussian Processes: One can choose any parametric or non parametric model to carry out the DriftFit routine. In this section we will briefly introduce our implementation. Our DriftFit routine is based on the work in ?; ?(?; ?) that uses GPs to estimate the drift of SDEs from observations. We refer to this routine as GPDriftFit. We can restate the regression problems given by Equation 10 and its reversed counterpart in the following form:

We now have the relevant ingredients to carry out IPML as specified in Algorithm 2. The computational cost of IPML with a Gaussian Process is detailed in Appendix I.5. The majority of the computation is spent fitting the GP in DriftFit at each iteration and scales as a function of the time discretization used as well as the size of $\pi_{0}$ and $\pi_{1}$.

Related Methodology

In this section we carry out a conceptual comparison with two pre-existing numerical approaches for solving the static SBP. While our goal is to solve the dynamic SBP, the solution of the static SBP can be used to construct that of the dynamic SBP (?), and this connection is central to our discussion. We would also like to highlight that an algorithm akin to IPML has been proposed concurrently and independently by (?), the main difference with our algorithm is that they estimate the drifts of the SDEs using neural networks score matching while we use using Gaussian processes and MLE based ideas.

Once the problem has been discretized as described, the Sinkhorn-Knopp algorithm (?) can be applied directly to fit an optimal discrete transport map (and discrete SBP potentials $\phi,\hat{\phi}$) between the two distributions. Recent work (?) has showed empirical success in forming a continuous approximation of the SBP potentials using the logsumexp formula(?), however, it still remains to formally analyse the accuracy of the logsumexp potentials.

The method proposed by (?) is perhaps the most similar approach to our approach and consists of iterating two coupled density estimation objectives (see Appendix H) fitted at the marginals until convergence. While conceptually similar to our approach, there are 3 key differences. As with Sinkhorn-Knopp based methods, their approach aims to solve the static SBP. Once converged it requires further approximation to estimate the optimal drift. The coupled maximum likelihood formulation of the static half bridges in (?) is based on un-normalized density estimation with respect to the SBP potentials $\phi,\hat{\phi}$. The coupling of $\phi,\hat{\phi}$ in these objective does not directly admit the application of modern methods in density estimation since it does not allow us to freely parametrise MLE estimators for the boundaries thus neural density estimators such as (?; ?; ?) cannot be taken advantage of to circumvent the computation of the partition function. Regression problems are ubiquitous in machine learning (ML) and thus why we believe that this formulation can be very impactful as it allows us to leverage all these methods from ML. Experimentally, regression methods have been observed to scale better to high dimensional problems than density estimation methods, we observed similar evidence as we were unable to scale up the DDSB method beyond 2 dimensions.

Additionally, to compute the normalizing term in the DDSB objective, we have to estimate a multidimensional integral that is not taken with respect to a probability distribution. This poses a difficult challenge in high dimensions. For a more detailed commentary, see Appendix H. Note that the method by (?) uses importance sampling to estimate these quantities which performs poorly beyond 2 dimensions.

It worth noting that we are interested in a method that can obtain a dynamic interpolation between two distributions. A downside of solving the static bridge either by Sinkhorn or (?) is that it does not directly provide us with an estimate of the optimal dynamics. In order to obtain an estimate of the optimal drift we require a series of approximations to estimate the integrals in 82, 83, thus every time we evaluate the optimal drift using these approaches we have to simulate the prior SDE ${\mathcal{O}}(N+M)$ times. This makes the run-time of obtaining the drift and dynamical interpolation expensive, since for each Euler step we take we have to simulate another SDE and backpropagate through it to evaluate the drift.

Finally we would like to highlight that a series of modern approaches to generative modelling (?; ?; ?; ?) motivate both empirically and theoretically the gain in accuracy obtained in generative modelling tasks when using a dynamical approach rather than a static one.

Numerical Experiments

In this section, we demonstrate the capability of IPML to solve the Schrödinger bridge while efficiently incorporating priors on a range of different tasks from synthetic experiments to embryo cells.Code supporting experiments can be found at https://github.com/AforAnonyMeta/IPML-2548.

The first experiment considered is a simple alignment experiment where $\pi_{0}$ and $\pi_{1}$ are either unimodal or bimodal Gaussian distributions (see Appendix I.1 for exact details on the distributions). In Table 1 we compare the accuracy of the fitted marginals with our implementation of the Data-Driven Schrödinger Bridge (DDSB) by (?). The scoring metrics used are the Earth mover’s distance (EMD) as well as a Kolmogorov-Smirnov (KS) statistic on both sample sets.

We were unable to get DDSB to work well when $\pi_{0}$ and $\pi_{1}$ were distant from each other. In these distant settings, DDSB collapses the mass of the marginals to a single data point (see Appendix I.1), as a result for the unimodal experiment we had to set $\gamma=100$ for the DDSB approach to yield sensible results. We can observe IPML obtains better marginals overall and at a lower value of $\gamma=1$.

We carried out 2D experiments with our approach to show the ability of our method to diffuse from a simple uni-modal distribution to a multi-modal distribution as in Figure 2 where our learned bridge successfully splits. Furthermore, we can visually observe how the forward and backwards trajectories are mirror images of each other as expected.

2 2D Double Well Experiments

In the double well experiment, we illustrate how to incorporate an arbitrary functional prior and learn the distribution over paths connecting $\pi_{0}$ and $\pi_{1}$. In order to encode prior information, we experiment with the potential well illustrated in Figure 1 and 9 (Figure 9 can be found in the Appendix). The boundary distributions $\pi_{1},\pi_{0}$ are taken to be Gaussian distributions centred at the centre of each well respectively (See Appendix I.2 for the experiment specification and IPML parameters).

The motivation behind this experiment is to show that the SBP with this prior follows low energy (according to the well’s potential function) trajectories for configurations of particles sampled at the wells. Intuitively, we can expect the learned trajectories to avoid the high energy peak located at ${\bm{x}}=(0,0)$ and go via the “passes” on either side. Note that if we estimated the optimal transport (OT) geodesics between $\pi_{0}$ and $\pi_{1}$ or similarly ran IPML with a Brownian motion prior, the learned optimal trajectories would go right through the middle, which is the highest energy path between wells.

3 Finite sample/iteration convergence

Theorem 1 provides us with asymptotic guarantees. However, it does not extend to the finite sample and discretisation case. To highlight the importance of finite effects on IPML, we carried out this analysis empirically. In Figure 4 we plot an empirical estimate of the error term in the control formulation of the SB (Eq. 5). This term is effectively the mean squared error between the learned drift and the prior drift (gradient field of the well). We analyze this metric for different values of $N$ (number of samples) and $\Delta_{t}$ (discretization factor). We observe that IPML quickly reaches a low error valley, then, the cumulative error from the successive finite sample MLE can be observed and the drift starts to slowly deviate from the prior, this motivates early stopping. As $N$ is increased, IPML achieves lower error faster and deviates less from the prior in later iterations. Finally, we observe placing the drift prior via the GP has a significant effect in improving the error and its convergence. Additionally in Appendix F we detail how this question could be approached from a theoretical perspective while underlining its significance and difficulty as illustrated by the lack of such analyses in related algorithms (?; ?).

4 Single Cell - Embryo Body (EB) Data set

We perform an experiment on an embryoid body scRNA-seq time course (?). Single-cell RNA sequencing enables accurate identification of cells at specific time-points, however, all cells are destroyed by measurement. This prevents modelling single-cell trajectory and instead we rely on modelling the data-manifold at discrete time points. The datasets consist of 5 time points illustrated in Figure 10. To evaluate the performance of the algorithms, we fit the models at the endpoints ($T=1,5$) and predict the intermediate frames. The metric used is the Earth mover’s distance between the data at intermediate frames and the predicted distribution. We evaluate the performance at the endpoints by considering the prediction of the forward model at $T=5$ and the backward model at $T=1$.

We compare the performance with two methods. The first one, TrajectoryNet (?), uses continuous normalizing flows with a soft constraint based on optimal transport. The second, leverages the McCann interpolant (?; ?) to interpolate the discrete OT solution; this corresponds to the linear interpolation induced by the transport map.

The results are summarized in Table 2. In most frames, IPML outperforms TrajectoryNet and performs similarly to OT. As the noise (volatility) in the trajectories goes to 0, IPML theoretically converges to the OT solution with linear geodesics. So without any additional prior information, we would not expect IPML to outperform OT. However, when dealing with finite data, the DriftFit procedure allows for some non-linearity in the trajectories if it improves the fit. As a result, we do see some differences in the convex hull displayed in Figure 5 where we observe a better coverage of the single cell observations. We hypothesize this explains the improvement in performance vs OT for frame 4. The performance of IPML may be further improved by incorporating domain-specific knowledge as a prior. It would outperform OT in those settings.

5 Motion Capture

In this experiment, we demonstrate how IPML can be used to model human motion from sensor data.Data from The CMU Graphics Lab Motion Capture Database funded by the NSF (http://mocap.cs.cmu.edu). The motion corresponds to a basketball movement where the subject raises both arms simultaneously as illustrated in Figure 6 and each sensor corresponds to an oriented angle. We focus on modelling the right shoulder and elbow where the starting distribution corresponds to the leftmost image and the ending one the rightmost. This results in a 4 dimensional space as we model both position and velocity for each sensor. We compare the fit of IPML using a Brownian and a 2nd order linear ODE (Langevin) prior. The experimental details can be found in I.4. As illustrated in figure 6, IPML is able to approximately model the dynamics using both priors. The Brownian prior displays noisy trajectories as expected in contrast to the Langevin prior that, by construction, smooths out the predicted positions (due to the 2nd order term). We observe that the Langevin prior approximates the step function nature of the true trajectory more closely than the Brownian prior, additionally, we can see that it also produces slightly better alignments.

Limitations and Opportunities

In this work, we propose to use GPs to estimate the drift, however, IPML could also be used with a parametric function estimator (e.g. a neural network). This could be useful with high-dimensional data where GPs may underperform. The main advantage of using a GP is its capacity to incorporate functional priors via the mean function. This can be useful in applications such as molecular dynamics where a potential function may be available. In contrast, implementing functional regularization in a neural network would require approximating a non-trivial high-dimensional integral to estimate the mean squared error between the parametric estimator and the functional prior.

A particular useful extension would be to adapt the SBP to work with more general forms of volatility functions that are not constant. This can be used, for example, in enforcing positivity constraints on a stochastic process via a geometric Brownian motion prior; this has applications in modelling biological signals such as transcription factors (?). Work in this direction would require extending the theory of SBPs to non constant volatility functions.

Another promising setting is when multiple frames of data are available rather than just two boundary conditions. The IPFP algorithm trivially adapts to multiple constraints (?) rather than just initial and terminal distributions. Future work could explore experiments similar to the one presented in (?) where multiple frames are considered during training and the performance is measured using a leave-one-out procedure.

We have presented IPML, a method to solve the Schrödinger bridge for arbitrary diffusion priors. We presented theoretical results guaranteeing convergence in the limit of infinite data. We devised a practical application of the algorithm using Gaussian Processes and presented several experiments on a variety of problems from synthetic to biological data. The approach opens up opportunities in science, where oftentimes, prior knowledge about the temporal evolution of a process has been developed but needs to be combined with data-driven methods to scale up to modern problems.

2 Acknowledgements

We would like to thank Abdul Fatir Ansari for insightful discussion and helpful remarks.

References

Appendix

Appendix A Assumptions Across Proofs

For abstraction purposes we will list the set of assumptions assumed across all results in this appendix:

All SDEs considered have $L-$Lipchitz diffusion coefficients as well as satisfying linear growth.

The optimal drifts are elements of a compact space and thus satisfy the HJB equations. Note for the proposes of Theorem 1 this can be relaxed using notions of $\Gamma-$convergence.

Appendix B Brief Introduction to the Schrödinger System and Potentials

The Schrödinger system and its potentials are mentioned when introducing some of the results and connections of the full Schrödinger bridge problem. In this section we will provide a brief introduction to how the system arises from the original static Schrödinger bridge. For brevity we will use ${\bm{x}}={\bm{x}}(0)$ and ${\bm{y}}={\bm{x}}(1)$ to denote the boundaries.

The static Schrödinger bridge problem can be derived from the full dynamic bridge by marginalising out the dynamics via the Disintegration Theorem and focusing on the problem only concerning what is happening at the boundaries. The static SBP is formulated as:

which via performing the appropriate variations wrt to $q$ leads to the optimal solution:

which when relabeling the terms containing the Lagrange multipliers we obtain:

furthermore if we relabel the potentials to indicate the times they correspond to we arrive at the Schrödinger system:

The above functional system is refered to as the Schrödinger system and the Schrödinger potentials are given by $\phi,\hat{\phi}$. Furthermore the time interpolates for the potentials can be obtained by the propagation’s:

For a more rigorous and extensive introduction please see (?).

Appendix C Disintegration Theorem - Product Rule for Measures

In this section, we present the Disintegration Theorem in the context of probability measures, which serves as the extension of the product rule to measures that do not admit the traditional product rule. Furthermore we will provide a proof for a direct lemma of the Disintegration Theorem that is more analogous to the standard product rule. Like the product rule these theorems are essential for decomposing and manipulating path measures and thus is needed for most results pertaining to the dynamic SBP.

(Disintegration Theorem for continuous probability measures):

$\pi_{i}:Z\rightarrow Z_{i}$ is a measurable function known as the canonical projection operator (i.e. $\pi_{x}(z_{x},z_{y})=z_{x}$ and $\pi^{-1}_{x}(z_{x})=\{y|\pi_{x}(z_{x})=z\}$),

A direct consequence of the above instance of the disintegration theorem is, with $f({\bm{x}},{\bm{y}})={\bm{1}}_{A_{x}\times A_{y}}({\bm{x}},{\bm{y}})$,

We can see that, in the context of probability measures, the above is effectively analogous to the product rule.

We now have the required ingredients to show the following:

which, via the Radon-Nikodym theorem, implies

Appendix D Proof Sketches For Half Bridges

In this section we provide a proof sketch for the closed form solution of the half bridges as well as a proof for Observation 1.

The forward half bridge admits the following solution:

Via the disintegration theorem, we have the following decomposition of KL:

Then, we can sample from the solution to the following half bridges:

via simulating trajectories following the SDEs

(Sketch) W.l.o.g., Consider the decomposition of the KL divergence that follows from the disintegration Theorem (Appendix C):

Appendix E Proof Sketch for Reverse-MLE Consistency

The proof sketch for Theorem 1 will show how the likelihood converges in the large data and small time step limit to an optimisation of the KL divergence between two reverse time diffusions, from here one can use the standard arguments to show this quantity is minimised when the two measures describe the same stochastic process or equivalently when the drifts are equal.

Normalizing the time reversed likelihood with the true discretised backwards SDE density does not affect the maximum likelihood estimate. That is :

The term -\ln{\mathcal{N}}\left({\bm{x}}^{(n)+}_{t_{k}-\Delta t}\Big{|}{\bm{x}}^{(n)+}_{t_{k}}-\Delta t{\bm{u}}^{-}_{0}\left({\bm{x}}^{(n)+}_{t_{l}},1-t_{k}\right),\gamma\Delta t\right) does not depend on ${\bm{u}}$ and thus is an additive constant. ∎

Then carrying out maximum likelihood estimation of the time reversed drift ${{\bm{u}}_{0}}^{-}$ on time reversed samples:

and ${\bm{u}}^{-}$ is the drift of our estimator SDE:

The density $p({\bm{x}},t)$ is differentiable with respect to ${\bm{x}}$.

The optimal drift lies in a compact space (This can be relaxed using notions of $\Gamma-$convergence).

The prior drift coefficient is $L-$Lipchitz and satisfies linear growth.

For the interest of brevity let ${\bm{u}}^{(n)-}_{t}={\bm{u}}^{-}\left({\bm{x}}^{(n)+}_{t_{k}},1-t_{k}\right)$. Taking logs and applying Lemma 5 to Equation 31 yields:

Where $q({\mathbf{x}}_{t_{T}})$ represents the terminal distribution of the forward SDE (i.e. ${\mathbf{x}}^{+}(1)\sim\pi$). Now we can equivalently write the above expression in terms of the time reversed samples (i.e. ${\bm{x}}_{t_{i}}^{-}={\bm{x}}_{t_{n-i}}^{+}$):

where $\Delta{\bm{x}}^{(n)-}={\bm{x}}^{(n)-}_{t_{k}}-{\bm{x}}^{(n)-}_{t_{k}-\Delta t}$. Now we can consider the limit of Equation 37:

We can write the sum over the time grid as a stochastic integral if we express the inner terms using the continuous time approximation from Lemma 6:

From (?; ?; ?) it follows that $\hat{{\bm{x}}}^{(n)-}(t)$ is a semi-martingale (w.r.t. to the backwards filtration see (?; ?)) then in the limit the stochastic integrals are taken with respect to the true time reversed stochastic process adapted to the backwards filtration $({\mathcal{F}}^{-}_{i})_{i\in T}$ (see Theorem 2.13 in (?)):

where by Lemma 6 each random function ${\mathbf{x}}^{(n)-}(t)$ is sampled i.i.d from the SDE:

we can now apply the weak law of large numbers (WLLN):

Now using that the log RN derivative between the estimator SDE and the true SDE is given by Girsanov’s theorem (?; ?; ?):

Thus using $\frac{1}{N}\mathcal{L}({\bm{u}})$ to denote the negative normalized log-likelihood we have shown the following pointwise (i.e. for ${\bm{u}}\in\mathcal{B}$) convergence in probability:

Then If $\mathcal{B}$ , is compact plus additional continuity and boundedness assumptions on $\mathcal{L}({\bm{u}})$ a stronger form of uniform convergence in probability can be attained:

where ${\bm{u}}^{*}_{N}$ is the MLE estimate of the dual drift at $N$ samples. Which implies:

See Chapter 4.5 of (?) for a more detailed discussion on the required assumptions of $\mathcal{L}$. ∎

Note that the above proof holds for discretised samples from the original SDE, however we have been unable to extend it to approximate sampling schemes such as EM. We believe that the following result motivates the possibility that Theorem 1 holds when the samples are obtained from a consistent scheme such as EM:

(Convergence of discrete time reversal) The time reversal of discrete Euler-Mayurama samples :

converges in probability to the solutions of the time reversed diffusion:

where ${\mathbf{x}}^{-}(t)={\mathbf{x}}^{+}(1-t)$

First let’s consider the continuous time step-wise approximation induced by the the EM samples:

From (?) it is a well known result that under the standard regularity assumptions on the drift ${\bm{u}}({\bm{x}},t)$ (i.e. Lipchitz continuity in $t$ and and ${\bm{x}}$) that the above approximation converges in probability to the SDE solution ${\mathbf{x}}^{+}(t)$ as $\Delta t\rightarrow 0$.

Observing that that the time reversal of the above corresponds to the approximation induced by the reverse samples we now consider the continuous time reversed approximation:

Using $\hat{{\mathbf{x}}}(t)$ converges in probability to ${\mathbf{x}}^{+}(t)$ we have that $\forall t\in,\;\epsilon>0$:

we carry out the following substitution $s=1-t$ then $\forall s\in,\;\epsilon>0$:

which completes the proof. Note that showing that reversing backwards samples reversed to a forward direction converges in probability follows the same sketch structure. ∎

A potential proof strategy would be to try and exploit the strong convergence properties of the Euler scheme to show that the stochastic integral from Theorem 1 also converges in the case of Euler samples. This remains an interesting question for future work.

Appendix F Towards a Finite Sample Analysis of Approximate IPFP Schemes

We use the term approximate IPFP schemes for methodologies such as the one we present in this paper (i.e. IPML) where steps 5,6 of IPFP (Algorithm 1) are replaced with inexact approximations. In this section we will present a rough sketch that takes the first step towards formally analysing approximate IPFP schemes in the finite sample / discretisation regime. This will serve to illustrate many of the challenges that still remain in this analysis as well as provide some initial results.

The exact IPFP at the $i^{th}$ iteration can be bounded from above as:

The approximate IPFP projection operators $\hat{{\mathcal{P}}}^{\pm}_{\pi}:{\mathcal{X}}\rightarrow{\mathcal{X}}$ are $K-$Lipchitz

and the approximate finite sample projection error can be bounded by a constant:

Then it follows that the approximate IPFP iteration error can be bounded from above by:

when ${\mathcal{P}}^{\pm}_{\pi}$ are non-expansive operators ($K=1$), and:

For notational simplicity we will relabel the half bridges in Algorithm 1 as projection operators ${\mathcal{P}}^{\pm}_{\pi}:{\mathcal{X}}\rightarrow{\mathcal{X}}$ in a function space ${\mathcal{X}}$. This leads to the following iterates:

And for the approximate IPFP we have the iterates:

Furthermore to simplify the analysis further we will consider the composition of the two projection operators (i.e. ${{\mathcal{P}}_{\pi_{0,1}}}={\mathcal{P}}^{+}_{\pi_{0}}\circ{\mathcal{P}}^{-}_{\pi_{1}}$ )combined into a single operator (for both exact and approximate projections):

Now we will proceed to analyse the first term :

Now we can expand the recurrence until the first iteration, yielding:

When the projection operator ${\hat{{\mathcal{P}}}_{\pi_{0,1}}}$ is non expansive this gives the bound:

Note that the above sketch is more of a strategy to outline the challenges and steps required for this analysis rather than a proof itself. We believe that the formal finite sample analysis on approximate IPFP schemes merits its own separate work and is thus not the focus of this work however in order to highlight its importance we have given a template towards this formal analysis following a strategy similar to (?). This analysis allows us to understand what conditions are needed to be shown and how they will affect the finite sample convergence rates:

We will now proceed to discuss each of the assumptions that are required as elements of the above analysis strategy.

The exact IPFP at the $i^{th}$ iteration can be bounded from above as:

The approximate IPFP projection operators $\hat{{\mathcal{P}}}^{\pm}_{\pi}:{\mathcal{X}}\rightarrow{\mathcal{X}}$ are $K-$Lipchitz . Specifically we are interested in the cases where they may be non-expansive and contractive.

This assumption is specific to the approximation method used to estimate the half-bridges. It requires going into rigorous details on the nature of the approximation and imposing suitable regularity assumptions. Not that ideally we want the operator to be contractive or at the very least no expansive since for Lipchitz constants greater than one the cumulative error would potentially accumulate exponentially.

It is interesting to note that in the contractive case we can attain the bound:

where the error term is constant with respect to the iteration number. Additionally the bound in Equation 60 also trivially applies.

The approximate finite sample projection error can be bounded by a constant:

This should typically be something that one shows however as seen in (?) it is something that is often assumed in these proof strategies. The reason for assuming the result is that the projection varies depending on how we approximate the KL minimisation’s, in our case we combine our proposed IPML with drift estimation via Gaussian processes. The error of the drift estimation (i.e. $\epsilon_{\Delta T,N}$) via this approach is not straightforward to obtain and is still an area of active research see (?; ?) for more details. As mentioned earlier this is not the focus of our work thus it is reasonable for us to assume that the drift estimation machinery proposed by prior authors is sound and should be possible to bound with a reasonable error.

In this section we provide the closed form predictive formulas used by our method to estimate the drift. Our GP formulation for the half bridge approximations yield the following mean per dimension $d$ estimates of the drift:

Similar to (?; ?), we are not making use of the predictive variancesUsing the predictive mean as an estimate for the drift can also be interpreted as a form of kernel ridge regression under the empirical risk minimisation framework. Instead, we simply use the predictive mean as an estimate of the drift and subsequently use that estimate to perform the EM method, thus effectively we could interpret this approach as a form of kernel ridge regression under the empirical risk minimisation framework.

A careful reader may enquire if the decoupled drift parametrisation we have used imposes any limitations on the generality of the drifts that our method can estimate in contrast to an approach that couples/correlates drift dimensions. While it may seem counter-intuitive, correlating the GP outputs (i.e. coupling the drift dimensions) is more restrictive than the decoupled approach we have taken. This point is discussed and motivated in (?; ?), where it is illustrated that the coupling of the drift outputs correspond to regularising the RKHS hypothesis space. So in fact the decoupled approach we take in this work is more flexible (less limited) than coupling the drift outputs as discussed in (?; ?). It is however important to note that the regularisation effects of coupling may be desirable in certain physical systems where we wish to impose a constraint, however this would come at a large computational cost. Finally note that our method parametrises the drift $[{\bm{u}}]_{i}=f_{i}({\bm{x}},t)$ with a function that depends on all dimensions of the dynamical process ${\mathbf{x}}(t)$, that is each drift coordinate depends on the entire input space and thus no simplifying assumptions have been made in this setting.

Appendix G Estimates Required by the Sinkhorn-Knop Algorithm

In order to adapt the Sinkhorn-Knop algorithm to a general prior we require the following considerations:

The Sinkhorn-Knop algorithm produces discrete potentials which can be interpolated using the logsumexp formula. However how do we go from these static potentials to time dependant trajectories ?

From (?) we can obtain expressions for the optimal drift as a function of the time extended potentials:

Note that in order to estimate the drift we would have to simulate the SDE prior every time we want to evaluate the drift which itself will be run in another SDE simulation to generate optimal trajectories.

Appendix H Approximations Required by (Pavon, Tabak, and Trigila 2018)

The approach in (?) iterates the following two objectives.

Where $\hat{\phi}_{0},\phi_{1}$ are parametric functions aimed at estimating the potentials, and:

Note that in Equation 86 the outer most integral is not taken with respect to a probability distribution and thus does not admit standard approximations. The authors in (?) propose the method of importance sampling (?), however this does not scale well to higher dimensions. This is contrasting to our approach where all integrals are expectations with respect to the empirical distribution and the SDEs being fitted.

Appendix I Experiments

In this section we provide in depth details regarding the parameter configurations of our approach plus our experimental setup. For all the experiments below, unless stated otherwise, we used the exponential kernel with a lengthscale set at the default value of 1. We ran IPML for 5 iterations with a discretization factor $\Delta_{t}=0.01$ and $\gamma=1$. The iteration number (5) was selected by observing that the algorithm converged after that number in all experiments considered.

For the unimodal boundary distributions we used:

Here we will illustrate a common failure case of the approach in (?). That is when the distributions $\pi_{0}$ and $\pi_{1}$ are distant from each other the methodology proposed in (?) breaks for suitable values of $\gamma$ (i.e. $\gamma=1,2,3...$). Using the same marginals $\pi_{0},\pi_{1}$ as in Experiment 1 we re-train the method by (?) with $\gamma=1$. In Figure 8 we can see the results of this experiment. We can see the marginals learned by DDSB collapse all the mass in a very small low density region for the true marginals. We found this phenomena to occur often and harder to overcome in 2 dimensional experiments and thus did not pursue further comparisons with this method. Note we normalised the delta spike of the DDSB marginals to have a range in $$so we could compare to the true density, its actual value is on the order of$10^{11}$ which confirms that the model is collapsing to a point mass (dirac delta function).

I.2 Well experiment

We used the following potential to model the double wells:

furthermore we used the boundary distributions:

we select this as the marginal prior since its samples cover the space over which the wells is defined well, and helps estimate the backwards drift of the prior more accurately. As mentioned in the paper, using different kernesl as well as lengthscales resulted in slightly different behavior on the well experiment where the trajectory would not split as illustrated in figure 9. The splitting pattern was obtained using an exponential kernel with a lower lengthscale ($0.25$) underlying the need for careful consideration of the kernel and its parameters.

We can observe that when using EQ kernel, the trajectories do not cross both passes but instead choose one. We conjecture that this is due to the fitted predictive means having a preference for simpler functions that send nearby particles in the same direction. We empirically verify this conjecture by experimenting with alternative kernels to EQ (i.e. exponential) that model a wider class of functions. We believe that these alternative kernels are less prone to discouraging splitting since, when viewed within the kernel ridge regression framework, they span an RKHS of functions that do not have the smoothness constraint of the EQ kernel.

I.3 Cell experiment

The data used for this experiment was taken directly from the github repository published by (?). PCA is first applied to the data and the first 5 components are selected. We used the code available at https://github.com/KrishnaswamyLab/TrajectoryNet/tree/master/TrajectoryNet to reproduce the performance of TrajectoryNet as well as the optimal transport baseline.

I.4 Motion experiment

For this experiment we use a second order SDE prior (i.e. Langevin dynamics). This can be encoded in first order SDE using the companion/controllable form of a dynamical system:

For simplicity we set the dampening factor $\beta$ to and $\gamma=0.5$. As for the spring matrix we explored the following decoupled and coupled forms:

Where $k$ defines the frequency of the solution via $\omega=\sqrt{k/m}$ in our setting we ignore the mass (imagine is at absorbed in the matrix ${\bm{K}}$ and work directly with setting $\omega$. Since the movement we select moves from rest to a high position and does not return to its starting point / or repeat itself we set $\omega\approx\pi$ that corresponds to $0.5$Hz which describes half a revolution. We visually confirmed with the true trajectories that this seems like a reasonable prior.

For the Brownian motion prior that we use as a baseline comparison in Figure I.4 we explored several values of $\gamma$ and selected the best performing one (visually) of $\gamma=0.3$. Note we do pick a smaller value of $\gamma$ for the Brownian motion and this is partly to compensate that the Langevin prior trajectories will always look smooth by construction thus we pick a not so noisy Brownian motion that still gives good results.

Note that the volatility term here is singular and in our mocap application setting it is exactly the vector:

Where the dimension $d$ is given by the number of sensors we use to fit the motion. A careful reader may wonder weather the SBP machinery still applies to such a sparsely structured diffusion and in fact it does. Following Theorem 4 of (?) we can see how the Radon-Nikodym derivative is finite and can be expressed in an almost identical fashion to the original controlled SBP formulation and thus justifying the existence of the SBP in this scenario as well as the application of the IPML algorithm and its derivatives.

I.5 Computational resources

The main computational cost of IPML comes from fitting a Gaussian Process to model the drift (DriftFit). Assuming the cardinality of $|\pi_{0}|=|\pi_{1}|=N$ and a discretization factor $\Delta_{t}$, at each iteration, IPML requires fitting a Gaussian Process on $\frac{N}{\Delta_{t}}$ samples. It is a well-known fact that GPs have a cubic time complexity making the costs of the DriftFit subroutine $O((\frac{N}{\Delta_{t}})^{3})$. This could be scaled down by using GP approximations (Nystrom) and was successful in our preliminary results. The computational costs of the SDESolve sub-routine in comparison is $O((\frac{N}{\Delta_{t}})^{2})$ due to the GP predictions costs. Finally, the memory complexity of the algorithm is $O((\frac{N}{\Delta_{t}})^{2})$ due to the GP fitting in DriftFit. The running time of IPML on the machine described in the section below for the well experiment is around 5 minute. For the cell experiment, the running time is between 1 and 2 hours depending on the discretization factor chosen.

The experiments are performed on Compute Canada clusters. Specifically, they ran on a CPU cluster composed of Intel CPUs. Each node contains 20 Intel Skylake cores (2.4GHz, AVX512), for a total of 40 cores per node and 202GB of RAM. The computational costs for the toy experiments (Well and 1D) was quite small, however, the RAM consumption of the cell experiment was significant due to the use of a Gaussian Process.