[论文解读] Denoising MCMC for Accelerating Diffusion-Based Generative Models

Beomsu Kim, Jong Chul Ye|arXiv (Cornell University)|Sep 29, 2022

Generative Adversarial Networks and Image Synthesis被引用 4

一句话总结

本文提出去噪MCMC（DMCMC），一种用于基于得分的生成模型的正交加速方法，通过马尔可夫链蒙特卡洛（MCMC）从数据与扩散时间的乘积空间中采样，随后通过反向S/ODE积分实现去噪。通过从数据流形附近开始，DMCMC大幅降低了反向S/ODE积分的计算成本，在CIFAR10上实现3.86（≈10 NFE）的SOTA FID得分，在CelebA-HQ-256上实现6.99（≈160 NFE）的SOTA FID得分。

ABSTRACT

Diffusion models are powerful generative models that simulate the reverse of diffusion processes using score functions to synthesize data from noise. The sampling process of diffusion models can be interpreted as solving the reverse stochastic differential equation (SDE) or the ordinary differential equation (ODE) of the diffusion process, which often requires up to thousands of discretization steps to generate a single image. This has sparked a great interest in developing efficient integration techniques for reverse-S/ODEs. Here, we propose an orthogonal approach to accelerating score-based sampling: Denoising MCMC (DMCMC). DMCMC first uses MCMC to produce samples in the product space of data and variance (or diffusion time). Then, a reverse-S/ODE integrator is used to denoise the MCMC samples. Since MCMC traverses close to the data manifold, the computation cost of producing a clean sample for DMCMC is much less than that of producing a clean sample from noise. To verify the proposed concept, we show that Denoising Langevin Gibbs (DLG), an instance of DMCMC, successfully accelerates all six reverse-S/ODE integrators considered in this work on the tasks of CIFAR10 and CelebA-HQ-256 image generation. Notably, combined with integrators of Karras et al. (2022) and pre-trained score models of Song et al. (2021b), DLG achieves SOTA results. In the limited number of score function evaluation (NFE) settings on CIFAR10, we have $3.86$ FID with $\approx 10$ NFE and $2.63$ FID with $\approx 20$ NFE. On CelebA-HQ-256, we have $6.99$ FID with $\approx 160$ NFE, which beats the current best record of Kim et al. (2022) among score-based models, $7.16$ FID with $4000$ NFE. Code: https://github.com/1202kbs/DMCMC

研究动机与目标

• 通过降低反向S/ODE积分的计算成本，加速基于得分的扩散模型采样。
• 通过在多模态、高维数据分布中实现跨不连通模态的遍历，解决传统MCMC方法效率低下的问题。
• 开发一种即插即用的框架，将MCMC与反向S/ODE积分器结合，实现更快、更高质量的采样生成。
• 在CIFAR10和CelebA-HQ-256等标准基准上，以最少的得分函数评估次数（NFE）实现SOTA FID得分。

提出的方法

• 提出去噪MCMC（DMCMC），通过MCMC从数据与扩散时间（方差）的乘积空间中采样，以探索靠近数据流形的高密度区域。
• 使用反向S/ODE积分器对MCMC采样结果进行去噪，起始时间接近零（低噪声），与从最大噪声起始相比，显著缩短了积分区间。
• 提出去噪拉姆齐吉布斯（DLG），作为DMCMC的实用实例，交替使用拉姆齐动力学进行数据更新和吉布斯采样进行噪声水平预测。
• 采用预训练的噪声条件得分网络和轻量级噪声水平分类器，实现在几乎不增加额外计算成本的前提下高效且可扩展的采样。
• 将扩散过程视为时变得分函数，允许从任意接近零的初始时间t_n开始反向S/ODE积分。
• 优化超参数，如步长η、总NFE以及分配给去噪的NFE占比（n_den/n），以平衡采样质量与多样性。

实验结果

研究问题

• RQ1 在数据与扩散时间的乘积空间中进行MCMC采样，是否能实现基于得分的生成模型中反向S/ODE积分的更快、更高效？
• RQ2 通过MCMC实现从低噪声状态开始反向S/ODE积分，是否能显著减少所需的得分函数评估次数（NFE），同时保持或提升采样质量？
• RQ3 一种结合MCMC与反向S/ODE积分器的即插即用框架，是否能在多样化的生成建模基准上实现SOTA性能？
• RQ4 超参数如拉姆齐步长（η）、总NFE以及去噪占比（n_den/n）如何影响DMCMC框架中采样质量与多样性的权衡？

主要发现

• 在CIFAR10上，DLG仅使用约10次NFE即达到SOTA FID 3.86，优于此前在相同NFE预算下实现4.17 FID的方法。
• 在CelebA-HQ-256上，DLG以约160次NFE实现SOTA FID 6.99，超越此前使用不同得分模型在4,000次NFE下达到的7.16 FID最佳结果。
• 该方法通过从接近零的噪声水平起始，显著缩短了平均反向S/ODE积分区间的长度，实现更快收敛且额外计算成本可忽略。
• 消融实验表明，当去噪占比（n_den/n）处于中间值（非接近0或1）时性能最优，且采用与维度无关的位移参数（κ = η/√d）可实现跨数据集的有效超参数迁移。
• 噪声水平分类器的计算成本可忽略不计，使得加速效果在额外FLOPs方面近乎免费。
• 该框架具有通用性与即插即用特性：可适配任意MCMC方法、VE过程及反向S/ODE积分器，展现出广泛的适用性与可扩展性。