Practical and Asymptotically Exact Conditional Sampling in Diffusion Models
Authors: Luhuan Wu
Brian L. Trippe
Christian A Naesseth
David Blei
John Patrick Cunningham
Venue: NeurIPS 2023 poster (2023)
Keywords: diffusion models; conditional sampling; sequential monte carlo methods; generative models; protein design
TL;DR: The paper presents the Twisted Diffusion Sampler (TDS), an efficient algorithm for exact sampling in diffusion models, with applications to protein design.

Abstract: 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. Code can be found in https://github.com/blt2114/twisted_diffusion_sampler

Practical and Asymptotically Exact Conditional Sampling in Diffusion Models
Luhuan Wu" Columbia University 1w28270columbia. edu
Brian L. Trippe" Columbia University
Christian A. Naesseth
University of Amsterdam
blt21148columbie. edu
c. d..nesseet@uva.nn
David M. Blei
John P. Cunningham
Columbia University
Columbia University
david. blei@columiaa edu
jpc21818columbia. edu
Abstract
Diffusion models have been successful range of conditional generation tasks including molecular design and text-to-image generation. However, these achieve- ments have primarily depended on task-specific conditional training or error-prone heuristic approximations. Ideally, conditional generation method should provide exact samples for broad range of conditional distributions without requiring task-specific training. To this end, we introduce the Twisted Diffusion Sampler, TDS. TDS sequential Monte Carlo (SMC) algorithm that targets the conditional distributions of diffusion models through simulating set of weighted particles. The main idea to use twisting, an SMC technique that enjoys good computational efficiency, to incorporate heuristic approximations without compromising asymp- totic exactness. We first find in simulation and in conditional image generation tasks that TDS provides computational statistical trade-off, yielding more accu- rate approximations with many particles but with empirical improvements over heuristics with few two particles. We then turn motif- scaffolding, core task in protein design, using TDS extension to Riemannian diffusion models; on benchmark tasks. TDS allows flexible conditioning criteria and often outperforms the state-of- state-of-the-arr, conditionally trained model.2
Introduction
Conditional sampling is essential primitive in the machine learning toolkit. One begins with generative model that parameterizes distribution pa (x)) on data æ, and then augments the model to include information y in a joint distribution Pa(a, implies conditional distribution (x y). from which desired outputs are sampled.
(=))))()) xr). This joint distribution then
For example, in protein design, (z) can represent a distribution of physically realizable protein structures, a substructure that imparts desired biochemical function, and samples from po vv are then physically realizable structures that contain the substructure of interest [e.g. 34, 37].
class of generative models that have demonstrated success in conditional gen- eration tasks [18, 27, 33, 37J. They parameterize distributions pa po(x)) through an iterative refinement process that builds up data from noise. When diffusion model is used for conditional generation, this refinement process is modified to account for conditioning at each step [9, 19, 27, 33).
Diffusion models
"Equal contribution, order by coin flip. *Code: attps://githur http:///ww.....//aeeeeed.dddddd.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

