【NIPS 2024】One-Step Diffusion Distillation through Score Implicit Matching

type

status

date

slug

summary

原文摘要翻译

尽管扩散模型在许多生成任务中表现出色，但它们需要大量的采样步骤才能生成逼真的样本。这促使研究社区开发有效的方法，将预训练的扩散模型提炼成更高效的模型。然而，这些方法通常仍需要少量步骤的推断过程，或者其性能显著低于原始模型。在本文中，我们提出了一种名为 Score Implicit Matching (SIM) 的新方法，用于将预训练的扩散模型提炼成单步生成器模型，同时几乎保持与原始模型相同的样本生成能力。此外，该方法是无数据需求的，在提炼过程中不需要任何训练样本。

该方法的核心在于，尽管传统的基于得分的损失（score-based loss）对于生成器模型而言难以直接最小化，但在某些条件下，我们可以高效地计算扩散模型与生成器之间一类广泛的基于得分的散度（score-based divergences）的梯度。在经验验证中，SIM 在单步生成器上表现出强劲性能：在 CIFAR10 数据集上，无条件生成的 FID 为 2.06，类别条件生成的 FID 为 1.96。此外，通过将 SIM 应用于领先的基于 transformer 的扩散模型，我们提炼出了一种用于文本到图像（T2I）生成的单步生成器，在无性能下降的情况下实现了 6.42 的美学评分，明显优于其他单步生成器，包括 SDXL-TURBO（5.33）、SDXL-LIGHTNING（5.34）和 HYPER-SDXL（5.85）。我们将随论文发布这一适用于工业的单步 transformer 基础的 T2I 生成器。

1. 解决了什么问题？

这是一个蒸馏出来的扩散生成模型，在生成速度和生成质量取得很棒的平衡

通过一种特殊的方法来蒸馏模型，使得扩散模型仅需 “单步（One-Step）” 即可生成高质量的图片

所谓特殊：

“Though these studies have contributed to the community in both theoretical and empirical aspects with applicable single-step generator models, their theories are built upon specific divergences, namely the KullbackLeibler divergence and the Fisher divergence, which potentially restrict the distillation performances. A more general framework for understanding and improving diffusion distillation is still lacking.”

蒸馏过程还不需要额外的训练数据

2. 核心思想&创新点

生成数据的本质，是模拟这种数据的分布。生成模型的核心目标是学习目标数据的分布（通常称为 真实数据分布，记为），并能够从中采样出与真实数据分布相似的样本。现实世界中的数据（如图像、文本、音频等）并不是随机生成的，而是遵循某种分布。例如：

自然图像的数据分布可能包含某些特定的统计特性（如相邻像素值的局部相关性）。

文本数据分布可能反映语言的语法规则和上下文语义关系。

数据分布的本质是：某种规律，让某些数据更可能出现，而其他数据不太可能出现。

如果我们将所有可能的数据点（如像素值、单词序列）看作一个空间，则数据分布可以用一个概率密度函数来描述：

是一个非负值，表示在数据空间中的某个点出现的概率。
目标是用生成模型去学习这个分布，使得生成的样本与原始数据的分布一致。

如何评判你学习得到的分布，它到底好不好？那自然是看你学得的分布与真实数据分布一致程度如何。如何判断一致成程度？算一下“两个函数之间的距离”就行。

2.1 为什么要用“距离”来评估分布一致性？

(1) 数据分布 和模型分布

目标：生成模型的训练目的是让模型分布尽可能接近真实数据分布。

问题：但我们无法直接看到这两个分布的具体形式，尤其是通常是未知的（我们只能通过训练数据的样本来间接推测它）。

(2) “分布之间的距离”

为了评估和的一致性，我们需要定义一种度量方式，用来衡量两个分布之间的差异。这种度量方式可以是散度（divergence）或距离（distance），它们是数学上的工具，用于描述两个分布之间的“远近”。

(3) 通过距离引导优化

在训练生成模型时，我们会以这种距离作为优化目标，通过最小化距离来不断调整，使其更接近。

