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.
翻译:扩散模型在一系列条件生成任务中取得了成功,包括分子设计和文本到图像生成。然而,这些成就主要依赖于任务特定的条件训练或容易出错的启发式近似方法。理想情况下,条件生成方法应当能够为广泛的条件分布提供精确样本,且无需任务特定的训练。为此,我们提出了扭曲扩散采样器(Twisted Diffusion Sampler,TDS)。TDS是一种序列蒙特卡洛(SMC)算法,通过模拟一组加权粒子来逼近扩散模型的条件分布。其核心思想是利用扭曲技术——一种具有良好计算效率的SMC技术——来融入启发式近似,同时不损害渐近精确性。我们首先在仿真和条件图像生成任务中发现,TDS提供了一种计算统计权衡:使用更多粒子时可获得更精确的近似,但即使仅使用两个粒子,其经验性能也优于启发式方法。随后,我们将TDS扩展至黎曼扩散模型,并将其应用于蛋白质设计中的核心任务——基序支架设计。在基准测试案例中,TDS支持灵活的条件设定标准,并且通常优于现有最优方法。