# Control, Transport and Sampling: Towards Better Loss Design

Qijia Jiang
and David Nabergoj Lawrence Berkeley National Laboratory,
<EMAIL_ADDRESS>of Ljubljana, Faculty of Computer and Information
Science<EMAIL_ADDRESS>
###### Abstract

Leveraging connections between diffusion-based sampling, optimal transport,
and optimal stochastic control through their shared links to the Schrödinger
bridge problem, we propose novel objective functions that can be used to
transport $\nu$ to $\mu$, consequently sample from the target $\mu$, via
_optimally controlled_ dynamics. We highlight the importance of the pathwise
perspective and the role various optimality conditions on the path measure can
play for the design of valid training losses, the careful choice of which
offer numerical advantages in practical implementation.

##  1 Introduction

Traditionally, the task of sampling from un-normalized densities is largely
delegated to MCMC methods. However, modern machine learning developments in
optimal transport and generative modeling have greatly expanded the toolbox we
have available for performing such tasks, which can further benefit from
powerful advancements in deep learning. By reducing the problem to performing
empirical risk minimization with neural networks, they hold promise especially
for sampling from high-dimensional and multimodal distributions compared to
MCMC-based alternatives. In this work, we propose novel training objectives
that are amenable to tractable approximations for sampling (without access to
data from the target, as typically considered in the MCMC setup), and
demonstrate their numerical advantages compared to several related methods
that are also built around forward-backward SDEs through time-reversals.
Before elaborating on these connections and synthesis in Section 2, we briefly
summarize our contributions below.

  * •

We present a transport/control based sampler using Schrödinger bridge (SB)
idea (generalizable to extended state space and nonlinear prior), which
compared to MCMC-based approaches has two benefits: (1) alleviate
metastability without being confined to make local moves in the state space;
(2) based on interpolating path, it comes with the ability to provide low-
variance, unbiased normalizing constant estimates using the trajectory
information (c.f. Section 3.2).

  * •

Unlike previously studied Schrödinger-bridge-based [7] and diffusion-based
samplers [3, 31, 25], we explicitly enforce _optimality_ and _uniqueness_ of
the trajectory, which also ensure that the dynamics reaches target in _finite_
time. Compared to IPF-based approaches for solving SB [32], our joint training
approach of the forward and backward controls is not susceptible to prior
forgetting, at the same time amenable to importance sampling correction for
bias coming from neural network training (e.g., approximation error). The
algorithm and setup we consider deviates from the more well-studied generative
model use case of SB [9, 22, 27].

  * •

In contrast to previously proposed losses that follow optimal trajectory for
sampling from the target [32, 18], our (discretized) training objectives have
the numerical advantage of (1) vanishing variance of the stochastic gradient
at the optimal control; (2) expense the need for evaluating expensive
Laplacian terms (c.f. Section 3.1). These heavily rely on the path measure
based representations of our losses and estimators and generalize much beyond
the half bridge case considered in [34]. Comparisons of the performance are
given in Section 4 w.r.t alternative loss proposals.

##  2 General Framework

In this section we put several recently proposed (and diffusion-related)
methods in context, rendering them as special instantiations of a more
unifying _path-wise_ picture, which will in turn motivate the need for a
control-based approach that involves designing an optimal path exactly
interpolating between two distributions. We loosely follow the framework put
forth in [32] for parts of the exposition.

###  2.1 Setup

We are interested in sampling from $p_{\text{target}}(x)=\mu(x)/Z$ by
minimizing certain tractable loss, assuming sampling from
$p_{\text{prior}}(z)=\nu(z)$ is easy, but unlike in generative modeling, even
though analytical expression for $\mu$ is readily available, we do not have
access to data from it that can be used to learn the score function. In this
sense, in terms of “transport mapping", the two sides are crucially not
symmetric. Below we introduce forward-backward SDEs for formalizing such
transitions. Pictorially, given a base drift $f$, we have the “sampling
process" (the two processes are time reversals of each other):

$\nu(z)\mathrel{\mathop{\mathrel{\begin{subarray}{c}\xrightarrow{\rule{42.67912pt}{0.0pt}}\\\\[-3.87495pt]
\xleftarrow{\rule{42.67912pt}{0.0pt}}\end{subarray}}}^{\mathrm{\mathbb{P}^{\nu,f+\sigma
u}}}_{\mathrm{\mathbb{P}^{\mu,f+\sigma v}}}}\mu(x)$ (1)

for tunable controls $u,v$ and terminal marginals $\nu,\mu$. Written as SDEs,
they become (the drifts $u,v$ here are not independent)

$\displaystyle dX_{t}=[f_{t}(X_{t})+\sigma
u_{t}(X_{t})]dt+\sigma\overrightarrow{dW_{t}},\;X_{0}\sim\nu\Rightarrow(X_{t})_{t\in[0,T]}\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}\,,$ (2) $\displaystyle dX_{t}=[f_{t}(X_{t})+\sigma
v_{t}(X_{t})]dt+\sigma\overleftarrow{dW_{t}},\;X_{T}\sim\mu\Rightarrow(X_{t})_{t\in[0,T]}\sim\overleftarrow{\mathbb{P}}^{\mu,f+\sigma
v}\,.$ (3)

Operationally, (2)-(3) denote (picking $f=0$, and $\\{z_{t}\\}$ a sequence of
i.i.d standard Gaussians for illustration)

$\displaystyle X_{t}$ $\displaystyle=X_{0}+\int_{0}^{t}\sigma
u_{s}(X_{s})ds+\int_{0}^{t}\sigma\overrightarrow{dW_{s}}\;\;\Rightarrow
X_{t+h}\approx X_{t}+h\sigma u_{t}(X_{t})+\sqrt{h}\sigma z_{t},\quad
X_{0}\sim\nu$ $\displaystyle X_{t}$ $\displaystyle=X_{T}-\int_{t}^{T}\sigma
v_{s}(X_{s})ds-\int_{t}^{T}\sigma\overleftarrow{dW_{s}}\;\;\Rightarrow
X_{t-h}\approx X_{t}-h\sigma v_{t}(X_{t})+\sqrt{h}\sigma z_{t},\quad
X_{T}\sim\mu$

where the forward (the usual Itô’s) and the backward integrals indicate
different endpoints at which we make the approximation. For processes
$(Y_{t})_{t},(Z_{t})_{t}$,

$\int_{0}^{T}a_{t}(Y_{t})\,\overrightarrow{dZ_{t}}\approx\sum_{i}a_{t_{i}}(Y_{t_{i}})(Z_{t_{i+1}}-Z_{t_{i}})\,,\quad\int_{0}^{T}a_{t}(Y_{t})\,\overleftarrow{dZ_{t}}\approx\sum_{i}a_{t_{i+1}}(Y_{t_{i+1}})(Z_{t_{i+1}}-Z_{t_{i}})\,,$
(4)

which in particular implies the martingale property

$\mathbb{E}_{X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}}\left[\int_{0}^{t}a_{s}(X_{s})\,\overrightarrow{dW_{s}}\right]=0\,,\quad\mathbb{E}_{X\sim\overleftarrow{\mathbb{P}}^{\mu,f+\sigma
v}}\left[\int_{T-t}^{T}a_{s}(X_{s})\,\overleftarrow{dW_{s}}\right]=0\,.$ (5)

Forward/backward integral can be converted through

$\sigma^{2}\int_{0}^{T}(\nabla\cdot
a_{t})(Y_{t})\,dt+\int_{0}^{T}a_{t}(Y_{t})\,\overrightarrow{dZ_{t}}=\int_{0}^{T}a_{t}(Y_{t})\,\overleftarrow{dZ_{t}}=\int_{0}^{T}\hat{a}_{t}(Y_{t})\,\overrightarrow{d\hat{Z}_{t}}\,,$
(6)

which will be used repeatedly throughout. Above hat denotes $t\rightarrow
T-t$.

######  Remark 1 (Nelson’s identity).

The following relationship between drifts for the SDE (2)-(3) is well known:
$\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}=\overleftarrow{\mathbb{P}}^{\mu,f+\sigma v}$ iff
$\overrightarrow{\mathbb{P}}^{\nu,f+\sigma u}_{T}=\mu$, and $\sigma
v_{t}(X_{t})=\sigma
u_{t}(X_{t})-\sigma^{2}\nabla\log(\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}_{t})=\sigma
u_{t}(X_{t})-\sigma^{2}\nabla\log(\overleftarrow{\mathbb{P}}^{\mu,f+\sigma
v}_{t})$ for all $t\in[0,T]$. This is used to give a more transparent
derivation of the likelihood ratio in Lemma 3.

###  2.2 Goal and Related Approaches

#### Forward KL

We would like the two path measures (with the specified two end-point
marginals) to agree progressing in either direction. The methods in [31, 34]
propose to set up a reference process with a similar structure as (1) with
$f=0$:

$\nu(z)\mathrel{\mathop{\mathrel{\begin{subarray}{c}\xrightarrow{\rule{42.67912pt}{0.0pt}}\\\\[-3.87495pt]
\xleftarrow{\rule{42.67912pt}{0.0pt}}\end{subarray}}}^{\mathrm{\mathbb{P}^{\nu,\sigma
r}}}_{\mathrm{\mathbb{P}^{\eta,\sigma v}}}}\eta(x)$ (7)

where $r$ is the drift for a reference process:

$dX_{t}=\sigma
r_{t}(X_{t})\,dt+\sigma\overrightarrow{dW_{t}},\;X_{0}\sim\nu\Rightarrow(X_{t})_{t\in[0,T]}\sim\overrightarrow{\mathbb{P}}^{\nu,\sigma
r}$ (8)

and $\eta=\overrightarrow{\mathbb{P}}^{\nu,\sigma r}_{T}$. Via Girsanov’s
theorem, the method amounts to minimizing loss of the following type:

$\displaystyle\mathcal{L}_{KL}(u)$
$\displaystyle=\mathbb{E}_{\overrightarrow{\mathbb{P}}^{\nu,\sigma
u}}\Big{[}\log\Big{(}\frac{d\overrightarrow{\mathbb{P}}^{\nu,\sigma
u}}{d\overleftarrow{\mathbb{P}}^{\mu,\sigma
v}}\Big{)}\Big{]}=\mathbb{E}_{\overrightarrow{\mathbb{P}}^{\nu,\sigma
u}}\Big{[}\log\Big{(}\frac{d\overrightarrow{\mathbb{P}}^{\nu,\sigma
u}}{d\overrightarrow{\mathbb{P}}^{\nu,\sigma
r}}\frac{d\overleftarrow{\mathbb{P}}^{\eta,\sigma
v}}{d\overleftarrow{\mathbb{P}}^{\mu,\sigma
v}}\Big{)}\Big{]}=\mathbb{E}_{\overrightarrow{\mathbb{P}}^{\nu,\sigma
u}}\Big{[}\log\Big{(}\frac{d\overrightarrow{\mathbb{P}}^{\nu,\sigma
u}}{d\overrightarrow{\mathbb{P}}^{\nu,\sigma r}}\frac{d\eta}{d\mu}\Big{)}+\log
Z\Big{]}$
$\displaystyle=\mathbb{E}_{X\sim\overrightarrow{\mathbb{P}}^{\nu,\sigma
u}}\Big{[}\int_{0}^{T}\frac{1}{2}\|u_{s}(X_{s})-r_{s}(X_{s})\|^{2}ds+\log\left(\frac{d\eta}{d\mu}\right)(X_{T})+\log
Z\Big{]}\,.$ (9)

This suggests initializing from $\nu$ to estimate the loss (9) with the
current control $u^{\hat{\theta}}$ by simulating (2), followed by gradient
descent to optimize the $\hat{\theta}$-parameterized control policy and
iterating between the two steps can be a viable strategy, that could identify
$u^{\theta^{*}}$ eventually and used to run (2) to draw samples from $\mu$.
Note $\eta,\nu$ are simple distributions we have the freedom to pick, along
with $r(\cdot)$ in (7). The judicious choice of the common $v$ and $\nu$ in
(1) and (7) allow (9) to take on a control-theoretic interpretation [3], where
the cost (9) is composed of a running control cost and a terminal cost.

For a concrete example consider
$\nu=\mathcal{N}(0,\sigma^{2}I),r_{t}(x)=-x/2\sigma$ in (8), then process (7)
is simply OU in equilibrium, i.e., $\nu=\eta$ [31]. The purpose of introducing
the reference process (7) is to fix the backward drift $v$ in (1) so that
$\min_{u}\,\mathcal{L}_{KL}(u)$ from (9) is unique. However, this example
illustrates that in general unless $T\rightarrow\infty$, $\mathcal{L}_{KL}(u)$
in (9) can’t be minimized to $0$ since $\overleftarrow{\mathbb{P}}^{\mu,\sigma
v}_{0}\neq\nu$ under the OU process unless $T\rightarrow\infty$.

######  Remark 2.

The loss (9) enforce uniqueness (and correct marginals $\nu,\mu$) if minimized
to $0$, but doesn’t impose optimality of the interpolating trajectory in any
way. The approach in [7] that relies on training two drifts essentially
amounts to solving
$\min_{u,v}\,D_{KL}(\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}||\overleftarrow{\mathbb{P}}^{\mu,f+\sigma v})$ jointly, and as shown in
[25], although ensure correct marginals, has non-unique minimizers.

#### Reverse KL

In diffusion generative modeling, score-matching-based loss [14] can be seen
as minimizing the reverse KL over $s=(u-v)/\sigma$ by running the backward
process (3) to estimate the loss using samples from $\mu$. Using the Radon-
Nikodym derivative in Lemma 3, it gives ($C$ is independent of $s$ below)

$\displaystyle\mathcal{L}_{KL}(s)$
$\displaystyle=\mathbb{E}_{\overleftarrow{\mathbb{P}}^{\mu,\sigma
v}}\Big{[}\log\Big{(}\frac{d\overleftarrow{\mathbb{P}}^{\mu,\sigma
v}}{d\overrightarrow{\mathbb{P}}^{\nu,\sigma
v+\sigma^{2}s}}\Big{)}\Big{]}=\mathbb{E}_{\overleftarrow{\mathbb{P}}^{\mu,\sigma
v}}\Big{[}\int_{0}^{T}\frac{\sigma^{2}}{2}\|s_{t}(X_{t})\|^{2}\,dt+\sigma^{2}\int_{0}^{T}\nabla\cdot
s_{t}(X_{t})\,dt\Big{]}+C$
$\displaystyle=\mathbb{E}_{\overleftarrow{\mathbb{P}}^{\mu,\sigma
v}}\Big{[}\int_{0}^{T}\frac{\sigma^{2}}{2}\|s_{t}(X_{t})-\nabla_{x}\log
p_{t}^{\mu,\sigma v}(X_{t}|X_{T}=x_{T})\|^{2}\,dt\Big{]}+C^{\prime}\,.$ (10)

Here one fixes the backward drift $v$ in
$\overleftarrow{\mathbb{P}}^{\mu,\sigma v}$ so that $\nu$ is easy to sample
from (e.g., an OU process for one side of (1)). In the last transition above,
we used the integration by parts identity
$\mathbb{E}_{\rho_{t}}[\int_{0}^{T}s_{t}^{\top}\nabla\log\rho_{t}\,dt]=-\mathbb{E}_{\rho_{t}}[\int_{0}^{T}\nabla\cdot
s_{t}\,dt]\,.$ In practice crucially there will be an irreducible loss since
the terminals $\nu$ and $\overleftarrow{\mathbb{P}}_{0}^{\mu,\sigma v}$ don’t
match exactly for any finite $T$, but the dynamics
$\overrightarrow{\mathbb{P}}^{\nu,\sigma v+\sigma^{2}s}$ still makes sense as
for two processes $(p_{t})_{t},(q_{t})_{t}$ with different initializations
that share the same drift $\sigma v+\sigma^{2}s^{*}$,
$\partial_{t}\,D_{KL}(p_{t}\|q_{t})=-\mathbb{E}_{p_{t}}[\|\nabla\log\frac{p_{t}}{q_{t}}\|^{2}]\leq
0$ contracts (c.f. Remark 17), although we have not been quantitative about
the decay rate. It is worth noting that such approach doesn’t require re-
generating trajectories as (9) does.

We therefore see that either the approach of (9) or (10) relies on fixing some
aspect of the process involving (1) to restore uniqueness of the loss
$\mathcal{L}_{KL}$, but the choice is mostly out of convenience. In both
cases, it results in an _one-parameter loss_ , and the minimizer is not
affected by the unknown constant $Z$. Successfully optimizing
$D_{KL}(\overrightarrow{\mathbb{P}}^{\nu,\sigma
u}\|\overleftarrow{\mathbb{P}}^{\mu,\sigma v})$ or
$D_{KL}(\overleftarrow{\mathbb{P}}^{\mu,\sigma
v}\|\overrightarrow{\mathbb{P}}^{\nu,\sigma u})$ to zero error will imply that
$u$ pushes $\nu$ to $\mu$ and $v$ vice versa, mimicking a noising/denoising
reversible procedure. In what follows in Section 3, we deviate from these
perspectives by adopting a control formulation that does not rely on mixing of
stochastic processes for transporting between two distributions.

##  3 Methodology

As we saw from Section 2, (1) relying on mixing property of diffusion process
can make the trajectory rather long; (2) there are many degrees of freedom in
transporting $\nu$ to $\mu$, neither of which is desirable for training
purpose. This gives us the motivation to turn to losses based on path measures
that can enforce a canonical choice. We will adopt an “optimal control"
perspective and leverage special properties of the SB problem to come up with
valid control objectives for this effort.

Consider over path space $\mathcal{C}([0,T];\mathbb{R}^{d})$, given a
reference measure $Q$, the _constrained_ optimization problem

$P^{*}=\arg\min_{P_{0}=\nu,P_{T}=\mu}\,D_{KL}(P\|Q)$ (11)

where we assume $Q$ admits the SDE representation (this can be thought of as a
prior)

$dX_{t}=f_{t}(X_{t})\,dt+\sigma dW_{t}\,,\quad X_{0}\sim\nu\,,$ (12)

which is a slight generalization of the classical case where typically $f=0$.
Furthermore the path measure $P$ is assumed to correspond to the following SDE
and we are interested in finding the optimal control $\phi$ for

$dX_{t}=[f_{t}(X_{t})+\sigma^{2}\nabla\phi_{t}(X_{t})]\,dt+\sigma
dW_{t}\,,\quad X_{0}\sim\nu$ (13)

that solves (11). Various perspectives on the Schrödinger Bridge problem
(11)-(12), which we heavily leverage in the next part for designing our
losses, are included in Appendix A, along with related methods for solving the
problem in different setups than what we consider in Appendix B.

###  3.1 Training Loss Proposal

Recall our focus is on solving the regularized optimal transport problem (11)
between $\nu$ and $\mu$, by learning the controls on the basis of samples from
$\nu$. Out of many joint couplings with correct marginals (i.e., transport
maps), the choice of a reference process will select a particular trajectory
$\rho_{t}$ between $\nu$ and $\mu$. The execution of this plan crucially
hinges on two ingredients: (1) general backward / forward change of measure
formula given in Lemma 3; (2) properties of the optimal drifts for SB from
Section A, some of which can be exploited for training the controls.
Challenge, as emphasized before, is we need to be able to estimate the
resulting loss and ensuring that successful optimization guarantees
convergence to the unique solution dictated by the SB. Proposition 1 below
serves as our main result, where we show that adding regularization can
accomplish these. All proofs are deferred to Appendix C.

######  Proposition 1 (Regularized Control Objective).

For the problem of (11), the following losses are valid:

  1. (a)

$\arg\min_{\nabla\phi,\nabla\psi}\,D_{KL}(\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}\|\overleftarrow{\mathbb{P}}^{\mu,f-\sigma^{2}\nabla\psi})+\lambda\cdot\mathbb{E}_{X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}\big{[}\int_{0}^{T}\frac{\sigma^{2}}{2}\|\nabla\phi_{t}(X_{t})\|^{2}dt\big{]}$

  2. (b)

$\arg\min_{\nabla\phi,\nabla\psi}\,\text{Var}_{\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}\left[\log\left(\frac{\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}{\overleftarrow{\mathbb{P}}^{\mu,f-\sigma^{2}\nabla\psi}}\right)\right]+$

