用新方法让扩散模型一步生成,速度更快且更全面。
One-step Diffusion Models with $f$-Divergence Distribution Matching
- 提出基于f-散度的统一框架,优化单步生成分布匹配
- 使用Jensen-Shannon散度在ImageNet64和COCO上达当前最优
- 适合追求高速生成且注重多样性的应用开发者
扩散模型采样依赖缓慢的迭代过程,限制了其在交互式应用中的部署。为加速生成,现有方法通过变分得分蒸馏将多步模型压缩为单步生成器,但采用反向KL散度进行分布匹配,存在模式丢失问题。本文提出新的f-散度最小化框架f-distill,涵盖多种散度,可在模式覆盖与训练方差间权衡。推导出教师与学生分布间f-散度的梯度形式,其由得分差与密度比决定的加权函数调控,该函数在使用非强模式寻求散度时自然强调教师分布高密度区域。反向KL作为特例包含于本框架。实验证明,前向KL与Jensen-Shannon散度在图像生成任务中优于现有最佳方法,尤其使用Jensen-Shannon散度时,在ImageNet64上实现当前最优单步生成性能,并在MS-COCO零样本文生图任务中取得领先效果。
原文摘要 · Abstract (English)
Sampling from diffusion models involves a slow iterative process that hinders their practical deployment, especially for interactive applications. To accelerate generation speed, recent approaches distill a multi-step diffusion model into a single-step student generator via variational score distillation, which matches the distribution of samples generated by the student to the teacher's distribution. However, these approaches use the reverse Kullback-Leibler (KL) divergence for distribution matching which is known to be mode seeking. In this paper, we generalize the distribution matching approach using a novel $f$-divergence minimization framework, termed $f$-distill, that covers different divergences with different trade-offs in terms of mode coverage and training variance. We derive the gradient of the $f$-divergence between the teacher and student distributions and show that it is expressed as the product of their score differences and a weighting function determined by their density ratio. This weighting function naturally emphasizes samples with higher density in the teacher distribution, when using a less mode-seeking divergence. We observe that the popular variational score distillation approach using the reverse-KL divergence is a special case within our framework. Empirically, we demonstrate that alternative $f$-divergences, such as forward-KL and Jensen-Shannon divergences, outperform the current best variational score distillation methods across image generation tasks. In particular, when using Jensen-Shannon divergence, $f$-distill achieves current state-of-the-art one-step generation performance on ImageNet64 and zero-shot text-to-image generation on MS-COCO. Project page: https://research.nvidia.com/labs/genair/f-distill
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。