Conditional sampling within generative diffusion models

Recall that we aim to sample $\pi(\cdot{\;|\;}y)$, and assume that we can sample the joint $\pi_{X,Y}$. Denote by $q_{t{\;|\;}s}(\cdot{\;|\;}v)$ the conditional distribution of $U(t)$ conditioned on $U(s)=v$ for $t>s$. The idea is now to construct the reversal in such a way that

for all $x\in\mathbb{R}^{d}$. If we in addition assume that the dimensions of $X$ and $Y$ match, then an immediate construction that satisfies the requirement above is

where the terminal and initial jointly follow $\pi_{X,Y}$ by marginalising out the intermediate path. Thus, sampling $\pi(\cdot{\;|\;}y)$ simplifies to simulating the SDE in Equation (6) with initial value $U(0)=y$. Directly finding such an SDE is hard, but we can leverage Anderson’s construction to form it as the reversal of a forward equation

Now we have a few handy methods to explicitly construct the forward equation above. One notable example is via Doob’s $h$-transform (Rogers and Williams, 2000) by choosing $X(0)=X$ and

where $\mu$ is the drift function of another (arbitrary) reference SDE

and the $h$-function $h_{T}(y,x,t)\coloneqq p_{Z(T){\;|\;}Z(t)}(y{\;|\;}x)$ is the conditional density function of $Z(T)$ with condition at $t$. With this choice, we immediately have $X(T)=Y$, and thus the desired $\bigl{(}X(0),X(T)\bigr{)}\sim\pi_{X,Y}$ is satisfied. However, we have to remark a caveat that the process $t\mapsto X(t)$ in Equation (7) marginally is not a Markov process, since its SDE now invokes $Y$ which is correlated with the process across time. But on the other hand, this problem can be solved by an extension $\mathop{}\!\mathrm{d}Y(t)=0$, $Y(0)=Y$, such that $(X(t),Y(t))$ is jointly Markov. We can then indeed learn the reversal in Equation (6) by Anderson’s construction. This approach was exploited by Somnath et al. (2023) who coined the term “aligned Schrödinger bridge”. But we note that the approach does not solve the DSB problem in Equation (4): In DSB we aim to find the optimal coupling under a path constraint, but in this approach, the coupling is already fixed by the given $\pi_{X,Y}$.

The generative conditional sampling with Doob’s bridging is summarised as follows.

From the algorithm above we see that the approach is as straightforward as the standard generative sampling for unconditional distributions, with an additional involvement of $Y$ in the forward and reverse equations; It adds almost no computational costs. Moreover, the algorithm does not assume $T\to\infty$ unlike many other generative samplers. However, the algorithm suffers from a few non-trivial problems due to the use of Doob’s transform to exactly pin two points. First, Doob’s $h$-function is mostly not available in closed form when the reference process is nonlinear. Although it is still possible to simulate Equation (7) without knowing $h_{T}$ explicitly (Schauer et al., 2017), and to approximate $h_{T}$ (Bågmark et al., 2022; Baker et al., 2024). This hinders the use of denoising score matching which is a computationally efficient estimator, as the conditional density $p_{X(t){\;|\;}X(s)}$ of the forward equation is consequently intractable. Second, the function $h_{T}(\cdot,\cdot,T)$ at time $T$ is not defined, and moreover, the forward drift is rather stiff around that time. This induces numerical errors and instabilities that are hard to eliminate. Finally, the approach assumes that the dimensions of $X$ and $Y$ are the same, which limits the range of applications. Recently, Denker et al. (2024) generalise the method by exploiting the likelihood $\pi_{Y{\;|\;}X}$ in such a way that the two points do not have to be exactly pinned. This can potentially solve the problems mentioned above.

Recall again that we aim to find a pair of forward and reversal SDEs such that the reversal satisfies Equation (5). In the previous section, we have exploited Doob’s transform for the purpose, but the approach comes with several problems due to the use of Doob’s $h$-transform for exactly pinning the data random variables $X$ and $Y$. In this section, we present a generalised forward-reverse SDE that does not require exact pinning.