$\displaystyle\text{Var}_{X\sim\overrightarrow{\mathbb{P}}^{\nu,f}}\Big{(}\psi_{T}(X_{T})-\psi_{0}(X_{0})+\int_{0}^{T}(-\frac{\sigma^{2}}{2}\|\nabla\psi_{t}\|^{2}+\nabla\cdot
f_{t}-\sigma^{2}\Delta\psi_{t})(X_{t})\,dt-\sigma\int_{0}^{T}\nabla\psi_{t}(X_{t})^{\top}dW_{t}\Big{)}$
or
$\displaystyle\text{Var}_{X\sim\overrightarrow{\mathbb{P}}^{\nu,f}}\Big{(}\phi_{T}(X_{T})-\phi_{0}(X_{0})+\frac{\sigma^{2}}{2}\int_{0}^{T}\|\nabla\phi_{t}\|^{2}(X_{t})\,dt-\sigma\int_{0}^{T}\nabla\phi_{t}(X_{t})^{\top}dW_{t}\Big{)}$

  3. (c)

$\arg\min_{\phi,\psi}\,\text{Var}_{X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}\left((\phi_{T}+\psi_{T}-\log\mu)(X_{T})\right)+\text{Var}_{X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}\left((\phi_{0}+\psi_{0}-\log\nu)(X_{0})\right)+$

$\displaystyle\text{Var}_{X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}\Big{(}\phi_{T}(X_{T})-\phi_{0}(X_{0})-\frac{\sigma^{2}}{2}\int_{0}^{T}\|\nabla\phi_{t}\|^{2}(X_{t})\,dt-\sigma\int_{0}^{T}\nabla\phi_{t}(X_{t})^{\top}dW_{t}\Big{)}+$
$\displaystyle\text{Var}_{X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}\Big{(}\psi_{T}(X_{T})-\psi_{0}(X_{0})-\int_{0}^{T}\Big{(}\frac{\sigma^{2}}{2}\|\nabla\psi_{t}\|^{2}+\nabla\cdot(\sigma^{2}\nabla\psi_{t}-f_{t})+\sigma^{2}\nabla\psi_{t}^{\top}\nabla\phi_{t}\Big{)}(X_{t})\,dt$
$\displaystyle-\sigma\int_{0}^{T}\nabla\psi_{t}(X_{t})^{\top}dW_{t}\Big{)}$

  4. (d)

$\arg\min_{\nabla\phi_{t},\nabla\log\rho_{t}}\,D_{KL}(\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi_{t}}\|\overleftarrow{\mathbb{P}}^{\mu,f+\sigma^{2}\nabla\phi_{t}-\sigma^{2}\nabla\log\rho_{t}})+\lambda\cdot\mathbb{E}_{\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}\big{[}\int_{0}^{T}\frac{\sigma^{2}}{2}\|\nabla\phi_{t}(X_{t})\|^{2}\,dt\big{]}$

In particular, if the loss is minimized to $0$ for (b) and (c), the resulting
$\nabla\phi^{*}$ and the corresponding backward drift $\nabla\psi^{*}$ solve
the Schrödinger Bridge problem (11) from $\nu$ to $\mu$, that can in turn be
used for sampling from $\mu$ by running (13). In all cases,
$X_{0}=x_{0}\sim\nu$ is assumed given as initial condition.

Both (a) and (d) are guided by “minimum-action" principle w.r.t a reference
(i.e., minimum control energy spent). (b) bases itself on a reformulation of
the HJB PDE involving the optimal control, and (c) is grounded in the FBSDE
system for SB optimality [7] (c.f. Remark 14 for the equivalence between
particular SDEs and path measures). In all cases, the objective is a _two-
parameter loss_ that also allows us to recover the score
$\sigma^{2}\nabla\log\rho_{t}$ from the returned solution. Note that variance
is taken w.r.t the uncontrolled process in (b) and w.r.t the controlled
process in (c). Losses (a) and (d), being conversion from the constrained
problem (and each other), will need $\lambda$ to be picked relatively small so
that the first part of the objective becomes $0$ to identify the unique SB
solution, but we establish a bound on the optimal objective value in Remark 8.

The loss in (c) is different from log-variance divergence over path space
$\text{Var}_{\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}[\log(\frac{d\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}{d\overleftarrow{\mathbb{P}}^{\mu,f-\sigma^{2}\nabla\psi}})]\,,$
which will _not_ guarantee finding the optimal path, whereas the objective in
(c) incorporates the dynamics of _two_ controlled (and coupled) dynamics w.r.t
$\overrightarrow{\mathbb{P}}^{\nu,f}$ that we know how to characterize
optimality for from results in Section A (c.f. additional discussion in
Section 4.1 and the proof of Proposition 1). In fact, another way to view the
variance regularizers in loss (b) and (c) is through the SDE representation of
the controls from Lemma 4 in Appendix C and observe that the variance
condition precisely encodes the optimally-controlled dynamical information. It
is important to note that we are evaluating the change of $\phi,\psi$ on a
particular trajectory, rather than tracking the cumbersome evolution of
density over the full $\mathbb{R}^{d}$ space, as what the PDEs (43) may
suggest. The proof of these results heavily use the factorization property
(46) satisfied by the optimal coupling, and exploit different ways to encode
the optimality condition.

######  Remark 3 (Stochastic gradient w.r.t controls at optimality).

In [25], the authors show that log-variance has the advantage of having
variance of gradient = 0 at the optimal $\phi^{*},\psi^{*}$, which is not true
for $D_{KL}$ in general, and has consequence for gradient-based updates such
as Algorithm 1. Similar argument applies to the variance regularizer we
consider (e.g., loss (b)). For this, we look at the Gâteaux derivative of the
variance function $V$ in an arbitrary direction $\tau$ since from the chain
rule

$\frac{\delta}{\delta\phi}V(\phi,\psi;\tau):=\frac{d}{d\epsilon}\bigg{|}_{\epsilon=0}V(\phi+\epsilon\tau,\psi)\Rightarrow\partial_{\theta_{i}}V(\phi_{\theta},\psi_{\gamma})=\frac{\delta}{\delta\phi}\bigg{|}_{\phi=\phi_{\theta}}V(\phi,\psi_{\gamma};\partial_{\theta_{i}}\phi_{\theta})\,.$
(14)

Now if $\hat{V}$ is the Monte-Carlo estimate of the variance of a random
quantity $g(\cdot)$, it is always the case that (we use
$\frac{\delta}{\delta\phi}(\cdot)_{\tau}$ to denote derivative in the $\tau$
direction)

$\frac{\delta}{\delta\phi}\hat{V}(\phi,\psi;\tau)=\frac{\delta}{\delta\phi}(\hat{\mathbb{E}}[g(\phi,\psi)^{2}]-\hat{\mathbb{E}}[g(\phi,\psi)]^{2})_{\tau}=2\hat{\mathbb{E}}[g(\phi,\psi)\frac{\delta}{\delta\phi}g(\phi,\psi)_{\tau}]-2\hat{\mathbb{E}}[g(\phi,\psi)]\frac{\delta}{\delta\phi}\hat{\mathbb{E}}[g(\phi,\psi)]_{\tau}\,,$

therefore if $g(\phi^{*},\psi^{*})=c$ almost surely for every i.i.d sample,
such as the regularizer in loss (b), the derivative w.r.t the control $\phi$
in direction $\partial_{\theta_{i}}\phi_{\theta}$ is $0$ at optimality,
implying
$\text{Var}(\partial_{\theta_{i}}\hat{V}(\phi_{\theta^{*}},\psi_{\gamma^{*}}))=0$
using (14).

It is natural to ask if one could replace the variance regularizer
$\text{Var}(\cdot)$ with a moment regularizer $\mathbb{E}[|\cdot|^{2}]$ in
e.g., loss (c). However, while the variance is oblivious to constant shift,
the moment loss will require knowledge of the normalizing constant $Z$ of the
target $\mu$ to make sense. In the case of loss (b), using (5) and Lemma 4, an
alternative proposal based on moment regularizer for e.g., $\phi$ can be

$\mathbb{E}_{X\sim\overrightarrow{\mathbb{P}}^{\nu,f}}\Big{[}\Big{|}\phi_{T}(X_{T})-\phi_{0}(X_{0})+\frac{\sigma^{2}}{2}\int_{0}^{T}\|\nabla\phi_{t}\|^{2}(X_{t})\,dt\Big{|}^{2}\Big{]}=:\mathbb{E}[\bar{g}(\phi)^{2}]\,.$

However, the Gâteaux derivative of the loss in this case takes the form of

$\frac{\delta}{\delta\phi}(\hat{\mathbb{E}}[\bar{g}(\phi)^{2}])_{\tau}=2\hat{\mathbb{E}}[\bar{g}(\phi)\frac{\delta}{\delta\phi}\bar{g}(\phi)_{\tau}]\,,$

which means that while $\mathbb{E}[\bar{g}(\cdot)]=0$ vanish for the optimal
control $\phi^{*}$ in expectation, with stochastic gradient estimate
$\hat{\mathbb{E}}[\bar{g}(\cdot)]\neq 0$ generally at $\phi^{*}$, hence
yielding non-vanishing variance and posing challenges for optimization. For
our variance regularizer in loss (b) where $g(\phi^{*})=0$ almost surely for
all trajectories, it implies zero-variance and will therefore identify the
optimal $\phi^{*}$ even with mini-batches. In terms of the actual minimum
value of the loss, our loss (b) obeys $\hat{V}[g(\phi^{*})]=0$ whereas using
Itô’s isometry
$\hat{\mathbb{E}}[\bar{g}(\phi^{*})^{2}]=\hat{\mathbb{E}}[\int\sigma^{2}\|\nabla\phi^{*}_{t}\|^{2}\,dt]\neq
0$ is much less predictable for monitoring the performance of the final
control. This is somewhat similar to [35] where the authors observe that
including random $0$-mean martingale terms is important for variance reduction
in a different context.

######  Remark 4 (Comparisons).

In [32], the authors propose
$\int_{0}^{T}\mathbb{E}\big{|}\partial_{t}\phi+f^{\top}\nabla\phi+\frac{\sigma^{2}}{2}\Delta\phi+\frac{\sigma^{2}}{2}\|\nabla\phi\|^{2}\big{|}(X_{t})dt$
as the HJB regularizer (c.f. (35)) on top of
$D_{KL}(\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}\|\overleftarrow{\mathbb{P}}^{\mu,f-\sigma^{2}\nabla\psi})$
as the training loss, inspired by PINN [24]. As is clear from the proof in
Proposition 1, it is equivalent to our loss (b) (see also related result in
[19]). Combining (35)-(36) we can also get another HJB for the backward drift
$\nabla\psi$, which read as $\partial_{t}\psi+f^{\top}\nabla\psi+\nabla\cdot
f-\frac{\sigma^{2}}{2}\Delta\psi-\frac{\sigma^{2}}{2}\|\nabla\psi\|^{2}=0\,.$
By trading a PDE constraint for a SDE one (based on likelihood ratio of path
measures), we can avoid evaluating the divergence term for the discrete
updates. It remains to note the additional benefit of “sticking to the
landing" (c.f. Remark 3). We refer to Section 4 for numerical illustration.

In Appendix D, we consider the extension of such training methodologies to the
case of second-order dynamics in the augmented $(X,V)$ space à la under-damped
Langevin.

###  3.2 Discretization and Practical Implementation

We discuss implementation of our losses and estimators here. In practice, with
imperfect controls from the training procedure, one can perform importance
sampling to correct for the bias / improve on the estimate. We also draw a
connection between the optimal $Z$ estimator and the optimal controls in
Remark 15.

######  Proposition 2 (Importance Sampling).

The following can be used to give unbiased estimate of the normalizing
constant for $p_{\text{target}}$:

  1. (1)

For the optimal $\phi^{*},\psi^{*}$, with
$X_{t}\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi^{*}}$,

$Z=\frac{\mu(X_{T})}{\nu(X_{0})}\exp\Big{(}\frac{\sigma^{2}}{2}\int_{0}^{T}\nabla\cdot(\nabla\phi_{t}^{*}-\nabla\psi_{t}^{*})(X_{t})\,dt+\int_{0}^{T}\nabla\cdot
f_{t}(X_{t})dt\Big{)}\,.$

  2. (2)

For any potentially suboptimal $\phi,\psi$, with
$X_{t}\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}$,

$Z=\mathbb{E}\Big{[}\exp\Big{(}-\frac{\sigma^{2}}{2}\int_{0}^{T}\|\nabla\phi_{t}+\nabla\psi_{t}\|^{2}+\nabla\cdot(f_{t}-\sigma^{2}\nabla\psi_{t})dt-\sigma\int_{0}^{T}\nabla\phi_{t}+\nabla\psi_{t}dW_{t}-\log\frac{\nu(X_{0})}{\mu(X_{T})}\Big{)}\Big{]}\,.$

With a sub-optimal control $\nabla\phi$ and
$X_{t}\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}$, re-
weighting can be used to get an unbiased estimator of a statistics
$g\colon\mathbb{R}^{d}\rightarrow\mathbb{R}$ as
$\mathbb{E}_{\phi}[g(X_{T})w^{\phi}(X_{T})]/\mathbb{E}_{\phi}[w^{\phi}(X_{T})]=\mathbb{E}_{\phi^{*}}[g(X_{T})]=\mathbb{E}_{p_{\text{target}}}[g]$
for

$w^{\phi}(X)=\frac{d\mu(X_{T}^{\phi})}{d\nu(X_{0}^{\phi})}\exp\Big{(}\int_{0}^{T}\sigma^{2}\Delta\phi_{t}-\frac{\sigma^{2}}{2}\Delta\log\rho_{t}+\nabla\cdot
f_{t}\,dt\Big{)}\,.$

As in [32, Proposition F.1], it is possible to trade divergence term for a
backward integral when dealing with path integrals using (6). Thanks to the
fact that our various estimators and regularizers are built upon path
measures, the following lemma generalizes this idea and provides a recipe for
estimating the regularizer from Proposition 1 and the normalizing constant
from Proposition 2 with discrete-time updates. We work out loss (c) from
Proposition 1 below, since most of the other parts are straightforward to
adapt – we detail them in Section 4.1.

######  Lemma 1 (Discretized Loss and Estimator).

For $X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}$, the last
part of loss $(c)$ on $\psi$ can be estimated as ($\text{Var}_{N}$ denotes
empirical estimate of variance using $N$ samples)

$\text{Var}_{N}\Big{(}\psi_{K}(X_{K+1}^{i})-\psi_{0}(X_{0}^{i})+\frac{1}{2\sigma^{2}h}\sum_{k=0}^{K-1}\|X_{k+1}^{i}-X_{k}^{i}-f_{k}(X_{k}^{i})h\|^{2}-\|X_{k}^{i}-X_{k+1}^{i}+(f_{k+1}-\sigma^{2}\nabla\psi_{k+1})(X_{k+1}^{i})h\|^{2}\Big{)}\,.$

The importance-weighted $Z$-estimator from Proposition 2 can be approximated
as

$\frac{1}{N}\sum_{i=1}^{N}\exp\Big{(}\log\frac{\mu(X_{K}^{i})}{\nu(X_{0}^{i})}+\frac{1}{2}\sum_{k=0}^{K-1}\|z_{k}^{i}\|^{2}-\frac{1}{2\sigma^{2}h}\|X_{k}^{i}-X_{k+1}^{i}+(f_{k+1}-\sigma^{2}\nabla\psi_{k+1})(X_{k+1}^{i})h\|^{2}\Big{)}\,.$
(15)

In both cases for $i\in[N]$ independently, $z_{k}^{i}\sim\mathcal{N}(0,I)$,
$X_{k+1}^{i}=X_{k}^{i}+(f_{k}(X_{k}^{i})+\sigma^{2}\nabla\phi_{k}(X_{k}^{i}))h+\sigma\sqrt{h}\cdot
z_{k}^{i}\,.$

######  Remark 5 (Optimal $\log Z$ estimator).

As observed in [32], discretizations with backward integral have the
additional benefit of giving ELBO for the normalizing constant $Z$ of
$p_{\text{target}}$ with the estimator

$\hat{Z}=\frac{\mu(X_{K})q^{v}(X_{0:K-1}|X_{K})}{\nu(X_{0})p^{u}(X_{1:K}|X_{0})}$
(16)

in the sense that

$\mathbb{E}_{\nu(X_{0})p^{u}(X_{1:K}|X_{0})}[\log\hat{Z}]\leq\log(\mathbb{E}_{\nu(X_{0})p^{u}(X_{1:K}|X_{0})}[\hat{Z}])=\log\left[\int\mu(X_{K})\,dX_{K}\right]=\log(Z)\,.$

This aspect is exploited in Appendix E when comparing the $\log Z$ estimator,
since note our estimator (15) from Lemma 1 can also be understood as a ratio
of two discrete chains similar to (16). It is reminiscent of the philosophy
adopted in annealed importance sampling (AIS) with extended target, but in our
case, the backward kernel $q^{v}(\cdot)$ is the _time-reversal_ of the forward
one, which can be shown to be the optimal transition kernel minimizing the
variance of the resulting evidence estimate [12]. In addition, the forward
kernel $p^{u}(\cdot)$ in our case follows an optimal trajectory that ensures
the chain reaches the _exact_ target $\mu$ rapidly.

Putting everything together, we give the final proposed algorithm below.

Algorithm 1 Control Objective Training for Sampling

Initial draw $(X_{0}^{i,(0)})_{i=1}^{N}\in\mathbb{R}^{d}\sim\nu$ independent,
initial controls $\phi^{(0)},\psi^{(0)}$

Un-normalized density $\mu$, drift $f$, num of time steps $K$, num of
iterations $T$

for $t=0,\cdots,T-1$ do

Simulate (17) with current control $(\nabla\phi^{(t)}_{k})_{k=0,\cdots,K}$,
obtain $(X_{k}^{n,(t)})_{k=0,\cdots,K}$ for $n=1,\cdots,N$

Estimate the loss (18) $+$ discretized regularizer (c.f. Lemma 1 & Section
4.1) and the gradient w.r.t the two parameterized controls using the samples
$(X_{k}^{n,(t)})_{k=0,\cdots K,n=1,\cdots,N}$

Gradient update on the parameters to obtain $\nabla\phi^{(t+1)}$ and
$\nabla\psi^{(t+1)}$

end for

return $X_{K}^{1,(T)},\cdots,X_{K}^{N,(T)}$ as $N$ samples from $\mu$ with
their importance weights $w^{\phi^{(T)}}(X_{K}^{n,(T)})$ (90), together with
the weighted $Z$ estimator (15)

##  4 Numerics and Comparisons

In this section, we instantiate our main contributions (Proposition 1 and
Lemma 1) and offer numerical evidence on their advantages compared to existing
proposals for solving an optimal trajectory problem in a typical MCMC setup.
We highlight that our algorithm does not use gradient information from the
target $\nabla\log\mu$ as e.g., Langevin would.

###  4.1 Algorithm specification

Below for simplicity we pick the reference process to be a Brownian motion
with $f=0$ and $\lambda$ is a parameter that we tune for best performance, but
in theory any $\lambda>0$ would work. This aspect deviates from other
constrained formulation of the problem.

  1. (1)

PINN-regularization [32]: for $i=1,\cdots,n$, and $Z\sim\mathcal{N}(0,I)$
independently draw in parallel

$x_{k+1}^{i}=x_{k}^{i}+\sigma^{2}h\nabla\phi(x_{k}^{i},kh)+\sigma\sqrt{h}Z_{k}^{i},\,x_{0}^{i}\sim\nu$
(17)

for $k=0,\cdots,K$ with stepsize $h=c/(K+1)$ for some $c\geq 1$. Using these
trajectories $\\{x_{k}^{i}\\}$, minimize over $\phi$ and $s$,

$\displaystyle\frac{1}{K+1}\text{Var}_{n}\Big{[}\log\frac{\nu(x_{0}^{i})}{\mu(x_{K+1}^{i})}+\sum_{k=0}^{K}\frac{1}{2\sigma^{2}h}\|x_{k}^{i}-x_{k+1}^{i}+\sigma^{2}h\left(\nabla\phi(x_{k+1}^{i},(k+1)h)-\nabla
s(x_{k+1}^{i},(k+1)h)\right)\|^{2}$
$\displaystyle-\sum_{k=0}^{K}\frac{1}{2\sigma^{2}h}\|x_{k+1}^{i}-x_{k}^{i}-\sigma^{2}h\nabla\phi(x_{k}^{i},kh)\|^{2}\Big{]}$
(18)
$\displaystyle+\lambda\cdot\frac{1}{n}\sum_{i=1}^{n}\Big{[}\sum_{k=0}^{K}\Big{|}\partial_{t}\phi(x_{k}^{i},kh)+\frac{\sigma^{2}}{2}\Delta\phi(x_{k}^{i},kh)+\frac{\sigma^{2}}{2}\|\nabla\phi(x_{k}^{i},kh)\|^{2}\Big{|}\cdot
h\Big{]}\,.$ (19)

