添加链接
link之家
链接快照平台
  • 输入网页链接,自动生成快照
  • 标签化管理网页链接
相关文章推荐
成熟的椰子  ·  Flutter:scrollControll ...·  12 月前    · 
不要命的饼干  ·  hibernate - ...·  1 年前    · 
扩散模型蒸馏技术综述

扩散模型蒸馏技术综述

讲一下这篇论文:《A Comprehensive Survey on Knowledge Distillation of Diffusion Models》

论文地址: A Comprehensive Survey on Knowledge Distillation of Diffusion Models

扩散模型(Diffusion Models,简称DMs),也被称为基于分数的扩散模型,利用神经网络来指定分数函数。与大多数其他概率模型不同,DMs直接对分数函数进行建模,使其在参数化方面更加灵活,并且在概率建模方面具有潜在的高表达能力。DMs可以学习底层分布的细粒度知识,即边缘分数函数。因此,探索如何蒸馏DMs的知识并充分利用其潜力是一个关键的研究方向。我们的目标是提供一个全面易懂的现代DM蒸馏方法综述,从介绍DMs开始,讨论将其蒸馏为神经向量场所涉及的挑战。我们还概述了现有关于将DMs蒸馏为随机和确定性隐式生成器的研究工作。最后,我们回顾了加速扩散采样算法作为一种无需训练的蒸馏方法。我们的教程适用于具有基本生成模型理解的个人,他们希望应用DM的蒸馏技术或从事该领域的研究项目。

自回归模型(ARMs): 建模显式的似然函数

深度生成建模的目标是训练能够生成高度逼真数据样本的神经参数化模型。已经提出了许多深度生成模型来实现这一目标,每个模型都从不同的角度出发。一般来说,生成模型通过最小化概率差异或度量来表达和逼近底层数据分布的某些充分特征。似然函数,即底层分布的对数密度,是一种常用的分布特征。自回归模型(Auto-regressive models,ARMs)[1,2,3] 是使用神经网络对对数似然函数进行参数化的典型模型,即概率密度函数的对数,并学习匹配底层数据的对数似然。ARMs使用KL散度最小化进行训练。假设 p_d 表示底层数据分布,我们只能获取一致的样本,即 x\sim p_d 。ARMs通过严格的顺序将一系列神经网络输出相加,以显式表达模型似然函数的条件因子分解,即

p_\theta(x) = f_\theta(x^{(1)})+ \Sigma_{i=2}^D f_\theta(x^{(i)}|x^{(1)},...,x^{(i-1)}).

这里, x^{(i)} 表示点 x\in \mathbb{R}^D 的第 i 个坐标,遵循一定的严格顺序。如果 x 是连续的,且在欧几里得空间中,项 f_\theta(x^{(i)}|x^{(1)},...,x^{(i-1)}) 通常通过具有均值和方差的高斯分布的条件分布实现,其中均值和方差是从神经网络的输出中获得的,如下所示:

f_\theta(x^{(i)}|x^{(1)},...,x^{(i-1)}) = \log \mathcal{N}(x^{(i)}; \mu_\theta(x^{(1)},..,x^{(i-1)}), \Sigma_\theta (x^{(1)},..,x^{(i-1)})).

这里, \mathcal{N}(x;\mu,\Sigma) 表示具有均值 \mu 和协方差矩阵 \Sigma 的多元分布的密度函数。术语 \mu^{(i)}_\theta \Sigma^{(i)}_\theta 通常实现为一些神经网络的输出。对于离散数据,公式中的条件分布通常通过某些神经网络的软最大化输出来实现,表达式如下所示:

f_\theta(x^{(i)}|x^{(1)},...,x^{(i-1)}) = \operatorname{Softmax} (g_\theta(x^{(1)},..,x^{(i-1)})).

这里, g_\theta(.) 是一个输出与 x^{(i)} 的离散范围具有相同维度的神经网络。

p_d p_\theta 之间的KL散度定义为:

\mathcal{D}{KL}(p_d,p_\theta) :=\mathbb{E}_{x\sim p_d} \big[ \log p_d(x) - \log p_\theta(x) \big]= \mathbb{E}_{p_d}\big[\log p_d(x) \big] - \mathbb{E}_{p_d} \big[\log p_\theta(x)\big].

由于上式的第一项不依赖于参数 \theta ,通过最小化 p_d p_\theta 之间的KL散度来训练ARMs等价于

\theta^* = \arg\min_\theta \mathcal{D}{KL}(p_d,p_\theta) = \arg\max_\theta \mathbb{E}_{p_d} \big[\log p_\theta(x) \big].

问题中的最优解 \theta^* 也被称为参数 \theta 的最大似然估计(MLE)。项 \mathcal{L}(\theta) = \mathbb{E}_{p_d} \big[\log p_\theta(x) \big] 是期望似然。关于ARMs需要强调两个重要的问题。首先,ARMs显式地对数似然函数进行建模,这限制了实现的灵活性。其次,ARMs在生成算法中需要严格的顺序,这使得从ARMs中采样在计算上效率低下。

能量模型(EBMs):从严格的似然函数到非归一化的似然函数

EBMs(能量基模型)是一类生成模型,其工作原理是通过定义能量函数为每个数据样本分配一个得分。EBMs的基本原则是能量值越低,对应的数据样本的概率越高。最初,EBMs的设计目标是对显式似然进行建模,其中能量函数与负对数似然成比例。然而,这种归一化限制了神经网络的全部潜力。例如,连续值ARMs的实现方式将条件分布限制为多元高斯分布,这可能不适用于真实世界的数据分布。即使神经网络具有无限的表达能力,这种约束也会阻止ARMs完全匹配数据分布。相反,能量基模型(EBMs)通过使用无约束的神经网络来表示和匹配数据的势函数(即非归一化密度函数的对数),克服了这种归一化问题。形式上,分布 p 的势函数被定义为满足方程

e^{E(x)}/\int e^{E(x)}dx = p(x).

直观上,EBMs中的势函数 E(x) 表示密度 p(x) 的非归一化部分的对数。由于势函数本身不需要归一化,EBMs使用神经网络 E_\theta 来表示模型的势函数。EBMs的最基本的训练策略也是使用最大似然训练。具体来说,由于EBM引导的分布具有形式 p_\theta(x) = e^{E_\theta(x)}/Z_\theta ,其中 Z_\theta = \int e^{E_\theta(x)}dx 是归一化常数,往往难以计算,在大多数情况下被视为难以处理的。EBM引导的分布的期望似然写为

\mathcal{L}(\theta) = \mathbb{E}_{p_d} \big[\log p_\theta(x) \big] = \mathbb{E}_{p_d} \big[E_\theta(x)\big] - \log Z_\theta.

对期望似然进行参数梯度计算,我们有

\frac{\partial}{\partial\theta}\mathcal{L}(\theta) = \mathbb{E}_{p_d} \big[\frac{\partial}{\partial\theta}E_\theta(x)\big] - \frac{\partial}{\partial\theta} \log Z_\theta

= \mathbb{E}_{p_d} \big[\frac{\partial}{\partial\theta}E_\theta(x)\big] - \frac{1}{Z_\theta} \frac{\partial}{\partial\theta} Z_\theta

= \mathbb{E}_{p_d} \big[\frac{\partial}{\partial\theta}E_\theta(x)\big] - \frac{1}{Z_\theta} \frac{\partial}{\partial\theta} \int e^{E_\theta(x)}dx

= \mathbb{E}_{p_d} \big[\frac{\partial}{\partial\theta}E_\theta(x)\big] - \frac{1}{Z_\theta} \int \big[ \frac{\partial}{\partial\theta} E_\theta(x)\big] e^{E_\theta(x)}dx

= \mathbb{E}_{p_d} \big[\frac{\partial}{\partial\theta}E_\theta(x)\big] - \int \big[ \frac{\partial}{\partial\theta} E_\theta(x)\big] \frac{e^{E_\theta(x)}}{Z_\theta}dx

= \mathbb{E}_{p_d} \big[\frac{\partial}{\partial\theta}E_\theta(x)\big] - \mathbb{E}_{p_\theta} \big[ \frac{\partial}{\partial\theta} E_\theta(x)\big].

