SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Proposition 4.2 (Optimality w.r.t.  $\kappa_{\mathrm{boundary}}$ ). Given the reference path measure  $\bar{\pi}$  driven by the backward policy  $\widehat{\mathbf{z}}_t$  from boundary  $\mu_{t_{i+1}}$  in the reverse time direction, the optimal path measure in the forward time direction of the following problem
$$
\min  _ {\pi} \mathcal {J} (\pi) := \sum_ {i = 0} ^ {N - 1} K L \left(\pi_ {t _ {i}: t _ {i + 1}} | \bar {\pi} _ {t _ {i}: t _ {i + 1}}\right), \quad s. t \quad \pi \in \left\{\int \pi_ {t _ {i}: t _ {i + 1}} \mathrm {d} \mathbf {m} _ {t _ {i + 1}} = \mu_ {t _ {i}}, \int \mu_ {t _ {i}} \mathrm {d} \mathbf {v} _ {t _ {i}} = \rho_ {t _ {i}} \right\}
$$
$$
i s: \quad \pi_ {t _ {i}: t _ {i + 1}} ^ {*} = \frac {\rho_ {t _ {i}} \bar {\pi} _ {t _ {i}: t _ {i + 1}}}{\int \bar {\pi} _ {t _ {i}: t _ {i + 1}} \mathrm {d} \mathbf {m} _ {t _ {i + 1}} \mathrm {d} \mathbf {v} _ {t _ {i}}}.
$$
When  $\pi_{t_i:t_{i + 1}}\equiv \pi_{t_i:t_{i + 1}}^*$ , the following equations need to hold  $\forall t\in [t_i,t_{i + 1}]$ :
$$
\left\| \mathbf {z} _ {t} + \widehat {\mathbf {z}} _ {t} - g \nabla_ {\mathbf {v}} \log \hat {p} _ {t} \right\| _ {2} ^ {2} = 0, \tag {9a}
$$
$$
p _ {t _ {i}} \left(\mathbf {v} _ {t _ {i}} \mid \mathbf {x} _ {t _ {i}}\right) \equiv \hat {q} _ {t _ {i}} \left(\mathbf {v} _ {t _ {i}} \mid \mathbf {x} _ {t _ {i}}\right), \tag {9b}
$$
where  $\hat{p}_t$  and  $\hat{q}_t$  denote the marginal density and conditional velocity distribution of the reference path measure at time  $t$ , respectively.
Proof. See appendix.B.5
Remark 4.3. When the ground truth distributions of velocity  $\gamma_{t_i}$  are available, one can simply sample from  $\gamma_{t_i}$  since the joint distribution  $\mu_t$  is available in this case. In order to matching the reference path measure in KL divergence sense, one needs to match both the intermediate path measure eq.9a and the boundary condition eq.9b. In the traditional two-boundary SB case, matching the boundary condition is often disregarded due to either having a predefined data distribution or a tractable prior. However, in our specific case, as the velocity is not predefined, it becomes imperative to address this issue and optimize it through the application of Langevin dynamics.
4.3 Optimization in set  $\kappa_{\mathrm{bridge}}$
The formulation of optimization under  $\kappa_{\mathrm{bridge}}$  is similar to the previous section but differs by the boundary condition (eq.10b):
Proposition 4.4 (Optimality w.r.t.  $\kappa_{\mathrm{bridge}}$ ). Given the reference path measure  $\bar{\pi}$  driven by the backward policy  $\widehat{\mathbf{z}}_t$  from boundary  $\mu_{t_N}$  in the reverse time direction, the optimal path measure in the forward time direction of the following problem
$$
\min  _ {\pi} \mathcal {J} (\pi) := \sum_ {i = 0} ^ {N - 1} K L \left(\pi_ {t _ {i}: t _ {i + 1}} | \bar {\pi} _ {t _ {i}: t _ {i + 1}}\right), \quad s. t \quad \pi \in \mathcal {K} _ {b r i d g e} = \left\{\cap_ {i = 1} ^ {N - 1} \mathcal {K} _ {t _ {i}} ^ {2} \right\}
$$
$$
i s: \pi_ {t _ {0}: t _ {N}} ^ {*} = \frac {q _ {t _ {0}} \bar {\pi} _ {t _ {0} : t _ {N}}}{\int \bar {\pi} _ {t _ {0}: t _ {N}} \mathrm {d} \mathbf {m} _ {t _ {N}} \mathrm {d} \mathbf {v} _ {t _ {0}}}.
$$
when  $\pi_{t_0:t_N}\equiv \pi_{t_0:t_N}^*$ , the following equations need to hold  $\forall t\in [t_0,t_N]$ :
$$
\left\| \mathbf {z} _ {t} + \widehat {\mathbf {z}} _ {t} - g \nabla_ {\mathbf {v}} \log \hat {p} _ {t} \right\| _ {2} ^ {2} = 0 \tag {10a}
$$
$$
p _ {t _ {0}} \left(\mathbf {v} _ {t _ {0}}, \mathbf {x} _ {t _ {0}}\right) \equiv \hat {q} _ {t _ {0}} \left(\mathbf {v} _ {t _ {0}}, \mathbf {x} _ {t _ {0}}\right) \tag {10b}
$$
Proof. See appendix.B.6
Conceptually, the above optimization objective with  $\kappa_{\mathrm{bridge}}$  constraint aims at finding a continuous path measure close to reference path measure  $\bar{\pi}$  while any intermediate marginals constraints will not be considered. The boundary condition of reference path measure in the next iteration  $p_{t_0}(\mathbf{v}_{t_0},\mathbf{x}_{t_0})$  is determined by eq.10b. Fortunately, the empirical samples from this distribution are available, though the analytic representation of the distribution  $\hat{q}_{t_0}(\mathbf{v}_{t_0},\mathbf{x}_{t_0})$  is unknown. Hence we can utilize these samples as empirical sources from boundary distribution  $\hat{q}_{t_0}(\mathbf{v}_{t_0},\mathbf{x}_{t_0})$  for the next BP. For further explanation and intuition, one can find it in Appendix.G
4.4 Parameterization and Training Objective Function
Inspired by the success of prior work [14], we parameterize path measure  $\pi$  by forward policy  $\mathbf{z}_t^\theta$  or backward policy  $\widehat{\mathbf{z}}_t^\phi$  combined with one of constraints in  $\kappa_{\mathrm{boundary}}$  or  $\kappa_{\mathrm{bridge}}$  (see Fig.8 in Appendix for visualization). We adopt Euler-Maruyama discretization and denote the timestep as  $\delta_t$ . Notably, eq.9b and eq.10b can be implied by minimizing phase space NLL in Prop.3.1. This leads to the following objective function, termed as phase space mean matching objective, which will be used to train neural networks that represent  $\mathbf{z}_t^\theta$  and  $\widehat{\mathbf{z}}_t^\phi$  after time discretization:
$$
\mathcal {L} _ {M M} = \mathbb {E} \left[ | | \delta_ {t} \mathbf {z} _ {t} ^ {\theta} (\mathbf {m} _ {t + \delta_ {t}}) + \delta_ {t} \widehat {\mathbf {z}} _ {t + \delta_ {t}} ^ {\phi} (\mathbf {m} _ {t + \delta_ {t}}) - (\mathbf {m} _ {t} + \delta_ {t} \mathbf {z} _ {t} ^ {\theta} - \mathbf {m} _ {t + \delta_ {t}}) | | ^ {2} \right].
$$
The velocity boundary condition for the reference path measure in the succeeding BP is encoded in eq.9b or eq.10b, but the representation of conditional distribution eq.9b is not clear. We leverage the favorable property of SB to parameterize and sample from such distribution.
Figure 4: Validation of our DMSB model on complex GMM synthetic dataset. The velocity and position of the same sample correspond to the same shade level. Upper: Samples' evolution in the position space. Bottom: Learnt samples' evolution in the velocity space.
Proposition 4.5 ([27, 28]). If  $\pi^{\theta}$  and  $\bar{\pi}^{\phi}$  shares same path measure, then
$$
\tilde {p} _ {t _ {i}} ^ {\theta , \phi} \left(\mathbf {v} _ {t _ {i}}, \mathbf {x} _ {t _ {i}}\right) \equiv q _ {t _ {i}} ^ {\phi} \left(\mathbf {v} _ {t _ {i}}, \mathbf {x} _ {t _ {i}}\right) \propto q _ {t _ {i}} ^ {\phi} \left(\mathbf {v} _ {t _ {i}} \mid \mathbf {x} _ {t _ {i}}\right), \quad \text {w h e r e :} \quad \nabla_ {\mathbf {v}} \log \tilde {p} _ {t} ^ {\theta , \phi} = \left(\mathbf {z} _ {t} ^ {\theta} + \widehat {\mathbf {z}} _ {t} ^ {\phi}\right) / g. \tag {11}
$$
Prop.4.5 suggests that one can use  $p_{t_i}(\mathbf{v}_{t_i}|\mathbf{x}_{t_i}) \coloneqq \hat{p}_{t_i}^{\theta ,\phi}$  to imply condition (9b) and obtain samples from such distribution by simulating Langevin dynamics. Namely, we first sample position from ground truth  $\mathbf{x}_{t_i}\sim \rho_{t_i}$ , and then sample  $\mathbf{v}_{t_i}\sim \tilde{p}_t^{\theta ,\phi}$  using eq.11. One can further adopt the same regularization [29] to enforce the condition of Prop.4.5.
4.5 Training Scheme
Here we introduce the scheme to traverse BI (see Fig.2). In one BI, all constraints must be iterated once. For the sake of  $\mathcal{L}_{\mathrm{MM}}$ , the reference path measure should be induced by opposite direction. A single BI cannot be recursively repeated due to the conflict of reference path measure direction. For example (see Fig.2), at the end of  $d$ -th BI,  $\bar{\pi}$  is yielded by forward policy while the first BP of  $d$ -th BI is also optimizing forward policy which violates  $\mathcal{L}_{\mathrm{MM}}$ . Instead, we reschedule the optimization order. Specifically, in  $(d + 1)$ -th BI, we optimize backward policy at the first BP and the last BP.
**Setup:** We test DMSB on 2D synthetic datasets and real-world scRNA-seq dataset [30]. We choose state of the art algorithms MIOFlow [9] and NLSB [11] as our baselines. We tune both models to the best of our hardware capacity. We choose Sliced-Wasserstein Distance (SWD)[31] and Maximum Mean Discrepancy (MMD)[32] together with visualization as our criterion. The detailed setup of training and evaluation can be found in Appendix.C.
Synthetic Datasets: The Petal [9] and Gaussian Mixture Model (GMM) dataset are simple yet challenging, as they mimic natural dynamics arising in cellular differentiation, including bifurcations and merges. We compare our algorithm with MIOFlow in Fig.3. DMSB can infer trajectories aligned with ground truth distribution more faithfully at timesteps when snapshots are taken.
In GMM experiments (see Fig.4), we choose standard Gaussian at initial and terminal time steps while four-modal GMM and eight-modal GMM are placed at intermediate time steps. Besides good position trajectory, it is almost serendipity that DMSB can also learn the reasonable velocity trajectory without any access to ground truth velocity information. This paves the way for our later velocity estimation for the RNAsc dataset. scRNA-seq Dataset: The emergence of single-cell profiling technologies has facilitated the acquisition of high-resolution single-cell data, enabling the characterization of individual cells at distinct developmental states [7]. However, because the cell population is eliminated after the measurement, one may only gather statistical data for single samples at particular timesteps, which neither preserves
Figure 3: Comparison with MIOFlow and ground truth on challenging petal dataset. DMSB is able to generate trajectories whose time marginal matches ground truth faithfully and outperforms prior work. Time is indicated by colors.
any correlations over time nor provides access to the ground truth trajectory. The diversity of embryonic stem cells after development from embryoid bodies, which comprises mesoderm, endoderm,
Table 3: Numerical result of MMD and SWD on 100 dimensions single-cell RNA-seq dataset and results for leaving out (LO) marginals at different observation. DMSB outperforms prior work by a large margin for both metrics and all leave-out case. See Appendix.4 for Results over 3 seeds.
Algorithm	MMD ↓	SWD ↓
w/o LO	LO-t1	LO-t2	LO-t3	w/o LO	LO-t1	LO-t2	LO-t3
NLSB[10]	0.66	0.38	0.37	0.37	0.54	0.55	0.54	0.55
MIOFlow[8]	0.23	0.23	0.90	0.23	0.35	0.49	0.72	0.50
DMSB(ours)	0.03	0.04	0.04	0.04	0.20	0.20	0.19	0.18
neuroectoderm, and neural crest in 27 days, is demonstrated by the scRNA-seq dataset. The snapshot of cells are collected between  $(t_0$ : day 0 to 3,  $t_1$ : day 6 to 9,  $t_2$ : day 12 to 15,  $t_3$ : day 18 to 21,  $t_4$ : day 24 to 27). Snapshot data are prepossessed by the quality control [30] and then projected to feature space by principal component analysis (PCA). We inherit processed data from [8]. We validate DMSB on 5-dim and 100-dim PCA space to show superior performance on high-dimension problems compared with baselines. We further show that DMSB can estimate better velocity distribution compared with baselines when the ground truth is absent during training and testing.
We testify the performance of our model by computing MMD and SWD with full snapshots and when one of snapshots is left out (LO). We postpone the comparison of all the models on 5-d RNA space to the appendix (see Fig.9 and Table.6) because the problem is relatively simple and all models can infer accurate trajectory. Table.3 summarizes the average MMD and SWD between estimated marginal and ground truth over different snapshot timesteps. DMSB outperforms prior work by a large margin in high (100) dimensional scenarios. The visualization (Fig.5) in PCA space further justifies the numerical result and highlights the variety and quality of the samples produced by DMSB.
Figure 5: Comparison of population-level dynamics on 100-dimensional PCA space at the moment of observation for scRNA-seq data using MIOFlow, NLSB, and DMSB. We display the plot of the first 6 principle components (PC). Baselines can only learn the trajectory's fundamental trend, whereas DMSB can match the target marginal along the trajectory across different dimensions. The right figure shows Kernel Density Estimation [33] of samples generated by DMSB and ground truth at  $t_3$  and  $t_4$ . The generated samples for all timesteps and comparison with baseline are in Appendix.F.
Interestingly, Fig.4 demonstrates that DMSB can reconstruct reasonable evolution of the velocity distribution which was not accessible to the algorithm. We further validate such property in 100-D RNAsc dataset. During the training and testing, all the models do not have access to the ground truth velocity. We run the experiments of 100-D and 5-D RNAsc datasets and average the discrepancy between ground truth velocity and estimated velocity over snapshot time. The numerical values are listed in the Table.7 and Table.6. The plot of velocity and position can be found in Fig.9 and Fig.10. The plot illustrates that while all models are capable of learning reasonable trajectories, only DMSB has the ability to estimate a plausible velocity distribution. This property holds even for
100-D RNA dataset (see Fig.5,11,12). This is notable, despite the velocity estimated by DMSB does not perfectly match the ground truth, because it should be noted that the proposed phase space SDE and the optimality of OT are artificial and may not necessarily represent the actual RNA evolution. Moreover, as individual evolutions cannot be tracked, possibilities such as  $\{\mathrm{A}\rightarrow \mathrm{A},\mathrm{B}\rightarrow \mathrm{B}\}$  versus  $\{\mathrm{A}\to \mathrm{B},\mathrm{B}\to \mathrm{A}\}$  can not be discerned, which renders exact velocity recovering almost impossible.
In this paper, we propose DMSB, a scalable algorithm that learns the trajectory which fits the different marginal distributions over time. We extend the mean matching objective to phase space which enables efficient mSB computing. We propose a novel training scheme to fit the mean matching objective without violating BI which is the root of solving mmmSB problem. We demonstrate the superior result of DMSB compared with the existing algorithms.
A main limitation of this work is, the rate of convergence to the actual mmmSB has not been quantified after neural network approximations are introduced. Even though [15] theoretically analyzed the convergence of mean matching iteration, supporting its outstanding performance [14], the iteration still fails to converge to the actual SB [34] precisely due to practical neural network estimation errors accumulating over BI. However, recent work [35] shows the convergence of SB when training error exists. In addition, DMSB cannot simulate the process with death and birth of cells which can be potentially described as unbalanced optimal transport [36].
This research was supported by the ARO Award W911NF2010151, and the DoD Basic Research Office AwardHQ00342110002.
[35] Yu Chen, Wei Deng, Shikai Fang, Fengpei Li, Nicole Tianjiao Yang, Yikai Zhang, Kashif Rasul, Shandian Zhe, Anderson Schneider, and Yuriy Nevmyvaka. Provably convergent schr\* odinger bridge with applications to probabilistic time series imputation. arXiv preprint arXiv:2305.07247, 2023.
[36] Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. The most likely evolution of diffusing and vanishing particles: Schrodinger bridges with unbalanced margins. SIAM Journal on Control and Optimization, 60(4):2016-2039, 2022.
[37] Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. arXiv e-prints, pages arXiv-2101, 2021.
[38] Jiongmin Yong and Xun Yu Zhou. Stochastic controls: Hamiltonian systems and HJB equations, volume 43. Springer Science & Business Media, 1999.
[39] Kenneth Caluya and Abhishek Halder. Wasserstein proximal algorithms for the schrödinger bridge problem: Density control with nonlinear drift. IEEE Transactions on Automatic Control, 2021.
[40] Ioannis Exarchos and Evangelos A Theodorou. Stochastic optimal control via forward and backward stochastic differential equations and importance sampling. Automatica, 87:159-165, 2018.
[41] Tianrong Chen, Ziyi O Wang, Ioannis Exarchos, and Evangelos Theodorou. Large-scale multiagent deep fbsdes. In International Conference on Machine Learning, pages 1740-1748. PMLR, 2021.
[42] Tianrong Chen, Ziyi Wang, and Evangelos A Theodorou. Deep graphic fbsdes for opinion dynamics stochastic control. In 2022 IEEE 61st Conference on Decision and Control (CDC), pages 4652-4659. IEEE, 2022.
[43] Hopf Eberhard. The partial differential equation  $\mathrm{ut} + \mathrm{uux} = \mu \mathrm{xx}$ . Communications on Pure and Applied Mathematics, 3(3):201-230, 1950.
[44] Julian D Cole. On a quasi-linear parabolic equation occurring in aerodynamics. Quarterly of applied mathematics, 9(3):225-236, 1951.
[45] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017.
[46] Francisco Vargas, Pierre Thodoroff, Neil D Lawrence, and Austen Lamacraft. Solving schrödinger bridges via maximum likelihood. arXiv preprint arXiv:2106.02081, 2021.
B Proof in §3 and §4
Before stating our proofs, we provide the assumptions used throughout the paper. These assumptions are adopted from stochastic analysis for SGM [27, 37, 38], SB [39], and FBSDE [40-42].
(i)  $\mu_{t_i}$  with finite second-order moment for all  $t_i$
(ii)  $\pmb{f}$  and  $g$  are continuous functions, and  $|g(t)|^2 > 0$  is uniformly lower-bounded w.r.t.  $t$ .
(iii)  $\forall t\in [0,T]$ , we have  $\nabla_{\mathbf{v}}\log p_t(\mathbf{m}_t,t),\nabla_{\mathbf{v}}\log \Psi (\cdot ,\cdot ,\cdot),\nabla_{\mathbf{v}}\log \widehat{\Psi} (\cdot ,\cdot ,\cdot),\mathbf{Z}(\cdot ,\cdot ,\cdot ;\theta)$ , and  $\widehat{\mathbf{Z}} (\cdot ,\cdot ,\cdot ;\phi)$  Lipschitz and at most linear growth w.r.t.  $\mathbf{x}$  and  $\mathbf{v}$ .
(iv)  $\Psi ,\widehat{\Psi}\in C^{1,2}.$
(v)  $\exists k > 0: p_t^{SB}(\mathbf{m}) = \mathcal{O}(\exp^{-\|\mathbf{m}\|_k^2})$  as  $\mathbf{m} \to \infty$ .
Assumptions (i) (ii) (iii) are standard conditions in stochastic analysis to ensure the existence-uniqueness of the SDEs; hence also appear in SGM analysis [37]. Assumption (iv) allows applications of Ito formula and properly defines the backward SDE in FBSDE theory. Finally, assumption (v) assures the exponential limiting behavior when performing integration by parts. w.o.l.g, we denote  $\pmb{f} = [\mathbf{v},\mathbf{0}]^{\top}$ .
B.1 Proof of Proposition.3.1
The results of the Prop.3.1 is part of results of Prop.B.4 which gives the results for both forward and backward likelihood objective.
Theorem B.1. The optimization problem
$$
\min  \int_ {0} ^ {1} \int \frac {1}{2} \| \mathbf {a} \| _ {2} ^ {2} \mu \mathrm {d} \mathbf {x} \mathrm {d} \mathbf {v} \mathrm {d} t, \tag {12}
$$
$$
s. t \quad \left\{ \begin{array}{l} \frac {\partial \mu (\mathbf {m} _ {t})}{\partial t} = - \nabla_ {\mathbf {m}} \cdot \left\{\left[ (\boldsymbol {f} + g \mathbf {u}) \right] \mu \right\} + \frac {1}{2} g ^ {2} \Delta_ {\mathbf {v}} \mu , \\ \mu_ {0} = p (0, \mathbf {x}, \mathbf {v}), \quad \mu_ {1} = p (T, \mathbf {x}, \mathbf {v}), \end{array} \right. \tag {13}
$$
will induce the coupled PDEs,
$$
\frac {\partial \mu (\mathbf {m} _ {t})}{\partial t} = - \nabla_ {\mathbf {m}} \cdot \left[ \left(\boldsymbol {f} + g ^ {2} \nabla_ {\mathbf {v}} \phi\right) \mu \right] + \frac {1}{2} g ^ {2} \Delta_ {\mathbf {v}} \mu , \tag {14}
$$
$$
\frac {\partial \phi (\mathbf {m} _ {t})}{\partial t} = - \frac {1}{2} \| g \nabla_ {\mathbf {v}} \phi \| _ {2} ^ {2} - \mathbf {v} ^ {\mathrm {T}} \nabla_ {\mathbf {x}} \phi - \frac {1}{2} g ^ {2} \Delta_ {\mathbf {v}} \phi , \tag {15}
$$
and the optimal control of the problem is
$$
\mathbf {a} ^ {*} = g \nabla_ {\mathbf {v}} \phi .
$$
Proof. One can write the Lagrange by introducing lagrangian multiplier  $\phi$ :
$$
\begin{array}{l} \mathcal {L} (\mu , \mathbf {a}, \phi) = \int_ {0} ^ {1} \int_ {\mathbb {R} ^ {n} \times \mathbb {R} ^ {n}} \frac {1}{2} \| \mathbf {a} \| _ {2} ^ {2} \mu \mathrm {d} \mathbf {x} \mathrm {d} \mathbf {v} \mathrm {d} t + \int_ {0} ^ {1} \int_ {\mathbb {R} ^ {n} \times \mathbb {R} ^ {n}} \phi \frac {\partial \mu}{\partial t} \mathrm {d} \mathbf {v} \mathrm {d} \mathbf {x} \mathrm {d} t \\ + \int_ {0} ^ {1} \int_ {\mathbb {R} ^ {n} \times \mathbb {R} ^ {n}} \phi \left\{- \frac {1}{2} g ^ {2} \Delta_ {\mathbf {m}} \mu + \nabla_ {\mathbf {m}} \cdot [ (\boldsymbol {f} + g \mathbf {u})   \mu ] \right\} d \mathbf {v} d \mathbf {x} d t \\ = \int_ {0} ^ {1} \int_ {\mathbb {R} ^ {n} \times \mathbb {R} ^ {n}} \frac {1}{2} \| \mathbf {a} \| _ {2} ^ {2} \mu \mathrm {d} \mathbf {x} \mathrm {d} \mathbf {v} \mathrm {d} t - \int_ {0} ^ {1} \int_ {\mathbb {R} ^ {n} \times \mathbb {R} ^ {n}} \mu \frac {\partial \phi}{\partial t} \mathrm {d} \mathbf {v} \mathrm {d} \mathbf {x} \mathrm {d} t \\ + \int_ {0} ^ {1} \int_ {\mathbb {R} ^ {n} \times \mathbb {R} ^ {n}} \phi \nabla_ {\mathbf {m}} \cdot [ (\boldsymbol {f} + g \mathbf {u}) \mu ] - \phi \left[ \frac {1}{2} g ^ {2} \Delta_ {\mathbf {m}} \mu \right] d \mathbf {v} d \mathbf {x} d t \\ = \int_ {0} ^ {1} \int_ {\mathbb {R} ^ {n} \times \mathbb {R} ^ {n}} \frac {1}{2} \| \mathbf {a} \| _ {2} ^ {2} \mu \mathrm {d} \mathbf {x} \mathrm {d} \mathbf {v} \mathrm {d} t - \int_ {0} ^ {1} \int_ {\mathbb {R} ^ {n} \times \mathbb {R} ^ {n}} \mu \frac {\partial \phi}{\partial t} \mathrm {d} \mathbf {v} \mathrm {d} \mathbf {x} \mathrm {d} t \\ + \int_ {0} ^ {1} \int_ {\mathbb {R} ^ {n} \times \mathbb {R} ^ {n}} - \nabla_ {\mathbf {m}} \phi^ {\mathsf {T}} [ (\boldsymbol {f} + g \mathbf {u}) ] \mu - \mu \left[ \frac {1}{2} g ^ {2} \Delta_ {\mathbf {m}} \phi \right] d \mathbf {v} d \mathbf {x} d t \\ = \int_ {0} ^ {1} \int_ {\mathbb {R} ^ {n} \times \mathbb {R} ^ {n}} \left\{\frac {1}{2} \| \mathbf {a} \| _ {2} ^ {2} - \frac {\partial \phi}{\partial t} - \mathbf {v} ^ {\mathsf {T}} \nabla_ {\mathbf {x}} \phi - \frac {1}{2} g ^ {2} \Delta_ {\mathbf {v}} \phi - g \nabla_ {\mathbf {v}} \phi^ {\mathsf {T}} \mathbf {a} \right\} \mu d \mathbf {v} d \mathbf {x} d t \\ \end{array}
$$
By taking the minimization within the bracket, The optimal control is,
$$
\mathbf {a} ^ {*} = g \nabla_ {\mathbf {v}} \phi
$$
By Plugging it back, the optimality of the aforementioned problem is presented as:
$$
\frac {\partial \mu (\mathbf {m} _ {t})}{\partial t} = - \nabla_ {\mathbf {v}} \cdot \left[ \left(\boldsymbol {f} + g ^ {2} \nabla_ {\mathbf {v}} \phi\right) \mu \right] + \frac {1}{2} g ^ {2} \Delta_ {\mathbf {v}} \mu ,
$$
$$
\frac {\partial \phi (\mathbf {m} _ {t})}{\partial t} = - \frac {1}{2} \| g \nabla_ {\mathbf {v}} \phi \| _ {2} ^ {2} - \mathbf {v} ^ {\mathsf {T}} \nabla_ {\mathbf {x}} \phi - \frac {1}{2} g ^ {2} \Delta_ {\mathbf {v}} \phi ,
$$
Theorem B.2. The optimal forward and backward processes are represented as:
$$
\mathrm {d} \mathbf {m} _ {t} = \left[ \boldsymbol {f} + g \mathbf {u} _ {t} ^ {f *} \right] \mathrm {d} t + g (t) \mathrm {d} \mathbf {w} _ {t} \quad (\text {f o r w a r d}) \tag {16}
$$
$$
\mathrm {d} \mathbf {m} _ {s} = \left[ \boldsymbol {f} + g \mathbf {u} _ {s} ^ {b *} \right] \mathrm {d} t + g (t) \mathrm {d} \mathbf {w} _ {s} \quad (\text {B a c k w a r d}) \tag {17}
$$
in which  $\mathbf{f} = [\mathbf{v},\mathbf{0}]^{\top}$ . Optimal control is expressed as,
$$
\mathbf {u} _ {t} ^ {f *} := \mathbf {Z} _ {t} \equiv \left( \begin{array}{l} \mathbf {0} \\ \mathbf {z} _ {t} \end{array} \right) \equiv \left( \begin{array}{c} \mathbf {0} \\ g \nabla_ {\mathbf {v}} \log \Psi_ {t} \end{array} \right) \tag {18}
$$
$$
\mathbf {u} _ {t} ^ {b *} := \widehat {\mathbf {Z}} _ {t} \equiv \left( \begin{array}{c} \mathbf {0} \\ \widehat {\mathbf {z}} _ {t} \end{array} \right) \equiv \left( \begin{array}{c} \mathbf {0} \\ g \nabla_ {\mathbf {v}} \log \widehat {\Psi} _ {t} \end{array} \right) \tag {19}
$$
where  $\Psi$  and  $\hat{\Psi}$  are the solution of following PDEs,
$$
\frac {\partial \Psi_ {t}}{\partial t} = - \frac {1}{2} g ^ {2} \Delta_ {\mathbf {v}} \Psi_ {t} - \nabla_ {\mathbf {x}} \Psi_ {t} ^ {\mathsf {T}} \mathbf {v}
$$