The first term (18) is an estimate of the log-variance divergence between the
two path measures. Above $\phi(x,t),s(x,t)$ are two neural networks that take
$t\in\mathbb{R}$ and $x\in\mathbb{R}^{d}$ as inputs and maps to $\mathbb{R}$.
We repeat (17) and (19) several times, and compute statistics using the latest
samples $\\{x_{K+1}^{i}\\}_{i=1}^{n}$.

  2. (2)

Variance-regularization (loss (b) from Proposition 1): Simulate trajectories
(17) as before, but additionally simulate $\\{y_{k}^{i}\\}_{k=0}^{K}$ as
follows and cache them:

$y_{k+1}^{i}=y_{k}^{i}+\sigma\sqrt{h}Z_{k}^{i},\,y_{0}^{i}\sim\nu\,.$ (20)

Minimize over $\phi,s$ the following discretized loss

$\displaystyle\eqref{eqn:log_divergence}+\frac{\lambda}{K+1}\cdot$
$\displaystyle\text{Var}_{n}\Big{[}\phi(y_{K+1}^{i},(K+1)h)-\phi(y_{0}^{i},0)+\frac{1}{2\sigma^{2}h}\sum_{k=0}^{K}\|y_{k+1}^{i}-y_{k}^{i}-\sigma^{2}\nabla\phi(y_{k}^{i},kh)h\|^{2}-\|y_{k+1}^{i}-y_{k}^{i}\|^{2}\Big{]}.$
(21)

Alternate between (17) and (21) several times. Above $\text{Var}_{n}$ denotes
the empirical variance across the $n$ trajectories $\\{y_{k}^{i}\\}_{i=1}^{n}$
of the quantity inside $[\cdot]$. The loss (21), as the proof of Proposition 1
shows, comes from the fact that along the prior
$X\sim\overrightarrow{\mathbb{P}}^{\nu,f}$,

$\log\Big{(}\frac{d\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi_{t}}}{d\overrightarrow{\mathbb{P}}^{\nu,f}}\Big{)}(X)=\int_{0}^{T}-\frac{\sigma^{2}}{2}\|\nabla\phi_{t}\|^{2}dt+\int_{0}^{T}\sigma\nabla\phi_{t}^{\top}\,dW_{t}=\phi_{T}(X_{T})-\phi_{0}(X_{0})\,,$
(22)

and we discretized the Radon-Nikodym derivative similar to how it was done in
the KL divergence
$D_{KL}(\overrightarrow{\mathbb{P}}^{\nu,\sigma^{2}\nabla\phi}\|\overleftarrow{\mathbb{P}}^{\mu,-\sigma^{2}\nabla\psi})$
(see Lemma 1 for a similar derivation).

  3. (3)

Instead of the regularizer (21), another discretization of condition (22) can
be a TD-like regularizer similar in spirit to [18]:

$\eqref{eqn:log_divergence}+\lambda\cdot\frac{h}{n}\sum_{i=1}^{n}\sum_{k=0}^{K}\Big{|}\phi(y_{k+1}^{i},(k+1)h)-\phi(y_{k}^{i},kh)+\frac{\sigma^{2}h}{2}\|\nabla\phi(y_{k}^{i},kh)\|^{2}-\sigma\sqrt{h}\nabla\phi(y_{k}^{i},kh)^{\top}Z_{k}^{i}\Big{|}\,,$
(23)

where the $Z_{k}^{i}$’s are re-used from (20). The loss above can be justified
with Lemma 4.

######  Remark 6 (Backward drift).

In all (19), (21), (23), we parameterize the forward / backward drift as
$\sigma^{2}\nabla\phi,\sigma^{2}\nabla\phi-\sigma^{2}\nabla s$. A
regularization on the backward drift involving $s$ is also possible for PINN
(19):

$\eqref{eqn:log_divergence}+\lambda\cdot\frac{h}{n}\sum_{i=1}^{n}\sum_{k=0}^{K}\Big{|}(\partial_{t}s-\partial_{t}\phi)(x_{k}^{i},kh)+\frac{\sigma^{2}}{2}\left[(\Delta\phi-\Delta
s)(x_{k}^{i},kh)-\|(\nabla s-\nabla\phi)(x_{k}^{i},kh)\|^{2}\right]\Big{|}\,.$
(24)

And similarly for variance loss (21), which will read as

$\displaystyle\eqref{eqn:log_divergence}$
$\displaystyle+\frac{\lambda}{K+1}\cdot\text{Var}_{n}\Big{[}(s-\phi)(y_{K+1}^{i},(K+1)h)-(s-\phi)(y_{0}^{i},0)+$
$\displaystyle\sum_{k=0}^{K}\frac{1}{2\sigma^{2}h}\Big{(}\|y_{k+1}^{i}-y_{k}^{i}\|^{2}-\|y_{k}^{i}-y_{k+1}^{i}+\sigma^{2}h(\nabla\phi-\nabla
s)(y_{k+1}^{i},(k+1)h)\|^{2}\Big{)}\Big{]}$ (25)

and TD loss (23) becomes

$\displaystyle\eqref{eqn:log_divergence}$
$\displaystyle+\lambda\cdot\frac{1}{n}\sum_{i=1}^{n}\sum_{k=0}^{K}h\cdot\Big{|}(s-\phi)(y_{k+1}^{i},(k+1)h)-(s-\phi)(y_{k}^{i},kh)$
$\displaystyle-\frac{\sigma^{2}h}{2}\|(\nabla
s-\nabla\phi)(y_{k}^{i},kh)\|^{2}-\sigma^{2}(\Delta
s-\Delta\phi)(y_{k}^{i},kh)+\sigma\sqrt{h}(\nabla\phi-\nabla
s)(y_{k}^{i},kh)^{\top}Z_{k}^{i}\Big{|}\,.$ (26)

But these should witness similar behavior as the regularizer on the forward
control that we experiment with (using same trajectories). Notice crucially
there is no divergence term in (25) in contrast to (26).

  4. (4)

Separately-controlled loss (loss (c) from Proposition 1): Simulate (17) as
before, with the $n$ trajectories $\\{x_{k}^{i}\\}$, minimize over $\phi,\psi$
the following discretized loss (c.f. Lemma 1):

$\displaystyle\text{Var}_{n}\Big{[}\psi(x_{K+1}^{i},(K+1)h)+\phi(x_{K+1}^{i},(K+1)h)-\log\mu(x_{K+1}^{i})\Big{]}$
(27)
$\displaystyle+\text{Var}_{n}\Big{[}\psi(x_{0}^{i},0)+\phi(x_{0}^{i},0)-\log\nu(x_{0}^{i})\Big{]}+$
$\displaystyle\frac{\lambda}{K+1}\text{Var}_{n}\Big{[}\psi(x_{K+1}^{i},(K+1)h)-\psi(x_{0}^{i},0)+\frac{1}{2\sigma^{2}h}\sum_{k=0}^{K}\|x_{k+1}^{i}-x_{k}^{i}\|^{2}-\|x_{k}^{i}-x_{k+1}^{i}-\sigma^{2}h\nabla\psi(x_{k+1}^{i},(k+1)h)\|^{2}\Big{]}$
$\displaystyle+\frac{\lambda}{K+1}\text{Var}_{n}\Big{[}\phi(x_{0}^{i},0)-\phi(x_{K+1}^{i},(K+1)h)+\frac{1}{2\sigma^{2}h}\sum_{k=0}^{K}\|x_{k+1}^{i}-x_{k}^{i}\|^{2}-\|x_{k+1}^{i}-x_{k}^{i}-\sigma^{2}h\nabla\phi(x_{k}^{i},kh)\|^{2}\Big{]}.$

Notice that it consists of $4$ variance terms, and we alternate between
simulating (17) and updating $\phi,\psi$ from (27). For this loss, we
parameterize the forward / backward drift as
$\sigma^{2}\nabla\phi,-\sigma^{2}\nabla\psi$ so it’s more symmetric. One can
also consider SDE-based discretization for the last 2 terms but will incur
additional Laplacian and divergence terms as suggested by Lemma 4, (68)-(69).

#### Comparison between discretized losses

We emphasize that the discretized variance regularizers (21) and (25) wouldn’t
be available without the path-wise stochastic process perspective. The FBSDE
view (Lemma 4) will naturally lend to TD-like regularizers (23) and (26)
similar to [18], but the crucial differences are: (1) we base the estimate on
reference measure $\overrightarrow{\mathbb{P}}^{\nu,f}$ therefore there’s no
need to differentiate through the generated trajectory when optimize the loss
over $\nabla\phi$ or $\nabla\psi$; (2) [18] consider dynamics with mean-field
interaction and different discretization. In loss (27) we _separately_ impose
optimality condition on

$\log\left(\frac{d\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi_{t}}}{d\overrightarrow{\mathbb{P}}^{\nu,f}}\right)(X)\quad\text{and}\quad\log\left(\frac{d\overrightarrow{\mathbb{P}}^{\nu,f}}{d\overleftarrow{\mathbb{P}}^{\mu,f-\sigma^{2}\nabla\psi_{t}}}\right)(X)\,,$

for $X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi_{t}}$ to
take factorized forms, as opposed to looking at divergence metrics on their
sum
$\log\left(\frac{d\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}{d\overleftarrow{\mathbb{P}}^{\mu,f-\sigma^{2}\nabla\psi}}\right)$
only, effectively erasing the “SB optimality enforcement" part. The two
discretized losses (21)/(25) and (27) both enjoy the variance reduction
property as elaborated in Remark 3, without the need to evaluate expensive
Laplacian terms (as the PDE-based PINN approach (19)/(24) or TD approach (26)
would require).

###  4.2 Result

In our experiments, we picked the following 4 targets as benchmark as they
capture different properties of distributions that can affect the training of
the controls. These properties (multimodality, spatially varying curvature)
pose challenges for MCMC methods and are ideal examples to demonstrate the
differences between the regularizers.

Figure 1: Weighted marginal for double well

| PINN | Variance | TD | SC
---|---|---|---|---
Normal | $6.317\,639\,350\,891\,113$ | $2.882\,527\,828\,216\,552\,7$ | $2.013\,974\,666\,595\,459$ | $1.304\,690\,361\,022\,949\,2$
Funnel | $8.854\,369\,163\,513\,184$ | $3.560\,049\,533\,843\,994$ | $3.550\,189\,018\,249\,511\,7$ | $3.390\,867\,233\,276\,367$
GMM | $5.802\,513\,122\,558\,594$ | $8.346\,628\,189\,086\,914$ | $4.097\,251\,892\,089\,844$ | $4.393\,266\,677\,856\,445$
Double well | $12.274\,190\,902\,709\,961$ | $5.391\,733\,169\,555\,664$ | $4.436\,683\,654\,785\,156$ | $1.761\,485\,099\,792\,480\,5$

Table 1: $-\log Z$ estimator (lower is better)

Additional supporting numerical experiments and details can be found in
Appendix E. Across our experiments, we observe that our separately controlled
loss (27) is generally much better compared to the PINN loss (19) that tends
to exhibit mode-seeking behavior, while TD (23) and variance regularizer (21)
can sometimes be comparable. In practice, we also observe that the separately
controlled loss is less sensitive to tuning parameters, and the training loss
trajectory is often smoother.

##  5 Conclusion

We exploited the connections between diffusion generative modeling, stochastic
control and optimal transport, with the goal of sampling from high-
dimensional, complex distributions in mind. This is orthogonal to MCMC-based
approaches, and is accomplished by a more “learning-driven" methodology on the
optimal control / drift that can be trained with a suitable control objective.
While there are a lot of flexibility in the design (that also makes the case
for SB more compelling than existing pathwise samplers mentioned in Section
2.2), we have illustrated that careful choice is needed for both numerical
implementation and generalizability of the framework.

More generally, the mapping between target $p_{\text{target}}(x)$ and control
$\phi(x_{t},t)$ can also be thought of as a form of operator learning mapping
between functions in infinite dimensional space, for which universal
approximation theorems have emerged recently for various neural network
architectures [10]. More specifically, one can view the task as learning the
solution operator of the coupled PDEs (35)-(36). We believe quantitative
estimates for its parameter and data complexity is an interesting future
direction to pursue. Exploiting such methods in the context of molecular
dynamics simulation for sampling transition path is also exciting and
development is already underway [13].

## References

  * [1] Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden.  Stochastic interpolants: A unifying framework for flows and diffusions.  arXiv preprint arXiv:2303.08797, 2023.
  * [2] Jean-David Benamou and Yann Brenier.  A computational fluid mechanics solution to the Monge-Kantorovich mass transfer problem.  Numerische Mathematik, 84(3):375–393, 2000.
  * [3] Julius Berner, Lorenz Richter, and Karen Ullrich.  An optimal control perspective on diffusion-based generative modeling.  arXiv preprint arXiv:2211.01364, 2022.
  * [4] Kenneth F Caluya and Abhishek Halder.  Wasserstein proximal algorithms for the Schrödinger bridge problem: Density control with nonlinear drift.  IEEE Transactions on Automatic Control, 67(3):1163–1178, 2021.
  * [5] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud.  Neural ordinary differential equations.  Advances in neural information processing systems, 31, 2018.
  * [6] Tianrong Chen, Guan-Horng Liu, Molei Tao, and Evangelos Theodorou.  Deep momentum multi-marginal Schrödinger bridge.  Advances in Neural Information Processing Systems, 36, 2023.
  * [7] Tianrong Chen, Guan-Horng Liu, and Evangelos Theodorou.  Likelihood Training of Schrödinger Bridge using Forward-Backward SDEs Theory.  In International Conference on Learning Representations, 2021.
  * [8] Yongxin Chen, Tryphon T Georgiou, and Michele Pavon.  Stochastic control liaisons: Richard Sinkhorn meets Gaspard Monge on a Schrödinger Bridge.  SIAM Review, 63(2):249–313, 2021.
  * [9] Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet.  Diffusion Schrödinger bridge with applications to score-based generative modeling.  Advances in Neural Information Processing Systems, 34:17695–17709, 2021.
  * [10] Tim De Ryck and Siddhartha Mishra.  Generic bounds on the approximation error for physics-informed (and) operator learning.  Advances in Neural Information Processing Systems, 35:10945–10958, 2022.
  * [11] Tim Dockhorn, Arash Vahdat, and Karsten Kreis.  Score-Based Generative Modeling with Critically-Damped Langevin Diffusion.  In International Conference on Learning Representations (ICLR), 2022\.
  * [12] Arnaud Doucet, Will Grathwohl, Alexander G Matthews, and Heiko Strathmann.  Score-based diffusion meets annealed importance sampling.  Advances in Neural Information Processing Systems, 35:21482–21494, 2022.
  * [13] Lars Holdijk, Yuanqi Du, Ferry Hooft, Priyank Jaini, Berend Ensing, and Max Welling.  Stochastic Optimal Control for Collective Variable Free Sampling of Molecular Transition Paths.  Advances in Neural Information Processing Systems, 36, 2023.
  * [14] Aapo Hyvärinen and Peter Dayan.  Estimation of non-normalized statistical models by score matching.  Journal of Machine Learning Research, 6(4), 2005.
  * [15] Hicham Janati, Boris Muzellec, Gabriel Peyré, and Marco Cuturi.  Entropic optimal transport between unbalanced Gaussian measures has a closed form.  Advances in neural information processing systems, 33:10468–10479, 2020.
  * [16] Christopher Jarzynski.  Nonequilibrium equality for free energy differences.  Physical Review Letters, 78(14):2690, 1997.
  * [17] Christian Léonard.  A survey of the Schrödinger problem and some of its connections with optimal transport.  Discrete and Continuous Dynamical Systems-Series A, 34(4):1533–1574, 2014.
  * [18] Guan-Horng Liu, Tianrong Chen, Oswin So, and Evangelos Theodorou.  Deep generalized Schrödinger bridge.  Advances in Neural Information Processing Systems, 35:9374–9388, 2022.
  * [19] Nikolas Nüsken and Lorenz Richter.  Interpolating between BSDEs and PINNs: deep learning for elliptic and parabolic boundary value problems.  arXiv preprint arXiv:2112.03749, 2021.
  * [20] Marcel Nutz.  Introduction to entropic optimal transport.  Lecture notes, Columbia University, 2021.
  * [21] Michele Pavon, Giulio Trigila, and Esteban G Tabak.  The Data-Driven Schrödinger Bridge.  Communications on Pure and Applied Mathematics, 74(7):1545–1573, 2021.
  * [22] Stefano Peluchetti.  Diffusion Bridge Mixture Transports, Schrödinger Bridge Problems and Generative Modeling.  arXiv preprint arXiv:2304.00917, 2023.
  * [23] Gabriel Peyré and Marco Cuturi.  Computational optimal transport: With applications to data science.  Foundations and Trends in Machine Learning, 11(5-6):355–607, 2019\.
  * [24] Maziar Raissi, Paris Perdikaris, and George E Karniadakis.  Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations.  Journal of Computational physics, 378:686–707, 2019.
  * [25] Lorenz Richter, Julius Berner, and Guan-Horng Liu.  Improved sampling via learned diffusions.  arXiv preprint arXiv:2307.01198, 2023.
  * [26] Erwin Schrödinger.  Über die umkehrung der naturgesetze.  Sitzungsberichte der preussischen Akademie der Wissenschaften, physikalische mathematische Klasse, 1931.
  * [27] Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet.  Diffusion Schrödinger bridge matching.  Advances in Neural Information Processing Systems, 36, 2023.
  * [28] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.  Score-Based Generative Modeling through Stochastic Differential Equations.  In International Conference on Learning Representations, 2020.
  * [29] Sep Thijssen and HJ Kappen.  Path integral control and state-dependent feedback.  Physical Review E, 91(3):032104, 2015.
  * [30] Belinda Tzen and Maxim Raginsky.  Theoretical guarantees for sampling and inference in generative models with latent diffusions.  In Conference on Learning Theory, pages 3084–3114. PMLR, 2019.
  * [31] Francisco Vargas, Will Sussman Grathwohl, and Arnaud Doucet.  Denoising Diffusion Samplers.  In The Eleventh International Conference on Learning Representations, 2022.
  * [32] Francisco Vargas and Nikolas Nüsken.  Transport, Variational Inference and Diffusions: with Applications to Annealed Flows and Schrödinger Bridges.  arXiv preprint arXiv:2307.01050, 2023.
  * [33] Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence.  Solving Schrödinger bridges via maximum likelihood.  Entropy, 23(9):1134, 2021.
  * [34] Qinsheng Zhang and Yongxin Chen.  Path Integral Sampler: A Stochastic Control Approach For Sampling.  In International Conference on Learning Representations, 2021.
  * [35] Mo Zhou, Jiequn Han, and Jianfeng Lu.  Actor-critic method for high dimensional static Hamilton–Jacobi–Bellman partial differential equations based on neural networks.  SIAM Journal on Scientific Computing, 43(6):A4043–A4066, 2021.

##  Appendix A Various Perspectives on Schrödinger Bridge Problem

Most of these can be traced out in [17, 8], which has its roots in statistical
mechanics [26] (and in modern terms, closely related to large-deviation
results via Sanov’s theorem). For us however, the following perspectives will
be more fruitful.

(1) Mixture of pinned diffusions (weights given by $\pi^{*}$) can be seen by
disintegration of path measure:

$D_{KL}(P\|Q)=D_{KL}(P_{0T}\|Q_{0T})+\int
D_{KL}(P^{zx}\|Q^{zx})\,dP_{0T}(z,x)\,.$ (28)

Since the only constraints are on the two end-points (11) is therefore
equivalent to (29), by choosing $P^{zx}=Q^{zx}$. The optimal solution takes
the form $P^{*}=\pi^{*}Q^{zx}$, which means one can sample from
$(z,x)\sim\pi^{*}\in\Pi_{\nu,\mu}$, and sample from the bridges conditioned on
the end-points at $t=0,T$.

(2) It has the interpretation of entropy-regularized optimal transport [23]
when $f=0$: in terms of static formulation because of (28), if the reference
$Q$ is simply the Wiener process (accordingly $r(x|z)\propto
e^{-\frac{1}{2T}\|x-z\|^{2}}$),

$\displaystyle\pi^{*}(z,x)$
$\displaystyle=\arg\min_{\pi_{z}=\nu,\pi_{x}=\mu}\;D_{KL}(\pi(z,x)\|r(z,x))$
(29)
$\displaystyle=\arg\min_{\pi_{z}=\nu,\pi_{x}=\mu}\;\mathbb{E}_{z\sim\nu}[D_{KL}(\pi(x|z)\|r(x|z))]+D_{KL}(\nu(z)\|r(z))$
$\displaystyle=\arg\min_{\pi_{z}=\nu,\pi_{x}=\mu}\int\frac{1}{2}\|x-z\|^{2}\pi(z,x)\,dxdz+T\int\pi(z,x)\log\pi(z,x)\,dxdz$
$\displaystyle=\arg\min_{\pi_{z}=\nu,\pi_{x}=\mu}\int\frac{1}{2}\underbrace{\|x-z\|^{2}}_{\mathcal{C}(z,x)}\pi(z,x)\,dxdz+T\underbrace{\int\pi(z,x)\log\frac{\pi(z,x)}{\nu(z)\otimes\mu(x)}\,dxdz}_{D_{KL}(\pi(z,x)\|\nu(z)\otimes\mu(x))}\,,$

