Simulating Diffusion Bridges with Score Matching

It can be shown that the time-reversed bridge process $Z^{\star}=(Z_{t}^{\star})_{t\in[0,T]}=(X_{T-t}^{\star})_{t\in[0,T]}$ satisfies

with drift function $b(t,z_{t})=-f(T-t,z_{t})+\Sigma(T-t,z_{t})s(T-t,z_{t})+\nabla\cdot\Sigma(T-t,z_{t})$, another standard Brownian motion $B=(B_{t})_{t\in[0,T]}$, $s(t,x_{t})=s^{\star}(t,x_{t})-\nabla\log h(t,x_{t})$, and $\nabla\cdot\Sigma(t,x)=(\sum_{j=1}^{d}\partial_{x_{j}}\Sigma^{1,j}(t,x),\ldots,\sum_{j=1}^{d}\partial_{x_{j}}\Sigma^{d,j}(t,x))$ is the divergence of $\Sigma$. Here $s^{\star}(t,x_{t})=\nabla\log p^{\star}(t,x_{t})$ denotes the score of the marginal density $p^{\star}(t,x_{t})$ of the diffusion bridge process $X_{t}^{\star}$ at time $t\in(0,T)$. We refer readers to Haussmann and Pardoux (1986), Millet et al. (1989), and Cattiaux et al. (2023) for conditions under which the representation in (4) holds. By the Markov property, we have the relation

as $h(t,x_{t})=p(T,x_{T}\mid t,x_{t})$. This implies that $s(t,x_{t})=\nabla\log p(t,x_{t}\mid 0,x_{0})$ is simply the score of the transition density of $X$.

Exploiting this backward time representation to derive diffusion bridge approximations is also highly non-trivial due to the intractability of the transition density $\eta(t,x_{t})=p(t,x_{t}\mid 0,x_{0})$, which is now characterized by the forward Kolmogorov equation $\partial_{t}\eta(t,x_{t})+(\mathcal{F}\eta)(t,x_{t})=0$, with initial condition at time $0$ given by the Dirac measure at $x_{0}$ and $\mathcal{F}$ denotes the Fokker–Planck operator of $X$ (Stroock and Varadhan, 1997). Numerical resolution of $\eta$ using partial differential equation solvers also suffers from the same difficulties as (3). A key observation is that (4) can be understood as first setting $Z_{0}^{\star}=x_{T}$ to satisfy the terminal constraint, and then evolving $Z^{\star}$ using the time-reversal of (1). Due to the influence of the score $s(t,x_{t})$, the process will end at the initial constraint $Z_{T}^{\star}=x_{0}$ by construction. This connection between simulation of a diffusion bridge and time-reversal of its original diffusion process will form the basis of our score approximation. We refer readers to Section B of the Supplementary Material for an alternative and more elementary argument to establish this connection.

We introduce a variational formulation to learn the time-reversal of (1), involving path measures on the space of continuous functions from $[0,T]$ to $\mathbb{R}^{d}$, equipped with the cylinder $\sigma$-algebra. Consider the time-reversed process $Z=(Z_{t})_{t\in[0,T]}$ satisfying

with drift function $b_{\phi}(t,z_{t})=-f(T-t,z_{t})+\Sigma(T-t,z_{t})s_{\phi}(T-t,z_{t})+\nabla\cdot\Sigma(T-t,z_{t})$ that mimics the form of $b(t,z_{t})$ in (4). Here $s_{\phi}:[0,T]\times\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ represents a function approximation of the score $s(t,x_{t})$ that depends on parameters $\phi\in\Phi$ to be optimized. We shall measure the score approximation error as

where $E^{x_{0}}$ denotes expectation with respect to the path measure $\mathbb{P}^{x_{0}}$ induced by $(X_{t})_{t\in[0,T]}$ in (1) with $X_{0}=x_{0}$, and $\|\cdot\|_{A}$ denotes the Euclidean norm weighted by a positive definite $A\in\mathbb{R}^{d\times d}$.

