PaperlayerView exact revision on arXiv

Source-linked AI summary

Practical and Asymptotically Exact Conditional Sampling in Diffusion Models

Luhuan Wu, Brian L. Trippe, Christian A. Naesseth, David M. Blei, John P. Cunningham

arXiv:2306.17775v2stat.MLcs.LGq-bio.BM

TL;DR

Conditional diffusion generation commonly depends on task-specific training or heuristic approximations whose accuracy may be limited. The paper introduces TDS, an SMC sampler that uses twisting and weighting to target diffusion-model conditionals asymptotically exactly. Across image and protein-design tasks, TDS improves over heuristics and often outperforms the state of the art, while requiring additional computation and problem-dependent twisting choices.

  • Problem

    Existing conditional diffusion methods require task-specific paired training data or heuristic approximations, motivating exact conditional sampling without retraining for each task.

  • Method

    TDS uses sequential Monte Carlo with twisted proposals and corrective weighting to sample conditional distributions from unconditional diffusion models.

  • Results

    TDS improves empirically over heuristic approaches on image tasks and achieves higher success rates than the state-of-the-art conditionally trained model on short-scaffold protein motif-scaffolding.

  • Takeaways & Limitations

    TDS supports flexible conditioning across image, inpainting, and protein motif-scaffolding problems while retaining asymptotic exactness.

  • Takeaways & Limitations

    TDS requires additional computation for multiple particles, with the optimal particle count and twisting quality depending on the problem and conditioning information.

Abstract

from arXiv · show

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 test cases, TDS allows flexible conditioning criteria and often outperforms the state of the art.

1 Introduction

Conditional sampling enables generation from distributions conditioned on desired information, but existing diffusion-model approaches require task-specific training or rely on heuristics with uncertain accuracy. TDS uses SMC and twisting to target diffusion-model conditionals asymptotically exactly across multiple tasks.

  • Conditional sampling augments a generative distribution pθ(x) with information y to define and sample from pθ(x | y).
  • Conditional training requires paired data and a task-specific model for each new conditioning task.Image inpainting and class-conditional generation can share the same unconditional image distribution but require separate conditional training sets and models.
  • Heuristic methods flexibly reuse unconditional diffusion models, but their approximation to exact conditionals is unclear and can produce inconsistent inpainting outputs.The paper identifies this concern as especially important for domains requiring accurate conditionals.
  • TDS is a sequential Monte Carlo algorithm that simulates weighted particles to provide asymptotically exact conditional sampling from diffusion models.Its proposals and weighting mechanisms operate over the sequential structure of diffusion models.
  • TDS applies across conditional generation problems, extends to Riemannian diffusion models, improves empirically over heuristics, and outperforms a conditionally trained model on short-scaffold protein motif-scaffolding.

2 Background: Diffusion models and sequential Monte Carlo

Diffusion models generate data by reversing a noise-adding process, while SMC approximates sequential target distributions with weighted particles. With suitable proposals and weights, the particle approximation becomes arbitrarily accurate as the particle count grows.

  • Diffusion models: A diffusion model generates x0 by iteratively refining noisy states xt from an initial noise state xT.The resulting distribution of x0 is the marginal of a length-T Markov chain.
  • Diffusion models: The diffusion model’s reverse transitions approximate the forward process’s reverse conditionals, yielding pθ(x0) ≈ q(x0) when the approximations are accurate.
  • Diffusion models: A score network can be represented through a denoising network trained to approximate E_q[x0 | xt].Tweedie’s formula relates this denoising estimate to the score used in reverse diffusion transitions.
  • Sequential Monte Carlo: SMC generates K particles across T steps using proposals, weighting functions, and resampling to approximate a sequence of target distributions.Intermediate targets can represent progressively richer posteriors, with the final target incorporating all observations.
  • Sequential Monte Carlo: As K approaches infinity, weighted particle measures converge to the target under the SMC construction, giving arbitrarily low approximation error when the final target matches the desired distribution.

3 Twisted Diffusion Sampler: SMC sampling for diffusion model conditionals