where the first term is nothing but the definition of the Wasserstein-2
distance (and optimal transport with quadratic cost in the sense of
Kantorovich). The second entropy term favors independent coupling
$\nu\otimes\mu$. This is a Lagrangian description of the transport. Different
choices of $r(x|z)$ transition will induce different transport costs. The
objective above can also be written as

$\arg\min_{\pi_{z}=\nu,\pi_{x}=\mu}\;D_{KL}\left(\pi(z,x)\,\|\,e^{-\mathcal{C}(z,x)/2T}\nu(z)\otimes\mu(x)\right)\,.$

In terms of dynamical formulation, via Girsanov’s theorem on path measure,
with the controlled dynamics as

$dX_{t}=[f_{t}(X_{t})+\sigma u_{t}(X_{t})]\,dt+\sigma dW_{t}\,,$

(11) can be reformulated as a _constrained_ problem over $u$:

$P^{*}=\arg\min_{P}\,\mathbb{E}_{X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}}\left[\frac{1}{2}\int_{0}^{T}\|u_{t}(X_{t})\|^{2}dt\,\Big{|}\,\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}_{T}=\mu\right]\,.$ (30)

Or a _regularized_ Benamou-Brenier fluid-dynamics [2] analogy of the optimal
transport

$\displaystyle\inf_{\rho,v}\;\int_{\mathbb{R}^{d}}\int_{0}^{T}\left[\frac{1}{2}\|v_{t}(X_{t})-\bar{v}_{t}(X_{t})\|^{2}+\frac{\sigma^{4}}{8}\left\|\nabla\log\frac{\rho_{t}(X_{t})}{\bar{\rho}_{t}(X_{t})}\right\|^{2}\right]\rho_{t}(X_{t})\,dtdx$
(31) $\displaystyle\text{s.t.}\;\frac{\partial\rho_{t}}{\partial
t}+\nabla\cdot(\rho_{t}v_{t})=0,\,\rho_{0}=\nu,\rho_{T}=\mu$ (32)

where above $\bar{v}_{t}=f_{t}-\frac{\sigma^{2}}{2}\nabla\log\bar{\rho}_{t}$
is the velocity field of the prior process, and we see the penalization
results in an additional relative Fisher information term. Since $v^{*}$ is a
velocity field in the continuity equation, it means a deterministic evolution
(i.e., ODE) as

$\dot{X}_{t}=v_{t}(X_{t}),\,X_{0}\sim\nu$

will have $X_{t}\sim\rho_{t}$, the optimal entropic interpolation flow, which
gives an Eulerian viewpoint.

(3) Optimal control views the problem as steering $\nu$ at $t=0$ to $\mu$ at
$t=T$ with minimal control effort. The value function (i.e., optimal cost-to-
go)

$V(x,t):=\min_{u}\,\mathbb{E}_{u}\left[\frac{1}{2}\int_{t}^{T}\|u_{s}(X_{s})\|^{2}ds\,\Big{|}\,X_{t}^{u}=x,X_{T}^{u}\sim\mu\right]$
(33)

with the expectation taken over the stochastic dynamics

$dX_{t}^{u}=[f_{t}(X_{t}^{u})+\sigma u_{t}(X_{t}^{u})]dt+\sigma
dW_{t},\;X_{0}^{u}\sim\nu$ (34)

should satisfy the Hamilton-Jacobi-Bellman equation via the dynamical
programming principle

$\frac{\partial V(x,t)}{\partial t}+f_{t}(x)^{\top}\nabla
V(x,t)+\frac{\sigma^{2}}{2}\Delta V(x,t)-\frac{\sigma^{2}}{2}\|\nabla
V(x,t)\|^{2}=0\,,$ (35)

where the optimal control $u_{t}^{*}(X_{t})=-\sigma\nabla V(X_{t},t)$ and
gives the unique solution $(\rho_{t}^{u})_{t\geq 0}$ solving

$\frac{\partial\rho_{t}^{u}}{\partial
t}=-\nabla\cdot(\rho_{t}^{u}(f_{t}-\sigma^{2}\nabla
V_{t}))+\frac{\sigma^{2}}{2}\Delta\rho_{t}^{u},\quad\rho^{u}_{0}\sim\nu,\,\rho^{u}_{T}\sim\mu\,.$
(36)

The two _coupled_ PDEs, Fokker-Planck (36) and HJB (35) are also the KKT
optimality condition of

$\displaystyle\inf_{\rho,u}\;\int_{\mathbb{R}^{d}}\int_{0}^{T}\frac{1}{2}\|u_{t}(X_{t})\|^{2}\rho_{t}(X_{t})\,dtdx$
(37) $\displaystyle\text{s.t.}\;\frac{\partial\rho_{t}}{\partial
t}+\nabla\cdot(\rho_{t}(f_{t}+\sigma
u_{t}))=\frac{\sigma^{2}}{2}\Delta\rho_{t},\,\rho_{0}=\nu,\rho_{T}=\mu$ (38)

where again the optimal $u^{*}=-\sigma\nabla V$ is of gradient type. Above the
Laplacian is responsible for the diffusion part, and (37)-(38) is related to
(31)-(32) via a change of variable. One might try to design schemes by forming
the Lagrangian for the above (37)-(38), and solve the resulting saddle-point
problem, but this deviates somewhat from our pathwise narrative. Indeed [1,
Theorem 3.13] proposed min-max based formulation that may be harder to train
in practice compared to our joint minimization approach.

######  Remark 7 (Langevin).

Compared to the Langevin SDE $dX_{t}=-\nabla f(X_{t})dt+\sqrt{2}dW_{t}$, which
only involves forward-evolving density characterization and reaches
equilibrium as $T\rightarrow\infty$, the controlled SDE (34) is time-
inhomogeneous and involves two PDEs (35)-(36). Langevin also has a backward
Kolmogorov evolution for the expectation of a function $g$: let
$V(x,t)=\mathbb{E}[g(X_{T})|X_{t}=x]$ we have $\partial_{t}V(x,t)-\nabla
f(x)^{\top}\nabla V(x,t)+\Delta V(x,t)=0$ with $V(x,T)=g(x)$, and is _de-
coupled_ from the Fokker-Planck equation
$\partial_{t}\rho_{t}-\nabla\cdot(\rho_{t}\nabla f)-\Delta\rho_{t}=0$ with
$\rho_{0}=\nu$ in this case.

######  Remark 8 (Bound on optimal objective).

Notice that in fact

$\min_{u}\;\mathbb{E}_{u}\left[\int_{0}^{T}\frac{1}{2}\|u_{t}(X_{t})\|^{2}dt-\log\frac{p_{\text{target}}}{Q_{T}}(X_{T})\right]\geq
0\,,$ (39)

which means we know a lower bound on the minimum of our losses (a) and (d) in
Proposition 1. This holds since for the controlled process $P$,

$D_{KL}(P\|Q)=\mathbb{E}_{u}\left[\int_{0}^{T}\frac{1}{2}\|u_{t}(X_{t})\|^{2}dt\right]=\mathbb{E}_{u}\left[\frac{1}{2}\int_{0}^{T}\|u_{t}(X_{t})\|^{2}dt-\log\frac{p_{\text{target}}}{Q_{T}}(X_{T})\right]+D_{KL}(p_{\text{target}}\|Q_{T})\,,$

which gives the claim by the data processing inequality
$D_{KL}(p_{\text{target}}\|Q_{T})\leq D_{KL}(P\|Q)$. A special instance is
when $\nu$ is $\delta_{0}$ a fixed Dirac delta and the reference process is
simply Brownian motion $f_{t}=0$. In this case, the optimal drift $u^{*}$ is
known to be the Föllmer drift [30, 34] and can be computed explicitly as
(picking $T=1$ for simplicity below)

$\displaystyle u_{t}^{*}(x)$
$\displaystyle=\arg\min_{u}\;\mathbb{E}_{X\sim\overrightarrow{\mathbb{P}}^{\nu,\sigma
u}}\Big{[}\int_{0}^{1}\frac{1}{2}\|u_{s}(X_{s})\|^{2}ds+\log\left(\frac{Q_{1}}{p_{\text{target}}}\right)(X_{1})\Big{]}$
(40)
$\displaystyle=\nabla\log\mathbb{E}_{X}\left[\frac{d\mu}{dQ_{1}}(X_{1})\,\Big{|}\,X_{t}=x\right]$
(41)
$\displaystyle=\nabla\log\mathbb{E}_{z\sim\mathcal{N}(0,\sigma^{2}I)}\left[\frac{d\mu}{d\mathcal{N}(0,\sigma^{2}I)}(x+\sqrt{1-t}z)\right]\,,$

where the expectation in (41) is taken w.r.t the reference measure – Wiener
process in this case; and it follows that $v_{t}^{*}(x)=x/t$ by Doob’s
$h$-transform. This is technically speaking a half-bridge where one can show
$P_{t}^{*}(X_{t})=Q(X_{t}|X_{T})p_{\text{target}}(X_{T})$ and equality in (39)
holds exactly. (40) gives an intuitive explanation of the optimally-controlled
process, with the first part corresponding to running cost and the second part
to terminal cost.

(4) The non-negative functions $\phi,\psi$ (which are closely related to the
dual potentials $f,g$ from (46) below) yield optimal forward/backward drifts:
It holds that the optimal curve admits the representation

$\log\rho_{t}=\log\phi_{t}+\log\psi_{t}\quad\text{for all }t$ (42)

and solving the boundary-coupled linear PDE system on the control

$\frac{\partial\phi_{t}}{\partial
t}=-\frac{\sigma^{2}}{2}\Delta\phi_{t}-\nabla\phi_{t}^{\top}f_{t},\quad\frac{\partial\psi_{t}}{\partial
t}=\frac{\sigma^{2}}{2}\Delta\psi_{t}-\nabla\cdot(\psi_{t}f_{t})\quad\text{for
$\phi_{0}\psi_{0}=p_{\text{prior}},\,\phi_{T}\psi_{T}=p_{\text{target}}$}$
(43)

gives two SDEs for the optimal curve in (11):

$\displaystyle dX_{t}$
$\displaystyle=[f_{t}(X_{t})+\sigma^{2}\nabla\log\phi_{t}(X_{t})]dt+\sigma\overrightarrow{dW_{t}},\;X_{0}\sim\nu$
(44) $\displaystyle dX_{t}$
$\displaystyle=[f_{t}(X_{t})-\sigma^{2}\nabla\log\psi_{t}(X_{t})]dt+\sigma\overleftarrow{dW_{t}},\;X_{T}\sim\mu\,.$
(45)

Note that equations (44)-(45) are time reversals of each other and obey
Nelson’s identity thanks to (42). The transformation (42)-(43) that involves
$(\rho_{t}^{*},u_{t}^{*})=(\rho_{t}^{*},-\sigma\nabla
V_{t})\mapsto(\phi_{t},\psi_{t})$ is a typical $\log\leftrightarrow\exp$ Hopf-
Cole change-of-variable from (35)-(36). These PDE optimality results can be
found in [4].

######  Remark 9 (Boundary condition).

In rare cases, such as those from Föllmer drift (41), the PDE system (43) (or
(35)-(36)) may decouple and the drift can be expressed analytically. But the
boundary conditions of (35)-(36) is atypical for control problems, which
generally would have (36) specified at initial time and runs forward, with
(35) specified at the terminal and runs backward. The PDE dynamics of (43)
also take similar form as those from Remark 7 but with slightly unusually
coupled boundary conditions.

(5) Factorization of the optimal coupling: in fact it is always the case that

$\frac{d\pi^{*}}{dr}(X_{0},X_{T})=e^{f(X_{0})}e^{g(X_{T})}\quad
r\text{-}a.s.\,.$ (46)

Moreover under mild conditions if there exists $\pi,f,g$ for which such
decomposition holds and $\pi_{0}=\nu,\pi_{T}=\mu$, $\pi$ must be optimal –
such condition (46) is _necessary and sufficient_ for characterizing the
solution to the SB problem. (46) together with (28) give that (which can also
be thought of as re-weighting on the path space)

$\frac{dP^{*}}{dQ}(X_{0:T})=e^{f(X_{0})}e^{g(X_{T})}\quad Q\text{-}a.s.\,,$
(47)

where the $f$ and $g$ obey the Schrödinger system

$e^{f(X_{0})}\int r(X_{0T})e^{g(X_{T})}dX_{T}=p_{\text{prior}}(X_{0})\,,\quad
e^{g(X_{T})}\int e^{f(X_{0})}r(X_{0T})dX_{0}=p_{\text{target}}(X_{T})\,.$ (48)

The $\phi,\psi$ in (44)-(45) can also be expressed as a conditional
expectation:

$\phi_{t}(x)\stackrel{{\scriptstyle(\\#)}}{{=}}\int
e^{g(X_{T})}r(X_{T}|X_{t}=x)\,dX_{T}\stackrel{{\scriptstyle(*)}}{{=}}\int_{\mathbb{R}^{d}}\phi_{T}(X_{T})r(X_{T}|X_{t}=x)\,dX_{T}\,,$
$\psi_{t}(x)\stackrel{{\scriptstyle(\\#)}}{{=}}\int
e^{f(X_{0})}r(X_{t}=x,X_{0})\,dX_{0}\stackrel{{\scriptstyle(*)}}{{=}}\int_{\mathbb{R}^{d}}\psi_{0}(X_{0})r(X_{t}=x|X_{0})\,dX_{0}$

for $t\in[0,T]$, where the consistency of

$e^{f(X_{0})}r(X_{0})=\psi_{0}(X_{0}),\,e^{g(X_{T})}=\phi_{T}(X_{T})\;\Leftrightarrow\;\pi^{*}(X_{0},X_{T})=\psi_{0}(X_{0})r(X_{T}|X_{0})\phi_{T}(X_{T})$

can be verified by (48), (42) and $(\\#)$. This representation of
$\phi_{t},\psi_{t}$ as the conditional expectation $(*)$ can be seen with the
Feynman-Kac formula on (43) as well. But in general, the transition kernel of
the un-controlled process $r(\cdot)$, or more importantly the terminals
$\phi_{T},\psi_{0}$ are not available analytically for implementation of $(*)$
as such. But for example, in the simple case of $f=0$, $(*)$ simply imply that
$\psi,\phi$ describe the evolution of density corresponding to the following
SDE: $dY_{t}=\overrightarrow{dW_{t}},Y_{0}\sim\psi_{0}$ and
$dY_{t}=\overleftarrow{dW_{t}},Y_{T}\sim\phi_{T}$.

######  Remark 10.

Such factorization property (46) is maintained by Sinkhorn algorithm (that
updates $f,g$ as iteration proceeds), which is well-known to converge to the
optimal coupling obeying (48) eventually.

(6) Equation (44) means the optimal $v_{t}^{*}$ in the continuity equation
(32) is
$f_{t}+\sigma^{2}\nabla\log\phi_{t}-\frac{\sigma^{2}}{2}\nabla\log\rho_{t}=f_{t}+\frac{\sigma^{2}}{2}\nabla\log\frac{\phi_{t}}{\psi_{t}}$,
which gives the ODE probability flow for this dynamics

$dX_{t}=f_{t}(X_{t})+\frac{\sigma^{2}}{2}(\nabla\log\phi_{t}(X_{t})-\nabla\log\psi_{t}(X_{t}))\,dt\,,\quad
X_{0}\sim\nu\,,$ (49)

in the sense that the time marginals
$\rho_{t}^{\eqref{eqn:ode}}=\rho_{t}^{\eqref{eqn:forward_sb}}=\rho_{t}^{\eqref{eqn:backward_sb}}=P_{t}^{\eqref{eqn:bridge_problem}}$
all agree. Moreover, using the instantaneous change of variables formula [5],
we also have

$\log\rho_{T}(X_{T})=\log\rho_{0}(X_{0})-\int_{0}^{T}\nabla\cdot
f_{t}(X_{t})\,dt-\frac{\sigma^{2}}{2}\int_{0}^{T}\nabla\cdot(\nabla\log\phi_{t}(X_{t})-\nabla\log\psi_{t}(X_{t}))\,dt\,,$
(50)

which will be useful in Proposition 2 for estimating the normalizing constant.
Note that in (50) both the density and the point at which we are evaluating is
changing.

######  Remark 11 (Minimum control energy).

The stochastic control formulation makes it clear that the trajectory we are
trying to recover is a meaningful one in the sense of minimal effort. If one
were to switch the order of $P$ and $Q$ in (11), the optimal control problem
becomes (e.g., for $f=0$)

$\displaystyle\inf_{u}\;\mathbb{E}\left[\int_{0}^{T}\frac{1}{2}\|u_{t}(X_{t})\|^{2}dt\right]$
$\displaystyle\text{s.t.}\;dX_{t}=\sigma u_{t}(X_{t})dt+\sigma
dW_{t},\,X_{0}\sim\nu,X_{T}\sim\mu$

for the expectation taken over the reference process, instead of the
controlled state density $\rho_{t}^{u}$, which is not very intuitive. The
slightly non-conventional aspect of this control problem is the fixed terminal
constraint.

To briefly summarize, all these different viewpoints explore the deep
connections between PDEs (controls) and SDEs (diffusions) in one way or
another.

##  Appendix B Additional Related Work: Natural attempts for solving SB

In this section we discuss several natural attempts with the intention of
adapting the SB formalism for sampling from un-normalized densities.

  * •

Iterative proportional fitting (IPF) / Sinkhorn performs iterative projection
as

$P^{(1)}=\arg\min_{Q\in\mathcal{P}(\nu,\cdot)}\;D_{KL}(Q\|P^{(0)})=\frac{P^{(0)}\nu}{P^{(0)}_{0}},\quad
P^{(0)}=\arg\min_{Q\in\mathcal{P}(\cdot,\mu)}\;D_{KL}(Q\|P^{(1)})=\frac{P^{(1)}\mu}{P^{(1)}_{T}}$
(51)

i.e., only the end point differ and one solves half-bridges using drifts
learned from the last trajectory transition rollout, initializing from either
samples from $\mu(x)$ or $\nu(z)$. But since we don’t have samples from $\mu$,
neither the score function nor the succeeding IPF updates/refinements can be
implemented. In fact, the first iteration of the IPF proposal in [9, 33]
precisely corresponds to the score-based diffusion proposal in [28]. However,
there is a connection between IPF and a path space EM implementation
formulated in terms of drifts that rely on change in the KL direction: recall
under mild assumptions

$\arg\min_{\phi}\;D_{KL}(\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}\|\overleftarrow{\mathbb{P}}^{\mu,f+\sigma^{2}\nabla\psi})=\arg\min_{\phi}\;D_{KL}(\overleftarrow{\mathbb{P}}^{\mu,f+\sigma^{2}\nabla\psi}\|\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi})\,.$

Using this fact, it is shown in [32] that coordinate descent on the objective
$\min_{\phi,\psi}\,D_{KL}(\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}\|\overleftarrow{\mathbb{P}}^{\mu,f+\sigma^{2}\nabla\psi})$,
when initializing at $\phi=0$ (i.e., the prior), is a valid strategy for
solving (11).

######  Lemma 2 (Optimization of drifts of EM [32]).

The alternating scheme initialized with $\phi_{0}=0$ converge to $P^{*}$
(i.e., the coupling at time $t=0,T$ is optimal):

$\displaystyle\psi_{n}$
$\displaystyle=\arg\min_{\psi}\;D_{KL}(\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi_{n-1}}\|\overleftarrow{\mathbb{P}}^{\mu,f+\sigma^{2}\nabla\psi})$
(52) $\displaystyle\phi_{n}$
$\displaystyle=\arg\min_{\phi}\;D_{KL}(\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}\|\overleftarrow{\mathbb{P}}^{\mu,f+\sigma^{2}\nabla\psi_{n}})\,.$
(53)

Moreover, both updates are implementable assuming samples from $\nu$ is
available, which resolves non-uniqueness of the two parameter loss
$\min_{\phi,\psi}\,D_{KL}(\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}\|\overleftarrow{\mathbb{P}}^{\mu,f+\sigma^{2}\nabla\psi})$
in an algorithmic manner (fixing one direction of the drift at a time).

######  Remark 12 (Benefit of joint training).

Lemma 2 above also shows that the prior only enters in the first step,
therefore as it proceeds, the prior influence tends to be ignored as error
accumulates – this aspect is different from our loss proposals in Proposition
1. Another advantage of our joint minimization procedure compared to an
iterative sequential IPF method is that while the equivalence of EM and IPF
will no longer hold for neural network models not expressive enough (therefore
the convergence does not immediately carry over), one could still estimate the
normalizing constant / sample from the target with a sub-optimal control
learned from the losses proposed in Proposition 1 with e.g., important weights
(c.f. (91)).

  * •

Some alternatives to solve SB do not require analytical expression for $\mu$:
diffusion mixture matching [22, 27] tilts the product measure $\nu\otimes\mu$
towards optimality gradually by learning a slightly different term than the
score, and alternate between such Markovian and reciprocal projection. Data-
driven bridge [21] aims at setting up a fixed point recursion on the SB system
(48) for finding the optimal $\phi^{*},\psi^{*}$, but both rely on (1) the
availability of samples from $\mu$ (as well as $\nu$) to estimate various
quantities for implementation; (2) _tractable_ bridge distribution/Markov
kernel for the reference $Q(\cdot|X_{0},X_{T})$. In fact a different
initialization from the reference $Q$ using method in [27] will reduce to IPF.

  * •

Naively one may hope to reparametrize (44)-(45) as (with the time-reversal
condition baked in)

$\min_{u}\,D_{KL}(\overrightarrow{\mathbb{P}}^{\nu,\frac{\sigma^{2}}{2}\nabla\log\rho_{t}+\nabla
u_{t}}\|\overleftarrow{\mathbb{P}}^{\mu,\nabla
u_{t}-\frac{\sigma^{2}}{2}\nabla\log\rho_{t}})$

to learn the drift $\nabla u$ (this is somewhat akin to annealed flow, and
expected to have unique solution). This way, when the two path measures are
optimized to agree, the resulting forward & backward drift necessarily obey
Nelson’s identity for the desired curve $\rho_{t}$. However, this only applies
if we have access to the score $\nabla\log\rho_{t}$, for
$(\rho_{t})_{t\in[T]}$ the SB curve of density we wish to follow interpolating
from $\nu$ to $\mu$. Learning both $\rho$ and $\nabla u$ will become similar
to learning $\phi$ and $\psi$, as the approaches we study in Section 3.1.

  * •

The work of [4] investigated the case where the reference process has a
gradient drift (i.e., $f=-\nabla U$) and reduce the optimal control task to
solving a high-dimensional PDE subject to initial-value constraint (c.f. Eqn
(33) and (47) therein). However, solving PDEs is largely regarded to be
computationally more demanding than simulating SDEs.

##  Appendix C Missing Calculations

We give the likelihood ratio calculation Lemma 3 below, which is closely
related to those from [25, 32].

######  Lemma 3 (Forward / Backward path-space likelihood ratio).

Written solely in terms of drifts, the KL divergence
$D_{KL}(\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}\|\overleftarrow{\mathbb{P}}^{\mu,f+\sigma v})$ becomes