为了计算关于参数 \theta 的期望对数似然梯度,我们需要从 EBM 产生的分布 x \sim p_\theta 中获取一致的样本,该分布是一个非归一化分布。幸运的是,有几种 MCMC 算法能够生成这样的样本,包括 [4]、[5]、[6]、[7] 和 [8] 提出的算法。通过将梯度公式与 MCMC 算法相结合,可以使用最大似然估计来训练 EBM。能量基模型(EBMs)和自回归模型(ARMs)之间的关键区别在于它们如何利用神经网络。EBMs 使用无约束的神经网络来建模潜在函数,而 ARMs 将神经网络作为显式条件密度的组成部分。这种方法上的差异使得 EBMs 可以充分利用神经网络的表达能力,避免了归一化要求所带来的限制。

SBMs and DMs: 从势能函数到得分函数

在前面的章节中,我们介绍了能量基模型(EBM)及其训练和采样方法。为了近似数据潜力函数,EBMs使用神经网络,并利用学习到的潜力函数进行马尔可夫链蒙特卡洛(MCMC)算法生成样本。在众多可用的MCMC方法中,Langevin动力学(LD)是首选,因为它易于实现,并且在较弱条件下具有良好的性能。设 p(x) 为可微的密度函数,LD通过随机微分方程定义如下:

dX_t = \frac{1}{2}\nabla_{X_t} \log p(X_t)dt + dW_t, ~p^{(0)}=p_0, ~t\in [0,\infty].

Langevin动力学有两个显著特点。首先,当 t\to\infty 时,LD的边际分布可以收敛到 p(x) ,无论初始分布 p_0 如何,只要 p(x) 满足一定条件。其次,在LD的模拟过程中只需要潜力函数的梯度,即分数函数 \nabla_x \log p(x) 。因此,即使分布 p(x) 未经过归一化,使用潜力函数 \log p(x) 进行LD模拟仍会生成有效的样本作为归一化分布。得到启发于LD的基于分数的模型(SBMs)利用神经网络训练神经分数函数 S_\theta 以匹配基础数据分布。如果神经分数函数能够很好地匹配数据分数函数 S_d(x):= \nabla_x \log p_d(x) ,则可以通过在LD模拟中用SBM的神经分数函数 S_\theta(x) 替换分数函数 \nabla_x \log p(x) 来获得SBM的样本。

现在,让我们转向SBM的训练策略。由于 S_\theta 直接表达了模型的基础分数函数,模型的潜力函数 p_\theta 变得难以处理。因此,前面介绍的用于训练EBM和ARM的KL散度在训练SBM时不起作用。相反,SBM使用费舍尔散度(Fisher divergence)进行训练,这是一种只需要模型的分数函数的概率散度。具体而言,两个分布 p(x) q(x) 之间的费舍尔散度定义为:

\mathcal{D}_{F}(p,q) := \mathbb{E}_{x\sim p} ||\nabla_x \log p(x) - \nabla_x \log q(x)||^2_2.

p_d 表示数据分布, S_d(x)=\nabla_x \log p_d(x) 表示数据分数函数。数据分布与SBM诱导的分布 p_\theta 之间的费舍尔散度可以写成:

\mathcal{D}_{F}(p_d, p_\theta) = \mathbb{E}_{p_d} \frac{1}{2}||S_d(x) - S_\theta(x)||_2^2

= \mathbb{E}_{p_d}\big[ \frac{1}{2}||S_d(x) ||_2^2 \big] + \mathbb{E}_{p_d}\big[ \frac{1}{2}||S_\theta(x) ||_2^2 \big] - \mathbb{E}_{p_d}\big[ \langle S_d(x), S_\theta(x) \rangle \big].

在实践中,费舍尔散度的第一项不依赖于参数 \theta ,因此可以忽略。第三项可以等价地表示为一个无关数据分数的形式: \mathbb{E}_{p_d}\big[ \langle S_d(x), S_\theta(x) \rangle \big] = -\mathbb{E}_{p_d} \sum_{d=1}^D \frac{\partial s_\theta^{(d)}(x)}{\partial x^{(d)}}.

其中, s_\theta^{(d)}(x) 表示模型分数函数的第 d 个分量, x^{(d)} 表示输入向量的第 d 个分量。因此,最小化费舍尔散度等价于最小化一个易处理的目标函数:

\mathcal{L}_{SM}(\theta) \colon= \mathbb{E}_{p_d} \big[\frac{1}{2}||S_\theta(x)||_2^2 + \mathbb{E}_{p_d} \sum_{d=1}^D \frac{\partial s_\theta^{(d)}(x)}{\partial x^{(d)}} \big].

方程中的优化问题被称为分数匹配(Score Matching)。然而,通过神经网络计算梯度项 \langle \nabla_x, S_\theta(x) \rangle 可能需要大量内存,在处理高维数据时会带来挑战。为了克服这个限制,已经提出了几种方法来提高分数匹配的效率,例如sliced score matching, Finite-difference score matching, 和 denoising score matching。其中开创性的工作denoising score matching提出了一种称为去噪分数匹配(Denoising Score Matching,DSM)的目标函数,它不需要计算数据梯度,而是使用以下目标函数:

\mathcal{L}_{DSM}(\theta) \colon= \mathbb{E}_{x\sim p_d, \tilde{x}\sim p(\tilde{x}|x) } ||S_\theta(\tilde{x}) - \nabla_{\tilde{x}} \log p(\tilde{x}|x)||^2_2.

在上述段落中,目标函数 \mathcal{L}_{DSM}(\theta) 使用了扰动核函数 p(\tilde{x}|x) 进行最小化,该函数在采样时高效且具有明确的表达式。对于扰动核函数,常见的选择是具有噪声方差 \Sigma 的高斯分布 \mathcal{N}(\tilde{x};x,\sigma^2 \mathbf{I}) 。通过最小化这个目标函数,我们可以获得数据分布 \tilde{p}_d(\tilde{x}) 的近似,如果扰动核函数与单位矩阵接近,则该近似与原始数据分布类似。然而,对于高维数据使用LD进行采样可能存在问题,因为数据集中在高维空间中的某个低维流形中。作为解决方案,基于分数的扩散模型使用多个或连续索引的扰动核函数来改善学习和采样。这种方法可以通过生成更准确、更具代表性的原始数据分布样本来提高SBM的性能。

扩散模型:多层SBMs

与使用单个评分网络的SBMs不同,基于评分的扩散模型(DMs)采用了更高级的方法,利用多层或连续索引的评分网络 S_\theta(x,t) 。此外,DMs不使用单个扰动核,而是使用由随机微分方程引发的一族条件转移核来扰动数据。考虑正向扩散SDE

dX_t = F(X_t,t)dt + G(t)dW_t,~X_0\sim p_d^{(0)}=p_d,

其中 W_t 是一个维纳过程。设 p_t(x_t|x_0) 表示正向扩散的条件转移核, p_d^{(t)} 表示扩散时间 t 的边际分布,初始为 p_d^{(0)} = p_d 。在扩散模型的文献中,两个特殊的正向扩散,即方差保持(VP)扩散和方差爆炸(VE)扩散备受青睐。

VP扩散: VP扩散的形式为 dX_t = -\frac{1}{2}\beta(t)X_t dt + \sqrt{\beta(t)}dW_t, ~t\in [0,T], 其中 \beta(t) 是预定义的时间表函数。VP扩散的条件转移核具有显式表达式 p(x_t|x_0) = \mathcal{N}(x_t;\sqrt{\alpha_t}x_0;(1-\alpha_t)\mathbf{I}), 其中 \alpha_t = e^{-\int_0^t \beta(s)ds} 。利用公式,时间 t 的样本模拟是高效的,只需对初始样本进行缩放并添加高斯噪声 X_t = \sqrt{\alpha_t} X_0 + \sqrt{1-\alpha_{t}}\epsilon, ~X_0\sim p_0. 这里 \epsilon\sim \mathcal{N}(0,\mathbf{I}) 是与 X_0 大小相同的标准正态向量。利用公式,获得 X_t 是廉价的,因为我们不需要对扩散过程进行顺序模拟。VP扩散的另一个优点是,在宽松条件下,VP扩散可以将任意初始分布 p_0 转换为标准多变量高斯分布 \mathcal{N}(0, \mathbf{I}) 。VP扩散可能是DM文献中最广泛使用的正向扩散方法之一。