TDS formulates conditional diffusion sampling as SMC and combines tractable twisting with corrective weights. The resulting sampler targets pθ(x0 | y), converges asymptotically under regularity conditions, and supports inpainting and Riemannian diffusion extensions.

  • TDS embeds the conditioning likelihood into a joint diffusion model and targets the conditional distribution pθ(x0 | y).
  • The diffusion model’s Markov structure factorizes the extended conditional distribution into a form suitable for SMC.
  • Naive importance sampling can require exponentially many particles when pθ(x0 | y) is too dissimilar from pθ(x0).
  • Optimal twisting would permit exact conditional samples with one particle, but its required twisting functions are not readily sampleable.
  • TDS approximates optimal twisting by evaluating the conditioning likelihood at the denoising estimate x̂θ(xt), making the twisting function tractable.This uses one denoising-network call rather than an intractable integral over many calls.
  • Twisted proposals and weighting functions correct proposal errors while preserving convergence to the desired conditional target.At t = 0, the intermediate target reduces to pθ(x0:T | y).
  • Under regularity conditions, the particle approximation converges setwise to pθ(x0 | y) as K approaches infinity.
  • Alternative twisting functions extend TDS to inpainting and flexible conditioning beyond differentiable, strictly positive likelihoods.

4 Related work

Prior diffusion conditional-generation methods use task-specific training, heuristic guidance, or restricted SMC formulations. TDS retains heuristic guidance in proposals while using additional particles and weights to reduce approximation error and broaden applicability.

  • Prior conditional-generation approaches either require task-specific training or use unqualified approximations that can perform poorly in practice.
  • Gradient guidance: Gradient guidance approximations recover TDS with one particle for specified likelihood-based conditioning.
  • Gradient guidance: TDS uses gradient guidance in proposals while adding particles to provide a mechanism for eliminating approximation error.
  • Replacement method: The replacement method is limited to inpainting and provides no accuracy guarantees, with possible inconsistency between observed and inpainted regions.
  • SMC samplers: SMC-Diff provides asymptotically accurate inpainting samples but is limited to inpainting and assumes the learned diffusion model exactly matches the forward noising process.
  • SMC samplers: MCGdiff avoids denoising-network gradients but applies only to linear inverse problems, whereas TDS supports a broader conditioning setup.

5 Simulation study and conditional image generation

The study evaluates TDS across synthetic conditional distributions and image-generation tasks, finding accuracy improvements with more particles and favorable empirical comparisons with heuristic methods. TDS also produces qualitatively good, class-consistent high-dimensional image samples, though its inference cost grows with particle count.

  • Simulation study: TDS, TDS-IS, IS, and Gradient Guidance are evaluated across smooth-likelihood, inpainting, and degrees-of-freedom conditioning cases.TDS applies to all three settings, whereas several comparison methods apply only to subsets.
  • Computational considerations: TDS inference costs increase linearly with particle count, while conditional training can provide faster amortized inference.TDS also incurs a constant-factor cost for backpropagation through the denoising network.
  • Simulation study: O(1/K) mean-square estimation error decreases with particle count for TDS and applicable comparison samplers in the tested settings.The convergence appears as slopes of −1 on log-log plots, consistent with standard SMC theory.
  • Simulation study: Gradient Guidance has significant estimation error and does not improve with many particles in the three synthetic conditioning settings.It is the only other method besides TDS that applies to all three settings.
  • Conditional image generation: On MNIST, TDS produces authentic class-7 digits, while Gradient Guidance samples show artifacts and often do not resemble the target digit.For fixed particle count, TDS and TDS-IS outperform Gradient Guidance and IS; K = 64 yields nearly perfect accuracy for SMC samplers.
  • Conditional image generation: On ImageNet at 256×256×3 dimensions, TDS with K = 16 produces qualitatively good samples that capture the requested class label.The experiment uses a pretrained diffusion model and classifier.

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

The protein-design case study applies TDS to motif scaffolding with flexible motif placement and pose handling. On a 24-problem benchmark, TDS and RFdiffusion perform comparably overall, with TDS stronger on smaller scaffolds.

  • Motif-scaffolding problem: Motif scaffolding seeks protein structures whose overall scaffold stabilizes a functional motif.The task is framed as conditional generation from a distribution of designable protein structures.
  • Motif-scaffolding problem: Existing motif-scaffolding methods require specifying motif locations in the primary sequence, which can require expert knowledge and trial and error.The general motif-scaffolding problem remains open despite experimentally validated structures for certain motifs.
  • TDS extension: TDS is applied to FrameDiff, a Riemannian diffusion model representing protein backbones as N rigid bodies in SE(3)^N.Each manifold element encodes a residue’s rotation matrix and translation.
  • TDS extension: TDS averages over possible motif masks to avoid requiring a pre-specified motif placement within the scaffold.The method treats motif indices as a set of possible masks and applies Equation (14) across them.
  • Ablation study: Accounting for motif locations and rotations is necessary for non-zero success in the reported ablation, with success increasing as these degrees of freedom expand.The single-problem ablation used 1,000 motif locations and 100 rotations in one setting.
  • Benchmark evaluation: On 24 benchmark problems, TDS and RFdiffusion have comparable performance, with each exceeding the other on 11 cases and both reaching 0% on two.TDS uses K = 8, twist scale = 2, 100 rotations, and 1,000 motif locations.
  • Benchmark evaluation: TDS provides higher success rates on smaller scaffolds, while the performance division is primarily explained by total scaffold length.The authors suspect this shift reflects properties of the underlying diffusion models.

