SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Rearrange and let  $C = \log |2\pi \sigma_{t + 1}^2 I|^{-1 / 2} / |2\pi \tilde{\sigma}_{t + 1}^2 I|^{-1 / 2}$
$$
= \frac {- 1}{2} \left[ \sigma_ {t + 1} ^ {- 2} \| \hat {\mu} - x ^ {t} \| ^ {2} - \tilde {\sigma} _ {t + 1} ^ {- 2} \| \hat {\mu} _ {\psi} - x ^ {t} \| ^ {2} \right] + C
$$
Expand and rearrange  $\| \hat{\mu}_{\psi} - x^{t}\|^{2} = \| \hat{\mu} -x^{t}\|^{2} + 2\langle \hat{\mu}_{\psi} - \hat{\mu},\hat{\mu} -x^{t}\rangle +\| \hat{\mu}_{\psi} - \hat{\mu}\|^{2}$
$$
= \frac {- 1}{2} \left[ \left(\sigma_ {t + 1} ^ {- 2} - \tilde {\sigma} _ {t + 1} ^ {- 2}\right) \| \hat {\mu} - x ^ {t} \| ^ {2} - 2 \tilde {\sigma} _ {t + 1} ^ {- 2} \langle \hat {\mu} _ {\psi} - \hat {\mu}, \hat {\mu} - x ^ {t} \rangle - \tilde {\sigma} _ {t + 1} ^ {2} \| \hat {\mu} _ {\psi} - \hat {\mu} \| ^ {2} \right] + C
$$
Let  $C' = C - \frac{1}{2}\tilde{\sigma}_{t + 1}^2\|\hat{\mu}_\psi - \hat{\mu}\|^2$  and rearrange. Note that  $\|\hat{\mu}_\psi - \hat{\mu}\|^2 < \infty$  by Assm. (c).
$$
= \frac {- 1}{2} (\sigma_ {t + 1} ^ {- 2} - \tilde {\sigma} _ {t + 1} ^ {- 2}) \| \hat {\mu} - x ^ {t} \| ^ {2} + \tilde {\sigma} _ {t + 1} ^ {- 2} \langle \hat {\mu} _ {\psi} - \hat {\mu}, \hat {\mu} - x ^ {t} \rangle + C ^ {\prime}
$$
Apply Cauchy-Schwarz
$$
\leq \frac {- 1}{2} (\sigma_ {t + 1} ^ {- 2} - \tilde {\sigma} _ {t + 1} ^ {- 2}) \| \hat {\mu} - x ^ {t} \| ^ {2} + \tilde {\sigma} _ {t + 1} ^ {- 2} \| \hat {\mu} _ {\psi} - \hat {\mu} \| \cdot \| \hat {\mu} - x ^ {t} \| + C ^ {\prime}
$$
Upper-bounding using that  $\max_x\frac{-a}{2} x^2 +bx = \frac{b^2}{2a}$  for  $a = \sigma_{t + 1}^{-2} - \tilde{\sigma}_{t + 1}^{-2} > 0$  by Assm. (d).
$$
\begin{array}{l} \leq \frac {1}{2} \frac {\left(\tilde {\sigma} _ {t + 1} ^ {- 2} \| \hat {\mu} _ {\psi} - \hat {\mu} \|\right) ^ {2}}{\sigma_ {t + 1} ^ {- 2} - \tilde {\sigma} _ {t + 1} ^ {- 2}} + C ^ {\prime} \\ = \frac {\tilde {\sigma} _ {t + 1} ^ {- 4}}{2 (\sigma_ {t + 1} ^ {- 2} - \tilde {\sigma} _ {t + 1} ^ {- 2})} \| \hat {\mu} _ {\psi} - \hat {\mu} \| ^ {2} + C ^ {\prime} \\ \end{array}
$$
Note that  $\hat{\mu}_{\psi} = \hat{\mu} + \sigma_{t+1}^{2}\nabla_{x^{t+1}}\log \tilde{p}_{\theta}(y \mid x^{t+1})$ .
$$
\begin{array}{l} = \frac {\tilde {\sigma} _ {t + 1} ^ {- 4}}{2 \left(\sigma_ {t + 1} ^ {- 2} - \tilde {\sigma} _ {t + 1} ^ {- 2}\right)} \sigma_ {t + 1} ^ {4} \| \nabla_ {x ^ {t + 1}} \log \tilde {p} _ {\theta} (y \mid x ^ {t + 1}) \| ^ {2} + C ^ {\prime} \\ \leq C ^ {\prime \prime}. \\ \end{array}
$$
The final line follows from Assumption (c), that the gradients of the twisting functions are bounded. The above derivation therefore provides that each  $w_{t}$  is bounded, concluding the proof.
B Riemannian Twisted Diffusion Sampler
This section provides additional details on the extension of TDS to Riemannian diffusion models introduced in Section 3. We first introduce the tangent normal distribution. We then provide with background on Riemannian diffusion models, which we parameterize with the tangent normal. Then we describe the extension of TDS to these models. Finally we show how Algorithm 1 modifies to this setting.
The tangent normal distribution. Just as the Gaussian is the parametric family underlying Euclidean diffusion models in Equation (3), the tangent normal (see e.g. [4, 8]) underlies generation in Riemannian diffusion models so we review it here.
We take the tangent normal to be the distribution is implied by a two step procedure. Given a variable  $x$  in the manifold, the first step is to sample a variable  $\bar{y}$  in  $\mathcal{T}_x$ , the tangent space at  $x$ ; if  $\{h_1, \ldots, h_D\}$  is an orthonormal basis of  $\mathcal{T}_x$  one may generate
$$
\bar {y} = \mu + \sum_ {d = 1} ^ {D} \sigma \epsilon_ {d} \cdot h _ {d},
$$
with  $\epsilon_{d} \stackrel{i.i.d.}{\sim} \mathcal{N}(0,1)$ , and  $\sigma^2 > 0$  a variance parameter. The second step is to project  $\bar{y}$  back onto the manifold to obtain  $y = \exp_x\{\bar{y}\}$  where  $\exp_x\{\cdot\}$  denotes the exponential map at  $x$ . The resulting distribution of  $y$  is denoted as  $\mathcal{T}\mathcal{N}_x(\mu, \sigma^2)$ . By construction,  $\mathcal{T}_x$  is a Euclidean space with its origin at  $x$ , so when  $\| \mu \|_2 = 0$ ,  $\mathcal{T}\mathcal{N}_x(\mu, \sigma^2)$  is centered on  $x$ . And since the geometry of a Riemannian manifold is locally Euclidean, when  $\sigma^2$  is small the exponential map is close to the identity map and the tangent normal is essentially a narrow Gaussian distribution in the manifold at  $x$ . Finally, we use  $\mathcal{T}\mathcal{N}_x(y; \mu, \sigma^2)$  to denote the density of the tangent normal evaluated at  $y$ .
Because this procedure involves a change of variables from  $\bar{y}$  to  $y$  (via the exponential map), to compute the tangent normal density one computes
$$
\mathcal {T N} _ {x} (y; \mu , \sigma^ {2}) = \mathcal {N} (\exp_ {x} ^ {- 1} \{y \}; \mu , \sigma^ {2}) \left| \frac {\partial}{\partial y} \exp_ {x} ^ {- 1} \{y \} \right|
$$
where  $\exp_x^{-1}\{y\}$  is the inverse of the exponential map (from the manifold into  $\mathcal{T}_x$ ), and  $\left|\frac{\partial}{\partial y}\exp_x^{-1}\{y\}\right|$  is the determinant of the Jacobian of the exponential map; when the manifold lives in a higher dimensional subset of  $R$ , we take  $\left|\frac{\partial}{\partial y}\exp_x^{-1}\{y\}\right|$  to be the product of the positive singular values of the Jacobian.
Riemannian diffusion models and the tangent normal distribution. Riemannian diffusion models proceeds through a geodesic random walk [8]. At each step  $t$ , one first samples a variable  $\bar{x}^t$  in  $\mathcal{T}_{x^{t+1}}$ ; if  $\{h_1, \ldots, h_D\}$  is an orthonormal basis of  $\mathcal{T}_{x^{t+1}}$  one may generate
$$
\bar {x} ^ {t} = \sigma_ {t + 1} ^ {2} s _ {\theta} (x ^ {t + 1}) + \sum_ {d = 1} ^ {D} \sigma_ {t + 1} \epsilon_ {d} \cdot h _ {d},
$$
with  $\epsilon_d \stackrel{i.i.d.}{\sim} \mathcal{N}(0,1)$ . One then projects  $\bar{x}^t$  back onto the manifold to obtain  $x^t = \exp_{x^{t+1}}\{\bar{x}^t\}$  where  $\exp_x\{\cdot\}$  denotes the exponential map at  $x$ .
This is equivalent to sampling  $x^{t}$  from a tangent normal distribution as
$$
x ^ {t} \sim p (x ^ {t} \mid x ^ {t + 1}) = \mathcal {T N} _ {x ^ {t + 1}} \left(x ^ {t}; \sigma_ {t + 1} ^ {2} s _ {\theta} (x ^ {t + 1}), \sigma_ {t + 1} ^ {2}\right).
$$
TDS for Riemannian diffusion models. To extend TDS, appropriate analogues of the twisted proposals and weights are all that is needed. For this extension we require that the diffusion model is also associated with a manifold-valued denoising estimate  $\hat{x}_{\theta}$  as will be the case when, for example,  $s_{\theta}(x^{t},t)\coloneqq \nabla_{x^{t}}\log q(x^{t} \mid x^{0} = \hat{x}_{\theta})$  for  $\hat{x}_{\theta} = \hat{x}_{\theta}(x^{t},t)$ . In contrast to the Euclidean case, a relationship between a denoising estimate and a computationally tractable score approximation may not always exist for arbitrary Riemannian manifolds; however for Lie groups when the forward diffusion is the Brownian motion, tractable score approximations do exist [39, Proposition 3.2].
For the case of positive and differentiable  $p_{y|x^0}(y|x^0)$ , we again choose twisting functions  $\tilde{p}_{\theta}(y \mid x^t) \coloneqq p_{y|x^0}(y|\hat{x}_{\theta}(x^t))$ .
Next are the inpainting and inpainting with degrees of freedom cases. Here, assume that  $x^0$  lives on a multidimensional manifold (e.g.  $SE(3)^N$ ) and the unmasked observation  $y = x_{\mathbf{M}}^0$  with  $\mathbf{M} \subset \{1, \dots, N\}$  on a lower-dimensional sub-manifold (e.g.  $SE(3)^{|{\bf M}|}$ , with  $|{\bf M}| < N$ ). In this case, twisting functions are constructed exactly as in Section 3.3, except with the normal density in Equation (13) replaced with a Tangent normal as
$$
\tilde {p} _ {\theta} (y \mid x ^ {t}) = \mathcal {T N} _ {\hat {x} _ {\theta} (x ^ {t}) _ {\bf M}} (y; 0, \bar {\sigma} _ {t} ^ {2}).
$$
For all cases, we propose the twisted proposal as
$$
\tilde {p} _ {\theta} \left(x ^ {t} \mid x ^ {t + 1}, y\right) = \mathcal {T N} _ {x ^ {t + 1}} \left(x ^ {t}; \sigma_ {t + 1} ^ {2} s _ {\theta} \left(x ^ {t + 1}, y\right), \tilde {\sigma} _ {t + 1} ^ {2}\right) \tag {17}
$$
where as in the Euclidean case  $s_{\theta}(x^t,y) = s_{\theta}(x^t) + \nabla_{x^t}\log \tilde{p} (y\mid x^t)$ .
Weights at intermediate steps are computed as in the Euclidean case (Equation (11)):
$$
\begin{array}{l} 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)} \\ = \frac {\mathcal {T N} _ {x ^ {t + 1}} (x ^ {t} ; \sigma_ {t + 1} ^ {2} s _ {\theta} (x ^ {t + 1}) , \sigma_ {t + 1} ^ {2}) \tilde {p} _ {\theta} (y \mid x ^ {t})}{\tilde {p} _ {\theta} (y \mid x ^ {t + 1}) \mathcal {T N} _ {x ^ {t + 1}} (x ^ {t} ; \sigma_ {t + 1} ^ {2} s _ {\theta} (x ^ {t + 1} , y) , \tilde {\sigma} _ {t + 1} ^ {2})}. \\ \end{array}
$$
While the proposal and target contribute identical Jacobian determinant terms that cancel out, they remain in the twisting functions.
Adapting the TDS algorithm to the Riemannian setting. To translate the TDS algorithm to the Riemannian setting we require only two changes.
The first is on Algorithm 1 Line 6. Here we assume that the unconditional score is computed through the transition density function:
$$
s _ {\theta} \left(x ^ {t + 1}\right) := \nabla_ {x ^ {t + 1}} \log q _ {t + 1 | 0} \left(x ^ {t + 1} \mid \hat {x} _ {\theta}\right) \quad \text {f o r} \hat {x} _ {\theta} = \hat {x} _ {\theta} \left(x ^ {t + 1}\right).
$$
Note that the gradient above ignores the dependence of  $\hat{x}_{\theta}$  on  $x^{t + 1}$ .
The conditional score approximation in Line 6 is then replaced with
$$
\tilde {s} _ {k} \leftarrow \tilde {p} _ {\theta} (\cdot \mid x _ {k} ^ {t + 1}, y) := \nabla_ {x _ {k} ^ {t + 1}} \log q _ {t + 1 | 0} (x _ {k} ^ {t + 1} \mid \hat {x} _ {\theta}) + \nabla_ {x _ {k} ^ {t + 1}} \log \tilde {p} _ {k} ^ {t + 1} \quad \text {f o r} \hat {x} _ {\theta} = \hat {x} _ {\theta} (x _ {k} ^ {t + 1}).
$$
Notably,  $\tilde{s}_k$  is a  $\mathcal{T}_{x_k^{t + 1}}$ -valued Riemannian gradient.
The second change is to make the proposal on Algorithm 1 Line 7 a tangent normal, as defined in eq. (17).
Training with conditioning information. Several approaches involve training a neural network to directly sample from a conditional diffusion models. These approaches include (i) conditional training with embeddings of conditioning information, e.g. for denoising images [30], and text-to-image generation [27], (ii) conditioning training with a subset of the state space, e.g. for protein design [37], and image inpainting [29], (iii) classifier-guidance [9, 33], an approach for generating samples of a desired class, which requires training a time-dependent classifier to approximate, for each  $t$ ,  $p_{\theta}(y \mid x^t)$ , and training such a time-dependent classifier may be inconvenient and costly), and (iv) classifier-free guidance [17], an approach that builds on classifier-guidance without training a classifier, and instead trains a diffusion model with class information as additional input.
Langevin and Metropolis-Hastings steps. Some prior work has explored using Markov chain Monte Carlo steps in sampling schemes to better approximate conditional distributions. For example unadjusted Langevin dynamics [19, 33] or Hamiltonian Monte Carlo [11] in principle permit asymptotically exact sampling in the limit of many steps. However these strategies are only guaranteed to target the conditional distributions of the joint model under the (unrealistic) assumption that the unconditional model exactly matches the forward process, and in practice adding such steps can worsen the resulting samples, presumably as a result of this approximation error [21]. By contrast, TDS does not require the assumption that the learned diffusion model exactly matches the forward noising process.
D Empirical results additional details
D.1 Synthetic diffusion models on two dimensional problems
Forward process. Our forward process is variance preserving (as described in Appendix A) with  $T = 100$  steps and a quadratic variance schedule. We set  $\sigma_t^2 = \sigma_{\mathrm{min}}^2 + (\frac{t}{T})^2 \sigma_{\mathrm{max}}^2$  with  $\sigma_{\mathrm{min}}^2 = 10^{-5}$  and  $\sigma_{\mathrm{max}}^2 = 10^{-1}$ .
Unconditional target distributions and likelihood. We evaluate the different methods with two different unconditional target distributions:
Figure D: Errors of conditional mean estimations with 2 SEM error bars averaged over 25 replicates on mixture of Gaussians unconditional target. TDS applies to all three tasks and provides increasing accuracy with more particles.
1. A bivariate Gaussian with mean at  $\left(\frac{1}{2},\frac{1}{2}\right)$  and covariance 0.9 and
2. A Gaussian mixture with three components with mixing proportions [0.3, 0.5, 0.2], means  $[(1.54, -0.29), (-2.18, 0.57), (-1.09, -1.40)]$ , and 0.2 standard deviations.
We evaluate on conditional distributions defined by the three likelihoods described in Section 5.1.
Figure D provides results analogous to those in Figure 1 but with the mixture of Gaussians unconditional target. In this second example we evaluate a with a variation of the inpainting degrees of freedom case wherein we consider  $y = 1$  and  $\mathcal{M} = \{[1], [2]\}$ , so that  $p_{y|x^0}(y|x^0) = \delta_y(x_1^0) + \delta_y(x_2^0)$ .
D.2 Image conditional generation experiments
We perform extensive ablation studies on class-conditional generator task and inpainting task using the small-scale MNIST dataset  $(28\times 28\times 1$  dimensions) in Appendix D.2.1 and Appendix D.2.2. And we show TDS's applicability to higher-dimensional datasets, namely CIFAR10  $(32\times \times 3$  dimensions) and ImageNet256  $(256\times 256\times 3$  dimensions) in Appendix D.2.3.
D.2.1 Class-conditional generation on MNIST
Set-up. For MNIST, we set up a diffusion model using the variance preserving framework. The model architecture is based on the guided diffusion codebase<sup>5</sup> with the following specifications: number of channels = 64, attention resolutions = "28,14,7", number of residual blocks = 3, learn sigma (i.e. to learn the variance of  $p_{\theta}(x^{t-1} \mid x^t)$ ) = True, resblock updown = True, dropout = 0.1, variance schedule = "linear". We trained the model for 60k epochs with a batch size of 128 and a learning rate of  $10^{-4}$  on 60k MNIST training images. The model uses  $T = 1,000$  for training and  $T = 100$  for sampling.
The classifier used for class-conditional generation and evaluation is a pretrained ResNet50 model. This classifier is trained on the same set of MNIST training images used by diffusion model training.
In addition, we include a variation called TDS-truncate that truncates the TDS procedure at  $t = 10$  and returns the prediction  $\hat{x}_{\theta}(x^{10})$ .
Sample plots. To supplement the sample plot conditioned on class 7 in Figure 2a, we present samples conditioned on each of the remaining 9 classes and from other methods. We observe that samples from Gradient Guidance have noticeable artifacts, whereas the other 4 methods produce authentic and correct digits. However, most samples from IS or TDS-IS are identical due to the collapse of importance weights. By contrast, samples from TDS and TDS-truncate have greater diversity, with the latter exhibiting slightly more variations
Ablation study on twist scales. We consider exponentiating and re-normalizing twisting functions by a twist scale  $\gamma$ , i.e. setting new twisting functions to  $\tilde{p}_{\theta}(y \mid x^t; \gamma) \propto p_{y|x^0}(y|\hat{x}_{\theta}(x^t))^{\gamma}$ . In particular, when  $t = 0$ , we set  $p_{y|x^0}(y|x^0; \gamma) \propto p_{y|x^0}(y|x^0)^{\gamma}$ . This modification suggests that the targeted conditional distribution is now
$$
p _ {\theta} \left(x ^ {0} \mid y; \gamma\right) \propto p _ {\theta} \left(x ^ {0}\right) p _ {y \mid x ^ {0}} \left(y \mid x ^ {0}\right) ^ {\gamma}.
$$
Downloaded from https://github.com/VSehwag/minimal-diffusion
Figure E: MNIST class-conditional generation. 16 randomly chosen conditional samples from 64 particles in a single run, given class  $y$ . From top to bottom,  $y = 0, 1, 2, 3, 4, 5, 6, 8, 9$ .
By setting  $\gamma > 1$ , the classification likelihood becomes sharper, which is potentially helpful for twisting the samples towards a specific class. The TDS algorithm (and likewise TDS-IS and Gradient Guidance) still apply with this new definition of twisting functions. The use of twist scale is similar to the guidance scale introduced in the literature of Classifier Guidance, which is used to multiply the gradient of the log classification probability [9].
In Figure F, we examine the effect of varying twist scales on classification accuracy of TDS, TDS-IS and Gradient Guidance. We consider two ways to evaluate the accuracy. First, classification accuracy computed by a neural network classifier, where the evaluation setup is the same as in Section 5.2. Second, the human-rated classification accuracy, where a human (one of the authors) checks if a generated digit has the right class and does not have artifacts. Since human evaluation is expensive, we only evaluate TDS, TDS-IS (both with  $K = 64, 2$ ) and Gradient Guidance. In each run, we randomly sample one particle out of  $K$  particles according to the associated weights. We conduct 64 runs for each class label, leading to a total of 640 samples for human evaluation.
Figure Fa depicts the classification accuracy measured by a neural network classifier. We observe that in general larger twist scale improves the classification accuracy. For TDS and TDS-IS, the improvement is more significant for smaller number of particles  $K$  used.
Figure Fb depicts the human-rated accuracy. In this case, we find that larger twist scale is not necessarily better. A moderately large twist scale  $(\gamma = 2,3)$  generally helps increasing the accuracy, while an excessively large twist scale  $(\gamma = 10)$  decreases the accuracy. An exception is TDS with  $K = 64$  particles, where any  $\gamma > 1$  leads to worse accuracy compared to the case of  $\gamma = 1$ . Study on twist scale aside, we find that using more particles  $K$  help improving human-rated accuracy (recall that Gradient Guidance is a special case of TDS with  $K = 1$ ): given the same twist scale, TDS with  $K = 1,2$  or 64 has increasing accuracy. In addition, both TDS and TDS-IS with  $K = 64$  and  $\gamma = 1$
Twist scale=1
Twist scale=2
Twist scale=3
Twist scale=4
Twist scale=10
(a) TDS  $(K = 64)$ . # of good-quality samples counted by human is 62, 61, 63, 60, 62, from left to right.
Twist scale=1
(b) TDS  $(K = 2)$ . # of good-quality samples counted by human 49, 51, 51, 46, 45, from left to right.
Twist scale=1
(c) TDS-IS  $(K = 64)$ . # of good-quality samples counted by a human is 62, 61, 64, 64, 61, from left to right.
Twist scale=1
(d) TDS-IS  $(K = 2)$ . # of good-quality samples counted by a human is 49, 52, 53, 52, 48, from left to right.
Twist scale=1
(e) Gradient Guidance. # of good-quality samples counted by a human is 31, 38, 41, 39, 34, from left to right.
Figure G: MNIST class-conditional generation: random samples selected from 64 random runs conditioned on class 6 under different twist scales. Top to bottom: TDS ( $K = 64$ ), TDS ( $K = 2$ ), TDS-IS ( $K = 64$ ), TDS-IS ( $K = 2$ ) and Gradient Guidance. In general, moderately large twist scales improve the sample quality. However, overly large twist scale (e.g. 10) would distort the digit shape with more artifacts, though retaining useful features that may allow a neural network classifier to identify the class. 24
Figure H: ESS trace avg. over 100 runs  $(K = 64)$
have almost perfect accuracy. The effect of  $K$  on human-rated accuracy is consistent with previous findings with neural network classifier evaluation in Section 5.2.
We note that there is a discrepancy on the effects of twist scales between neural network evaluation and human evaluation. We suspect that when overly large twist scale is used, the generated samples may fall out of the data manifold; however, they may still retain features recognizable to a neural network classifier, thereby leading to a low human-rated accuracy but a high classifier-rated accuracy. To validate this hypothesis, we present samples conditioned on class 6 in Figure G. For example, in Figure Ge, Gradient Guidance with  $\gamma = 1$  has 31 good-quality samples out of 64, and the rest of the samples often resemble the shape of other digits, e.g. 3,4,8; and Gradient Guidance with  $\gamma = 10$  has 34 good-quality samples, but most of the remaining samples resemble 6 with many artifacts.
Effective sample size. Effective sample size (ESS) is a common metric used to diagnose the performance of SMC samplers, which is defined as  $\left(\sum_{k=1}^{K} w_k^t\right)^2 / \left(\sum_{k=1}^{K} (w_k^t)^2\right)$  for  $K$  weighted particles  $\{x_k^t, w_k^t\}_{k=1}^K$ . Note that ESS is always bounded between 0 and  $K$ .
Figure H depicts the ESS trace comparison of TDS, TDS-IS, and IS. TDS has a general upward trend of ESS approaching  $K = 64$ . Though in the final few steps ESS of TDS drops by a half, it is still higher than that of TDS-IS and IS that deteriorates to around 1 and 6 respectively. In practice, one can use the TDS-truncate variant introduced above to avoid the ESS drop in the final few steps.
D.2.2 Image inpainting on MNIST
The inpainting task is to sample images from  $p_{\theta}(x^0 \mid x_{\mathbf{M}}^0 = y)$  given observed part  $x_{\mathbf{M}}^0 = y$ . Here we segment  $x^0 = [x_{\mathbf{M}}^0, x_{\mathbf{M}}^0]$  into observed dimensions  $\mathbf{M}$  and unobserved dimensions  $\mathbf{M}$ , as described in Section 3.3.
In this experiment, we consider two types of observed dimensions: (1)  $\mathbf{M} =$  "half", where the left half of an image is observed, and (2)  $\mathbf{M} =$  "quarter", where the upper left quarter of an image is observed.
We run TDS, TDS-IS, Gradient Guidance, SMC-Diff, and Replacement method to inpaint 10,000 validation images. We also include TDS-truncate that truncates the TDS procedure at  $t = 10$  and returns  $\hat{x}_{\theta}(x^{10})$ .
Twisting function variance schedule. We use a flexible variance schedule  $\{\hat{\sigma}_t^2\}$  to extend the twisting functions in eq. (13) to
$$
\tilde {p} _ {\theta} (y \mid x ^ {t}, \mathbf {M}) := \mathcal {N} (y \mid \hat {x} _ {\theta} (x ^ {t}) _ {\mathbf {M}}, \hat {\sigma} _ {t} ^ {2}).
$$
Ideally,  $\hat{\sigma}_t^2$  should match  $\mathrm{Var}_{p_\theta}[y\mid x^t,\mathbf{M}]$ . In TDS, we choose  $\hat{\sigma}_t^2 = \frac{(\bar{\sigma}_t^2 / \bar{\alpha}_t)\cdot\tau^2}{\bar{\sigma}_t^2 / \bar{\alpha}_t + \tau^2}$  where  $\tau^2 = 0.12$  is the estimated sample variance of the training data (averaged over pixels). This choice is inspired from Song et al. [32, Appendix A.3]: if the data distribution is  $q(x^{0}) = \mathcal{N}(0,\tau^{2})$ , and  $q(x^{t}\mid x^{0}) = \mathcal{N}(x^{t}\mid \sqrt{\bar{\alpha}_{t}} x^{0},\sigma_{t}^{2})$ , then by Bayes rule  $\mathrm{Var}_q[x^0\mid x^t] = \frac{(\bar{\sigma}_t^2 / \bar{\alpha}_t)\cdot\tau^2}{(\bar{\sigma}_t^2 / \bar{\alpha}_t) + \tau^2}$ .
Changing  $\hat{\sigma}_t^2$  has a similar effect as using a twist scale to exponentiate the twisting functions in Appendix D.2.1. In particular, for a Gaussian distribution, exponentiating the distribution corresponds to re-scaling the variance:  $\mathcal{N}(x\mid \mu ,(\sigma /\sqrt{\gamma})^2)\propto \mathcal{N}(x\mid \mu ,\sigma^2)^\gamma$
Metrics. We consider 3 metrics. (1) We use the effective sample size (ESS) to compare the particle efficiency among different SMC samplers (namely TDS, TDS-IS, and SMC-Diff).