Multi-marginal Schrodinger bridges

TL;DR

This paper extends the classical Schrödinger Bridge Problem to second-order particle dynamics with partial temporal data, formulating it as an optimal control problem with a Fisher information term, and connecting it to measure-valued splines in Wasserstein space.

Contribution

It introduces a novel second-order Schrödinger bridge formulation with partial observations, incorporating Fisher information and linking to measure-valued splines in Wasserstein space.

Findings

Derived a time-symmetric optimal control formulation with Fisher information.

Established a connection to measure-valued splines and Sinkhorn-like algorithms.

Demonstrated potential applications in signal processing and data science.

Abstract

We consider the problem to identify the most likely flow in phase space, of (inertial) particles under stochastic forcing, that is in agreement with spatial (marginal) distributions that are specified at a set of points in time. The question raised generalizes the classical Schrodinger Bridge Problem (SBP) which seeks to interpolate two specified end-point marginal distributions of overdamped particles driven by stochastic excitation. While we restrict our analysis to second-order dynamics for the particles, the data represents partial (i.e., only positional) information on the flow at {\em multiple} time-points. The solution sought, as in SBP, represents a probability law on the space of paths this closest to a uniform prior while consistent with the given marginals. We approach this problem as an optimal control problem to minimize an action integral a la Benamou-Brenier, and derive a time-symmetric formulation that includes a Fisher information term on the velocity field. We underscore the relation of our problem to recent measure-valued splines in Wasserstein space, which is akin to that between SBP and Optimal Mass Transport (OMT). The connection between the two provides a Sinkhorn-like approach to computing measure-valued splines. We envision that interpolation between measures as sought herein will have a wide range of applications in signal/images processing as well as in data science in cases where data have a temporal dimension.

Equations

$${\rho }_1(x_1)\ne {\int }_{{\mathbb{R}}^n}q(t_0,x_0,t_1,x_1){\rho }_0(x_0)d{x}_0,$$

$$q(s,x,t,y)=(2\pi {)}^{-n/2}(t-s{)}^{-n/2}\exp \left(-\frac{1}{2}\frac{\Vert x-y{\Vert }^2}{(t-s)}\right)$$

$$dx$$

$$dv$$

$$\frac{\partial \mu }{\partial t}+v\cdot {\nabla }_x\mu -\frac{1}{2}{\Delta }_v\mu =0,$$

$$H(\mathcal{P},\mathcal{Q})=\int d\mathcal{P}\log \frac{d\mathcal{P}}{d\mathcal{Q}}$$

$$dx$$

$$dv$$

$$H(\mathcal{P},\mathcal{Q})={\mathbb{E}}_{\mathcal{P}}\left\{{\int }_0^1\Vert a(t){\Vert }^{2}dt\right\}={\int }_0^1\int \Vert a{\Vert }^2\mu dxdvdt.$$

$$\min$$

$${\int }_0^1\int \Vert a{\Vert }^2\mu dxdvdt$$

$${\int }_0^1\int ⟨\hat{a},{\nabla }_v\log \mu ⟩\mu dxdvdt=\int \{{\mu }_1\log {\mu }_1-{\mu }_0\log {\mu }_0\}dxdv,$$

$$\min$$

$$\inf$$

$$dx$$

$$dv$$

$$\min$$

$$\min J(\pi ):=\sum _{i=0}^{N-1}KL({\pi }_{i,i+1}\mid e^{-C_{i,i+1}/ϵ})$$

$$\int {\pi }_{i,i+1}d{x}_{i+1}d{v}_{i+1}={\mu }_i,i=0,\dots ,N-1$$

$$\int {\pi }_{i,i+1}d{x}_{i}d{v}_i={\mu }_{i+1},i=0,\dots ,N-1$$

$$\int {\mu }_{i}d{v}_i={\rho }_i,i=0,\dots ,N$$

$$C_{i,i+1}(x_i,v_i,x_{i+1},v_{j+1})$$

$$K_0$$

$$K_N$$

$$K_i$$

$${\pi }^n=P_{K_{i_n}}^{KL}({\pi }^{n-1}),n=1,2,3,\dots$$

$$P_K^{KL}(\bar{\pi }):={\mathrm{argmin}}_{\pi \in K}KL(\pi \mid \bar{\pi }).$$

$$P_{K_0}:$$

$$P_{K_N}:$$

$${\pi }_{i-1,i}$$

$${\pi }_{i,i+1}$$

Full text

1 Introduction