VE扩散: VE扩散的形式为 dX_t = \sqrt{\frac{d \sigma^2(t)}{dt}}dW_t, ~t\in [0, T]. W_t 是一个独立的维纳过程。VE扩散的转移核为 p(x_t|x_0) = \mathcal{N}(x_t; x_0, \sigma(t) \mathbf{I}). 类似于VP扩散,VE扩散的边际样本是廉价的,可以使用以下公式绘制 X_t = X_0 + \sigma(t) \epsilon,~X_0\sim p_0 这里 \epsilon\sim \mathcal{N}(0,\mathbf{I}) 是一个标准高斯向量。

VE和VP扩散过程在各种任务的扩散模型中已经成功使用。然而,近年来,已经提出了一些新的扩散过程,这些扩散过程不仅改善了DM的性能,而且还针对特定任务进行了设计。

训练方法:

DM通过在每个时间步 t 上最小化带有扰动核 p_t(.|.) 的DSM的加权组合来进行训练。更具体地, DM的训练目标可以表示为:

\mathcal{L}_{WDSM}(\theta) = \int{t=0}^T w(t) \mathcal{L}_{DSM}^{(t)}(\theta)dt.

其中 \mathcal{L}_{DSM}^{(t)}(\theta) = \mathbb{E}_{x_0\sim p_d^{(0)}, x_t|x_0 \sim p_t(x_t|x_0)} ||S_\theta(x_t,t) - \nabla_{x_t}\log p_t(x_t|x_0)||_2^2. 通过最小化目标函数,连续索引的评分网络 S_\theta(x,t) 能够匹配正向扩散过程的边际评分函数。

在某些文献中,DSM目标被重新表述为一种通过学习预测添加噪声来训练DM的 噪声预测 目标。具体来说,以VP的转移核为例,如果使用重新参数化技术获得 x_t x_0 ,则梯度项可以写成

\nabla_{x_t} \log p_t(x_t|x_0) = \nabla_{x_t} \bigg[ -\frac{1}{2(1-\alpha_t)} ||x_t - \sqrt{\alpha_t}x_0 ||_2^2 \bigg]\

= \frac{1}{1-\alpha_t} (x_t - \sqrt{\alpha_t} x_0) = \frac{1}{1-\alpha_t} \sqrt{1-\alpha_t}\epsilon\

= \frac{1}{\sqrt{1-\alpha_t}} \epsilon

将梯度项 和DSM目标结合起来,VP扩散的DSM目标可以重新表述为

\mathcal{L}_{DSM}(\theta) = \mathbb{E}_{x_0\sim p_d, \epsilon\sim \mathcal{N}(0,\mathbf{I})} \frac{w(t)}{1-\alpha_t} ||\sqrt{1-\alpha_t}S_\theta(x_t,t) - \epsilon||_2^2

= \mathbb{E}_{x_0\sim p_d, \epsilon\sim \mathcal{N}(0,\mathbf{I})} \frac{1}{1-\alpha_t} ||\epsilon_\theta(x_t,t) - \epsilon||_2^2

DSM可以使用噪声预测目标进行重新表述,其中一个神经网络 \epsilon_\theta(x_t,t) 被训练用于使用加了噪声的数据样本 x_t 和评分网络 S_\theta(x_t,t) 来预测添加的噪声 \epsilon 。这个重新表述涉及到评分网络的修改版本,表示为 \epsilon_\theta(x_t,t) \colon= \sqrt{1-\alpha_t} S_\theta(x_t,t) ,以及扩散过程的重新表达式 x_t = \sqrt{\alpha_t} x_0 + \sqrt{1-\alpha_t} \epsilon 。类似的表述也可以应用于VE扩散过程,详细内容可以在附录中找到。

采样策略

DSM中训练的评分网络可以在各种应用中发挥作用,其中最直接的应用之一是设计一种采样策略来近似底层数据分布。这种机制背后的基本概念是存在一个反向SDE,它具有与正向扩散过程相同的边际分布,

dX_t = [F(X_t,t) - G^2(t)\nabla_{x_t} \log p^{(t)}(X_t)]dt + G(t)d\bar{W}t, t\in [T,0], X_T\sim p_d^{(T)}.

此外,还发现了一个ODE,它与正向SDE具有相同的边际分布,

dX_t = [F(X_t,t) - \frac{1}{2}G^2(t)\nabla{x_t} \log p^{(t)}(X_t)]dt.

反向SDE和ODE都依赖于真实的边际评分函数 \nabla_{x_t}\log p_d^{(t)}(x_t) 。然而,通过用学习到的神经评分函数 S_\theta(x,t) 替代真实的边际评分函数,可以得到DM的生成SDE和ODE。此外,在某些情况下,生成ODE的概念可以扩展到神经连续时间归一化流模型。通过使用学习到的评分函数 S_\theta(x,t) ,可以通过采样SDE或ODE的数值解来从DM中进行采样。已经开发了许多使用先进数值技术的实际算法,以提高DM的生成性能或在最小损失性能的情况下提高采样效率 [9, 10, 11, 12]。

DMs的成功

自[13]、[14] 和 [15]的开创性工作以来,扩散模型(DMs)已成为生成建模的领先方法,在各个领域得到广泛应用,包括神经图像合成和编辑[16, 17, 18, 19, 20]、音频和分子合成 [21, 22]、图像分割 [23],以及视频或三维物体生成 [24, 25, 26] 等。DMs在性能上取得了显著的改进,可以通过对数据集(如CIFAR10)上的无条件Frechet Inception Score(FID)[27] 的减少趋势进行观察,从25.32 [28] 减少到1.97 [29] 。

扩散蒸馏: 知识蒸馏(knowledge distillation)的概念旨在创建更小、更高效的模型,同时保持准确性,在各个研究领域取得了巨大成功 [30, 31] 。特别是,从预训练分类器中蒸馏知识已经导致了具有可比准确性、减小模型大小和提高推理效率的模型 [32] 。鉴于扩散模型(DMs)在许多应用中取得的成功,越来越多的研究者对从这些模型中蒸馏知识以创建更小、更高效版本产生了兴趣。扩散蒸馏的一个重要动机是显著加快采样速度,目前这一速度受到了大量神经函数评估的限制。为了提高DMs的推理效率,研究者们正在探索将从DMs中学到的知识蒸馏到高效的采样机制中,例如直接的隐式生成器或更少步骤的向量场。通过这样做,他们已经成功地创建了学生模型,进一步提高了推理效率,并且在性能损失最小的情况下实现了可比的生成性能。一些经过蒸馏的学生模型只需要少于10次神经函数评估,但仍能提供与较大模型相当的生成性能。

扩散蒸馏也为建立DMs与其他生成模型(如隐式生成模型和归一化流)之间的联系提供了一种手段。通过DMs与其他模型之间的知识转移,研究者可以研究它们之间的微观连接,并探索它们在未来生成建模研究中的潜力。

本文对扩散蒸馏策略的现有研究进行了全面的综述。我们的综述分为三个主要类别:扩散到场(diffusion-to-field,D2F)蒸馏,扩散到生成器(diffusion-to-generator,D2G)蒸馏和无需训练(training-free,TF)蒸馏。每个类别包含了具有相似设置和方法的研究。除了我们的分类,我们还在本综述的其余部分讨论了扩散蒸馏的更广泛主题。

扩散模型蒸馏

扩散模型蒸馏(Diffusion-to-Field Distillation)的方法旨在解决确定性采样方法在生成随机模拟数据(DMs)中的低效率问题,通过将生成的常微分方程(ODE) 蒸馏为另一个生成的向量场,从而减少生成相似样本所需的 NFE 数量。该方法可以分为两类:输出蒸馏和路径蒸馏。输出蒸馏旨在教导一个学生向量场来复制 DM 的确定性采样方法的输出结果。而路径蒸馏则旨在生成一种具有比教师ODE更好的路径特性的学生ODE。输出蒸馏和路径蒸馏可以结合使用,以改善教师ODE的路径特性和模拟效率。


输出蒸馏

首先,我们回顾生成的常微分方程(ODE):

dX_t = [F(X_t,t)-\frac{1}{2}G^2(t)\nabla_{X_t} \log p^{(t)}(X_t)]dt,~t\in [T,0], ~p^{(T)}=p_T