Let $\mathbb{Q}_{\phi}^{x_{0}}$ be the path measure induced by the parameterised forward process $(Z_{T-t})_{t\in[0,T]}$. We consider minimizing the Kullback–Leibler divergence $\textsc{kl}(\mathbb{P}^{x_{0}}|\mathbb{Q}_{\phi}^{x_{0}})=E^{x_{0}}\{\log\mathrm{d}\mathbb{P}^{x_{0}}/\mathrm{d}\mathbb{Q}_{\phi}^{x_{0}}(X)\}$, where $\mathrm{d}\mathbb{P}^{x_{0}}/\mathrm{d}\mathbb{Q}_{\phi}^{x_{0}}$ denotes the Radon–Nikodym derivative of $\mathbb{P}^{x_{0}}$ with respect $\mathbb{Q}_{\phi}^{x_{0}}$. The following result gives an upper bound of this objective and shows that the process $Z$ will end at the initial constraint $Z_{T}=x_{0}$ if the score approximation error is finite.

Under Novikov’s condition, an exponential integrability assumption, Girsanov’s theorem shows that $\textsc{kl}(\mathbb{P}^{x_{0}}|\mathbb{Q}_{\phi}^{x_{0}})=e(\phi)/2$ (Le Gall, 2016a, Theorem 5.22, Theorem 4.13). We do not provide here any assumption on $s_{\phi}(t,x_{t})$ to guarantee that $e(\phi)<\infty$. Previous work by Delyon and Hu (2006) and Schauer et al. (2017) have studied necessary conditions for the law of a proposal bridge process and the law of the diffusion bridge to be equivalent. These findings suggest that the behaviour of $b_{\phi}(t,z_{t})$ should be $(x_{0}-z_{t})/t$ as $t\rightarrow 0$, which can be used to guide our parameterization of the score approximation. For example, we consider $s_{\phi}(t,x_{t})=\Sigma^{-1}(t,x_{t})(x_{0}-x_{t})/t+v_{\phi}(t,x_{t})$ where $v_{\phi}:[0,T]\times\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ denotes a bounded function.

Although Proposition 1 clearly relates the approximation of the time-reversal of (1) to the approximation of the score $s(t,x_{t})$, its form is not amenable to optimization. The following result gives an alternative and practical expression by adapting the idea of denoising score matching (Vincent, 2011) to our setting.

Therefore minimizing the Kullback–Leibler divergence $\textsc{kl}(\mathbb{P}^{x_{0}}|\mathbb{Q}_{\phi}^{x_{0}})$ is equivalent to minimizing the loss function $L(\phi)$. This allows us to circumvent the intractable score $s(t,x_{t})$ by working with $g(t_{m-1},x_{t_{m-1}},t,x_{t})$, the score of the transition density $p(t,x_{t}\mid t_{m-1},x_{t_{m-1}})$. Although the latter is also intractable, approximations can be made when the sub-interval $[t_{m-1},t_{m}]$ is sufficiently small. For example, under the Euler–Maruyama scheme (Kloeden and Platen, 1992, p. 340) with stepsize $\delta t=T/M$,

Hence the loss $L(\phi)$ can be approximated and minimized using stochastic gradient algorithms by simulating time-discretized paths under (1). The minimal loss of $L(\phi)=-C$, achieved when $s_{\phi}(t,x_{t})=s(t,x_{t})$ $\mathbb{P}^{x_{0}}$-almost surely, is unknown as the constant $C$ is typically intractable. After obtaining the score approximation $s_{\phi}$, we can simulate a proposal bridge $Z$ from (6) with $Z_{0}=x_{T}$ and correct it using importance sampling or independent Metropolis–Hastings. Time-discretization considerations and proposal correction procedures are detailed in Section C of the Supplementary Material.