In 1931/32, in an attempt to gain insights into the stochastic nature of quantum mechanics, Schrödinger [22, 23] raised the following question regarding a system of a large number of classical independent identically distributed (i.i.d.) Brownian particles. He hypothesized that this “cloud” of particles is observed to have (empirical) distributions $\rho_{0}(x_{0})$ and $\rho_{1}(x_{1})$ at two points in time $t_{0}=0$ and $t_{1}=1$, respectively, and further that $\rho_{1}(x_{1})$ differs from what is dictated by the law of large numbers, i.e., that

[TABLE]

where

[TABLE]

denotes the Brownian transition probability kernel. Schrödinger then sought to find the “most likely” evolution for the cloud of particles to have transitioned from $\rho_{0}$ to $\rho_{1}$. In the language or large deviation theory (which was not in place at the time), Schrödinger’s question amounts to seeking a probability law on the path space that is in agreement with the two marginals while being the closest to the Brownian prior in the sense of relative entropy [12]. The solution is known as the Schrödinger bridge since the law “bridges” the two given end-point marginal distributions.

Renewed interest in the Schrödinger Bridge Problem (SBP) has been fueled by its connections to the Monge-Kantorovic Optimal Mass Transport (OMT) and a wide range aplications in image analysis, stochastic control, and physics [18, 19, 6, 13, 4, 9, 21]. More specifically, SBP, seen as a suitable regularization of OMT, provides a natural model for uncertainty in the transport of distributions as well as a valuable computational tool for interpolating distributional data.

In this work we consider a natural generalization of the Schrödinger bridge theory to address the situation where the data consist of possible partial marginal distributions at various points in time. Thus, we postulate a similar experiment with stochastic particles. However, in contrast to the standard SB theory, we concieve these particles to obey second order stochastic differential equations with Brownian stochastic forcing that accounts for random acceleration along trajectories in phase space (see Section 2). This new setting connects with recent results on measure-valued splines [3, 1, 8], which are general notions of splines on the Wasserstein space of measures. In this short paper, we provide a summary of the theory. A more detailed account will appear in a forthcoming publication (in preparation).

2 Multi-marginal Schrödinger bridges for inertial particles

Suppose we are given a large number of independent inertial particles driven by white noise, that is, they follow the dynamics

[TABLE]

where $dw$ denotes standard Brownian motion and $(x,\,v)$ is the phase space ($x$ denoting position and $v$ velocity). The flow of probability densities $\mu_{t}(x,v)$ in phase space obeys the Fokker-Planck equation

[TABLE]

with initial condition $\mu_{0}$ at time $t=0$. The law of large numbers dictates that, when the number of particles is large enough, the distributions of the particles will be closed to $\mu_{t}$. The data of the problem we consider will, as in the standard SB problem, be inconsistent with the Fokker-Planck equation, and viewed as a “atypical/rare” event. In standard SB setting, where only two end-point marginals are specified, the “most likely” evolution amounts to an adjustment of the Fokker-Planck equation by adding a suitable drift term to match the two marginals. For inertial particles driven by white noise, the generator is hypoelliptic [16] and the SB theory carries over to matching marginals in phase space.

Throughout, we suppose that we have only access to position $x$ and that empirical marginals $\rho_{0}=\operatorname{Proj}_{x}\mu_{0},\rho_{1}=\operatorname{Proj}_{x}\mu_{1}$ represent projections, accordingly. The extra degree of freedom, since $\mu$’s are partially specified, make the corresponing multi-marginal problem nontrivial. Specifically, we seek the most like paths the particles have taken that match positional distributions $\rho_{0},\rho_{1},\cdots,\rho_{N}$ at times $0=t_{0}<t_{1}<\cdots<t_{N}=1$. To this end, we let ${\mathcal{Q}}$ be the law of (1) (on path space) and we let ${\mathcal{P}}$ be the any other law. We seek the minimizer of

[TABLE]

over all laws ${\mathcal{P}}$ that are consistent with the marginals $\rho_{0},\rho_{1},\cdots,\rho_{N}$.

To guarantee the boundedness of the relative entropy $H$ between ${\mathcal{P}}$ and the prior process ${\mathcal{Q}}$, ${\mathcal{P}}$ has to be of the form

[TABLE]

where $a$ is a suitable drift that may depend on the current and past values of the process state. Invoking Girsanov’s theorem [14, 17, 11, 7], we obtain that

[TABLE]

Thus, we arrive at the optimal control formulation

[TABLE]