The core of the idea is to introduce an auxiliary variable $V\in\mathbb{R}^{d_{y}}$ in the joint $\bigl{(}U(0),U(T),V\bigr{)}$, such that it facilities sampling $q_{T,0{\;|\;}Y}(\cdot{\;|\;}y)$, and that the marginal at $T$ satisfies $q_{T{\;|\;}V}(\cdot{\;|\;}y)=\pi(\cdot{\;|\;}y)$ (cf. Equation (5)). More precisely, the reversal that we desire is

starting at any reference $\pi_{\mathrm{ref}}$, so that $\int q_{T{\;|\;}0,V}(\cdot{\;|\;}u,y)\pi_{\mathrm{ref}}(\mathop{}\!\mathrm{d}u{\;|\;}y)=q_{T{\;|\;}V}(\cdot{\;|\;}y)=\pi(\cdot{\;|\;}y)$. With this reversal, if we can sample $U(0)\sim\pi_{\mathrm{ref}}(\cdot{\;|\;}y)$, then consequently $U(T)\sim\pi(\cdot{\;|\;}y)$ throughout the SDE. Define the forward equation that is paired with the reversal by

where the reference measure $\pi_{\mathrm{ref}}=p_{T,Y}$ is the joint distribution of $(X(T),Y)$. The Doob’s bridging is a special case, in the way that we have coerced $X(T)=Y$ as the reference. Also note that we can extend Equations (9) and (10) to Markov processes with $\mathop{}\!\mathrm{d}Y(t)=0$ in the same way as in Doob’s bridging. To ease later discussions, let us first summarise the conditional sampling by the joint bridging method as follows.

The remaining question is how to identify the generative sampler to use in Algorithm 3.2. Now looking back at Equations (1) and (2), we see that the algorithm is essentially an unconditional generative sampler that operates on the joint space of $X$ and $Y$. Hence, we can straightforwardly choose either Anderson’s construction or DSB to obtain the required reversal in Equation (9). With Anderson’s approach, the reversal’s drift function becomes

where we see that the marginal score in Equation (3) now turns into a conditional one. This conditional score function is the major hurdle of the construction, as the conditional score is often intractable. Song et al. (2021) utilise the decomposition $\operatorname{\nabla\!}_{u}\log p_{t{\;|\;}Y}(u{\;|\;}v)=\operatorname{\nabla\!}_{u}\log p_{t}(u)+\operatorname{\nabla\!}_{u}\log p_{Y{\;|\;}t}(v{\;|\;}u)$, where the two terms on the right-hand side can be more accessible. Take image classification for example, where $X$ and $Y$ stand for the image and label, respectively. The generative score $\operatorname{\nabla\!}_{u}\log p_{t}(u)$ is learnt via standard score matching, while the likelihood score $\operatorname{\nabla\!}_{u}\log p_{Y{\;|\;}t}(v{\;|\;}u)$ can be estimated by training an image classifier with categorical $p_{Y{\;|\;}t}$ for all $t$. However, this is highly application dependent, and developing a generic and efficient estimator for the conditional score remains an active research question (see, e.g., approximations in Chung et al., 2023; Song et al., 2023).

Another problem of using Anderson’s construction is again the common assumption $T\to\infty$, as we need to sample the conditional reference $p_{T{\;|\;}V}=\pi_{\mathrm{ref}}(\cdot{\;|\;}y)$. For example, Luo et al. (2023) choose $a(x,Y,t)=\theta\,(Y-x)$, so that $\pi_{\mathrm{ref}}(\cdot{\;|\;}y)\approx\mathrm{N}(\cdot{\;|\;}y,\sigma^{2})$ is approximately a stationary Gaussian. Under this choice, the resulting forward equation is akin to that of Doob’s bridging but instead pins between $X$ and a Gaussian-noised $Y$.

As an alternative to Anderson’s construction, we can employ DSB to search for a pair of Equations (9) and (10) that exactly bridges between $\pi_{X,Y}$ and a given $\pi_{\mathrm{ref}}$, within a finite horizon $T$. Formally, we are looking for a process

whose path measure $\mathbb{P}$ minimises $\mathrm{KL}(\cdot\,\|\,\mathbb{Q})$, where $\mathbb{Q}$ is the measure of a reference process