$\mathcal{L}_{KL}(u,v)=\mathbb{E}_{X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}}\left[\int_{0}^{T}\frac{1}{2}\|u_{t}(X_{t})-v_{t}(X_{t})\|^{2}-\nabla\cdot(f_{t}+\sigma
v_{t})(X_{t})dt+\log\frac{d\nu(X_{0})}{d\mu(X_{T})}\right]+\log Z\,.$ (54)

And the log-variance divergence over the same two path measures can be
evaluated to be

$\text{Var}_{X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}}\left[\int_{0}^{T}\frac{1}{2}\|u_{t}(X_{t})-v_{t}(X_{t})\|^{2}-\nabla\cdot(f_{t}+\sigma
v_{t})(X_{t})dt+\log\frac{d\nu(X_{0})}{d\mu(X_{T})}+\int_{0}^{T}(u_{t}-v_{t})(X_{t})\,\overrightarrow{dW_{t}}\right]\,.$

The reverse one $D_{KL}(\overleftarrow{\mathbb{P}}^{\mu,f+\sigma
v}\|\overrightarrow{\mathbb{P}}^{\nu,f+\sigma u})$ becomes up to constant

$\mathbb{E}_{X\sim\overleftarrow{\mathbb{P}}^{\mu,f+\sigma
v}}\left[\int_{0}^{T}\frac{1}{2}\|v_{t}(X_{t})-u_{t}(X_{t})\|^{2}+\nabla\cdot(f_{t}+\sigma
u_{t})(X_{t})dt+\log\frac{d\mu(X_{T})}{d\nu(X_{0})}\right]\,.$

###### Proof of Lemma 3.

We give a few equivalent expressions for the Radon-Nikodym derivative.

Argument due to [32]: Starting with Proposition 2.2 of [32], in the general
case of reference

$dX_{t}=\sigma
r^{+}_{t}(X_{t})dt+\sigma\overrightarrow{dW_{t}},\;X_{0}\sim\Gamma_{0}\quad\quad
dX_{t}=\sigma
r^{-}_{t}(X_{t})dt+\sigma\overleftarrow{dW_{t}},\;X_{T}\sim\Gamma_{T}$ (55)

where $\overrightarrow{\mathbb{P}}^{\Gamma_{0},\sigma
r^{+}}=\overleftarrow{\mathbb{P}}^{\Gamma_{T},\sigma r^{-}}$, it is shown
$\overrightarrow{\mathbb{P}}^{\nu,f+\sigma u}$ almost surely,

$\displaystyle\log\left(\frac{d\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}}{d\overleftarrow{\mathbb{P}}^{\mu,f+\sigma
v}}\right)(X)=\log(\frac{d\nu}{d\Gamma_{0}})(X_{0})-\log(\frac{d\mu}{d\Gamma_{T}})(X_{T})+\log
Z$ $\displaystyle+\frac{1}{\sigma^{2}}\int_{0}^{T}(f_{t}+\sigma u_{t}-\sigma
r_{t}^{+})(X_{t})\left(\overrightarrow{dX_{t}}-\frac{1}{2}(f_{t}+\sigma
u_{t}+\sigma r_{t}^{+})(X_{t})dt\right)$ (56)
$\displaystyle-\frac{1}{\sigma^{2}}\int_{0}^{T}(f_{t}+\sigma v_{t}-\sigma
r_{t}^{-})(X_{t})\left(\overleftarrow{dX_{t}}-\frac{1}{2}(f_{t}+\sigma
v_{t}+\sigma r_{t}^{-})(X_{t})dt\right)\,.$ (57)

One can use (6) to convert the backward integral to a forward one with an
additional divergence term

$\displaystyle D_{KL}(\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}\|\overleftarrow{\mathbb{P}}^{\mu,f+\sigma
v})=\mathbb{E}_{X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}}\left[\log\left(\frac{d\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}}{d\overleftarrow{\mathbb{P}}^{\mu,f+\sigma v}}\right)(X)\right]$
$\displaystyle=\mathbb{E}_{X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}}\bigg{[}\log(\frac{d\nu}{d\Gamma_{0}})(X_{0})-\log(\frac{d\mu}{d\Gamma_{T}})(X_{T})+\log
Z$ $\displaystyle+\frac{1}{2\sigma^{2}}\int_{0}^{T}(f_{t}+\sigma u_{t}-\sigma
r_{t}^{+})(X_{t})^{\top}(f_{t}+\sigma u_{t}-\sigma r_{t}^{+})(X_{t})dt$
$\displaystyle-\frac{1}{\sigma^{2}}\int_{0}^{T}(f_{t}+\sigma v_{t}-\sigma
r_{t}^{-})(X_{t})^{\top}(\frac{1}{2}f_{t}+\sigma
u_{t}-\frac{\sigma}{2}v_{t}-\frac{\sigma}{2}r_{t}^{-})(X_{t})dt-\int_{0}^{T}\nabla\cdot(f_{t}+\sigma
v_{t}-\sigma r_{t}^{-})(X_{t})dt\bigg{]}$ (58)
$\displaystyle+\mathbb{E}_{X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}}\left[\int_{0}^{T}(u_{t}-r_{t}^{+})(X_{t})-(v_{t}-r_{t}^{-})(X_{t})\,\overrightarrow{dW_{t}}\right]$

and the last term vanishes because of (5). Therefore we see that there are 2
boundary terms, and 3 extra terms corresponding to the forward/backward
process. By picking $\gamma^{+},\gamma^{-}=0$, and Lebesgue base measure for
$\Gamma_{0},\Gamma_{T}$, we get

$\displaystyle D_{KL}(\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}\|\overleftarrow{\mathbb{P}}^{\mu,f+\sigma v})$
$\displaystyle=\mathbb{E}_{X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}}\left[\int_{0}^{T}\frac{1}{2}\|u_{t}(X_{t})-v_{t}(X_{t})\|^{2}-\nabla\cdot(f_{t}+\sigma
v_{t})(X_{t})dt\right]$ (59)
$\displaystyle+\mathbb{E}_{X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}}\left[\log\frac{d\nu(X_{0})}{d\mu(X_{T})}+\int_{0}^{T}(u_{t}-v_{t})(X_{t})\,\overrightarrow{dW_{t}}\right]+\log
Z$ (60)

which agrees with Proposition 2.3 in [25] up to a conventional sign in $v$,
and is also the same as the ELBO loss in [7, Theorem 4].

Argument using Remark 1: Another way to show (54) is to start with Nelson’s
identity and apply Girsanov’s theorem with (2) and time reversal of (3). Using
the chain rule for the KL between two forward processes we get

$\displaystyle D_{KL}$
$\displaystyle(\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}\|\overleftarrow{\mathbb{P}}^{\mu,f+\sigma
v})=\mathbb{E}_{\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}}\left[\log\left(\frac{d\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}}{d\overleftarrow{\mathbb{P}}^{\mu,f+\sigma v}}\right)(X)\right]$
$\displaystyle=\mathbb{E}_{\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}}\left[\log\left(\frac{d\nu}{d\overleftarrow{\mathbb{P}}_{0}^{\mu,f+\sigma
v}}\right)(X_{0})+\frac{1}{2}\int_{0}^{T}\|u_{t}(X_{t})-(v_{t}+\sigma\nabla\log\rho^{\mu,f+\sigma
v}_{t})(X_{t})\|^{2}dt\right]\,.$ (61)

To see this is equivalent to (59)-(60), we use Fokker-Planck on the forward
process with drift $v_{t}+\sigma\nabla\log\rho^{\mu,f+\sigma v}_{t}$ to reach

$\displaystyle\mathbb{E}_{\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}}\left[\log\left(\frac{d\mu}{d\overleftarrow{\mathbb{P}}_{0}^{\mu,f+\sigma
v}}\right)(X)-\log
Z\right]=\mathbb{E}_{\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}}\left[\log\left(\frac{d\overleftarrow{\mathbb{P}}_{T}^{\mu,f+\sigma
v}}{d\overleftarrow{\mathbb{P}}_{0}^{\mu,f+\sigma v}}\right)(X)\right]$
$\displaystyle=\mathbb{E}_{\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}}\left[\int_{0}^{T}-\nabla\cdot(f_{t}+\sigma
v_{t})(X_{t})+\sigma(u_{t}-v_{t})(X_{t})^{\top}\nabla\log\rho^{\mu,f+\sigma
v}_{t}(X_{t})-\frac{\sigma^{2}}{2}\|\nabla\log\rho^{\mu,f+\sigma
v}_{t}(X_{t})\|^{2}dt\right]$ (62)

as by Itô’s lemma for the process (2) with drift $f+\sigma u$, (denote
$\overleftarrow{\mathbb{P}}_{t}^{\mu,f+\sigma v}$ as $\rho_{t}$)

$\displaystyle\int_{0}^{T}\partial_{t}\log\overleftarrow{\mathbb{P}}_{t}^{\mu,f+\sigma
v}(X_{t}^{f+\sigma u})\,dt$
$\displaystyle=\int_{0}^{T}\frac{1}{\rho_{t}}\left(-\nabla\cdot(\rho_{t}(f+\sigma
v))+\frac{\sigma^{2}}{2}\Delta\rho_{t}+\nabla\rho_{t}^{\top}(f+\sigma
u)\right)dt+\sigma\int_{0}^{T}\frac{\nabla\rho_{t}^{\top}}{\rho_{t}}dW_{t}+\int_{0}^{T}\frac{\sigma^{2}}{2}\Delta\log\rho_{t}\,dt$
$\displaystyle=\int_{0}^{T}-\frac{\nabla\rho_{t}^{\top}}{\rho_{t}}(f_{t}+\sigma
v_{t}+\frac{\sigma^{2}}{2}\nabla\log\rho_{t}-f_{t}-\sigma
u_{t})-\nabla\cdot(f_{t}+\sigma
v_{t}+\frac{\sigma^{2}}{2}\nabla\log\rho_{t})+\frac{\sigma^{2}}{2}\Delta\log\rho_{t}\,dt+\int_{0}^{T}\sigma\frac{\nabla\rho_{t}^{\top}}{\rho_{t}}\,dW_{t}\,,$

which upon simple re-arranging and taking expectation over $X^{f+\sigma
u}_{t}\sim\overrightarrow{\mathbb{P}}_{t}^{\nu,f+\sigma u}$ give (62). Now
adding up the previous two displays (61) and (62) using

$\mathbb{E}_{\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}}\left[\log\left(\frac{d\nu}{d\overleftarrow{\mathbb{P}}_{0}^{\mu,f+\sigma
v}}\right)\right]=\mathbb{E}_{\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}}\left[\log\left(\frac{d\nu}{d\mu}\right)+\log\left(\frac{d\mu}{d\overleftarrow{\mathbb{P}}_{0}^{\mu,f+\sigma
v}}\right)\right]$

finishes the proof of the expression (54). The log-variance claim follows
easily from this.

Reverse KL using Remark 1: Symmetrically, we can use Girsanov’s theorem to
compute the reverse KL divergence between two backward processes as:

$\displaystyle D_{KL}$ $\displaystyle(\overleftarrow{\mathbb{P}}^{\mu,f+\sigma
v}\|\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u})=\mathbb{E}_{\overleftarrow{\mathbb{P}}^{\mu,f+\sigma
v}}\left[\log\left(\frac{d\overleftarrow{\mathbb{P}}^{\mu,f+\sigma
v}}{d\overrightarrow{\mathbb{P}}^{\nu,f+\sigma u}}\right)(X)\right]$
$\displaystyle=\mathbb{E}_{\overleftarrow{\mathbb{P}}^{\mu,f+\sigma
v}}\left[\log\left(\frac{d\mu}{d\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}_{T}}\right)(X_{T})+\frac{1}{2}\int_{0}^{T}\|v_{t}(X_{t})-(u_{t}-\sigma\nabla\log\rho_{t}^{\nu,f+\sigma
u})(X_{t})\|^{2}dt\right]\,.$

We decompose

$\mathbb{E}_{\overleftarrow{\mathbb{P}}^{\mu,f+\sigma
v}}\left[\log\left(\frac{d\mu}{d\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}_{T}}\right)\right]=\mathbb{E}_{\overleftarrow{\mathbb{P}}^{\mu,f+\sigma
v}}\left[\log\left(\frac{d\mu}{d\nu}\right)+\log\left(\frac{d\nu}{d\overrightarrow{\mathbb{P}}^{\nu,f+\sigma
u}_{T}}\right)\right]\,.$

Now note the second term is nothing but
$-\int_{0}^{T}\partial_{t}\log\overrightarrow{\mathbb{P}_{t}}^{\nu,f+\sigma
u}(X_{t})\,dt$ for $X_{t}\sim\overleftarrow{\mathbb{P}}^{\mu,f+\sigma v}$, and
we compute using Fokker-Planck and backward Itô’s lemma to reach (denote
$\overrightarrow{\mathbb{P}_{t}}^{\nu,f+\sigma u}$ as $\rho_{t}$)

$\displaystyle\int_{0}^{T}\partial_{t}\log\overrightarrow{\mathbb{P}_{t}}^{\nu,f+\sigma
u}(X_{t}^{f+\sigma v})\,dt$
$\displaystyle=\int_{0}^{T}\frac{1}{\rho_{t}}\left(-\nabla\cdot(\rho_{t}(f+\sigma
u))+\frac{\sigma^{2}}{2}\Delta\rho_{t}+\nabla\rho_{t}^{\top}(f+\sigma
v)\right)dt+\sigma\int_{0}^{T}\frac{\nabla\rho_{t}^{\top}}{\rho_{t}}\overleftarrow{dW_{t}}-\int_{0}^{T}\frac{\sigma^{2}}{2}\Delta\log\rho_{t}\,dt$
$\displaystyle=\int_{0}^{T}-\frac{\nabla\rho_{t}^{\top}}{\rho_{t}}(f_{t}+\sigma
u_{t}-\frac{\sigma^{2}}{2}\nabla\log\rho_{t}-f_{t}-\sigma
v_{t})-\nabla\cdot(f_{t}+\sigma
u_{t}-\frac{\sigma^{2}}{2}\nabla\log\rho_{t})-\frac{\sigma^{2}}{2}\Delta\log\rho_{t}\,dt+\int_{0}^{T}\sigma\frac{\nabla\rho_{t}^{\top}}{\rho_{t}}\,\overleftarrow{dW_{t}}$
$\displaystyle=\int_{0}^{T}\sigma(v_{t}-u_{t})\nabla\log\rho_{t}+\frac{\sigma^{2}}{2}\|\nabla\log\rho_{t}\|^{2}-\nabla\cdot(f_{t}+\sigma
u_{t})\,dt+\int_{0}^{T}\sigma\frac{\nabla\rho_{t}^{\top}}{\rho_{t}}\,\overleftarrow{dW_{t}}\,,$

which by using (5) and putting everything together allow us to finish. ∎

Immediately follows is our main result: Proposition 1, along with auxiliary
Lemma 4 from Section 3.1.

######  Remark 13 (Explicit expression).

In Proposition 1, up to constants,
$D_{KL}(\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}\|\overleftarrow{\mathbb{P}}^{\mu,f-\sigma^{2}\nabla\psi})$
can be written as for
$X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}$,

$\mathbb{E}\Big{[}\int_{0}^{T}\frac{\sigma^{2}}{2}\|\nabla\phi_{t}(X_{t})+\nabla\psi_{t}(X_{t})\|^{2}-\nabla\cdot(f_{t}-\sigma^{2}\nabla\psi_{t})(X_{t})dt+\log\frac{\nu(X_{0})}{\mu(X_{T})}+\sigma\int_{0}^{T}(\nabla\phi_{t}+\nabla\psi_{t})(X_{t})\,\overrightarrow{dW_{t}}\Big{]}\,,$
(63)

and the log-variance divergence between the same two path measures has
$\text{Var}[\cdot]$ in place of $\mathbb{E}[\cdot]$. Moreover
$D_{KL}(\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi_{t}}\|\overleftarrow{\mathbb{P}}^{\mu,f+\sigma^{2}\nabla\phi_{t}-\sigma^{2}\nabla\log\rho_{t}})$
evaluates to

$\mathbb{E}_{\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}\Big{[}\log\frac{\nu(X_{0})}{\mu(X_{T})}+\int_{0}^{T}\frac{\sigma^{2}}{2}\|\nabla\log\rho_{t}(X_{t})\|^{2}+\sigma^{2}\nabla\cdot(\nabla\log\rho_{t}-\nabla\phi_{t})(X_{t})-\nabla\cdot
f_{t}(X_{t})\,dt\Big{]}\,.$ (64)

Both are straightforward consequences of Lemma 3.

###### Proof of Proposition 1.

The first term (63)-(64) simply involves writing $D_{KL}$ in terms of the
(current) controls using Lemma 3, up to constants. One can regularize to
remove the constraints in (30) while still ensuring the right marginals via

$\mathcal{L}(\nabla\phi,\nabla\psi):=D_{KL}(\overrightarrow{\mathbb{P}}^{\nu,\sigma^{2}\nabla\phi}\|\overleftarrow{\mathbb{P}}^{\mu,-\sigma^{2}\nabla\psi})+\lambda
R(\nabla\psi)\quad\text{or}\quad
D_{KL}(\overrightarrow{\mathbb{P}}^{\nu,\sigma^{2}\nabla\phi}\|\overleftarrow{\mathbb{P}}^{\mu,-\sigma^{2}\nabla\psi})+\lambda^{\prime}R^{\prime}(\nabla\phi)$

where $R(\cdot)$ can either utilize the variational formulation as in (a) or
the optimality condition on the control as in (b). In both cases, the first
term is a $\rho$ constraint (time reversal consistency with correct terminal
marginals), and the second one a control constraint (enforce optimality).

