[论文解读] 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 Science被引用 13

一句话总结

本论文提出 Twisted Diffusion Sampler (TDS)，一种基于 SMC 的方法，用于对扩散模型进行渐近精确的条件采样，使用 twisting 来融入条件信息而无需任务特定训练，并在 inpainting、class-conditional MNIST 和 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.

研究动机与目标

• 在不需要任务特定训练或大量启发式方法的情况下，推动从无条件扩散模型进行条件采样的研究。
• 开发一个实用的顺序蒙特卡洛（SMC）方法，通过 twisting 来靶向 p_theta(x^0|y) 。
• 将 TDS 扩展到 inpainting 和 Riemannian diffusion 模型，以覆盖更广的条件任务。
• 在 MNIST inpainting 和 class conditioning，以及蛋白质基序-支架成像任务上，展示 TDS 相对于启发式方法的经验优势。

提出的方法

• 通过扩展扩散序列和 SMC，将条件采样从 p_theta(x^0|y) 表述为问题。
• 引入基于 p_theta(y|x^t) 的 twisting 函数，该函数由去噪估计 hat{x}_theta(x^t) 来近似。
• 定义 twisted proposals 和 weights，以在保持渐近精确性的同时逼近最优的 SMC 选择。
• 提供对 inpainting 和 Riemannian diffusion 模型的扩展，包括自由度条件等。
• 以非正式的方式证明在正则性条件下，粒子近似收敛到 p_theta(x^0|y)。

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.

实验结果

研究问题

• RQ1 TDS 是否能够在不需要任务特定条件训练的情况下，给出从 p_theta(x^0|y) 的渐近精确样本？
• RQ2 在准确性和粒子效率方面，TDS 相对于启发式方法和条件训练在不同任务中的表现如何？
• RQ3 如何将 TDS 扩展到 inpainting、自由度条件和 Riemannian diffusion 模型？

主要发现

• TDS 能够为具有条件信息的任务从扩散模型中得到渐近精确的条件样本。
• 在经验上，TDS 即使使用少量粒子，在 MNIST 类条件生成和 inpainting 设置中通常优于启发式方法。
• TDS 展示了对蛋白质基序-支架成像的灵活条件设定能力，并且在这些问题上可以超越最先进的条件模型。
• TDS 展现出有利的计算-统计权衡，随着粒子数量增加，精度提升（实验中呈 O(1/K) 速率）。
• 对 inpainting 和 Riemannian 流形的扩展保持渐近精确性并扩大了适用性。

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