The difference to the standard Schrödinger bridge problem lies in the constraint (5d); multiple marginals are partially specified. The existence and uniqueness of the solution follow from the fact that it is a strongly convex optimization problem on path space measures. The argument is similar to the standard argument in SB theory [12, 18, 19, 5] and will be presented in an extended version of the paper. By utilizing the Fokker-Planck equation, we can rewrite the above as

[TABLE]

Here $\int\mu_{t_{i}}(x,v)dv=\rho_{i}(x)$ since $\operatorname{Proj}_{x}(\mu_{t_{i}})=\rho_{i}$.

Now let $\hat{a}=a-\frac{1}{2}\nabla_{v}\log\mu$, then the diffusion term in (6c) can be absorbed into the convection terms, and then the cost becomes

[TABLE]

Direct calculation yields

[TABLE]

which only depends on the two end distributions. Thus, we need to consider

[TABLE]

A similar formulation with a Fisher information term has been studied in [25, 6, 13, 20] for standard Schrödinger bridge problems. However, in the standard setting, the term $\int\{\mu_{1}\log\mu_{1}\!-\!\mu_{0}\log\mu_{0}\}dxdv$ can be dropped since the full-state marginal distributions are already specified. It important to note that, compared to (6), (8) is time symmetric.

3 Connections to measure-valued splines

A variational formulation [15] of splines going through $\{x_{1},x_{2},\ldots,x_{N}\}$ in Euclidean space is given by

[TABLE]

where the minimization is taken over all twice-differentiable trajectories that satisfy the constraints. This formation has been generalized to the Wasserstein space of measures [3, 1]. In particular, the fluid dynamic formulation for measure-valued splines in [3] with marginals $\rho_{0},\rho_{1},\cdots,\rho_{N}$ at $0=t_{0}<t_{1}<\cdots<t_{N}=1$ reads

[TABLE]

We note that the above formulation is almost the same as (6) except for a missing diffusion term in the constraint. This resembles the relation between standard Schrödinger bridges and optimal mass transport [18, 19, 5].

Indeed, a zero-noise limit argument follows. If we replace the dynamics (1) of the inertial particles by

[TABLE]

then the multi-marginal SB problem becomes

[TABLE]

The “slowed down” formulation (11) reduces to (6) when we take the limit $\epsilon\rightarrow 0$. Therefore, we establish the measure-valued spline as a zero-noise limit of a multi-marginal SB, and SB as a regularized version of measure-valued spline. Rigorous proof of these conclusions will be presented in a forthcoming paper.

4 Algorithms

The Sinkhorn [24, 10] algorithm is a natural iterative scheme in SB problems. Due to its efficiency and simplicity, it has became a workhorse for data science applications of optimal transport [10]. In this section, we develop a Sinkhorn-type algorithm for multi-marginal Schrödinger bridge problems. The relation established in Section 3 implies that the same algorithm can also be used to approximate measure-valued splines.

Using a measure decomposition argument, the Schrödinger problem (11) can be rewritten as

[TABLE]

where $KL(\alpha|\beta)=\int\alpha\log\frac{\alpha}{\beta}-\alpha+\beta$ and

[TABLE]

The optimization variables are joint distributions on the consecutive time points over the phase space. Since the cost is a summation of relative entropies and the constraints are convex, a natural algorithm is that of Bregman projections [2].

Define convex constraint sets

[TABLE]

for $i=1,\ldots,N-1$, then the Bregman iterative projection becomes

[TABLE]

where $i_{n}$ enumerates $\{0,1,\ldots,N\}$ repeatedly. The projection operator is

[TABLE]

These projections can be derived via Lagrangian method. The projections to $K_{0},K_{N}$ are easier and the projections to $K_{1},\ldots,K_{N-1}$ are more involved. Specifically,

[TABLE]

whereas for $P_{K_{i}},~{}i=1,\ldots,N-1$:

[TABLE]

In real implementation, we need to discretize the phase space over a grid. After discretization, the algorithm only involves matrix multiplication, pointwise-division, multiplication, square root, and therefore can be parallelized easily. The linear convergence rate is guarantee by the property of Bregman projections [2]. Our algorithm should be compared to that developed in [1]. A major difference is that our algorithm doesn’t require discretization over the time domain.

5 Conclusion

We considered a natural extension of the Schrödinger bridge problems to multi-marginal partially observable setting. We focused on inertial particles, but more general dynamics can be examined similarly. We discussed the physical meaning, stochastic control formulation and several other aspects of the problems. Just like in the standard SB problem, it has a natural relation to the measure-valued spline theory. An efficient algorithm was also developed, which makes ready for possible applications. We envision that this line of research is going to spark interest in optimal transport theory and application with multiple time points.