Progressive Distillation for Fast Sampling of Diffusion Models
Tim Salimans、Jonathan Ho Google Research, Brain team
Abstract
Diffusion models have recently shown great promise for generative modeling, outperforming GANs on perceptual quality and autoregressive models at density estimation. A remaining downside is their slow sampling time: generating high quality samples takes many hundreds or thousands of model evaluations. Here we make two contributions to help eliminate this downside: First, we present new parameterizations of diffusion models that provide increased stability when using few sampling steps. Second, we present a method to distill a trained deterministic diffusion sampler, using many steps, into a new diffusion model that takes half as many sampling steps. We then keep progressively applying this distillation procedure to our model, halving the number of required sampling steps each time. On standard image generation benchmarks like CIFAR-10, ImageNet, and LSUN, we start out with state-of-the-art samplers taking as many as 8192 steps, and are able to distill down to models taking as few as 4 steps without losing much perceptual quality; achieving, for example, a FID of 3.0 on CIFAR-10 in 4 steps. Finally, we show that the full progressive distillation procedure does not take more time tha
中文速览
扩散模型(diffusion models)生成质量很高,但采样太慢——动辄需要数百甚至数千次神经网络推理才能出一张图,严重限制了实际应用。为此,研究者提出了"渐进式蒸馏"(progressive distillation)方法:先训练一个用很多步采样的"教师"模型,再反复将其行为蒸馏进"学生"模型,每轮蒸馏把所需步数砍一半,如此循环直到只需寥寥几步。为了让蒸馏在极低信噪比下也稳定可行,他们还设计了新的模型参数化方式(如预测"v"而非直接预测噪声),避免了原有ε预测在低信噪比时数值不稳定的问题。实验表明,在CIFAR-10、ImageNet、LSUN等标准基准上,从最多8192步的最先进采样器出发,蒸馏到仅需4步时图像质量几乎没有明显下降,例如在CIFAR-10上4步即可达到FID 3.0;更重要的是,整个渐进式蒸馏流程的总耗时不超过从头训练原始模型的时间,让扩散模型在训练和推理两端都变得高效实用。
原文 arXiv:2202.00512;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2202.00512v2