SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Deep Multi-Marginal Momentum Schrödinger Bridge
Tianrong Chen, Guan-horng Liu, Molei Tao, Evangelos A. Theodorou
Georgia Institute of Technology, USA
{tianrong.chen,ghliu,mtao,evangelos.theodorou}@gatech.edu
It is a crucial challenge to reconstruct population dynamics using unlabeled samples from distributions at coarse time intervals. Recent approaches such as flow-based models or Schrödinger Bridge (SB) models have demonstrated appealing performance, yet the inferred sample trajectories either fail to account for the underlying stochasticity or are unnecessarily rigid. In this article, We extend the approach in [1] to operate in continuous space and propose Deep Momentum Multi-Marginal Schrödinger Bridge (DMSB), a novel computational framework that learns the smooth measure-valued spline for stochastic systems that satisfy position marginal constraints across time. By tailoring the celebrated Bregman Iteration and extending the Iteration Proportional Fitting to phase space, we manage to handle high-dimensional multi-marginal trajectory inference tasks efficiently. Our algorithm outperforms baselines significantly, as evidenced by experiments for synthetic datasets and a real-world single-cell RNA sequence dataset. Additionally, the proposed approach can reasonably reconstruct the evolution of velocity distribution, from position snapshots only, when there is a ground truth velocity that is nevertheless inaccessible.
We consider the multi-marginal trajectory inference problem, which pertains to elucidating the dynamics and reactions of indiscernible individuals, given static snapshots of them taken at sporadic time points. Due to the inability of tracking each individual, one considers the evolution of the statistical distribution of the population instead. This problem received considerable attention, and associated applications appear in various scientific areas such as estimating cell dynamics [2, 3], predicting meteorological evolution [4], and medical healthcare statistics tracking [5]. [6, 7] constructed an energy landscape that best aligned with empirical observations using neural network. [8, 9] learn regularized Neural ODE [10] to encode such potential landscape. Notably, in the aforementioned work, the trajectory of samples is represented in a deterministic way. In contrast, [11, 12] employ Schrödinger Bridge (SB) to determine the most likely evolution of samples between marginal distributions when individual sample trajectories are also affected by environmental stochasticity. Yet, these approaches scale poorly w.r.t. the state dimension due to specialized neural network architectures and computational frameworks.
SB can be viewed as a solution to the entropy-regularized optimal transport problem. SB seeks a nonlinear SDE that yields a straight path measure between two arbitrary distributions. The straightness is implied by achieving optimality of minimizing transportation costs (i.e. 2-Wasserstein distance  $(W_{2}))$ . We note SB is often related to Score-based Generated Model (SGM), both of which can be used for generative modeling by constructing certain Stochastic Differential Equation (SDE) that links data distribution and a tractable prior distribution (i.e. 2 marginals). SGM accomplishes the generative task by first diffusing data to prior through a pre-specified linear SDE, during which a neural network is also learned to approximate the score function. Then this score approximator is used to reverse this diffusion process, and consequently establish the generation. Critically-damped
Table 1: Comparison between different models in terms of optimality and boundary distributions  $p_0$  and  $p_1$ . Our DMSB extends standard SB, which generalizes SGM beyond Gaussian priors, to phase space, similar to CLD. However, unlike CLD, DMSB jointly learns the phase space distributions, i.e.,  $p_{\theta}(x,v) = p_{\mathcal{A}}(x)q_{\theta}(v|x)$  and  $p_{\phi}(x,v) = p_{\mathcal{B}}(x)q_{\phi}(v|x)$ . In other words, DMSB infers the underlying phase state dynamics given only state distributions.
Models	Optimality	p0(·)	p1(·)
SGM [20]	X	pA(x)	N(0, Σ)
CLD [13]	X	pA(x) ⊗ N(0, Σ)	N(0, Σ) ⊗ N(0, Σ)
SB [14]	↓W2→ kinks	pA(x)	pB(x)
DMSB (ours)	↓W2→ smooth	pA(x)qθ(v|x)	pB(x)qφ(v|x)
Langevin Diffusion (CLD) [13] extends the SGM SDE to the phase space by introducing an auxiliary velocity variable with a tractable Gaussian distribution at both the initial and terminal time. The resulting trajectory in the position space becomes smoother, as stochasticity is only injected into the velocity space, and the empirical performance and sample efficiency are enhanced due to the structure of the critical damped SDE. The connection between SGM and SB has been elaborated in [14, 15] and scalable mean matching Iterative Proportional Fitting algorithm (IPF) is proposed to estimate SB efficiently in high dimensional cases. Applications of SB, such as image-to-image transformation [16, 17], RNA trajectory inference [11], solving Mean Field Game[18], Riemannian interpolation [19], demonstrate the effectiveness of SB in various domains.
In this work, we start with SB in phase space (termed momentum SB, mSB in short), and then further investigate mSB with multiple empirical marginal constraints present in the position space, which was formulated as multi-marginal mSB (mmmSB) in [1]. This circumvents the need for expensive space discretization which does not scale well to high dimensions. We also address the challenge of intricate geometric averaging in continuous space setup by strategically partitioning and reorganizing the constraint sets. Furthermore, we enhance the algorithm's computational efficiency by incorporating the method of half-bridge IPF. The optimality of transportation cost in SB leads to straight trajectories, and if one solves N 2-marginal SB problems and connect the resulting trajectories to match N+1 marginals, the connected trajectories will have kinks at all connection points. On the contrary, in mmmSB, the optimality of transportation cost leads to a smooth measure-spline over the state space that also interpolates the empirical marginals. Therefore, this approach is highly suitable for problems originated from physical systems and/or those that should have smooth trajectories, such as trajectory inference in single-cell RNA sequencing. Our research will emphasize on solving mmmSB efficiently in high-dimensions (thus the approach will differ from that in the seminal work [1]; see Sec.4). The differences between our algorithm and prior work are demonstrated in Table.1, and the main contributions of our work are fourfold:
- We extend the mean matching IPF to phase space allowing for scalable mSB computing.
- We introduce and tailor the Bregman Iteration [21] for mmmSB which makes it compatible with the phase space mean matching objective, thus the efficient computation is activated for high dimensional mmmSB.
- We show how to overcome the challenge of sampling the velocity variable when it is not available in training data, which enhances the applicability of our model.
- We show the performance of proposed algorithm DMSB on toy datasets which contains intricate bifurcations and merge. On realistic high-dimension (100-D) single-cell RNA-seq (scRNA-seq) datasets, DMSB outperforms baselines by a significant margin in terms of the quality of the generated trajectory both visually and quantitatively. We show that DMSB is able to capture reasonable velocity distribution compared with ground truth while other baselines fail.
2.1 Dynamical Schrödinger Bridge problem
Dynamical Schrödinger Bridge problem has been extensively studied in the past few decades. The objective of the SB problem is to solve the following optimization problem:
$$
\min  _ {\pi \in \Pi \left(\rho_ {0}, \rho_ {T}\right)} D _ {K L} (\pi | | \xi), \tag {1}
$$
where  $\pi \in \Pi (\rho_0,\rho_T)$  belongs to a set of path measures with its marginal densities at  $t = 0$  and  $T$  being  $\rho_0$  and  $\rho_{T}$ .  $\xi$  is the reference path measure (i.e., [14] sets  $\xi$  as Wiener process from  $\rho_0$ ). The optimality of the problem (1) is characterized by a set of PDEs (3).
Theorem 2.1 ([22]). The optimal path measure  $\pi$  in the problem (1) is represented by forward and backward stochastic processes
$$
\mathrm {d} \mathbf {x} _ {t} = \left[ 2 \nabla_ {\mathbf {x}} \log \Psi_ {t} \right] \mathrm {d} t + \sqrt {2} \mathrm {d} \mathbf {w} _ {t}, \quad \mathbf {x} _ {0} \sim \rho_ {0}, \tag {2a}
$$
$$
\mathrm {d} \mathbf {x} _ {t} = \left[ - 2 \nabla_ {\mathbf {x}} \log \widehat {\Psi} _ {t} \right] \mathrm {d} t + \sqrt {2} \mathrm {d} \widehat {\mathbf {w}} _ {t}, \mathbf {x} _ {T} \sim \rho_ {T}. \tag {2b}
$$
in which  $\Psi, \widehat{\Psi} \in C^{1,2}$  are the solutions to the following coupled PDEs,
$$
\frac {\partial \Psi_ {t}}{\partial t} = - \Delta \Psi_ {t}, \quad \frac {\partial \widehat {\Psi} _ {t}}{\partial t} = \Delta \widehat {\Psi} _ {t} \tag {3}
$$
$$
s. t. \Psi (0, \cdot) \widehat {\Psi} (0, \cdot) = \rho_ {0} (\cdot), \Psi (T, \cdot) \widehat {\Psi} (T, \cdot) = \rho_ {T} (\cdot),
$$
The stochastic processes of SB in (2a) and (2b) are equivalent in the sense of  $\forall t\in [0,T],p_t^{(2a)}\equiv p_t^{(2b)}\equiv p_t^{SB}$ . Here  $p_t^{SB}$  stands for the marginal distribution of SB at time  $t$ , which also represents the marginal density of stochastic process induced by either of Eq.2. The potentials  $\Psi_{t}$  and  $\widehat{\Psi}_{t}$  explicitly represent the solution of Fokker-Plank Equation (FPE) and Hamilton-Jacobi-Bellman equation (HJB) after exponential transform [14] where FPE describes the evolution of samples density and HJB represents for the optimality of Eq.1. Furthermore, the marginal density also obeys a factorization of  $p_t^{SB} = \Psi_t\widehat{\Psi}_t$ . Such rich structures of SB will later on be used to construct the log-likelihood objective (Thm.B.1) and Langevin sampler for velocity (\$4.4).
To solve SB, prior work have primarily used the half-bridge optimization technique, also known as Iterative Proportional Fitting (IPF), in which one iteratively solves the optimization problem with one of the two boundary conditions [14, 15, 23],
$$
\pi^ {(d + 1)} := \underset {\pi \in \Pi (\cdot , \rho_ {1})} {\arg \min } D _ {K L} (\pi | | \pi^ {(d)}) \quad \rightleftharpoons \quad \pi^ {(d + 2)} := \underset {\pi \in \Pi (\rho_ {0}, \cdot)} {\arg \min } D _ {K L} (\pi | | \pi^ {(d + 1)}) \tag {4}
$$
with initial path measure  $\pi^{(0)}\coloneqq \xi$ . By repeatedly iterating over aforementioned optimizations until the algorithm converges, the SB solution will be attained as  $\pi^{SB}\equiv \lim_{d\to \infty}\pi^{(d)}$  [24]. In addition, [25] shows that the drift term in SB problem can also be interpreted as the solution Stochastic Optimal Control (SOC) problem by having optimal control policy  $\mathbf{z}^* = 2\nabla_{\mathbf{x}}\log \Psi (t,\mathbf{x}_t)$ :
$$
\mathbf {z} ^ {*} (\mathbf {x}) \in \underset {\mathbf {z} \in \mathcal {Z}} {\arg \min} \mathbb {E} \left[ \int_ {0} ^ {T} \frac {1}{2} \| \mathbf {z} _ {t} \| ^ {2} \mathrm {d} t \right] \quad s. t \quad \left\{ \begin{array}{l l} \mathrm {d} \mathbf {x} _ {t} = \mathbf {z} _ {t} \mathrm {d} t + \sqrt {2} \mathrm {d} \mathbf {w} _ {t} \\ \mathbf {x} _ {0} \sim \rho_ {0}, \quad \mathbf {x} _ {1} \sim \rho_ {T}. \end{array} \right.
$$
This formulation will be used later on for constructing phase space likelihood objective function in §3. Regarding solving the half-bridge problem, abundant results exist in the literature for the vanilla SB described above [14, 15, 23], but we will be solving a different SB problem; see Prop.4.1 for formulation and §.4 for a solution.
2.2 Bregman Iterations for Multiple Constraints
Bregman iteration [21] can be viewed as a multiple marginal generalization of IPF, and it is widely used to solve entropy regularized optimal transport problem [1] with multiple constraints. The algorithm can efficiently solve problems in the form of,
$$
\inf  _ {\pi \in \mathcal {K}} K L (\pi | \xi),
$$
Figure 1: A summary of various SB problems and corresponding algorithms. The toy example in the 3rd row illustrates that vanilla SB determines 'straight' paths (modulo fluctuations due to noise) between pairwise empirical marginals, while our multi-marginal momentum SB approach establishes a smooth measure-spline between marginals in the position space (albeit still stochastic, the path is smooth between any pair of adjacent 2 marginals, because noise is added to velocity, and the path is also smooth across different pairs of adjacent 2 marginals per design.
where  $\mathcal{K}$  is the intersection of multiple closed convex constraint sets  $\mathcal{K}_l$ :  $\mathcal{K} = \cap_{l=1}^{L} \mathcal{K}_l$ . Bregman Projection (BP) is defined as optimization w.r.t one of the constraint  $\mathcal{K}_l$ ,
$$
P_{\mathcal{K}_{l}}^{KL}(\xi):= \operatorname *{arg  min}_{\pi \in \mathcal{K}_{l}}KL(\pi |\xi),
$$
and  $d$ -th Bregman Iteration (BI) is recursively computing BP over all the constraints in  $\mathcal{K}$ :
$$
\forall 0 <   n \leq L, \quad \pi^ {(d, n)} := P _ {\mathcal {K} _ {l} ^ {n}} ^ {K L} (\pi^ {(d, n - 1)}),
$$
The initial condition for  $(d + 1)$ -th BI is  $\pi^{(d + 1,0)} = \pi^{(d,L)}$ . Under certain conditions (see e.g., [24]), one has that  $\pi^{(d,L)}$  converges to the unique solution:
$$
\pi^ {(d, L)} \rightarrow P _ {\mathcal {K}} ^ {K L} (\xi) \quad \mathrm {a s} \quad d \rightarrow + \infty
$$
Remark 2.2. One BI traverses all constraints via multiple BPs, and each BP solves an optimization problem with one constraint. One can notice that the BI will become the aforementioned IPF procedure solving SB problem (1) by defining  $L = 2$ ,  $\mathcal{K}_1 = \Pi(\rho_0, \cdot)$ ,  $\mathcal{K}_2 = \Pi(\cdot, \rho_1)$ .
Table 2: Mathematical notation.
Notation	Definition	Notation	Definition
x	position variable	ρ	position distribution ρ(x)
v	velocity variable	γ	velocity Distribution γ(v)
m	concatenation of [x, v]T	μ	distribution of μ(x, v)
3 Momentum Schrödinger Bridge
We first describe how to conduct half-bridge IPF training in the phase space, which can be used to solve momentum SB (mSB) problem with two marginals constraints. This scalable phase space half-bridge technique will then be applied to multi-marginal cases (Sec.4). Fig.1 demonstrates how we develop an algorithm based on [14]. Notations used in following sections are listed in Table.2. mSB extends SB problem to phase space, which consists of both position and velocity.
We will first consider boundary distributions that depend on both  $\mathbf{x}$  and  $\mathbf{v}$ , although eventually we will use this as a module to find transport maps between two distributions that only depend on position  $\mathbf{x}$ , as velocity  $\mathbf{v}$  is an auxiliary variable artificially introduced for obtaining smooth transport. Conceptually, as an entropy regularized optimal transport problem, SB tries to obtain the straightest path between empirical marginals of positions  $\mathbf{x}$  with additive noise, but mSB aims at finding the smooth interpolation between empirical marginals of  $\mathbf{x}$  [26] conditioned on boundary velocity distributions (see Fig.1). Such smooth measure-valued splines in the position space are obtained by the optimization problem in the phase space [1]:
$$
\min  _ {\pi \in \Pi (\mu_ {0}, \mu_ {T})} K L (\pi | \xi) \quad s. t \quad \pi = \operatorname {L a w} (\mathbf {x}, \mathbf {v}): \underbrace {\left( \begin{array}{c} \mathrm {d} \mathbf {x} _ {t} \\ \mathrm {d} \mathbf {v} _ {t} \end{array} \right)} _ {\mathrm {d} \mathbf {m} _ {t}} = \underbrace {\left( \begin{array}{c} \mathbf {v} _ {t} \\ \mathbf {0} \end{array} \right)} _ {\mathbf {f} (\mathbf {v}, t)}   \mathrm {d} t + \underbrace {\left( \begin{array}{c c} \mathbf {0} & \mathbf {0} \\ \mathbf {0} & g _ {t} \end{array} \right)} _ {\mathbf {g} (t)} \underbrace {\left( \begin{array}{c} \mathbf {0} \\ \mathbf {z} _ {t} \end{array} \right)} _ {\mathbf {Z} (t)}   \mathrm {d} t + \underbrace {\left( \begin{array}{c c} \mathbf {0} & \mathbf {0} \\ \mathbf {0} & g _ {t} \end{array} \right)} _ {\mathbf {g} (t)}   \mathrm {d} \mathbf {w} _ {t},
$$
Similar to Theorem 2.1, one can derive a set of PDEs using the potential functions  $\Psi(t,\mathbf{x},\mathbf{v})$  and  $\widehat{\Psi}(t,\mathbf{x},\mathbf{v})$ , and subsequently apply IPF procedure to solve the problem. The formulation of the phase space PDE can be found in Appendix.B.2. Such PDE representation of mSB results in a straightforward yet innovative log-likelihood training that enables efficient optimization of the IPF.
Proposition 3.1 (likelihood bound). The half-bridge IPF in phase space
$$
\pi^ {(d + 1)} := \underset {\pi \in \Pi (\mu_ {0}, \cdot)} {\arg \min } D _ {K L} (\pi | | \pi^ {(d)}) \quad \leftrightarrows \quad \pi^ {(d + 2)} := \underset {\pi \in \Pi (\cdot , \mu_ {T})} {\arg \min } D _ {K L} (\pi | | \pi^ {(d + 1)})
$$
represents the bound of the likelihood and gives approximate likelihood training:
$$
\mathbf {Z} _ {t} := \underset {\mathbf {Z} _ {t}} {\arg \min } - \log p (\mathbf {m} _ {0}, 0) \quad \rightleftharpoons \quad \widehat {\mathbf {Z}} _ {t} := \underset {\widehat {\mathbf {Z}} _ {t}} {\arg \min } - \log p (\mathbf {m} _ {T}, T).
$$
$$
{ w h e r e } { \log p ( \mathbf { m } _ { 0 } , 0 ) \propto \int _ { 0 } ^ { T } \mathbb { E } _ { \hat { \mathbf { m } } _ { t } } \left[ \frac { 1 } { 2 } \| \widehat { \mathbf { z } } _ { t } + \mathbf { z } _ { t } - g \nabla _ { \mathbf { v } } \log \hat { p } _ { t } \| ^ { 2 } \right] \mathrm { d } t . }
$$
$$
a n d \widehat {\mathbf {m}} _ {t} s a m p l e s f r o m: \quad \mathrm {d} \widehat {\mathbf {m}} _ {t} = \left[ \boldsymbol {f} - \mathbf {g} \widehat {\mathbf {Z}} _ {t} \right] \mathrm {d} t + \mathbf {g} (t) \mathrm {d} \mathbf {w} _ {t}, \quad \widehat {\mathbf {m}} _ {T} \sim \mu_ {T} \tag {5}
$$
$\widehat{\mathbf{Z}}_t \triangleq \left( \begin{array}{c} \mathbf{0} \\ \widehat{\mathbf{z}}_t \end{array} \right)$  and  $\widehat{p}_t$  is the density of path measure induced by eq.5 at time  $t$ . A similar result for  $\log (\mathbf{m}_T, T)$  can be obtained in a similar derivation.
Proof. See Appendix B.1.
Remark 3.2. After optimizing  $\widehat{\mathbf{Z}}_t$ , the reference path measure becomes eq.5, which implies  $\pi \in \Pi(\cdot, \mu_T)$ , i.e., the constraint in half-bridge IPF is satisfied. A path measure  $\pi$  is induced by either  $\mathbf{Z}_t$  or  $\widehat{\mathbf{Z}}_t$ . As being mentioned in Remark.2.2. One half-bridge IPF is basically one BP and one IPF is one BI. Prop.3.1 provides a convenient way to perform one BP in the form of  $\pi := \arg \min_{\pi \in \mathcal{K}_l} D_{KL}(\pi ||\bar{\pi})$  by maximizing log-likelihood given constraint  $\mathcal{K}$  and reference path measure  $\bar{\pi}$ .
Prop.3.1 provides an alternative way to conduct the BI which will be heavily used in mmmSB §3, and it is computationally efficient after parameterizing and discretization (§4.4).
4 Deep Momentum Multi-Marginal Schrödinger Bridge
We first state the problem formulation of momentum multi-marginal Schrödinger Bridge (mmmSB). Different from previous two marginals case, we consider the scenario where  $N + 1$  probability measures  $\mu_{t_i}$  are lying at time  $t_i$ . In addition, velocity distributions are not necessarily known.
Proposition 4.1 ([1]). The dynamical mmmSB with multiple marginal constraints reads:
$$
\min  _ {\pi} \mathcal {J} (\pi) := \sum_ {i = 0} ^ {N - 1} K L \left(\pi_ {t _ {i}: t _ {i + 1}} \mid \xi_ {t _ {i}: t _ {i + 1}}\right), \quad s. t \quad \pi \in \mathcal {K} := \cap_ {i = 0} ^ {N} \mathcal {K} _ {t _ {i}} \tag {6}
$$
Figure 2: The procedure details the Bregman Iteration (BI) employed in DMSB. The gray and blue blocks represent the BP step performed under  $\kappa_{\mathrm{boundary}}$  constraint for forward and backward policies, respectively. The red block signifies the BP step executed under the  $\kappa_{\mathrm{bridge}}$  constraint. Algorithms for training and sampling can be found in Appendix.D.
where:  $\mathcal{K}_{t_0} = \left\{\int \pi_{t_0:t_1}\mathrm{d}\mathbf{m}_{t_1} = \mu_{t_0},\int \mu_{t_0}\mathrm{d}\mathbf{v}_{t_0} = \rho_{t_0}\right\}$
$$
\mathcal {K} _ {t _ {N}} = \left\{\int \pi_ {t _ {N - 1}: t _ {N}} \mathrm {d} \mathbf {m} _ {t _ {N - 1}} = \mu_ {t _ {N}}, \int \mu_ {t _ {N}} \mathrm {d} \mathbf {v} _ {t _ {N}} = \rho_ {t _ {N}} \right\}
$$
$$
\mathcal {K} _ {t _ {i}} = \left\{\int \pi_ {t _ {i}: t _ {t + 1}} \mathrm {d} \mathbf {m} _ {t _ {i + 1}} = \mu_ {t _ {i}}, \int \pi_ {t _ {i - 1}: t _ {i}} \mathrm {d} \mathbf {m} _ {t _ {i - 1}} = \mu_ {t _ {i}}, \int \mu_ {t _ {i}} \mathrm {d} \mathbf {v} _ {t _ {i}} = \rho_ {t _ {i}} \right\}, \tag {7}
$$
and  $\mathcal{K}$  is the intersection of close convex set of  $\mathcal{K}_{t_i}$ .
The problem described in Prop.4.1 can be solved by classical BI algorithm integrated with Sinkhorn method [1]. However, due to the curse of dimensionality and unfavorable geometric explicit solution, the BP cannot be applied in high-dimensional and continuous state space directly. To tackle these difficulties, we parameterize the forward and backward policies  $\mathbf{z}_t$  and  $\widehat{\mathbf{z}}_t$  by a pair of neural networks. We further decouple and resemble the constraints by which it enables the scalable likelihood IPF and avoids the geometric averaging issue under mmmSB context.
4.1 Decoupling and Reassembling Constraints
We decompose the constraint set (7) by
$$
\mathcal {K} _ {t _ {i}} = \cap_ {r = 0} ^ {2} \mathcal {K} _ {t _ {i}} ^ {r}, \quad \text {w h e r e} \quad \begin{array}{l} \mathcal {K} _ {t _ {i}} ^ {0} = \left\{\int \pi_ {t _ {i}: t _ {t + 1}} \mathrm {d} \mathbf {m} _ {t _ {i + 1}} = \hat {\mu} _ {t _ {i}}, \int \hat {\mu} _ {t _ {i}} \mathrm {d} \mathbf {v} _ {t _ {i}} = \rho_ {t _ {i}} \right\} \\ \mathcal {K} _ {t _ {i}} ^ {1} = \left\{\int \pi_ {t _ {i - 1}: t _ {i}} \mathrm {d} \mathbf {m} _ {t _ {i - 1}} = \mu_ {t _ {i}}, \int \mu_ {t _ {i}} \mathrm {d} \mathbf {v} _ {t _ {i}} = \rho_ {t _ {i}} \right\} \\ \mathcal {K} _ {t _ {i}} ^ {2} = \left\{\int \pi_ {t _ {i}: t _ {t + 1}} \mathrm {d} \mathbf {m} _ {t _ {i + 1}} = \int \pi_ {t _ {i - 1}: t _ {i}} \mathrm {d} \mathbf {m} _ {t _ {i - 1}} \right\}. \end{array} \tag {8}
$$
One can notice that the  $\mathcal{K}_{t_i}^0$  and  $\mathcal{K}_{t_i}^1$  share similar structure as simpler boundary marginal conditions  $\mathcal{K}_{t_0}$  and  $\mathcal{K}_{t_N}$ , hence we can get rid of the notorious geometric averaging (see §4 in [1]). Notably, this type of constraint provides an opportunity to utilize Proposition 3.1 for optimization, but the joint distribution of  $\mathbf{x}$  and  $\mathbf{v}$  is still absent. We classify the constraints into two categories:
$$
\mathcal {K} _ {\text {b o u n d a r y}} = \left\{\cap_ {i = 1} ^ {N - 1} \mathcal {K} _ {t _ {i}} ^ {r} \cap \mathcal {K} _ {t _ {0}} \cap \mathcal {K} _ {t _ {N}} | \forall r \in \{0, 1 \} \right\}, \quad \mathcal {K} _ {\text {b r i d g e}} = \left\{\cap_ {i = 1} ^ {N - 1} \mathcal {K} _ {t _ {i}} ^ {2} \right\}.
$$
By following BI (§2.2), we execute optimization w.r.t. (6) while projecting the solution to subset of  $\mathcal{K}_{\mathrm{boundary}}$  or  $\mathcal{K}_{\mathrm{bridge}}$  iteratively. The sketch can be found in Fig.2. The next sections will provide more details on obtaining the joint distribution  $\mu$  and optimizing within each constraint set.
Hereafter, we only demonstrate the optimization for forward policy  $\mathbf{z}_t$  given reference path measure  $\bar{\pi}$  driven by fixed backward policy  $\widehat{\mathbf{z}}_t$ . The procedure can be applied for the  $\widehat{\mathbf{z}}_t$  and vice versa.
4.2 Optimization in set  $\kappa_{\mathrm{boundary}}$
We first show how to optimize forward policy  $\mathbf{z}_t$  w.r.t. objective function (6) given the reference path measure  $\bar{\pi}$  driven by fixed backward policy  $\widehat{\mathbf{z}}_t$  under one subset of  $\mathcal{K}_{\mathrm{boundary}}$ .