[论文解读] Score-Based Generative Modeling with Critically-Damped Langevin Diffusion

Tim Dockhorn, Arash Vahdat+1 more|arXiv (Cornell University)|Dec 13, 2021

Generative Adversarial Networks and Image Synthesis被引用 35

一句话总结

引入临界阻尼 Langevin 扩散（CLD）用于基于分数的生成模型，通过在数据中加入速度变量并学习在数据条件下的速度分数，展示了更好的综合质量和更快、更平滑的采样方法。

ABSTRACT

Score-based generative models (SGMs) have demonstrated remarkable synthesis quality. SGMs rely on a diffusion process that gradually perturbs the data towards a tractable distribution, while the generative model learns to denoise. The complexity of this denoising task is, apart from the data distribution itself, uniquely determined by the diffusion process. We argue that current SGMs employ overly simplistic diffusions, leading to unnecessarily complex denoising processes, which limit generative modeling performance. Based on connections to statistical mechanics, we propose a novel critically-damped Langevin diffusion (CLD) and show that CLD-based SGMs achieve superior performance. CLD can be interpreted as running a joint diffusion in an extended space, where the auxiliary variables can be considered "velocities" that are coupled to the data variables as in Hamiltonian dynamics. We derive a novel score matching objective for CLD and show that the model only needs to learn the score function of the conditional distribution of the velocity given data, an easier task than learning scores of the data directly. We also derive a new sampling scheme for efficient synthesis from CLD-based diffusion models. We find that CLD outperforms previous SGMs in synthesis quality for similar network architectures and sampling compute budgets. We show that our novel sampler for CLD significantly outperforms solvers such as Euler--Maruyama. Our framework provides new insights into score-based denoising diffusion models and can be readily used for high-resolution image synthesis. Project page and code: https://nv-tlabs.github.io/CLD-SGM.

研究动机与目标

• 通过重新审视扩散动态以降低去噪复杂度来推动改进 SGMs。
• 提出一个前向扩散过程，用速度变量来增强数据以加速达到平衡。
• 推导一个只学习给定数据条件下的速度分数 p_t(v_t|x_t) 的分数匹配目标，从而简化学习。
• 为从基于 CLD 的 SGMs 高效采样开发定制的 SDE 求解器（SSCS）。
• 在 CIFAR-10 上展示经验增益，在相似架构和预算下实现更好的 FID。

提出的方法

• 定义带有 x_t 和 v_t 的联合数据-速度扩散以及将它们耦合的哈密顿分量。
• 将扩散设为临界阻尼 Langevin（Gamma^2 = 4M）以优化收敛到平衡。
• 推导一个面向 nabla_{v_t} log p_t(v_t|x_t) 的分数匹配目标（式 (6)）。
• 采用混合分数匹配（HSM）通过对初始速度进行边际化并利用可处理的扰动核（式 (8)）实现高效训练。
• 使用混合分数参数化，利用近似正态分布的速度以简化分数目标（式 (9)）。
• 引入对称分割 CLD 采样器（SSCS），基于 Strang 分割的定制 SDE 求解器，以改进采样相对于欧拉-马里尤玛法。

实验结果

研究问题

• RQ1 通过 CLD 将数据与速度变量联合起来，是否能改善基于扩散的去噪与 SGM 的合成质量？
• RQ2 学习集中在给定数据条件下的速度分数 p_t(v_t|x_t) 是否能简化训练并提升模型平滑性？
• RQ3 基于 CLD 的采样器（SSCS）在样本质量和效率方面与标准 EM/欧拉方法相比如何？
• RQ4 在具有类似架构和计算预算的前提下，相对于先前的扩散模型，在 CIFAR-10 上的实证增益如何？
• RQ5 CLD 对超参数是否鲁棒，且能否扩展到更高分辨率数据？

主要发现

• CLD 基于 SGMs 相较于 VPSDE 基模型，在许多 t 上产生更平滑的神经分数网络，表明学习目标更易达成。
• 在 CIFAR-10 上，CLD 基于模型在相似架构和预算下实现了近似于最先进的合成质量，FID 最低可达到 2.23（SDE）和 2.25（概率流）。
• CLD 采样器（SSCS）在采样效率与质量上显著优于欧拉-马里尤玛。
• 该方法获得的 NLL 上限约为 3.31，与近期基于扩散的工作相当，尽管并非以最大化似然为主要目标。
• 辅助的速度增强和哈密顿分量加速了混合与扩散达到平衡的过程，使高质量、可扩展的图像合成成为可能。