17 天前

扩散模型本质上是一步生成器

Bowen Zheng, Tianming Yang
扩散模型本质上是一步生成器
摘要

扩散模型(Diffusion Models, DMs)在图像生成及其他领域取得了显著成功。通过基于已训练好的得分模型(score model),利用SDE/ODE求解器沿轨迹进行精细采样,DMs能够生成质量极高的图像。然而,这种高精度采样通常需要多步迭代,计算成本较高。为解决这一问题,基于实例的蒸馏方法被提出,旨在通过一个结构更简单的学生模型模仿结构更复杂的教师模型,从而从DM中蒸馏出单步生成器。然而,我们的研究揭示了这类方法存在一个固有局限:教师模型因具有更多采样步骤和参数量,其优化过程往往收敛于与学生模型不同的局部极小值,导致学生模型在复现教师行为时性能受限。为克服这一问题,我们提出一种新颖的分布蒸馏方法,该方法引入了一种专有的分布损失(distributional loss)。该方法在显著减少训练图像数量的前提下,超越了当前最优(SOTA)性能表现。此外,我们发现扩散模型在不同时间步长下,其各层激活程度存在显著差异,这表明模型本身具备单步生成图像的内在潜力。在分布蒸馏过程中,冻结大部分卷积层,可有效释放这一内在能力,进一步提升生成性能。所提方法在CIFAR-10(FID 1.54)、AFHQv2 64×64(FID 1.23)、FFHQ 64×64(FID 0.85)以及ImageNet 64×64(FID 1.16)等基准数据集上均取得了当前最优结果,且效率极高。绝大多数实验仅使用500万张训练图像,在8块A100 GPU上训练6小时内即完成,展现出强大的训练效率与生成性能。