In scenarios where one is interested in multiple pairs of conditioned states $(x_{0},x_{T})$, we can extend the above methodology to avoid having to learn multiple score approximations as follows. We let the score approximation $s_{\phi}(t,x_{t})$ in (6) also depend on the initial condition $x_{0}$, and average the Kullback–Leibler objective $\textsc{kl}(\mathbb{P}^{x_{0}}|\mathbb{Q}_{\phi}^{x_{0}})$ with a distribution $p_{0}(\mathrm{d}x_{0})$ on $\mathbb{R}^{d}$ that can be sampled from. By applying the arguments of Proposition 2 conditionally on $X_{0}=x_{0}$, we obtain a loss function given by averaging (8) over $p_{0}(\mathrm{d}x_{0})$, which can be minimized using time-discretization and stochastic gradient algorithms. In this setting, we parameterize the score approximation as $s_{\phi}(t,x_{t},x_{0})=\Sigma^{-1}(t,x_{t})(x_{0}-x_{t})/t+v_{\phi}(t,x_{t},x_{0})$, where $v_{\phi}(t,x_{t},x_{0})$ is now a function that also depends on the initial condition $x_{0}$.

It is instructive to consider the time-reversal of (4), which gives

with drift function $f^{\star}(t,x_{t})=-b(T-t,x_{t})+\Sigma(t,x_{t})s^{\star}(t,x_{t})+\nabla\cdot\Sigma(t,x_{t})$. Using the form of $b(t,z_{t})$ and the relation $s(t,x_{t})=s^{\star}(t,x_{t})-\nabla\log h(t,x_{t})$, we can rewrite $f^{\star}$ as

and recover the Doob’s $h$-transform in (2). Although this is to be expected as the reversal of the time-reversed process should recover the original process, it forms the basis of our approximation of (2).

After obtaining an approximation of $s(t,x_{t})$, another iteration of our methodology can be used to obtain a function approximation of $s^{\star}(t,x_{t})$. For brevity, this is detailed in Section 2.3 of the Supplementary Material. Plugging in both approximations in (10) then gives an approximation of $X^{\star}$, which could be of interest in algorithms where one requires a forward time representation of the proposal bridge process (Lin et al., 2010; Del Moral and Murray, 2015). However, this comes at the cost of learning two approximations and some accumulation of errors. Recent work by Baker et al. (2024) have extended our approach to directly approximate $\nabla\log h(t,x_{t})$, which is not necessarily the score of a probability density.

As our methodology allows one to employ any function approximator, we harness the flexibility of neural networks and the ease of implementation using modern software to approximate score functions. We adopt the parameterization of the score approximation in Section 2.2 with $v_{\phi}$ defined by a neural network. The choice of neural network and stochastic gradient algorithm is detailed in Section D.2 of the Supplementary Material. Optimization times ranged between several seconds to a few minutes on a simple desktop machine and can be reduced with hardware accelerators. As such computational overheads are marginal when deploying proposal bridge processes within an independent Metropolis–Hastings algorithm with many iterations or an importance sampler with many samples, we focus on assessing the quality of our proposals in settings where existing proposal methods are unsatisfactory. The performance measures considered are the acceptance rate of independent Metropolis–Hastings, and in the case of an importance sampling estimator $\hat{p}(T,x_{T}\mid 0,x_{0})$ of $p(T,x_{T}\mid 0,x_{0})$, the effective sample size proportion, and the variance var$\{\log\hat{p}(T,x_{T}\mid 0,x_{0})\}$ or the mean squared error ${E}[\{\log\hat{p}(T,x_{T}\mid 0,x_{0})-\log p(T,x_{T}\mid 0,x_{0})\}^{2}]$ when the true transition density is known. More implementation details are described in Section D of the Supplementary Material. These measures were computed using $1024$ samples or iterations, and $100$ independent repetitions of each method. We benchmark our approximations of the backward diffusion bridge (BDB) and forward diffusion bridge (FDB) in (4) and (9) against the Clark–Delyon–Hu (CDH) proposal bridge studied by Clark (1990) and Delyon and Hu (2006), the forward diffusion (FD) of Pedersen (1995), the guided proposal bridge (GPB) by Schauer et al. (2017), and the modified diffusion bridge (MDB) of Durham and Gallant (2002). Given the difficulty of comparing the wide range of methods for diffusion bridges in a completely fair manner, as their strengths and weaknesses can depend on the specificities of the problem under consideration, we note that our objective is merely to illustrate a new avenue to improve the construction of proposal bridge processes. A Python package to reproduce all numerical results is available online²²2https://github.com/jeremyhengjm/DiffusionBridge.

