[论文解读] Multistep Consistency Models

Jonathan Heek, Emiel Hoogeboom+1 more|arXiv (Cornell University)|Mar 11, 2024

Complex Systems and Decision MakingDecision Sciences被引用 3

一句话总结

本文提出多步一致性模型（Multistep Consistency Models），统一了一致性模型与TRACT，通过在单步一致性与完整扩散采样之间插值，实现了性能突破。通过在2–8步内共享参数地训练分段一致性模型，该方法在8步内实现了SOTA级FID分数（ImageNet64上为1.4，ImageNet128上为2.1），在保持高速采样速度的同时，达到了标准扩散模型的生成质量。

ABSTRACT

Diffusion models are relatively easy to train but require many steps to generate samples. Consistency models are far more difficult to train, but generate samples in a single step. In this paper we propose Multistep Consistency Models: A unification between Consistency Models (Song et al., 2023) and TRACT (Berthelot et al., 2023) that can interpolate between a consistency model and a diffusion model: a trade-off between sampling speed and sampling quality. Specifically, a 1-step consistency model is a conventional consistency model whereas a $\infty$-step consistency model is a diffusion model. Multistep Consistency Models work really well in practice. By increasing the sample budget from a single step to 2-8 steps, we can train models more easily that generate higher quality samples, while retaining much of the sampling speed benefits. Notable results are 1.4 FID on Imagenet 64 in 8 step and 2.1 FID on Imagenet128 in 8 steps with consistency distillation, using simple losses without adversarial training. We also show that our method scales to a text-to-image diffusion model, generating samples that are close to the quality of the original model.

研究动机与目标

• 弥合标准扩散模型与低步数一致性模型之间的性能差距。
• 通过多步推理实现采样速度与生成质量之间的平滑权衡。
• 通过泛化TRACT与一致性蒸馏，提升一致性建模的训练稳定性和样本质量。
• 开发一种确定性采样器（aDDIM），以校正积分误差，减少低步数生成中的模糊现象。
• 在ImageNet128等复杂基准测试以及文生图生成任务中展示优异性能。

提出的方法

• 该方法将扩散过程划分为预定义的段，为每一段独立训练一致性模型，同时在所有段之间共享参数。
• 将一致性训练（CT）与一致性蒸馏（CD）泛化至多步设置，实现更平滑且更易学习的去噪轨迹。
• 该方法引入步长调度退火与一致性建模中的同步丢弃，以提升训练稳定性。
• 提出一种新型确定性采样器——修正DDIM（aDDIM），通过放大噪声预测以校正积分误差，降低低步数采样中的模糊现象。
• 该框架支持从1步一致性模型到完整扩散模型（无穷步）的插值，其中8步版本的性能与标准扩散模型相当。
• 在蒸馏过程中，该方法使用aDDIM从预训练的扩散模型中训练出高质量的16步一致性模型。

实验结果

研究问题

• RQ1 一致性模型能否被扩展至多步采样，以在保持速度优势的同时提升生成质量？
• RQ2 采用共享参数的多步一致性训练是否优于单步一致性或标准扩散模型？
• RQ3 像aDDIM这样的确定性采样器能否在ImageNet128等复杂数据集上，以极少步数实现具有竞争力的FID分数？
• RQ4 步长调度退火对低步数多步一致性模型性能有何影响？
• RQ5 多步一致性模型能否在文生图生成任务中，以极少采样步数达到标准扩散模型的性能？

主要发现

• 所提出的多步一致性模型在ImageNet64上实现了SOTA级FID 1.4，在ImageNet128上实现了FID 2.1，且仅使用8次采样步数。
• 在4次采样步数下，模型在ImageNet64上达到1.6的FID，在ImageNet128上达到2.3的FID，展现出极低推理成本下的强大性能。
• 该方法仅用8步即达到标准扩散模型性能（如100步DDIM），有效弥合了质量差距。
• 步长调度退火显著提升了低步数模型的性能，且随着步数增加，收益递减。
• aDDIM采样器使ImageNet128上实现具有竞争力的FID分数成为可能，该基准此前缺乏可靠的确定性低步数基线。
• 在文生图生成任务中，16步蒸馏模型在使用相同随机种子时，生成的样本几乎与原始100步DDIM模型无法区分。