Title: Practical and Asymptotically Exact Conditional Sampling in Diffusion Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: diffusion models; conditional sampling; sequential monte carlo methods; generative models; 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

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. 2

1 Introduction

Conditional sampling is an essential primitive in the machine learning toolkit. One begins with a generative model that parameterizes a distribution p θ ( x ) on data x , and then augments the model to include information y in a joint distribution p θ ( x, y ) = p θ ( x ) p ( y | x ) . This joint distribution then implies a conditional distribution p θ ( x | y ) , from which desired outputs are sampled.

For example, in protein design, p θ ( x ) can represent a distribution of physically realizable protein structures, y a substructure that imparts a desired biochemical function, and samples from p θ ( x | y ) are then physically realizable structures that contain the substructure of interest [e.g. 34, 37].

Diffusion models are a class of generative models that have demonstrated success in conditional generation tasks [18, 27, 33, 37]. They parameterize distributions p θ ( x ) through an iterative refinement process that builds up data from noise. When a diffusion model is used for conditional generation, this refinement process is modified to account for conditioning at each step [9, 19, 27, 33].

One approach is to incorporate the conditioning information into training [e.g. 9, 27]; one either modifies the unconditional model to take y as input or trains a separate conditional model to predict y from partially noised inputs. However, conditional training requires (i) assembling a large set of paired examples of the data and conditioning information ( x, y ) , and (ii) designing and training a task-specific model when adapting to new conditioning tasks. For example, image inpainting and class-conditional image generation can be both formalized as conditional sampling problems based on the same (unconditional) image distribution p θ ( x ) ; however, the conditional training approach requires training separate models on two curated sets of conditioning inputs.

To avoid conditional training, a separate line of work uses heuristic approximations that directly operate on unconditional diffusion models: once an unconditional model p θ ( x ) is trained, it can be flexibly combined with various conditioning criteria to generate customized outputs. These approaches have been applied to inpainting problems [24, 33, 34], and other inverse problems [1, 6, 19, 32]. But it is unclear how well these heuristics approximate the exact conditional distributions they are designed to mimic; for example on inpainting tasks they often fail to return outputs consistent with both the conditioning information and unconditional model [40]. These concerns are critical in domains that require accurate conditionals. In molecular design, for example, even a small approximation error could result in atomic structures with chemically implausible bond distances.

This paper develops a practical and asymptotically exact method for conditional sampling from an unconditional diffusion model. We use sequential Monte Carlo (SMC), a general tool for asymptotically exact inference in sequential probabilistic models [5, 10, 25]. SMC simulates an ensemble of weighted trajectories, or particles , through a sequence of proposals and weighting mechanisms. These weighted particles then form an asymptotically exact approximation to a desired target distribution.