7 Discussion

The discussion presents TDS as an asymptotically exact conditional sampler that combines heuristic twisting with particle-based correction. It highlights broad applicability and protein-design gains, while noting compute, approximation, and twisting-function limitations.

  • Summary: TDS is presented as a practical and asymptotically exact conditional sampler for diffusion models.The paper reports effectiveness and flexibility on image generation and inpainting, plus improved performance on short-scaffold motif scaffolding.
  • Limitations: TDS requires additional computation for multiple particles, and its optimal particle count depends on the problem.Efficiency also depends on how closely twisting functions approximate exact conditionals.
  • Extensions: TDS can use alternative twisting functions for inpainting and more flexible conditioning specifications.The algorithm remains applicable with new twisting-function definitions, including adaptations to variance-preserving diffusion models.
  • Caveats: The final twisting choice used for some inpainting tasks introduces error relative to the exact conditional target.Alternative final proposals and weights are described to maintain asymptotic exactness.

A.5 Asymptotic accuracy of TDS – additional details and full theorem statement

TDS is asymptotically accurate under sufficient regularity conditions on its twisting functions, proposals, and weighting functions. Its particle approximation converges to the diffusion model’s conditional distribution as the number of particles increases.

  • Sufficient conditions: Theorem 2 gives sufficient conditions under which TDS provides arbitrarily accurate estimates as the particle count increases.The analysis characterizes conditions on the model and twisting functions, including boundedness and proposal-variance requirements.
  • Convergence guarantee: Under regularity conditions, the weighted particle measure converges setwise to pθ(x0 | y) as K approaches infinity.The convergence holds with probability one for every measurable set through the SMC convergence result.
  • Practical assumptions: The assumptions are readily satisfied in common applications, including by setting the final twisting function equal to the likelihood.The paper gives smoothness, positivity, compactness, bounded-gradient, and proposal-variance conditions as practical sufficient conditions.
  • Proof strategy: The proof first identifies the final target’s x0 marginal as pθ(x0 | y), then establishes particle convergence by applying an SMC theorem.The target identification follows by substituting the weights, canceling proposal and twisting terms, and applying Bayes’ rule.
  • Proof strategy: Bounded weighting functions follow from bounded twisting ratios, bounded gradients of twisting functions, and proposal variances larger than model variances.The paper notes that slightly inflating proposal variances may satisfy the variance condition without markedly changing sampler behavior.

B Riemannian Twisted Diffusion Sampler

The Riemannian extension adapts TDS to manifold-valued diffusion models by replacing Euclidean Gaussian components with tangent-normal constructions and manifold-valued denoising estimates. The resulting design preserves the twisted-proposal and weighting framework while requiring geometry-aware score and proposal operations.

  • Tangent-normal construction: A tangent-normal sample is generated in the tangent space and mapped back to the manifold through the exponential map.The tangent-space variable uses a Gaussian construction, while exp_x projects it onto the manifold.
  • Tangent-normal construction: The tangent-normal density includes the inverse exponential map and the Jacobian determinant of the exponential map.These terms account for the change of variables between tangent-space coordinates and manifold points.
  • Riemannian TDS: Riemannian TDS uses manifold-valued denoising estimates and tangent-normal analogues of twisted proposals and weights.The extension applies when a tractable score approximation can be formed, including certain Lie-group Brownian-motion settings.
  • Manifold conditioning: For inpainting on manifolds, twisting functions are constructed as in the Euclidean case with the normal density replaced by a tangent-normal density.This supports conditioning observed manifold coordinates while sampling the unobserved dimensions.
  • Algorithmic changes: The Riemannian algorithm changes the conditional score computation and replaces the Euclidean proposal with a tangent-normal proposal.The conditional score is a tangent-space-valued Riemannian gradient, and proposal-target Jacobian terms cancel in the weights.