2.2 如何衡量两个分布之间的距离？

衡量两个分布的距离有很多方法，以下是常见的几种：

(1) Kullback-Leibler 散度（KL 散度）

定义：

其中 $P$ 是目标分布（真实分布 $p_{\text{data}}$），$Q$ 是模型分布 $p_\theta$。

直观解释：

KL 散度衡量的是：如果我们用模型分布来近似真实分布，会有多少信息损失。

性质：

KL 散度对的高概率区域更敏感，而对低概率区域的差异不敏感。
KL 散度是不对称的，即。

(2) 基于分数函数的距离

定义：在一些生成模型中，我们并不直接比较分布，而是比较其分数函数（score function，即分布对数梯度）：

直观解释：

分数函数表示分布的“梯度信息”，基于分数的距离可以更高效地衡量分布差异。

2.3 在生成模型中如何具体评估分布距离？

不同生成模型中，分布距离的评估方式可能略有不同：

(1) 最大似然估计

在基于概率密度的生成模型（如 VAE）中，训练目标是最大化数据在模型分布下的似然：

最大似然目标等价于最小化 KL 散度。

(2) 对抗训练（如 GANs）

在生成对抗网络（GANs）中，目标是通过对抗训练使生成分布和真实分布更接近。

原始 GAN 使用的是 JS 散度的近似。

Wasserstein GAN 使用 Wasserstein 距离作为优化目标。

(3) 扩散模型

在扩散模型中，分布距离通常通过分数匹配（score matching）来间接优化：

这里的目标是使模型的分数函数逼近真实数据分布的分数函数。

现状仍是“A more general framework for understanding and improving diffusion distillation is still lacking.” 于是，这篇工作提出了 Score Implicit Matching (SIM)，一种能保持高生成质量的扩散模型蒸馏框架。它的核心放在减小蒸馏前后模型生成质量的差异 ，自然可以联想到上面用各种散度来衡量两种分布的差异的方法，用一种散度来衡量模型蒸馏前后的差异——也是本文的做法：

we propose a wide and flexible class of score-based divergences between the (intractable) score function of the generator model and that of the original diffusion model,

但是，这篇文章指出，从前的方法用来计算蒸馏前后模型差异的散度过于单一，会潜在的限制蒸馏的表现：

For instance, Luo et al. [42], Yin et al. [81] have studied the algorithms that minimize the KL divergence between teacher and one-step student models. Zhou et al. [92] have explored distilling with Fisher divergences, resulting in impressive empirical performances. Though these studies have contributed to the community in both theoretical and empirical aspects with applicable single-step generator models, their theories are built upon specific divergences, namely the Kullback-Leibler divergence and the Fisher divergence, which potentially restrict the distillation performances.

这个论点的提出引出了该文的突破点——能够用任意的距离函数来衡量两者的差异：

To do so, we propose a wide and flexible class of score-based divergences between the (intractable) score function of the generator model and that of the original diffusion model, for arbitrary distance functions between the two score functions.

而在采取“The Pseudo-Huber distance”这款距离函数时，蒸馏效果最佳，表现突出。作者还在Appendix

2.4 相关工作

扩散蒸馏 (Diffusion Distillation) 是一个研究领域，旨在通过教师扩散模型降低生成成本。它主要包含以下三种蒸馏方法：

轨迹蒸馏 (Trajectory Distillation)：

此方法通过减少扩散模型的生成步骤，训练学生模型模仿扩散模型的生成过程。它有以下变体：

直接蒸馏 (Direct Distillation) ([38, 14]) 和 渐进式蒸馏 (Progressive Distillation) ([60, 47])：这些方法通过从噪声输入中预测更少噪声的数据来实现。
基于一致性的蒸馏方法 (Consistency-based Methods) ([67, 28, 65, 35, 16])：这些方法最小化自一致性度量 (self-consistency metric)，但需要真实数据样本来进行训练。

分布匹配 (Distributional Matching)：