为了简化讨论,我们考虑最简单的Euler-Maruyama(EM)离散化方法,它通过顺序更新来求解ODE:

X_{t_i} = X_{t_{i+1}} + F(X_{t_{i+1}},t_{i+1})-\frac{1}{2}G^2(t_{i+1})\nabla_{X_{t_{i+1}}} \log p^{(t_{i+1})}(X_{t_{i+1}}), ~i=N,...,1

在从DMs中进行采样的情况下,直接使用EM方法进行数值求解会导致计算效率低下。这是因为EM离散化误差随着步长的增加而显著增加,导致生成的样本质量较差。为了解决这个问题,提出了一种名为输出蒸馏的替代方法。该方法涉及训练一个学生神经网络,使用更大的步长学习ODE的输出。具体来说,假设步长$\Delta t$不太小,我们训练一个连续索引的学生神经映射 S_\phi^{(stu)}(x,t) 来近似ODE输出的变化。这有助于通过减少生成竞争性样本所需的NFE数量来提高从DMs中进行采样的计算效率。更准确地说,我们训练学生连续索引的神经映射 S_\phi^{(stu)}(x,t) 来近似教师ODE在时间 t t-\Delta t 之间的输出变化:

\Delta X \colon= X_{t-\Delta t} - X_{t} = \int_{s=t}^{t-\Delta t} [F(X_s,s)-\frac{1}{2}G^2(s)\nabla_{X_s} \log p^{(s)}(X_s)]ds.

通过使用足够小的步长,可以使用数值方法处理方程中的残差。为了克服这种方法的计算效率问题,引入了一个学生网络 S_\phi^{(stu)} ,它作为一个经过蒸馏的时间相关向量场,近似了方程中的非线性残差。通过利用神经网络的表达能力,实证研究表明,经过适当设计的扩散到场技术可以产生一个新的采样ODE,该ODE只需要一个NFE,但在生成性能(以Fretchet Inception Distance(FID)衡量)方面仍然具有可比较的性能。值得注意的是,在接下来的章节中,除非另有说明,我们使用 \theta 表示教师模型(DMs)的参数,使用 \phi 表示学生模型的参数。

[33] 提出了一种知识蒸馏(KD)策略,将DDIM采样器蒸馏为一个只需要一个NFE的高斯模型。DDIM是一种典型的VP扩散模型的确定性采样策略。假设 \epsilon_\theta(x,t) 是DM中的噪声预测网络,如方程所讨论的,DDIM采样器是一个确定性采样器,它按顺序更新

X_{t_{i-1}} = \sqrt{\frac{\alpha_{t_{i-1}}}{\alpha_{t_i}}}(X_{t_i} - \sqrt{1-\alpha_{t_{i}}}\epsilon_\theta(X_{t_i})) + \sqrt{1-\alpha_{t_{i-1}}}\epsilon_\theta(X_{t_i})

DDIM是一种确定性采样策略,已经证明可以显著减少生成样本所需的NFE数量,仅需要100个NFE即可维持良好的生成性能。与生成ODE方法相比,DDIM被广泛认为是一种更优越的确定性采样策略。在他们的工作中,[33] 提议将条件高斯模型作为学生生成模型 p_{stu}(x_0|x_T) = \mathcal{N}(x_0; f_\phi(x_T), \mathbf{I}). 神经网络 f_\phi(.) 的输入和输出维度与数据相同。为了实现这一点,作者选择了与教师DM中得分网络相同的结构作为 f_\phi 的体系结构。设 \operatorname{DDIM}(.) 表示由DDIM引起的确定性映射,他们取

p_{teacher}(x_0|x_t) = \mathcal{N}(x_0; \operatorname{DDIM}(x_T),\mathbf{I})

为了训练学生模型,他们提议最小化学生模型和DDIM采样器之间的条件KL散度

\mathcal{L}(\phi) = \mathbb{E}{x_T\sim \mathcal{N}(0,\mathbf{I})} \mathcal{D}{KL}\bigg[p_{teacher}(x_0|x_T), p_{stu}(x_0|x_T)\bigg]

= \mathbb{E}_{x_T\sim \mathcal{N}(0,\mathbf{I})} \mathcal{D}_{KL}\bigg[\frac{1}{2} ||f_\phi(x_T) - \operatorname{DDIM}(x_T)||_2^2\bigg]

[33] 提出的学生模型采样策略很简单。它涉及从均值为0、协方差为单位矩阵的正态分布中抽取高斯随机变量 x_T 。然后,通过将 x_T 传递给神经网络 f_\phi 来获取学生模型的均值向量 x_0 。这导致一个只需要一个NFE的采样模型,其FID为9.39,而教师生成ODE的FID为4.16。尽管这种方法在考虑扩散模型的知识蒸馏方面提供了一种初始步骤,但由于需要生成DDIM或其他ODE采样器的最终输出,它在计算效率上存在不足,因为计算单个批次的训练数据需要数百个NFE。

[34] 提出的渐进蒸馏(PD)策略旨在通过学习教师DM的确定性采样策略的两步预测,训练一个只需要一半NFE数量的学生神经网络。教师扩散模型被离散化为 N 个时间戳,表示为 \mathcal{T} = {t_i}_{\{i=0, 1, ..., N-1\}} ,而 \mathcal{T}'={0,2,4, ...} 表示 \mathcal{T} N/2 个偶数时间戳。学生网络用 f_\phi(x,t) 表示,DDIM方法从教师DM的一个时间戳 t_j 更新到另一个时间戳 t_i 的更新用 \operatorname{DDIM}(x,t_j,t_i) 表示,其中 i\leq j 。PD通过最小化以下目标训练学生网络:

\mathcal{L}(\phi) = \mathbb{E}_{x_0\sim p_d, i\sim Unif(\mathcal{T}'), \epsilon\sim \mathcal{N}(0,\mathbf{I})} ||f_\phi(\tilde{x}, t_i) -\operatorname{DDIM}(\tilde{x}, t_i, t_{i-2})||_2^2,

其中 \tilde{x} = \sqrt{\alpha_{t_i}}x_0 + \sqrt{1-\alpha_{t_i}}\epsilon 是VP扩散在时间 t_i 的向前扩散数据。通过最小化PD目标,学生网络学习输出教师模型的两步预测,从而将总的NFE数量减半。在将学生模型训练到准确预测教师模型的两步采样策略后,它取代了教师模型,并进一步训练新的学生模型,将采样步骤的数量再减半。论文作者在实现渐进蒸馏(PD)方法时,使用了与教师模型的得分网络相同的UNet架构和DDIM方法作为初始教师采样策略。他们的结果显示,连续的PD轮次可以将所需的向量场减少到仅需要4个NFEs,使其比教师扩散的ODE采样器高效率提高250倍,而以FID测量的生成性能下降了5%。PD和KD都使用教师网络的架构通过从多步教师采样方法中学习来蒸馏出少数步采样方法。这两种方法都通过最小化教师网络的多步预测和学生网络的单步预测之间的 L^2 误差来实。PD和KD之间的关键区别在于,PD逐步减少所需的函数评估数量,而KD直接训练最终预测的一步学生模型。因此,KD可以被视为将教师采样策略的完整时间戳在单轮中减少为一个的极端PD方法。

[35] 提出了一种两阶段蒸馏策略,以解决从无分类器引导的条件扩散模型(例如GLIDE、DALL \cdot E-2、Stable Diffusion和Imagen)中蒸馏知识的挑战。关键挑战在于在保持无分类器引导机制的同时,从教师扩散模型中传递知识。这个机制成功地训练了单个扩散模型同时学习条件和无条件分布。通过在扩散模型中引入条件输入,将条件知识融入扩散模型中,并通过用 None 输入替换条件输入来学习无条件知识。在该策略的第一阶段,训练一个带有无分类器引导输入的学生条件扩散模型,从教师扩散模型中学习知识。学生模型采用可学习参数的神经网络实现,接受输入 x 和条件上下文输入 c 。为简化起见,以下讨论中将忽略符号 c 。通过最小化目标函数来训练学生模型的输出与教师模型输出的匹配:

\mathcal{L}(\phi_1) = \mathbb{E}_{w\sim p_w, t\sim U[0,1], x\sim p_d} \bigg[\lambda(t)||f_{\phi_1}(x_t,t,w) - \hat{x}^w_\theta(x_t) ||_2^2 \bigg]

这里的 \lambda(t) 表示预定义的加权函数。

\hat{x}^w_\theta(x_t) = (1+w)\hat{x}_{c,\theta}(x_t) - w \hat{x}_{\theta}(x_t),x_t\sim p_t(x_t|x_0), p_w = U[w_{min}, w_{max}]

值得注意的是,第一阶段的蒸馏仅通过学生模型的额外输入 w 引入了无分类器引导,没有采用减少NFE(Number of Function Evaluations,函数评估数量)和与效率相关的蒸馏方法。第二阶段采用了 [34] 提出的渐进扩散(PD)策略,显著减少之前训练的带有无分类器引导输入的学生模型的扩散步骤数量。这种两阶段蒸馏方法用于蒸馏像素空间和潜空间的无分类器引导条件扩散模型,适用于各种任务,包括类别条件生成、文本引导图像生成、文本引导图像到图像的转换和图像修复。事实上,在ImageNet-64x64数据集上,经过蒸馏的学生模型在像素空间类别条件生成的实验中,甚至比教师模型取得了更好的Fréchet Inception Distance(FID)分数。

在最近的工作中,[36] 提出了一种名为分类器特征蒸馏(Classifier-based Feature Distillation,CFD)的特征空间蒸馏方法,用于处理图像数据。这种方法的主要动机是解决直接对齐教师模型输出和少步骤学生模型输出中的像素非常困难的问题。相反,他们训练学生网络在经过预训练分类器提取的特征空间中与教师模型的多步输出对齐。该方法采用了与PD技术类似的蒸馏策略。

为了实现这一目标,他们提出最小化学生模型一步输出的预测概率分布(经过Softmax函数)与教师模型的多步输出之间的KL散度。他们还发现,加入熵和多样性正则化等附加项可以提高蒸馏性能。在实现中,他们采用了与 [34] 相同的扩散模型和学生模型,并使用DenseNet-201作为分类器。他们的蒸馏方法在CIFAR10上实现了FID 3.80的学生模型,仅使用4个NFE,低于 [34] 中的DP(Diffusion PDE)的实现。

[37] 也采用了渐进扩散(PD)策略,对基于预训练扩散TTS模型的快速文本转语音(TTS)生成模型进行蒸馏。他们引入了方差预测器和声谱图去噪器来改进专门用于TTS应用的学生模型架构。

与模仿扩散模型的生成ODE输出不同,[38] 提出了最小化生成ODE的自洽函数的差异来实现输出蒸馏。他们随机扩散一个真实数据样本,并模拟生成ODE的几个步骤,得到另一个位于同一ODE路径上的噪声数据。他们将这两个噪声样本输入到学生模型中,并最小化输出之间的差异,以确保生成ODE的自洽性。他们将他们提出的模型命名为自洽模型(Consistency Model,CM)。CM可以被视为另一种利用生成ODE的自洽性进行蒸馏的输出蒸馏方法。

对于输出蒸馏,学生网络被训练为最小化其输出与相应时间点上教师ODE输出之间的差异。输出蒸馏在教师ODE具有相对简单形式(例如对角ODE)或难以训练的情况下尤其有用。在这些情况下,学生网络可以通过较少的NFE高效地模仿教师ODE的输出。经过训练的学生网络可以用于从扩散模型中进行高效采样。

值得注意的是,输出蒸馏仅考虑ODE在离散时间点上的输出,因此可能无法完全捕捉到扩散模型的路径特性。因此,对于路径特性重要的应用(如图像合成和编辑),应使用路径蒸馏,下一部分将对其进行讨论。

路径蒸馏

路径蒸馏是另一种技术,通过让学生神经网络模仿教师模型的多步输出来改进扩散模型(DMs)的采样策略。与之相反,路径蒸馏旨在改进DMs的采样策略以获得更好的性质。一些研究人员认为正向(和反向)扩散过程在数据空间中创建了连接数据分布和先验分布的曲线。因此,路径蒸馏的重点是将扩散生成式 SDE 或 ODE 优化为更有效的直线版本,其具有比教师模型更高效的采样路径。路径蒸馏和输出蒸馏的关键区别在于,路径蒸馏更关注优化现有教师模型的采样策略,而输出蒸馏则专注于教导学生模型学习教师模型的跳过输出,而不改变学生模型的采样路径和机制。

Reflow 方法,由 [39] 提出,是一种路径蒸馏方法,旨在通过学生模型修改预训练的教师神经ODE来提高生成速度。学生模型通过最小化输出和从预训练模型得到的数据样本及其输出的插值之间的 L^2 损失的时间平均来拉直教师模型的路径。更具体地说,令 p_T 为初始分布,则目标函数为:

\mathcal{L}(\phi) = \mathbb{E}_{x_T\sim p_T, x_0\sim p_{teacher}(x_0|x_T), t\sim U[0,T]} \bigg[ ||x_T - x_0 - f_\phi(\frac{t}{T} x_T + \frac{T-t}{T} x_0, t)||_2^2 \bigg]

这里的 p_{teacher}(x_0|x_T) 可以是任意教师模型,无论它是SDE还是ODE。在由 [39] 提出的Reflow蒸馏中,如果假设教师模型是DDIM采样器,则 p_{teacher}(x_0|x_T) = \delta(x_0=\operatorname{DDIM}(x_T)) 。Reflow蒸馏的优势在于可以进行多轮重复,进一步拉直教师模型的路径。此外,Reflow可以与逐步蒸馏结合使用,其中Reflow策略首先用于拉直教师模型的路径,然后再使用逐步蒸馏使学生模型更高效。最后,作者提供了修正和未修正模型的ODE路径的数值比较结果。


在他们的工作中,[40] 将 Reflow 技术应用于三维点云生成问题。他们提出了一个三阶段的过程,可以构建一个能够生成高质量点云的生成式ODE。在第一阶段的训练中,他们通过最小化目标函数来训练一个教师生成式ODE f_\theta

\mathcal{L}(\theta) = \mathbb{E}_{x_0\sim p_d, x_T\sim \mathcal{N}(0,\mathbf{I}), t\sim U[0,T]} \bigg[ ||f_\theta(x_t,t) - (x_0-x_T)||^2_2 \bigg]

其中 x_t = (t/T)x_0 + (T-t/T)x_T 是时间 t 的插值点。在第二阶段,他们采用 Reflow 策略进一步“拉直”教师模型的路径。在第三阶段,他们使用一个学生模型 f_\phi ,通过最小化以下目标函数,将多步教师模型提炼成单步学生模型:

\mathcal{L}(\phi) = \mathbb{E}_{x_T\sim p_d} \operatorname{Dist}\bigg[x_T + f_\phi(x_T,T), x_0\bigg]

这里的 x_0 \sim p_{teacher}(x_0|x_T) 是从教师模型中获得的。 \operatorname{Dist}(.,.) 表示两个点之间的某个距离函数。在他们的实现中,他们使用 Chamfer 距离来衡量两个点云之间的距离。

[41] 提出了一种目标函数,可以导致具有更小曲率的正向过程,从而对截断误差更不敏感。[42]、[43] 和 [44] 进一步研究了扩散模型生成路径的改进。[45] 提出了另一种路径蒸馏的观点,学习一个能够生成路径的映射算子。[46] 还提出了另一种路径蒸馏方法,通过微调教师ODE以根据正向路径最小化某个IPM。[47] 提出了通过最小化生成路径的每个边缘分布与数据之间的最大平均差异(MMD)来微调生成路径。

扩散到场的蒸馏仍然是高效扩散模型领域的一个活跃研究领域。其目标是利用从扩散模型和生成式SDE/ODE中获得的知识,将一个具有更快采样速度的学生模型提炼出来,同时保持与教师模型相当的生成性能水平。

Diffusion到生成器的蒸馏

与之前讨论的扩散到场的蒸馏不同,扩散到生成器(D2G)蒸馏是另一类重要的蒸馏方法。D2G蒸馏的主要目标是将扩散模型学习到的分布知识转移到一个高效的生成器中。与D2F蒸馏不同,D2G蒸馏通常涉及具有与数据空间不同潜在维度的隐式生成器。此外,根据具体应用,确定性生成器和随机生成器都被视为学生生成器。D2G蒸馏的训练目标通常采用与扩散模型训练目标类似的形式,而不是D2F蒸馏中常用的基于均方误差的目标函数。

蒸馏确定性生成器

越来越多的关注被放在将确定性生成器(例如神经半径场)作为学生模型进行蒸馏,以进一步应用预训练的大规模扩散模型。更具体地说,预训练的文本到图像扩散模型被发现对于学习与给定文本提示相关的神经半径场非常有用。神经半径场(NeRF)[48] 是一种三维物体,它使用多层感知器(MLP)将网格坐标映射到颜色和密度等体积属性。给定相机参数,如视角的角度,渲染算法输出一个二维图像,即三维 NeRF 场景的视图投影。从给定的视角,NeRF 可视为一个确定性的二维图像生成器,其 MLP 的参数是可学习的。

由于构建神经半径场(NeRF)所需的数据有限,研究人员开始探索蒸馏方法,以获得与给定文本提示相关的NeRF。Poole等人的开创性工作 [26] 提出了一种称为分数蒸馏采样(SDS)的方法,将2D文本到图像扩散模型蒸馏成3D NeRF。与传统的NeRF构建方法不同,传统方法需要目标3D物体的多个视角图像,而基于文本的NeRF构建缺少了3D物体和多个视角图像。SDS方法通过最小化固定视角下NeRF生成的图像与扩散模型的损失函数来优化NeRF。为了避免直接优化扩散模型的损失函数所带来的计算开销,研究人员提出通过省略Unet雅可比项来近似蒸馏目标。具体而言,从具有可学习参数 \phi 的固定视角下渲染的NeRF图像表示为 x=g(\phi) ,前向扩散样本表示为 x_t = \sqrt{\alpha_t}x_0 + \sqrt{1-\alpha_t}\epsilon 。训练的文本条件扩散模型表示为 \epsilon_\theta(x,t,c) ,其中 t c 表示时间戳和文本提示。SDS使用梯度来更新参数 \phi ,更新规则如下:

\operatorname{Grad}(\phi) = \frac{\partial}{\partial \phi} \mathcal{L}(\phi) = \mathbb{E}_{t,\epsilon,x=g(\phi)} \bigg[w(t)(\epsilon_\theta(x_t,t,c) - \epsilon)\frac{\partial x}{\partial \phi} \bigg]

Poole等人的工作在从文本提示生成3D NeRF的任务上取得了显著的成果,这进一步推动了将扩散模型蒸馏为确定性生成器的研究。一些研究将SDS方法扩展到其他应用领域,如 [49, 50, 51, 52, 53] 所报道。特别是 [49] 对SDS算法在文本到3D生成方面的可调参数进行了深入研究。同时,[50] 提出了一种两阶段的优化策略,进一步提升了SDS在高分辨率图像上的性能。他们的方法的第一阶段类似于 [26] 的方法,可以得到低分辨率的NeRF。第二阶段将第一阶段训练得到的NeRF放大到更高的分辨率,并进行微调以获得更好的性能。

扩散模型蒸馏成功应用于确定性生成器,特别是在3D生成建模领域,已被广泛认可。然而,开发性能更好、更高效的蒸馏策略的研究仍然较少,这使得它成为一个热门的研究课题,需要进一步探索。

蒸馏随机生成器

在本节中,我们将讨论针对随机生成器(也称为隐式生成模型)的蒸馏策略。这些模型在生成建模中被广泛使用,与确定性生成器不同的是,它们使用神经变换将潜在向量映射到具有随机性的数据样本。随机生成器在过去的十年中取得了巨大的成功 [54, 55, 56] ,其优点包括快速推断速度、低推断内存和轻量级模型。将扩散模型蒸馏成随机生成器的动机是为了实现极高的推断效率。

[57] 提出了Diff-Instruct,这是一个统一的框架,用于将预训练的扩散模型的知识蒸馏到通用生成器,包括随机生成器和确定性生成器。他们使用另一个由预训练扩散模型初始化的扩散模型,沿着扩散过程跟踪生成器输出分布的边际评分函数。具体而言,设 \bm{s}_{q^{(t)}} 表示预训练扩散模型, \bm{s}_{\phi}(\bm{x}_t,t) 表示用于生成器的学习扩散模型,通过梯度公式更新生成器的参数:

\operatorname{Grad}(\theta)=\int_{t=0}^T w(t)\mathbb{E}_{\bm{z}\sim p_z, \bm{x}_0 = g_\theta(\bm{z}),\atop \bm{x}_t|\bm{x}_0 \sim q_t(\bm{x}_t|\bm{x}_0)}\big[ \bm{s}_{\phi}(\bm{x}_t,t) - \bm{s}_{q^{(t)}}(\bm{x}_t)\big]\frac{\partial \bm{x}_t}{\partial\theta}\mathrm{d}t.

Diff-Instruct建立在坚实的理论基础上,最小化了他们提出的积分Kullback-Leibler散度,以将扩散过程纳入蒸馏方案中。Diff-Instruct不仅基于坚实的理论基础,而且在实证性能方面表现出色,在ImageNet64基准测试上实现了最先进的单步蒸馏,并具有更高的效率。

加速采样算法作为扩散模型蒸馏

扩散模型具有训练和采样过程分离的独特特点。在训练过程中,扩散模型不需要采样,在采样过程中,可以对其进行灵活的设计和改进。通常,扩散模型是在所有离散或连续的噪声水平上进行训练的,但是从中采样并不一定需要查询所有的扩散时间级别。最近的研究表明,使用一个较小的扩散时间级别子集可以显著加速采样过程,并且所需的非平衡推导次数更少。这些加速采样算法可以看作是扩散模型的广义蒸馏,其目标是训练一个大模型但使用一个小模型。我们可以将这些算法进一步分为无需训练和基于训练的加速算法,作为扩散蒸馏的两个独立的类别。

为了训练扩散模型,需要训练一个神经评分网络,以匹配数据扩散过程的边际评分函数。正如 [15] 所指出的,反向SDE 和 ODE 与正向扩散共享相同的边际分布,如果进行反向模拟,可以用于采样。因此,反向ODE或SDE被用作从扩散模型中进行采样的起点。具体而言,最简单的反向SDE模拟使用欧拉-马鲁雅马离散化方法,其公式为

X_{t+1} = X_t - \big[ F(X_t,t)-G^2(t) S_\theta(X_t,t) \big] \Delta t + G(t)\Delta t, t= T,...,1.

其中 S_\theta(x,t) 表示经过训练的评分网络,即教师模型。方程中的简单采样器通过模拟反向SDE来更新一批样本,按顺序考虑所有的扩散时间级别。然而,这种使用所有时间级别的采样策略在计算效率上存在问题。在接下来的章节中,我们将介绍基于训练和无需训练的加速采样算法。这些算法旨在通过利用较小的扩散时间级别子集来提高扩散模型的采样速度。

基于训练的加速算法

选择适当的扩散时间级别是加速采样算法的一个关键问题。

之前的研究已经表明,可以通过仅使用扩散时间级别的子集构建更高效的采样器,并且其性能与朴素采样器相当。 [58] 通过提出使用Fréchet Inception Distance(FID)度量作为目标来解决这个问题。他们学习了一个调度器模型,用于选择采样过程中要使用的步骤子集,从而提高了效率。

为了改善扩散模型的生成性能,[59] 提出了训练额外的协方差网络的方法。其思想是捕捉潜在空间分布的完整协方差矩阵,而不是像原始扩散模型那样假设对角协方差。通过这样做,模型能够更好地捕捉数据的复杂相关结构,并生成更逼真的样本。作者在各种数据集上取得了有 promising 的结果,并显示他们的方法在对角协方差的基线扩散模型上表现优越。

[60] 提出的方法涉及使用似然比估计模型(本质上是一个判别器)来学习真实评分函数和模型评分函数之间的差异。然后将这个差异与DM的评分函数结合起来,通过学习扩散过程的微分,得到一个无偏的评分函数。这样得到的改进型DM预计在较少的非平衡推导次数下能够提供更好的生成性能。

无需训练的加速算法

无需训练的算法旨在通过使用更少的扩散时间级别来实现与较多时间级别相当的生成性能。无需训练的加速算法的最显著部分是设计更快速的采样算法,无需训练新的参数化模型,只需要通过查询预训练的DM进行推理。这类算法的大部分内容围绕着生成性SDE或ODE的不同数值求解器展开。

首先,我们从去噪扩散概率模型(DDPM)[14] 开始。DDPM实现了VP扩散的离散化版本,并使用离散化的采样方案:

x_{t_{i-1}} = \frac{1}{\alpha_{t_i}}(x_{t_i}-\frac{1-\alpha_{t_i}}{\sqrt{1-\bar{\alpha}_{t_i}}}\epsilon\theta(x_{t_i},t_{i})) + \sigma_{t_i} z_i

其中 \alpha_{t_i} = 1-\beta_{t_i} \bar{\alpha}{t_i} = \prod_{j=1}^i \alpha_{t_i}\approx e^{-\int_{s=0}^{t_i} e^{\beta(s)}ds} 是积分的离散实现, \sigma_{t_i} 是任意的 \sigma 水平, z_i 是独立的高斯噪声。DDPM使用了所有的 1000 个扩散时间级别(即 {t_1,...,t_{1000}} )进行采样。

[9] 通过将DDPM的推导重新改为非马尔科夫正向扩散模型,得到了一系列新的采样器:

x_{t_{i-1}} = \sqrt{\alpha_{t_{i-1}}}\bigg[ \frac{x_{t_i} - \sqrt{1-\alpha_{t_i}}\epsilon_\theta(x_{t_i}, t_i) }{\sqrt{\alpha_{t_i}}} \bigg] + \sqrt{1-\alpha_{t_{i-1}} - \sigma_{t_i}^2}\cdot \epsilon_\theta(x_{t_i}, t_i) + \sigma_{t_i} z_{t_i}

其中 \alpha_{t_i} 的含义与 DDPM 采样算法中的相同。 \sigma_{t_{i}} 是一个自由的超参数,用于控制 DDIM 采样器中随机性的强度。当 \sigma=\sqrt{(1-\alpha_{t_{i-1}})/(1-\alpha_{t_i})}\sqrt{1-\alpha_{t_i}/\alpha_{t_{i-1}}} 时,DDIM 采样器与 DDPM 采样器一致。当 \sigma_{t_i} = 0 时,采样器是确定性的。DDIM 采样器需要与 DDPM 采样器相同的预训练DM,并且通过查询更少的扩散时间级别进行采样,已经显示出能够保持生成性能。

[10] 的工作表明,DDIM 采样器中存在一个最优反向方差 \sigma_{t_i} ,它最小化了 DDIM 正反向马尔可夫链之间的KL散度。此外,他们推导出了这个最优方差的显式表达式,形式如下:

\sigma_n^{*2} = \lambda_n^2 + \bigg[\sqrt{\frac{\bar{\beta}_n}{\alpha_n}} - \sqrt{\bar{\beta}_{n-1} - \lambda_n^2}\bigg]^2\cdot \bigg[1-\bar{\beta}_n \mathbb{E}_{q_n(x_n)}\frac{||\nabla_{x_n}\log q_n(x_n)||_2^2}{d} ) \bigg]

他们声称使用这种最优方差的 DDIM 采样器可以改进采样性能,而不论预训练扩散模型是什么。实际上,最优方差是用蒙特卡洛方法估计的,借助于预训练的DMs,形式如下:

\Gamma_n = \frac{1}{M} \sum_{m=1}^M \frac{||s_n(x_{n,m})||_2^2}{d},x_{n,m}\sim^{iid} q_n(x_n)

\hat{\sigma}^2_n = \lambda_n^2 + \bigg[ \sqrt{\frac{\bar{\beta}_n}{\alpha_n}} - \sqrt{\bar{\beta}_{n-1} - \lambda^2_n} \bigg]^2 (1-\bar{\beta}_n \Gamma_n).

[11]、[61] 和 [62] 发现了VP反向ODE的半线性结构:

\frac{dX_t}{dt} = F(t)X_t - \frac{1}{2} G^2(t)\nabla_{X_t} \log p^{(t)}(X_t)

其中 F(t) = -\frac{1}{2}\beta(t) G(t)=\sqrt{\beta(t)} 。对于VP扩散。他们在ODE上应用了指数积分器技术,进一步简化了VP反向ODE为:

x_t = \frac{\alpha_t}{\alpha_s}X_s - \alpha_t \int_s^t \frac{d\lambda_\tau}{d\tau} \frac{\sigma_\tau}{\alpha_\tau}\epsilon_\theta(X_\tau, \tau)d\tau

其中 \lambda_t\colon= \log (\alpha_t/\sigma_t) 是对数信噪比函数。[62]进一步提出了一种变量变换技巧和Taylor展开的简化算法,并获得了用于模拟反向ODE和SDE的高阶求解器。

[63] 提出了高阶SDE求解器作为EM离散化的替代方法,并展示了在FID方面的改进采样性能。在此基础上,[29] 建议对DMs进行更好的神经预调节,并使用二阶Heun离散化来模拟反向ODE和SDE。他们的方法在DM的生成性能方面取得了新的记录,FID得分为1.79。

其他研究工作还探索了使用先进数值求解器来加速DMs的采样过程,例如使用Runge-Kutta方法 [64] 或线性多步法 [65]。这些无需训练的加速算法在改进采样效率方面取得了很大的成功,使得使用较少的扩散时间级别也能够保持较高的生成性能。它们提供了一种有效的方式来减少计算成本,并提高训练和推理的速度。

总结

总的来说,扩散模型蒸馏技术在提高模型效率和效果方面显示出巨大的潜力。扩散到场模型蒸馏方法能够将扩散模型转化为计算效率高的场模型,从而实现更快速的推断和降低内存需求。扩散到生成器模型蒸馏方法可以将高容量扩散模型的知识转移给生成器模型,实现更高效的采样。加速算法旨在通过训练基方法学习选择部分扩散步骤或训练无需方法利用先进的数值求解器来减少采样所需的扩散步骤。

尽管这些技术已经显示出了良好的结果,但仍然存在一些挑战需要解决。改进扩散到场模型蒸馏方法的效率和效果,例如开发更好的架构和优化策略,可以进一步提高扩散模型在实际场景中的适用性。对于扩散到生成器模型蒸馏,探索新的蒸馏目标和技术可能会实现更有效的知识转移。此外,加速算法可以从数值求解器和优化技术的进展中受益,以实现更快速和更准确的采样。

总体而言,扩散模型蒸馏是一个活跃的研究领域,具有推动生成建模领域发展的巨大潜力。对这些技术的持续研究和创新将有助于扩散模型在各个领域的广泛应用和实际使用。


References

[1] Graves, Alex. "Generating sequences with recurrent neural networks." arXiv preprint arXiv:1308.0850 (2013).

[2] Van Den Oord, Aäron, Nal Kalchbrenner, and Koray Kavukcuoglu. "Pixel recurrent neural networks." International conference on machine learning. PMLR, 2016.

[3] Jozefowicz, Rafal, et al. "Exploring the limits of language modeling." arXiv preprint arXiv:1602.02410 (2016).

[4] Robert, Christian P., George Casella, and George Casella. Monte Carlo statistical methods. Vol. 2. New York: Springer, 1999.

[5] Hastings, W. Keith. "Monte Carlo sampling methods using Markov chains and their applications." (1970): 97-109.

[6] Roberts, Gareth O., and Jeffrey S. Rosenthal. "Optimal scaling of discrete approximations to Langevin diffusions." Journal of the Royal Statistical Society: Series B (Statistical Methodology) 60.1 (1998): 255-268.

[7] Xifara, Tatiana, et al. "Langevin diffusions and the Metropolis-adjusted Langevin algorithm." Statistics & Probability Letters 91 (2014): 14-19.

[8] Neal, Radford M. "MCMC using Hamiltonian dynamics." Handbook of markov chain monte carlo 2.11 (2011): 2.

[9] Song, Jiaming, Chenlin Meng, and Stefano Ermon. "Denoising diffusion implicit models." arXiv preprint arXiv:2010.02502 (2020).

[10] Bao, Fan, et al. "Analytic-dpm: an analytic estimate of the optimal reverse variance in diffusion probabilistic models." arXiv preprint arXiv:2201.06503 (2022).

[11] Liu, Luping, et al. "Pseudo numerical methods for diffusion models on manifolds." arXiv preprint arXiv:2202.09778 (2022).

[12] Zhao, Wenliang, et al. "UniPC: A Unified Predictor-Corrector Framework for Fast Sampling of Diffusion Models." arXiv preprint arXiv:2302.04867 (2023).

[13] Sohl-Dickstein, Jascha, et al. "Deep unsupervised learning using nonequilibrium thermodynamics." International Conference on Machine Learning. PMLR, 2015.

[14] Ho, Jonathan, Ajay Jain, and Pieter Abbeel. "Denoising diffusion probabilistic models." Advances in Neural Information Processing Systems 33 (2020): 6840-6851.

[15] Song, Yang, et al. "Score-based generative modeling through stochastic differential equations." arXiv preprint arXiv:2011.13456 (2020).

[16] Nichol, Alexander Quinn, and Prafulla Dhariwal. "Improved denoising diffusion probabilistic models." International Conference on Machine Learning. PMLR, 2021.

[17] Dhariwal, Prafulla, and Alexander Nichol. "Diffusion models beat gans on image synthesis." Advances in Neural Information Processing Systems 34 (2021): 8780-8794.

[18] Ramesh, Aditya, et al. "Hierarchical text-conditional image generation with clip latents." arXiv preprint arXiv:2204.06125 (2022).

[19] Saharia, Chitwan, et al. "Photorealistic text-to-image diffusion models with deep language understanding." Advances in Neural Information Processing Systems 35 (2022): 36479-36494.

[20] Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.

[21] Hoogeboom, Emiel, et al. "Equivariant diffusion for molecule generation in 3d." International Conference on Machine Learning. PMLR, 2022.

[22] Chen, Nanxin, et al. "Wavegrad: Estimating gradients for waveform generation." arXiv preprint arXiv:2009.00713 (2020).

[23] Baranchuk, Dmitry, et al. "Label-efficient semantic segmentation with diffusion models." arXiv preprint arXiv:2112.03126 (2021).

[24] Ho, Jonathan, et al. "Imagen video: High definition video generation with diffusion models." arXiv preprint arXiv:2210.02303 (2022).

[25] Molad, Eyal, et al. "Dreamix: Video diffusion models are general video editors." arXiv preprint arXiv:2302.01329 (2023).

[26] Poole, Ben, et al. "Dreamfusion: Text-to-3d using 2d diffusion." arXiv preprint arXiv:2209.14988 (2022).

[27] Heusel, Martin, et al. "Gans trained by a two time-scale update rule converge to a local nash equilibrium." Advances in neural information processing systems 30 (2017).

[28] Song, Yang, and Stefano Ermon. "Generative modeling by estimating gradients of the data distribution." Advances in neural information processing systems 32 (2019).

[29] Karras, Tero, et al. "Elucidating the design space of diffusion-based generative models." arXiv preprint arXiv:2206.00364 (2022).

[30] Hinton, Geoffrey, Oriol Vinyals, and Jeff Dean. "Distilling the knowledge in a neural network." arXiv preprint arXiv:1503.02531 (2015).

[31] Oord, Aaron, et al. "Parallel wavenet: Fast high-fidelity speech synthesis." International conference on machine learning. PMLR, 2018.

[32] Touvron, Hugo, et al. "Training data-efficient image transformers & distillation through attention." International conference on machine learning. PMLR, 2021.

[33] Luhman, Eric, and Troy Luhman. "Knowledge distillation in iterative generative models for improved sampling speed." arXiv preprint arXiv:2101.02388 (2021).

[34] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models." arXiv preprint arXiv:2202.00512 (2022).

[35] Meng, Chenlin, et al. "On distillation of guided diffusion models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

[36] Sun, Wujie, et al. "Accelerating Diffusion Sampling with Classifier-based Feature Distillation." arXiv preprint arXiv:2211.12039 (2022).

[37] Huang, Rongjie, et al. "Prodiff: Progressive fast diffusion model for high-quality text-to-speech." Proceedings of the 30th ACM International Conference on Multimedia. 2022.

[38] Song, Yang, et al. "Consistency models." arXiv preprint arXiv:2303.01469 (2023).

[39] Liu, Xingchao, Chengyue Gong, and Qiang Liu. "Flow straight and fast: Learning to generate and transfer data with rectified flow." arXiv preprint arXiv:2209.03003 (2022).

[40] Wu, Lemeng, et al. "Fast Point Cloud Generation with Straight Flows." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

[41] Lee, Sangyun, Beomsu Kim, and Jong Chul Ye. "Minimizing Trajectory Curvature of ODE-based Generative Models." arXiv preprint arXiv:2301.12003 (2023).

[42] Chen, Ricky TQ, and Yaron Lipman. "Riemannian flow matching on general geometries." arXiv preprint arXiv:2302.03660 (2023).

[43] Li, Lingxiao, Samuel Hurault, and Justin Solomon. "Self-consistent velocity matching of probability flows." arXiv preprint arXiv:2301.13737 (2023).

[44] Albergo, Michael S., and Eric Vanden-Eijnden. "Building normalizing flows with stochastic interpolants." arXiv preprint arXiv:2209.15571 (2022).

[45] Zheng, Hongkai, et al. "Fast Sampling of Diffusion Models via Operator Learning." arXiv preprint arXiv:2211.13449 (2022).

[46] Fan, Ying, and Kangwook Lee. "Optimizing ddpm sampling with shortcut fine-tuning." arXiv preprint arXiv:2301.13362 (2023).

[47] Aiello, Emanuele, Diego Valsesia, and Enrico Magli. "Fast Inference in Denoising Diffusion Models via MMD Finetuning." arXiv preprint arXiv:2301.07969 (2023).

[48] Mildenhall, Ben, et al. "Nerf: Representing scenes as neural radiance fields for view synthesis." Communications of the ACM 65.1 (2021): 99-106.

[49] Wang, Haochen, et al. "Score jacobian chaining: Lifting pretrained 2d diffusion models for 3d generation." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

[50] Lin, Chen-Hsuan, et al. "Magic3d: High-resolution text-to-3d content creation." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

[51] Deng, Congyue, et al. "Nerdi: Single-view nerf synthesis with language-guided diffusion as general image priors." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

[52] Xu, Dejia, et al. "NeuralLift-360: Lifting An In-the-wild 2D Photo to A 3D Object with 360° Views." arXiv e-prints (2022): arXiv-2211.

[53] Singer, Uriel, et al. "Text-to-4d dynamic scene generation." arXiv preprint arXiv:2301.11280 (2023).

[54] Goodfellow, Ian, et al. "Generative adversarial networks." Communications of the ACM 63.11 (2020): 139-144.

[55] Karras, Tero, et al. "Analyzing and improving the image quality of stylegan." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2020.

[56] Brock, Andrew, Jeff Donahue, and Karen Simonyan. "Large scale GAN training for high fidelity natural image synthesis." arXiv preprint arXiv:1809.11096 (2018).

[57] Luo, Weijian, et al. "Diff-Instruct: A Universal Approach for Transferring Knowledge From Pre-trained Diffusion Models." arXiv preprint arXiv:2305.18455 (2023).

[58] Watson, Daniel, et al. "Learning fast samplers for diffusion models by differentiating through sample quality." International Conference on Learning Representations. 2022.

[59] Bao, Fan, et al. "Estimating the optimal covariance with imperfect mean in diffusion probabilistic models." arXiv preprint arXiv:2206.07309 (2022).

[60] Kim, Dongjun, et al. "Refining Generative Process with Discriminator Guidance in Score-based Diffusion Models." arXiv preprint arXiv:2211.17091 (2022).

[61] Zhang, Qinsheng, and Yongxin Chen. "Fast sampling of diffusion models with exponential integrator." arXiv preprint arXiv:2204.13902 (2022).

[62] Lu, Cheng, et al. "Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps." arXiv preprint arXiv:2206.00927 (2022).

[63] Jolicoeur-Martineau, Alexia, et al. "Gotta go fast with score-based generative models." The Symbiosis of Deep Learning and Differential Equations. 2021.

[64] Wizadwongsa, Suttisak, and Supasorn Suwajanakorn. "Accelerating guided diffusion sampling with splitting numerical methods." arXiv preprint arXiv:2301.11558 (2023).

[65] Lu, Cheng, et al. "Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models." arXiv preprint arXiv:2211.01095 (2022).

编辑于 2023-06-07 19:06 ・IP 属地加拿大