Consider (1) with linear drift function $f(t,x_{t})=-2x_{t}$ and identity diffusion coefficient $\sigma(t,x_{t})=I_{d}$. The transition density and score function are explicitly known and used as ground truth. The first and second rows of Fig. 1 illustrates the impact of the time horizon $T$ and dimension $d$ on algorithmic performance with the constraints $x_{0}=x_{T}=(1,\ldots,1)$. To study how performance degrades when we condition on rare states under the diffusion process $\eqref{eqn:SDE}$, in the third row of Fig. 1, we set the initial state as $x_{0}=1$ and vary where the terminal state $x_{T}$ is in the tails of the transition density $p(T,x_{T}\mid 0,x_{0})$ in multiples of its standard deviation.

We omit the guided proposal bridge in this example as it performs perfectly when its auxiliary process is the Ornstein–Uhlenbeck process. Our proposed methods offer substantial improvements over other methods without exploiting the correct parameterization. When comparing our forward and backward diffusion bridge approximations, we notice some accumulation of error, which is to be expected as the forward process is constructed using an additional score approximation. Given this observation, we will only consider the backward process in the following.

We consider a special case of an interest rates model in Aït-Sahalia and Lo (1998), defined by (1) with $f(t,x_{t})=4/x_{t}-x_{t}$ and $\sigma(t,x_{t})=1$. This specification admits a tractable transition density and score function as ground truth, given in terms of modified Bessel functions and its derivative. For each $T$, we first learn a single score approximation to handle multiple conditioned states $(x_{0},x_{T})$ by minimizing the loss in (8) averaged over initial states $X_{0}=x_{0}$ from the gamma distribution with shape $5$ and rate $2$. We construct guided proposal bridges based on a first-order Taylor approximation of $f$ at the stable stationary point $x^{\star}=2$. We plot our score approximation error for $T=1$ in Fig. 2, and examine how algorithmic performance depends on $T$ and $(x_{0},x_{T})$ in Fig. 3. Our backward diffusion bridge approximation performs well for all $T$ considered and when conditioning requires the process to move away from the stationary point (first and third columns). Its performance is similar to guided proposal bridge and significantly better than other existing methods. While the performance of modified diffusion bridge typically degrades with $T$, the results of forward diffusion and Clark–Delyon–Hu bridge depend on the specific conditioned states.

We end with a cell differentiation and development model from Wang et al. (2011). Cellular expression levels $X_{t}=(X_{t,1},X_{t,2})$ of two genes are modelled by (1) with $f=(f_{1},f_{2})$, where $f_{i}(t,x_{t})=x_{t,i}^{4}/(2^{-4}+x_{t,i}^{4})+2^{-4}/(2^{-4}+x_{t,j}^{4})-x_{t,i}$, for $(i,j)\in\{(1,2),(2,1)\}$, describe self-activation, mutual inhibition and inactivation respectively, and $\sigma(t,x_{t})=\sigma_{X}I_{2}$ captures intrinsic or external fluctuations. We consider the cell development from the undifferentiated state of $x_{0}=(1,1)$ to a differentiated state $x_{T}=(x_{T,1},x_{T,2})$ defined by another stable fixed point of $f$ satisfying $x_{T,1}<x_{T,2}$. We employ an auxiliary Ornstein–Uhlenbeck process in the guided proposal bridge and optimize its parameters using the Kullback–Leibler divergence considered in Schauer et al. (2017). Fig. 4 displays the performance of all proposal methods to approximate this diffusion bridge. For this application, we observe improvement over guided proposal bridge when the diffusion coefficient $\sigma_{X}$ is smaller, and significant improvement over other methods when the time horizon is long. These findings are consistent with numerical results obtained in Baker et al. (2024) for the same model.