该方法关注学生模型的生成分布与教师扩散模型分布的对齐，包括以下两类：

对抗训练方法 (Adversarial Training Methods) ([75, 76])：这些方法需要真实数据来进行扩散模型蒸馏。
基于散度最小化的方法：这一研究方向通过最小化散度（如 KL 散度 [81]）来实现，例如 Diff-Instruct (DI) [44, 81]，以及通过最小化 Fisher 散度（如 Score Identity Distillation (SiD) [92]）来实现。许多此类方法通常不需要真实数据。

尽管 SIM 从 SiD 和 DI 获得了一些灵感，但 SIM 与 SiD 和 DI 之间存在显著差异：

SIM 不仅提供了扎实的数学基础，有助于深入理解扩散蒸馏，还在使用不同距离函数方面表现出极大的灵活性。
SIM 在使用特定距离函数（例如 Pseudo-Huber 距离）时展现了强大的实验性能。

其他方法 (Other Methods)：

如 算子学习 (Operator Learning) ([85]) 和 ReFlow ([36]) 等方法为蒸馏提供了不同的见解。

此外，许多研究在将扩散蒸馏扩展到一步文本到图像生成及其他领域中取得了卓越成就 ([39, 49, 68, 81, 91])。

3. 核心方法

经过上面的介绍，可以自然的想到，作者提出一种更好的计算蒸馏前后扩散模型所学得的数据分布散度的方法，拿来作为蒸馏的优化目标。

3.1 优化目标

首先定义了一个用于计算两个 score function 之间散度的公式，任意两个数据分布之间的散度都可以用这个式子来算：

表示在时间区间 [0,T] 上，分布 p 和 q 的一种距离度量。

它衡量的是 p 和 q 在整个时间区间的分数函数（score functions）之间的差异。

对时间 t 从 0 到 T 进行积分，表示在整个时间区间上的累积误差。

权重函数，表示每个时间步 t 对整体距离的贡献权重。

它可以调整时间步之间的影响程度，例如更关注生成过程的某些阶段。

表示对随机变量从分布中采样的期望。

距离函数衡量和之间的差异。

这个公式自然可以用来作为 Loss Function。优化过程中，自然是会让其对参数求导，但这显然是计算量过于巨大。于是作者推导出了定理3.1 来简化损失函数，让其能高效的求导、收敛。该定理是分数散度优化的核心结果，为隐式生成模型（特别是扩散蒸馏模型）的训练提供了理论基础。

定理 3.1（分数散度梯度定理，Score-Divergence Gradient Theorem）

如果分布满足一些轻微的正则性条件，那么对于任意的分数函数，以下公式对所有参数都成立：

这里的关键观察是，我们将其左侧的分数函数梯度的不可解部分替换为右侧的一个更易计算的分数函数，后者可以通过一个单独的近似网络更容易地完成。这一定理可以通过分数投影恒等式（score-projection identity [69, 92]）证明，该恒等式最初被引入用于将去噪分数匹配与去噪自编码器联系起来。然而，证明定理 3.1 的关键在于，恰当地选择 θ 参数的依赖性或独立性，并在适当位置停止梯度运算。

论文的核心贡献是通过引入停止梯度操作符和分数投影恒等式，将不可解的梯度优化问题转化为可解形式。这一过程依赖于对分数函数、条件分布以及梯度操作的精确理解，并通过定理 3.1 提供了理论支持。

现在我们可以揭示用于蒸馏的目标函数了

其中。根据定理 3.1，这种替代的损失函数的梯度与原始损失函数的梯度是相同的，并且不需要访问分数网络的梯度。

在实践中，我们可以使用另一个在线扩散模型来逐点逼近生成器模型的分数函数，这一方法在之前的工作中也被采用。

作者将最小化目标函数的蒸馏方法命名为 Score Implicit Matching (SIM)，因为该学习过程隐式地将隐式学生模型的不可解边际分数函数 s_{p_θ,ₜ}(.) 与预训练扩散模型的显式分数函数 s_{qₜ}(.) 相匹配。