The premise of this work is to recognize that the sequential structure of diffusion models permits the application of SMC for sampling from conditional distributions p θ ( 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 is intractable, we effectively approximate it with recent heuristic approaches to conditional sampling [eg. 19], and correct the errors by the weighting mechanisms. The resulting algorithm maintains asymptotic exactness to p θ ( x | y ) , and empirically it can outperform heuristics alone even with just two particles.

We summarize our contributions: (i) We propose a practical SMC algorithm, Twisted Diffusion Sampler or TDS, for asymptotically exact conditional sampling from diffusion models; (ii) We show that TDS applies to a range of conditional generation problems, and extends to Riemannian manifold diffusion models; (iii) On MNIST inpainting and class-conditional generation tasks we demonstrate TDS's empirical improvements beyond heuristic approaches; and (iv) On protein motif-scaffolding problems with short scaffolds TDS provides greater flexibility and achieves higher success rates than the state-of-the-art conditionally trained model.

2 Background: Diffusion models and sequential Monte Carlo

Diffusion models. A diffusion model generates a data point x 0 by iteratively refining a sequence of noisy data x t , starting from pure noise x T . This procedure parameterizes a distribution of x 0 as the marginal of a length T Markov chain

where p ( x T ) is an easy-to-sample noise distribution, and each p θ ( x t -1 | x t ) is the transition distribution defined by the ( T -t ) th refinement step.

Diffusion models p θ are fitted to match a data distribution q ( x 0 ) from which we have samples. To achieve this goal, a forward process q ( x 0 ) ∏ T t =1 q ( x t | x t -1 ) is set to gradually add noise to the data, where q ( x t | x t -1 ) = N ( x t ; x t -1 , σ 2 ) , and σ 2 is a positive variance. To fit a diffusion model, one finds θ such that p θ ( x t -1 | x t ) ≈ q ( x t -1 | x t ) , which is the reverse conditional of the forward process. If this approximation is accomplished for all t , and if Tσ 2 is big enough that q ( x T ) may be approximated as q ( x T ) ≈ N (0 , T σ 2 ) =: p ( x T ) , then we will have p θ ( x 0 ) ≈ q ( x 0 ) .

In particular, when σ 2 is small enough then the reverse conditionals of q are approximately Gaussian,

where q ( x t ) = ∫ q ( x 0 ) q ( x t | x 0 ) dx 0 and ∇ x t log q ( x t ) is known as the (Stein) score [33]. To mirror eq. (2), diffusion models parameterize p θ ( x t -1 | x t ) via a score network s θ ( x t , t )

When s θ ( x t , t ) is trained to approximate ∇ x t log q ( x t ) , we have p θ ( x t -1 | x t ) ≈ q ( x t -1 | x t ) .

Notably, approximating the score is equivalent to learning a denoising neural network ˆ x θ ( x t , t ) to approximate E q [ x 0 | x t ] . The reason is that by Tweedie's formula [12, 28] ∇ x t log q ( x t ) = ( E q [ x 0 | x t ] -x t ) /tσ 2 and one can set s θ ( x t , t ) := (ˆ x θ ( x t , t ) -x t ) /tσ 2 . The neural network ˆ x θ ( x t , t ) may be learned by denoising score matching (see [e.g. 18, 35]). For the remainder of paper we drop the argument t in ˆ x θ and s θ when it is clear from context.

Appendix A generalizes the formulation above to diffusion process formulations that are commonly used in practice.

Sequential Monte Carlo. Sequential Monte Carlo (SMC) is a general tool to approximately sample from a sequence of distributions on variables x 0: T , terminating at a final target of interest [5, 10, 14, 25]. SMC approximates these targets by generating a collection of K particles { x t k } K k =1 across T steps of an iterative procedure. The key ingredients are proposals r T ( x T ) , { r t ( x t | x t +1 ) } T -1 t =0 and weighting functions w T ( x T ) , { w t ( x t , x t +1 ) } T -1 t =0 . At the initial step T , one draws K particles of x T k ∼ r T ( x T ) and sets w T k := w T ( x T k ) , and sequentially repeats the following for t = T -1 , . . . , 0 :

propose x t k ∼ r t ( x t | x t +1 k ) , k = 1 , · · · , K

The proposals and weighting functions together define a sequence of intermediate target distributions,

where L t is a normalization constant. A classic example that SMC applies is the state space model [25, Chapter 1] that describes a distribution over a sequence of latent states x 0: T and observations y 0: T . Each intermediate target ν t is constructed to be the posterior p ( x 0: T | y t : T ) given the first T -t +1 observations, and the final target p ( x 0: T | y 0: T ) is the posterior given all observations.

The defining property of SMC is that the weighted particles at each t form discrete approximations ( ∑ K k ′ =1 w t k ′ ) -1 ∑ K k =1 w k t δ x t k ( x t ) (where δ is a Dirac measure) to ν t ( x t ) that become arbitrarily accurate in the limit that many particles are used [5, Proposition 11.4]. So, by choosing r t and w t so that ν 0 ( x 0 ) matches the desired distribution, one can guarantee arbitrarily low approximation error in the large compute limit.

3 Twisted Diffusion Sampler: SMC sampling for diffusion model conditionals

Consider conditioning information y associated with a given likelihood function p y | x 0 ( y | x 0 ) . We embed y in a joint model over x 0: T and y as p θ ( x 0: T , y ) = p θ ( x 0: T ) p y | x 0 ( y | x 0 ) , where y and x 1: T are conditionally independent given x 0 . Our goal is to sample from the conditional p θ ( x 0 | y ) .

In this section, we develop Twisted Diffusion Sampler (TDS), a practical SMC algorithm targeting p θ ( x 0 | y ) . First, we describe how the Markov structure of diffusion models permits a factorization of an extended conditional distribution to which SMC applies. Then, we show how a diffusion model's denoising predictions support the application of twisting , an SMC technique in which one uses proposals and weighting functions that approximate the 'optimal' ones. Lastly, we extend TDS to certain 'inpainting' problems where p y | x 0 ( y | x 0 ) is not smooth, and to Riemannian diffusion models.

3.1 Conditional diffusion sampling as an SMC procedure

The Markov structure of the diffusion model permits a factorization that is recognizable as the final target of an SMC algorithm. We write the conditional distribution, extended to include x 1: T , as

Algorithm 1: Twisted Diffusion Sampler (TDS)

with the desired marginal, p θ ( x 0 | y ) . Comparing the diffusion conditional of eq. (5) to the SMC target of eq. (4) suggests SMC can be used.

For example, consider a first attempt at an SMC algorithm. Set the proposals as r T ( x T ) = p ( x T ) and r t ( x t | x t +1 ) = p θ ( x t | x t +1 ) for 1 ≤ t ≤ T, and weighting functions as w T ( x T ) = w t ( x t , x t +1 ) = 1 for 1 ≤ t ≤ T and w 0 ( x 0 , x 1 ) = p y | x 0 ( y | x 0 ) .

Substituting these choices into eq. (4) results in the desired final target ν 0 = p θ ( x 0: T | y ) with normalizing constant L 0 = p θ ( y ) . As a result, the associated SMC algorithm produces a final set of K samples and weights { x 0 k ; w 0 k } K k =1 that provides an asymptotically accurate approximation P K ( x 0 | y ) := ( ∑ K k w 0 k ) -1 ∑ K k =1 w 0 k δ x 0 k ( x 0 ) to the desired p θ ( x 0 | y ) .

The approach above is simply importance sampling with proposal p θ ( x 0: T ) ; with all intermediate weights set to 1, one can skip resampling steps to reduce the variance of the procedure. Consequently, this approach will be impractical if p θ ( x 0 | y ) is too dissimilar from p θ ( x 0 ) as only a small fraction of unconditional samples will have high likelihood: the number of particles required for accurate estimation of p θ ( x 0: T | y ) is exponential in KL [ p θ ( x 0: T | y ) ∥ p θ ( x 0: T ) ] [3].

3.2 Twisted diffusion sampler

Twisting is a technique in the SMC literature intended to reduce the number of particles required for good approximation accuracy [15, 16]. Loosely, it introduces a sequence of twisting functions that modify the naive proposals and weighting functions, so that the resulting intermediate targets are closer to the final target of interests. We refer the reader to Naesseth et al. [25, Chapter 3.2] for background on twisting in SMC.

Optimal twisting. Consider defining the twisted proposals r ∗ t by multiplying the naive proposals r t described in Section 3.1 by p θ ( y | x t ) as

The factors p θ ( y | x t ) are the optimal twisting functions because they permit an SMC sampler that draws exact samples from p θ ( x 0: T | y ) even when run with a single particle.

To see that a single particle is an exact sample, by Bayes rule, the proposals in eq. (6) reduce to

As a result, if one samples x T ∼ r ∗ T ( x T ) and x t ∼ r ∗ t ( x t | x t +1 ) for t = T -1 , . . . , 0 , by (i) the law of total probability and (ii) the chain rule of probability, one obtains x 0 ∼ p θ ( x 0 | y ) as desired.

However, we cannot readily sample from each r ∗ t because p θ ( y | x t ) is not analytically tractable. The challenge is that p θ ( y | x t ) = ∫ p y | x 0 ( y | x 0 ) p θ ( x 0 | x t ) dx 0 depends on x t through p θ ( x 0 | x t ) . The latter in turn requires marginalizing out x 1 , . . . , x t -1 from the joint density p θ ( x 0: t -1 | x t ) , whose form depends on t calls to the neural network ˆ x θ .

Tractable twisting. To avoid this intractability, we approximate the optimal twisting functions by

which is the likelihood function evaluated at ˆ x θ ( x t ) , the denoising estimate of x 0 at step t from the diffusion model. This tractable twisting function ˜ p θ ( y | x t ) is the key ingredient needed to define the Twisted Diffusion Sampler (TDS, Algorithm 1). We motivate and develop its components below.

The approximation in eq. (8) offers two favorable properties. First, ˜ p θ ( y | x t ) is computable because it depends on x t only though one call to ˆ x θ , instead of an intractable integral over many calls as in the case of optimal twisting. Second, ˜ p θ ( y | x t ) becomes an increasingly accurate approximation of p θ ( y | x t ) , as t decreases and p θ ( x 0 | x t ) concentrates on E p θ [ x 0 | x t ] , which ˆ x θ is fit to approximate; at t = 0 , where we can choose ˆ x θ ( x 0 ) = x 0 , we obtain ˜ p θ ( y | x 0 ) = p y | x 0 ( y | x 0 ) .

We next use eq. (8) to develop a sequence of twisted proposals ˜ p θ ( x t | x t +1 , y ) , to approximate the optimal proposals p θ ( x t | x t +1 , y ) in eq. (7). Specifically, we define twisted proposals as

is an approximation of the conditional score, s θ ( x t +1 , y ) ≈ ∇ x t +1 log p θ ( x t +1 | y ) , and ˜ σ 2 is the proposal variance. For simplicity one could choose ˜ σ 2 = σ 2 to match the variance of p θ ( x t | x t +1 ) .

Equation (9) builds on previous works (e.g. [33, 31]) that seek to approximate the reversal of a conditional forward process. The gradient in eq. (10) is computed by back-propagating through ˆ x θ ( x t +1 ) . We further discuss this technique, which has been used before [e.g. in 19, 32], in Section 4.

Twisted targets and weighting functions. Because ˜ p θ ( x t | x t +1 , y ) will not in general coincide with the optimal twisted proposal p θ ( x t | x t +1 ) , we must introduce non-trivial weighting functions to ensure the resulting SMC sampler converges to the desired final target. In particular, we define twisted weighting functions as

and w T ( x T ) := ˜ p θ ( y | x T ) . The weighting functions in eq. (11) recover the optimal, constant weighting functions if all other approximations at play are exact.

These tractable twisted proposals and weighting functions define intermediate targets that gradually approach the final target ν 0 = p θ ( x 0: T | y ) . Substituting into eq. (4) each ˜ p θ ( x t | x t +1 , y ) in eq. (9) for r t ( x t | x t +1 ) and w t ( x t , x t +1 ) in eq. (11) and then simplifying we obtain the intermediate targets

The right-hand bracketed term in eq. (12) can be understood as the discrepancy of ν t from the final target ν 0 accumulated from step t to 0 (see Appendix A.3 for a derivation). As t approaches 0 , ˜ p θ ( y | x t ) improves as an approximation of p θ ( y | x t ) , and the t -term product inside the bracket consists of fewer terms - the latter accounts for the discrepancy between practical and optimal proposals. Finally, at t = 0 , because ˜ p θ ( y | x 0 ) = p θ ( y | x 0 ) by construction, eq. (12) reduces to ν 0 ( x 0: T ) = p θ ( x 0: T | y ) , as desired.

The TDS algorithm and asymptotic exactness. Together, the twisted proposals ˜ p θ ( x t | x t +1 , y ) and weighting functions w t ( x t , x t +1 ) lead to Twisted Diffusion Sampler , or TDS (Algorithm 1). While Algorithm 1 states multinomial resampling for simplicity, in practice other resampling strategies (e.g. systematic [5, Ch. 9]) may be used as well. Under additional conditions, TDS provides arbitrarily accurate estimates of p θ ( x 0 | 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 K ( x 0 ) = ( ∑ K k ′ w k ′ ) -1 ∑ K k =1 w k δ x 0 k ( x 0 ) denote the discrete measure defined by the particles and weights returned by Algorithm 1 with K particles. Under regularity conditions on the twisted proposals and weighting functions, P K ( x 0 ) converges setwise to p θ ( x 0 | y ) as K approaches infinity.

3.3 TDS for inpainting, additional degrees of freedom

The twisting functions ˜ p θ ( y | x t ) := p y | x 0 ( y | ˆ x θ ( x t )) we introduced above are one convenient option, but are sensible only when p y | x 0 ( y | ˆ x θ ( x t )) is differentiable and strictly positive. We now show how alternative twisting functions lead to proposals and weighting functions that address inpainting problems and more flexible conditioning specifications. In these extensions, Algorithm 1 still applies with the new definitions of twisting functions. Appendix A provides additional details, including the adaptation of TDS to variance preserving diffusion models [18].

Inpainting. Consider the case that x 0 can be segmented into observed dimensions M and unobserved dimensions ¯ M such that we may write x 0 = [ x 0 M , x 0 ¯ M ] and let y = x 0 M , and take p y | x 0 ( y | x 0 ) = δ y ( x 0 M ) . The goal, then, is to sample p θ ( x 0 | x 0 M = y ) = p θ ( x 0 ¯ M | x 0 M ) δ y ( x 0 M ) . Here we define the twisting function for each t > 0 as

and set twisted proposals and weights according to eqs. (9) and (11). The variance in eq. (13) is chosen as tσ 2 = Var p θ [ x t | x 0 ] for simplicity; in general, choosing this variance to more closely match Var p θ [ y | x t ] may be preferable. For t = 0 , we define the twisting function analogously with small positive variance, for example as ˜ p θ ( y | x 0 ) = N ( y ; x 0 M , σ 2 ) . This choice of simplicity changes the final target slightly; alternatively, the final twisting function, proposal, and weights may be chosen to maintain asymptotic exactness (see Appendix A.4).

Inpainting with degrees of freedom. We next consider the case when we wish to condition on some observed dimensions, but have additional degrees of freedom. For example in the context of motif-scaffolding in protein design, we may wish to condition on a functional motif y appearing anywhere in a protein structure, rather than having a pre-specified set of indices 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 a uniform prior on p ( M ) = 1 / |M| for each M ∈ M , and (iii) again embed this new variable into our joint model to define the degrees of freedom likelihood by p ( y | x 0 ) = ∑ M ∈M p θ ( M , y | x 0 ) = |M| -1 ∑ M ∈M p ( y | x 0 , M ) . Accordingly, we approximate p θ ( y | x t ) with the twisting function

with each ˜ p θ ( y | x t , M ) defined as in eq. (13). Notably, eqs. (13) and (14) coincide when |M| = 1 .

The sum in eq. (14) may be computed efficiently because each term depends on x t only through the same denoising estimate ˆ x θ ( x t ) , which must be computed only once. Since computation of ˆ x θ ( x t ) is the expensive step, the overall run-time is not significantly increased by using even large |M| .

3.4 TDS on Riemannian manifolds

TDS extends to Riemannian diffusion models on with little modification. Riemannian diffusion models [8] are structured like in the Euclidean case, but with conditionals defined by tangent normal distributions parameterized with a score approximation followed by a manifold projection step (see e.g. [4, 8]). When we assume that (as, e.g., in [39]) the model is associated with a denoising network ˆ x θ , twisting functions are also constructed analogously. For conditional tasks defined by likelihoods, ˜ p θ ( y | x t ) in eq. (8) applies. For inpainting (and by extension, degrees of freedom), we propose

where T N ˆ x θ ( x t ) M (0 , tσ 2 ) is a tangent normal distribution centered on ˆ x θ ( x t ) M . As in the Euclidean case, ˜ p θ ( x t | x t +1 , y ) is defined with conditional score approximation s θ ( x t , y ) = s θ ( x t ) + ∇ x t log ˜ p ( y | x t ) , 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 a general approach that incorporates conditioning information with gradients through the neural network ˆ x θ ( x t ) . For example, in an

SmoothLikelihood

Inpaint

Inpaint(DegreesofFreedom)

10-1

Error

10-1

10-1

TDS

square

10-2

TDS-IS

10-2

10-2

Gradient

Guidance

S

Mean

10-3

10-3

10-3

SMC-Diff

Replacement

10-4

10-4

10-4

101

102

103

104

101

102

103

104

101

102

103

104

#of Particles (K)

# of Particles (K)

#of Particles (K)

inpainting setting, Ho et al. [19] propose a Gaussian approximation N ( y | ˆ x θ ( x t ) M , α ) to p θ ( y | x t ) . This approximation motivates a modified transition distribution as in Equation (9), with the corresponding approximation to the conditional score as s θ ( x t , y ) = s θ ( x t ) -∇ x t ∥ y -ˆ x θ ( x t ) M ∥ 2 /α. This approach coincides exactly with TDS applied to the inpainting task when a single particle is used. Similar Gaussian approximations have been used by [6, 32] for other inverse problems.

Gradient guidance can also be used with non-Gaussian approximations; e.g. using p y | x 0 ( y | ˆ x θ ( x t )) ≈ p θ ( y | x t ) for a given likelihood p y | x 0 ( y | x 0 ) . This choice again recovers TDS with one particle.

Empirically, these heuristics can have unreliable performance, e.g. in image inpainting problems [40]. By comparison, TDS enjoys the benefits of gradient guidance by using it in proposals, while also providing a mechanism to eliminate approximation error by simulating additional particles.

Replacement method. The replacement method [33] is a method introduced for image inpainting using only unconditional diffusion models. The idea is to replace the observed dimensions of intermediate samples x t M , with a noisy version of observation x 0 M . However, it is a heuristic approximation and can lead to inconsistency between inpainted region and observed region [23]. Additionally, the replacement method applies only to inpainting problems. While recent work has extended the replacement method to linear inverse problems [e.g. 22], the approach provides no accuracy guarantees. It is unclear how to extend these methods to arbitrary differentiable likelihoods.

SMC samplers for diffusion models. Most closely related to the present work is SMC-Diff [34], which uses SMC to provide asymptotically accurate conditional samples for the inpainting problem. However, this prior work (i) is limited to the inpainting case, and (ii) provides asymptotic guarantees only under the assumption that the learned diffusion model exactly matches the forward noising process, which is rarely satisfied in practice. Also, SMC-Diff does not leverage twisting functions.

In concurrent work, Cardoso et al. [2] propose MCGdiff, an alternative SMC algorithm that uses the framework of auxiliary particle filtering [26] to provide asymptotically exact conditional inference. Compared to TDS, MCGdiff avoids computing gradients of the denoising network but applies only to linear inverse problems, with inpainting as a special case.

5 Simulation study and conditional image generation

We first test the dependence of the accuracy of TDS on the number of particles in synthetic settings with tractable exact conditionals in Section 5.1. Section 5.2 compares TDS to alternative approaches on class-conditional image generation, and an image inpainting experiment is included in Appendix D.2.2. See Appendix D for all additional details.

Our evaluation includes: (1) TDS ; (2) TDS-IS , an importance sampler that uses TDS's proposal; (3) IS , a naive importance sampler described in Section 3.1; and (4) Gradient Guidance . For inpainting-type problems we further include: (5) Replacement method ; and (6) SMC-Diff .

Each SMC sampler forms an approximation to p θ ( x 0 | y ) with K weighted particles ( ∑ K k ′ =1 w k ′ ) -1 · ∑ K k =1 w k δ x 0 k . Gradient guidance and replacement method are considered to form a similar particlebased approximation with K independent samples viewed as K particles with uniform weights.

Compute Cost: Compared to unconditional generation, TDS has compute cost that is (i) larger by a constant factor due to the need to 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 to be achieved.

By comparison, conditional training methods provide fast inference by amortization [13], in which most computation is done ahead of time. Hence they may be preferable for applications where sampling time computation is a primary concern. On the other hand, TDS may be preferable when the likelihood criterion is readily available (e.g. through an existing classifier on clean data) but training an amortized model poses challenges; for example, the labeled data, neural-network engineering expertise, and up-front compute resources required for amortization training can be prohibitive.

5.1 Applicability and precision of TDS in two dimensional simulations

We explore two questions in this section: (i) what sorts of conditioning information can be handled by TDS and other methods, and (ii) how does the precision of TDS depend on the number of particles?

To study these questions, we first consider an unconditional diffusion model p θ ( x 0 ) approximation of a bivariate Gaussian. For this choice, the marginals of the forward process are also Gaussian, and so we may define p θ ( x 0: T ) with an analytically tractable score function without neural network approximation. Consequently, we can analyze the performance without the influence of score network approximation errors. And the choice of a two-dimensional diffusion permits close approximation of exact conditional distributions by numerical integration that can then be used as ground truth.

We consider three test cases defining the conditional information: (1) Smooth likelihood: y is an observation of the Euclidean norm of x with Laplace noise, with p θ ( y | x 0 ) = exp {|∥ x 0 ∥ 2 -y |} / 2 . This likelihood is smooth almost everywhere. 3 (2) Inpainting: y is an observation of the first dimension of x 0 , with p θ ( y | x 0 , M = 0) = δ y ( x 0 0 ) . (3) Inpainting with degrees-of-freedom: y is a an observation of either the first or second dimension of x 0 , with M = { 0 , 1 } and p θ ( y | x 0 ) = 1 2 [ δ y ( x 0 0 ) + δ y ( x 0 1 )] . . In all cases we fix y = 0 and consider estimating E p θ [ x 0 | y ] .

Figure 1 reports the estimation error for the mean of the desired conditional distribution, i.e. ∥ ∑ w k x 0 k -E q [ x 0 | y ] ∥ 2 . TDS provides a computational-statistical trade-off: using more particles decreases mean square estimation error at the O (1 /K ) parametric rate (note the slopes of -1 in log-log scale) as expected from standard SMC theory [5, Ch. 11]. This convergence rate is shared by TDS, TDS-IS, and IS in the smooth likelihood case, and by TDS, SMCDiff and in the inpainting case; TDS-IS, IS and SMCDiff are applicable however only in these respective cases, whereas TDS applicable in all cases. The only other method which applies to all three settings is Gradient Guidance, which exhibits significant estimation error and does not improve with many particles.

5.2 Class-conditional image generation

Wenext study the performance of TDS on diffusion models with neural network approximations to the score functions. In particular, we study the class-conditional image generation task, which involves sampling an image from p θ ( x 0 | y ) ∝ p θ ( x 0 ) p y | x 0 ( y | x 0 ) , where p θ ( · ) is a pretrained diffusion model on images x 0 , y is a given image class, and p y | x 0 ( y |· ) is the classification likelihood. To assess the faithfulness of generation, we evaluate classification accuracy on predictions of conditional samples x 0 given y , made by the same classifier that specifies the likelihood. In all experiments, we follow the standard practice of returning the denoising mean on the final sample [18].

On the MNIST dataset, we compare TDS to TDS-IS, Gradient Guidance, and IS. Figure 2a compares the conditional samples of TDS and Gradient Guidance given class y = 7 . Samples from Gradient Guidance have noticeable artifacts, and most of them do not resemble the digit 7; by contrast, TDS produces authentic and correct digits.

Figure 2b presents an ablation study of the effect of # of particles K on MNIST. For all SMC samplers, more particles improve the classification accuracy, with K = 64 leading to nearly perfect accuracy. The performance of Gradient Guidance is constant with respect to K (in expectation). Notably, for fixed K , TDS and TDS-IS have comparable performance and outperform Gradient guidance and IS.

We next apply TDS to higher dimension datasets. Figure 2c shows samples from TDS ( K = 16) using a pre-trained diffusion model and a pretrained classifier on the ImageNet dataset ( 256 × 256 × 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.

TDS

Gradient Guidance

7

7

7

7

9

a

又

1

7

7

7

七

8

3

7

L

L

L

3

2

7

G

L

7

5

1

2

4

8

16

32

64

# of Particles (K)

20

60

100

Class. Accuracy (%)

TDS

TDS-IS

Gradient

Guidance

IS

TDS can be extended by exponentiating twisting functions with a twist scale . This extension is related to the existing literature of Classifier Guidance that considers re-scaling the gradient of the log classification probability. See Appendix D.2.1 for details and ablation study.

6 Case study in computational protein design: the motif-scaffolding problem

The biochemical functions of proteins are typically imparted by a small number of atoms, known as a motif , that are stabilized by the overall protein structure, known as the scaffold [36]. A central task in protein design is to identify stabilizing scaffolds in response to motifs expected to confer function. We here describe an application of TDS to this task, and compare TDS to the state-of-the-art conditionally-trained model, RFdiffusion. See additional details in Appendix E.

Given a generative model supported on designable protein structures p θ ( x 0 ) , suitable scaffolds may be constructed by solving a conditional generative modeling problem [34]. Complete structures are first segmented into a motif x 0 M and a scaffold x 0 ¯ M , i.e. x 0 = [ x 0 M , x 0 ¯ M ] . Putative compatible scaffolds are then identified by (approximately) sampling from p θ ( x 0 ¯ M | x 0 M ) [34].

While the conditional generative modeling approach to motif-scaffolding has produced functional, experimentally validated structures for certain motifs [37], the general problem remains open. Moreover, current methods for motif-scaffolding require one to specify the location of the motif within the primary sequence of the full scaffold; this choice can require expert knowledge and trial and error.

We hypothesized that improved motif-scaffolding could be achieved through accurate conditional sampling. To this end, we applied TDS to FrameDiff, a Riemannian diffusion model that parameterizes protein backbones as a collection of N rigid bodies (known as residues) in the manifold SE (3) N [39]. 4 Each of the N elements of SE (3) consists of a rotation matrix and a translation that parameterize the locations of the backbone atoms of each residue.

Likelihood, twisting, and degrees of freedom. The basis of our approach to the motif scaffolding problem is analogous to the inpainting case described in Section 3.3. We let p y | x 0 ( y | x 0 ) = δ y ( x 0 M ) , where y ∈ SE (3) M describes the coordinates of backbone atoms of an M = |M| residue motif. As such, to define twisting function we adopt the Riemannian TDS formulation described in Section 3.4.

To eliminate the requirement that the placement of the motif within the scaffold be pre-specified, we incorporate the motif placement as a degree of freedom. We (i) treat the indices of the motif within the chain as a mask M , (ii) let M be a set of possible masks of size equal to the length of the motif, and (iii) apply Equation (14) to average over these possible masks.

For some scaffolding problems, it is known that all motif residues must appear with contiguous indices. In this case we choose M to be the set of all possible contiguous masks. However, when motif residues are only known to appear in two or 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.

We similarly eliminate the global translation and rotation of the motif as a degree of freedom. The motivation is that restricting to one possible pose of the motif narrows the conditional distribution,

1

2

4

8

16

# Particles

0

5

10

15

20

25

Success rate (%)

1

10

10010000

1

10

100 1000

0.5 0.75 1.0

2.0

3.0

Twist Scale

(a) TDS motif-scaffolding success rate (test case 5IUS ) improves with more particles, and degrees of freedom, and twist-scale.

thereby making inference more challenging. For translation, we use a likelihood that is invariant to the motif's center-of-mass and placement in the final scaffold by choosing p y | x 0 ( y | x ) = δ Py ( Px M ) where P is a projection matrix that removes the center of mass of a vector [see e.g. 39, section 3.3]. For rotation, we average the likelihood across some number ( # Motif Rots. ) of possible rotations.

Ablation study. We first examine the impact of several parameters of TDS on success rate in an in silico self-consistency evaluation [34]. We begin with single problem ( 5IUS ) in the benchmark set introduced by [37], before testing on the full set. Figure 3a (Left) shows that success rate increases monotonically with the number of particles. Non-zero success rates in this setting required accounting for multiple motif locations; Figure 3a (Left) uses 1 , 000 possible motif locations and 100 rotations.

We tested the impact of the degrees of freedom by evaluating the success rate of TDS (K=1) with increasing motif locations and 100 rotations (Figure 3a Center Left), and increasing rotations and 1,000 locations (Figure 3a Center Right). The success rate was 0% without accounting for either degree of freedom, and increased with larger numbers of locations and rotations. We also explored including a heuristic twist scale as considered for image tasks (Section 5.2); in this case, the twist scale is a multiplicative factor on the logarithm of the twisting function. Figure 3a (Right) shows larger twist scales gave higher success rates on this test case, where we use 8 particles, 1,000 possible motif locations and 100 rotations. However, this trend is not monotonic for all problems (Figure P).

Evaluation on full benchmark. Wenext evaluate on TDS on a benchmark set of 24 motif-scaffolding problems [37] and compare to the previous state of the art, RFdiffusion. RFdiffusion operates on the same rigid body representation of protein backbones as FrameDiff. TDS is run with K=8, twist scale=2, and 100 rotations and 1 , 000 motif location degrees of freedom ( 100 , 000 combinations total).

Overall, TDS (applied to FrameDiff) and RFdiffusion have comparable performance (Figure 3b); each provides a success rate higher than the other in 11/24 cases; on two problems both methods have a 0% success rate (full results in Figure O). This performance is obtained despite the fact that FrameDiff, unlike RFdiffusion, is not trained to perform motif scaffolding. The division between problems on which each method performs well is primarily explained by total scaffold length, with TDS providing higher success rates on smaller scaffolds.

We suspect the shift in performance with scaffold length owes properties of the underlying diffusion models. First, long backbones generated unconditionally by RFdiffusion are designable with higher frequency than those generated by FrameDiff [39]. Second, unlike RFdiffusion, FrameDiff can not condition on the fixed motif sequence.

7 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.

A limitation of TDS is its requirement for additional computes to simulate multiple particles. While we observe improved performance with just two particles in some cases, the optimal number is problem dependent. Moreover, the computational efficiency depends on how closely the twisting functions approximate exact conditionals, which depends on the unconditional model and conditioning information. Lastly, choosing twisting functions for generic constraints may be challenging. Our future work will focus on addressing these limitations and improving the computational efficiency.

References

Arpit Bansal, Hong-Min Chu, Avi Schwarzschild, Soumyadip Sengupta, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Universal guidance for diffusion models. arXiv preprint arXiv:2302.07121 , 2023.

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.

Sourav Chatterjee and Persi Diaconis. The sample size required in importance sampling. The Annals of Applied Probability , 2018.

Gregory Chirikjian and Marin Kobilarov. Gaussian approximation of non-linear measurement models on Lie groups. In 53rd IEEE Conference on Decision and Control . IEEE, 2014.

Nicolas Chopin and Omiros Papaspiliopoulos. An introduction to sequential Monte Carlo . Springer, 2020.

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.

Justas Dauparas, Ivan Anishchenko, Nathaniel Bennett, Hua Bai, Robert J Ragotte, Lukas F 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 learning-based protein sequence design using ProteinMPNN. Science , 2022.

Valentin De Bortoli, Emile Mathieu, Michael John Hutchinson, James Thornton, Yee Whye Teh, and Arnaud Doucet. Riemannian score-based generative modelling. In Advances in Neural Information Processing Systems , 2022.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat GANs on image synthesis. Advances in Neural Information Processing Systems , 2021.

Arnaud Doucet, Nando De Freitas, and Neil James Gordon. Sequential Monte Carlo methods in practice . Springer, 2001.

Yilun Du, Conor Durkan, Robin Strudel, Joshua B Tenenbaum, Sander Dieleman, Rob Fergus, Jascha Sohl-Dickstein, Arnaud Doucet, and Will Grathwohl. Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and MCMC. In International Conference on Machine Learning , 2023.

Bradley Efron. Tweedie's formula and selection bias. Journal of the American Statistical Association , 2011.

Samuel Gershman and Noah Goodman. Amortized inference in probabilistic reasoning. In Proceedings of the Annual Meeting of the Cognitive Science Society , 2014.

Neil J Gordon, David J Salmond, and Adrian FM Smith. Novel approach to nonlinear/nonGaussian Bayesian state estimation. In IEE Proceedings F (Radar and Signal Processing) . IET, 1993.

Pieralberto Guarniero, Adam M Johansen, and Anthony Lee. The iterated auxiliary particle filter. Journal of the American Statistical Association , 2017.

Jeremy Heng, Adrian N Bishop, George Deligiannidis, and Arnaud Doucet. Controlled sequential Monte Carlo. Annals of Statistics , 2020.

Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598 , 2022.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems , 2020.

Jonathan Ho, Tim Salimans, Alexey A Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. In Advances in Neural Information Processing Systems , 2022.

John M. Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zídek, Anna Potapenko, Alex Bridgland, Clemens Meyer, Simon A A Kohl, Andy Ballard, Andrew Cowie, Bernardino Romera-Paredes, 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.

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.

Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. In Advances in Neural Information Processing Systems , 2022.

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 IEEE/CVF Conference on Computer Vision and Pattern Recognition , 2022.

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.

Christian A Naesseth, Fredrik Lindsten, and Thomas B Schön. Elements of sequential Monte Carlo. Foundations and Trends in Machine Learning , 2019.

Michael K Pitt and Neil Shephard. Filtering via simulation: Auxiliary particle filters. Journal of the American Statistical Association , 1999.

Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with CLIP latents. arXiv preprint arXiv:2204.06125 , 2022.

Herbert Robbins. An empirical Bayes approach to statistics. In 3rd Berkeley Symposium on Mathematical Statistics , 1956.

Chitwan Saharia, William Chan, Huiwen Chang, Chris Lee, Jonathan Ho, Tim Salimans, David Fleet, and Mohammad Norouzi. Palette: Image-to-image diffusion models. In ACMSIGGRAPH 2022 Conference Proceedings , 2022.

Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David J Fleet, and Mohammad Norouzi. Image super-resolution via iterative refinement. IEEE Transactions on Pattern Analysis and Machine Intelligence , 2022.

Yuyang Shi, Valentin De Bortoli, George Deligiannidis, and Arnaud Doucet. Conditional simulation using diffusion Schrödinger bridges. In Uncertainty in Artificial Intelligence , 2022.

Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. Pseudoinverse-guided diffusion models for inverse problems. In International Conference on Learning Representations , 2023.

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.

Brian L Trippe, Jason Yim, Doug Tischer, Tamara Broderick, David Baker, Regina Barzilay, and Tommi Jaakkola. Diffusion probabilistic modeling of protein backbones in 3D for the motif-scaffolding problem. In International Conference on Learning Representations , 2023.

Pascal Vincent. A connection between score matching and denoising autoencoders. Neural Computation , 2011.

Jue Wang, Sidney Lisanza, David Juergens, Doug Tischer, Joseph L Watson, Karla M Castro, Robert Ragotte, Amijai Saragovi, Lukas F Milles, Minkyung Baek, 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.

Joseph L. Watson, David Juergens, Nathaniel R. Bennett, Brian L. Trippe, Jason Yim, Helen E. Eisenach, Woody Ahern, Andrew J. Borst, Robert J. Ragotte, Lukas F. Milles, Basile I. M. Wicky, Nikita Hanikel, Samuel J. 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 S. Jaakkola, Frank DiMaio, Minkyung Baek, and David Baker. De novo design of protein structure and function with RFdiffusion. Nature , 2023.

Nick Whiteley and Anthony Lee. Twisted particle filters. The Annals of Statistics , 2014.

Jason Yim, Brian L Trippe, Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. SE(3) diffusion model with application to protein backbone generation. In International Conference on Machine Learning , 2023.

Guanhua Zhang, Jiabao Ji, Yang Zhang, Mo Yu, Tommi Jaakkola, and Shiyu Chang. Towards coherent image inpainting using denoising diffusion implicit models. In International Conference on Machine Learning , 2023.

Appendix

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 to the final step of TDS in inpainting and inpainting with degrees of freedom applications. And Appendix A.5 provides a full statement and proof of Theorem 1.

A.1 Diffusion models with non-constant noise variance and variance-preserving scaling

TDS as developed in Section 3 assumed noise with constant variance σ 2 added at each timestep. This choice corresponds to a particular discretization of a 'variance exploding' (VE) diffusion model [33]. Here we describe generalizations of TDS to non-constant variance schedules, and to variance preserving (VP) diffusion models [18]. We adopt the VP formulation in our experiments; though it introduces additional notation, it is known to work well in practice [33].

TDS algorithm for variance exploding diffusion models. VE diffusion models define the forward process q by

where σ 2 t is an increasing sequence of variances such that q ( x T ) ≈ N (0 , ¯ σ 2 T ) , with ¯ σ 2 t := ∑ t t ′ =1 σ 2 t ′ for t = 1 , · · · , T . And so one can set p ( x T ) = N (0 , ¯ σ 2 T ) to match q ( x T ) . Notably, Equation (16) implies the conditional q ( x t | x 0 ) = N ( x t ; x 0 , ¯ σ 2 t ) .

The reverse diffusion process p θ is parameterized as

where the score network s θ is modeled through a denoiser network ˆ x θ by s θ ( x t , t ) := (ˆ x θ ( x t , t ; θ ) -x t ) / ¯ σ 2 t . Note that the constant schedule is a special case where σ 2 t = σ 2 , and ¯ σ 2 t = tσ 2 for all t .

The TDS algorithm for general VE models is described in Algorithm 1, where tσ 2 in Line 6 is replaced by ¯ σ 2 t and σ 2 in Line 7 is replaced by σ 2 t .

Extension to variance preserving diffusion models. Another widely used diffusion framework is variance preserving (VP) diffusion models [18]. VP models define the forward process q by

where σ 2 t is a sequence of increasing variances chosen such that q ( x T ) ≈ N (0 , 1) , and so one can set p ( x T ) = N (0 , 1) . Define α t := 1 -σ 2 t , ¯ α t := ∏ t t ′ =1 α t ′ , and ¯ σ 2 t := 1 -¯ α t . Then the marginal conditional of eq. (16) is q ( x t | x 0 ) = N ( x t ; √ ¯ α t x 0 , ¯ σ 2 t ) . The reverse diffusion process p θ is parameterized as

where s θ is now defined through the denoiser ˆ x θ by s θ ( x t , t ) := ( √ ¯ α t ˆ x θ ( x t , t ; θ ) -x t ) / ¯ σ 2 t .

TDS in Algorithm 1 extends to VP models as well, where the conditional score approximation in Line 6 is changed to ˜ s k ← ( √ ¯ α t +1 ˆ x θ ( x t +1 ) -x t +1 ) / ¯ σ 2 t +1 + ∇ t +1 log ˜ p t +1 , and the proposal in

A.2 TDS parameters

Proposal variance. The proposal distribution in Line 7 of Algorithm 1 is associated with a variance parameter ˜ σ 2 . In general, this parameter can be dependent on the time step, i.e. replacing ˜ σ 2 by some ˜ σ 2 t +1 in Line 7. Unless otherwise specified, we set ˜ σ 2 t +1 := Var p θ [ x t | x t +1 ] the variance of the unconditional diffusion model (typically learned along with the score network during training).

Resampling strategy. The mulinomial resampling strategy in Line 4 of Algorithm 1 can be replaced by other strategies, see [25, Chapter 2] for an overview. In our experiments, we use the systematic resampling strategy.

In addition, one can consider setting an effective sample size (ESS) threshold (between 0 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 e.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 ν t ( x 0: T ) provide approximations to the final target p θ ( x 0: T | y ) that become increasingly accurate as t approaches 0. We obtain Equation (12) by substituting the proposal distributions in Equation (9) and weights in Equation (11) into Equation (4) and simplifying.

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 ( t = 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 ˜ p θ ( y | x 0 ) = p y | x 0 ( y | x 0 ) . This choice introduces error in the final target of TDS relative to the exact conditional p θ ( x 0: T | y ) .

For inpainting, to maintain asymptotic exactness one may instead choose the final proposal and weights as

One can verify the resulting final target is p θ ( x 0 | x 0 M = y ) according to eq. (4).

Similarly, for inpainting with degrees of freedom, one may define the final proposal and weight as

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 is increased and (ii) discuss when these conditions will hold in practice for the twisting function ˜ p θ ( y | x t ) introduced in Section 3.

We begin with a theorem providing sufficient conditions for asymptotic accuracy of TDS.

Theorem 2. Let p θ ( x 0: T ) be a diffusion generative model (defined by eqs. (1) and (3) ) with

with variances σ 2 1 , . . . , σ 2 T . Let ˜ p θ ( y | x t ) be twisting functions, and

be proposals distributions for t = 0 , . . . , T -1 , and let P K = ∑ K k =1 w 0 k δ x 0 k for weighted particles { ( x 0 k , w 0 k ) } K k =1 returned by Algorithm 1 with K particles. Assume

(a) the final twisting function is the likelihood, ˜ p θ ( y | x 0 ) = p y | x 0 ( y | x 0 ) ,

(b) the first twisting function ˜ p θ ( y | x T ) , and the ratios of subsequent twisting functions ˜ p θ ( y | x t ) / ˜ p θ ( y | x t +1 ) are positive and bounded,

(c) each log ˜ p θ ( y | x t ) with t > 0 is continuous and has bounded gradients in x t , and

(d) the proposal variances are larger than the model variances, i.e. for each t, ˜ σ 2 t > σ 2 t .

Then P K converges setwise to p θ ( x 0 | y ) with probability one, that is for every set A, lim K →∞ P K ( A ) = ∫ A p θ ( x 0 | y ) dx 0 .

The assumptions of Theorem 2 are readily satisfied in common applications.

Assumption (a) may be satisfied by construction by choosing ˜ p θ ( y | x 0 ) = p y | x 0 ( y | x 0 ) by defining ˆ x θ ( x 0 , t = 0) = x 0 .

Assumption (b) is satisfied if (i) p y | x 0 ( y | x ) is smooth in x and everywhere positive and (ii) ˜ p θ ( y | x t ) = p y | x 0 ( y | ˆ x θ ( x t )) where ˆ x θ ( x t ) takes values in some compact domain. An alternative sufficient condition for Assumption (b) is for p y | x 0 ( y | x ) to be positive and bounded away from zero; this latter condition will hold when, for example, p y | x 0 ( y | x ) is a classifier fit with regularization.

Assumption (c) is the strongest assumption. It will be satisfied, for example, if (i) ˜ p θ ( y | x t ) = p y | x 0 ( y | ˆ x θ ( x t )) , and (ii) p y | x 0 ( y | x ) and ˆ x θ ( x t ) are smooth in x and x t , with uniformly bounded gradients. While smoothness of ˆ x θ ( · , t ) can be encouraged by the use of skip-connections and regularization, ˆ x θ ( · , t ) may present sharp transitions, particularly for t close to zero.

Assumption (d), that the proposal variances satisfy ˜ σ 2 t > σ 2 t 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 ˜ σ 2 t = σ 2 t , but alternatively the assumption could be met by inflating each ˜ σ 2 t by some arbitrarily small δ without markedly impacting the behavior of the sampler.

Proof of Theorem 2: Theorem 2 characterizes a set of conditions under which SMC algorithms converge. We restate this result below in our own notation.

Theorem 3 (Chopin and Papaspiliopoulos [5] - Proposition 11.4) . Let { ( x 0 k , w 0 k ) } K k =1 be the particles and weights returned at the last iteration of a sequential Monte Carlo algorithm with K particles using multinomial resampling. If each weighting function w t ( x t , x t +1 ) is positive and bounded, then for every bounded, ν 0 -measurable function ϕ of x t

with probability one.

An immediate consequence of Theorem 3 is the setwise convergence of the discrete measures, ˆ P K = ∑ K k =1 w 0 k δ x 0 k . This can be seen by taking for each ϕ ( x ) = I [ x ∈ A ] for any ν 0 -measurable set A . The theorem applies both in the Euclidean setting, where each x t k ∈ R D , as well as the Riemannian setting.

We now proceed to prove Theorem 2.

Proof. To prove the theorem we show (i) the x 0 marginal final target ν 0 is p θ ( x 0 | y ) and then (ii) P K converges setwise to ν 0 .

We first show (i) by manipulating ν 0 in Equation (4) to obtain p θ ( x 0: T | y ) . From Equation (4) we first have

The final line reveals that once we marginalize out x 1: T we obtain ν 0 ( x 0 ) = p ( x 0 | y ) as desired.

We next show that P K converges to ν 0 with probability one by applying Theorem 3. To apply Theorem 3 it is sufficient to show that the weights at each step are upper bounded, as they are defined through (ratios of) probabilities and hence are positive. Since there are a finite number of steps T, it is enough to show that each w t is bounded. The inital weight is the initial twisting function, w T ( x T ) = ˜ p θ ( y | x T ) , 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 logweighting functions as

.

and show independently that log ˜ p θ ( y | x t ) / ˜ p θ ( y | x t +1 ) and log p θ ( x t | x t +1 ) /r t ( x t | x t +1 ) are bounded. The first term log ˜ p θ ( y | x t ) / ˜ p θ ( y | x t +1 ) is again bounded by Assumption (b), and we proceed to the second.

That log p θ ( x t | x t +1 ) /r t ( x t | x t +1 ) is bounded follows from Assumptions (c) and (d). First write with ˆ µ = x t +1 + σ 2 s ( x t +1 ) , and

for ˆ µ ψ = ˆ µ + σ 2 t +1 ∇ x t +1 log ˜ p θ ( y | x t +1 ) . The log-ratio then simplifies as

Rearrange and let C = log | 2 πσ 2 t +1 I | -1 / 2 / | 2 π ˜ σ 2 t +1 I | -1 / 2

Expand and rearrange ∥ ˆ µ ψ -x t ∥ 2 = ∥ ˆ µ -x t ∥ 2 +2 ⟨ ˆ µ ψ -ˆ µ, ˆ µ -x t ⟩ + ∥ ˆ µ ψ -ˆ µ ∥ 2

Let C ′ = C -1 2 ˜ σ 2 t +1 ∥ ˆ µ ψ -ˆ µ ∥ 2 and rearrange. Note that ∥ ˆ µ ψ -ˆ µ ∥ 2 < ∞ by Assm. (c).

Apply Cauchy-Schwarz

Upper-bounding using that max x -a 2 x 2 + bx = b 2 2 a for a = σ -2 t +1 -˜ σ -2 t +1 > 0 by Assm. (d).

t +1 t +1 Note that ˆ µ ψ = ˆ µ + σ 2 t +1 ∇ x t +1 log ˜ p θ ( y | x t +1 ) .

The final line follows from Assumption (c), that the gradients of the twisting functions are bounded. The above derivation therefore provides that each w t is bounded, concluding the proof.

B Riemannian Twisted Diffusion Sampler

This section provides additional details on the extension of TDS to Riemannian diffusion models introduced in Section 3. We first introduce the tangent normal distribution. We then provide with background on Riemannian diffusion models, which we parameterize with the tangent normal. Then we describe the extension of TDS to these models. Finally we show how Algorithm 1 modifies to this setting.

The tangent normal distribution. Just as the Gaussian is the parametric family underlying Euclidean diffusion models in Equation (3), the tangent normal (see e.g. [4, 8]) underlies generation in Riemannian diffusion models so we review it here.

We take the tangent normal to be the distribution is implied by a two step procedure. Given a variable x in the manifold, the first step is to sample a variable ¯ y in T x , the tangent space at x ; if { h 1 , . . . , h D } is an orthonormal basis of T x one may generate

with ϵ d i.i.d. ∼ N (0 , 1) , and σ 2 > 0 a variance parameter. The second step is to project ¯ y back onto the manifold to obtain y = exp x { ¯ y } where exp x {·} denotes the exponential map at x. The resulting distribution of y is denoted as T N x ( µ, σ 2 ) . By construction, T x is a Euclidean space with its origin at x, so when ∥ µ ∥ 2 = 0 , T N x ( µ, σ 2 ) is centered on x. And since the geometry of a Riemannian manifold is locally Euclidean, when σ 2 is small the exponential map is close to the identity map and the tangent normal is essentially a narrow Gaussian distribution in the manifold at x. Finally, we use T N x ( y ; µ, σ 2 ) to denote the density of the tangent normal evaluated at y .

Because this procedure involves a change of variables from ¯ y to y (via the exponential map), to compute the tangent normal density one computes

where exp -1 x { y } is the inverse of the exponential map (from the manifold into T x ), and ∣ ∣ ∣ ∂ ∂y exp -1 x { y } ∣ ∣ ∣ is the determinant of the Jacobian of the exponential map; when the manifold

∣ ∂ -1 ∣

lives in a higher dimensional subset of R, we take ∣ ∣ ∂y exp x { 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 ¯ x t in T x t +1 ; if { h 1 , . . . , h D } is an orthonormal basis of T x t +1 one may generate

with ϵ d i.i.d. ∼ N (0 , 1) . One then projects ¯ x t back onto the manifold to obtain x t = exp x t +1 { ¯ x t } where exp x {·} denotes the exponential map at x.

This is equivalent to sampling x t from a tangent normal distribution as

TDS for Riemannian diffusion models. To extend TDS, appropriate analogues of the twisted proposals and weights are all that is needed. For this extension we require that the diffusion model is also associated with a manifold-valued denoising estimate ˆ x θ as will be the case when, for example, s θ ( x t , t ) := ∇ x t log q ( x t | x 0 = ˆ x θ ) for ˆ x θ = ˆ x θ ( x t , t ) . In contrast to the Euclidean case, a relationship between a denoising estimate and a computationally tractable score approximation may not always exist for arbitrary Riemannian manifolds; however for Lie groups when the the forward diffusion is the Brownian motion, tractable score approximations do exist [39, Proposition 3.2].

For the case of positive and differentiable p y | x 0 ( y | x 0 ) , we again choose twisting functions ˜ p θ ( y | x t ) := p y | x 0 ( y | ˆ x θ ( x t )) .

Next are the inpainting and inpainting with degrees of freedom cases. Here, assume that x 0 lives on a multidimensional manifold (e.g. SE (3) N ) and the unmasked observation y = x 0 M with M ⊂ { 1 , . . . , N } on a lower-dimensional sub-manifold (e.g. SE (3) | M | , 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 a Tangent normal as

For all cases, we propose the twisted proposal as

where as in the Euclidean case s θ ( x t , y ) = s θ ( x t ) + ∇ x t log ˜ p ( y | x t ) .

Weights at intermediate steps are computed as in the Euclidean case (Equation (11)):

While the proposal and target contribute identical Jacobian determinant terms that cancel out, they remain in the twisting functions.

Adapting the TDS algorithm to the Riemannian setting. To translate the TDS algorithm to the Riemannian setting we require only two changes.

The first is on Algorithm 1 Line 6. Here we assume that the unconditional score is computed through the transition density function:

Note that the gradient above ignores the dependence of ˆ x θ on x t +1 .

The conditional score approximation in Line 6 is then replaced with

Notably, ˜ s k is a T x t +1 k -valued Riemannian gradient.

The second change is to make the proposal on Algorithm 1 Line 7 a tangent normal, as defined in eq. (17).

C Additional Related Work

Training with conditioning information. Several approaches involve training a neural network to directly sample from a conditional diffusion models. These approaches include (i) conditional training with embeddings of conditioning information, e.g. for denoising images [30], and text-toimage generation [27], (ii) conditioning training with a subset of the state space, e.g. for protein design [37], and image inpainting [29], (iii) classifier-guidance [9, 33], an approach for generating samples of a desired class, which requires training a time-dependent classifier to approximate, for each t, p θ ( y | x t ) , and training such a time-dependent classifier may be inconvenient and costly), and (iv) classifier-free guidance [17], an approach that builds on classifier-guidance without training a classifier, and instead trains a diffusion model with class information as additional input.

Langevin and Metropolis-Hastings steps. Some prior work has explored using Markov chain Monte Carlo steps in sampling schemes to better approximate conditional distributions. For example unadjusted Langevin dynamics [19, 33] or Hamiltonian Monte Carlo [11] in principle permit asymptotically exact sampling in the limit of many steps. However these strategies are only guaranteed to target the conditional distributions of the joint model under the (unrealistic) assumption that the unconditional model exactly matches 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 diffusion models on two dimensional problems

Forward process. Our forward process is variance preserving (as described in Appendix A) with T = 100 steps and a quadratic variance schedule. We set σ 2 t = σ 2 min +( t T ) 2 σ 2 max with σ 2 min = 10 -5 and σ 2 max = 10 -1 .

Unconditional target distributions and likelihood. We evaluate the different methods with two different unconditional target distributions:

SmoothLikelihood

Inpaint

Inpaint(DegreesofFreedom)

10-1

10-1

Error

10-1

TDS

lare

10-2

10-2

10-2

TDS-IS

S

enb

10-3

10-3

10-3

Guidance

Gradient

SMC-Diff

lean

10-4

10-4

10-4

Replacement

M

101

102

103

104

101

102

103

104

101

102

103

104

# ofParticles (K)

# of Particles (K)

# of Particles (K)

A bivariate Gaussian with mean at ( 1 2 , 1 2 ) and covariance 0 . 9 and

A Gaussian mixture with three components with mixing proportions [0 . 3 , 0 . 5 , 0 . 2] , means [(1 . 54 , -0 . 29) , ( -2 . 18 , 0 . 57) , ( -1 . 09 , -1 . 40)] , and 0 . 2 standard deviations.

We evaluate on conditional distributions defined by the three likelihoods described in Section 5.1.

Figure D provides results analogous to those in Figure 1 but with the mixture of Gaussians unconditional target. In this second example we evaluate a with a variation of the inpainting degrees of freedom case wherein we consider y = 1 and M = { [1] , [2] } , so that p y | x 0 ( y | x 0 ) = δ y ( x 0 1 )+ δ y ( x 0 2 ) .

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 × 1 dimensions) in Appendix D.2.1 and Appendix D.2.2. And we show TDS's applicability to higher-dimensional datasets, namely CIFAR10 ( 32 ×× 3 dimensions) and ImageNet256 ( 256 × 256 × 3 dimensions) in Appendix D.2.3.

D.2.1 Class-conditional generation on MNIST

Set-up. For MNIST, we set up a diffusion model using the variance preserving framework. The model architecture is based on the guided diffusion codebase 5 with the following specifications: number of channels = 64, attention resolutions = "28,14,7", number of residual blocks = 3, learn sigma (i.e. to learn the variance of p θ ( x t -1 | x t ) ) = True, resblock updown = True, dropout = 0.1, variance schedule = "linear". We trained the model for 60k epochs with a batch size of 128 and a learning rate of 10 -4 on 60k MNIST training images. The model uses T = 1 , 000 for training and T = 100 for sampling.

The classifier used for class-conditional generation and evaluation is a pretrained ResNet50 model. 6 This classifier is trained on the same set of MNIST training images used by diffusion model training.

In addition we include a variation called TDS-truncate that truncates the TDS procedure at t = 10 and returns the prediction ˆ x θ ( x 10 ) .

Sample plots. To supplement the sample plot conditioned on class 7 in Figure 2a, we present samples conditioned on each of the remaining 9 classes and from other methods. We observe that samples from Gradient Guidance have noticeable artifacts, whereas the other 4 methods produce authentic and correct digits. However, most samples from IS or TDS-IS are identical due to the collapse of importance weights. By contrast, samples from TDS and TDS-truncate have greater diversity, with the latter exhibiting slightly more variations

Ablation study on twist scales. We consider exponentiating and re-normalizing twisting functions by a twist scale γ , i.e. setting new twisting functions to ˜ p θ ( y | x t ; γ ) ∝ p y | x 0 ( y | ˆ x θ ( x t )) γ . In particular, when t = 0 , we set p y | x 0 ( y | x 0 ; γ ) ∝ p y | x 0 ( y | x 0 ) γ . This modification suggests that the targeted conditional distribution is now

TDS

TDS-truncate

TDS-IS

Gradient Guidance

IS

?

3 3 3 3

3 3 3 3 3 3 3 3 3 3 3 3 3 3

3 3 3 3 3 3

3

2

3

9

3

D

3

G

3

5

5

6

D

9

D

6

9

9

9

6

9

9

q

g

9

9

G

Q

1

2

4

8 16 32 64

# of Particles (K)

50

75

100

Class. Accuracy (%)

TDS

1

2

4

8 16 32 64

# of Particles (K)

TDS-IS

1

2

4

8 16 32 64

# of Particles (K)

Gradient Guidance

Twist Scale

1

2

3

4

10

1

2

3

4

10

Twist Scale

50

75

100

Human-rated

Accuracy (%)

TDS K=64

TDS K=2

TDS-IS K=2

Gradient Guidance

(b) Human-rated classification accuracy v.s. twist scale, for TDS ( K = 64 , 2 ), TDS-IS ( K = 64 , 2 ) and Gradient Guidance. Results are averaged over 640 randomly chosen samples with error bands indicating 2 standard errors. For TDS ( K = 64 ), increasing the twist scale generally decreases the human-rated accuracy. For remaining methods, a moderate increase in twist scale improves the human-rated accuracy; however, excessively large twist scale can hurt the accuracy.

By setting γ > 1 , the classification likelihood becomes sharper, which is potentially helpful for twisting the samples towards a specific class. The TDS algorithm (and likewise TDS-IS and Gradient Guidance) still apply with this new definition of twisting functions. The use of twist scale is similar to the guidance scale introduced in the literature of Classifier Guidance, which is used to multiply the gradient of the log classification probability [9].

In Figure F, we examine the effect of varying twist scales on classification accuracy of TDS, TDS-IS and Gradient Guidance. We consider two ways to evaluate the accuracy. First, classification accuracy computed by a neural network classifier , where the evaluation setup is the same as in Section 5.2. Second, the human-rated classification accuracy, where a human (one of the authors) checks if a generated digit has the right class and does not have artifacts. Since human evaluation is expensive, we only evaluate TDS, TDS-IS (both with K = 64 , 2 ) and Gradient Guidance. In each run, we randomly sample one particle out of K particles according to the associated weights. We conduct 64 runs for each class label, leading to a total of 640 samples for human evaluation.

Figure Fa depicts the classification accuracy measured by a neural network classifier. We observe that in general larger twist scale improves the classification accuracy. For TDS and TDS-IS, the improvement is more significant for smaller number of particles K used.

Figure Fb depicts the human-rated accuracy. In this case, we find that larger twist scale is not necessarily better. A moderately large twist scale ( γ = 2 , 3 ) generally helps increasing the accuracy, while an excessively large twist scale ( γ = 10 ) decreases the accuracy. An exception is TDS with K = 64 particles, where any γ > 1 leads to worse accuracy compared to the case of γ = 1 . Study on twist scale aside, we find that using more particles K help improving human-rated accuracy (recall that Gradient Guidance is a special case of TDS with K = 1 ): given the same twist scale, TDS with K = 1 , 2 or 64 has increasing accuracy. In addition, both TDS and TDS-IS with K = 64 and γ = 1

Twist scale=1

Twist scale=2

Twist scale=3

Twist scale=4

Twist scale=10

D

6

6

6

6

6

6

6

G

G

C

6

1

G

6

9

6

6

6

6

6

6

G

6

6

G

6

6

6

6

6

G

6

2

6

6

6

6

6

G

6

6

b

G

G

6

6

6

5

6

6

6

6

6

6

6

6

G

6

6

2

？

2

6

6

6

6

6

6

6

6

9

6

6

6

6

6

6

6

6

?

G

C

6

6

6

G

6

6

G

G

9

6

6

6

6

6

6

6

G

6

6

b

6

6

9

G

6

6

6

6

6

G

6

G

6

6

6

6

6

9

G

6

6

G

6

6

Twist scale=1

Twist scale=2

Twist scale=3

Twist scale=4

Twist scale=10

5

6

6

6

G

6

6

G

G

G

6

可

6

6

6

G

6

Q

小

G

6

G

G

6

6

?

G

G

e

6

6

G

6

6

6

6

G

6

6

6

6

6

G

?

6

6

6

6

9

G

G

6

6

6

6

6

6

Q

6

6

6

6

L

B

6

6

6

G

C

6

6

6

b

6

C

6

6

6

6

6

G

6

G

6

羽

G

G

6

6

6

6

6

6

G

日

G

6

G

6

G

6

6

6

9

飞

6

G

6

9

6

6

6

Q

6

G

6

6

G

6

6

G

6

6

6

G

9

6

6

6

6

6

6

B

6

6

6

G

6

1

6

6

6

6

?

Twist scale=1

Twist scale=2

Twist scale=3

Twist scale=4

Twist scale=10

6

6

9

6

C

C.

6

6

6

6

G

6

6

6

6

o

S

E.

6

E

G

6

G

6

G

G

6

G

G

G

G

6

6

6

6

b

6

6

6

9

9

6

6

白

6

6

6

6

6

6

b

6

G

6

G

9

9

6

6

9

6

6

6

6

6

D

6

6

6

G

6

6

6

G

6

6

6

の

6

6

6o

G

6

6

6

6

6

6

6

6

9

G

G

6

6

6

G

6

6

6

6

6

G

2

6

G

6

6

6

6

6

G

6

6

6

C

G

6

の

6

6

b

G

6

6

G

6

G

2

G

6

6

6

G

6

9

6

6

6

5

Twist scale=1

Twist scale=2

Twist scale=3

Twist scale=4

Twist scale=10

B

6

G

6

6

G

G

9

6

lo

D

G

8

c

&

6

6

le

6

o

6

6

6

6

(e

lo

6

G

O

lo

b

6

6

6

6

の

G

6

6

6

6

D

6

6

6

G

6

6

6

D

6

6

G

2

4

6

7

6

D

6

6

a

6

G

6

D

6

G

G

6

6

b

6

G

6

6

6

6

6

6

o

6

6

6

6

?

b

6

6

9

Q

G

4

6

G

4

6

G

6

6

C

22

6

6

6

6

6

b

Twist scale=1

Twist scale=2

Twist scale=3

Twist scale=4

Twist scale=10

B

R

G

6

B

G

19

C

G

盈

8

Q

le

Q

GQ

6

6

(e

8

G

6

6

6

G

9

G

6

G

6

6

6

2

6

6

6

A

6

9

6

6

6

S

q

6

G

4

日

C

6

6

6

D

Co

6

6

6

6

6

6

o

6

6

6

6

6

6

D

θ

6

G

2

9

6

6

100

75

50

25

0

Time Step t

0

32

64

Effect. Sample Size

TDS

TDS-IS

IS

have almost perfect accuracy. The effect of K on human-rated accuracy is consistent with previous findings with neural network classifier evaluation in Section 5.2.

We note that there is a discrepancy on the effects of twist scales between neural network evaluation and human evaluation. We suspect that when overly large twist scale is used, the generated samples may fall out of the data manifold; however, they may still retain features recognizable to a neural network classifier, thereby leading to a low human-rated accuracy but a high classifier-rated accuracy. To validate this hypothesis, we present samples conditioned on class 6 in Figure G. For example, in Figure Ge, Gradient Guidance with γ = 1 has 31 good-quality samples out of 64, and the rest of the samples often resamble the shape of other digits, e.g. 3,4,8; and Gradient Guidance with γ = 10 has 34 good-quality samples, but most of the remaining samples resemble 6 with many artifacts.

Effective sample size. Effective sample size (ESS) is a common metric used to diagnose the performance of SMC samplers, which is defined as ( ∑ K k =1 w t k ) 2 / ( ∑ K k =1 ( w t k ) 2 ) for K weighted particles { x t k , w t k } K k =1 . Note that ESS is always bounded between 0 and K .

Figure H depicts the ESS trace comparison of TDS, TDS-IS, and IS. TDS has a general upward trend of ESS approaching K = 64 . Though in the final few steps ESS of TDS drops by a half, it is still higher than that of TDS-IS and IS that deteriorates to around 1 and 6 respectively. In practice, one can use the TDS-truncate variant introduced above to avoid the ESS drop in the final few steps.

D.2.2 Image inpainting on MNIST

The inpainting task is to sample images from p θ ( x 0 | x 0 M = y ) given observed part x 0 M = y . Here we segment x 0 = [ x 0 M , x 0 ¯ M ] into observed dimensions M and unobserved dimensions ¯ M , as described in Section 3.3.

In this experiment, we 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 of an image is observed.

We run TDS, TDS-IS, Gradient Guidance, SMC-Diff, and Replacement method to inpaint 10,000 validation images. We also include TDS-truncate that truncates the TDS procedure at t = 10 and returns ˆ x θ ( x 10 ) .

Twisting function variance schedule. We use a flexible variance schedule { ˆ σ 2 t } to extend the twisting functions in eq. (13) to

Ideally, ˆ σ 2 t should match Var p θ [ y | x t , M ] . In TDS, we choose ˆ σ 2 t = (¯ σ 2 t / ¯ α t ) · τ 2 ¯ σ 2 t / ¯ α t + τ 2 where τ 2 = 0 . 12 is the estimated sample variance of the training data (averaged over pixels). This choice is inspired from Song et al. [32, Appendix A.3]: if the data distribution is q ( x 0 ) = N (0 , τ 2 ) , and q ( x t | x 0 ) = N ( x t | √ ¯ α t x 0 , σ 2 t ) , then by Bayes rule Var q [ x 0 | x t ] = (¯ σ 2 t / ¯ α t ) · τ 2 (¯ σ 2 t / ¯ α t )+ τ 2 .

Changing ˆ σ 2 t has a similar effect as using a twist scale to exponentiate the twisting functions in Appendix D.2.1. In particular, for a Gaussian distribution, exponentiating the distribution corresponds to re-scaling the variance: N ( x | µ, ( σ/ √ γ ) 2 ) ∝ N ( x | µ, σ 2 ) γ .

Metrics. We consider 3 metrics. (1) We use the effective sample size (ESS) to compare the particle efficiency among different SMC samplers (namely TDS, TDS-IS, and SMC-Diff).

(2) In addition, we ground the performance of a sampler in the downstream task of classifying a partially observed image x 0 M = y : we use a classifier to predict the class ˆ z of an inpainted image x 0 , and compute the accuracy against the true class z ∗ of the unmasked image ( z ∗ is provided in MNIST dataset). This prediction is made by the same classifier used in Section 5.2.

Consider K weighted particles { x 0 k ; w 0 k } K k =1 drawn from a sampler conditioned on x 0 M = y and assume the weights are normalized. We define the Bayes accuracy (BA) as

where ˆ z ( y ) is viewed as an approximation to the Bayes optimal classifier ˆ z ∗ ( y ) given by

(In eq. (18) we assume the classifier p y | x 0 ( ·| x 0 ) is the optimal classifier on full images.)

We also consider classification accuracy (CA) defined as the following

BA and CA evaluate different aspects of a sampler. BA is focused on the optimal prediction among multiple particles, whereas CA is focused on the their weighted average prediction.

Comparison results of different methods. Figure I 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 CA is lower, its BA is comparable to TDS. This result is 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 a high BA.

Ablation study on twisting function variance schedule. We compare the following three options:

DPS: ˆ σ t := 2 ∥ ˆ x θ ( x t ) M -y ∥ 2 σ 2 t , adapted from the DPS method [6, Algorithm 1],

Π GDM: ˆ σ t := σ 2 t / √ ¯ α t , adapted from the Π GDMmethod [32, Algorithm 1],

TDS: ˆ σ 2 t = (¯ σ 2 t / ¯ α t ) · τ 2 (¯ σ 2 t / ¯ α t )+ τ 2 where τ 2 = 0 . 12 .

Figure J shows the classification accuracy of TDS, TDS-IS and Gradient Guidance with different twist scale schemes. We find that our choice has similar performance to that of Π GDM, and outperforms DPS in most cases. Exceptions are when M = 'quarter' and for large K , TDS with twist scale choice of Π GDMor DPS has higher CA, as is shown in the left panel in Figure Jb.

D.2.3 Class-conditional generation on ImageNet and CIFAR-10

ImageNet. We compare TDS (with # of particles K = 1 , 16 ) to Classifier Guidance (CG) using the same unconditional model from https://github.com/openai/guided-diffusion/tree/ main . CG uses a classifier trained on noisy inputs taken from this repository as well. For TDS, we

100

75

50

25

0

Time Step t

0

32

64

Effect. Sample Size

(a)

M

= 'half'. ESS trace.

1

2

4

8

16

32

64

# Particles K

60

80

100

Class. Accuracy (%)

1

2

4

8

16

32

64

# Particles K

60

80

100

Bayes Accuracy (%)

TDS

TDS-truncate

TDS-IS

Gradient

Guidance

SMC-Diff

Replacement

(b)

M

= 'half'. Class. accuracy (left), and Bayes accuracy (right) v.s.

K

.

100

75

50

25

0

Time Step t

0

32

64

Effect. Sample Size

1

2

4

8

16

32

64

# Particles K

20

50

80

Class. Accuracy (%)

1

2

4

8

16

32

64

# Particles K

20

50

80

Bayes Accuracy (%)

TDS

TDS-truncate

TDS-IS

Gradient

Guidance

SMC-Diff

Replacement

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 100 sampling steps. TDS uses a twist scale of 10, and CG uses a guidance scale of 10. Notably, given a fixed class, TDS ( K = 16 ) generates correlated samples in a single SMC run, and TDS ( K = 1 ) and CG generate 16 independent samples.

In Figure K, we observe that TDS can faithfully capture the class and have comparable image quality to 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 on 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 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 ( K = 16 ).

CIFAR-10. We ran TDS ( K = 16 ) with the twist scale = 1 and 100 sampling steps using diffusion model from https://github.com/openai/improved-diffusion and classifier from https: //github.com/VSehwag/minimal-diffusion/tree/main . TDS generates faithful and diverse

1

2

4

8 16 32 64

# Particles K

85.0

87.5

90.0

92.5

95.0

Class. Accuracy (%)

TDS

1

2

4

8 16 32 64

# Particles K

85

90

95

TDS-IS

1

2

4

8 16 32 64

# Particles K

86

88

90

Gradient Guidance

Twist Scale

DPS

GDM

Ours

(a) M = 'half'. Classification accuracy under different twist scale schemes.

1

2

4

8

16 32 64

# Particles K

50

60

70

Class. Accuracy (%)

TDS

1

2

4

8

16 32 64

# Particles K

50

55

60

65

TDS-IS

1

2

4

8

16 32 64

# Particles K

50

52

54

Gradient Guidance

Twist Scale

DPS

GDM

Ours

(a) TDS (P=16)

(b) TDS (P=1)

(c) Classifier Guidance

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 in Figure N. We see there is a sudden ESS drop in the final stage for MINST, sometimes for CIFAR10, but usually not on ImageNet. 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 . In practice, one can consider early truncating the TDS procedure to prevent such ESS drops and promote particle diversity, as 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, a diffusion generative model described by Yim et al. [39]. FrameDiff parameterizes protein N -residue protein backbones as a collection of rigid bodies defined by rotations and translations as x 0 = [( R 1 , z 1 ) , . . . , ( R N , z N )] ∈

J

(a) airplane

(b) deer

(c) horse

(d) ship

(e) truck

(f) failure mode

100

75

50

25

0

Time step t

0

4

8

12

16

Effect. sample size

MNIST

100

75

50

25

0

Time step t

CIFAR-10

100

75

50

25

0

Time step t

Imagenet

SE (3) N . SE (3) is the special Euclidean group in three dimensions (a Riemannian manifold). Each R n ∈ SO (3) (the special orthogonal group in three dimensions) is a rotation matrix and each z n ∈ R 3 in a translation. Together, R n and z n describe how one obtains the coordinates of the three backbone atoms C , C α , 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 y = x 0 M ∈ SE (3) | M | for some M ⊂ { 1 , . . . , N } . FrameDiff is a continuous time diffusion model and includes the number of steps as a hyperparmeter; we use 200 steps in all experiments. We refer the reader to [39] for details on the neural network architecture, and details of the forward and reverse diffusion processes.

Twisting functions for motif scaffolding. To define a twisting function, we use a tangent normal approximation to p θ ( y | x t ) 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 Euclidean cases, using the variance preserving extension described in Appendix A. For the rotational components, represented in the special orthogonal group SO (3) , the tangent normal may be computed as described in Appendix B, using the known exponential map on SO (3) . In particular, we use as the twisting function

where y T m , ˆ x θ ( x t ) T M m ∈ R 3 represent the translations associated with the m th residue of the motif and its prediction from x t , and y R m , ˆ x θ ( x t ) R M m ∈ SO (3) represent the analogous quantities for the rotational component. Next, 1 -¯ α t = Var q [ x t | x 0 ] and ¯ σ 2 t is the time of the Brownian motion associated with the forward process at step t [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 in the limit that ¯ σ t 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 to include the rotation with a uniform prior, and write

and with p ( R ) 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 a Monte Carlo approximation defined by subsampling a finite number of rotations, R = { R k } (with |R| = # Motif Rots. in total). Altogether we have

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 so 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 a 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 ProteinMPNN [7] for step (i) and AlphaFold [20] on a single sequence (no multiple sequence alignment) for (ii).

In this evaluation we use ProteinMPNN [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 a 'success' as a generated backbone for which at least one of the 8 sequences has backbone atom with both scRMSD < 1 Å on the motif and scRMSD < 2 Å on the full backbone.

We benchmarked TDS on 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 FrameDiff. FrameDiff 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, 116 becomes 125 and 62-83 becomes 75.

As discussed in the main text, in our handling of the motif-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 'contig' 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.

6EXZ_med

6EXZ_short

1YCR

5TRV_med

6E6R_med

3IXT

6E6R_long

6E6R_short

2KL8

6EXZ_long

5TRV_short

4ZYP

5IUS

1QJG

7MRX_60

5WN9

1PRW

1BCF

7MRX_85

5TPN

7MRX_128

5TRV_long

5YUI

4JHW

10 0

10 1

10 2

Success Rate (%)

TDS (K=1)

RFdiffusion

102

(%)

Success Rate(

TDStwist-scale=1

101

TDStwist-scale=2

TDStwist-scale=4

100

1PRW

5TRV_short

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 was 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 RFdiffusion [37]. RFdiffusion is a diffusion model that uses the same backbone structure representation as FrameDiff. However, unlike FrameDiff, RFdiffusion trained with a mixture of conditional examples, in which a segment of the backbone is presented as input as a desired motif, and unconditional training examples in which the only a noised backbone is 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 structure prediction pretraining and limited fine-tuning, which may allow it to achieve good performance without having fit a smooth score-approximation.

6EXZ_med

5IUS

Size

8

Size

8

:Sample

Sample :

6

6

4

4

Effective S

Effective

2

2

0

0

0

100

200

0

100

200

Step (T-t)

Step (T-t)

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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\nQUESTIONS:\nIt 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.\n\nLIMITATIONS:\nSee weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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\n\nWEAKNESSES:\n- 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?\n\nQUESTIONS:\nSee weaknesses.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\nBased 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.\n\nQUESTIONS:\n1. ``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?\n\nLIMITATIONS:\nI would expect the authors address the weakness and questions aforementioned.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nI 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-19"}