SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Luhuan Wu*
lw2827@columbia.edu
Brian L. Trippe*
blt2114@columbia.edu
Christian A. Naesseth
c.a.naesseth@uva.nl
David M. Blei
david.blei@columbia.edu
John P. Cunningham
jpc2181@columbia.edu
Diffusion models have been successful on a range of conditional generation tasks including molecular design and text-to-image generation. However, these achievements have primarily depended on task-specific conditional training or error-prone heuristic approximations. Ideally, a conditional generation method should provide exact samples for a broad range of conditional distributions without requiring task-specific training. To this end, we introduce the Twisted Diffusion Sampler, or TDS. TDS is a sequential Monte Carlo (SMC) algorithm that targets the conditional distributions of diffusion models through simulating a set of weighted particles. The main idea is to use twisting, an SMC technique that enjoys good computational efficiency, to incorporate heuristic approximations without compromising asymptotic exactness. We first find in simulation and in conditional image generation tasks that TDS provides a computational statistical trade-off, yielding more accurate approximations with many particles but with empirical improvements over heuristics with as few as two particles. We then turn to motif-scaffolding, a core task in protein design, using a TDS extension to Riemannian diffusion models; on benchmark tasks, TDS allows flexible conditioning criteria and often outperforms the state-of-the-art, conditionally trained model.[2]
Conditional sampling is an essential primitive in the machine learning toolkit. One begins with a generative model that parameterizes a distribution  $p_{\theta}(x)$  on data  $x$ , and then augments the model to include information  $y$  in a joint distribution  $p_{\theta}(x,y) = p_{\theta}(x)p(y|x)$ . This joint distribution then implies a conditional distribution  $p_{\theta}(x|y)$ , from which desired outputs are sampled.
For example, in protein design,  $p_{\theta}(x)$  can represent a distribution of physically realizable protein structures,  $y$  a substructure that imparts a desired biochemical function, and samples from  $p_{\theta}(x \mid y)$  are then physically realizable structures that contain the substructure of interest [e.g. 34, 37].
Diffusion models are a class of generative models that have demonstrated success in conditional generation tasks [18, 27, 33, 37]. They parameterize distributions  $p_{\theta}(x)$  through an iterative refinement process that builds up data from noise. When a diffusion model is used for conditional generation, this refinement process is modified to account for conditioning at each step [9, 19, 27, 33].
*Equal contribution, order by coin flip.
$^{2}$ Code: https://github.com/blt2114/twisted_diffusionSampler
One approach is to incorporate the conditioning information into training [e.g. 9, 27]; one either modifies the unconditional model to take  $y$  as input or trains a separate conditional model to predict  $y$  from partially noised inputs. However, conditional training requires (i) assembling a large set of paired examples of the data and conditioning information  $(x, y)$ , and (ii) designing and training a task-specific model when adapting to new conditioning tasks. For example, image inpainting and class-conditional image generation can be both formalized as conditional sampling problems based on the same (unconditional) image distribution  $p_{\theta}(x)$ ; however, the conditional training approach requires training separate models on two curated sets of conditioning inputs.
To avoid conditional training, a separate line of work uses heuristic approximations that directly operate on unconditional diffusion models: once an unconditional model  $p_{\theta}(x)$  is trained, it can be flexibly combined with various conditioning criteria to generate customized outputs. These approaches have been applied to inpainting problems [24, 33, 34], and other inverse problems [1, 6, 19, 32]. But it is unclear how well these heuristics approximate the exact conditional distributions they are designed to mimic; for example on inpainting tasks they often fail to return outputs consistent with both the conditioning information and unconditional model [40]. These concerns are critical in domains that require accurate conditionals. In molecular design, for example, even a small approximation error could result in atomic structures with chemically implausible bond distances.
This paper develops a practical and asymptotically exact method for conditional sampling from an unconditional diffusion model. We use sequential Monte Carlo (SMC), a general tool for asymptotically exact inference in sequential probabilistic models [5, 10, 25]. SMC simulates an ensemble of weighted trajectories, or particles, through a sequence of proposals and weighting mechanisms. These weighted particles then form an asymptotically exact approximation to a desired target distribution.
The premise of this work is to recognize that the sequential structure of diffusion models permits the application of SMC for sampling from conditional distributions  $p_{\theta}(x \mid y)$ . We design an SMC algorithm that leverages twisting, an SMC technique that modifies proposals and weighting schemes to approach the optimal choices [15, 38]. While optimal twisting is intractable, we effectively approximate it with recent heuristic approaches to conditional sampling [eg. 19], and correct the errors by the weighting mechanisms. The resulting algorithm maintains asymptotic exactness to  $p_{\theta}(x \mid y)$ , and empirically it can outperform heuristics alone even with just two particles.
We summarize our contributions: (i) We propose a practical SMC algorithm, Twisted Diffusion Sampler or TDS, for asymptotically exact conditional sampling from diffusion models; (ii) We show that TDS applies to a range of conditional generation problems, and extends to Riemannian manifold diffusion models; (iii) On MNIST inpainting and class-conditional generation tasks we demonstrate TDS's empirical improvements beyond heuristic approaches; and (iv) On protein motif-scaffolding problems with short scaffolds TDS provides greater flexibility and achieves higher success rates than the state-of-the-art conditionally trained model.
2 Background: Diffusion models and sequential Monte Carlo
Diffusion models. A diffusion model generates a data point  $x^0$  by iteratively refining a sequence of noisy data  $x^t$ , starting from pure noise  $x^T$ . This procedure parameterizes a distribution of  $x^0$  as the marginal of a length  $T$  Markov chain
$$
p _ {\theta} \left(x ^ {0}\right) = \int p \left(x ^ {T}\right) \prod_ {t = 1} ^ {T} p _ {\theta} \left(x ^ {t - 1} \mid x ^ {t}\right) d x ^ {1: T}, \tag {1}
$$
where  $p(x^T)$  is an easy-to-sample noise distribution, and each  $p_{\theta}(x^{t - 1} \mid x^t)$  is the transition distribution defined by the  $(T - t)^{\mathrm{th}}$  refinement step.
Diffusion models  $p_{\theta}$  are fitted to match a data distribution  $q(x^0)$  from which we have samples. To achieve this goal, a forward process  $q(x^0) \prod_{t=1}^{T} q(x^t \mid x^{t-1})$  is set to gradually add noise to the data, where  $q(x^t \mid x^{t-1}) = \mathcal{N}(x^t; x^{t-1}, \sigma^2)$ , and  $\sigma^2$  is a positive variance. To fit a diffusion model, one finds  $\theta$  such that  $p_{\theta}(x^{t-1} \mid x^t) \approx q(x^{t-1} \mid x^t)$ , which is the reverse conditional of the forward process. If this approximation is accomplished for all  $t$ , and if  $T\sigma^2$  is big enough that  $q(x^T)$  may be approximated as  $q(x^T) \approx \mathcal{N}(0, T\sigma^2) \eqqcolon p(x^T)$ , then we will have  $p_{\theta}(x^0) \approx q(x^0)$ .
In particular, when  $\sigma^2$  is small enough then the reverse conditionals of  $q$  are approximately Gaussian,
$$
q \left(x ^ {t - 1} \mid x ^ {t}\right) \approx \mathcal {N} \left(x ^ {t - 1}; x ^ {t} + \sigma^ {2} \nabla_ {x ^ {t}} \log q \left(x ^ {t}\right), \sigma^ {2}\right), \tag {2}
$$
where  $q(x^{t}) = \int q(x^{0})q(x^{t}\mid x^{0})dx^{0}$  and  $\nabla_{x^t}\log q(x^t)$  is known as the (Stein) score [33]. To mirror eq. (2), diffusion models parameterize  $p_{\theta}(x^{t - 1}\mid x^t)$  via a score network  $s_\theta (x^t,t)$
$$
p _ {\theta} \left(x ^ {t - 1} \mid x ^ {t}\right) := \mathcal {N} \left(x ^ {t - 1}; x ^ {t} + \sigma^ {2} s _ {\theta} \left(x ^ {t}, t\right), \sigma^ {2}\right). \tag {3}
$$
When  $s_{\theta}(x^t, t)$  is trained to approximate  $\nabla_{x^t} \log q(x^t)$ , we have  $p_{\theta}(x^{t-1} \mid x^t) \approx q(x^{t-1} \mid x^t)$ .
Notably, approximating the score is equivalent to learning a denoising neural network  $\hat{x}_{\theta}(x^{t}, t)$  to approximate  $\mathbb{E}_q[x^0 \mid x^t]$ . The reason is that by Tweedie's formula [12, 28]  $\nabla_{x^t} \log q(x^t) = (\mathbb{E}_q[x^0 \mid x^t] - x^t) / t\sigma^2$  and one can set  $s_{\theta}(x^t, t) := (\hat{x}_{\theta}(x^t, t) - x^t) / t\sigma^2$ . The neural network  $\hat{x}_{\theta}(x^t, t)$  may be learned by denoising score matching (see [e.g. 18, 35]). For the remainder of paper we drop the argument  $t$  in  $\hat{x}_{\theta}$  and  $s_{\theta}$  when it is clear from context.
Appendix A generalizes the formulation above to diffusion process formulations that are commonly used in practice.
Sequential Monte Carlo. Sequential Monte Carlo (SMC) is a general tool to approximately sample from a sequence of distributions on variables  $x^{0:T}$ , terminating at a final target of interest [5, 10, 14, 25]. SMC approximates these targets by generating a collection of  $K$  particles  $\{x_k^t\}_{k=1}^K$  across  $T$  steps of an iterative procedure. The key ingredients are proposals  $r_T(x^T)$ ,  $\{r_t(x^t \mid x^{t+1})\}_{t=0}^{T-1}$  and weighting functions  $w_T(x^T)$ ,  $\{w_t(x^t, x^{t+1})\}_{t=0}^{T-1}$ . At the initial step  $T$ , one draws  $K$  particles of  $x_k^T \sim r_T(x^T)$  and sets  $w_k^T := w_T(x_k^T)$ , and sequentially repeats the following for  $t = T-1, \ldots, 0$ :
- resample  $\{x_{k}^{t + 1:T}\}_{k = 1}^{K}\sim \mathrm{Multinomial}(\{x_{k}^{t + 1:T}\}_{k = 1}^{K};\{w_{k}^{t + 1}\}_{k = 1}^{K})$
- propose  $x_{k}^{t} \sim r_{t}(x^{t} \mid x_{k}^{t + 1}), \quad k = 1, \dots, K$
- weight  $w_{k}^{t} \coloneqq w_{t}(x_{k}^{t}, x_{k}^{t + 1}), k = 1, \dots, K$
The proposals and weighting functions together define a sequence of intermediate target distributions,
$$
\nu_ {t} \left(x ^ {0: T}\right) := \frac {1}{\mathcal {L} _ {t}} \left[ r _ {T} \left(x ^ {T}\right) \prod_ {t ^ {\prime} = 0} ^ {T - 1} r _ {t ^ {\prime}} \left(x ^ {t ^ {\prime}} \mid x ^ {t ^ {\prime} + 1}\right) \right] \left[ w _ {T} \left(x ^ {T}\right) \prod_ {t ^ {\prime} = t} ^ {T - 1} w _ {t ^ {\prime}} \left(x ^ {t ^ {\prime}}, x ^ {t ^ {\prime} + 1}\right) \right] \tag {4}
$$
where  $\mathcal{L}_t$  is a normalization constant. A classic example that SMC applies is the state space model [25, Chapter 1] that describes a distribution over a sequence of latent states  $x^{0:T}$  and observations  $y^{0:T}$ . Each intermediate target  $\nu_t$  is constructed to be the posterior  $p(x^{0:T} \mid y^{t:T})$  given the first  $T - t + 1$  observations, and the final target  $p(x^{0:T} \mid y^{0:T})$  is the posterior given all observations.
The defining property of SMC is that the weighted particles at each  $t$  form discrete approximations  $(\sum_{k'=1}^{K} w_{k'}^t)^{-1} \sum_{k=1}^{K} w_t^k \delta_{x_k^t}(x^t)$  (where  $\delta$  is a Dirac measure) to  $\nu_t(x^t)$  that become arbitrarily accurate in the limit that many particles are used [5, Proposition 11.4]. So, by choosing  $r_t$  and  $w_t$  so that  $\nu_0(x^0)$  matches the desired distribution, one can guarantee arbitrarily low approximation error in the large compute limit.
3 Twisted Diffusion Sampler: SMC sampling for diffusion model conditionals
Consider conditioning information  $y$  associated with a given likelihood function  $p_{y|x^0}(y|x^0)$ . We embed  $y$  in a joint model over  $x^{0:T}$  and  $y$  as  $p_{\theta}(x^{0:T}, y) = p_{\theta}(x^{0:T})p_{y|x^0}(y|x^0)$ , where  $y$  and  $x^{1:T}$  are conditionally independent given  $x^0$ . Our goal is to sample from the conditional  $p_{\theta}(x^0 \mid y)$ .
In this section, we develop Twisted Diffusion Sampler (TDS), a practical SMC algorithm targeting  $p_{\theta}(x^0 \mid y)$ . First, we describe how the Markov structure of diffusion models permits a factorization of an extended conditional distribution to which SMC applies. Then, we show how a diffusion model's denoising predictions support the application of twisting, an SMC technique in which one uses proposals and weighting functions that approximate the "optimal" ones. Lastly, we extend TDS to certain " inpainting" problems where  $p_{y|x^0}(y|x^0)$  is not smooth, and to Riemannian diffusion models.
3.1 Conditional diffusion sampling as an SMC procedure
The Markov structure of the diffusion model permits a factorization that is recognizable as the final target of an SMC algorithm. We write the conditional distribution, extended to include  $x^{1:T}$ , as
$$
p _ {\theta} \left(x ^ {0: T} \mid y\right) = \frac {p _ {\theta} \left(x ^ {0 : T} , y\right)}{p _ {\theta} (y)} = \frac {1}{p _ {\theta} (y)} \left[ p \left(x ^ {T}\right) \prod_ {t = 0} ^ {T - 1} p _ {\theta} \left(x ^ {t} \mid x ^ {t + 1}\right) \right] p _ {y \mid x ^ {0}} \left(y \mid x ^ {0}\right) \tag {5}
$$
Algorithm 1: Twisted Diffusion Sampler (TDS)
1 for  $k = 1:K$    
2  $\begin{array}{rl} & x_k^T\sim p(x^T),w_k\gets \tilde{p}_k^T = \tilde{p}_\theta (y|x_k^T) / /\mathrm{initial~proposal~and~weight}\\ & \{x_{k}^{t + 1},\tilde{p}_{k}^{t + 1}\}_{k = 1}^{K}\sim \mathrm{Multinomial}\left(\{x_{k}^{t + 1},\tilde{p}_{k}^{t + 1}\}_{k = 1}^{K};\{w_{k}\}_{k = 1}^{K}\right) / /\mathrm{resample}\\ & \mathrm{for~}k = 1:K\\ & \tilde{s}_k = (\hat{x}_\theta (x_k^{t + 1}) - x_k^{t + 1}) / t\sigma^2 +\nabla_{x_k^{t + 1}}\log \tilde{p}_k^{t + 1} / /\mathrm{conditional~score~approx.}\\ & x_k^t\sim \tilde{p}_\theta (\cdot |x_k^{t + 1},y):= \mathcal{N}\left(x_k^{t + 1} + \sigma^2\tilde{s}_k,\tilde{\sigma}^2\right) / /\mathrm{proposal}\\ & \tilde{p}_k^t\gets \tilde{p}_\theta (y|x_k^t) / /\mathrm{twisting~function~(in~eqs.~(8),(13),~or~(14))}\\ & w_k\gets p_\theta (x_k^t |x_k^{t + 1})\tilde{p}_k^t /[ \tilde{p}_\theta (x_k^t |x_k^{t + 1},y)\tilde{p}_k^{t + 1}] / /\mathrm{weight} \end{array}$    
3   
4   
5   
6   
7   
8   
9
with the desired marginal,  $p_{\theta}(x^0 \mid y)$ . Comparing the diffusion conditional of eq. (5) to the SMC target of eq. (4) suggests SMC can be used.
For example, consider a first attempt at an SMC algorithm. Set the proposals as  $r_T(x^T) = p(x^T)$  and  $r_t(x^t \mid x^{t+1}) = p_\theta(x^t \mid x^{t+1})$  for  $1 \leq t \leq T$ , and weighting functions as  $w_T(x^T) = w_t(x^t, x^{t+1}) = 1$  for  $1 \leq t \leq T$  and  $w_0(x^0, x^1) = p_{y|x^0}(y|x^0)$ .
Substituting these choices into eq. (4) results in the desired final target  $\nu_{0} = p_{\theta}(x^{0:T}\mid y)$  with normalizing constant  $\mathcal{L}_0 = p_\theta (y)$ . As a result, the associated SMC algorithm produces a final set of  $K$  samples and weights  $\{x_k^0;w_k^0\}_{k = 1}^K$  that provides an asymptotically accurate approximation  $\mathbb{P}_K(x^0\mid y)\coloneqq (\sum_k^K w_k^0)^{-1}\sum_{k = 1}^K w_k^0\delta_{x_k^0}(x^0)$  to the desired  $p_{\theta}(x^{0}\mid y)$ .
The approach above is simply importance sampling with proposal  $p_{\theta}(x^{0:T})$ ; with all intermediate weights set to 1, one can skip resampling steps to reduce the variance of the procedure. Consequently, this approach will be impractical if  $p_{\theta}(x^0 \mid y)$  is too dissimilar from  $p_{\theta}(x^0)$  as only a small fraction of unconditional samples will have high likelihood: the number of particles required for accurate estimation of  $p_{\theta}(x^{0:T} \mid y)$  is exponential in KL  $[p_{\theta}(x^{0:T} \mid y) \| p_{\theta}(x^{0:T})]$  [3].
3.2 Twisted diffusion sampler
Twisting is a technique in the SMC literature intended to reduce the number of particles required for good approximation accuracy [15, 16]. Loosely, it introduces a sequence of twisting functions that modify the naive proposals and weighting functions, so that the resulting intermediate targets are closer to the final target of interests. We refer the reader to Naesseth et al. [25, Chapter 3.2] for background on twisting in SMC.
Optimal twisting. Consider defining the twisted proposals  $r_t^*$  by multiplying the naive proposals  $r_t$  described in Section 3.1 by  $p_\theta(y \mid x^t)$  as
$$
r _ {T} ^ {*} \left(x ^ {T}\right) \propto p \left(x ^ {T}\right) p _ {\theta} \left(y \mid x ^ {T}\right) \text {a n d} r _ {t} ^ {*} \left(x ^ {t} \mid x ^ {t + 1}\right) \propto p _ {\theta} \left(x ^ {t} \mid x ^ {t + 1}\right) p _ {\theta} \left(y \mid x ^ {t}\right) \quad \text {f o r} \quad 0 \leq t <   T. \tag {6}
$$
The factors  $p_{\theta}(y \mid x^t)$  are the optimal twisting functions because they permit an SMC sampler that draws exact samples from  $p_{\theta}(x^{0:T} \mid y)$  even when run with a single particle.
To see that a single particle is an exact sample, by Bayes rule, the proposals in eq. (6) reduce to
$$
r _ {T} ^ {*} \left(x ^ {T}\right) = p _ {\theta} \left(x ^ {T} \mid y\right) \text {a n d} r _ {t} ^ {*} \left(x ^ {t} \mid x ^ {t + 1}\right) = p _ {\theta} \left(x ^ {t} \mid x ^ {t + 1}, y\right) \quad \text {f o r} \quad 0 \leq t <   T. \tag {7}
$$
As a result, if one samples  $x^{T} \sim r_{T}^{*}(x^{T})$  and  $x^{t} \sim r_{t}^{*}(x^{t} \mid x^{t + 1})$  for  $t = T - 1, \ldots, 0$ , by (i) the law of total probability and (ii) the chain rule of probability, one obtains  $x^{0} \sim p_{\theta}(x^{0} \mid y)$  as desired.
However, we cannot readily sample from each  $r_t^*$  because  $p_{\theta}(y \mid x^t)$  is not analytically tractable. The challenge is that  $p_{\theta}(y \mid x^t) = \int p_{y|x^0}(y|x^0)p_{\theta}(x^0 \mid x^t)dx^0$  depends on  $x_t$  through  $p_{\theta}(x^0 \mid x^t)$ . The latter in turn requires marginalizing out  $x^1, \ldots, x^{t-1}$  from the joint density  $p_{\theta}(x^{0:t-1} \mid x^t)$ , whose form depends on  $t$  calls to the neural network  $\hat{x}_{\theta}$ .
Tractable twisting. To avoid this intractability, we approximate the optimal twisting functions by
$$
\tilde {p} _ {\theta} (y \mid x ^ {t}) := p _ {y \mid x ^ {0}} (y | \hat {x} _ {\theta} (x ^ {t})) \approx p _ {\theta} (y \mid x ^ {t}), \tag {8}
$$
which is the likelihood function evaluated at  $\hat{x}_{\theta}(x^{t})$ , the denoising estimate of  $x^0$  at step  $t$  from the diffusion model. This tractable twisting function  $\tilde{p}_{\theta}(y\mid x^{t})$  is the key ingredient needed to define the Twisted Diffusion Sampler (TDS, Algorithm 1). We motivate and develop its components below.
The approximation in eq. (8) offers two favorable properties. First,  $\tilde{p}_{\theta}(y \mid x^t)$  is computable because it depends on  $x^t$  only though one call to  $\hat{x}_{\theta}$ , instead of an intractable integral over many calls as in the case of optimal twisting. Second,  $\tilde{p}_{\theta}(y \mid x^t)$  becomes an increasingly accurate approximation of  $p_{\theta}(y \mid x^t)$ , as  $t$  decreases and  $p_{\theta}(x^0 \mid x^t)$  concentrates on  $\mathbb{E}_{p_{\theta}}[x^0 \mid x^t]$ , which  $\hat{x}_{\theta}$  is fit to approximate; at  $t = 0$ , where we can choose  $\hat{x}_{\theta}(x^0) = x^0$ , we obtain  $\tilde{p}_{\theta}(y \mid x^0) = p_{y|x^0}(y|x^0)$ .
We next use eq. (8) to develop a sequence of twisted proposals  $\tilde{p}_{\theta}(x^t \mid x^{t+1}, y)$ , to approximate the optimal proposals  $p_{\theta}(x^t \mid x^{t+1}, y)$  in eq. (7). Specifically, we define twisted proposals as
$$
\tilde {p} _ {\theta} \left(x ^ {t} \mid x ^ {t + 1}, y\right) := \mathcal {N} \left(x ^ {t}; x ^ {t + 1} + \sigma^ {2} s _ {\theta} \left(x ^ {t + 1}, y\right), \tilde {\sigma} ^ {2}\right), \tag {9}
$$
$$
\text {w h e r e} \quad s _ {\theta} \left(x ^ {t + 1}, y\right) := s _ {\theta} \left(x ^ {t + 1}\right) + \nabla_ {x ^ {t + 1}} \log \tilde {p} _ {\theta} (y \mid x ^ {t + 1}) \tag {10}
$$
is an approximation of the conditional score,  $s_{\theta}(x^{t + 1},y)\approx \nabla_{x^{t + 1}}\log p_{\theta}(x^{t + 1}\mid y)$ , and  $\tilde{\sigma}^2$  is the proposal variance. For simplicity one could choose  $\tilde{\sigma}^2 = \sigma^2$  to match the variance of  $p_{\theta}(x^t\mid x^{t + 1})$ .
Equation (9) builds on previous works (e.g. [33, 31]) that seek to approximate the reversal of a conditional forward process. The gradient in eq. (10) is computed by back-propagating through  $\hat{x}_{\theta}(x^{t + 1})$ . We further discuss this technique, which has been used before [e.g. in 19, 32], in Section 4.
Twisted targets and weighting functions. Because  $\tilde{p}_{\theta}(x^t \mid x^{t+1}, y)$  will not in general coincide with the optimal twisted proposal  $p_{\theta}(x^t \mid x^{t+1})$ , we must introduce non-trivial weighting functions to ensure the resulting SMC sampler converges to the desired final target. In particular, we define twisted weighting functions as
$$
w _ {t} \left(x ^ {t}, x ^ {t + 1}\right) := \frac {p _ {\theta} \left(x ^ {t} \mid x ^ {t + 1}\right) \tilde {p} _ {\theta} \left(y \mid x ^ {t}\right)}{\tilde {p} _ {\theta} \left(y \mid x ^ {t + 1}\right) \tilde {p} _ {\theta} \left(x ^ {t} \mid x ^ {t + 1} , y\right)}, \quad t = 0, \dots , T - 1 \tag {11}
$$
and  $w_{T}(x^{T}) \coloneqq \tilde{p}_{\theta}(y \mid x^{T})$ . The weighting functions in eq. (11) recover the optimal, constant weighting functions if all other approximations at play are exact.
These tractable twisted proposals and weighting functions define intermediate targets that gradually approach the final target  $\nu_{0} = p_{\theta}(x^{0:T}\mid y)$ . Substituting into eq. (4) each  $\tilde{p}_{\theta}(x^{t}\mid x^{t + 1},y)$  in eq. (9) for  $r_t(x^t\mid x^{t + 1})$  and  $w_{t}(x^{t},x^{t + 1})$  in eq. (11) and then simplifying we obtain the intermediate targets
$$
\nu_ {t} \left(x ^ {0: T}\right) \propto p _ {\theta} \left(x ^ {0: T} \mid y\right) \left[ \frac {\tilde {p} _ {\theta} \left(y \mid x ^ {t}\right)}{p _ {\theta} \left(y \mid x ^ {t}\right)} \prod_ {t ^ {\prime} = 0} ^ {t - 1} \frac {\tilde {p} _ {\theta} \left(x ^ {t ^ {\prime}} \mid x ^ {t ^ {\prime} + 1} , y\right)}{p _ {\theta} \left(x ^ {t ^ {\prime}} \mid x ^ {t ^ {\prime} + 1} , y\right)} \right]. \tag {12}
$$
The right-hand bracketed term in eq. (12) can be understood as the discrepancy of  $\nu_{t}$  from the final target  $\nu_{0}$  accumulated from step  $t$  to 0 (see Appendix A.3 for a derivation). As  $t$  approaches 0,  $\tilde{p}_{\theta}(y \mid x^{t})$  improves as an approximation of  $p_{\theta}(y \mid x^{t})$ , and the  $t$ -term product inside the bracket consists of fewer terms - the latter accounts for the discrepancy between practical and optimal proposals. Finally, at  $t = 0$ , because  $\tilde{p}_{\theta}(y \mid x^{0}) = p_{\theta}(y \mid x^{0})$  by construction, eq. (12) reduces to  $\nu_{0}(x^{0:T}) = p_{\theta}(x^{0:T} \mid y)$ , as desired.
The TDS algorithm and asymptotic exactness. Together, the twisted proposals  $\tilde{p}_{\theta}(x^t \mid x^{t+1}, y)$  and weighting functions  $w_t(x^t, x^{t+1})$  lead to Twisted Diffusion Sampler, or TDS (Algorithm 1). While Algorithm 1 states multinomial resampling for simplicity, in practice other resampling strategies (e.g. systematic [5, Ch. 9]) may be used as well. Under additional conditions, TDS provides arbitrarily accurate estimates of  $p_{\theta}(x^0 \mid y)$ . Crucially, this guarantee does not rely on assumptions on the accuracy of the approximations used to derive the twisted proposals and weights. Appendix A provides the formal statement with complete conditions and proof.
Theorem 1. (Informal) Let  $\mathbb{P}_K(x^0) = (\sum_{k'}^K w_{k'})^{-1}\sum_{k = 1}^K w_k\delta_{x_k^0}(x^0)$  denote the discrete measure defined by the particles and weights returned by Algorithm 1 with  $K$  particles. Under regularity conditions on the twisted proposals and weighting functions,  $\mathbb{P}_K(x^0)$  converges setwise to  $p_\theta (x^0\mid y)$  as  $K$  approaches infinity.
3.3 TDS for inpainting, additional degrees of freedom