Due to the fact that $\pi_{\mathrm{ref}}$ is now allowed to arbitrary, we can choose a $\pi_{\mathrm{ref}}$ such that $\pi_{\mathrm{ref}}(\cdot{\;|\;}y)$ is a reasonable approximation to $\pi(\cdot{\;|\;}y)$ that is easy to sample from (e.g., Laplace approximation).

Currently there are two (asymptotic) numerical approaches to solve for the (conditional) DSB in the two equations above. The approach developed by De Bortoli et al. (2021) aims to find a sequence $\mathbb{P}^{i}$ that converges to $\mathbb{P}$ as $i\to\infty$ via half-bridges

starting at $\mathbb{P}^{0}\coloneqq\mathbb{Q}$. The solution of each half-bridge can be approximated by conventional parameter estimation methods in SDEs. As an example, in Equation (15), suppose that $\mathbb{P}^{2\,i-1}$ is obtained and that we can sample from it, then we can approximate $\mathbb{P}^{2\,i}$ by parametrising the SDE in Equation (12) and then estimate its parameters via drift matching between $\mathbb{P}^{2\,i}$ and $\mathbb{P}^{2\,i-1}$. However, a downside of this approach is that the solution does not exactly bridge $\pi_{X,Y}$ and $\pi_{\mathrm{ref}}$ at any $i$ albeit the asymptoticity: it is either $\pi_{X,Y}$ or $\pi_{\mathrm{ref}}$ that is preserved, not the both. This problem is solved by Shi et al. (2023). They utilise the decomposition $\mathbb{P}=\pi^{\star}\,\mathbb{Q}_{{\;|\;}0,T}$, where $\pi^{\star}$ solves the optimal transport $\min_{\nu}\{\mathrm{KL}(\nu\,\|\,\mathbb{Q}_{0,T})\colon\nu_{0}=\pi_{X,Y},\nu_{T}=\pi_{\mathrm{ref}}\}$, and $\mathbb{Q}_{{\;|\;}0,T}$ stands for the measure of the reference process conditioning at times $0$ and $T$ (e.g., Doob’s $h$-transform). The resulting algorithm is then alternating between

starting at $\mathbb{P}^{0}\coloneqq\widetilde{\pi}\,\mathbb{Q}_{{\;|\;}0,T}$, where $\widetilde{\pi}$ is any coupling of $\pi_{X,Y}$ and $\pi_{\mathrm{ref}}$, for instance, the trivial product measure $\pi_{X,Y}\times\pi_{\mathrm{ref}}$. In the iteration above, $\mathrm{proj}_{\mathcal{M}}\bigl{(}\mathbb{P}^{2\,i-2}\bigr{)}$ approximates $\mathbb{P}^{2\,i-2}$ as a Markov process, since the previous $\mathbb{P}^{2\,i-2}$ is not necessarily Markovian. Moreover, Shi et al. (2023) choose a particular projection $\mathrm{proj}_{\mathcal{M}}\bigl{(}\mathbb{P}^{2\,i-2}\bigr{)}=\operatorname*{arg\,min\,}_{\mathbb{P}\in\mathcal{M}}\mathrm{KL}\bigl{(}\mathbb{P}^{2\,i-2}\,\|\,\mathbb{P}\bigr{)}$, such that their time-marginal distributions match, and therefore preserving the marginals $\mathbb{P}^{i}_{0}=\pi_{X,Y}$ and $\mathbb{P}^{i}_{T}=\pi_{\mathrm{ref}}$ at any $i$. In essence, this algorithm builds a sequence on couplings $\mathbb{P}^{i}_{0,T}$ that converges to $\pi^{\star}$, in conjunction with Doob’s bridging. Although the method has to simulate Doob’s $h$-transform $\mathbb{Q}_{{\;|\;}0,T}$ which can be hard for complex reference processes, the cost is marginal in the context.

Compared to the iteration in Equation (15), this projection-bridging method is particularly useful in the generative diffusion scenario, since we mainly care about the marginal distribution and not so much about the actual interpolation path. That said, at any iteration $i$, the solution, is already a valid bridge between $\pi_{X,Y}$ and $\pi_{\mathrm{ref}}$ at disposal for generative sampling, even though it is not yet a solution to the Schrödinger bridge.

