[论文解读] Elucidating the Design Space of Diffusion-Based Generative Models

Tero Karras, Miika Aittala+2 more|arXiv (Cornell University)|Jun 1, 2022

Generative Adversarial Networks and Image Synthesis被引用 307

一句话总结

本文提出了扩散模型的模块化设计空间，提出采样与训练改进，并在 CIFAR-10 和 ImageNet-64 上达到新的最先进 FID，同时实现更快的采样。它还展示了对先前工作中预训练模型的改进。

ABSTRACT

We argue that the theory and practice of diffusion-based generative models are currently unnecessarily convoluted and seek to remedy the situation by presenting a design space that clearly separates the concrete design choices. This lets us identify several changes to both the sampling and training processes, as well as preconditioning of the score networks. Together, our improvements yield new state-of-the-art FID of 1.79 for CIFAR-10 in a class-conditional setting and 1.97 in an unconditional setting, with much faster sampling (35 network evaluations per image) than prior designs. To further demonstrate their modular nature, we show that our design changes dramatically improve both the efficiency and quality obtainable with pre-trained score networks from previous work, including improving the FID of a previously trained ImageNet-64 model from 2.07 to near-SOTA 1.55, and after re-training with our proposed improvements to a new SOTA of 1.36.

研究动机与目标

• 通过将具体组成部分与选择分离，阐明基于扩散的生成模型的实际设计空间。
• 通过更高阶的求解器、优化的时间表和随机性分析提升采样效率与输出质量。
• 开发针对分数网络的原理性预处理以及端到端训练实践，以提升鲁棒性和性能。
• 在现有模型上展示模块化改进并在 CIFAR-10 和 ImageNet-64 上取得提升。

提出的方法

• 在一个通用的 ODE/SDE 框架中表达扩散模型并识别独立的设计组件。
• 为确定性采样应用一个二阶海因求解器，结合优化的时间表以及对曲率敏感的 σ(t) 和 s(t)。
• 引入带受控噪声注入（churn）的随机采样器，以分析采样中随机性的作用。
• 提出带 σ 相关跳跃和缩放（cin, cout, cskip, cnoise）以稳定训练的预处理分数网络 Dθ。
• 在训练期间优化损失权重和噪声分布（λ(σ), ptrain(σ)），并采用非泄漏数据增强以提升泛化能力。
• 展示训练改进带来新的最先进 FID，并展示更快的采样（NFE 减少）。

实验结果

研究问题

• RQ1 扩散模型中哪些独立设计选择会影响性能和采样速度？
• RQ2 当与训练解耦时，采样策略（确定性与随机性）如何在不同模型族中影响图像质量？
• RQ3 系统化的预处理和训练损失是否能在不同分辨率和数据集上提升鲁棒性和最终 FID？
• RQ4 在采样过程中，调度（σ(t)、s(t)）对 ODE 轨迹和去噪器引导的影响是什么？
• RQ5 模块化改进在多大程度上可以迁移到先前工作的预训练扩散模型？

主要发现

• 在 CIFAR-10（有条件）达到 1.79 的最先进 FID，未条件为 1.97，且采样更快（每张图像 35 次 Dθ 评估）。
• 在 ImageNet-64 上接近最先进的 1.55，使用先前模型，重新训练并应用提出的改进后为 1.36。
• 通过采用二阶海恩求解器、优化的 σ(t) 和 s(t)、以及更精细的时间步进调度，显示出显著的采样加速。
• 证明更好的训练预处理和损失加权（λ(σ), ptrain(σ)）以及非泄漏增强在各分辨率上都带来显著改进。
• 发现随机采样的收益依赖于模型设置，并且通过改进训练（确定性采样在某些情况下可超越随机采样）可以减少甚至消除。
• 通过将改进应用于多种模型族（VP/VE、DDPM/DDIM、ImageNet-64）并获得一致的增益，验证了改进的模块化。