D.2.1 Class-conditional generation on MNIST

MNIST experiments examine how particle count and twist scale affect classifier- and human-rated conditional generation quality. Moderate twist scales and more particles generally help, while excessive twisting can distort samples and expose evaluator disagreement.

  • Sample diversity: TDS and TDS-truncate produce greater diversity than importance sampling and TDS-IS, whose samples often become identical through importance-weight collapse.Gradient Guidance samples show noticeable artifacts, whereas the other methods generally produce authentic and correct digits.
  • Classifier evaluation: Larger twist scales generally improve neural-network classification accuracy, with stronger gains for TDS and TDS-IS at smaller K.The evaluation averages results over 1,000 random runs with error bands indicating 2 standard errors.
  • Human evaluation: Moderate twist scales γ = 2, 3 generally improve human-rated accuracy, whereas γ = 10 decreases it.For TDS with K = 64, any γ > 1 performs worse than γ = 1 under human evaluation.
  • Particle count: Increasing K from 1 to 2 or 64 improves TDS human-rated accuracy at a fixed twist scale.Gradient Guidance is treated as the K = 1 special case of TDS.
  • Sample quality: Excessive twist scale can distort digit shapes and add artifacts despite retaining features recognizable to a neural-network classifier.The paper suggests these samples may fall outside the data manifold, explaining lower human-rated but higher classifier-rated accuracy.
  • Effective sample size: With K = 64, TDS’s ESS generally rises toward 64 before dropping in the final steps, while TDS-IS and IS deteriorate to around 1 and 6.The TDS-truncate variant is proposed to avoid the final ESS drop.

D.2.2 Image inpainting on MNIST

The MNIST inpainting study evaluates particle samplers using ESS and downstream classification measures across different observed-image regions. TDS outperforms SMC-Diff on classification metrics despite lower ESS, while its twist-variance schedule generally matches ΠGDM and beats DPS.

  • Experimental setup: The study evaluates inpainting with either the left half or upper-left quarter of an image observed.These settings define the observed dimensions used for the partially observed MNIST images.
  • Metrics: The evaluation uses ESS, Bayes accuracy, and classification accuracy to assess particle efficiency and downstream conditional-image quality.Bayes accuracy focuses on the optimal prediction across particles, while classification accuracy focuses on the weighted average prediction.
  • Sampler comparison: TDS outperforms SMC-Diff in both classification accuracy metrics, although SMC-Diff has consistently higher ESS.The comparison includes classification accuracy and Bayes accuracy for partially observed images.
  • Metric interpretation: Gradient Guidance has lower classification accuracy but Bayes accuracy comparable to TDS because a few good particles can yield accurate optimal prediction.Bayes accuracy emphasizes the best prediction among particles, whereas classification accuracy emphasizes their weighted average prediction.
  • Variance schedule: The chosen twisting-function variance schedule performs similarly to ΠGDM and outperforms DPS in most cases.For quarter-image observations and large K, ΠGDM or DPS can yield higher classification accuracy than TDS’s default choice.

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

On ImageNet, TDS produces class-faithful samples with quality comparable to Classifier Guidance, while particle correlation reduces diversity. On CIFAR-10, TDS generally generates faithful and diverse images but can occasionally fail for the truck class.

  • The ImageNet comparison evaluates TDS with K = 1 and K = 16 against Classifier Guidance using 100 sampling steps.
  • TDS can faithfully capture ImageNet classes with image quality comparable to Classifier Guidance, but its samples are less diverse.
  • TDS and Classifier Guidance provide similar ImageNet classification accuracy, while Classifier Guidance has better FID and inception score.
  • TDS generates faithful and diverse CIFAR-10 images for different class labels, with an occasional incorrect sample for the class ‘truck’.

E.1 Additional Results

Additional motif-scaffolding results show that twist-scale effects vary across problems, effective sample size can trigger frequent resampling, and TDS application to RFdiffusion was unstable.

  • Twist-scale performance does not improve monotonically across all motif-scaffolding problems.
  • For 6EXZ-med, resampling was triggered 38 times, including 14 times in the final 25 steps; for 5UIS, it was triggered 63 times, including 13 times late.
  • Applying TDS to RFdiffusion produced numerically unstable gradients and frequently diverging trajectories even with one particle.
  • The motif-scaffolding ESS traces shown for 6EXZ-med and 5IUS use K = 8 and a resampling threshold of 0.5K.
Loading 2306.17775v2…