The previous joint bridging method (including Doob’s bridging as a special case) aims to find a reversal such that the conditional terminal is $q_{T{\;|\;}V}(\cdot{\;|\;}y)=\pi(\cdot{\;|\;}y)$. In essence, it thinks of sampling the target as simulating a reverse SDE. In this section, we show another useful insight, by viewing sampling the target as simulating a stochastic filtering distribution (Corenflos et al., 2024; Dou and Song, 2024; Trippe et al., 2023). Formally, this approach sets base on

and the core is the identity

where the integrand $p_{X(0){\;|\;}Y(0),Y_{(0,T]},X(T)}(\cdot{\;|\;}y,y_{(0,T]},x_{T})$ is the filtering distribution (in the reverse-time direction) with initial value $x_{T}$ and measurement $y_{[0,T]}\coloneqq[y,y_{(0,T]}]$. The identity describes an ancestral sampling that, for any given condition $y$, we first sample $\bigl{(}X(T),Y_{(0,T]}\bigr{)}\sim p_{X(T),Y_{(0,T]}{\;|\;}Y(0)}(\cdot{\;|\;}y)$ and then use it to solve the filtering problem of the reversal, hence the name “forward-backward path bridging”. However, since these two steps cannot be computed exactly, we show how to approximate them in practice in the following.

The distribution $p_{X(T),Y_{(0,T]}{\;|\;}Y(0)}(\cdot{\;|\;}y)$ is in general not tractable. In Dou and Song (2024) and Trippe et al. (2023) they make assumptions so that they can facilitate an approximation

More specifically, they assume that 1) the SDE in Equation (17) is separable (i.e., $a_{1}$ does not depend on $Y$ while $a_{2}$ does not depend on $X$), so that they can independently simulate $p_{Y_{(0,T]}{\;|\;}Y(0)}(\cdot{\;|\;}y)$ leaving away the part that depends on $X$, and that 2) the SDE is stationary enough for large $T$ such that $p_{X(T){\;|\;}Y_{[0,T]}}\approx p_{X(T){\;|\;}Y(T)}$. Furthermore, the (continuous-time) filtering distribution $p_{X(0){\;|\;}Y(0),Y_{(0,T]},X(T)}$ is intractable too. In practice we apply sequential Monte Carlo (SMC, Chopin and Papaspiliopoulos, 2020) to approximately sample from it. This induces error due to 3) using a finite number of particles. Since continuous-time particle filtering (Bain and Crisan, 2009) is particularly hard in the generative context, we often have to 4) discretise the SDE beforehand, which again introduce errors, though simulating SDEs without discretisation is possible (Beskos and Roberts, 2005; Blanchet and Zhang, 2020). These four points constitute the main sources of errors of this method.

The work in Corenflos et al. (2024) eliminates all the errors except for that of the SDE discretisation. The idea is based on extending the identity in Equation (18) as an invariant Markov chain Monte Carlo (MCMC) kernel $\mathcal{K}$ by

where we see that the kernel is indeed $\pi(\cdot{\;|\;}y)$-variant: $\mathcal{K}(p_{X(0){\;|\;}Y(0)})=p_{X(0){\;|\;}Y(0)}=\pi_{X{\;|\;}Y}$. Moreover, we can replace the filtering distribution with another MCMC kernel by, for instance, gradient Metropolis–Hasting (Titsias and Papaspiliopoulos, 2018), or more suitable in this context, the conditional sequential Monte Carlo (CSMC, Andrieu et al., 2010). The algorithm finally works as a Metropolis-within-Gibbs sampler as follows.

The Gibbs filtering-based conditional sampler above has significantly less errors compared to Trippe et al. (2023) and Dou and Song (2024) as empirically shown in Corenflos et al. (2024). In a sense, the Gibbs sampler has transformed the two major errors due to $T\to\infty$ and using a finite number of particles, to the increased autocorrelation in Algorithm 3.3. When $T$ is small, the samples are more correlated, but still statistically exact nonetheless. Moreover, the sampler does not assume the separability of the diffusion model, meaning that in the algorithm we are free to use the reversal construction of any kind (e.g., Anderson or DSB). On the other hand, Trippe et al. (2023) and Dou and Song (2024)’s approach does not instantly apply on DSBs which are rarely separable. The cost we pay for using the Gibbs sampler is mainly the implementation and statistical efficiency of the filtering MCMC in 3 of Algorithm 3.3. It is also an open question for simulating the filtering MCMC in continuous-time without discretising the SDE.