For (b), we give one direction of the argument for $\nabla\psi$ first. Let
$\overrightarrow{\mathbb{P}}^{\nu,f}$ denote the path measure associated with
$dX_{t}=f_{t}(X_{t})dt+\sigma dW_{t},X_{0}\sim\nu$, then
$\overrightarrow{\mathbb{P}}^{\nu,f}$ almost surely, the Radon-Nikodym
derivative between $\overrightarrow{\mathbb{P}}^{\nu,f}$ and the controlled
backward process is

$\log\left(\frac{d\overrightarrow{\mathbb{P}}^{\nu,f}}{d\overleftarrow{\mathbb{P}}^{\mu,f-\sigma^{2}\nabla\psi_{t}}}\right)(X)=\log\frac{d\nu(X_{0})}{d\mu(X_{T})}+\int_{0}^{T}\frac{\sigma^{2}}{2}\|\nabla\psi_{t}\|^{2}-\nabla\cdot(f-\sigma^{2}\nabla\psi_{t})dt+\sigma\int_{0}^{T}\nabla\psi_{t}^{\top}\,dW_{t}\,.$

In the case when the variance (taken along the prior
$X\sim\overrightarrow{\mathbb{P}}^{\nu,f}$)

$\text{Var}\left(\psi_{T}(X_{T})-\psi_{0}(X_{0})-\int_{0}^{T}\frac{\sigma^{2}}{2}\|\nabla\psi_{t}\|^{2}+\nabla\cdot
f-\sigma^{2}\Delta\psi_{t}(X_{t})dt-\int_{0}^{T}\sigma\nabla\psi_{t}^{\top}\,dW_{t}\right)=0\,,$

it implies that the random quantity is almost surely a constant independent of
the realization, and

$\log\left(\frac{d\overrightarrow{\mathbb{P}}^{\nu,f}}{d\overleftarrow{\mathbb{P}}^{\mu,f-\sigma^{2}\nabla\psi_{t}}}\right)(X)=\log\frac{d\nu(X_{0})}{d\mu(X_{T})}+\psi_{T}(X_{T})-\psi_{0}(X_{0})+\log
Z\,,$

from which using the factorization characterization (47), and the terminal
constraint
$\overleftarrow{\mathbb{P}}^{\mu,f-\sigma^{2}\nabla\psi_{t}}_{0}=\nu,\overleftarrow{\mathbb{P}}^{\mu,f-\sigma^{2}\nabla\psi_{t}}_{T}=\mu$
imposed by the first KL condition, concludes that
$\overleftarrow{\mathbb{P}}^{\mu,f-\sigma^{2}\nabla\psi_{t}}$ must be the
unique solution to the SB problem. The other direction on $\nabla\phi$ is
largely similar, and one can show using Girsanov’s theorem and the variance
condition that along the prior $X\sim\overrightarrow{\mathbb{P}}^{\nu,f}$,

$\displaystyle\log\left(\frac{d\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi_{t}}}{d\overrightarrow{\mathbb{P}}^{\nu,f}}\right)(X)$
$\displaystyle=\int_{0}^{T}-\frac{\sigma^{2}}{2}\|\nabla\phi_{t}\|^{2}dt+\int_{0}^{T}\sigma\nabla\phi_{t}^{\top}\,dW_{t}=\phi_{T}(X_{T})-\phi_{0}(X_{0})\,,$

to conclude in the same way. Note that evaluating the variance regularizer
only requires simulating from the reference process (12).

In (c), the first two parts enforce

$\phi_{T}(X_{T})+\psi_{T}(X_{T})=\log
p_{\text{target}}(X_{T})\quad\text{and}\quad\phi_{0}(X_{0})+\psi_{0}(X_{0})=\log
p_{\text{prior}}(X_{0})=\log\nu(X_{0})\,.$

Now for any $\phi_{t}$, the likelihood ratio along
$dX_{t}=(f+\sigma^{2}\nabla\phi_{t})dt+\sigma dW_{t},X_{0}\sim\nu$ is

$\log\left(\frac{d\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi_{t}}}{d\overrightarrow{\mathbb{P}}^{\nu,f}}\right)(X)=\int_{0}^{T}\frac{\sigma^{2}}{2}\|\nabla\phi_{t}\|^{2}dt+\int_{0}^{T}\sigma\nabla\phi_{t}^{\top}\,dW_{t}\,,$
(65)

which according to (47), has to be equal to

$\psi_{0}(X_{0})+\phi_{T}(X_{T})-\log\nu(X_{0})=-\phi_{0}(X_{0})+\phi_{T}(X_{T})\,\;\;\mathbb{P}^{\nu,f}\,a.s.\Rightarrow\mathbb{P}^{\nu,f+\sigma^{2}\nabla\phi}\,a.s.$

to be optimal, justifying

$\displaystyle\text{Var}_{X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}\left(\phi_{T}(X_{T})-\phi_{0}(X_{0})-\frac{\sigma^{2}}{2}\int_{0}^{T}\|\nabla\phi_{t}\|^{2}(X_{t})\,dt-\sigma\int_{0}^{T}\nabla\phi_{t}(X_{t})^{\top}\,dW_{t}\right)\,.$

In a similar spirit, for any $\psi_{t}$, along the same process
$X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi_{t}}$,

$\displaystyle\log\left(\frac{d\overrightarrow{\mathbb{P}}^{\nu,f}}{d\overleftarrow{\mathbb{P}}^{\mu,f-\sigma^{2}\nabla\psi_{t}}}\right)(X)$
$\displaystyle=\int_{0}^{T}\sigma^{2}\nabla\psi_{t}^{\top}\nabla\phi_{t}+\frac{\sigma^{2}}{2}\|\nabla\psi_{t}\|^{2}-\nabla\cdot(f_{t}-\sigma^{2}\nabla\psi_{t})dt+\int_{0}^{T}\sigma\nabla\psi_{t}^{\top}\,dW_{t}$
$\displaystyle+\log\left(\frac{d\nu(X_{0})}{d\mu(X_{T})}\right)+\log Z\,,$
(66)

which again using (47), has to be equal to

$-\psi_{0}(X_{0})-\phi_{T}(X_{T})+\log\nu(X_{0})=-\psi_{0}(X_{0})-\log
p_{\text{target}}(X_{T})+\psi_{T}(X_{T})+\log\nu(X_{0})$

to be optimal, yielding the anticipated variance regularizer. Additionally, in
order to verify the terminal constraint along
$\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi_{t}}$, it suffices to
sum up (65)-(66), which again due to the first two parts of the loss, gives
$0$, which necessarily impose time-reversal consistency
$\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi_{t}}/\overleftarrow{\mathbb{P}}^{\mu,f-\sigma^{2}\nabla\psi_{t}}=1$
a.s.

In (d), we use $\nabla\phi_{t}$ and $\nabla\log\rho_{t}$ as optimization
variables instead of the two drifts, and it follows from the dynamical
formulation (37)-(38). One can also read (64) as

$\displaystyle\partial_{t}\log\rho_{t}(X_{t})$
$\displaystyle=\frac{1}{\rho_{t}}(\partial_{t}\rho_{t}+\nabla\rho_{t}^{\top}\dot{X_{t}})$
$\displaystyle=\frac{1}{\rho_{t}}[-\nabla\cdot(\rho_{t}(f_{t}+\sigma^{2}\nabla\phi_{t}))+\frac{\sigma^{2}}{2}\Delta\rho_{t}]+\frac{1}{\rho_{t}}\nabla\rho_{t}^{\top}(f_{t}+\sigma^{2}\nabla\phi_{t}-\frac{\sigma^{2}}{2}\nabla\log\rho_{t})$
$\displaystyle=-\sigma^{2}\Delta\phi_{t}+\frac{\sigma^{2}}{2}\Delta\log\rho_{t}-\nabla\cdot
f_{t}\,,$ (67)

which upon integrating and using integration by parts together with (70) give
the loss. This first part establishes a particular relationship between the
two variables, and the second part enforces optimality among all curves
transporting between $\nu$ and $\mu$ driven by such dynamical system – one
could also replace this part with (31) if $\bar{v}_{t}$ from the reference
process is known. ∎

######  Lemma 4 (SDE correspondence to SB optimality).

We have for the optimal forward drift $\nabla\phi_{t}$ and
$X\sim\overrightarrow{\mathbb{P}}^{\nu,f}$ as in (12),

$d\phi_{t}(X_{t})=-\frac{\sigma^{2}}{2}\|\nabla\phi_{t}(X_{t})\|^{2}dt+\sigma\nabla\phi_{t}(X_{t})^{\top}\,\overrightarrow{dW_{t}}\,,$

analogously for the optimal backward drift $-\nabla\psi_{t}$, along
$X\sim\overrightarrow{\mathbb{P}}^{\nu,f}$,

$d\psi_{t}(X_{t})=\Big{[}\frac{\sigma^{2}}{2}\|\nabla\psi_{t}\|^{2}-\nabla\cdot
f_{t}+\sigma^{2}\Delta\psi_{t}\Big{]}(X_{t})dt+\sigma\nabla\psi_{t}(X_{t})^{\top}\,\overrightarrow{dW_{t}}\,.$

Moreover, along the controlled forward dynamics
$X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}$, the optimal
control $\phi,\psi$ satisfy

$\displaystyle
d\phi_{t}(X_{t})=\frac{\sigma^{2}}{2}\|\nabla\phi_{t}(X_{t})\|^{2}dt+\sigma\nabla\phi_{t}(X_{t})^{\top}\,\overrightarrow{dW_{t}}\,,$
(68) $\displaystyle
d\psi_{t}(X_{t})=\Big{[}\frac{\sigma^{2}}{2}\|\nabla\psi_{t}\|^{2}+\nabla\cdot(\sigma^{2}\nabla\psi_{t}-f_{t})+\sigma^{2}\nabla\phi_{t}^{\top}\nabla\psi_{t}\Big{]}(X_{t})dt+\sigma\nabla\psi_{t}(X_{t})^{\top}\overrightarrow{dW_{t}}\,.$
(69)

In the above, $\nabla\phi_{t},\nabla\psi_{t}$ refer to the optimal forward /
backward drift in the SDE

$dX_{t}=[f_{t}(X_{t})+\sigma^{2}\nabla\phi_{t}(X_{t})]dt+\sigma\overrightarrow{dW_{t}},\;X_{0}\sim\nu\,,$
$dX_{t}=[f_{t}(X_{t})-\sigma^{2}\nabla\psi_{t}(X_{t})]dt+\sigma\overleftarrow{dW_{t}},\;X_{T}\sim\mu\,,$

and $\phi_{T}(X_{T})+\psi_{T}(X_{T})=\log
p_{\text{target}}(X_{T}),\phi_{0}(X_{0})+\psi_{0}(X_{0})=\log\nu(X_{0})$ for
$X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}$.

###### Proof of Lemma 4.

Using Itô’s lemma, we have along the reference SDE
$dX_{t}=f_{t}(X_{t})dt+\sigma dW_{t},X_{0}\sim\nu$,

$d\phi_{t}=\left[\frac{\partial\phi_{t}}{\partial
t}+\nabla\phi_{t}^{\top}f_{t}+\frac{\sigma^{2}}{2}\Delta\phi_{t}\right]dt+\sigma\nabla\phi_{t}^{\top}\overrightarrow{dW_{t}}\,.$

Now deducing from (43), since the optimal $\phi_{t}$ solves the PDE for all
$(t,x)\in[0,T]\times\mathbb{R}^{d}$

$\partial_{t}\phi_{t}=-f_{t}^{\top}\nabla\phi_{t}-\frac{\sigma^{2}}{2}\Delta\phi_{t}-\frac{\sigma^{2}}{2}\|\nabla\phi_{t}\|^{2}\,,$

substituting the last display into the previous one gives the result.
Analogously, along the same reference process with Itô’s lemma,

$d\psi_{t}=\left[\frac{\partial\psi_{t}}{\partial
t}+\nabla\psi_{t}^{\top}f_{t}+\frac{\sigma^{2}}{2}\Delta\psi_{t}\right]dt+\sigma\nabla\psi_{t}^{\top}\overrightarrow{dW_{t}}$

and using the fact (43) that the optimal $\psi_{t}$ solves for all
$(t,x)\in[0,T]\times\mathbb{R}^{d}$

$\partial_{t}\psi_{t}=-\nabla\psi_{t}^{\top}f_{t}-\nabla\cdot
f_{t}+\frac{\sigma^{2}}{2}\left(\|\nabla\psi_{t}\|^{2}+\Delta\psi_{t}\right)$

and plugging into the previous display yields the claim. In both of the PDE
derivations above, we used the fact that for any
$g\colon\mathbb{R}^{d}\rightarrow\mathbb{R}$,

$\frac{1}{g}\nabla^{2}g=\nabla^{2}\log g+\frac{1}{g^{2}}\nabla g\nabla
g^{\top}\Rightarrow\frac{1}{g}\Delta g=\Delta\log g+\|\nabla\log g\|^{2}$ (70)

by taking trace on both sides.

The second part of the lemma statement, where $X_{t}$ evolves along the
optimally controlled SDE, follows from [7, Theorem 3] up to a change of
variable. Notice the sign change and the absence of the cross term in the
dynamics for $\phi_{t},\psi_{t}$ when
$X\sim\overrightarrow{\mathbb{P}}^{\nu,f}$ vs.
$X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}$. ∎

######  Remark 14 (Equivalence between SDE, PDE, Path measure).

Adding (68)-(69) up and integrating over time, we get that along optimally
controlled forward trajectory, $\log
p_{\text{target}}(X_{T})-\log\nu(X_{0})=\int_{0}^{T}\frac{\sigma^{2}}{2}\|\nabla\phi_{t}(X_{t})+\nabla\psi_{t}(X_{t})\|^{2}+\nabla\cdot(\sigma^{2}\nabla\psi_{t}-f_{t})(X_{t})dt+\sigma\int_{0}^{T}(\nabla\phi_{t}+\nabla\psi_{t})(X_{t})\,\overrightarrow{dW_{t}}$,
exactly matching the KL objective (63)
$D_{KL}(\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}\|\overleftarrow{\mathbb{P}}^{\mu,f-\sigma^{2}\nabla\psi})=0$.
This gives an SDE alternative representation of
$\partial_{t}\log\rho^{*}_{t}(X_{t})$ in (67). The summation, however,
introduces ambiguity as any potential $+C$ shift in one that’s cancelled by
$-C$ in another will result in the same integral equality.

Proposition 2 and Lemma 1 are stated in Section 3.2, whose proof we give
below.

###### Proof of Proposition 2.

The optimal $\nabla\phi_{t},\nabla\psi_{t}$ allow us to calculate $\log Z$ for
$p_{\text{target}}$ using (63) as (can be used without expectation, or with
expectation and ignore the last zero-mean martingale term)

$\displaystyle\log Z$
$\displaystyle=\mathbb{E}_{\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}\left[-\frac{\sigma^{2}}{2}\int_{0}^{T}\|\nabla\phi_{t}(X_{t})+\nabla\psi_{t}(X_{t})\|^{2}dt+\int_{0}^{T}\nabla\cdot(f_{t}(X_{t})-\sigma^{2}\nabla\psi_{t}(X_{t}))dt-\log\frac{d\nu(X_{0})}{d\mu(X_{T})}\right]$
$\displaystyle+\mathbb{E}_{\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}\left[-\sigma\int_{0}^{T}(\nabla\phi_{t}+\nabla\psi_{t})(X_{t})\,\overrightarrow{dW_{t}}\right]=:\mathbb{E}[-S]\,.$

Alternatively using (50), with the optimal $\nabla\phi^{*},\nabla\psi^{*}$,
since $Z$ is independent of $X_{T}$,

$-\log
Z=\log\nu(X_{0})-\frac{\sigma^{2}}{2}\int_{0}^{T}\nabla\cdot(\nabla\log\phi_{t}^{*}-\nabla\log\psi_{t}^{*})(X_{t})\,dt-\int_{0}^{T}\nabla\cdot
f_{t}(X_{t})dt-\log\mu(X_{T})\,,$ (71)

in which case the estimator is exact with $X_{t}$ following (49).

In general for imperfect control, since
$D_{KL}(\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}\|\overleftarrow{\mathbb{P}}^{\mu,f-\sigma^{2}\nabla\psi})>0$,
$\log Z$ will only be lower bounded by $\mathbb{E}[-S]$. Using Lemma 3
however,

$\displaystyle
1=\mathbb{E}_{\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}\left[\left(\frac{d\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}{d\overleftarrow{\mathbb{P}}^{\mu,f-\sigma^{2}\nabla\psi}}\right)^{-1}\right]$
$\displaystyle=\mathbb{E}_{\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}\left[\exp\left(-\frac{\sigma^{2}}{2}\int_{0}^{T}\|\nabla\phi_{t}+\nabla\psi_{t}\|^{2}+\nabla\cdot(f_{t}-\sigma^{2}\nabla\psi_{t})dt-\sigma\int_{0}^{T}(\nabla\phi_{t}+\nabla\psi_{t})dW_{t}-\log\frac{\nu}{\mu}\right)\frac{1}{Z}\right]$
(72) $\displaystyle=:\mathbb{E}[\exp(-S^{\prime})/Z]\,,$

giving
$Z=\mathbb{E}_{\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}[\exp(-S^{\prime})]$
as unbiased estimator using any (potentially sub-optimal) controls
$\nabla\phi,\nabla\psi$.

For the importance sampling, we use path weights suggested by the terminal
requirement that $X_{T}^{\phi^{*}}\sim\mu$, therefore we let (using (67))

$\displaystyle w^{\phi}(X^{\phi}_{T})$
$\displaystyle=\frac{d\mu}{d\mathbb{P}_{X_{T}^{\phi}}}(X_{T}^{\phi})=\frac{d\mu(X_{T}^{\phi})}{d\nu(X_{0}^{\phi})}\frac{d\nu(X_{0}^{\phi})}{d\mathbb{P}_{X_{T}^{\phi}}(X_{T}^{\phi})}$
$\displaystyle=\frac{d\mu(X_{T}^{\phi})}{d\nu(X_{0}^{\phi})}\exp\left(-\int_{0}^{T}\partial_{t}\log\rho^{\phi}_{t}(X_{t}^{\phi})\,dt\right)$
$\displaystyle=\frac{d\mu(X_{T}^{\phi})}{d\nu(X_{0}^{\phi})}\exp\left(\int_{0}^{T}\sigma^{2}\Delta\phi_{t}-\frac{\sigma^{2}}{2}\Delta\log\rho_{t}^{\phi}+\nabla\cdot
f_{t}\,dt\right)\,.$

Indeed, this choice guarantees
$\mathbb{E}_{\phi}[g(X_{T}^{\phi})w^{\phi}(X_{T}^{\phi})]=\int g(x)\mu(x)dx$
for any function $g$ on the terminal variable $X_{T}$ generated with the
suboptimal control $\phi$. It remains to normalize the weights (i.e., account
for the constant $Z$) via dividing by
$\mathbb{E}_{\phi}[w^{\phi}(X_{T}^{\phi})]=Z$. In the case of perfect
controls, all samples will have equal weight. ∎

######  Remark 15 (optimal control $\Leftrightarrow$ optimal estimator).

Notice (71) bears resemblance to Jarzynski’s identity in the context of free
energy calculation [16], since the term that’s been integrated is essentially
$\partial_{t}\log\rho_{t}(X_{t})$. In general one always has the importance
sampling identity

$1=\mathbb{E}_{\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}\left[\left(\frac{d\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}{d\overleftarrow{\mathbb{P}}^{\mu,f-\sigma^{2}\nabla\psi}}\right)^{-1}\right]\,.$

With the optimal controls $\phi^{*},\psi^{*}$, we see that since inside the
square bracket of (72), we have

$\exp\left(\phi_{0}^{*}(X_{0})-\phi_{T}^{*}(X_{T})+\psi_{0}^{*}(X_{0})-\psi_{T}^{*}(X_{t})-\log\frac{\nu(X_{0})}{\mu(X_{T})}-\log(Z)\right)=1$

holds deterministically (although the path is random), the $Z$-estimator is
optimal when optimally controlled in the sense that it’s unbiased and has zero
variance (it is almost surely a constant and equality holds not only in
expectation). This property is desirable for the Monte-Carlo estimates that we
employ. Such optimal control $\Leftrightarrow$ optimal sampling equivalence
also appears in [29] for path integral control problem that involves a _single
control_. This choice of $\phi^{*},\psi^{*}$ also admits the interpretation of
minimizing the $\chi^{2}$-divergence between path measures:

$\chi^{2}(\overleftarrow{\mathbb{P}}^{\mu,f-\sigma^{2}\nabla\psi}\|d\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi})=\text{Var}_{\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}\left[\frac{\overleftarrow{\mathbb{P}}^{\mu,f-\sigma^{2}\nabla\psi}}{\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}\right]\,.$

