SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The twisting functions  $\tilde{p}_{\theta}(y|x^t) \coloneqq p_{y|x^0}(y|\hat{x}_\theta(x^t))$  we introduced above are one convenient option, but are sensible only when  $p_{y|x^0}(y|\hat{x}_\theta(x^t))$  is differentiable and strictly positive. We now show how alternative twisting functions lead to proposals and weighting functions that address inpainting problems and more flexible conditioning specifications. In these extensions, Algorithm 1 still applies with the new definitions of twisting functions. Appendix A provides additional details, including the adaptation of TDS to variance preserving diffusion models [18].
Inpainting. Consider the case that  $x^0$  can be segmented into observed dimensions  $\mathbf{M}$  and unobserved dimensions  $\bar{\mathbf{M}}$  such that we may write  $x^0 = [x_{\mathbf{M}}^0, x_{\bar{\mathbf{M}}}^0]$  and let  $y = x_{\mathbf{M}}^0$ , and take  $p_{y|x^0}(y|x^0) = \delta_y(x_{\mathbf{M}}^0)$ . The goal, then, is to sample  $p_\theta(x^0 \mid x_{\mathbf{M}}^0 = y) = p_\theta(x_{\bar{\mathbf{M}}}^0 \mid x_{\mathbf{M}}^0) \delta_y(x_{\mathbf{M}}^0)$ . Here we define the twisting function for each  $t > 0$  as
$$
\tilde {p} _ {\theta} (y \mid x ^ {t}, \mathbf {M}) := \mathcal {N} \left(y; \hat {x} _ {\theta} (x ^ {t}) _ {\mathbf {M}}, t \sigma^ {2}\right), \tag {13}
$$
and set twisted proposals and weights according to eqs. (9) and (11). The variance in eq. (13) is chosen as  $t\sigma^2 = \mathrm{Var}_{p_\theta}[x^t \mid x^0]$  for simplicity; in general, choosing this variance to more closely match  $\mathrm{Var}_{p_\theta}[y \mid x^t]$  may be preferable. For  $t = 0$ , we define the twisting function analogously with small positive variance, for example as  $\tilde{p}_\theta(y \mid x^0) = \mathcal{N}(y; x_{\mathbf{M}}^0, \sigma^2)$ . This choice of simplicity changes the final target slightly; alternatively, the final twisting function, proposal, and weights may be chosen to maintain asymptotic exactness (see Appendix A.4).
Inpainting with degrees of freedom. We next consider the case when we wish to condition on some observed dimensions, but have additional degrees of freedom. For example in the context of motif-scaffolding in protein design, we may wish to condition on a functional motif  $y$  appearing anywhere in a protein structure, rather than having a pre-specified set of indices  $\mathbf{M}$  in mind. To handle this situation, we (i) let  $\mathcal{M}$  be a set of possible observed dimensions, (ii) express our ambivalence in which dimensions are observed as  $y$  using randomness by placing a uniform prior on  $p(\mathbf{M}) = 1 / |\mathcal{M}|$  for each  $\mathbf{M} \in \mathcal{M}$ , and (iii) again embed this new variable into our joint model to define the degrees of freedom likelihood by  $p(y \mid x^0) = \sum_{\mathbf{M} \in \mathcal{M}} p_\theta(\mathbf{M}, y \mid x^0) = |\mathcal{M}|^{-1} \sum_{\mathbf{M} \in \mathcal{M}} p(y \mid x^0, \mathbf{M})$ . Accordingly, we approximate  $p_\theta(y \mid x^t)$  with the twisting function
$$
\tilde {p} _ {\theta} (y \mid x ^ {t}) := | \mathcal {M} | ^ {- 1} \sum_ {\mathbf {M} \in \mathcal {M}} \tilde {p} _ {\theta} (y \mid x ^ {t}, \mathbf {M}), \tag {14}
$$
with each  $\tilde{p}_{\theta}(y \mid x^t, \mathbf{M})$  defined as in eq. (13). Notably, eqs. (13) and (14) coincide when  $|\mathcal{M}| = 1$ .
The sum in eq. (14) may be computed efficiently because each term depends on  $x^t$  only through the same denoising estimate  $\hat{x}_{\theta}(x^{t})$ , which must be computed only once. Since computation of  $\hat{x}_{\theta}(x^{t})$  is the expensive step, the overall run-time is not significantly increased by using even large  $|\mathcal{M}|$ .
3.4 TDS on Riemannian manifolds
TDS extends to Riemannian diffusion models on with little modification. Riemannian diffusion models [8] are structured like in the Euclidean case, but with conditionals defined by tangent normal distributions parameterized with a score approximation followed by a manifold projection step (see e.g. [4, 8]). When we assume that (as, e.g., in [39]) the model is associated with a denoising network  $\hat{x}_{\theta}$ , twisting functions are also constructed analogously. For conditional tasks defined by likelihoods,  $\tilde{p}_{\theta}(y \mid x^{t})$  in eq. (8) applies. For inpainting (and by extension, degrees of freedom), we propose
$$
\tilde {p} _ {\theta} (y \mid x ^ {t}) = \mathcal {T N} _ {\hat {x} _ {\theta} \left(x ^ {t}\right) _ {\mathrm {M}}} (y; 0, t \sigma^ {2}), \tag {15}
$$
where  $\mathcal{T}\mathcal{N}_{\hat{x}_{\theta}(x^t)_{\mathbf{M}}} (0, t\sigma^2)$  is a tangent normal distribution centered on  $\hat{x}_{\theta}(x^t)_{\mathbf{M}}$ . As in the Euclidean case,  $\tilde{p}_{\theta}(x^t \mid x^{t+1}, y)$  is defined with conditional score approximation  $s_{\theta}(x^t, y) = s_{\theta}(x^t) + \nabla_{x^t} \log \tilde{p}(y \mid x^t)$ , which is computable by automatic differentiation. Appendix B provides details.
There has been much recent work on conditional generation using diffusion models. These prior works demand either task specific conditional training, or involve unqualified approximations and can suffer from poor performance in practice. We discuss additional related work in Appendix C.
Gradient guidance. We use gradient guidance to refer to a general approach that incorporates conditioning information with gradients through the neural network  $\hat{x}_{\theta}(x^{t})$ . For example, in an
Figure 1: Errors of conditional mean estimations with 2 SEM error bars averaged over 25 replicates. TDS applies to all three tasks and provides increasing accuracy with more particles.
inpainting setting, Ho et al. [19] propose a Gaussian approximation  $\mathcal{N}(y\mid \hat{x}_{\theta}(x^{t})^{\mathbf{M}},\alpha)$  to  $p_{\theta}(y\mid x^{t})$  This approximation motivates a modified transition distribution as in Equation (9), with the corresponding approximation to the conditional score as  $s_\theta (x^t,y) = s_\theta (x^t) - \nabla_{x^t}\| y - \hat{x}_\theta (x^t)_M\| ^2 /\alpha .$  This approach coincides exactly with TDS applied to the inpainting task when a single particle is used. Similar Gaussian approximations have been used by [6, 32] for other inverse problems.
Gradient guidance can also be used with non-Gaussian approximations; e.g. using  $p_{y|x^0}(y|\hat{x}_\theta (x^t)) \approx p_\theta (y\mid x^t)$  for a given likelihood  $p_{y|x^0}(y|x^0)$ . This choice again recovers TDS with one particle.
Empirically, these heuristics can have unreliable performance, e.g. in image inpainting problems [40]. By comparison, TDS enjoys the benefits of gradient guidance by using it in proposals, while also providing a mechanism to eliminate approximation error by simulating additional particles.
Replacement method. The replacement method [33] is a method introduced for image inpainting using only unconditional diffusion models. The idea is to replace the observed dimensions of intermediate samples  $x_{\mathbf{M}}^{t}$ , with a noisy version of observation  $x_{\mathbf{M}}^{0}$ . However, it is a heuristic approximation and can lead to inconsistency between inpainted region and observed region [23]. Additionally, the replacement method applies only to inpainting problems. While recent work has extended the replacement method to linear inverse problems [e.g. 22], the approach provides no accuracy guarantees. It is unclear how to extend these methods to arbitrary differentiable likelihoods.
SMC samplers for diffusion models. Most closely related to the present work is SMC-Diff [34], which uses SMC to provide asymptotically accurate conditional samples for the inpainting problem. However, this prior work (i) is limited to the inpainting case, and (ii) provides asymptotic guarantees only under the assumption that the learned diffusion model exactly matches the forward noising process, which is rarely satisfied in practice. Also, SMC-Diff does not leverage twisting functions.
In concurrent work, Cardoso et al. [2] propose MCGdiff, an alternative SMC algorithm that uses the framework of auxiliary particle filtering [26] to provide asymptotically exact conditional inference. Compared to TDS, MCGdiff avoids computing gradients of the denoising network but applies only to linear inverse problems, with inpainting as a special case.
5 Simulation study and conditional image generation
We first test the dependence of the accuracy of TDS on the number of particles in synthetic settings with tractable exact conditionals in Section 5.1. Section 5.2 compares TDS to alternative approaches on class-conditional image generation, and an image inpainting experiment is included in Appendix D.2.2. See Appendix D for all additional details.
Our evaluation includes: (1) TDS; (2) TDS-IS, an importance sampler that uses TDS's proposal; (3) IS, a naive importance sampler described in Section 3.1; and (4) Gradient Guidance. For inpainting-type problems we further include: (5) Replacement method; and (6) SMC-Diff.
Each SMC sampler forms an approximation to  $p_{\theta}(x^0 \mid y)$  with  $K$  weighted particles  $(\sum_{k'=1}^{K} w_{k'})^{-1}$ .  $\sum_{k=1}^{K} w_k \delta_{x_k^0}$ . Gradient guidance and replacement method are considered to form a similar particle-based approximation with  $K$  independent samples viewed as  $K$  particles with uniform weights.
Compute Cost: Compared to unconditional generation, TDS has compute cost that is (i) larger by a constant factor due to the need to backpropagate through the denoising network when computing the conditional score approximation and (ii) linear in the number of particles. As a result, the compute cost at inference cost is potentially large relative for accurate inference to be achieved.
By comparison, conditional training methods provide fast inference by amortization [13], in which most computation is done ahead of time. Hence they may be preferable for applications where sampling time computation is a primary concern. On the other hand, TDS may be preferable when the likelihood criterion is readily available (e.g. through an existing classifier on clean data) but training an amortized model poses challenges; for example, the labeled data, neural-network engineering expertise, and up-front compute resources required for amortization training can be prohibitive.
5.1 Applicability and precision of TDS in two dimensional simulations
We explore two questions in this section: (i) what sorts of conditioning information can be handled by TDS and other methods, and (ii) how does the precision of TDS depend on the number of particles?
To study these questions, we first consider an unconditional diffusion model  $p_{\theta}(x^0)$  approximation of a bivariate Gaussian. For this choice, the marginals of the forward process are also Gaussian, and so we may define  $p_{\theta}(x^{0:T})$  with an analytically tractable score function without neural network approximation. Consequently, we can analyze the performance without the influence of score network approximation errors. And the choice of a two-dimensional diffusion permits close approximation of exact conditional distributions by numerical integration that can then be used as ground truth.
We consider three test cases defining the conditional information: (1) Smooth likelihood:  $y$  is an observation of the Euclidean norm of  $x$  with Laplace noise, with  $p_{\theta}(y \mid x^0) = \exp \{||x^0||_2 - y|\} / 2$ . This likelihood is smooth almost everywhere. (2) Inpainting:  $y$  is an observation of the first dimension of  $x^0$ , with  $p_{\theta}(y \mid x^0, \mathbf{M} = 0) = \delta_y(x_0^0)$ . (3) Inpainting with degrees-of-freedom:  $y$  is an observation of either the first or second dimension of  $x^0$ , with  $\mathcal{M} = \{0, 1\}$  and  $p_{\theta}(y \mid x^0) = \frac{1}{2} [\delta_y(x_0^0) + \delta_y(x_1^0)]$ . In all cases we fix  $y = 0$  and consider estimating  $\mathbb{E}_{p_{\theta}}[x^0 \mid y]$ .
Figure 1 reports the estimation error for the mean of the desired conditional distribution, i.e.  $\| \sum w_k x_k^0 - \mathbb{E}_q[x^0 \mid y] \|_2$ . TDS provides a computational-statistical trade-off: using more particles decreases mean square estimation error at the  $O(1 / K)$  parametric rate (note the slopes of  $-1$  in log-log scale) as expected from standard SMC theory [5, Ch. 11]. This convergence rate is shared by TDS, TDS-IS, and IS in the smooth likelihood case, and by TDS, SMCDiff and in the inpainting case; TDS-IS, IS and SMCDiff are applicable however only in these respective cases, whereas TDS applicable in all cases. The only other method which applies to all three settings is Gradient Guidance, which exhibits significant estimation error and does not improve with many particles.
5.2 Class-conditional image generation
We next study the performance of TDS on diffusion models with neural network approximations to the score functions. In particular, we study the class-conditional image generation task, which involves sampling an image from  $p_{\theta}(x^0 \mid y) \propto p_{\theta}(x^0) p_{y|x^0}(y|x^0)$ , where  $p_{\theta}(\cdot)$  is a pretrained diffusion model on images  $x^0$ ,  $y$  is a given image class, and  $p_{y|x^0}(y|\cdot)$  is the classification likelihood. To assess the faithfulness of generation, we evaluate classification accuracy on predictions of conditional samples  $x^0$  given  $y$ , made by the same classifier that specifies the likelihood. In all experiments, we follow the standard practice of returning the denoising mean on the final sample [18].
On the MNIST dataset, we compare TDS to TDS-IS, Gradient Guidance, and IS. Figure 2a compares the conditional samples of TDS and Gradient Guidance given class  $y = 7$ . Samples from Gradient Guidance have noticeable artifacts, and most of them do not resemble the digit 7; by contrast, TDS produces authentic and correct digits.
Figure 2b presents an ablation study of the effect of # of particles  $K$  on MNIST. For all SMC samplers, more particles improve the classification accuracy, with  $K = 64$  leading to nearly perfect accuracy. The performance of Gradient Guidance is constant with respect to  $K$  (in expectation). Notably, for fixed  $K$ , TDS and TDS-IS have comparable performance and outperform Gradient guidance and IS.
We next apply TDS to higher dimension datasets. Figure 2c shows samples from TDS  $(K = 16)$  using a pre-trained diffusion model and a pretrained classifier on the ImageNet dataset  $(256 \times 256 \times 3$  dimensions). These samples are qualitatively good and capture the class label. Appendix D.2.3 provides more samples, comparison to Classifier Guidance [9], and results for the CIFAR-10 dataset.
This likelihood is smooth except at the point  $x = (0,0)$ .
Gradient Guidance
Figure 2: Image class-conditional generation task.
TDS can be extended by exponentiating twisting functions with a twist scale. This extension is related to the existing literature of Classifier Guidance that considers re-scaling the gradient of the log classification probability. See Appendix D.2.1 for details and ablation study.
6 Case study in computational protein design: the motif-scaffolding problem
The biochemical functions of proteins are typically imparted by a small number of atoms, known as a motif, that are stabilized by the overall protein structure, known as the scaffold [36]. A central task in protein design is to identify stabilizing scaffolds in response to motifs expected to confer function. We here describe an application of TDS to this task, and compare TDS to the state-of-the-art conditionally-trained model, RFdiffusion. See additional details in Appendix E.
Given a generative model supported on designable protein structures  $p_{\theta}(x^0)$ , suitable scaffolds may be constructed by solving a conditional generative modeling problem [34]. Complete structures are first segmented into a motif  $x_{\mathbf{M}}^{0}$  and a scaffold  $x_{\overline{\mathbf{M}}}^{0}$ , i.e.  $x^0 = [x_{\mathbf{M}}^0, x_{\overline{\mathbf{M}}}^0]$ . Putative compatible scaffolds are then identified by (approximately) sampling from  $p_{\theta}(x_{\overline{\mathbf{M}}}^{0} \mid x_{\mathbf{M}}^{0})$  [34].
While the conditional generative modeling approach to motif-scaffolding has produced functional, experimentally validated structures for certain motifs [37], the general problem remains open. Moreover, current methods for motif-scaffolding require one to specify the location of the motif within the primary sequence of the full scaffold; this choice can require expert knowledge and trial and error.
We hypothesized that improved motif-scaffolding could be achieved through accurate conditional sampling. To this end, we applied TDS to FrameDiff, a Riemannian diffusion model that parameterizes protein backbones as a collection of  $N$  rigid bodies (known as residues) in the manifold  $SE(3)^{N}$  [39]. Each of the  $N$  elements of  $SE(3)$  consists of a rotation matrix and a translation that parameterize the locations of the backbone atoms of each residue.
Likelihood, twisting, and degrees of freedom. The basis of our approach to the motif scaffolding problem is analogous to the inpainting case described in Section 3.3. We let  $p_{y|x^0}(y|x^0) = \delta_y(x_{\mathbf{M}}^0)$ , where  $y \in SE(3)^M$  describes the coordinates of backbone atoms of an  $M = |\mathcal{M}|$  residue motif. As such, to define twisting function we adopt the Riemannian TDS formulation described in Section 3.4.
To eliminate the requirement that the placement of the motif within the scaffold be pre-specified, we incorporate the motif placement as a degree of freedom. We (i) treat the indices of the motif within the chain as a mask  $\mathbf{M}$ , (ii) let  $\mathcal{M}$  be a set of possible masks of size equal to the length of the motif, and (iii) apply Equation (14) to average over these possible masks.
For some scaffolding problems, it is known that all motif residues must appear with contiguous indices. In this case we choose  $\mathcal{M}$  to be the set of all possible contiguous masks. However, when motif residues are only known to appear in two or more possibly discontinuous blocks, the number of possible placements can be too large and we choose  $\mathcal{M}$  by randomly sampling at most some maximum number (# Motif Locs.) of masks.
We similarly eliminate the global translation and rotation of the motif as a degree of freedom. The motivation is that restricting to one possible pose of the motif narrows the conditional distribution,
<sup>4</sup>https://github.com/jasonkyuyim/se3_diffusion
Figure 3: Protein motif-scaffolding case study results
Scaffold size	TDS & FrameDiff	RF diffusion
<100 res.	9	3
≥ 100 res.	2	8
Overall	11	11
(b) # problems with higher success rate
thereby making inference more challenging. For translation, we use a likelihood that is invariant to the motif's center-of-mass and placement in the final scaffold by choosing  $p_{y|x^0}(y|x) = \delta_{Py}(Px_{\mathbf{M}})$  where  $P$  is a projection matrix that removes the center of mass of a vector [see e.g. 39, section 3.3]. For rotation, we average the likelihood across some number (# Motif Rots.) of possible rotations.
Ablation study. We first examine the impact of several parameters of TDS on success rate in an in silico self-consistency evaluation [34]. We begin with single problem (5IUS) in the benchmark set introduced by [37], before testing on the full set. Figure 3a (Left) shows that success rate increases monotonically with the number of particles. Non-zero success rates in this setting required accounting for multiple motif locations; Figure 3a (Left) uses 1,000 possible motif locations and 100 rotations.
We tested the impact of the degrees of freedom by evaluating the success rate of TDS  $(\mathrm{K} = 1)$  with increasing motif locations and 100 rotations (Figure 3a Center Left), and increasing rotations and 1,000 locations (Figure 3a Center Right). The success rate was  $0\%$  without accounting for either degree of freedom, and increased with larger numbers of locations and rotations. We also explored including a heuristic twist scale as considered for image tasks (Section 5.2); in this case, the twist scale is a multiplicative factor on the logarithm of the twisting function. Figure 3a (Right) shows larger twist scales gave higher success rates on this test case, where we use 8 particles, 1,000 possible motif locations and 100 rotations. However, this trend is not monotonic for all problems (Figure P).
Evaluation on full benchmark. We next evaluate on TDS on a benchmark set of 24 motif-scaffolding problems [37] and compare to the previous state of the art, RFdiffusion. RFdiffusion operates on the same rigid body representation of protein backbones as FrameDiff. TDS is run with  $K = 8$ , twist scale=2, and 100 rotations and 1,000 motif location degrees of freedom (100,000 combinations total).
Overall, TDS (applied to FrameDiff) and RFdiffusion have comparable performance (Figure 3b); each provides a success rate higher than the other in 11/24 cases; on two problems both methods have a  $0\%$  success rate (full results in Figure O). This performance is obtained despite the fact that FrameDiff, unlike RFdiffusion, is not trained to perform motif scaffolding. The division between problems on which each method performs well is primarily explained by total scaffold length, with TDS providing higher success rates on smaller scaffolds.
We suspect the shift in performance with scaffold length owes properties of the underlying diffusion models. First, long backbones generated unconditionally by RFdiffusion are designable with higher frequency than those generated by FrameDiff [39]. Second, unlike RFdiffusion, FrameDiff can not condition on the fixed motif sequence.
We propose TDS, a practical and asymptotically exact conditional sampling algorithm for diffusion models. We compare TDS to other approaches and demonstrate the effectiveness and flexibility of TDS on image class conditional generation and inpainting tasks. On protein motif-scaffolding problems with short scaffolds, TDS outperforms the current (conditionally trained) state of the art.
A limitation of TDS is its requirement for additional computes to simulate multiple particles. While we observe improved performance with just two particles in some cases, the optimal number is problem dependent. Moreover, the computational efficiency depends on how closely the twisting functions approximate exact conditionals, which depends on the unconditional model and conditioning information. Lastly, choosing twisting functions for generic constraints may be challenging. Our future work will focus on addressing these limitations and improving the computational efficiency.
We thank Hai-Dang Dau, Arnaud Doucet, Joe Watson, and David Juergens for helpful discussion, Jason Yim for discussions and code, and David Baker for additional guidance.