Let us assume that we are given a pair of forward-backward diffusion models (using arbitrary reversal construction) that targets $\pi_{X}$ as in Equations (1) and (2). For simplicity, let us change to work with the models at discrete times $k=0,1,2,\ldots,N$ via

where $p_{0}(\cdot)=\pi_{X}(\cdot)=q_{N}(\cdot)$, and $q_{0}(\cdot)=\pi_{\mathrm{ref}}(\cdot)=p_{N}(\cdot)$. Since we can sample $\pi_{X}$ and evaluate the likelihood $\pi(y{\;|\;}\cdot)$, it is natural to think of using importance sampling to sample $\pi(\cdot{\;|\;}y)$. That is, if $\{X^{j}\}_{j=1}^{J}\sim\pi_{X}$ then $\{(w_{j},X^{j})\}_{j=1}^{J}\sim\pi(\cdot{\;|\;}y)$ with weight $w_{j}=\pi(y{\;|\;}X^{j})$. However, the weights will barely be informative, as the prior samples are very unlikely samples from the posterior distribution in generative applications. Hence, in practice we generalise importance sampling within Feynman–Kac models (Chopin and Papaspiliopoulos, 2020) allowing us to effectively sample the target with sequential Monte Carlo (SMC) samplers. Formally, we define the generative Feynman–Kac model by

where $\{M_{k{\;|\;}k-1}\}_{k=0}^{N}$ and $\{G_{k}\}_{k=0}^{N}$ are Markov kernels and potential functions, respectively, and $\ell_{N}$ is the marginal likelihood. Instead of applying naive importance sampling, we apply SMC to simulate the model, crucially, the target $Q_{N}$, as in the following algorithm.

Given only the marginal constraints (e.g., at $Q_{N}$ and $Q_{0}$), the generative Feynman–Kac model is not unique. A trivial example is by letting $M_{0}=q_{0}$, $M_{k{\;|\;}k-1}=q_{k{\;|\;}k-1}$, $G_{0}=G_{1}=\cdots=G_{N-1}\equiv 1$, and $G_{N}(u_{k},u_{k-1})=\pi(y{\;|\;}u_{k})$, where Algorithm 4.4 recovers naive importance sampling which is not useful in reality. The purpose is thus to design $M$ and $G$ so that $Q_{k}$ continuously anneals to $\pi(\cdot{\;|\;}y)$ for $k=0,1,\ldots,N$ while keeping the importance samples in Algorithm 4.4 effective. Wu et al. (2023) and Janati et al. (2024) show a particularly useful system of Markov kernels and potentials by setting

where $l_{k}(y,u_{k})\coloneqq\int\pi(y{\;|\;}u_{N})\,q_{N{\;|\;}k}(u_{N}{\;|\;}u_{k})\mathop{}\!\mathrm{d}u_{N}$, and define $l_{N}(y,u_{N})\coloneqq\pi(y{\;|\;}u_{N})$ that reduces to the target’s likelihood. Indeed, if we substitute Equation (22) back into Equation (21), the marginal constraint $Q_{N}$ is satisfied. This choice has at least two advantages that makes it valuable in the generative diffusion context. First, this construction is locally optimal, in the sense that if we also fix $G_{k}(u_{k},u_{k-1})\equiv 1$ then $M_{k{\;|\;}k-1}(u_{k}{\;|\;}u_{k-1})=q_{k{\;|\;}k-1,Y}(u_{k}{\;|\;}u_{k-1},y)\propto l_{k}(y,u_{k})\,q_{k{\;|\;}k-1}(u_{k}{\;|\;}u_{k-1})$ which is exactly the Markov transition of the conditional reversal in Equation (11). Consequently, all samples in Algorithm 4.4 are equally weighted, and hence the SMC sampler is perfect. However, simulating the optimal Markov kernel is hard, as $l_{k}(y,u_{k})$ is intractable (except when $k=N$). This gives rise to another advantage: even if we replace the exact $l_{k}$ by any approximation $\widehat{l}_{k}$ for $k=0,1,\ldots,N-1$, the Feynman–Kac model remains valid. Namely, the weighted samples of Algorithm 4.4 converge to $\pi(\cdot{\;|\;}y)$ in distribution as $J\to\infty$ even if $\{l_{k}\}_{k=0}^{N-1}$ is approximate. However, the quality of the approximation $\widehat{l}_{k}$ significantly impacts the effective sample size of the algorithm, and recently there have been numerous approximations proposed.