One approach is to incorporate the conditioning information into training [e.g. 9, 27]; one either modifies the unconditional model take y as input or trains a separate conditional model to predict y from partially noised inputs. However, conditional training requires (i) assembling large set of paired examples of the data and conditioning information (I. 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 po();; however, the conditional training approach requires training separate models on two curated sets of conditioning inputs.
To avoid conditional training, separate line of work uses heuristic approximations that directly operate on unconditional diffusion models: once an unconditional model pe(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 asymptoti- cally exact inference in sequential probabilistic models [5, 10, 25]. SMC simulates an ensemble of weighted trajectories, 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 pe(x 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 intractable, we effectively approximate it with recent heuristic approaches to conditional sampling leg. 19], and correct the errors by the weighting mechanisms. The resulting algorithm maintains asymptotic exactness to (1 I y), and empirically it can outperform heuristics alone even with just two particles.
We summarize our contributions: (i We propose practical SMC algorithm, Twisted Diffusion Sampler or TDS, for asymptotically exact conditional sampling from diffusion models; (ii) We show that TDS applies to 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 sfate-of-the- the-art conditionally trained model.
2 Background: Diffusion models and sequential Monte Carlo
Diffusion models. A diffusion model generates a data point x by iteratively refining sequence of noisy data x', starting from pure noise x" This procedure parameterizes a distribution of x0 as the marginal of a length T Markov chain
/ p(x
T
po(x=))
II po(t=11
where p(x") ) is an easy-to-sample noise distribution, and each po (x=-1 | x²) is the transition distribu- tion defined by the (T - t)th refinement step.
Diffusion models po are fitted to match a data distribution g(x²) from which we have samples. To achieve this goal, a forward process ((x²) Ie 7(0² -;'1 is set gradually add noise to the data, where q(a² | -N (f/:af-1, - -o2), )²),, and o2 is a positive variance. To fit a diffusion model, one finds 0 such that (1 at) 2 q1="-1 x²), which is the reverse conditional of the forward process. If this approximation is accomplished for all f. and if To is big enough that g(xT) may be approximated as q(x") NN To²) =: p(øT). ) then we will have pe (2") z q(x²).
In particular, when ²² is small enough then the reverse conditionals of g are approximately Gaussian,
q(ا² x²)= = N (x' + o² o²Vz. log )0²)
2

---

where q(x") = I q(x)))))) )q(=' d²) and , log g(x") is known as the (Stein) score [33]. To mirror eq. (2), diffusion models parameterize pe pz x²) via a score network se(x",t)
po(a'~1 xx)) == N (x'=); ; + o² so (x',(()))))),)",) ,)).
(3)
When se t) is trained to approximate v. log q(x), we have (2 xx') 2 q(x"-1 ( x*).
Notably, approximating the score is equivalent learning denoising neural network (o(, t) to approximate Eg(x0) x²). The reason is that by Tweedie's formula [12, 28] Va log g(x")
(Eggzz (x" x²] - x²), //o and one can set 80 (0(2", 1) := (āg (a", () x //)/, The neural network Co(x", may be learned by denoising score matching (see [e.g. 18, 35)). For the remainder of paper
we drop the argument t in 2g and sy when 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 sequence of distributions variables 2.:T, terminating at final target of interest [5, 10, 14, 25]. SMC approximates these targets by generating a collection of K particles {=fe)==1 across T 1/
It= T-1 and weighting functions wT(x"") {wr(x', (x",x**))) ")[=0 At the initial step T, one draws K particles of T (x(x"" and sets w := =-((x)) and sequentially repeats the following for = I,...,0.
steps of an iterative procedure. The key ingredients are proposals rr (x="), {re(=" )x+
Ikk
resample
(xt+ (xt
Ik=1 ~ Multinomial(f2)" (1 | fuok tt 1k=1
ot ~ rx(" 1 ) k k= 1=1....1 K b+)),
propose
weight
wa = we 4x
k k=
The proposals and weighting functions together define a sequence of intermediate target distributions,
T-1
T-1
ve ) = C. rr (x")
II re |21
[wr(x")
II
we (x x +1
(4)
e'met
where C, is a normalization constant. A classic example that SMC applies is the state space model 125, Chapter I] that describes a distribution over sequence of latent states 0:T and observations D.T Each intermediate target V4 is constructed be the posterior p(x"-7 y yd given the first T + observations, and the final target p(x".7 (:T)) is the posterior given all observations.
The defining property of SMC is that the weighted particles each I form discrete approximations (E w - "' deg (a*) (where ö is Dirac measure) 14 (xx) that become arbitrarily accurate in the limit that many particles are used [5, Proposition 11.4]. So, by choosing fr and w, so that Po (=²) 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 given likelihood function Py)="((<>>) We embed y in joint model over o: and po u) pe (1 yy where and 1:T conditionally independent given z", Our goal is to sample from the conditional (=(2 yy)
In this section, we develop Twisted Diffusion Sampler (TDS), a practical SMC algorithm targeting Po (o("" y). First, we describe how the Markov structure of diffusion models permits factorization of an extended conditional distribution to which SMC applies. Then, we show how diffusion model's denoising predictions support the application of /wisting,
SMC technique in which one uses proposals and weighting functions that approximate the "optimal" ones. Lastly, we extend TDS to certain "inpainting" problems where Py|= (y/ (z²) 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 1:.
(x y)
02 ()
p(x"))
II 0=
(o( x(41) Potao (2)a²)
(5)
Po
Po (v(
3

---

Algorithm 1: Twisted Diffusion Sampler (TDS)
for k = 11:10
2/z =p(x"),
wk + p = Po xf) // initial proposal and weight
3 for 1. 1= T -,,.....,
(xkt) .1k )k== ~ Multinomial (fatt) Pik " Me= (w))kn) // resample
for k = 1: K
S = (fo(x( ++1 -
1+1 =//////t +V x log st+1
// conditional score approx.
# ~ 6. ,,) = N (1t + of
sk. 0²)
// proposal
ph + po x // twisting function (in eqs. (8), (13), or (14))
wi - po(2 Db (2 | .y)kk* // weight
with the desired marginal, po 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 rt (r(x)
p(x"" and re(x x x(+1) = po(x* x(+1) for 1 S t T, and weighting functions as wT(x") = wr(x", (,',,x = for 1stsT and wo(x",x') = (yl=),
Substituting these choices into eq. (4) results in the desired final target un p=("0 | ()) with normalizing constant Co = po()). As result, the associated SMC algorithm produces final set of K samples and weights {al; w =Jk=1 that provides an asymptotically accurate approximation (2 v)) = ( w Ex= 2g 2g (I") to the desired po(x:" v))
The approach above is simply importance sampling with proposal (o: 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 pa (xx y) is too dissimilar from p=(2)) only a small fraction of unconditional samples will have high likelihood: the number of particles required for accurate estimation of (x: y) is exponential in KL [po ON y) || pe (x [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 rf by multiplying the naive proposals Ft described in Section 3.1 by pe x²) as
(a o p(x?" )oo (y)=²)) and rt (a" (2 x po ")ooyl/z)))))))))))))))))))))))) for
0 05 1 <
(6)
The factors poly) x²) are the optimal twisting functions because they permit an SMC sampler that draws exact samples from (= y)) even when run with a single particle.
D:T
To see that a single particle is an exact sample, by Bayes rule, the proposals in eq. (6) reduce to
r+-xx = pe (xT yy) and rī (x² | = pe (x" | v) for
<<t<T.
(7)
As result, if one samples x" r (x*" and xt ~ (x x=-1) for = - 1,... O, by (i) the law of total probability and (ii) the chain rule of probability, one obtains dd ~ po (2 3) as desired.
However, we cannot readily sample from each r because Py(y) x²)) is not analytically tractable. The challenge is that Po(y x²) = Py|=" (y) (y]x=)) )po (x x x²)dx" depends on 2, through po xx²). ). The latter in turn requires marginalizing out x=1 w'...., 2-1 from the joint density po (2: x²), whose form depends on calls to the neural network So.

---

Tractable twisting. To avoid this intractabilityy we approximate the optimal twisting functions by
(8) which is the likelihood function evaluated at (o())) the denoising estimate of x0 at step & from the diffusion model. This tractable twisting function Po ( x²) the key ingredient needed to define the Twisted Diffusion Sampler (TDS, Algorithm 1). We motivate and develop its components below.
( 2²) Pylmo (y) = poly x²),
The approximation in eq. (8) offers two favorable properties. First, (v x²)) is computable because it depends on x' only though one call to žg, instead of an intractable integral over many calls as in the case of optimal twisting. Second, po (y/ x²) becomes an increasingly accurate approximation of (y x²), as decreases and pe p= x²) concentrates on Ep [x x²), which ăe is fit to approximate; at t = O, where we can choose (2(x)) ) x², we obtain (y x=)) = Pyl= (v|x²)
We next use eq. (8) to develop a sequence of twisted proposals po (z" | | y). to approximate the optimal proposals pa xff11 y) in eq. (7). Specifically, we define twisted proposals
Po(x" | y) = N (a": (a": xt + o0 ,x)
(9)
where so(xtt1, y) == s0(x + V t log Po x²t1)
(10)
an approximation of the conditional score, sz v) se g log pe (x y), and 8² is the proposal variance. For simplicity one could choose d² 0² to match the variance of po(a" |
Equation (9) builds on previous works (e-g. [33, 31J) that seek to approximate the reversal of a conditional forward process. The gradient in eq. (10) is computed by back- nack-propagating through 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 pa (x' | with the optimal twisted proposal Po (x 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
3) will not in general coincide
(x x (p x²)
(2", (21)
t= 0, O T-11
(11)
Po (o(+))fu (xx' xx(+1 .y"
and (x = Po (y )²)) 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 My (x. [ y). Substituting into eq. (4) each py (x' 1x't1,0 y) in eq. (9)
for re (x (f+) and wx (x' + in eq. (11) and then simplifying we obtain the intermediate targets
Po(y x²)
Po (a="
|xx 11++1 .))
(x(T x po (x v)
(12)
Po xx)
Po (x"
The right-hand bracketed term in eq. (12) can be understood as the discrepancy of 14 from the final target My accumulated from step & to 0 (see Appendix A.3 for derivation). As approaches 0, poly | w²)) improves as an approximation of po x²), 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 = O, because (o x²))= pa | æ²) by construction, eq. (12) reduces to vo(x"T)= = po (x. y), as desired.
The TDS algorithm and asymptotic exactness. Together, the twisted proposals pe (xf x y) and weighting functions (a'', xx+1) lead to Twisted Diffusion Sampler, TDS (Algorithm 1). While Algorithm states multinomial resampling for simplicity, in practice other resampling strategies (c.g. systematic [5, Ch. 9I) may be used as well. Under additional conditions, TDS provides arbitrarily accurate estimates of (1 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 P (x(x=) = -EK -' wad (x") denote the discrete measure defined by the particles and weights returned by Algorithm I with K particles. Under regularity conditions on the twisted proposals and weighting functions, IP (x()) converges setwise to po(2 1 as K approaches infinity.
5

---

3.3 TDS for inpainting, additional degrees of freedom
The twisting functions Po x²) :- Pvj (y (x')) we introduced above are one convenient option, but are sensible only when Py د (y) 2g (z")) 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 provides additional details, including the adaptation of TDS to variance preserving diffusion models [18].
Inpainting. Consider the case that œ can be segmented into observed dimensions M and unobserved dimensions M such that we may write x² = |>* M and let TM. and take Palno (y|z²) = (x)) The goal, then, is sample pe(x | y) po =21) (2 Here we define the twisting function for each t > 0 as
Po x M) := N (1: ( (MM")),.t))) to?)
(13)
and set twisted proposals and weights according to eqs. (9) and (11). The variance in eq. (13) is chosen as to² = Var pa [x=' =0) for simplicity; in general, choosing this variance more closely match Var,. [y/ x²| may be preferable. For t = 0, we define the twisting function analogously with small positive variance, for example as Po [y x²) = Ny: x o²). 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 M in mind. To handle this situation, we (i) let M be a set of possible observed dimensions, (ii) express our ambivalence in which dimensions are observed as y using randomness by placing uniform prior on P(M) = 1/|MM] for each M E M,, and (iii) again embed this new variable into our joint model to define the degrees
of freedom likelihood by p(y x²)) = ΣMEM Pe PM, y x0) = [M] EMEM p(y x .)) Accordingly, we approximate po (y x²)) with the twisting function
poly x²) = |M
po(y x² MM) MEM
Σ
(14)
with each Py x²,, ,M) defined as in eq. (13). Notably, eqs. (13) and (14) coincide when IM 1.
The sum in eq. (14) may be computed efficiently because each term depends on x' only through the same denoising estimate (xx), which must be computed only once. Since computation of i Ex(2") is the expensive step, the overall run-time is not significantly increased by using even large [.M..
3.4 TDS 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 manifold projection step (see e.g. [4. 8)). When we assume that (as, e.g., in [39]) the model is associated with denoising network fe. twisting functions are also constructed analogously. For conditional tasks defined by likelihoods, Po(y x²) in eq. (8) applies. For inpainting (and by extension, degrees of freedom), we propose
( x²) = TN (z (y: 0, to²),
(15)
where TN (o("") (a (0, to²) is tangent normal distribution centered on fo(x") As in the Eu- clidean case, Po (x' x=++11)) v) is defined with conditional score approximation sp(a", .y) sa(x²) + Va log ply xt), which is computable by automatic differentiation. Appendix B provides details.
4 Related work
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 general approach that incorporates conditioning information with gradients through the neural network fz(x"). For example, in an

---

Smooth Likelinood
Inpaiit
Inpaint (Degrees of Freedom)
10-1
Mian Square Error
10
10-2
TDS-5
10
70
Gradient
Guidence
10-3
10"3
10"3-
9MC-DIf
Replacement
10
10"1
10-1
10a
10a
10a
10
10a
10a
10
10*
10a
10a
10a
10
of Porticles (K)
# of Particles (K)
Particles (K)
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 N(y (1)) a:) to pa x²). This approximation motivates modified transition distribution as in Equation (9), with the corre- sponding approximation to the conditional score as sø , y) = s
V 2g (x (x') 1/ll//a This approach coincides exactly with TDS applied to the inpainting task when 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 -on-Saasiinn approximations; e.g. using Pala (v fo((")) x po(y ) for given likelihood Pyl= (y) )²) 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 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 zMM- with a noisy version of observation "M- However, it is a heuristic approximation and can lead 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 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 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 [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 set- tings 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, naive importance sampler described in Section 3.1; and (4) Gradient Guidance. For inpainting problems we further include: (5) Replacement method; and (6) SMC-Diff.
Each SMC sampler forms an approximation to pe (2 y) with K weighted particles (EK= ) ΣK_ wxd.g- Gradient guidance and replacement method are considered to form a similar particle- based approximation with K independent samples viewed K particles with uniform weights.
Compute Cost: Compared to unconditional generation, TDS has compute cost that (i) larger by a constant factor due to the need to backpropogate 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 be achieved.
7

---

By comparison, conditional training methods provide fast inference by amortization [13], in which most computation 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-firont 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 pa (2²) approximation of a bivariate Gaussian. For this choice, the marginals of the forward process are also Gaussian, and so we may define po (a:"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: (I) Smooth likelihood: y is an observation of the Euclidean norm of with Laplace noise, with po(y x=9) exp(|||=912> y/}/2
This likelihood is smooth almost everywhere." (2) Inpainting: y is an observation of the first
dimension of z², with ( I², M 0) = 6y (=g) (3) Inpainting with degrees-of- ofreoedoooorro:: is
observation of either the first or second dimension of z", with M = {0, 1} and pe(y x²)
(x)) + (x.)].. In all cases we fix y = O and consider estimating Ep. [x, [y]
Figure 1 reports the estimation error for the mean of the desired conditional distribution, i.c. E Eg 1/112- TDS provides a computational-staine statistical trade-off: using more par- ticles 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 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 classcconditional image generation task, which involves sampling an image from (z y) x po(x=") >p)}=0 (y(=="), where (o(-) is pretrained diffusion model on images z.", is given image class, and Py 00 (vl-) is the classification likelihood. To assess the faithfulness of generation, we evaluate classification accuracy on predictions of conditional samples x 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 P pre-trai diffusion model and pretrained classifier on the ImageNet dataset (256 x 256 x 3 dimensions). These samples are qualitatively good and capture the class label. Appendix D.2.3 provides more samples, comparision to Classifier Guidance [9], and results for the CIFAR- 10 dataset.
This likelihood is smooth except the point (0, 0).
8

---

Class. Accuracy (%6
100
7
7
7
7
9
2
2
7
7
7
7
7
7
7 2
8
3
TDS4S
Gradient Guidance
7
7
7
3
7
0
7
7
7
8
4
4
32 # of Particles
TDS
Gradient Guidance
(a) MNIST: samples by TDS (K = 64)
(h) MNIST: classification accuracy vs.
(c) ImageNet: samples by TDS (K = 16) given class "brambling' from 4 independent SMC runs.
and Gradient Guidance given class '7'
K. Results are averaged over 1,000 runs with error bars denoting stan- dard errors.
TDS samples are randomly selected out
of 64 particles in single SMC
Figure 2: Image class-conditional generation task.
TDS can be extended by exponentiating twisting functions with twist scale. This extension is related to the existing literature Classifier Guidance that considers re-scaling the gradient of the log classification probability. See Appendix D.2.1 for details and ablation study.
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 motif, that 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 application of TDS this task, and compare TDS to the state-of-the-art conditionally-trained model, RFdiffusion. See additional details in Appendix E.
Given generative model supported designable protein structures pe (2xx), suitable scaffolds may
be constructed by solving conditional generative modeling problem [34]. Complete structures
first segmented into motif dM and a scaffold z i.c. xx [0 Ma Putative compatible
scaffolds are then identified by (approximately) sampling from po x [34].
While the conditional generative modeling approach to scaffo-caffoliing has produced functional, experimentally validated structures for certain motifs [37]. the general problem remains open. More- over, current methods for motif-scaffolding require one to specify the location of the motif within the primary sequence of the full scaffold; this choice require expert knowledge and trial and error.
We hypothesized that improved soaff-sccafolding could be achieved through accurate conditional sampling. To this end, we applied TDS to FrameDiff, Riemannian diffusion model that param- eterizes protein backbones as a collection of rigid bodies (known as residues) in the manifold SE(3) [39].* Each of the N elements of S SE3) consists of rotation matrix and 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 Pyle (y)==)) = 6, (>)) where y E SE(3) describes the coordinates of backbone atoms of an = |M residue motif. As such, to define twisting function 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 M, (ii) let. M be 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, is known that all motif residues must appear with contiguous indices. In this case we choose M to be the set of all possible contiguous masks. However, when two more possibly discontiguous blocks, the number of possible placements can be too large and we choose M by randomly sampling at most some maximum number (# Motif Locs. )) of masks.
motif residues are only known to appear
We similarly eliminate the global translation and rotation of the motif as a degree freedom. The motivation is that restricting one possible pose of the motif narrows the conditional distribution,
'nttps://githur http:///tthuu..cmmmasso.eess htt.l//cc//ooyyyui///eeeee. diffuusiio

---

25
Success rate (%)
20-
Scaffold
&
15-
size
FrameDiff
diffusion
10-
<100 res.
9
3
2. 100 res.
2
8
Overall
16
10 100 10000
10 100 1000
0.50.75 1. 20 300
Particles
Motif Locs.
Motif Rots.
Twist Scale
(b) # problems with higher success rate
(a) TDS motif-scaffolding success (test 5IUS) improves with more particles, and degrees of freedom, and twist-scale.
Figure 3: Protein motif-scaffolding case study results
thereby making inference more challenging. For translation, we use likelihood that is invariant to the motif's center-of-mass and placement in the final scaffold by choosing Pyl= (u|z) - 6p (PPM) where P is a projection matrix that removes the center of mass of 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 in silico self-consistency evaluation [34]. We begin with single problem (5IUS) in the benchmark set introduced by [37], before testing 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 TDS (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 multiplicative factor 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 protein backbones as Frfme 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 FrfmeDifff, 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 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 REdiffusion, FrameDiff can not condition on the fixed motif sequence.
Discussion
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.
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 challenging. Our future work will focus addressing these limitations and improving the computational efficiency.
10

---

Acknowledgements
We thank Hai-Dang Dau, Arauu Doucet, Joe Watson, and David Juergens for helpful discussion, Jason Yim for discussions and code, and David Baker for additional guidance.
References
[1] Arpit Bansal, Mong-Min Chu, Avi Schwarzschild, Soumyadip Sengupta, Micah Goldblum. Jonas Geiping, and Tom Goldstein. Universal guidance for diffusion models. arXiv preprint arXiv:2302.0722/, 2023.
[2] Gabriel Cardoso, Yazid Janati El Idrissi, Sylvain Le Corff, and Eric Moulines. Monte Carlo guided diffusion for Bayesian linear inverse problems. arXiv preprint arXiv:2308.07983 2023.
[3] Sourav Chatterjee and Persi Diaconis. The sample size required in importance sampling. The Annals of Applied Probability, 2018.
[4] Gregory Chirikjian and Marin Kobilarov. Gaussian approximation of non-linear measurement models on Lie groups. In 53rd IEEE Conference on Decision and Control, IEFE, 2014.
[5] Nicolas Chopin and Omiros Papaspiliopoulos. An introduction to sequential Monte Carlo. Springer, 2020.
[6] Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Marc L Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. International Conference on Learning Representationss 2023.
[7] Justas Dauparas, Ivan Anishchenko, Nathaniel Bennett, Hua Bai, Robert Ragotte, Lukas Milles, Basile IM Wicky, Alexis Courbet, Rob J de Haas, Neville Bethel, , Philip J Leung, Timothy Huddy, Sam Pellock, Doug Tischer, F Chan, Brian Koepnick, H Nguyen, Alex Kang, B Sankaran, Asim K. Bera, Neil P. King, and David Baker. Robust deep bearning-based protein sequence design using ProteinMPNN. Science, 2022.
[8] Valentin De Bortoli, Emile Mathieu, Michael John Hutchinson, James Thornton, Yee Whye Teh, and Arnaud Doucet. Riemannian score-based generative modelling. Advances Neural Information Processing Systems, 2022.
[9] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat GANs on image synthesis. Advances in Neural Information Processing Systems, 2021.
[10] Arnaud Doucet, Nando De Freitas, and Neil James Gordon. Sequential Monte Carlo methods in practice. Springer, 2001.
[11] Yilun Du, Conor Durkan, Robin Strudel, Joshua B Tenenbaum, Sander Dieleman, Rob Fergus, Jascha Sohl-Dickstein, Arnaud Doucet, and Will Grathwohl. Reduce, reuse, recycle: Composi- tional generation with energy-based diffusion models and MCMC. In International Conference on Machine Learning, 2023.
[12] Bradley Efron. Tweedie's formula and selection bias. Journal of the American Statistical Association, 2011.
[13] Samuel Gershman and Noah Goodman. Amortized inference in probabilistic reasoning. In Proceedings of the Annual Meeting of the Cognitive Science Society, 2014.
[14] Neil J Gordon, David Salmond, and Adrian FM Smith. Novel approach to nonlinear/non- Gaussian Bayesian state estimation. In IEE Proceedings F (Radar and Signal Processing). IET, 1993.
[15] Pieralberto Guarniero, Adam M Johansen, and Anthony Lec. The iterated auxiliary particle filter. Journal of the American Statistical Association, 2017.
[16] Jeremy Heng, Adrian Bishop, George Deligiannidis, and Arnaud Doucet. Controlled sequen- tial Monte Carlo, Annals of Statistics, 2020.

---

[17] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv:2207.15595 2022.
arXiv preprint
[18] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, 2020.
[19] Jonathan Ho, Tim Salimans, Alexey A Gritsenko, William Chan, Mohammad Norouzi, and David Fleet. Video diffusion models. In Advances in Neural Information Processing Systems, 2022.
[20] John M. Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, Alex Bridgland, Clemens Meyer, Simon A A Kohl, Andy Ballard, Andrew Cowie, Bernardino Poreee-, Stanislav Nikolov, Rishub Jain, Jonas Adler. Trevor Back, Stig Petersen, David A. Reiman, Ellen Clancy, Michal Zielinski, Martin Steinegger, Michalina Pacholska, Tamas Berghammer, Sebastian Bodenstein, David Silver, Oriol Vinyals, Andrew W. Senior, Koray Kavukcuoglu, Pushmeet Kohli, and Demis Hassabis. Highly accurate protein structure prediction with AlphaFold. Nature, 2021.
[21] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. In Advances in Neural Information Processing Systems, 2022.
[22] Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. In Advances in Neural Information Processing Systems, 2022.
[23] Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool. Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, 2022.
[24] Chenlin Meng. Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun-Yan Zhu, and Stefano Ermon. SDEdit: Guided image synthesis and editing with stochastic differential equations. In International Conference on Learning Representations, 2021.
[25] Christian Naesseth, Fredrik Lindsten, and Thomas Schön. Elements of sequential Monte Carlo. Foundations and Trends in Machine Learning. 2019.
[26] Michael K Pitt and Neil Shephard. Filtering via simulation: Auxiliary particle filters. Journal of the American Statistical Association, 1999.
[27] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with CLIP latents. arXiv preprint NKXv:22044001155, 2022.
[28] Herbert Robbins. An empirical Bayes approach to statistics. In 3rd Berkeley Symposium on Mathematical Statistics, 1956.
[29] Chitwan Saharia, William Chan, Huiwen Chang, Chris Lee, Jonathan Ho, Tim Salimans, David Fleet, and Mohammad Norouzi. Palette: Image-to-image diffusion models. In ACM SIGGRAPH 2022 Conference Proceedings, 2022.
[30] Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David Fleet, and Mohammad Norouzi. Image super-resolution via iterative refinement. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2022.
[31] Yuyang Shi, Valentin De Bortoli, George Deligiannidis, and Arnaud Doucet. Conditional simulation using diffusion Schrödinger bridges. In Uncertainty in Artificial Intelligence, 2022.
[32] Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. Pseudoinverse-guiee diffusion models for inverse problems. In International Conference on Learning Representations, 2023.
[33] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2020.
12

---

[34] Brian L Trippe, Jason Yim. Doug Tischer, Tamara Broderick, David Baker, Regina Barzilay, and Tommi Jaakkola. Diffusion probabilistic modeling protein backbones 3D for the motif-scaffolding problem. In International Conference on Learning Representations, 2023.
[35] Pascal Vincent. A connection between score matching and denoising autoencoders. Neural Computation, 2011.
[36] Jue Wang, Sidney Lisanza, David Juergens, Doug Tischer, Joseph L Watson, Karla M Castro, Robert Ragotte, Amijai Saragovi, Lukas F Milles, Minkyung Back, Ivan Anishchenko, Wei Yang. Derrick R Hicks, Marc Exposit, Thomas Schlichthaerle, Jung-Ho Chun, Nathaniel Dauparas, Justas Bennett, Basile I M Wicky, Andrew Muenks, Frank DiMaio, Bruno Correia, Sergey Ovchinnikov, and David Baker. Scaffolding protein functional sites using deep learning. Science, 2022.
[37] Joseph L. Watson, David Juergens, Nathaniel Bennett, Brian Trippe, Jason Yim, Helen Eisenach, Woody Ahern, Andrew J. Borst, Robert J. Ragotte, Lukas F. Milles, Basile I. M. Wicky, Nikita Hanikel, Samuel Pellock, Alexis Courbet, William Sheffler, Jue Wang, Preetham Venkatesh, Isaac Sappington, Susana Vázquez Torres, Anna Lauko, Valentin De Bortoli, Emile Mathieu, Sergey Ovchinnikov, Regina Barzilay, Tommi Jaakkola, Frank DiMaio, Minkyung Baek, and David Baker. De novo design of protein structure and function with RFdiffusion. Nature, 2023.
[38] Nick Whiteley and Anthony Lee. Twisted particle filters. The Annals of Statistics, 2014.
[39] Jason Yim, Brian L Trippe, Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. SE(3) diffusion model with application protein backbone generation. In International Conference on Machine Learning, 2023.
[40] Guanhua Zhang, Jiabao Ji, Yang Zhang, Mo Yu, Tommi Jaakkola, and Shiyu Chang. Towards co- herent image inpainting using denoising diffusion implicit models. In International Conference on Machine Learning, 2023.
13

---

Appendix
diffusion models define the forward A Twisted Diffusion Sampler additional details
14
B Riemannian Twisted Diffusion Sampler
18
Additional Related Work
20
Empirical results additional details.
E Motif-scaffolding application details.
28
A Twisted Diffusion Sampler additional details
In this section we provide additional details on TDS. Appendix A.1 describes its generalization to alternative common formulations of diffusion models. Appendix A.2 describes the choices of proposal variance and resampling strategy. Appendix A.3 provides a derivation of Equation (12). Appendix A.4 describes modifications the final step of TDS in inpainting and inpainting with degrees of freedom applications. And Appendix A.5 provides a full statement and proof Theorem 1.
A.1 Diffusion models with non-constant noise variance and variance-preserving scaling
TDS developed in Section 3 assumed noise with constant variance o² added each timestep. This choice corresponds to a particular discretization of a "variance exploding" (VE) diffusion model [33]. Here describe generalizations of TDS to non-constant variance schedules, and variance preserving (VP) diffusion models [18]. We adopt the VP formulation in our experiments; though introduces additional notation, is known work well in practice [33].
TDS algorithm for variance exploding diffusion models. process by
1x 1:
II q((" x'-1),
g(0² x x2-1)) =
(16)
where of is an increasing sequence of variances such that g(x²') R ((.,o))) with of = Ei= of
for = 1.... T. And so one can set p(x"") = ((,,7))) to match ((x") Notably, Equation (16)
implies the conditional q(=" x²") = Nax x of).
The reverse diffusion process pa is parameterized as
pa (z = p()") II po (x xx),
Pa Pa
2 - + +z)s0(2',t))
where the score network 98 is modeled through denoiser network igg by sæ (x = (ag(x' : t; t) o0) Note that the constant schedule special case where of 0², and o to for all
The TDS algorithm for general VE models described Algorithm 1, where to? in Line is replaced by of and o² in Line 7 is replaced by of
Extension to variance preserving diffusion models. Another widely used diffusion framework is variance preserving (VP) diffusion models [18]. VP models define the forward process by
(ax 1:T x²)= = II g(a" xt-1),
q(*² x1-1) = Na' v afx' 1-1 .o,)
where of is a sequence of increasing variances chosen such that g(x²) NO0, 1), and so can set p(x-)= = A/ 1). Define 04 = 1 07, 0g == Ti ov,, and o? == 1 ã.. Then the marginal conditional of eq. (16) q(2" 20)) NNxx; vara" " The reverse diffusion process pe is parameterized as
of).
Nix'=1; x' of Vas Vau
P:T == p(x")
p0 xz'), po(z--1 | x²)
so(x", .tlftt)),o ,t),
14

---

where is now defined through the denoiser ig by sø(z', t) == (v", t; x²) )0??
TDS in Algorithm extends to VP models as well, where the conditional score approximation in Line 6 is changed to < (( = /// + V z log Night1 and the proposal in +1 Line 7 is changed to 2t po(: 1 +1 .9) == N (arlt
vau+1 8k. 02
A.2 TDS parameters
Proposal variance. The proposal distribution in Line of Algorithm is associated with variance parameter o² In general, this parameter can be dependent on the time step, i.e. replacing o² by some 0i-1 in Line Unless otherwise specified, we set 0[+1 == Varp, [a [x*+]] the variance of the unconditional diffusion model (typically learned along with the score network during training).
Resampling strategy. by other strategies, see [25, Chapter 2] for an overview. In our experiments, we use the systematic resampling strategy.
The mulinomial resampling strategy in Line 4 of Algorithm can be replaced
In addition, one consider setting effective sample size (ESS) threshold (between and 1), and only when the ESS is smaller than this threshold, the resampling step is triggered. ESS thresholds for resampling are commonly used to improve efficiency of SMC algorithms [see c.g. 25, Chapter 2.2.2]. but for simplicity we use TDS with resampling at every step unless otherwise specified.
A.3 Derivation of Equation (12)
Equation (12) illustrated how the extended intermediate targets (x::TT) provide approximations to the final target po y) that become increasingly accurate as f approaches We obtain Equation (12) by substituting the proposal distributions in Equation (9) and weights in Equation (11)
into Equation (4) and simplifying.
(x" 1x 11) (1) x² Po x=+1 (00 ot x
ve(x::T)
p(a")
II
po(a"
ilu x+)
II e't
P==
Rearrange and cancel po and po terms.
=(pla*) plaa)) II (o(a" (1)
T-1
poly .²4) II Fo | |
Group pe terms by chain rule of probability.
= po f:T (oo( x²) II
pola" 1 .1)
1'=0
Apply Bayes' rule and note that po (v x = po x²).
x²) x²)
1-1
(2 I. (v)
II
Po
p(v
===
C'11
Note that (x 0. v) = pa
(2) II po(a"
|oft1 v).
Poly x²)
t-1
Po(a" +1
= = 0:T
())
p(u
x²) M=0
(a 1
The final line is the desired expression in Equation (12).
A.4 Inpainting and degrees of freedom final steps for asymptotically exact target
The final 0) twisting function for inpainting and inpainting with degrees of freedom described in Section 3.3 do not satisfy the assumption of Theorem 2 that po(y) x²) Py|= (y)=.)) This choice introduces error in the final target of TDS relative to the exact conditional po(2 y).
15

---

For inpainting, to maintain asymptotic exactness one may instead choose the final proposal and weights
po x0 y; M) == MM
M xx) and wn(x", ,x')
=
One can verify the resulting final target is po Similarly, for impainting with degrees of freedom, one may define the final proposal and weight as (o M xt M)
#M = y) according to eq. (4).
Fo(x" |a M) = Σ
EM'EM Po Po x! M') by (o
x²)) and wo(x" x
MEM
A.5 Asymptotic accuracy of TDS additional details and full theorem statement
In this section we (i) characterize sufficient conditions on the model and twisting functions under which TDS provides arbitrarily accurate estimates as the number of particles increased and (ii) discuss when these conditions will hold in practice for the twisting function (y x²) introduced in Section 3.
We begin with theorem providing sufficient conditions for asymptotic accuracy of TDS.
Theorem 2. Let po ( be diffusion generative model (defined by eqs. (1) and (3)) with
po(x' x=+1)== N (x" 1x+|
)),
with variances of....,," .². Let po(y x²) be twisting functions, and
re(a' [ N (x* t + 1041100 (2) + Vatt1 log ((+1)),00 be proposals distributions for = O...., T and let Px Ek u q for weighted particles
{(2,, b returned by Algorithm I with K particles. Assume
(a) the final twisting function is the likelihood, po (y 2² = Poleo
(b) the first twisting function Poly x²), and the ratios of subsequent twisting functions poly x²)) x')/po(y) x((1) are positive and bounded,
(c) each log poly (²)) with t continuous and has bounded gradients in x². and
(d) the proposal variances are larger than the model variances, i.e. for each t, o7 > of.
Then PK converges setwise to pe(x y) with probability one, that is for every set A, 1im (A) JAPo(> yydx"
The assumptions of Theorem are readily satisfied in common applications.
Assumption (a) may be satisfied by construction by choosing (o |00) = Pulwo (v) by
defining &g (2 P, 0) = x
Assumption (b) is satisfied (i) Pyla (ylæz) is smooth in x and everywhere positive and (ii) po(y x²) Pylx" (xl()(=*)) where fe(x*) takes values in some compact domain. An alternative sufficient condition for Assumption (b) is for Pyla (y]z) to be positive and bounded away from zero; this latter condition will hold when, for example, Py x (2/z) is classifier fit with regularization.
Assumption (c) is the strongest assumption. will be satisfied, for example, if (i) Poly x²) = Pyj (y)lās (x*)), and (ii) Pola" (y/zz) and (e(x)) are smooth in and æf, with uniformly bounded gradients. While smoothness of &e (:(.))) can be encouraged by the use of skip-connections and regularization, io(-, t) may present sharp transitions, particularly for t close zero.
Assumption (d), that the proposal variances satisfy o7 > of is likely not needed for the result to hold. However, this assumption permits usage of existing SMC theoretical results in the proof; in practice, our experiments use o = =o but alternatively the assumption could be met by inflating each oう by some arbitrarily small 8 without markedly impacting the behavior of the sampler.
16

---

Theorem 2 characterizes set of conditions under which SMC algorithms converge. We restate this result below in our own notation.
Proof of Theorem 2:
Theorem 3 (Chopin and Papaspiliopoulos [5] Proposition 11.4). Ler {< w be the particles and weights returned at the last iteration of a sequential Monte Carlo algorithm with particles using multinomial resampling. If each weighting function (2', x)+1) is positive and bounded, then for every bounded, uy-measurable function @ of x!
#R((=))= = - p())))((=")
lim K K-000
with probability one.
An immediate consequence of Theorem the setwise convergence of the discrete measures, PK = Σ w This can be seen by taking for each ((x) I[ E A] for any 1y Vo-measurable set A. The theorem applies both in the Euclidean setting, where each 2 E RD,, well as the Riemannian setting.
We now proceed to prove Theorem 2.
Proof. To prove the theorem we show (i) the x² marginal final target vo pa(a") [y) and then (ii) Fx converges setwise to Vo.
We first show by manipulating My in Equation (4) to obtain po (x:21 y). From Equation (4) first have
vo(x
II
wT(x""
we (x", sett1) II x=
Lo
1=0
Substitute in weights from eq. (11).
1
p(æ"
T-1
II t=0
Po x²)
Po(a* | Io x²)
Eo
[= Po (y creet 2o(+1)
Rearrange and re terms, and cancel out re terms.
T-1 II t=
re (a 1 (2 x²)
p(x*)
(o(x* s polul x² II
co
Do x |اااي= + tzo
t=u
Collapse Pa terms.
Polu x²) II
Co co (
Bolu | x²))
Po wt1l)
==
Cancel out Pa terms.
= Co po(o (o x²)
Recognize Po(y | x0) = (y)z) by Assm, (a) and apply Bayes' rule with Co = po(y). = =p v)-
The final line reveals that once we marginalize out 1:7 we obtain 10(=))= = p(x0) v) as desired.
We next show that Px converges to Ye with probability one by applying Theorem 3. To apply Theorem it sufficient to show that the weights each step are upper bounded, as they are defined through (ratios of) probabilities and hence are positive. Since there are finite number of steps T, it is enough to show that each we is bounded. The inital weight is the initial twisting function, wT(x") py x a²),, which is bounded by Assumption (b). So we proceed to intermediate weights To show that the weighting functions at subsequent steps are bounded, we decompose the log- weighting functions as
log ",(a", )+ = log po x!))
Po(y x(+1] + log lo | | re(x x(+1}¹
17

---

and show independently that log py [ x and logpo(z') x²* 1/rr(x" | x are
bounded. The first term log poly x²)) p/o) (y xx+t) is again bounded by Assumption (b), and we proceed to the second.
That log po /rrtt)/r((a' 1 is bounded follows from Assumptions (c) and (d). First write
po(a* | = N (a' A. (
with j = w/+1 02-180(x0 ), and
(a ( N (x') Av, 11
for P& = ft + 07t1 V Vnt+i log (3 x 11+1).. The log-ratio then simplifies as
log Polx' (xtt1
log 12 11 1/2 -207((207+1)
rx x4+1)
1 - -20 [a - x² =²
Rearrange and let C = log | /12 11
-1 [o - - -x'1' x'113] +C
2
Expand and rearrange || x'11 = |1 x²11 + 2(p io f x²) + 11 -ill?
= 2 | - - x' x²12 22,22,2 (Ae - ju. i - x²) - ,2+ 1111. - -a|131 + C
Let C'=C 1
3 ~| and rearrange. Note that - 2 by Assm. (c).
-1
2 (0++1 ,, - x'1
3 11 - u. x+)++++
Apply Cauchy-Schwer
(o - f - - 3
- ja|l . 11ā - x²11 + C"
2
Upper-bounding using that max
-a x² br
b²
for -2 - - > 0 by Assm. (d).
0
2
2a
( - =ll)2
$
+ C"
2 0.41
-
0,41
0,41
1:p - Al| + C
2(0,71 2
0
Note that jup = p:
+
0 nt log 5o (y x t11
3,41
- loggo wt C"
2(0 0.71)
CC'
The final line follows from Assumption (c), that the gradients of the twisting functions are bounded. The above derivation therefore provides that each is bounded, concluding the proof.
B Riemannian Twisted Diffusion Sampler
This section provides additional details 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 modifies to this setting.
The tangent normal distribution. Just as the Gaussian is the parametric family underlying Eu- clidean diffusion models in Equation (3), the tangent normal (see e.g. [4, 8)) underlies generation in Riemannian diffusion models so review it here.
18

---

We take the tangent normal to be the distribution is implied by two step procedure. Given a variable in the manifold, the first step is to sample a variable ÿ in Tz. the tangent space z; if ( ................................ h is an orthonormal basis of Ta one may generate
y p
oed hds
with Aid. (00, 1), and 0² >> 0 variance parameter. The second step is to project back onto the manifold to obtain y = exp, op {ū} where exp, {.} denotes the exponential map at æ. The resulting distribution of is denoted as T.N (u, o²). By construction, Tz is Euclidean space with its origin at x, so when || = 0, T=N o²) is centered on æ. And since the geometry of a Riemannian manifold is locally Euclidean, when o² is small the exponential map is close to the identity map and the tangent normal is essentially a narrow Gaussian distribution the manifold at æ. Finally, we use TN p1, 2 to denote the density of the tangent normal evaluated at y-
Because this procedure involves a change of variables from ÿ to y (via the exponential map), to compute the tangent normal density one computes
i
T o²) > (y); o²)
expp"(yy) {4)
dy
where expe [y} is the inverse the exponential map (from the manifold into Tw), and By cxpx {v} the determinant of the Jacobian of the exponential map; when the manifold
lives higher dimensional subset of R, we take by exp= {y) 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 øf in Tgur; (fi,..., h is an orthonormal basis of Tzes1 one may generate
D
x' = =00110((x0t)
Σ
OetiEd has
d=1
with cd sid. (0, 1). One then projects x' back onto the manifold to obtain of = exp2111 (7*) where exp.t-} denotes the exponential map at æ.
This is equivalent to sampling x² from a tangent normal distribution as
x² ~ p(=" (ttl) = TN 2011
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 manifold-valese
denoising estimate åæ as will be the case when, for example, sa(x", ((a",, t) := v, log q(x* x = 2g) for ie = Le(x", .t) In contrast to the Euclidean case, relationship between denoising estimate and a computationally tractable score approximation may not always exist for arbitrary Riemannian manifolds; however for Lie groups when the the forward diffusion is the Brownian motion, tractable score approximations exist [39, Proposition 3.2].
For the case of positive and differentiable Pyl= ()) we again choose twisting functions Po(y x²) = Pyluo (ylāa(a")
Next are the inpainting and inpainting with degrees of freedom cases. Here, assume that z lives on multidimensional manifold (e.g. SE(3) N) and the unmasked observation y 2 with M c (1...., N) on lower-dimensional sub-manifold (e.g. MM with 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 Tangent normal
poly | (o
For all cases, we propose the twisted proposal as po xt+1, ,y)=TNat (x" o 0 .3), 021)
(17)
19

---

where as in the Euclidean case sa(x", y) = sa(x²) + Vaa p(y x²).
Weights at intermediate steps computed in the Fuclidean case (Equation (11)):
po (pel)) x²) No(y xx+ x (10 x
Wy (a", (
TN 2- (s"; 02t1 so = 1 poly x²)
po(y IT gt+h (x'; s0 y), 0
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 Line 6. Here we assume that the unconditional score is computed through the transition density function:
so so(x == 2et+ log 10
ie)) for Po (x(
Note that the gradient above ignores the dependence of dhe on x1+1 The conditional score approximation in Line is then replaced with
5% < (- x
++1 .) = V x log ( ++1 [Ta) + V at log 2
for de = (z (1t1))
Notably, 5% is T valued Riemannian gradient.
The second change to make the proposal on Algorithm Line a tangent normal, as defined in eq. (17).
C Additional Related Work
Training with conditioning information. Several approaches involve training a neural network 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-
[9, 33], an approach for generating samples of a desired class, which requires training a time-dependent classifier approximate, for each f, po (y/ x²), and training such time-dependent classifier may be inconvenient and costly), and (iv) classifier-free guidance [17], an approach that builds on classifier-guidane without training a classifier, and instead trains diffusion model with class information additional input.
Langevin and Cetropolis-Hastiine 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 asymp- totically 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 th 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 difflusion models on two dimensional problems
Forward process. Our forward process is variance preserving (as described in Appendix A) with
T = 100 steps and quadratic variance schedule. We set of = o + 2or max with o 10-5 and o max 1
Unconditional target distributions and likelihood. different unconditional target distributions:
We evaluate the different methods with two
20

---

Smooth Likelinood
Inpaint
Inpaint (Degrees Freedom)
10-1
10
10
10-2
10 2
10
TDS
TDS-S
10-3
10"2
10 --
Gradient
10
SMG-DIf
10-1
10
Hep
10a
10
101
10
10
10
Porticles
of Particles (K)
Particles
Figure D: Errors of conditional mean estimations with 2 SEM error bars averaged over 25 replicates on mixture Gaussians unconditional target. TDS applies to all three tasks and provides increasing accuracy with more particles.
1. A bivariate Gaussian with mean at (5 4) and covariance 0.9 and
2. A Gaussian mixture with three components with mixing proportions [0.3, 0.5, 0.2), means [(1.4, 0.229, (-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 uncon- ditional target. In this second example we evaluate a with a variation the inpainting degrees of freedom case wherein consider y 1 and M {[1), [2]}, so that Py a (y)x²)) ( (x2)-
D.2 Image conditional generation experiments
We perform extensive ablation studies on class-conditional generatoin task and inpainting task using the small-scale MNIST dataset (28 28 x dimensions) in Appendix D.2.1 and Appendix D.2.2. And we show TDS's upplicability to
datasets, namely CIFAR10 (32 x x3 dimensions) and ImageNet256 (256 x 256 x 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 the guided diffusion codebase' with the following specifications: number of channels 64, attention resolutions = "28.14,7" number of residual blocks 3. learn sigma (i.e. learn the variance of Pe pa x²))= True, resblock updown True, dropout 0. 0, variance schedule = "linear". We trained the model for 60k epochs with batch size of 128 and 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 Mlass-conditionn generation and evaluation is pretrained ResNet50 model. This classifier is trained on the same set of MNIST training images used by diffusion model training.
In addition we include variation called TDS-truncate that truncates the TDS procedure at = 10 and returns the prediction de (1:
Sample plots. To supplement the sample plot conditioned on class 7 in Figure 2a, we present samples conditioned on each of the remaining 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 or TDS-IS are identical due to 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 twist seale 7/, i.e. setting new twisting functions to Po =f:t) o Pyj= (y) (z'))". In particular, when = 0, we set Py "y/ax" (ylx";>) x (y) This modification suggests that the targeted conditional distribution now
p=(a0 7: po ) (y)x0)).
'Downloaded from https: wwapthhwb..mm/SSSbwe///nnnnaaala diffusio
21

---

8
3
O
6 3
7
/
2 3
1 6
4
/
3
0
2
2
2
2
2 2 2
2
2
2
2
2
2
2
Z
3
2
2
R
2
2
2
2
3
2
2
2
2
R 2
2
2
8
2
2
R
2
3
2
2
7
2
2
2
2
2
9
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
3
2
2
2
2
3
3
3
3
3
3
3
3
3
3
3
3
3
2
2
2
3
3
3
3
3
3
3
3
3
3
3
3
3
0
2
9
3
3
3
3
3
3
3
3
3
3
3
3
3
O
D
3
3
3
3
3
3
3
3
3
3
3
3
3
2
3
3
3
3
4
4
4
4
/
+
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
0
L
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
Q
4 4
4
4
4
4
4
4
4
4
4
4
Y
4
4
4
4
4
L
4
4
4
4
4
5
5 5
5 5
5
0 Q
5 7 4 2
9
5 5555
5 5555
5
7
8 2
5 5
5 5555
2
5 5
5 5
5
5
5 5
2
4F
0
6
6
6
6
6
6
6
6
6
6
3 6
6
0
6
6
6
2
6
6
6
6
6
6
6
6
6
6
6
O
0
6
6
6
6
6
6
6
9
6
6
9
9
6
6
6
6
&
6
6
6
6
6
6
6
6
4
6
6
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
&
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
6
6
2
8
X
8
8
8
8
8
&
&
8
T
X
X
R
9
9
9
9
9
9
9
9
0
3
0
9
9
9
9
9
9
9
9
7
7
9
9
9
0
6
9
9
4
9 9
9
9
9
9
9
9
7
9
9
9
9
9
0
7
7
9
9
9
9
9
9
9
9
9
9
0 9
4
2
9
9
9
9
TDS
TDS-truncate
TDS-IS
Gradient Guidance
1S
Figure E: MNIST class-conditional generation. 16 randomly chosen conditional samples from 64
particles in single run, given class y. From top to bottom, y = 0, 2, 3, 4, 5, 6, 8, 8.9.
22

---

TDS
TDS-IS
Gradient Guidance
Class. Accuracy (%)
100
Twist Scale
1
2
3
50
64 of Particles (K) # of Particles (K) # of Particles (K)
64
(a) Classification accuracy (measured by neural network classifier) v.s. number particles K, under different twist scales. Results are averaged over 1,000 random runs with error bands indicating 2 standard errors. Across the panels we see that for all the methods, the larger the twist scale (i.e. the darker the line color). the higher the classification accuracy. For TDS and TDS-IS, this improvement is more significant for smaller value of K,
100
Accuracy (%)
Nuuann-rates
TDS K=64
75
TDS Km2
TDS-IS K=64
TDS-IS K=2
50
Gradient Guidance
2
Twist Scale
(b) Human- rated classification accuracy v.s. twist scale, for TDS (K = 2), TDS-IS (K = 64, and Gradient Guidance. Results are averaged over 640 randomly chosen samples with error bands indicating 2 standard emrors. For TDS (K' 64), increasing the twist scale generally decreases the human-rated accuracy. For remaining methods, moderate increase in twist scale improves the human-rated accuracy; however, excessively large twist scale can hurt the accuracy.
Figure F: MNIST class-conditional generation: classification accuracy under different twist scales, computed by a neural network classifier (top panel) and a human (bottom panel).
By setting 7 1, the classification likelihood becomes sharper, which is potentially helpful for twisting the samples towards 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 human (one of the authors) checks if generated digit has the right class and does not have artifacts, Since human evaluation is expensive, only evaluate TDS, TDS-IS (both with K 64, 2) and Gradient Guidance. In each run, randomly sample one particle out of K particles according to the associated weights. We conduct 64 total of 640 samples for human evaluation.
runs for each class label, leading
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. moderately large twist scale (y = 2, 3) generally helps increasing the accuracy, while an excessively large twist scale (y 10) decreases the accuracy. An exception is TDS with K 64 particles, where any 7 > leads to worse accuracy compared to the case 7 = 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, 64 has increasing accuracy. In addition, both TDS and TDS-IS with K 64 and y
23

---

6
9
9
9
9
6
G
9
6
9
6
6
6
Twist scale=1 1
Twist sca scale=2 2
Twist 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.
6
6
6
Q 4
6
4
4
C
4
6 6
6
4
6
f
6 G 6 6 6 G L
5 G
6
6
6
6
e
6
G
C
0
<
8
6
6 6
6
6
6 P
6
6
G 6
8
C
6
4
6
6
G
6
-
6
6
6
9
9
6
6
@
6
b
6
6 7
6
&
6
6
6 6 6 6
6
7 G 6
G
O
L
6
6 6
é 6 G
6
G
A
G C
G
G
G
&
G
G 6
G
Q 6
6 6
6
& &
G
6
6
R
C
6
6
G
&
A
R
G
L
&
6 G 6
6
6
G
C
G
6
6
Twist scale=1 1
Twist scale=2
Twist 3cale=3
Twist scale=4
Twist scale=10 10
(b) TDS (K'
2). of good-quality samples counted by human 49, 51, 51, 46, 45, from left to right.
6 _
G
G
6
6
4
6
G
6
6
6
6 6
b
A
6
6
6
4
6
6
6
6
6
6
6
6
G
6
6
6
2 &
G
G
6
6
6 6
6
6
6
6
G
6 G
6 6 6 6
6
6
6
è
6
L
G
6
L 6
6
L
G
J
6 6
6
G
G
6
6
G 6
6
G
b
O
6 C
G
G
G
6
G
6
6
G
6
6
6
&
6
G
G
G
6
6
G
6
G
G
6
G
6
6
6
6
6
&
6
&
6
6
G
6
6
6
6
&
b
6
6
C
6
6
Twist scale=1 1
Twist scale scale=2 2
Twist 3
Twist scale=4 =4 4
Twist scale=10 10
(c) TDS-IS (K- 64). # of good-quality samples counted by human is 62, 61, 64, 64, 61, from left to right.
6
6
6
b
A
6
4
6
4
O
4 8
0
6
6
6
0
6 6
C
6
6
@
6
6
0
6
6 G
6
K
6
6
_
R
6
6
6
G
6
6
6
9
(
G
6
6
6
6
G
6
6
6
6
6
G
6
G
6
6
&
6 6 9
0
O
G
6
6
9
6
6
9
6
3
6
&
&
6
L
e 6 b
6
9
P
6
C
O
6 6 G 6
& 6
6
6
&
6 6 6
6
G
O
G
L
6
6 g 6
6
6
6
6
D
6
Q
l
G
6
G
6
6
4
6
4 6 6
6
6
6
G
6
6
6
Lo
Twist scale=1 1
Twist 2cale=2
Twist 3
Twist scale= =4 4
Twist scale=10
(d) TDS-IS (K- 2). # of good-quality samples counted by a human is 49, 52, 53, 52, 48, from left to right.
0
3
6
6
6
6
6
G
A
6
M
6
Q
6
4
O
6 6
6
C
Ø
0
D
6
6
6
6
0
C
6
A
O
L
8
C
6
6
5
6
6
G S
Q
6
C
Q
G
&
G
P
2
6
9
K
6
G
R
6
G
G
G
6
6
G
C
O
6
O
4
G
L
6
&
L
G
6
6
6
6
6
6
G
6
G
G
L
Q
G
6
L
J
Z
G
G
6
g
6
9
b
6
6
6
D
O
G
G
6
G
J
6
G
Twist scale=1 1
Twist scale=2
Twist 3
Twist scale=4
Twist scale=10
(e) Gradient Guidance. # of good 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 (c.g. 10) would distort the digit shape with more artifacts, though retaining useful features that may allow a neural network classifier to identity the class.
24

---

Effect. Sample Size
64
TDS
32
TDS-IS
15
100
75
50
25
Time Step
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 discrepancy 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 network classifier, thereby leading to a low human-rated accuracy but a high classifier-rated accuracy. To validate this hypothesis,
neural
present samples conditioned on class 6 in Figure G. For example, in has 31 good-quality samples out of 64, and the rest of the
Figure Ge, Gradient Guidance with samples often resamble the shape of other digits, c.g. 3,4.8; and Gradient Guidance with y = 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 (E (E for K weighted particles (rt w/ 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 general upward trend of ESS approaching K = 64. Though in the final few steps ESS of TDS drops by half, it still higher than that of TDS-IS and IS that deteriorates around and respectively. can use the TDS-truncate variant introduced above to avoid the ESS drop in the final few steps.
practice, one
D.2.2 Image inpainting on MNIST
The inpainting task is to sample images from (o 0 segment _ [= . into observed dimensions M and unobserved dimensions M, described in Section 3.3.
x y) given observed part 2fr y. Here we
In this experiment, consider two types of observed dimensions: (1) M "half", where the left half of an image is observed, and (2) M "quarter", where the upper left quarter 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 10 and returns (.)
Twisting function variance schedule. twisting functions in eq. (13) to
We use flexible variance schedule {o?} to extend the
Po x² M) N (t(x)))))))))) 6,).
Ideally, of should match Var If, M). In TDS, choose oo a² 2// where 72 0.12 the estimated sample variance the training data (averaged over pixels). This choice is inspired from Song et al. [32, Appendix A.3]: if the data distribution is g(x²)) NV0, r2), and q(2² |
(5? (03/a,)) )
Na(a' Varaa of), then by Bayes rule Var Var,
xx]
[01/0.)472-
25

---

Changing 07 has a similar effect using a twist scale to exponentiate the twisting functions in Appendix D.2.1. particular, for a Gaussian distribution, exponentiating the distribution corresponds to re-scaling the variance: (x (t.
x (x 12,
Metrics. We consider 3 metries. (1) We use the effective sample size (ESS) compare the particle efficiency among different SMC samplers (namely TDS, TDS-IS, and SMC-Diff).
(2) In addition, ground the performance of sampler in the downstream task of classifying partially observed image "M y: we use a classifier to predict the class 2 of an inpainted image zº. and compute the accuracy against the true class z" of the unmasked image (=" provided in MNIST dataset). This prediction is made by the same classifier used in Section 5.2.
K Consider K weighted particles {=1: = drawn from sampler conditioned on I'M y and assume the weights are normalized. We define the Bayes accuracy (BA) as
1 (&(y)= = 2"}, with 2(y) = arg max
u upp(2;
==1:10
where {y) is viewed as an approximation to the Bayes optimal classifier 2* (y) given by
/ (2 | y)p(z; x Idx".
x (y) == arg max po x y) = arg max
(18)
z= 1:10
2=1: 1:10
(In eq. (18) we assume the classifier Pylzo (: (x²) is the optimal classifier on full images.)
(3) We also consider classification accuracy (CA) defined as the following
K
= =²"}, with 2(22) == arg max (2|=4).
2=1:10
BA and CA evaluate different aspects of a sampler. BA focused on the optimal prediction among multiple particles, whereas CA is focused on the their weighted average prediction.
Comparison results of difffernt methods. Figure depicts the ESS trace, BA and CA for different samplers. The overall observations are similar to the observations in the class-conditional generation task in Section 5.2, except that SMC-Diff and Replacement methods are not available there.
Replacement has the lowest CA and BA across all settings. Comparing TDS to SMC-Diff, we find that SMC-Diff's ESS is consistently greater than TDS; however, SMC-Diff is outperformed by TDS in terms of both CA and BA.
We also note that despite Gradient Guidance's 's CA is lower, its BA is comparable to TDS. This result due to that as long as Gradient Guidance generates a few good-quality samples out of K particles, the optimal prediction can be accurate, thereby resulting in high BA.
Ablation study on twisting function variance schedule.
We compare the following three options:
1. DPS: o == 21 )m vll2a?, adapted from the DPS method [6, Algorithm 1],
2. IIDM: न. = 0 Von, adapted from the IIGDM method [32, Algorithm 1].
3. TDS: 0f = (a? 7 where 7 = 0.12.
(3//a.)+>*
Figure J shows the classification accuracy of TDS, TDS-IS and Gradient Guidance with different twist scale schemes. We find that DPS in most cases. Exceptions are when M "quarteer" and for large K, TDS with twist scale choice of IIGDM or DPS has higher CA, as shown in the left panel in Figure Jb.
our choice has similar performance to that IIGMM, and outperforms
D.2.3 Class-conditional generation on ImageNet and CIFAR-10
ImageNet. We compare TDS (with # of particles K 16) to Classifier Guidance (CG) using the same unconditional model from https:// h/:::///tttttt com/ coom/opeaii// openai/guided-de diffeddiddfffioooo /tre/ main. CG uses classifier trained on noisy inputs taken from this repository as well. For TDS,
26

---

Effect. Sample Size
64
100
Bayes Accuracy (%0
100
TDS
Class. Accuracy
IDS-truncate
TDS I5
Gradient
Guidance
SMC-Diff
75
32
Time Step
# Particles K
Particles K
(a) M "half". ESS trace.
(b) M "half". Class. accuracy (left), and Bayes accuracy (right) v.s. K.
Effect. Sample Size
64
Class. (%6
80
Bayes Accuracy (%6)
80
TDS
TDS truncate
32
50
50
TDS-15
Gradient
SMC-Diff
20
20
Replacement
100
75
50
25
8
32
8
16 32 64
Step
#f Particles
Particles
(c) M = "quarter". trace.
(d) M "quarter". Class. accuracy (left), and Bayes accuracy (right) v.s.
Figure MNIST image inpainting. Results for observed dimension M "half" shown in the top panel, and M = "quarter" in the bottom panel. (i) Left column: ESS traces are averaged over 100 different images inpainted by TDS, TDS-IS, and SMC-Diff, all with K 64 particles. TDS's ESS is generally increasing untill the final 20 steps where drops to around 1, suggesting significant particle collapse in the end of generation trajectory. TDS-IS's ESS is always around 1. SMC-Diff has higher particle efficiency compared to TDS. (ii) Right column: Classification accuracy and Bayes accuracy averaged over 10,000 images. In general increasing the number of particles K would improve the performance of all samplers. TDS and TDS-truncate have the highest accuracy among all given the same K. (iii) Finally, comparison of top and bottom panels shows that in a harder inpainting problem where M "quarter", TDS's has higher ESS but lower CA and BA.
Method
Classification Accuracy 1
FID
Inception Score 1
TDS (K = 16) 99.90%
26.65
64.03
TDS (K = I)
99.63%
26.05
44.45
Classifier Guidance 99.17%
14.03
100.33
Unconditional model
n/a
26.21
39.70
Table 1: ImageNet. Comparison of sample quality. The top three methods are evaluated on 16k samples generated with 100 steps. The unconditional model performance is provided in [9], which is evaluated on 50k samples generated with 250 steps.
use the same classifier evaluated at timestep = to mimic standard trained classifier. We generate 16 images for each of the 1000 class labels, using 100 sampling steps. TDS uses twist scale of 10, and CG uses guidance scale of 10. Notably, given fixed class, TDS (K 16) generates correlated samples in single SMC run, and TDS (K = 1) and CG generate independent samples.
In Figure K, we observe that TDS can faithfully capture the class and have comparable image quality CG's, although with less diversity than CG and TDS (K 1).
Figure L shows more samples given randomly selected classes. We also reported results of the unconditional model from the original paper [9] that are evaluated 50k samples with 250 sampling steps in Appendix D.2.3. TDS and CG provide similar classification accuracy. TDS has similar FIDs compared to the unconditional model and better inception score. CG's FID and inception better than TDS. We suspect this difference attributed to the sample correlation (and hence less diversity) within particles in single of TDS (K 16).
We ran TDS (K = 16) with the twist scale and 100 sampling steps using diffusion model from https://githen http://ttthb.ccooo/ooeome // //github. com
CIFAR-10.
com wttp////////pcee// cowm//penaiiimmpooeddddttttt
and classifier from https: TDS generates faithful and diverse
27

---

TDS
TDS-15
Gradient Guidance
Class. Accuracy (%)
95
95.0
90
Twist Scale
DPS NGDM
92.5
90.0
90
88
Ours
87.5
86
85.0
85
4 8 16
4 8 16
3264
8 16 32 64
# Particles K
Particles
Particles
(a) M = "half". Classification accuracy under different twist scale schemes.
TDS
TDS-IS
Gradient Guidance
Class. Accuracy (%)
70
65
54
Twist Scale
DP5 NGDM
60
60
52
Ours
55
50
50
50
2
8 16 32
64
2
8 16 32
64
2
8 16 32 64
Particles
Particles
Particles
(b) M "quarter". Classification accuracy under different twist scale schemes.
Figure J: MNIST image inpainting: Ablation study on twist scales. Top: observed dimensions M "half"" Bottom: M "quarter". In most case, the performance of our choice of twist scale similar to that of IIGDM, and is better compared to DPS.
(a) TDS (P=16)
(b) TDS (P=1)
(c) Classifier Guidance
Figure K: 256x 256 ImageNet. Samples from TDS and Classifier Guidance given class 'brambling'.
images, see Figure M. However, we found TDS can occasionally generate incorrect samples for the class 'truck' (the last panel in Figure M).
Finally, we present the ESS trace plot of the three image datasets Figure N. We see there is sudden ESS drop in the final stage for MINST, sometimes for CIFAR10, but usually not ImageNet. Mechanically, the drop in ESS implics a large discrepancy between the final and intermediate target distributions. We suspect such discrepancies might arise from irregularities in the denoising network near = O. In practice, one can consider early truncating the TDS procedure to prevent such ESS drops and promote particle diversity, is studied in Appendices D.2.1 and D.2.2.
E Motif-scaffolding application details
Unconditional model of protein backbones. We here use the FrameDiff, diffusion generative model described by Yim [39]. FrameDiff parameterizes protein N-residue protein backbones as collection of rigid bodies defined by rotations and translations as zd = [[R] 21),..., (Rv, zN)]
28

---

Figure L: x 256 ImageNet. 15 random class samples. Top to bottom: TDS(P=16), TDS(P=1), Classifier Guidance.
(a) airplane
(h) deer
(c) horse
(d) ship
(c) truck
(f) failure mode
Figure M: CIFARIO. Samples from TDS with K = 16 particles given different class labels (within a single SMC run). In most cases TDS generate authentic and diverse images. The last panel presents failure mode for the class 'truck'.
MNIST
CIFAR-10
Imagenet
Effect. sample size
16
8
4
100
25 Time step t
75
50
100
75 50
25
100
25 Time step t
75
50
Time step t
Figure N: example effective sample size traces of TDS (K 16, 100 sampling steps) on MNIST, CIFAR-10 and ImageNet models, respectively.
SE(3)N, SE(3) is the special Euclidean group three dimensions (a Riemannian manifold), Each R2 E SO(3) (the special orthogonal group in three dimensions) is a rotation matrix and each on E R translation. Together, R,, and 2n describe how one obtains the coordinates of the three backbone atoms C. Ca, and N for each residue by translating and rotating the coordinates of an idealized residue with C. carbon and the origin. The conditioning information is then a motif = E S F1(3) [M] for some M C [I....,NN N. FrameDiff is a continuous time diffuusion model and includes the number of steps as a hyperparmeter; details on the neural network architecture, and details of the forward and reverse diffusion processes.
use 200 steps in all experiments. We refer the reader to [39] for
Twisting functions for motif scaffolding. To define a twisting function, we tangent normal approximation to pe (y) x²) as introduced in eq. (15). In this case the tangent normal factorizes across each residue and across the rotational and translational components. The translational components are represented in R²3 and so are treated as in the previous Fuclidean cases, using the variance preserving extension described in Appendix A. For the rotational components, represented in the special orthogonal group (3)),, the tangent normal may be computed as described in Appendix B, using the known exponential map on SO(3). particular, use the twisting function
M
Po(y o M)
II
N 2o(xx')' (x, T.. So (n R (y 0, o²),
29

---

where Vin (z)xx)) Mm. E R represent the translations associated with the m residue of the motif and its prediction from x², and yn (=0) SO)3) represent the analogous quantities for the rotational component. Next, 1 ag Var, y a =0) and of is the time of the Brownian motion associated with the forward process at step [39]. For further simplicity, our implementation further approximates log density of the tangent normal on rotations using the squared Frobenius norm of the difference between the rotations associated with the motif and the denoising prediction (as in [37]), which becomes exact the limit that de approaches 0 but avoids computation of the inverse exponential map and its Jacobian.
Motif rotation and translation degrees of freedom. We similarly seek to eliminate the rotation and translation of the motif as a degree of freedom. We again represent our ambivalence about the rotation of the motif with randomness, augment our joint model include the rotation with uniform prior, and write
Pyle (g(l=) =
/ (y) (,)22, R) for (yl= R) = 6ny(2")
and with P() the uniform density (Haar measure) on SO(3), and Ry represents rotating the rigid bodies described by y by R. For computational tractability, we approximate this integral with Monte Carlo approximation defined by subsampling a finite number of rotations, R = { [Re) (with [R] =
Motif Rots. in total). Altogether have
Po = |
[M
Σ RER MEM
Σ
Po(Ry |2
The subsampling above introduces the number of motif locations and rotations subsampled as a hyperparameter. Notably the conditional training approach does not immediately address these degrees of freedom, and prior work randomly sampled a single value for these variables [34, 36, 37].
Evaluation details. We use self-consistency evaluation approach for generated backbones [34] that (i) uses fixed backbone sequence design (inverse folding) to generate putative amino acid sequence to encode the backbone, (ii) forward folds sequences to obtain backbone structure predictions, and (iii) judges the quality of initial backbones by their agreement (or self-consistency) with predicted structures. We inherit the specifics of our evaluation and success criteria set-up following [37], including using PPONeinMPNN [7] for step (i) and AlphaFold [20] on single sequence (no multiple sequence alignment) for (ii).
In this evaluation we use PrNNeinMPNN [7] with default settings to generate 8 sequences for each sampled backbone. Positions not indicated as resdesignable in [37, Methods Table 9] are held fixed. We use AlphaFold [20] for forward folding. We define "success" generated backbone for which at least one of the 8 sequences has backbone atom with both scRMSD < Á on the motif and scRMSD < 2 Á on the full backbone.
We benchmarked TDS 24/25 problems in the benchmark set introduced [37, Methods Table 9]. A 25th problem (6VW1) is excluded because it involves multiple chains, which cannot be represented by FifmeDifff. FifmeDiff requires specifying a total length of scaffolds. In all replicates, we fixed the scaffold the median of the Total Length range specified by Watson et al. [37, Methods Table 9]. For example, 1 16 becomes 125 and 62-83 becomes 75.
As discussed in the main text, in our handling of the motif pooif-placement degree of freedom we restrict the number of possible placements considered for discontiguous motifs. To select these placements we (i) restrict to masks which place the motif indices in a pre-specified order that we do not permute and do not separate residues that appear contiguously in source PDB file. and (ii) when there are still too many possible placements sub-sample randomly to obtain the set of masks M of at most some maximum size (# Motif Locs.). ). We do not enforce the spacing between segments specified in the "contii" description described by [37].
Resampling. In all motif-scaffolding experiments, we use systematic resampling. An effective sample size threshold is used K/2 in all cases, that is, we trigger resampling steps only when the effective sample size falls below this level.
30

---

Success Rate (%)
10a
TDS (K=8)
TDS (K=1)
RFdiffusion
101
10"
6EXZ. med
6EXZ. short
1YCR
STRV med
GEGR mad
31XT
6E6R long
6EGR short
2KL8
GEMZ. GEXZ_long long
STRV short
4ZYP
SIUS
1QGG
7MRX 60
5WN9
IPRW
1BCF
7MRX 85
5TPN
7MRX MMRX_128
STRV Jong
SYUI
4/HW
Figure O: Results on full motif-scaffolding benchmark set. The y-axis is the fraction of successes across least 200 replicates. Error bars are 22 standard errors of the mean. Problems with scaffolds shorter than 100 residues are bolded. TDS (K=8) outperforms the state-of-the-art (RFdiffusion) on most problems with short scaffolds.
Success Rate (%)
TDS twist-scale=1
10a
TDS twist-scale=2
wisst-sa
10a
3IXT
STRV short
1016
18cF
IPRW
Figure Increasing the twist scale has a different impact across motif-scaffolding benchmark problems.
E.1 Additional Results
Impact of twist scale on additional motifs. Figure 3a showed monotonically increasing success rates with the twist scale. However, this trend does not hold for every problem. Figure P demonstrates this by comparing the success rates with different twist-scales on five additional benchmark problems.
Effective sample size varies by problem. Figure Q shows two example effective sample size traces over the course of sample generation. For 6EXZ-med resampling was triggered 38 times (with 14 in the final 25 steps), and for 5UIS resampling triggered 63 times (with 13 in the final 25 steps). The traces are representative of the larger benchmark.
Application of TDS to RFdiffusion: We also tried applying TDS to REdiffusion [37]. RFdiffusion a diffusion model that uses the same backbone structure representation as FrameDiff. However, unlike FifmeDifff, RFdiffusion trained with mixture of conditional examples, in which a segment of the backbone is presented as input as desired motif, and unconditional training examples in which the only noised backbone provided. Though in our benchmark evaluation provided the motif information as explicit conditioning inputs, we reasoned that TDS should apply to RFdiffusion as well (with conditioning information not provided as input). However, we were unable to compute numerically stable gradients (with respect to either the rotational or translational components of the backbone representation): this lead to twisted proposal distributions that were similarly unstable and trajectories that frequently diverged even with one particle. We suspect this instability owes to RFdiffusion's 's structure prediction pretraining and limited fine-tuning, which may allow achieve good performance without having fit a smooth score- aporappproxiitioon.
31

---

6EXZ med
5IUS
Effective Sample Size
Effective Sample Size
8
0
100 Step (T-t)
200
0
100 Step (T-t)
200
Figure Q: Effective sample size traces for two motif-scaffolding examples (Left) 6EXZ-med and (Right) 5IUS. both case K and a resampling threshold of 0.5 is used. Dashed red lines indicate resampling times.
32


Reviews:
[{"review_id": "GsKFYiowsP", "reviewer": "Reviewer_8aww", "summary": "This paper addresses the challenge conditioning in unconditionally-trained diffusion models. The most successful approaches often require explicitly training on conditional data. This paper frames sampling from such conditionals as an SMC procedure and proposes Twisted Diffusion Sampler (TDS), a method derived from the twisting technique in SMC literature. TDS uses a classifier trained on the clean, non-noisy data to construct a sub-optimal proposal distribution for sampling from the diffusion process. This proposal is then used in the twisted SMC framework to get weighted particles that more closely approximate the conditional distribution of interest.\n\nThe authors prove that their proposed SMC method targets the correct distribution and is asymptotically exact. Furthermore, they show how to extend TDS to work for various inpainting problems and Reimannian diffusion models.\n\nFinally, the authors empirically verify TDS's correctness on a simple synthetic diffusion model with tractable score functions. Then they show its effectiveness on more realistic tasks of class-conditional MNIST image generation and inpainting. They lastly show TDS achieves state-of-the-art results on motif-scaffolding problem in protein design.", "questions": "1. I am not sure if what the authors refer to \"reconstruction guidance\" is correct. Quoting from the related work section \"the use of denoising estimate $\\hat{x}_0(x^t, t)$ to form an approximation $\\tilde{p}_\\theta(x^t | x^{t+1}, y)$ to $p_\\theta(x^t | x^{t+1}, y)$ is called _reconstruction guidance_.\" I am not sure if this is correct. As far as I know, reconstruction guidance is proposed in [1] and improves the replacement method [2] which is for inpainting applications. Moreover, the other two citations are [3,4] (citation numbers 5, 25 in the paper) which are unrelated to reconstruction guidance (this might be related to the question 3 below though)\n2. In the related work section (App. B), on line 523, there is a \"[33]\" citation which does not exist. The reference 29 in the paper ([5] here) seems to be what it meant to refer to. This paper, however, does not do reconstruction guidance as far as I can tell.\n3. Looking at the rest of the related work section, seems like most citation numbers are incorrect. For example citations 26 and 18 on lines 524 and 528, respectively do not seem to refer to the correct paper. I suggest double-checking the citations.\n4. In section 3.3, it is not clear to me how Eq. (18) is derived and whether it is a convenient likelihood that often works in practice or it is a general result. I did not see an explanation in the appendices either. I think it would be helpful to clarify this.\n5. In the experiments section\n   1. \"Guidance\" is defined as \"TDS\" with one sample. If I understand it correctly, it would be equivalent to \"TDS-IS\" with one sample. Is this correct? If so, I expect it to be outperformed by \"TDS-IS\" with more samples while the middle and right panel of Figure 1 shows otherwise. Additionally, in Figure 2(a), I do not understand why would \"TDS-IS\" show less diversity than \"guidance\". Why would more particles destroy the diversity?\n   2. I am wondering if the authors have insights on why in Figure 2(b) ESS jumps back up very close to t=0.\n6. Minor issues\n   - Section 3 and subsection 3.2 have the same name.\n   - The background section on diffusion models explains diffusion processes without scaling (e.g. the variance-preserving process).It is worth mentioning in the background section that not all diffusion models follow this exact framework.\n   - I think it's nicer if Eq. 11 is written to define $\\tilde{p}_\\theta(x^{t-1} | x^{t}, y)$ instead of $\\tilde{p}_\\theta(x^t | x^{t+1}, y)$ to be similar to Eq. 3.\n7. Typos\n   - Line 9: the -> that\n   - Line 99: provide -> provides\n   - Line 148: eq. -> Eq.\n   - Line 153: I might be wrong, but I guess $s_\\theta(x^t, y) = \\nabla_{x^t} \\log q(x^t, y)$ should actually be $s_\\theta(x^t) = \\nabla_{x^t} \\log q(x^t)$. Otherwise, the assumption of $\\tilde{p}_\\theta(y | x^t) = p_\\theta(y | x^t)$ becomes unnecessary for Eqs. (13,14).\n   - Line 157: $w_T(x^T) := \\tilde{p}_\\theta(x^T)$ -> $w_T(x^T) := \\tilde{p}_\\theta(y | x^T)$\n   - Linr 218: sytematic -> systematic\n   - Lien 570: systematica -> systematic\n   - In the second line of equations under line 578: $p(y | x^{t'})$ -> $p(y | x^{t})$\n\n[1] Jonathan Ho, Tim Salimans, Alexey A Gritsenko, William Chan, Mohammad Norouzi, David J Fleet.\n\"Video diffusion models\".\nNeurIPS 2022.\n\n[2] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, Ben Poole.\n\"Score-based generative modeling through stochastic differential equations\".\nICLR 2020.\n\n[3] Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Marc L Klasky, Jong Chul Ye.\n\"Diffusion posterior sampling for general noisy inverse problems\".\nICML 2023.\n\n[4] Jue Wang, Sidney Lisanza, David Juergens, Doug Tischer, Joseph L Watson, Karla M Castro, Robert Ragotte, Amijai Saragovi, Lukas F Milles, Minkyung Baek, et al.\n\"Scaffolding protein functional sites using deep learning\".\nScience 2022.\n\n[5] Guanhua Zhang, Jiabao Ji, Yang Zhang, Mo Yu, Tommi Jaakkola, and Shiyu Chang. \"Towards coherent image inpainting using denoising diffusion implicit models\".\narXiv preprint, 2023.", "limitations": "The main limitation of the proposed method is the additional sampling cost due to multiple SMC particles which is mentioned in the discussion section of the paper. Otherwise, I do not see other limitations.", "rating": 5, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The central contribution of the paper is introducing the twisted SMC framework for sampling from diffusion models. While the twisted SMC is not novel, applying it to diffusion models and demonstrating use cases is significant. Moreover, the extension of TDS to some settings other than class-conditional framework was very interesting to me.\n\nIn terms of clarity, this paper is mostly clearly written. There are some typos and clarifications required that I have listed in the questions section.", "weaknesses": "I think the main weakness of the paper lies in the experimental section. First, the experiments are rather small scale. It would be nice to see how their method perform on larger scales standard tasks such as class-conditional CIFAR or ImageNet generation. Second, I expected to see more baselines in the experiments. The paper is motivated by pointing out that many existing methods require additional training. I believe it such methods should be included as baseline. For example, classifier-guidance with a proper classifier trained on noisy data or classifier-free guidance would be good baselines to compare against. Admittedly, they might outperform TDS, particularly with smaller K. However, it would be a good upper bound to have. Moreover, it gives insights on the tradeoff between a large K or training a separate model.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their very careful and detailed review.  We are glad they found it to be significant and interesting.  We found the questions and weaknesses noted helpful and believe we have addressed them below.\n\n**Improving the experimental validation:** As noted in the response to all reviewers, we have implemented the suggestion to include class-conditional CIFAR and ImageNet generation, and find TDS to be competitive with classifier guidance (albeit at larger inference-time cost).  As the reviewer suggested, classifier guidance does indeed outperform TDS with just a single particle but performs similarly by Classification accuracy (see table 1 in high-level reply).\n\n**Replies to questions:**\n1. Reconstruction guidance uses the reconstruction estimate $\\hat x_0(x_t)$ as a proxy for the true conditional mean $E[x_0|x_t]$. Eq 7 in VDM paper [15] demonstrates the use of reconstruction guidance gradient obtained by backpropagating through the denoising network. In the inpainting case, TDS with P=1 coincides with the reconstruction guidance. We will make this connection more clear in the revised paper. \n2. The reviewer is correct that these references were wrong. We apologize for our mistake, which resulted from an accidental reordering of references between when we compiled and submitted our main text and our appendix files.\n3. See (2) above.\n4.  We thank the reviewer for pointing out that Eq.18 was unclear.  The final proposal and weights in Eq.18 are chosen to ensure that the final target is the exact conditional distribution of interest.  The intuition is that as defined in Eq.18,  tilde $p (x_0 | x_1, y) = p (x_0 | x_1, y)$, and as a result reweighting is not needed.  By substituting equation 18 into the expression for $\\nu_0$ in equation 16, we obtain the desired target.  We will clarify this in our revision. \n5. This question touches on both TDS-IS and ESS plots:\n  - *TDS- IS.*  We thank the reviewer for pointing out this surprising behavior.  We have investigated it and discovered that it owes to a bug in the implementation of our simulations; in short, we used an incorrect and unstable implementation of the final step weights. Our attached PDF presents the results for inpainting (Figure 4) where we find that the results agree with the reviewer's expectations.  Our revision will include corrected version of the remaining three simulations with this bug-fix.\nAs for why TDS-IS provides lower diversity as compared to guidance in Figure 2a, this is in essence a bias-variance trade-off.  TDS-IS explicitly weights and resamples based on the likelihood approximation at the last step. This resampling replicates some samples and eliminates others, which reduces diversity giving some samples that are identical to one another. By contrast, guidance generates independent samples and does not use importance weighting and resampling, hence resulting in a higher diversity but worse average quality. \n - *ESS jumps near t=0.*  This is a good question which we have wondered about as well.  We report additional ESS traces in Figure 5 showing that this jump happens irregularly; we see it regularly for MINST, sometimes for CIFAR10, but usually not on ImageNet (Figure 5A, 5B, and 5C).  Mechanically, the drop in ESS implies a large discrepancy between the final and intermediate target distributions.  We suspect such discrepancies might arise from irregularities in the denoising network near t=0, and will add a brief discussion of this behavior and its implications on particle diversity and the value of truncating resampling in the final steps.\n6. This question touches on section titles, diffusion model background and time-indices.\n - *section 3 and 3.2 titles.*  We thank the reviewer for pointing this out.  We will provide more specific headings: “3: Twisted Diffusion Sampler: SMC sampling for diffusion model conditionals\" and “3.2: Twisting functions and convergence for smooth likelihoods”.\n - *diffusion model background.*  We thank the reviewer for pointing out this confusion.  We omit the scaling factor to simplify demonstration, and include the discussion of  variance preserving and variance exploding models in appendix D. We will add a reference to the appendix in the revised background section. \n - *time indices.*  We appreciate the suggestion.  However we prefer to keep this choice of indexing to better agree with the indexing we use for sequential Monte Carlo.  Unfortunately, this cannot instead be resolved by changing the indexing of diffusion models in equation (3) without requiring $x_{t+1}$ as an argument of the denoiser and score model throughout.\n7. We thank the reviewer for identifying these typos.  These were all correctly identified and will be resolved in the revision."}, {"review_id": "gYrgDj3TTi", "reviewer": "Reviewer_UG6i", "summary": "This paper proposes a practical approach for achieving asymptotically exact inference from diffusion models through exact conditional sampling in terms of Sequential Monte Carlo (SMC) . This discovers the connection between SMC and diffusion models, and one of the key feature is to approximate the optimal twisting functions by som tractable twisting alternatives. The proposed method is extended to tackle inpainting, inpainting with degrees of freedom problems, and make diffusion model compatible on Riemannian manifolds. The experiments involve synthetic diffusion models, class-conditional sample generation on MNIST dataset, and motif-scaffolding problem, which is relevant to protein design.", "questions": "1. ``Choosing weighting function $w_T(x^T) = w_t(x^t, x^{t+1}) = 1$ for $t = 1, \\dots, T$'' constantly appearing in the paper to leave the proposal functions only in the equation. Are you assuming no resampling here and no particle is being discarded? If so, please make it more clear.\n2. For weighting function on line 95, $w_k^t$ is a function related to $x_k^t, x_k^{t+1}$ and no $y$ is involved, so what is the point of not including $y$ here? When twisted weighting function is defined on equation 15, it is related to $y$ by $\\tilde{p}_\\theta(y | x^t)$ though.\n3. On line 141, ``recall that $\\hat{x_0}(x^t) = \\mathbb{E}_{p_\\theta}[x^0 | x^t]$ if $p_\\theta$ is optimized to exactly match the true distribution $q$''. What if the assumption does not held perfectly? Is there any statistical analysis to justify this?\n4. On equation 10, ``the approximated optimal twisting functions is $\\tilde{p}_\\theta(y | x^t) := p(y; \\hat{x}_0(x^t, t))$''. Is there any analysis (if this idea is from [5], please ignore this question about analysis) or experiment result for not denoising $x^t$ back to $x^0$ for comparison?\n5. Is the first experiment two-dimensional? What if it is high-dimensional?\n6. What's the computational cost for various number of particles? Is there any variance reported for different number of particles for each experiment?", "limitations": "I would expect the authors address the weakness and questions aforementioned.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The inspiration of the paper is good, which tries to achieve exact conditional sampling from diffusion models by making connection to SMC, and extend to Riemannian diffusion models.\n3. (More like an) Ablation study in the experiment portion to examine whether the proposed method works better.", "weaknesses": "Based on understandings,\n1. The novelty of this paper is not clear to me. The strong and close connection between diffusion models and the SMC technique is not crystal clear to me. Even though the reverse process with T diffusion steps from the diffusion models and sampling across T steps in terms of SMC could be related somehow, the significance of this connection is not apparent. The authors also point out that Classifier guidance[1] requires additional classifier training, but no computational cost and sample quality comparison is provided; Classifier-free guidance [2] requires classifier information as an additional input but again, none of the aforementioned result is presented and class information is still needed by the proposed TDS method (from my understanding but could be wrong). Equation 10, which is referred to as key insight in the paper, not been explained well, and actually has been proposed in [5] with even further clear explanation. The previous paper is not cited either. The title is also not informative enough and the emphasis on SMC is not evident.\n2. Lack of real benchmark for the first two experiments, particularly the examination on the MNIST dataset. I am also concerned that the \"Guidance\" method performs significantly worse than the other four methods without any justification, and no pre-trained diffusion model is served as a benchmark here. Three experiments are all competing to themselves. It would be more convincing for the image synthesis experiments if they were conducted on more diverse and challenging datasets such as CIFAR-10, instead of solely relying on the MNIST dataset, which consists of grayscale images with only 10 classes.\n3. Not sure why but the ``Related work`` is added in the appendix. MNIST class-conditional generation experiment part is in the paper, the inpainting is put in the appendix as well (which is mentioned in the ``abstract`` and ``conclusion``). I do not think it is pretty common but I would suggest to move them into the main paper from appendix.\n4. There is a big room to improve for writing. Too many typos in both grammar and math expressions, not being cited properly, i.e., Tweedie's formula in the ``background``; unclear descriptions, i.e., the diffusion models described in the ``background`` and ``methodology`` are not referred to VE or VP diffusion models until the appendix, and mathematical notations appear without proper explanation, i.e. $\\nu_t(x^{t:T})$ the first occurrence on line 94, where its meaning is not clarified. The lack of coherence between sentences and paragraphs further hampers the clarity of the paper.\n\n[1]: Prafulla Dhariwal and Alexander Nichol. Diffusion models beat GANs on image synthesis. Advances in Neural Information Processing Systems, 2021.\n[2]: Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598, 2022.\n[3]: Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Marc L Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In International Conference on Learning Representations, 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed review and are glad they appreciate the inspiration for the method.  We hope the new clarifications and benchmarks described in the below provide improved support for the progress we have made in this direction.\n\n__Weaknesses:__\n1. The first comment on weakness has several components that we address independently:\n - **Significance between the connection between SMC and diffusion models.** The significance is that it allows us to use heuristic approximations to conditional sampling in diffusion models to implement a practical and asymptotically accurate conditional sampling procedure (See e.g. lines 51-57). \n - **Comparison to conditional approaches.** We note that class information is required for TDS since the goal is to generate samples given a class. The advantage of TDS is that it does not require additional training of a classifier on noisy inputs diffused at various steps (as in classifier-guidance), or a conditional diffusion model trained on class information (as in classifier-free guidance). TDS can operate on a pre-trained unconditional diffusion model and a standard classifier trained on clean inputs. See also our response to weakness (2).\n  - **Novelty of equation 10.** With regard to Eq.10, the insight is that such approximations (previously suggested in [5] and also [15,21]) can enter an SMC procedure as a “twisting function”. While we have discussed this connection in line 151 as well as in related works line 519, we will make it more clear in the main text. \n - **Title.** We appreciate the suggestion to change the title, as we recognize that the substance of our contributions build more heavily on existing SMC methods than one might have thought.  **Our new proposed title is:  \"Practical and Asymptotically Exact Conditional Sampling in Diffusion Models\"**\n2. At the reviewer’s suggestion we now include applications to CIFAR10 and ImageNet, and compare to a conditionally-trained classifier guidance baseline, and believe the results (described in the high-level response) strengthen the paper.\n3. We appreciate the suggestion and will move the discussion of related work into the main text.  We will either additionally move the image inpainting results into the main text (space allowing) or strike the references to these experiments from the abstract and conclusion.\n4. We appreciate the comments on writing and will add appropriate references and discussions of the connections to VE and VP models to the main text.\n\n__Questions:__\n1. Yes. When weights are uniformly one, one can skip resampling steps to reduce the variance of the procedure (see footnote on page4). \n2. We thank the reviewer for pointing out this omission. We have left “y” as an implicit input, but will make its omission explicit in the revision.\n3. Yes, this is a great question.  If it does not hold perfectly, the distribution of the samples returned by TDS has some deviation from the exact conditional distribution.  A primary contribution is to show that this deviation may be reduced arbitrarily by increasing the number of particles (Theorem 1).\n4. This approximation does indeed come from reference [5] (inheriting from [15] before it).  See also our response to weakness 1.  Not denoising back to x_0 is an interesting idea, but we suspect it would not provide a sensible result because noisy states would be very dramatically out of distribution for most likelihoods as compared to the denoised predictions.\n5. The first experiment is indeed two dimensional.  This experiment was chosen to be as small as possible to be illustrative; we intended the 784-dimensional MINST experiments to illustrate a higher dimensional toy problem, and hope that our ImageNet experiments further address concerns about higher dimensional problems.\n6. The compute cost increases linearly with the number of particles, and the inverse variance increases linearly with the number of particles.  This rate is a consequence of standard SMC results and can be seen from the slope of -1 (on log-log-scale) in our simulation experiments."}, {"review_id": "UH1vbgSi6D", "reviewer": "Reviewer_bGjC", "summary": "The paper proposes an SMC algorithm to draw conditional samples form a diffusion model. Specifically, they wish to draw samples p(x0|y) given a diffusion model p(x0) and likelihood p(y|x0). Existing techniques to do so rely on expensive training of conditional diffusion models or heuristics which do not sample from the correct conditional distribution. ", "questions": "See weaknesses.", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Theoretically well-founded method. The paper does a good job of describing the flaws of commonly used methods for conditional generation and proposes a technique to target the true conditional distribution without needing task-specific training.\n- State-of-the-art results for protein motif scaffolding, which is a very reasonable application of the proposed exact targeting of conditional distributions", "weaknesses": "- For some of the experiments (such as inpainting), a reasonable baseline to compare against would be that proposed in Section 3.1 of VDM (https://arxiv.org/abs/2204.03458). This performs conditional generation with an unconditional diffusion model using a heuristic that they show improves on naive guidance (although doesn't have theoretical guarantees like the proposed SMC method). A comparison against this would make it clearer when Twisted SMC is indeed the best option.\n- Section 4.3 is hard to read and lacking in plots or concrete results. Adding something like Figure J from the appendix would make the results more convincing for readers who do not venture to the appendix to search for plots.\n- How accurate is the inferred conditional distribution when less sampling steps are used? With progressive distillation or more modern ODE/SDE solvers it is now common to sample from diffusion models with tens of integration steps. The performance of SMC presumably degrades when less steps are used, since each individual step must be bigger. Can the authors comment on this or quantify this effect? E.g. how would Twisted SMC perform relative to the baselines if the number of diffusion steps used for motif-scaffolding was halved?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments and are glad they found the paper to be theoretically well-founded, and that they appreciate the state-of-the-art results.  We believe we thoroughly address the noted limitations in the below.\n\n**VDM as a baseline:** As we noted in our high-level reply, the method we labeled as “guidance” exactly coincides with “reconstruction guidance” as proposed in VDM.  We thank the reviewer for revealing that this term was unclear, and will be certain it is clarified in our revision.\n\n**Section 4.3 results and readability:**  We appreciate the suggestion to provide more concrete results and will move Figure J into the main text to accomplish this as suggested.  Additionally, we will add a brief preamble paragraph to section 4.3 to summarize the evaluation and state-of-the-art results.\n\n**Accuracy with fewer steps:** This is a great question.  We comment first on empirics and then on methodological considerations. With 50 or 100 steps rather than 200 steps performance seems to degrade to varying extents depending on the particular motif (see Table below).  The performance drop is statistically significant (by Fisher exact test, p<0.05) in only one of three cases (3IXT). \n\nHowever, the interpretation of this result is complicated by the fact that changing the number steps (or analogously using a different ODE/SDE integrator) slightly modifies the sampling distribution of the unconditional model, and therefore also slightly modifies conditional distributions.  Consequently, exactly quantifying the impact on accuracy is ill-posed because it is relative to a moving target.  We expect however that, though one can apply TDS with any number of sampling steps, effective sample sizes may be worse when step-sizes are large.  In fact, we conjecture that the KL between intermediate targets is linear in the step size, and therefore the number of particles needed could be exponential.  However, we have yet to prove this result and leave it to future work.  We additionally note that the performance of other conditional generation methods may also in some cases degrade when fewer steps are used.\n\n*Table:* Number of motif-scaffolding successes (out of 50 runs) with different numbers of diffusion steps with 8 particles.\n\n| Motif       | 200 steps | 100 steps | 50 steps |\n| ----------- | --------- | --------- | -------- |\n| 1QJG        | 3 / 50    | 1 / 50    | 1 / 50   |\n| 3IXT        | 48 / 50   | 43 / 50   | 42 / 40  |\n| 5TRV_short  | 20 / 50   | 9 / 50    | 8 / 50   |\n\nThese results above are from a preliminary exploration (conducted before submission) of the impact of the number of steps.  In response to the reviewers comments we intend to rerun these experiments to more comprehensively understand the impact of the number of steps.\n"}, {"review_id": "MHv7NZO3fm", "reviewer": "Reviewer_BFYS", "summary": "This paper focuses on solving inverse problems using diffusion based probabilistic models. More precisely, it is only assumed that one has access to a diffusion model for the prior distribution and a likelihood, so that no additional training is needed. The aim of the present paper is to provide an asymptotically exact method. \n\nThe method builds up on the idea of reconstruction guidance which aims at approximating the score involved in the diffused posterior,  $\\nabla \\log p_t(x_t | y)$ which is equal to $\\nabla \\log p(y|x_t) + \\nabla \\log p_t(x_t)$ where the second term is given by the prior diffusion model. The first term is however intractable but can be approximated by noting that $p(y | x_t) = \\int p(y|x_0) p(x_0 | x_t) dx_0$. This integral is then approximated by simply taking the mean a posteriori i.e. $p(y | x_t) \\approx p(y | E(x_0 | x_t))$ which  itself can be approximated using Tweedie's formula and the prior score.\n\nThe authors use this idea to obtain particle approximations of the marginal of the backward diffusion of the posterior. They do so using twisting functions. This allows them to define a principled SMC samplers for the target distribution. By standard SMC results, the obtained particle approximation of the posterior converges to the posterior of the diffusion model. The extension to Riemannian manifolds is also provided. \n", "questions": "It is claimed in the appendix that the assumptions required to prove the convergence are mild. However, assuming that the ratio of the twisting function (assumption (b), line 588-589) is bounded is not a mild assumption since it doesn't hold in the simplest cases. It is claimed that this holds if the prediction $\\hat{x}_0(x_t)$ is compactly supported but I am not sure how this is possible even if the data distribution is compactly supported. Indeed, the marginals of the forward process $q_t$ are the result of the convolution of the data distribution with a Gaussian kernel so that it cannot be compactly supported. Since  $\\hat{x}_0(x_t)$ contains the score, it cannot be compactly supported too. Next, assuming that the score is bounded is also not reasonable. Finally, assuming that $p(y; x^0)$ is bounded away from zero works but this does not hold for Gaussian linear inverse problems, which is quite unfortunate. \n\nIn short, both assumptions (b) and (c) are strong assumptions and it should be stated explicitely. \n\nThere is a typo in assumption (c); its not the gradient that should be continuous with bounded gradients. ", "limitations": "See weaknesses.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The contribution of this paper is original. Unlike other methods trying to solve inverse problems, this method is theoretically grounded and is guaranteed to be reliable in relatively complex problems. This is for example not the case of DPS [1] which fails to sample the posterior even in this simplest Bayesian settings. \n\n- I also acknowledge that the paper is well written, I enjoyed reading it. \n\n- The numerical experiments are sound and the application to protein design is interesting. ", "weaknesses": "- While I find the approximations of the optimal twisting functions reasonable, I am not sure if this is the best idea if we take into account the computational cost. Indeed, sampling from the transition requires computing the gradient of the score network with respect to the input. In large scale applications (like images with 3x256x256 dimensions) this severly limits the number of particles that can be used, since the gradient will be computed for each particle. At the end of the day, the number of particles used is very important even if the transition kernels and weights are near optimal; if the posterior is highly multimodal then we inevitably need many particles in order to populate each mode. SMCDiff on the other hand should work better on these problems since one can use a larger number of particles given a fixed computational budget. \n\n- The comparison with SMCDiff is in my opinion unfair. SMCDiff performs particle filtering only on the unobserved part of the state and does not require computing the gradient of the score network. Therefore the computational time and memory cost of TDS are much larger than SMCDiff. This is not taken into account in the numerical experiments. Furthermore, it seems that SMCDiff outperforms TDS in terms of ESS. I understand that SMCDiff relies on a stringent assumption that does not hold in practice, but in my opinion it is as reasonable as saying that the instrumental kernel and weights proposed for TDS are a good approximation. In short, i am not really convinced that TDS is better than SMCDiff for inpainting tasks. For noisy inverse problems of course the comparison is not relevant since SMCDiff is not designed for such problems. \n\n- As far as I can tell, the impact of the dimension of the observation is not discussed. As is widely known in the SMC community, this makes the weights degenerate. I believe that a discussion on this matter should be added so that readers are not misled into thinking that in high dimensional problems this method will provide $K$ diversified samples from the posterior. It should be emphasized that this method works best on low dimensional problems.\n\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful comments and are glad they found the paper to be original, well-written and to have sound theoretical and numerical support.  We believe that our new experiments and clarifications below thoroughly addresses the weaknesses noted.\n\n**Computational cost.**  Compared to other approaches, we believe the proposed algorithm can present a favorable trade-off between computational cost and accuracy.  But we agree with the reviewer that the proposed algorithm is not the best idea when both the following hold: (1) repeated, fast generation at inference time is of primary concern and (2) designing and training an effective conditional model is an option.  We summarize some comparisons both from our new experiments and submission:\n* For the class-conditional ImageNet experiment described in our high-level response, we estimate the compute spent on training a noisy classifier required by Classifier Guidance was approximately 330 gpu hours ([they](https://github.com/openai/guided-diffusion/tree/main) report 300,000 training iterations with a batch size of 256, and in our hands one training iteration with a batch size of 8 takes 0.1244 GPU seconds).   By contrast, generation of a sample using TDS with 16 particles requires 6 GPU minutes, 0.03% of the conditional training time providing comparable accuracy.  Moreover, TDS does not require the engineering time to assemble a labeled dataset and design the time-dependent model.  However, because inference time with classifier guidance is lower, the total compute cost would lower if many (e.g. > 3,300) samples were to be generated.  The reviewer is correct that parallelizing across many particles is not as straightforward for ImageNet due to high dimensionality, but could in principle be scaled across multiple GPUs or computed in sequence (we used the latter approach for our experiments).\n* Relative to SMCDiff, in our experiments the efficiency improvement more than compensates for the cost of each particle.  The compute cost of TDS is roughly 2.2X higher  than that of SMCDiff, but in Figure H we show that TDS with 2 particles outperforms SMCDiff with even 64 particles. In this situation, TDS is >10x faster and provides better results.  The observation about effective sample size of SMCDiff relative to TDS is interesting.  We are surprised by this result given the worse empirical performance of SMCDiff by classification accuracy and do not have a good explanation for why these metrics do not agree.\n\n**Impact of dimension:** Our promising empirical results on MNIST (784 dimensional), protein design (\\~600 dimensional, varying by test-case), CIFAR 10 (\\~3K dimensional) and ImageNet (196K dimensional) suggest TDS can work well even on high dimensional problems. \nNotably, in the ImageNet case, despite that the particles lack diversity in global-level features, there remain variations in local patches (see Figure 2).  \nIndeed these high-dimensional, empirical successes are not typical for SMC methods, which can have notoriously bad dimension dependence.  We suspect this good performance in high dimensions owes to the quality of the proposals distributions obtained with our twisting functions.\n\n\n**Suitability of the assumptions of Theorem 1:**\nWe thank the reviewer for their thoughtful engagement with the assumptions of our theorem.  We will replace the word “mild” with a discussion of conditions, as we agree that our claim of “mildness” in the main text was overly blase. However, we disagree with the claim that our conditions “don’t hold in the simplest cases” and hope the points below will satisfy the reviewer and (once added to our revision) future readers.\n* **Compact support of $\\hat x_0(x_t)$:** We expect the range of $\\hat x_0(x_t)$ to exist within some compact set for two reasons.  First, it is common practice to truncate denoising predictions to within the range of the data; for example in image diffusion models by clipping values between the maximum and minimum pixel intensities values.  Second, because $\\hat x_0(x_t)$ is trained to approximate $E[x_0 | x_t]$ which exists within the convex hull of the population distribution; for example, the (mean centered) protein structures have bounded support because they contain a finite number of atoms whose distances are constrained by fixed bond lengths, and so we always observe denoising predictions within the bounds of some maximum size. We do not assume that the $x_t$’s for $t>0$ have compact support (which indeed is not satisfied in general).\n* **Bounded score:** Our assumption is not that the (unconditional) score is bounded (which indeed would not be satisfied), but that the gradient of the log likelihood approximation $\\nabla_{x_t} \\log \\tilde p(y|x_t)$ is bounded.  As we noted on line 601, this is our “strongest” assumption but will be satisfied if $\\hat x_0(x_t)$ and $p(y;x_0)$ are smooth in $x_t$ and $x_0$ respectively.  Though this condition is difficult to check, we suspect this assumption will be satisfied for typical denoising networks and classification models when trained with regularization to encourage smoothness.\n* **Applicability to linear inverse problems:** The reviewer is correct that assumption c does not hold in these cases, and we will make this more explicitly.  We had separated the likelihood case into a separate first section (3.2) in which we presented our theorem, and described the inverse problems case subsequently for this reason.  Generalizing our theory to cover this case as well is of interest for future work, but requires moving away from the standard SMC theory – although SMC methods are commonly used with weight functions that are not bounded (as in the inverse problems case) most existing theory requires this assumption.\n\nWe thank the reviewer for identifying our typo in our statement of assumption (c).   We will correct it.  Thank you very much again for your careful reading!"}]

Author rebuttal:
We thank the reviewers for their detailed comments and suggestions. We are pleased that the reviewers found our method to be “practical”, “original”, “theoretically well-founded”, and to give “state-of-the-art” results in protein design. We believe we have addressed all suggested weaknesses, and have improved the submission both by adding several new experiments, and through revisions to the text.  Notably we apply TDS to 256x256 ImageNet class-conditional generation and find TDS a viable alternative to Classifier Guidance which requires expensive training of a classifier on noisy inputs. 

Before describing these changes we provide additional high-level context and summarize our contributions.

Our submission proposes an algorithm, TDS,  for accurately estimating conditional distributions implied by diffusion models useful in cases when it is desirable to trade off sampling-time with accuracy. For example, the development of protein-based drugs might involve (1) developing and training a diffusion model of protein structures, (2) sampling candidate protein structures from the model, and (3) experimentally and clinically validating candidate proteins.  While steps (1) and (3) might require years, sampling (2) takes only seconds.  The proposed algorithm would apply in this second step. On several full-scale protein tasks, we show that TDS significantly outperforms the previous state of the art, a conditionally-trained diffusion model [26].  Moreover, we show the algorithm allows one to target conditional distributions to which previous methods do not even apply. In addition, we evaluated TDS on synthetic problems and well-studied computer vision benchmarks.

We now move to shared concerns with discussions of new experiments with figures in the attached PDF. 

__Higher-dimensional image problems and comparison to baselines.__
Reviewers 8aww, UG6i and BFYS noted that our submission had not explored performance on high-dimensional problems.  To this end we have included additional experiments on class-conditional generation on CIFAR10 and 256x256 ImageNet. 
- ImageNet: we compare TDS (with # of particles P=1,16) to Classifier Guidance (CG) using the same unconditional model from [here](https://github.com/openai/guided-diffusion/tree/main). CG uses a classifier trained on noisy inputs. For TDS, we use the same classifier evaluated at timestep = 0 to mimic a standard trained classifier. We generate 16 images for each of the 1000 class labels, using a guidance scale of 10 and 100 sampling steps. Notably, given a fixed class, TDS(P=16) generates correlated samples in a single SMC run, and TDS(P=1) and CG generate 16 independent samples. 
TDS can faithfully capture the class and have comparable image quality to CG’s (Figure 2, class: bramblings), although with less diversity than CG and TDS(P=1). Figure 3 shows more samples given randomly selected classes. We also reported results of the unconditional model from [1] that are evaluated on 50k samples with 250 sampling steps (Table 1). TDS and CG provide similar classification accuracy. TDS has similar FIDs compared to the unconditional model and better inception score. CG’s FID and inception score are better than TDS. We suspect this difference is attributed to the sample correlation (and hence less diversity) within particles in a single run of TDS(P=16). 
- CIFAR10: we ran TDS (P=16) with guidance scale = 1 and 100 sampling steps (Figure 1) using diffusion model from [here](https://github.com/openai/improved-diffusion) and classifier from [here](https://github.com/VSehwag/minimal-diffusion/tree/main). TDS generates faithful and diverse images.  However, we found TDS can occasionally generate off-the-manifold samples for the class ‘truck’ (Figure 1F).

__Reconstruction guidance, video diffusion models (VDM) and baselines terminology.__ Reviewer bGjC suggested comparing TDS to the reconstruction guidance approach presented in “Video Diffusion Models” [15], and reviewer 8aww identified that our description of this baseline (first introduced in [15]) was unclear.  We compared TDS to reconstruction guidance in our submission under the name “guidance”, a choice of terminology we will clarify in our revision.  

In brief, we used “guidance” to describe a generalization of “reconstruction guidance” to classification problems and inpainting with degrees of freedom (App. B); in the inpainting case, it exactly coincides with reconstruction guidance [15].  Our approach may be seen as wrapping this heuristic approximation in an SMC sampler to improve accuracy by adding additional compute (lines 53-55).  See our reply to reviewer bGjC for details.

__Compute Cost.__ We will include the below discussion in our revised manuscript. 
- Classifier guidance vs TDS:  In Classifier Guidance [8] one trains a noise-level dependent classifier.  One therefore incurs a large, up-front computational cost to train the classifier. It amortizes the inference problem and allows conditional samples to be generated in a single trajectory and so is fast at inference time.  By comparison, TDS demands compute that is (1) linear in the number of particles used and (2) higher by a constant factor due to backpropagating through the denoising network.  In the ImageNet model, TDS takes 0.34s to generate one particle at a time step, and Classifier Guidance takes 0.15s (results averaged over 100 samples on a V100 GPU). So TDS requires 120% more gpu time in this instance. Notably, training a noisy classifier required by Classifier Guidance can take around 330 gpu hours.

- TDS vs RFDiffusion on motif-scaffolding:  TDS(P=8) is faster than the state of the art baseline (RFDiffusion).  On a 100 residue test case (1QJG) on an A4000 gpu, run-time was 80 seconds for TDS, as compared to 150 seconds for RFdiffusion.  In both cases the models used 200 diffusion time-steps.  The slower speed of RFdiffusion owes to its use of a larger neural network.

Meta-review: I am thrilled to convey my recommendation for the acceptance of your paper to be presented at NeurIPS. Your work has impressed both the reviewers and myself with its quality and significance, and I am confident that it will make a valuable contribution to the NeurIPS conference.

Your responsiveness to the feedback provided during the review process reflects your commitment to advancing the field and strengthening your contribution. Your willingness to engage with the reviewers' insights demonstrates your dedication to delivering a high-quality presentation that will resonate with the conference attendees.

NeurIPS is a prestigious platform that attracts the brightest minds in the field, and your paper's acceptance adds to the conference's reputation for excellence. I am genuinely excited to see your work presented and discussed among peers who share your passion for pushing the boundaries of knowledge.
Decision: Accept (poster)