[Paper Review] Practical and Asymptotically Exact Conditional Sampling in Diffusion Models

Luhuan Wu, Brian L. Trippe+3 more|UvA-DARE (University of Amsterdam)|Jun 30, 2023

Machine Learning in Materials Science13 citations

TL;DR

The paper introduces Twisted Diffusion Sampler (TDS), an SMC-based method for asymptotically exact conditional sampling from diffusion models, using twisting to incorporate conditioning without task-specific training, and demonstrates empirical gains across inpainting, class-conditional MNIST, and protein motif-scaffolding.

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

Motivation & Objective

• Motivate conditional sampling from unconditional diffusion models without task-specific training or heavy heuristics.
• Develop a practical sequential Monte Carlo (SMC) approach using twisting to target p_theta(x^0|y).
• Extend TDS to inpainting and Riemannian diffusion models for broader conditioning tasks.
• Demonstrate empirical advantages of TDS over heuristic methods on MNIST inpainting and class conditioning, and on protein motif-scaffolding.

Proposed method

• Formulate conditional sampling from p_theta(x^0|y) via an extended diffusion sequence and SMC.
• Introduce twisting functions based on p_theta(y|x^t) approximated by the denoising estimate hat{x}_theta(x^t).
• Define twisted proposals and weights to approximate optimal SMC choices while maintaining asymptotic exactness.
• Provide extensions to inpainting and to Riemannian diffusion models, including degrees-of-freedom conditioning.
• Prove (informally) convergence of the particle approximation to p_theta(x^0|y) under regularity conditions.

Figure 1: Errors of conditional mean estimations with 2 SEM error bars averaged over 25 replicates. TDS applies to all three tasks and provides increasing accuracy with more particles.

Experimental results

Research questions

• RQ1 Can TDS provide asymptotically exact samples from p_theta(x^0|y) without task-specific conditioning training?
• RQ2 How does TDS compare to heuristic methods and conditional training in terms of accuracy and particle efficiency across tasks?
• RQ3 How can TDS be extended to inpainting, degrees-of-freedom conditioning, and Riemannian diffusion models?

Key findings

• TDS yields asymptotically exact conditional samples from diffusion models for tasks with conditioning information.
• Empirically, TDS with few particles often outperforms heuristic approaches in MNIST class-conditional generation and inpainting settings.
• TDS demonstrates flexible conditioning for protein motif-scaffolding and can surpass state-of-the-art conditional models on these problems.
• TDS shows a favorable computational-statistical trade-off, with accuracy improving as the number of particles increases (O(1/K) rate in trials).
• Extensions to inpainting and Riemannian manifolds preserve asymptotic exactness and widen applicability.

(a) Approximate conditional samples of each method with 64 particles given class $y=7$ .