######  Remark 16 (Special case).

In the case of half-bridge from Dirac delta [34], there is a different
decomposition one can do. For any $\phi$ and
$X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}$,

$\displaystyle w^{\phi}(X)$
$\displaystyle=\frac{d\mathbb{P}_{X^{\phi^{*}}}}{d\mathbb{P}_{X^{\phi}}}(X)=\frac{d\mathbb{P}_{X^{\phi^{*}}}}{d\mathbb{P}_{X^{r}}}(X)\frac{d\mathbb{P}_{X^{r}}}{d\mathbb{P}_{X^{\phi}}}(X)$
$\displaystyle=\frac{dp_{\text{target}}}{d\mathbb{P}_{X^{r}_{T}}}(X_{T})\frac{d\mathbb{P}_{X^{r}}}{d\mathbb{P}_{X^{\phi}}}(X)$
$\displaystyle=\exp\left(\log\frac{dp_{\text{target}}}{d\mathbb{P}_{X^{r}_{T}}}(X_{T})-\frac{1}{2\sigma^{2}}\int_{0}^{T}\|f_{t}+\sigma^{2}\nabla\phi_{t}\|^{2}dt-\frac{1}{\sigma}\int_{0}^{T}(f_{t}+\sigma^{2}\nabla\phi_{t})\,dW_{t}\right)\,,$

where we assumed the reference $r$ is drift-free:

$dX_{t}=\sigma dW_{t}\,,\quad X_{0}\sim\nu$

hence $X_{T}^{r}\sim\nu\ast\mathcal{N}(0,\sigma^{2}T\cdot
I)\overset{d}{=}y+\sigma\sqrt{T}z$ for $y\sim\nu$, and
$\mathbb{P}_{X^{r}_{T}}$ is easy to evaluate. But it crucially hinges on the
property stated in Remark 8 for the third equality.

###### Proof of Lemma 1.

We obtain the forward trajectory

$X_{k+1}^{i}=X_{k}^{i}+(f_{k}(X_{k}^{i})+\sigma^{2}\nabla\phi_{k}(X_{k}^{i}))h+\sigma\sqrt{h}\cdot
z_{k}^{i},\;z_{k}^{i}\sim\mathcal{N}(0,I)$ (73)

with Euler-Maruyama for each of the $i\in[N]$ samples. Rewriting the term (66)
using (6), (56)-(57) and proceeding with the approximation (4) give

$\displaystyle\frac{1}{\sigma^{2}}\int_{0}^{T}f_{t}(X_{t})\overrightarrow{dX_{t}}-\frac{1}{2\sigma^{2}}\int_{0}^{T}\|f_{t}(X_{t})\|^{2}dt-\frac{1}{\sigma^{2}}\int_{0}^{T}(f_{t}-\sigma^{2}\nabla\psi_{t})(X_{t})\overleftarrow{dX_{t}}+\frac{1}{2\sigma^{2}}\int_{0}^{T}\|(f_{t}-\sigma^{2}\nabla\psi_{t})(X_{t})\|^{2}dt$
$\displaystyle\approx\frac{1}{\sigma^{2}}\sum_{k=0}^{K-1}f_{k}(X_{k})^{\top}(X_{k+1}-X_{k})-\frac{1}{2\sigma^{2}}\|f_{k}(X_{k})\|^{2}h$
$\displaystyle-\frac{1}{\sigma^{2}}(f_{k+1}(X_{k+1})-\sigma^{2}\nabla\psi_{k+1}(X_{k+1}))(X_{k+1}-X_{k})+\frac{1}{2\sigma^{2}}\|f_{k+1}(X_{k+1})-\sigma^{2}\nabla\psi_{k+1}(X_{k+1})\|^{2}h$
$\displaystyle=\sum_{k=0}^{K-1}\frac{1}{2\sigma^{2}h}\|X_{k}-X_{k+1}+(f_{k+1}-\sigma^{2}\nabla\psi_{k+1})(X_{k+1})h\|^{2}-\sum_{k=0}^{K-1}\frac{1}{2\sigma^{2}h}\|X_{k+1}-X_{k}-f_{k}(X_{k})h\|^{2}\,,$

which is the estimator for the variance regularizer with $(X_{k})_{k}$
following (73).

It is also possible to avoid the divergence (e.g., Hutchinson’s trace
estimator) in (72) when discretizing by leveraging similar ideas. Direct
computation using (56)-(57) tell us

$Z\approx\frac{1}{N}\sum_{i=1}^{N}\exp\left(\log\frac{\mu}{\nu}-\sum_{k=0}^{K-1}\frac{1}{2\sigma^{2}h}\|X_{k}^{i}-X_{k+1}^{i}+(f_{k+1}-\sigma^{2}\nabla\psi_{k+1})(X_{k+1}^{i})h\|^{2}+\frac{1}{2}\|z_{k}^{i}\|^{2}\right)$

for the same Euler-Maruyama trajectory (73) as the estimator for the
normalizing constant, where we used that the additional term

$\sum_{k=0}^{K-1}\frac{1}{2\sigma^{2}h}\|X_{k+1}-X_{k}-(f_{k}+\sigma^{2}\nabla\phi_{k})(X_{k})h\|^{2}=\frac{1}{2}\sum_{k=0}^{K-1}\|z_{k}\|^{2}$

from the update. ∎

Lastly we present an auxiliary Lemma 2 from Section B and a word about
contraction for method (10).

###### Proof of Lemma 2.

We have up to a constant independent of $\phi,\psi$, using Lemma 3, (52)-(53)
give the loss

$\displaystyle\psi_{n}\leftarrow\arg\min_{\psi}\int_{0}^{T}\mathbb{E}_{X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi_{n-1}}}\left[\frac{\sigma^{2}\|\nabla\psi\|^{2}}{2}-\sigma^{2}\nabla\phi_{n-1}^{\top}\nabla\psi-\nabla\cdot(\sigma^{2}\nabla\psi)\right](X_{t})dt$
$\displaystyle\phi_{n}\leftarrow\arg\min_{\phi}\int_{0}^{T}\mathbb{E}_{X\sim\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^{2}\nabla\phi}}\left[\frac{\sigma^{2}}{2}\|\nabla\phi-\nabla\psi_{n}\|^{2}-\nabla\cdot(f+\sigma^{2}\nabla\psi_{n})\right](X_{t})dt+\mathbb{E}\left[\log\frac{\nu(X_{0})}{\mu(X_{T})}\right]$

where updating $\phi$ bears resemblance to (9), and updating $\psi$ is akin to
the score matching loss (10). The equivalence to (51) assuming $\phi,\psi$ is
expressive enough (i.e.,
$(\phi^{*},\psi^{*})=(\phi_{\theta^{*}},\psi_{\gamma^{*}})$) and $\phi_{0}=0$
is shown in [32] and relies on the fact that $\phi,\psi$ update (52)-(53)
solve respectively

$\nabla\psi_{n}=\nabla\phi_{n-1}-\nabla\log\rho_{t}^{\nu,f+\sigma^{2}\nabla\phi_{n-1}},\quad\nabla\phi_{n}=\nabla\psi_{n}-\nabla\log\rho_{t}^{\mu,f+\sigma^{2}\nabla\psi_{n}}\,,$

which if each update is minimized to $0$, would fix one end of the marginal
alternatively, while reversing the path-wise transition from the previous
step. It is known that IPF update (51) converges to $\pi^{*}$ as
$n\rightarrow\infty$. ∎

######  Remark 17 (Contraction).

In Section 2.2 we claimed that for two processes $(p_{t})_{t},(q_{t})_{t}$
with different initialization, but same drift $\sigma v+\sigma^{2}s$ where $s$
is the score function,
$\partial_{t}\,D_{KL}(p_{t}\|q_{t})=-\sigma^{2}/2\cdot\mathbb{E}_{p_{t}}[\|\nabla\log\frac{p_{t}}{q_{t}}\|^{2}]\leq
0$ contracts towards each other, when we run the generative process

$dX_{t}=[\sigma
v_{t}(X_{t})+\sigma^{2}s_{t}(X_{t})]\,dt+\sigma\overrightarrow{dW_{t}},\;X_{0}\sim\nu\neq\overleftarrow{\mathbb{P}}_{0}^{\mu,\sigma
v}\,,$

although the actual rate maybe slow. We give this short calculation here for
completeness:

$\displaystyle\partial_{t}\,$ $\displaystyle D_{KL}(p_{t}\|q_{t})$
$\displaystyle=\int(\partial_{t}p_{t}(x))\log\frac{p_{t}(x)}{q_{t}(x)}\,dx+\int
q_{t}(x)\left[\frac{\partial_{t}p_{t}(x)}{q_{t}(x)}-\frac{p_{t}(x)\partial_{t}q_{t}(x)}{q_{t}^{2}(x)}\right]\,dx$
$\displaystyle=\int p_{t}\left[\sigma
v_{t}+\sigma^{2}s_{t}-\frac{\sigma^{2}}{2}\nabla\log
p_{t}(x)\right]^{\top}\nabla\log\frac{p_{t}(x)}{q_{t}(x)}dx-\int\frac{p_{t}(x)}{q_{t}(x)}\partial_{t}q_{t}(x)dx$
$\displaystyle=\int p_{t}\left[\sigma
v_{t}+\sigma^{2}s_{t}-\frac{\sigma^{2}}{2}\nabla\log
p_{t}(x)\right]^{\top}\nabla\log\frac{p_{t}(x)}{q_{t}(x)}dx-\int\nabla\frac{p_{t}(x)}{q_{t}(x)}\left[\sigma
v_{t}+\sigma^{2}s_{t}-\frac{\sigma^{2}}{2}\nabla\log
q_{t}(x)\right]q_{t}(x)dx$ $\displaystyle=\int\left[\sigma
v_{t}+\sigma^{2}s_{t}-\frac{\sigma^{2}}{2}\nabla\log p_{t}(x)-\sigma
v_{t}-\sigma^{2}s_{t}+\frac{\sigma^{2}}{2}\nabla\log
q_{t}(x)\right]^{\top}\nabla\log\frac{p_{t}(x)}{q_{t}(x)}\cdot p_{t}(x)dx$
$\displaystyle=-\frac{\sigma^{2}}{2}\int\left\|\nabla\log\frac{p_{t}(x)}{q_{t}(x)}\right\|^{2}p_{t}(x)dx\leq
0\,.$

##  Appendix D Second-order Dynamics

The motivation of this section is to design a smoother trajectory for $X_{t}$
while still maintaining optimality (and therefore uniqueness) of the dynamics.
Suppose the reference $Q$ is given by the following augmented process with a
velocity variable $V_{t}\in\mathbb{R}^{d}$, and $(X_{t},V_{t})$ have as the
stationary distribution $\mathcal{N}(0,I)\otimes\mathcal{N}(0,I)$:

$\displaystyle dX_{t}$ $\displaystyle=V_{t}\,dt$ (74) $\displaystyle dV_{t}$
$\displaystyle=-X_{t}dt-\gamma
V_{t}dt+\sqrt{2\gamma}\,\overrightarrow{dW_{t}}$ (75)

We assume it is initialized at
$Z_{0}:=(X_{0},V_{0})\sim\nu\otimes\mathcal{N}(0,I)$ independent and would
like to enforce $Z_{T}:=(X_{T},V_{T})\sim\mu\otimes\mathcal{N}(0,I)$ at the
terminal. In this second-order case [11], to solve the optimal trajectory
problem, we formulate $P$ as the controlled process

$\displaystyle\inf_{\rho,u}\;\int_{\mathbb{R}^{2d}}\int_{0}^{T}\gamma\|u_{t}(Z_{t})\|^{2}\rho_{t}(Z_{t})\,dtdz$
(76) $\displaystyle\text{s.t.}\;dX_{t}=V_{t}\,dt$ (77)
$\displaystyle\quad\quad dV_{t}=-[X_{t}+2\gamma u_{t}(Z_{t})]dt-\gamma
V_{t}dt+\sqrt{2\gamma}\,\overrightarrow{dW_{t}},\;\rho_{0}(Z_{0})=\nu\otimes\mathcal{N}(0,I),\rho_{T}(Z_{T})=\mu\otimes\mathcal{N}(0,I)$
(78)

One could check via the Girsanov’s theorem that by picking

$\bar{b}_{t}=\begin{bmatrix}V_{t}\\\ -X_{t}-\gamma V_{t}\end{bmatrix},\quad
b_{t}=\begin{bmatrix}V_{t}\\\ -X_{t}-2\gamma u_{t}(Z_{t})-\gamma
V_{t}\end{bmatrix},\quad G_{t}=\begin{bmatrix}0&0\\\
0&\sqrt{2\gamma}I\end{bmatrix}$

since $\bar{b}_{t}-b_{t}\in\text{image}(G_{t})$, the KL divergence between the
two path measures $P$ and $Q$ is
$\mathbb{E}[\int_{0}^{T}\gamma\|u_{t}(Z_{t})\|^{2}dt]$. This is the SB goal of

$\min_{P_{0}\sim\nu\otimes\mathcal{N},P_{T}\sim\mu\otimes\mathcal{N}}D_{KL}(P\|Q)\,,$

which if minimized will identify the path with the minimum control effort
between the two time marginals. We note that this formulation is different
from the one in [6, Section 3], and is closer to an underdamped version of the
SB dynamics. We denote the corresponding controlled path measure (77)-(78) as
$\overrightarrow{\mathbb{P}}$.

The Radon-Nikodym likelihood ratio from Lemma 3 still applies, therefore to
impose path measure consistency, we introduce a backward control $o(Z_{t},t)$
for the process $\overleftarrow{\mathbb{P}}$:

$\displaystyle dX_{t}$ $\displaystyle=V_{t}\,dt$ (79) $\displaystyle dV_{t}$
$\displaystyle=-[X_{t}-2\gamma o_{t}(Z_{t})]dt-\gamma
V_{t}dt+\sqrt{2\gamma}\,\overleftarrow{dW_{t}}\,.$ (80)

The KL divergence
$D_{KL}(\overrightarrow{\mathbb{P}}\|\overleftarrow{\mathbb{P}})$ for the
augmented process becomes (from now on we will take $\gamma=2$ for simplicity
corresponding to critical damping)

$\mathbb{E}_{Z\sim\overrightarrow{\mathbb{P}}}\left[\int_{0}^{T}2\|u_{t}(Z_{t})+o_{t}(Z_{t})\|^{2}-\nabla\cdot\hat{b}_{t}(Z_{t})\,dt+\log\frac{\nu(X_{0})\otimes\mathcal{N}(V_{0})}{\mu(X_{T})\otimes\mathcal{N}(V_{T})}\right]+C$
(81)

for $\gamma=2$ and

$\hat{b}_{t}=\begin{bmatrix}V_{t}\\\ -X_{t}+2\gamma o_{t}(Z_{t})-\gamma
V_{t}\end{bmatrix}\,,$

with $\nabla\cdot\hat{b}_{t}(Z_{t})=-\gamma
d+2\gamma\text{div}_{V}(o_{t}(Z_{t}))$. We minimize over vector fields
$u(Z_{t},t),o(Z_{t},t)$ taking gradient forms to impose time reversal
consistency. Below we derive a corresponding PINN and a variance regularizer
for this dynamics, first on the forward drift followed by the backward drift.

#### Forward condition

Writing out the optimality condition for the optimization problem (76), we
have for a Lagrange multiplier $\lambda(Z_{t},t)$, since by continuity
equation the dynamics (77)-(78) can be rewritten as

$\displaystyle\partial_{t}\rho_{t}(Z_{t})+\nabla_{[X,V]}\cdot\left(\rho_{t}(Z_{t})\begin{bmatrix}V_{t}\\\
-X_{t}-4u_{t}(Z_{t})-2V_{t}-2\nabla_{V}\log\rho_{t}(Z_{t})\end{bmatrix}\right)=0\,,$
(82)

we seek to optimize the Lagrangian

$\displaystyle\int_{\mathbb{R}^{2d}}\int_{0}^{T}2\|u_{t}(Z_{t})\|^{2}\rho_{t}(Z_{t})+\lambda(Z_{t},t)\left[\partial_{t}\rho+\nabla_{[X,V]}\cdot\left(\rho_{t}(Z_{t})\begin{bmatrix}V_{t}\\\
-X_{t}-4u_{t}(Z_{t})-2V_{t}-2\nabla_{V}\log\rho_{t}(Z_{t})\end{bmatrix}\right)\right]\,dtdz\,.$

Performing integration by parts, we have

$\displaystyle\int_{\mathbb{R}^{2d}}\int_{0}^{T}\left\\{2\|u_{t}(Z_{t})\|^{2}-\partial_{t}\lambda(Z_{t},t)-\nabla_{[X,V]}\lambda(Z_{t},t)^{\top}\begin{bmatrix}V_{t}\\\
-X_{t}-4u_{t}(Z_{t})-2V_{t}\end{bmatrix}-2\Delta_{V}\lambda(Z_{t},t)\right\\}\rho_{t}(Z_{t})\,dtdz\,.$

Fixing $\rho$, the optimal control $u_{t}\in\mathbb{R}^{d}$ satisfies

$u_{t}(Z_{t})=-\nabla_{V}\lambda(Z_{t},t)\,.$ (83)

Plugging this choice back in, the optimal $\lambda(Z_{t},t)$ must satisfies
the following PDE:

$\displaystyle
2\|\nabla_{V}\lambda(Z_{t},t)\|^{2}-\partial_{t}\lambda(Z_{t},t)-\nabla_{X}\lambda(Z_{t},t)^{\top}V_{t}+\nabla_{V}\lambda(Z_{t},t)^{\top}(X_{t}-4\nabla_{V}\lambda(Z_{t},t)+2V_{t})-2\Delta_{V}\lambda(Z_{t},t)=0$
$\displaystyle\Rightarrow\partial_{t}\lambda(Z_{t},t)=-2\|\nabla_{V}\lambda(Z_{t},t)\|^{2}-\nabla_{[X,V]}\lambda(Z_{t},t)^{\top}\bar{b}_{t}-2\Delta_{V}\lambda(Z_{t},t)$
(84)

where the $\Delta_{V}\lambda(Z_{t},t)$ term is interpreted as
$\text{Trace}(\nabla^{2}_{V}\lambda(Z_{t},t))$. Turning to the variance loss,
use Itô’s lemma we get that along the reference SDE (74)-(75),

$d\lambda_{t}(Z_{t})=\left[\frac{\partial\lambda_{t}(Z_{t})}{\partial
t}+\nabla_{X}\lambda_{t}(Z_{t})^{\top}V_{t}-\nabla_{V}\lambda_{t}(Z_{t})^{\top}(X_{t}+2V_{t})+2\Delta_{V}\lambda_{t}(Z_{t})\right]dt+2\nabla_{V}\lambda_{t}(Z_{t})^{\top}dW_{t}\,,$

which by plugging in the PDE (84), we end up with

$d\lambda_{t}(Z_{t})=-2\|\nabla_{V}\lambda(Z_{t},t)\|^{2}\,dt+2\nabla_{V}\lambda(Z_{t},t)^{\top}dW_{t}\,.$

It implies the regularizer on the forward control as

$\text{Var}\left(\lambda(Z_{T},T)-\lambda(Z_{0},0)+\int_{0}^{T}2\|\nabla_{V}\lambda(Z_{t},t)\|^{2}dt-2\int_{0}^{T}\nabla_{V}\lambda(Z_{t},t)^{\top}dW_{t}\right).$
(85)

Once we have the optimal control $u(Z_{t},t)$ from (83), one can simply
simulate (77)-(78) from $Z_{0}$ to draw samples from $\mu$. We recommend a
splitting approach where one alternates between the Hamiltonian part (with
reversible symplectic leapfrog denoted by flow map $\Phi_{h}$), the $u_{t}$
drift part (with Euler), and the rest OU part (with exact simulation) for SDE
discretization.

#### Backward condition

Via a change of variable

$\log\rho_{t}(Z_{t})=\lambda_{t}(Z_{t})+\eta_{t}(Z_{t}),\quad
u_{t}(Z_{t})=-\nabla_{V}\lambda_{t}(Z_{t})\,,$

now (82) and (84) imply the PDE regularizer on the backward control
$o(Z_{t},t)=-\nabla_{V}\eta(Z_{t},t)$ should read as

$\partial_{t}\eta(Z_{t},t)=-\nabla_{[X,V]}\eta(Z_{t},t)^{\top}\bar{b}_{t}-\nabla\cdot\bar{b}_{t}+2\|\nabla_{V}\eta(Z_{t},t)\|^{2}+2\Delta_{V}\eta(Z_{t},t)\,,$
(86)

