[论文解读] Consistency Models

Yang Song, Prafulla Dhariwal+2 more|arXiv (Cornell University)|Mar 2, 2023

Generative Adversarial Networks and Image Synthesis被引用 24

一句话总结

引入将噪声直接映射到数据的一致性模型，用于快速的一步生成，支持可选多步细化和零-shot 编辑，并在 CIFAR-10、ImageNet-64 和 LSUN 上取得强劲结果。

ABSTRACT

Diffusion models have significantly advanced the fields of image, audio, and video generation, but they depend on an iterative sampling process that causes slow generation. To overcome this limitation, we propose consistency models, a new family of models that generate high quality samples by directly mapping noise to data. They support fast one-step generation by design, while still allowing multistep sampling to trade compute for sample quality. They also support zero-shot data editing, such as image inpainting, colorization, and super-resolution, without requiring explicit training on these tasks. Consistency models can be trained either by distilling pre-trained diffusion models, or as standalone generative models altogether. Through extensive experiments, we demonstrate that they outperform existing distillation techniques for diffusion models in one- and few-step sampling, achieving the new state-of-the-art FID of 3.55 on CIFAR-10 and 6.20 on ImageNet 64x64 for one-step generation. When trained in isolation, consistency models become a new family of generative models that can outperform existing one-step, non-adversarial generative models on standard benchmarks such as CIFAR-10, ImageNet 64x64 and LSUN 256x256.

研究动机与目标

• 在不牺牲样本质量的前提下，激发比扩散采样更快的速度。
• 提出从 PF ODE 路径到轨迹起点的自洽映射。
• 通过一次性或多步生成实现零-shot 数据编辑。
• 提供两种训练范式：来自预训练扩散模型的蒸馏与独立训练。

提出的方法

• 定义一致性函数 f，将任意 PF ODE 路径点映射到轨迹上的起点。
• 将 f 参数化为在 t = ε 的边界条件，使 f(x, ε) = x，并使用跳跃连接风格的形式。
• 通过一致性蒸馏进行训练，使用来自预训练扩散模型和经验 PF ODE 的配对（见 Eq. 3 与 Eq. 6）。
• 或者在没有扩散模型的情况下在孤立状态训练 f，使用基于对数据扰动的分数估计器的一致性训练损失。
• 通过在高斯样本 x_T 上评估 f 以获得 x_ε，从而实现一步采样；通过在额外时间点重复应用 f 可以实现多步采样，以用计算量换取质量。
• 展示零-shot 数据编辑能力，如修补、上色、去噪、插值和笔刷引导编辑。

实验结果

研究问题

• RQ1 从 PF ODE 路径到轨迹起点的自洽映射是否能够实现高质量的一步样本生成？
• RQ2 在标准数据集上，一致性蒸馏相对于渐进蒸馏和其他快速采样方法的对比如何？
• RQ3 一致性模型是否能在不进行任务特定训练的情况下实现零-shot 数据编辑？
• RQ4 哪些有效的训练策略（蒸馏 vs 孤立训练）和求解器选择能够实现最优性能？
• RQ5 就计算量和样本质量而言，单步生成与多步细化之间有哪些权衡？

主要发现

• 一致性模型在 CIFAR-10 与 ImageNet-64 上实现单步生成，其质量与蒸馏基线相比具有竞争力甚至更优。
• 作为蒸馏使用时，在多个数据集上的少步生成方面优于现有扩散蒸馏方法（例如 CIFAR-10 的一步/两步结果；ImageNet-64 的一步/两步结果）。
• 在孤立状态训练一致性模型可产生新型生成模型，在标准基准上可超越单步非对抗性基线和许多 GAN。
• 一致性模型支持广泛的零-shot 数据编辑任务，如修补、上色、超分辨、去噪和笔画引导编辑。
• 所提出的一致性蒸馏与一致性训练框架提供了稳定的收敛性和明确的理论基础，包括渐近误差界和在某些极限下的等价性结果。