Though diffusion models excel in image generation, their step-by-step denoising leads to slow generation speeds. Consistency training addresses this issue with single-step sampling but often produces lower-quality generations and requires high training costs. In this paper, we show that optimizing consistency training loss minimizes the Wasserstein distance between target and generated distributions. As timestep increases, the upper bound accumulates previous consistency training losses. Therefore, larger batch sizes are needed to reduce both current and accumulated losses. We propose Adversarial Consistency Training (ACT), which directly minimizes the Jensen-Shannon (JS) divergence between distributions at each timestep using a discriminator. Theoretically, ACT enhances generation quality, and convergence. By incorporating a discriminator into the consistency training framework, our method achieves improved FID scores on CIFAR10 and ImageNet 64$\times$64, retains zero-shot image inpainting capabilities, and uses less than $1/6$ of the original batch size and fewer than $1/2$ of the model parameters and training steps compared to the baseline method, this leads to a substantial reduction in resource consumption.
翻译:尽管扩散模型在图像生成中表现出色,但其逐步去噪过程导致生成速度缓慢。一致性训练通过单步采样解决了这一问题,但往往产生质量较低的生成结果,且训练成本较高。本文表明,优化一致性训练损失等价于最小化目标分布与生成分布之间的Wasserstein距离。随着时间步长的增加,上界会累积先前的一致性训练损失。因此,需要更大的批量大小来同时减少当前损失和累积损失。我们提出对抗一致性训练(ACT),该方法通过判别器直接最小化每个时间步长上分布之间的Jensen-Shannon散度。理论上,ACT提升了生成质量与收敛性。通过将判别器融入一致性训练框架,我们的方法在CIFAR10和ImageNet 64$\times$64上取得了更优的FID分数,保留了零样本图像修复能力,且与基线方法相比,所需批量大小不足原始的$1/6$,模型参数和训练步数均少于$1/2$,从而大幅降低了资源消耗。