and the corresponding SDE regularizer along the reference (74)-(75) becomes
(using $\nabla\cdot\bar{b}_{t}=-2d$):

$d\eta_{t}(Z_{t})=\left[2\|\nabla_{V}\eta(Z_{t},t)\|^{2}+4\Delta_{V}\eta(Z_{t},t)+2d\right]\,dt+2\nabla_{V}\eta(Z_{t},t)^{\top}dW_{t}\,.$

Moreover the boundary condition should match the marginal densities at $t=0,T$
in the sense that time reversal imposes
$u_{t}(Z_{t})+o_{t}(Z_{t})=-\nabla_{V}\log\rho_{t}(Z_{t})$ for all $t\in[T]$.
From (82) the deterministic ODE implementation of the dynamics with given
$(u_{t}(Z_{t}),o_{t}(Z_{t}))$ therefore becomes

$dZ_{t}=\begin{bmatrix}dX_{t}\\\ dV_{t}\end{bmatrix}=\begin{bmatrix}V_{t}\\\
-X_{t}-2V_{t}-2u_{t}(Z_{t})+2o_{t}(Z_{t})\end{bmatrix}dt\,.$ (87)

######  Remark 18 (Cost and evolution).

We would like to note that the reference process (74)-(75) still gives an
entropy regularized optimal transport objective with quadratic cost, since

$\displaystyle\arg\min_{P_{0}\sim\nu\otimes\mathcal{N},P_{T}\sim\mu\otimes\mathcal{N}}\,KL(P_{0T}\|Q_{0T})$
(88)
$\displaystyle=\arg\min_{P_{0}\sim\nu\otimes\mathcal{N},P_{T}\sim\mu\otimes\mathcal{N}}\,\mathbb{E}_{P_{0T}}[-\log
Q_{T|0}]-\int P(Z_{0},Z_{T})\log P(Z_{0},Z_{T})\,dZ_{0}dZ_{T}$
$\displaystyle=\arg\min_{P_{0}\sim\nu\otimes\mathcal{N},P_{T}\sim\mu\otimes\mathcal{N}}\,\mathbb{E}_{P_{0T}}[\underbrace{(Z_{T}-M_{0T}Z_{0})^{\top}D^{-1}_{T|0}(Z_{T}-M_{0T}Z_{0})}_{\mathcal{C}_{T}(Z_{0},Z_{T})}]-H(P_{0T})\,,$

where we used that the conditional distribution of $Z_{T}$ given $Z_{0}$ is
Gaussian, which we denote with mean $M_{0T}Z_{0}$ and covaraince $D_{T|0}$.
Using the fact that the reference is a linear SDE that admits analytical
solution, assuming we initialize as

$(X_{0},V_{0})\sim\mathcal{N}\left(\begin{pmatrix}\mu_{X}^{0}\\\
\mu_{V}^{0}\end{pmatrix},\begin{pmatrix}\sigma_{XX}^{0}\cdot I&0\\\
0&\sigma_{VV}^{0}\cdot I\end{pmatrix}\right)\,,$

the dynamics will remain Gaussian with

$\mathbb{E}[Z_{T}]=\begin{bmatrix}(T+1)\mu_{X}^{0}+T\mu_{V}^{0}\\\
-T\mu_{X}^{0}+(1-T)\mu_{V}^{0}\end{bmatrix}e^{-T}=\begin{bmatrix}(T+1)e^{-T}\cdot
I&Te^{-T}\cdot I\\\ -Te^{-T}\cdot I&(1-T)e^{-T}\cdot
I\end{bmatrix}\begin{pmatrix}\mu_{X}^{0}\\\
\mu_{V}^{0}\end{pmatrix}\rightarrow 0_{2d}$

and

$\text{Cov}[Z_{T}]=\begin{bmatrix}\sigma_{XX}^{T}\cdot I&\sigma_{XV}^{T}\cdot
I\\\ \sigma_{VX}^{T}\cdot I&\sigma_{VV}^{T}\cdot
I\end{bmatrix}e^{-2T}\rightarrow\begin{bmatrix}I&0\\\ 0&I\end{bmatrix}$

for

$\displaystyle\sigma_{XX}^{T}$
$\displaystyle=(T+1)^{2}\sigma_{XX}^{0}+T^{2}\sigma_{VV}^{0}+e^{2T}-1-2T-2T^{2}$
$\displaystyle\sigma_{XV}^{T}$
$\displaystyle=\sigma_{VX}^{T}=-T(T+1)\sigma_{XX}^{0}+T(1-T)\sigma_{VV}^{0}+2T^{2}$
$\displaystyle\sigma_{VV}^{T}$
$\displaystyle=T^{2}\sigma_{XX}^{0}+(1-T)^{2}\sigma_{VV}^{0}-2T^{2}+2T+e^{2T}-1$

and plugging in
$\mu_{X}^{0}=x_{0},\mu_{V}^{0}=v_{0},\sigma_{XX}^{0}=\sigma_{VV}^{0}=0$ gives
the desired $M_{0T}$ and $D_{T|0}$. This asymptotic behavior matches our
intuition about the equilibrium.

Note that the coupling
$P_{0T}(X_{0},V_{0},X_{T},V_{T})=\pi^{*}(X_{0},X_{T})\mathcal{N}(V_{0})\mathcal{N}(V_{T})$
for $\int\pi^{*}(X_{0},X_{T})dX_{T}=\nu(X_{0})$,
$\int\pi^{*}(X_{0},X_{T})dX_{0}=\mu(X_{T})$ is always valid, in which case the
objective will simplify to a quadratic involving $(X_{0},X_{T})$ only since
$\mathbb{E}_{P_{0T}}[V_{0}V_{T}]=\mathbb{E}_{P_{0T}}[V_{0}X_{T}]=\mathbb{E}_{P_{0T}}[X_{0}V_{T}]=0$
and the constraints already enforce $\mathbb{E}_{P_{0T}}[\|V_{T}\|^{2}]$ and
$\mathbb{E}_{P_{0T}}[\|V_{0}\|^{2}]$ to be fixed and
$\mathbb{E}_{P_{0T}}[X_{T}V_{T}]=\mathbb{E}_{P_{0T}}[X_{0}V_{0}]=0$. Similarly
the entropy term $H(\cdot)$ will also only involve $(X_{0},X_{T})$ in this
case. But generally the marginal over $(X_{0},X_{T})$ from solving the optimal
transport SB problem (88) will not correspond to $\pi^{*}(X_{0},X_{T})$ from
(29), as we show below.

######  Lemma 5 (Structure of the joint optimal coupling).

The joint coupling from (88) in general do not factorize over the $X$ and $V$
variables, unless $T\rightarrow\infty$, in which case the coupling over the
$X$ variables will favor the independent one, i.e.,
$\nu(X_{0})\otimes\mu(X_{T})$.

###### Proof.

The optimal regularized coupling always takes the form for some $f,g$ [20]

$\displaystyle\pi^{*}(Z_{0},Z_{T})$ $\displaystyle=e^{f(Z_{0})}\cdot
e^{g(Z_{T})}\cdot
Q(Z_{T}|Z_{0})\cdot\nu(X_{0})\otimes\mathcal{N}(V_{0})\otimes\mu(X_{T})\otimes\mathcal{N}(V_{T})$
$\displaystyle=e^{f_{1}(X_{0})+f_{2}(V_{0})}\cdot
e^{g_{1}(X_{T})+g_{2}(V_{T})}\cdot
Q(Z_{T}|Z_{0})\cdot\nu(X_{0})\otimes\mu(X_{T})\otimes\mathcal{N}(V_{0})\otimes\mathcal{N}(V_{T})$
$\displaystyle=[e^{f_{1}(X_{0})+g_{1}(X_{T})}\nu(X_{0})\otimes\mu(X_{T})]\cdot[e^{f_{2}(V_{0})+g_{2}(V_{T})}\mathcal{N}(V_{0})\otimes\mathcal{N}(V_{T})]\cdot
Q(Z_{T}|Z_{0})\,.$

While the rest of the terms may factorize, there will be cross terms between
e.g., $X_{T}-V_{0}$ and $V_{T}-X_{0}$ coming from $Q(Z_{T}|Z_{0})$ unless
$T\rightarrow\infty$. Although the cost $\mathcal{C}_{T}(Z_{0},Z_{T})$ is
still quadratic, since it approaches in that limit

$\min_{X_{T}\sim\mu,V_{T}\sim\mathcal{N}}\;\mathbb{E}[\|X_{T}\|^{2}+\|V_{T}\|^{2}]$

using Remark 18, which is fixed by the constraint, the second entropy term
will dominate and therefore the resulting optimal coupling will be the product
distribution. ∎

######  Remark 19 (Gaussian entropy regularized $\mathcal{W}_{2}$ transport).

In the more commonly studied case when the cost is simply
$\|Z_{T}-Z_{0}\|^{2}$, and if both $\nu$ and $\mu$ are Gaussian with arbitrary
covariance, using closed-form expression from [15, Theorem 1] we have that the
regularized optimal transport plan is Gaussian over
$\mathbb{R}^{2d}\times\mathbb{R}^{2d}$:

$\displaystyle\pi^{*}\sim\mathcal{N}\left(\begin{pmatrix}[\bar{\nu},0]^{\top}\\\
[\bar{\mu},0]^{\top}\end{pmatrix},\begin{pmatrix}A&C_{T}\\\
C_{T}^{\top}&B\end{pmatrix}\right)$

for

$\displaystyle A=\begin{bmatrix}\Sigma_{\nu}&0\\\ 0&I\end{bmatrix}\quad
B=\begin{bmatrix}\Sigma_{\mu}&0\\\ 0&I\end{bmatrix}\quad
C_{T}=\frac{1}{2}A^{1/2}(4A^{1/2}BA^{1/2}+T^{2}I)^{1/2}A^{-1/2}-\frac{T}{2}I\,.$

We recognize

$\frac{1}{2}A^{1/2}(4A^{1/2}BA^{1/2}+T^{2}I)^{1/2}A^{-1/2}=A^{1/2}(A^{1/2}BA^{1/2}+T^{2}/4\cdot
I)^{1/2}A^{-1/2}=(AB+T^{2}/4\cdot I)^{1/2}$

as the unique square root of $AB+T^{2}/4\cdot I$ with non-negative eigenvalues
when $A\succ 0$. Therefore in this case, the optimal entropy-regularized
coupling does take the form of
$P_{0T}^{*}(X_{0},V_{0},X_{T},V_{T})=\pi^{*}(X_{0},X_{T})\bar{\pi}^{*}(V_{0},V_{T})$
where $\pi^{*}$ and $\bar{\pi}^{*}$ are respectively the entropy-regularized
optimal coupling between the $X$-variables and the $V$-variables (i.e., the
structure of $C_{T}$ implies that $X$ and $V$ are uncorrelated).

#### Normalizing constant estimator

Importance-weighted $Z$-estimators can be derived from the controlled sampling
trajectory (77)-(78) using similar ideas as those in Lemma 1 by leveraging
likelihood ratio (81). One can check that the discretized $Z$ estimator can be
written as

$\displaystyle\frac{1}{N}\sum_{i=1}^{N}\exp\Big{(}\log\frac{\mu(X_{K})\otimes\mathcal{N}(V_{K})}{\nu(X_{0})\otimes\mathcal{N}(V_{0})}+$
(89)
$\displaystyle\sum_{k=0}^{K-1}\frac{1}{2(1-e^{-2h})}\|\hat{V}_{k+1}^{i}-e^{-h}V_{k}^{i}\|^{2}-\frac{1}{2(1-e^{-2h})}\|V_{k}^{i}-e^{-h}\\{[\Phi_{-h}(Z_{k+1}^{i})]_{V}+4h\nabla_{V}\eta_{k+1}(\Phi_{-h}(Z_{k+1}^{i}))\\}\|^{2}\Big{)}$

for $3$-stage splitting updates

$\hat{V}_{k+1}\sim\mathcal{N}(e^{-h}V_{k},(1-e^{-2h})I),\;\tilde{V}_{k+1}=\hat{V}_{k+1}+4h\cdot\nabla_{V}\lambda_{k}(X_{k},\hat{V}_{k+1}),\;Z_{k+1}=(X_{k+1},V_{k+1})=\Phi_{h}(X_{k},\bar{V}_{k+1})$

from $Z_{0}\sim\nu\otimes\mathcal{N}(0,I)$. Notice the last two parts are
purely deterministic and the first two parts act solely on variable $V$, which
explains the last line in (89).

##  Appendix E Numerical Experiments

We provide additional details on the numerical experiments in this section.

###  E.1 Benchmark and Metrics

The following $4$ targets are considered with different priors $\nu$:

  * •

2D standard Gaussian: $\mathcal{N}(x;0,I)$ with prior
$\nu(x)=\mathcal{N}(x;0,2I)$

  * •

Funnel:
$x_{0}\sim\mathcal{N}(0,3^{2}),x_{1}|x_{0}\sim\mathcal{N}(0,\exp(x_{0}))$ with
prior $\nu(x)=\mathcal{N}(x;0,2I)$

  * •

Gaussian mixture model with $9$ modes:
$\frac{1}{9}\sum_{i=1}^{9}\mathcal{N}(x;\mu_{i},I)\;\textrm{where}\;\\{\mu_{i}\\}_{i=1}^{9}=\\{-5,0,5\\}\times\\{-5,0,5\\}$
with prior $\nu(x)=\mathcal{N}(x;0,3.5^{2}I)$

  * •

Double well: $\mu(x)\propto-(x_{0}^{2}-2)^{2}-(x_{1}^{2}-2)^{2}$ with prior
$\nu(x)=\mathcal{N}(0,2I)$

For each of the benchmarks and $4$ losses, we plot the marginals and report
the following:

  * •

Absolute error in mean and relative error in standard deviation compared to
the ground truth using importance-weighted Monte-Carlo estimates (91)

  * •

The log normalizing constant $\log Z$ estimator (c.f. Lemma 1):

$\log\left(\frac{1}{n}\sum_{i=1}^{n}\frac{\mu(x_{K+1}^{i})}{\nu(x_{0}^{i})}\exp\left[\sum_{k=0}^{K}\frac{1}{2}\|Z_{k}^{i}\|^{2}-\frac{1}{2\sigma^{2}h}\|x_{k}^{i}-x_{k+1}^{i}+\sigma^{2}h(\nabla\phi(x_{k+1}^{i},(k+1)h)-\nabla
s(x_{k+1}^{i},(k+1)h))\|^{2}\right]\right)$

for the PINN, Variance and TD losses, and

$\log\left(\frac{1}{n}\sum_{i=1}^{n}\frac{\mu(x_{K+1}^{i})}{\nu(x_{0}^{i})}\exp\left[\sum_{k=0}^{K}\frac{1}{2}\|Z_{k}^{i}\|^{2}-\sum_{k=0}^{K}\frac{1}{2\sigma^{2}h}\|x_{k}^{i}-x_{k+1}^{i}-\sigma^{2}h\nabla\psi(x_{k+1}^{i},(k+1)h)\|^{2}\right]\right)$

for the separate control loss, both using the final trajectory
$\\{x_{k}^{i}\\}$ from (17).

#### Importance weighting

For the very last sampling SDE, we simulate (17) with the latest $\nabla\phi$,
but re-weight the $n$ samples $\\{x_{K+1}^{i}\\}_{i=1}^{n}$ each with
individual (un-normalized) weight

$w^{\phi}(x_{K+1}^{i})=\frac{\mu(x_{K+1}^{i})}{\nu(x_{0}^{i})}\exp\left(h\left[\sum_{k=0}^{K}\sigma^{2}\Delta\phi(x_{k}^{i},kh)-\frac{\sigma^{2}}{2}\Delta
s(x_{k}^{i},kh)\right]\right)$

for PINN, Var and TD losses. And

$w^{\phi}(x_{K+1}^{i})=\frac{\mu(x_{K+1}^{i})}{\nu(x_{0}^{i})}\exp\left(h\left[\sum_{k=0}^{K}\frac{\sigma^{2}}{2}\Delta\phi(x_{k}^{i},kh)-\frac{\sigma^{2}}{2}\Delta\psi(x_{k}^{i},kh)\right]\right)$
(90)

for the separate control before taking their average, i.e.,

$\hat{\mathbb{E}}_{p_{\text{target}}}[g]=\frac{\sum_{i=1}^{n}g(x_{K+1}^{i})w^{\phi}(x_{K+1}^{i})}{\sum_{i=1}^{n}w^{\phi}(x_{K+1}^{i})}$
(91)

for a summary statistics $g:\mathbb{R}^{d}\rightarrow\mathbb{R}$ we are
interested in. This estimator follows from Proposition 2 and is used for post-
processing with a potentially suboptimal control $\nabla\phi$, caused by e.g.,
approximation or optimization error from neural network training.

###  E.2 Result

For the experiments, we parameterize
$\phi,s:\mathbb{R}^{d}\times[0,c]\rightarrow\mathbb{R}$ as residual feed-
forward neural networks where $c$ is a hyperparameter. Each network first maps
the input vector to a hidden state with a linear transformation, which is then
propagated through several residual layers that maintain hidden state
dimensionality. Each residual layer $f$ receives as input a state-time pair
$(x,t)\in\mathbb{R}^{d+1}$ and outputs $f(x,t)=\sigma(W(x,t)+b)+x$, where
$\sigma$ is the ReLU activation function. The number of hidden layers and
their sizes vary by target distribution. Adam optimizer was used to train the
models with $\beta_{1}=0.9,\beta_{2}=0.999$, and weight decay $0.01$, where
batches of trajectories are used for several steps of gradient updates in each
epoch, before regenerating the $n$ trajectories and estimating the objective
for the next round of updates on the NN parameters. Across all experiments, we
initialize $\phi(x_{0}^{i},0)\approx\log\mu(x_{0}^{i})$. For fair comparison,
the training process is stopped when the loss stops noticeably decreasing.

For a comparison on the computation speed of the four regularizers: for the
standard normal target it took 624 seconds to train with PINN, 108 seconds
with the Variance regularizer, 106 seconds with TD, and 85 seconds with
Separate Control. Generating trajectories requires much less time than
evaluating the loss and its gradient. Computing each loss requires processing
all $K$ states across $n$ trajectories. Combined with the number of training
epochs and the number of updates per batch, we estimate the processing time
per trajectory state as approximately

$\frac{\texttt{training\\_time}}{\texttt{epochs}\times\texttt{trajectories}\times\texttt{updates\\_per\\_batch}\times
K}\,.$

The stated processing time is therefore roughly $6.2\cdot 10^{-3}$\mathrm{s}$$
for PINN, $3.6\cdot 10^{-7}$\mathrm{s}$$ for Variance, $3.5\cdot
10^{-7}$\mathrm{s}$$ for TD, and $2.8\cdot 10^{-7}$\mathrm{s}$$ for Separate
Control. The latter three are comparable and much faster than PINN whose
Laplacian computation adds an order of magnitude processing time. This
efficiency comparison is also independent of the target distribution.

For reproducibility, the anonymous Github repository can be found at the
following link:

https://anonymous.4open.science/r/diffusion-sampler-E4F7

where hyperparameters used for the experiments are listed in the corresponding
notebooks.

In Table 2, Table 3 and Figure 2, Figure 3, Figure 4 below we present the
simulation results. All experiments are run with a GeForce RTX 2080 GPU and an
AMD Ryzen 9 3900X CPU.

| PINN | Variance | TD | Separate control
---|---|---|---|---
Standard normal | $0.231\,533$ | $0.382\,497$ | $0.240\,691$ | $0.221\,310$
Funnel | $0.794\,941$ | $0.796\,652$ | $0.794\,814$ | $0.810\,974$
GMM | $0.171\,790$ | $0.088\,070$ | $0.134\,221$ | $0.032\,793$
Double well | $0.161\,144$ | $0.177\,658$ | $0.008\,204$ | $0.026\,290$
Table 2: Relative error of weighted empirical standard deviation (lower is better) | PINN | Variance | TD | Separate control
---|---|---|---|---
Standard normal | $0.123\,589$ | $0.037\,648$ | $0.013\,826$ | $0.030\,821$
Funnel | $0.786\,865$ | $0.295\,754$ | $0.255\,022$ | $0.154\,966$
GMM | $0.632\,307$ | $0.039\,433$ | $0.288\,841$ | $0.301\,473$
Double well | $0.234\,222$ | $0.736\,263$ | $0.165\,676$ | $0.361\,951$
Table 3: Absolute error of empirical mean (lower is better) Figure 2: Training
loss plot for GMM with TD and Separate Control

Figure 3: GMM marginal before (left) and after (right) importance weighting

Figure 4: Standard normal marginal (left) and Funnel marginal (right) with
importance weighting