Following Equation (22), Wu et al. (2023) choose to approximate $l_{k}$ by linearisation, and let the Markov kernel be that of a Langevin dynamic that approximately leaves $q_{k{\;|\;}k-1,Y}$ invariant. More specifically, they choose

where $m_{N}(u_{k})$ stands for any approximate sample from $q_{N{\;|\;}k}(\cdot{\;|\;}u_{k})$, and $\delta_{k}$ is the Langevin step size. Originally, Wu et al. (2023) use $m_{N}(u_{k})$ as a sample of the reversal at $N$ starting at $u_{k}$, which often experiences high variance, but it can be improved by using $m_{N}(u_{k})$ as the conditional mean of $q_{N{\;|\;}k}(\cdot{\;|\;}u_{k})$ via Tweedie’s formula (Song et al., 2023). See also Cardoso et al. (2024) for an alternative construction of the proposal. While Wu et al. (2023)’s construction of the Feynman–Kac model is empirically successful in a number of applications, its main problem lies in its demanding computation. In particular, the proposal asks to differentiate through $\widehat{l}_{k}$ usually containing a denoising neural network.

Janati et al. (2024) follow up the construction in Equation (22) and propose an efficient sampler based on partitioning the Feynman–Kac model into several contiguous sub-models, where each sub-model targets its next. As an example, when we use two blocks, $Q_{0:N}$ is divided into $Q^{(1)}_{0:j}$ and $Q^{(2)}_{j:N}$ for some $0\leq j\leq N$, where $Q^{(1)}_{0:j}$ starts at $q_{0}$ and targets $Q_{j}$, whereas $Q^{(2)}_{j:N}$ starts at $Q_{j}$ and targets $q_{N}$. Although the division looks trivial at a first glance, the catch is that within each sub-model we can form more efficient approximations to the Feynman–Kac Markov transition. More precisely, recall that in Equation (23), $\widehat{l}_{k}(y,u_{k})$ aims to approximate $\int\pi(y{\;|\;}u_{N})\,q_{N{\;|\;}k}(u_{N}{\;|\;}u_{k})\mathop{}\!\mathrm{d}u_{N}$, the accuracy of which depends on the distance $N-k$. Now inside the sub-model $Q^{(1)}_{0:j}$, the terminal is $j$ instead of $N$, where $j-k$ is mostly smaller than $N-k$ for $0\leq k\leq j$. In addition, instead of running the SMC in Algorithm 4.4, they also propose a variational approximation to the Markov transition of the Feynman–Kac model in order to directly sample the model. Note that although Janati et al. (2024) set roots on linear Gaussian likelihood $\pi(y{\;|\;}\cdot)$, their method can straightforwardly generalise to other likelihoods as well, with suitable Gaussian quadratures applied.

A computationally simpler bootstrap construction of the generative Feynman–Kac model is

where we simplify the Markov kernel by that of the unconditional reversal which is easier compared to $\widehat{q}_{k}(u_{k}{\;|\;}u_{k-1},y)$ in Equation (23). Additionally, the approximate $\widehat{l}_{k}$ is also simplified as the target likelihood evaluated at $u_{k}$ and a suitably scaled $\lambda_{k}\,y$ with $\lambda_{N}=1$ justified by the heuristics in Janati et al. (2024, Sec. 4.1). However, since the Markov proposal is farther from the optimal one, the resulting algorithm is statistically less efficient. When the problem of interests is not difficult, this bootstrap model may indeed be useful in practice (see our example in Section 5).