生成扩散模型漫谈(二十五):基于恒等式的蒸馏(上)
By 苏剑林 | 2024-05-01 | 39735位读者 |今天我们分享一下论文《Score identity Distillation: Exponentially Fast Distillation of Pretrained Diffusion Models for One-Step Generation》,顾名思义,这是一篇探讨如何更快更好地蒸馏扩散模型的新论文。
即便没有做过蒸馏,大家应该也能猜到蒸馏的常规步骤:随机采样大量输入,然后用扩散模型生成相应结果作为输出,用这些输入输出作为训练数据对,来监督训练一个新模型。然而,众所周知作为教师的原始扩散模型通常需要多步(比如1000步)迭代才能生成高质量输出,所以且不论中间训练细节如何,该方案的一个显著缺点是生成训练数据太费时费力。此外,蒸馏之后的学生模型通常或多或少都有效果损失。
有没有方法能一次性解决这两个缺点呢?这就是上述论文试图要解决的问题。
重现江湖 #
论文将所提方案称为“Score identity Distillation(SiD)”,该名字取自它基于几个恒等式(Identity)来设计和推导了整个框架,取这个略显随意的名字大体是想突出恒等式变换在SiD中的关键作用,这确实是SiD的核心贡献。
至于SiD的训练思想,其实跟之前在《从去噪自编码器到生成模型》介绍过的论文《Learning Generative Models using Denoising Density Estimators》(简称“DDE”)几乎一模一样,甚至最终形式也有五六分相似。只不过当时扩散模型还未露头角,所以DDE是将其作为一种新的生成模型提出的,在当时反而显得非常小众。而在扩散模型流行的今天,它可以重新表述为一种扩散模型的蒸馏方法,因为它需要一个训练好的去噪自编码器——这正好是扩散模型的核心。
接下来笔者用自己的思路去介绍SiD。假设我们有一个在目标数据集训练好的教师扩散模型$\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t,t)$,它需要多步采样才能生成高质量图片,我们的目标则是要训练一个单步采样的学生模型$\boldsymbol{x} = \boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z})$,也就是一个类似GAN的生成器,输入指定噪声$\boldsymbol{z}$就可以直接生成符合要求的图像。如果我们有很多的$(\boldsymbol{z},\boldsymbol{x})$对,那么直接监督训练就可以了(当然损失函数和其他细节还需要进一步确定,读者可以自行参考相关工作),但如果没有呢?肯定不是不能训,因为就算没有$\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t,t)$也能训,比如GAN,所以关键是怎么借助已经训练好的扩散模型提供更好的信号。
SiD及前作DDE使用了一个看上去很绕但是也很聪明的思路:
如果$\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z})$产生的数据分布跟目标分布很相似,那么拿$\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z})$生成的数据集去训练一个扩散模型$\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t,t)$的话,它也应该跟$\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t,t)$很相似?
初级形式 #
这个思路的聪明之处在于,它绕开了对教师模型生成样本的需求,也不需要训练教师模型的真实样本,因为“拿$\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z})$生成的数据集去训练一个扩散模型”只需要学生模型$\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z})$生成的数据(简称“学生数据”),而$\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z})$是一个单步模型,用它来生成数据时间上比较友好。
当然,这还只是思路,将其转换为实际可行的训练方案还有一段路要走。首先回顾一下扩散模型,我们采用《生成扩散模型漫谈(三):DDPM = 贝叶斯 + 去噪》的形式,我们使用如下方式对输入$\boldsymbol{x}_0$进行加噪:
\begin{equation}\boldsymbol{x}_t = \bar{\alpha}_t\boldsymbol{x}_0 + \bar{\beta}_t\boldsymbol{\varepsilon},\quad \boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})\end{equation}
换言之$p(\boldsymbol{x}_t|\boldsymbol{x}_0)=\mathcal{N}(\boldsymbol{x}_t;\bar{\alpha}_t\boldsymbol{x}_0,\bar{\beta}_t^2 \boldsymbol{I})$。训练$\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t,t)$的方式则是去噪:
\begin{equation}\boldsymbol{\varphi}^* = \mathop{\text{argmin}}_{\boldsymbol{\varphi}} \mathbb{E}_{\boldsymbol{x}_0\sim \tilde{p}(\boldsymbol{x}_0),\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}}(\bar{\alpha}_t\boldsymbol{x}_0 + \bar{\beta}_t\boldsymbol{\varepsilon},t) - \boldsymbol{\varepsilon}\Vert^2\right] \label{eq:d-real-data}\end{equation}
这里的$\tilde{p}(\boldsymbol{x}_0)$就是教师模型的训练数据。同样地,如果我们想用$\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z})$的学生数据一个扩散模型,那么训练目标是
\begin{equation}\begin{aligned}
\boldsymbol{\psi}^* =&\, \mathop{\text{argmin}}_{\boldsymbol{\psi}} \mathbb{E}_{\boldsymbol{x}_0^{(g)}\sim p_{\boldsymbol{\theta}}(\boldsymbol{x}_0^{(g)}),\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\psi}}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\varepsilon}\Vert^2\right] \\
=&\, \mathop{\text{argmin}}_{\boldsymbol{\psi}} \mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\psi}}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\varepsilon}\Vert^2\right]
\end{aligned}\label{eq:dloss}\end{equation}
这里$\boldsymbol{x}_t^{(g)}=\bar{\alpha}_t\boldsymbol{x}_0^{(g)} + \bar{\beta}_t\boldsymbol{\varepsilon}=\bar{\alpha}_t\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z}) + \bar{\beta}_t\boldsymbol{\varepsilon}$,是由学生数据加噪后的样本,学生数据的分布记为$p_{\boldsymbol{\theta}}(\boldsymbol{x}_0^{(g)})$;第二个等号用到了“$\boldsymbol{x}_0^{(g)}$直接由$\boldsymbol{z}$决定”的事实,所以对$\boldsymbol{x}_0^{(g)}$的期望等价于对$\boldsymbol{z}$的期望。现在我们有两个扩散模型,它们之间的差异一定程度上衡量了教师模型和学生模型生成的数据分布差异,所以一个直观的想法是通过最小化它们之间的差异,来学习学生模型:
\begin{equation}
\boldsymbol{\theta}^* = \mathop{\text{argmin}}_{\boldsymbol{\theta}} \mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2\right]\label{eq:gloss-1}\end{equation}
注意式$\eqref{eq:dloss}$的优化依赖于$\boldsymbol{\theta}$,所以当$\boldsymbol{\theta}$通过式$\eqref{eq:gloss-1}$发生改变时,$\boldsymbol{\psi}^*$的值也随之改变,因此式$\eqref{eq:dloss}$和式$\eqref{eq:gloss-1}$实际上需要交替优化,类似GAN一样。
点睛之笔 #
谈到GAN,有读者可能会“闻之色变”,因为它是出了名的容易训崩。很遗憾,上述提出的式$\eqref{eq:dloss}$和式$\eqref{eq:gloss-1}$交替训练的方案同样有这个问题。首先它理论上是没有问题的,问题出现在理论与实践之间的gap,主要体现在两点:
1、理论上要求先求出式$\eqref{eq:dloss}$的最优解,然后才去优化式$\eqref{eq:gloss-1}$,但实际上从训练成本考虑,我们并没有将它训练到最优就去优化式$\eqref{eq:gloss-1}$了;
2、理论上$\boldsymbol{\psi}^*$随$\boldsymbol{\theta}$而变,即应该写成$\boldsymbol{\psi}^*(\boldsymbol{\theta})$,从而在优化式$\eqref{eq:gloss-1}$时应该多出一项$\boldsymbol{\psi}^*(\boldsymbol{\theta})$对$\boldsymbol{\theta}$的梯度,但实际上在优化式$\eqref{eq:gloss-1}$时我们都只当$\boldsymbol{\psi}^*$是常数。
这两个问题非常本质,它们也是GAN训练不稳定的根本原因,此前论文《Revisiting GANs by Best-Response Constraint: Perspective, Methodology, and Application》也特意从第2点出发改进了GAN的训练。看上去,这两个问题哪一个都无法解决,尤其是第1个,我们几乎不可能总是将$\boldsymbol{\psi}$求到最优,这在成本上是绝对无法接受的,至于第2个,在交替训练场景下我们也没什么好办法获得$\boldsymbol{\psi}^*(\boldsymbol{\theta})$的任何有效信息,从而更加不可能获得它关于$\boldsymbol{\theta}$的梯度。
幸运的是,对于上述扩散模型的蒸馏问题,SiD提出了一个有效缓解这两个问题的方案。SiD的想法可谓非常“朴素”:既然$\boldsymbol{\psi}^*$取近似值和$\boldsymbol{\psi}^*$当成常数都没法避免,那么唯一的办法就是通过恒等变换,尽量消除优化目标$\eqref{eq:gloss-1}$对$\boldsymbol{\psi}^*$的依赖了。只要式$\eqref{eq:gloss-1}$对$\boldsymbol{\psi}^*$的依赖足够弱,那么上述两个问题带来的负面影响也能足够弱了。
这就是SiD的核心贡献,也是让人拍案叫绝的“点睛之笔”。
恒等变换 #
接下来我们具体来看做了什么恒等变换。我们先来看式$\eqref{eq:d-real-data}$,它的优化目标可以等价地改写成
\begin{equation}\begin{aligned}
&\, \mathbb{E}_{\boldsymbol{x}_0\sim \tilde{p}(\boldsymbol{x}_0),\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}}(\bar{\alpha}_t\boldsymbol{x}_0 + \bar{\beta}_t\boldsymbol{\varepsilon},t) - \boldsymbol{\varepsilon}\Vert^2\right] \\
=&\, \mathbb{E}_{\boldsymbol{x}_0\sim \tilde{p}(\boldsymbol{x}_0),\boldsymbol{x}_t\sim p(\boldsymbol{x}_t|\boldsymbol{x}_0)}\left[\left\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}}(\boldsymbol{x}_t,t) - \frac{\boldsymbol{x}_t - \bar{\alpha}_t \boldsymbol{x}_0}{\bar{\beta}_t}\right\Vert^2\right] \\
=&\, \mathbb{E}_{\boldsymbol{x}_0\sim \tilde{p}(\boldsymbol{x}_0),\boldsymbol{x}_t\sim p(\boldsymbol{x}_t|\boldsymbol{x}_0)}\left[\left\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}}(\boldsymbol{x}_t,t) + \bar{\beta}_t\nabla_{\boldsymbol{x}_t}\log p(\boldsymbol{x}_t|\boldsymbol{x}_0)\right\Vert^2\right]
\end{aligned}\end{equation}
根据《生成扩散模型漫谈(五):一般框架之SDE篇》的得分匹配相关结果,上述目标的最优解是$\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t,t)=-\bar{\beta}_t\nabla_{\boldsymbol{x}_t}\log p(\boldsymbol{x}_t)$,同理式$\eqref{eq:dloss}$的最优解是$\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)=-\bar{\beta}_t\nabla_{\boldsymbol{x}_t^{(g)}}\log p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)})$。此时式$\eqref{eq:gloss-1}$的目标函数可以等价地改写成
\begin{equation}\begin{aligned}
&\,\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2\right] \\[5pt]
=&\, \mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\left\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t), \boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) + \bar{\beta}_t\nabla_{\boldsymbol{x}_t^{(g)}}\log p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)})\right\rangle\right] \\[5pt]
=&\, \color{green}{\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\left\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t), \boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t)\right\rangle\right]} \\[5pt]
&\,\qquad+ \color{red}{\mathbb{E}_{\boldsymbol{x}_t^{(g)}\sim p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)})}\left[\left\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t), \bar{\beta}_t\nabla_{\boldsymbol{x}_t^{(g)}}\log p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)})\right\rangle\right]}
\end{aligned}\end{equation}
接下来要用到在《生成扩散模型漫谈(十八):得分匹配 = 条件得分匹配》证明过的一个恒等式,来化简上式的红色部分:
\begin{equation}\nabla_{\boldsymbol{x}_t}\log p(\boldsymbol{x}_t) = \mathbb{E}_{\boldsymbol{x}_0\sim p(\boldsymbol{x}_0|\boldsymbol{x}_t)}\left[\nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{x}_t|\boldsymbol{x}_0)\right] \label{eq:id}\end{equation}
这是由概率密度定义以及贝叶斯公式推出的恒等式,不依赖于$p(\boldsymbol{x}_t),p(\boldsymbol{x}_t|\boldsymbol{x}_0),p(\boldsymbol{x}_0|\boldsymbol{x}_t)$的形式。将该恒等式代入到红色部分,我们有
\begin{equation}\color{red}{\begin{aligned}
&\,\mathbb{E}_{\boldsymbol{x}_t^{(g)}\sim p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)})}\left[\left\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t), \bar{\beta}_t\nabla_{\boldsymbol{x}_t^{(g)}}\log p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)})\right\rangle\right] \\[5pt]
= &\,\mathbb{E}_{\boldsymbol{x}_t^{(g)}\sim p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)}),\boldsymbol{x}_0^{(g)}\sim p_{\boldsymbol{\theta}}(\boldsymbol{x}_0^{(g)}|\boldsymbol{x}_t^{(g)})}\left[\left\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t), \bar{\beta}_t\nabla_{\boldsymbol{x}_t^{(g)}} \log p(\boldsymbol{x}_t^{(g)}|\boldsymbol{x}_0^{(g)})\right\rangle\right] \\[5pt]
= &\,-\mathbb{E}_{\boldsymbol{x}_0^{(g)}\sim p_{\boldsymbol{\theta}}(\boldsymbol{x}_0^{(g)}),\boldsymbol{x}_t^{(g)}\sim p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)}|\boldsymbol{x}_0^{(g)})}\left[\left\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t), \frac{\boldsymbol{x}_t - \bar{\alpha}_t \boldsymbol{x}_0}{\bar{\beta}_t}\right\rangle\right] \\[5pt]
= &\,-\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\left\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t), \boldsymbol{\varepsilon}\right\rangle\right]
\end{aligned}}\end{equation}
跟绿色部分合并,就得到学生模型新的损失函数
\begin{equation}\begin{aligned}
\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\left\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t), \boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\varepsilon}\right\rangle\right]
\end{aligned}\label{eq:gloss-2}\end{equation}
这就是SiD的核心结果,原论文的实验结果显示它能够高效地实现蒸馏,而式$\eqref{eq:gloss-1}$则没有训练出有意义的结果。
相比式$\eqref{eq:gloss-1}$,上式$\eqref{eq:gloss-2}$出现$\boldsymbol{\psi}^*$的次数显然更少,也就是对$\boldsymbol{\psi}^*$的依赖更弱。此外,上式是基于最优解$\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)=-\bar{\beta}_t\nabla_{\boldsymbol{x}_t^{(g)}}\log p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)})$恒等变换而来的,也就是说相当于(部分地)预先窥见了$\boldsymbol{\psi}^*$的精确值,这也是它更优越的原因之一
其他细节 #
到目前为止,本文的推导基本上是原论文推导的重复,但出了个别记号上的不一致外,还有一些细节上的不同,下面简单澄清一下,以免读者混淆。
首先,论文的推导默认了$\bar{\alpha}_t=1$,这是沿用了《Elucidating the Design Space of Diffusion-Based Generative Models》一文的设置。然而尽管$\bar{\alpha}_t=1$很有代表性,并且能简化形式,但并不能很好地覆盖所有扩散模型类型,所以本文的推导保留了$\bar{\alpha}_t$。其次,论文的结果是以$\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t) = \frac{\boldsymbol{x}_t - \bar{\beta}_t \boldsymbol{\epsilon}(\boldsymbol{x}_t,t)}{\bar{\alpha}_t}$为标准给出的,这显然跟扩散模型常见的以$\boldsymbol{\epsilon}(\boldsymbol{x}_t,t)$为准不符,笔者暂时没有领悟到原论文的表述方式的优越所在。
最后,原论文发现损失函数$\eqref{eq:gloss-1}$实在太不稳定,往往对效果还起到负面作用,所以SiD最终取了式$\eqref{eq:gloss-1}$的相反数作为额外的损失函数,加权到改进的损失函数$\eqref{eq:gloss-2}$上,这在个别情形还能取得更优的蒸馏效果。至于具体实验细节和数据,读者自行翻阅原论文就好。
相比其他蒸馏方法,SiD的缺点是对显存的需求比较大,因为它同时要维护三个模型$\boldsymbol{\epsilon}_{\boldsymbol{\varphi}}(\boldsymbol{x}_t,t)$、$\boldsymbol{\epsilon}_{\boldsymbol{\psi}}(\boldsymbol{x}_t,t)$和$\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z})$,它们具有相同的体量,虽然并非同时进行反向传播,但叠加起来也使得总显存量翻了一倍左右。针对这个问题,SiD在正文末尾提出,未来可以尝试对预训练的模型加LoRA来作为额外引入的两个模型,以进一步节省显存需求。
延伸思考 #
笔者相信,对于一开始的“初级形式”,即式$\eqref{eq:dloss}$和式$\eqref{eq:gloss-1}$的交替优化,那么不少理论基础比较扎实并且深入思考过的读者都有机会想到,尤其是已经有DDE“珠玉在前”,推出它似乎并不是那么难预估的事情。但SiD的精彩之处是并没有止步于此,而是提出了后面的恒等变换,使得训练更加稳定高效,这体现了作者对扩散模型和优化理论非常深刻的理解。
同时,SiD也留下了不少值得进一步思考和探索的问题。比如,学生模型的损失$\eqref{eq:gloss-2}$的恒等化简到了尽头了吗?并没有,因为它的内积左边还有$\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)$,还可以用同样的方式进行化简。具体来说,我们有
\begin{equation}\begin{aligned}
&\,\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2\right] \\[5pt]
=&\,\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2 - 2\left\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t),\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\right\rangle + \Vert\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2\right] \\[5pt]
=&\,\mathbb{E}_{\boldsymbol{x}_t^{(g)}\sim p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)})}\left[
\scriptsize{\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2 - 2\left\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t),-\bar{\beta}_t\nabla_{\boldsymbol{x}_t^{(g)}}\log p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)})\right\rangle + \left\langle\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t),-\bar{\beta}_t\nabla_{\boldsymbol{x}_t^{(g)}}\log p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)})\right\rangle}
\right] \\[5pt]
\end{aligned}\end{equation}
这里的每一个$-\bar{\beta}_t\nabla_{\boldsymbol{x}_t^{(g)}}\log p_{\boldsymbol{\theta}}(\boldsymbol{x}_t^{(g)})$都可以用相同的恒等变换$\eqref{eq:id}$最终转化为单个$\boldsymbol{\varepsilon}$(但要注意$\Vert\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2=\langle\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t),\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\rangle$只能转换一个,不能都转),而式$\eqref{eq:gloss-2}$相当于只转了一部分,如果全部转会更好吗?因为没有实验结果,所以暂时不得而知。但有一个特别有意思的形式,就是只转换上面的中间部分的话,该损失函数可以写成
\begin{equation}\begin{aligned}
&\,\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2\right] \\[5pt]
=&\,\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2 - 2\left\langle\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t),\boldsymbol{\varepsilon}\right\rangle + \Vert\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2\right] \\[5pt]
=&\,\mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\varepsilon}\Vert^2 + \Vert\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}(\boldsymbol{x}_t^{(g)},t)\Vert^2\right] + \text{常数} \\[5pt]
\end{aligned}\label{eq:gloss-3}\end{equation}
这是学生模型,也就是生成器的损失,然后我们再对比学生数据去噪模型的损失$\eqref{eq:dloss}$:
\begin{equation}\boldsymbol{\psi}^* = \mathop{\text{argmin}}_{\boldsymbol{\psi}} \mathbb{E}_{\boldsymbol{z},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\Vert\boldsymbol{\epsilon}_{\boldsymbol{\psi}}(\boldsymbol{x}_t^{(g)},t) - \boldsymbol{\varepsilon}\Vert^2\right]\label{eq:dloss-1}\end{equation}
这两个式子联合起来看,我们可以发现学生模型实则在向教师模型看齐,并且试图远离学生数据所训练的去噪模型,形式上很像LSGAN,$\boldsymbol{\epsilon}_{\boldsymbol{\psi}}(\boldsymbol{x}_t^{(g)},t)$类似GAN的判别器,不同的地方是,GAN的判别器一般是两项损失相加而生成器是单项损失,SiD则反过来了。这其实体现了两种不同的学习思路:
1、GAN:一开始造假者(生成器)和鉴别者(判别器)都是小白,鉴别者不断对比真品和赝品来提供自己的鉴宝水平,造假者则通过鉴别者的反馈不断提高自己的造假水平;
2、SiD:完全没有真品,但有一个绝对权威的鉴宝大师(教师模型),造假者(学生模型)不断制作赝品,同时培养自己的鉴别者(学生数据训练的去噪模型),然后通过自家鉴别者跟大师的交流来提高自己造假水平。
可能有读者会问:为什么SiD中的造假者不直接向大师请教,而是要通过培养自己的鉴别者来间接获得反馈呢?这是因为直接跟大师交流的话,可能会出现的问题就是长期都只交流同一个作品的技术,最终只制造出了一种能够以假乱真的赝品(模式坍缩),而通过培养自己的鉴别者一定程度上就可以避免这个问题,因为造假者的学习策略是“多得到大师的好评,同时尽量减少自家人的好评”,如果造假者还是只制造一种赝品,那么大师和自家的好评都会越来越多,这不符合造假者的学习策略,从而迫使造假者不断开发新的产品而不是固步自封。
此外,读者可以发现,SiD整个训练并没有利用到扩散模型的递归采样的任何信息,换句话说它纯粹是利用了去噪这一训练方式所训练出来的去噪模型,那么一个自然的问题是:如果单纯为了训练一个单步的生成模型,而不是作为已有扩散模型的蒸馏,那么我们训练一个只具有单一噪声强度的去噪模型会不会更好?比如像DDE一样,固定$\bar{\alpha}_t=1$、$\bar{\beta}_t=\beta=\text{某个常数}$取训练一个去噪模型,然后用它来重复SiD的训练过程,这样会不会能够简化训练难度、提高训练效率?这也是一个值得进一步确认的问题。
文章小结 #
在这篇文章中,我们介绍了一种新的将扩散模型蒸馏为单步生成模型的方案,其思想可以追溯到前两年的利用去噪自编码器训练生成模型的工作,它不需要获得教师模型的真实训练集,也不需要迭代教师模型来生成样本对,而引入了类似GAN的交替训练,同时提出了关键的恒等变换来稳定训练过程,整个方法有颇多值得学习之处。
转载到请包括本文地址:https://kexue.fm/archives/10085
更详细的转载事宜请参考:《科学空间FAQ》
如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。
如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!
如果您需要引用本文,请参考:
苏剑林. (May. 01, 2024). 《生成扩散模型漫谈(二十五):基于恒等式的蒸馏(上) 》[Blog post]. Retrieved from https://kexue.fm/archives/10085
@online{kexuefm-10085,
title={生成扩散模型漫谈(二十五):基于恒等式的蒸馏(上)},
author={苏剑林},
year={2024},
month={May},
url={\url{https://kexue.fm/archives/10085}},
}
May 4th, 2024
请教一下苏神,式(4)为什么两个扩散模型的输入都是$(x_t^{(g)},t)$呢?第二个扩散模型$\epsilon_{\phi*}()$ ,我理解为是用来对$x_t^{(g)}$数据去噪的模型,第一个扩散模型$\epsilon_{\varphi*}()$ 我理解为是是对原始训练数据加噪后再去噪的模型。但是第一个扩散模型我没想清楚,$\epsilon_{\varphi*}()$ 既然针对原始数据加噪后的数据去噪的,那式(4)不应该传入的是$x_t$嘛?为什么还是$(x_t^{(g)},t)$呀?
现在要配对来训练学生模型啊,你让$\boldsymbol{\epsilon}_{\boldsymbol{\varphi}^*}$输入$\boldsymbol{x}_t$、$\boldsymbol{\epsilon}_{\boldsymbol{\psi}^*}$输入$\boldsymbol{x}_t^{(g)}$的话,它们怎么匹配呢...
May 23rd, 2024
特地去看了原文,感觉没太get到作者的意图。
直到“点睛之笔”前的$L^{(1)}$好像就是ProlificDreamer的VSD loss,作者故意用EDM举例来让两者公式有所不同,还特地引入了$\delta$这个符号。
中间拆分等价的“点睛之笔”看上去很妙,但是最后导出的损失函数$L^{(2)}$作者又说效果不好,然后做了一个$L^{(1)}$和$L^{(2)}$的插值,最后$\alpha=1$的特殊情况好像和作者讲的减少对$\psi$的依赖没有什么关系。
然后原文的式(24)看上去好像意思是说计算了$\phi$和$\psi$相对于各自输入的梯度,但其实在实际应用时基于训练开销的考虑一般直接把这部分梯度给忽略了。因为有条件生成的teacher是需要做CFG的,所以$\phi$要前向两次,在加上$\phi$前向一次,那么loss部分的梯度计算开销就是student本身的三倍了。
针对Kiro的评论,有几个关键点值得指出:
第一,SiD的code是在EDM的codebase上开发的,所以所有的推导也是基于$x_0$ prediction来写的。$x_0$ prediction 和 $\epsilon$ prediction 两者在理论上是等价的,在实际上也可以通过调整网络的输入和输出的preconditioning和loss的weighting消除实现层面的差别。这个点在EDM文章里面有详细的描述。
第二,$L^{(1)}$ 的梯度是 $\nabla_{\theta}L^{(1)}=2\sigma_t^{-2}(f_{\phi}(x_t,t)-
f_{\psi}(x_t,t))\color{red}{(\nabla_{x_t}f_{\phi}(x_t,t)- \nabla_{x_t}f_{\psi}(x_t,t))}\nabla_{\theta}x_t$, 这里的标红的部分是VSD loss的梯度里面不存在的。不管写成SiD里面$x_0$-prediction的形式,还是按照剑林的$\epsilon$的形式写,这个区别都显著存在。如果想进一步了解这一点,可以参考SiD文章里面关于与weijian的Diff-Instruct的区别和联系的相关讨论。
第三,$L^{(2)}$ 单独是能work的:Inception Score和主观视觉效果(Figure 7的第二列)都还不错,只是FID不怎么好。SiD最后推荐的loss是$L^{(2)}\color{red}{-}\alpha L^{(1)}$, $\alpha=1$ or $1.2$。注意这里$\alpha$前面有负号,所以和普通插值是有非常大的区别的。
第四,SiD不允许忽略那两个梯度的计算。对比最近的几个蒸馏的工作(比如Diff Instruct/SwiftBrush,DMD,ADD),其实SiD所需要的网络数量是一致的。大家都至少需要三个网络,teacher,student,和 fake score network/Discriminator for student。SiD的主要额外代价是需要计算$\nabla_{x_t}f_{\phi}(x_t,t)$ 和 $\nabla_{x_t}f_{\psi}(x_t,t)$。这个在imagenet 64x64实测出来会增加50%的内存消耗,增加36%的time per iteration。这些细节文章里面都有对应的描述。
第五,SiD的效果是超群的。因为它是data free,也没有需要额外加regression或者adversarial losses,所以实现起来也很直接。我们最近会release code和model checkpoints,欢迎大家使用。
SiD目前在做class conditional生成的时候是没有加CFG的。
我们的后续工作有用SiD来做Stable Diffusion的Distillation。这里CFG会非常关键。合适的使用CFG后,SiD做文字到高清图的生成效果非常棒。这个我们会最近放ArXiv。
期待你们的新论文。
额外一个问题:
因为我看论文中似乎没有标红部分是否保留的消融实验,想问一下这部分影响是不是很大?因为我今天在SD1.5上复现了一下($\alpha=1$),没有标红部分也是可以训练的,但是效果和VSD没有什么差别。但是增加这部分梯度后,训练效率变低了很多,所以暂时还给不出结论。
如果把这个部分stop gradient,SiD会reduce到Diff-instruct(和weijian的实现细节有些出入,但是理论上应该没本质区别)。Diff-instruct效果也挺不错,也是data free distillation,但是我们观察它的FID降到一定程度就下不去了,而SiD可以继续降。这个现象也许和KL loss本身的局限性应该有关,相关的讨论可以参考DMD(本质是Dif-instruct+regression loss)那个文章。另外我们在cifar10上观察到diff instruct训练到一定程度后可能会很快变差。有一种改进速度的可能性是先把那个gradient去掉不加,训练到一定程度在加回来。但是这个我们实验还没试。
我的实验显示score distillation的FID高主要是因为多样性差,而DMD加reg loss其实是在通过配对来增加多样性,而不是输出模式坍塌到单个峰,但是似乎你们的$L^{(2)}$训练出来的模型IS也是高的,似乎不是多样性问题。
总之非常感谢您的回复,我现在基本理清了思路,之前一直把忽略$\phi$和$\psi$的梯度当作是单纯的节约训练开销理解。
期待你们的开源和新论文。
很高兴能解答你的疑问。我们目前没有观察到SiD有mode collapse的问题(我们的IS,precision,recall指标都挺好,视觉效果也挺好),应该和明确避免用KL作为loss有关系。
已经进行了实践,当前来看训练稳定性很好。
其实我之前实践时发现SDXL的蒸馏稳定性要远比SD 1.5差,可能是Transformer层的数量或者图像分辨率导致。不知你们有没有进行相关尝试,看论文还是依旧以SD 1.5和SD 2.1为主。
算法框的第10行是不是写错了,好像还是写成了EDM形式。
Kiro: 很高兴你通过实践验证了SiD的稳定性。
SDXL的distill目前实验还没具体开展。等计算资源准备好后,我们希望把这方面实验继续做起来。
SiD-LSG算法的10行的确有typo,正确的应该是$x_t = a_t x_0+\sigma_t \epsilon_t$
我们也遇到多样性损失的问题了,SD 1.5 的 base 和 SiD 蒸馏的 checkpoint 似乎确实没有这个问题,但是我们蒸馏后出现了这个问题。
请问大佬现在有什么解决思路么。
你观察的mode collapse发生在什么情况下?用SiD蒸馏EDM/DDPM的pretrained的models,用SiD-LSG蒸馏stable diffusion我们都没有发现问题。
周老师您好,我们是在内部训练的一个 SD 1.5 的模型上蒸馏的。
我们对 400 个提示词都推理了 4 张图,SD 1.5 的 lpips 为 0.7710,SD 1.5 SiD-LSG 的 lpips 为 0.7232,视觉效果看上去多样性确实也没有影响。
但是在我们内部的模型上,lpips 骤降到了 0.3739,视角效果上多样性下降也很明显。
我们推测可能是由于训练的内部数据集不均衡导致的,因为出现了一些 case,例如 prompt 是“男客服”,但是推理出来的图片主体却是女生,而且四张非常相似。目前正在进行验证实验。
我们把SiD用到Stable Diffusion的蒸馏的文章放到ArXiv了:
https://arxiv.org/abs/2406.01561
我们目前zero shot FID可以做到8.15, 在已经发表或ArXiv的文章里面应该是最低的了。我们还在继续train,按照目前趋势看来应该还能更低
厉害,恭喜!
我看到新文章的介绍中采用了跟本文一样以$\epsilon$为准的表达形式,看起来确实舒服多了。顺便问一下两者有效果上的差异吗?
看起来SiD具备GAN的灵活性,训练起来又比GAN稳定,确实难能可贵。
谢谢 剑林!因为Stable Diffusion默认的是$\epsilon$ prediction, 所以我们实现的时候用的网络的backbone是基于$\epsilon$ prediction的:
noise_pred = unet(latents, init_timesteps, encoder_hidden_states=text_embeddings).sample
如果需要用基于$x_0$ prediction,那么调用diffuser里面的函数就行:
x0_pred = noise_scheduler.step(noise_pred, init_timesteps[0], latents,return_dict=True).pred_original_sample
我按照您的思路去理解了一下,其实关键就在于标红部分,这部分在$\psi$未完全收敛时是会导致平凡解的。我试了一下,确实生成器学习率稍大一点就直接输出黑色图像了,但是只要把$\psi$的梯度从输出复制到输入就可以稳定训练了,看来这个假设应该是成立的。
不过我还是没搞懂这个加权的原理,是单纯的实验得出结论么?
如果选取KL loss,那么梯度部分推导出来就不需要对true/fake score再求梯度。只需要它们的差乘上Generator的梯度就可以。
SiD选取的是把model-based explicit score matching loss (Fisher divergence)作为理论优化目标,然后通过3个score identities来找出合适的近似。最终使用的loss是结合理论推导和实践检验提出来的。我觉得可能有更好的逼近理论目标的loss,但是目前我们还没找到。
很有收获,向您学习!
既然作者@周名远|comment-24411亲自光临寒舍来指点迷津,我就不画蛇添足了。向两位学习。
May 23rd, 2024
流程上和这个好像,https://arxiv.org/abs/2311.18828
One-step Diffusion with Distribution Matching Distillation
所有generative models其实本质都是distribution matching。
GANs 是通过最小化Jensen-Shannon divergence(JSD) 来匹配数据分布和模型分布。
Diffusion GANs是通过JSD来匹配带噪的模型分布和带同样强度的噪音的数据分布。
Diff-Instruct可以理解为通过KL divergence来匹配带噪的teacher分布和带同样强度的噪音的student分布。
DMD这个文章可以理解为Diff-Instruct + 额外的Regresion loss。
Diff-Instruct 和 DMD都可以理解为通过对 $\mathbb{E}_{x_t\sim q(x_t|x_g),x_g\sim p_{\theta}(x_g)}KL(q(x_t)||p(x_t))$ 求梯度来更新generator。
SiD和这两个非常棒的工作的核心区别之一是它即不是用JSD也不是用KL,而是用我们提出的model-based explicit score matching loss (也是一种Fisher divergence)来匹配带噪的teacher分布和带噪的student的分布。
June 22nd, 2024
蒸馏Stable Diffusion的SiD-LSG算法和模型checkpoints已经上线了:
https://github.com/mingyuanzhou/SiD-LSG
June 23rd, 2024
苏神,最近在拜读您合作的一篇关于扩散模型暴露偏差的论文《ELUCIDATING THE EXPOSURE BIAS IN DIFFUSION MODELS》(ICLR2024),我们知道,扩散模型神经网络的本质是拟合“训练过程中x0到xt阶段添加的实际标准高斯噪声,网络接收xt和t预测输出一个噪声(应该逼近标准高斯),理论上来讲,对于一个训练好的扩散模型神经网络(一般是Unet),无论时间步t,网络预测输出的噪声应该尽可能逼近标准高斯,其二范数应该保持在一个恒定的水平,为什么在本文中如Figure2、Figure5和Figure6等展示的那样,网络预测输出的噪声二范数随着时间步t的减小有一个明显的下降趋势?https://arxiv.org/abs/2308.15321
June 25th, 2024
论文中实际采用了$L^\Delta$,鄙人没看出来这有什么道理,$L^1$效果再差也不至于直接反号训吧。不知道苏神有什么见解
有一定理由怀疑真正work的部分是diff-inst, $-L^1$作为正则,鼓励生成score差异较大的样本,使得生成器更多样化了。论文数据中,在训练初期会观察到一次FID激增的现象,可能是这次激增产生了退火的效果
不能这样理解。比如我们有两个损失$L_a, L_b$,$L_a$更优,$L_b$稍差,那么$L_a - \lambda L_b$为Loss其实也不能说明反优化$L_b$,而是应该理解为
$$L_a - \lambda L_b = (1 - \lambda) L_a - \lambda(L_b - L_a)$$
真正负优化的是$L_b$与$L_a$的差,由于$L_a$比$L_b$更优,所以它们的差意味着是“不好的部份”,需要负优化。
最近读了一下Adversarial Score Distillation, 虽然理论还是有一些粗糙的地方,但我认为比SiD更合理一些,-L1基本上充当了WGAN中的梯度惩罚项
剑林的博客非常清楚地指出了SiD的核心思想之一:降低用$\psi$来近似$\psi^*(\theta)$对理论目标函数所带来的偏差。在这个近似下,SiD在Proposition 4给出了明确的例子来证明$L^1$有可能会失败。
至于为什么$L^2-L^1$在实际中会明显优于只用$L^2$,这个还有待在理论上继续研究。我个人猜测这么做是补偿了忽略$\psi^*(\theta)$这一项和$\theta$的相关性(所以相关的梯度也忽略了)所带来的偏差。但是这个假想目前还很难通过理论推导来证明。
另外我想强调SiD是data free:它不需要使用真实的训练数据。而你提到的adversarial score distillation是典型的需要利用真实训练数据的算法,所以它本质上和SiD的loss并不冲突。
我理解论文表达的思路,也肯定实验结果。然而最优的设置下$\alpha$高达1.2,依照论文的逻辑这等价于以-0.2L1为目标,这显然与论文的核心score identity背道而驰。其次我认为score identity并不是一个合理的选择,score identity并不能引到生成器向数据高密度区域移动,因为直观上越低密度score越接近正太分布score,远离中心的区域是两个分布score最接近的区域,而介于两个分布中心的区域是score相差最大的区域,-L1反而更具合理性也就是您指的$\psi^*(\theta)$与$\theta$的相关性,基于这样的考虑-L1能够鼓励样本的多样性(可以想象模式塌缩时缺失模式区域score差是较大的,-L1将引导生成器填补这些缺失的模式)。而L2的梯度中包含diff-inst逆KL的梯度。我认为两个消融实验能更说明问题:一, 在使用-1.2L1的情况下,用diff-inst替换L2 二,充分训练$\psi^*(\theta)$验证score identity的合理性
cifar10上的实验表明SiD在$\alpha=1$附近效果都不错。我们并没有对这个参量做仔细的grid search,所以也不清楚$\alpha=1.2$是否是最优选择。
L1是对目标函数最肤浅的有偏估计,L2是对由score identity推导出的等价目标函数的有偏估计。L2和L1都是理想目标函数的有偏估计,所以L2-1.2L1并不等价于用-0.2L1为目标。
你的分析思考挺有意思,但是你提出的观点和想法我找不到相关理论和实验的支持。我鼓励你试试SiD的开源的code,说不定能通过实践来找出比SiD更好更合理的算法。你说的第一个实验你可以试试,如果性能上能够提升Diff-Instruct甚至媲美SiD的话,我会很乐意进一步探讨。你说的第二个实验感觉挺难,因为我们忽略的是 $\nabla_{\theta}\psi^*(\theta)$,这个近似对理论目标函数的梯度的偏差的影响和$\psi$是否接近$\psi^*(\theta)$的关系好像不是很明确。
尝试L1,L2加权是自然的思路,然而尝试L2-$\alpha$L1是匪夷所思的。我很好奇您的团队是如何完成这一发现的。这非常有趣,不知能否恳请您分享这一过程
后接 @苏剑林|comment-24796
July 8th, 2024
想请教一个问题,文中提到“造假者的学习策略是“多得到大师的好评,同时尽量减少自家人的好评””。
代码中的损失函数:loss = (y_real-y_fake)*((y_real-y)-alpha*(y_real-y_fake))/weight_factor,然后我按照这个损失函数给了几个例子。
1. 当 y_real 和 y_fake 都为1时,且 y 也为1,表示预测准确,此时损失 loss 为 0,不需要进一步优化;
2. 当 y_real 为1,但 y_fake 为3,尽管 y 也是1,即当 y 和 y_real 相同但跟 y_fake 不同时,loss 相对较大;
3. 当 y_real 为2, y_fake 为1,而 y 也是1时,即当 y 和 y_fake 相同但跟 y_real 不同时,loss 又相对较小,这岂不是希望 y 多跟 y_fake 靠近,也就是造假者多得到自家人的好评?
我没看过开源代码,看上去你这里指的是原论文的损失,如果是的话,原论文的损失是可正可负的,并且跟本文的$(11)$是不完全等价的,因此不能直接沿用$(11)$的解释。要判断往哪个方向靠近,单纯从原论文的损失可能做不到,可能要比较对应的梯度才行。
好的,谢谢苏老师,还想再请教两个问题。
1. 就是学生模型是一步推理的,老师模型和 fake score network 也是一步推理的,这块会不会有什么问题呢,毕竟老师模型推理一步的效果并不好。
2. fake score network 学习的是什么呢?fake score network 跟学生模型是交替学习的,首先 fake score network 训练的时候的 loss 是噪声预测的 MSE,这个好理解,但是学生模型训练的时候损失,不是很理解有什么直观的作用,是希望学生模型生成的内容分别跟老师模型和 fake score network 拉进么,但是如 1 所述,老师模型推理一步的效果并不好呀。
1. 教师模型提供的是带噪图像的score,并不是直接用于生成的。
2. fake score network学习的是由一步生成器(学生模型)生成的假数据的score。
3. fake score network 和教师模型联合起来指导生成器如何去优化
周老师您好,我来自百度,您的 SiD-LSG 在我们的场景下蒸馏的效果很好,非常感谢您的工作。
我们现在遇到了一些代码上的小问题,推理的时候,您在 generate_onestep 中调用的是自己实现的 sid_sd_sampler,我研究了一下,应该是跟 diffusers 的 pipeline 的 call 方法类似的。
但是,我在将 network_pkl 文件读取出来,然后再赋值给 pipeline 的 unet 时,推理出来的结果跟使用 generate_onestep 推理出来的结果并不一致。
我想问的是,这是因为 SiD-LSG 自定义的那个 sid_sd_sampler 的原因么?您是否有尝试过直接用 diffusers 的 pipeline 进行推理呢?
不好意思,刚读到您的消息!我们自己写的这个sid_sd_sampler和diffusers的pipeline没有本质区别。没调用那个pipeline主要还是调试code方便。如果您使用pipeline的话应该需要把初始的time设成625,而不是默认的999。有时间的时候我会去研究一下把我们的checkpoint存成safetensor并让他和pipeline匹配,当然也欢迎您提前弄好的话帮忙push到SiD的github
1、学生模型作为生成器来说是一步推理的,但教师模型和fake score network只是起到去噪器的作用,理论上都不需要DDPM等采样过程,所以这里并没有问题;
2、fake score network相当于用学生模型的生成数据集训练了一个扩散模型(去噪模型),通过跟教师模型(真实数据训练出来的去噪模型)对比,得出学生模型的生成结果与真实数据的差异程度。
好的,谢谢苏老师,结合您的文章推理了一遍公式,清晰了很多。
July 12th, 2024
帮忙新开一楼,接着@梁宇辰|comment-24779:
权重为负的“异想天开”,对于一般人来说确实“匪夷所思”,有什么头脑风暴过程可以分享学习的就更好了。@周名远|comment-24701
其实也没有什么特别的诀窍。我个人直觉认为,可以结合score matching loss和我们以前发表的diffusion GAN中多噪声尺度下进行分布匹配的思路来实现蒸馏。在这个信念的指引下,我们几乎动用了所有本地GPU来调试这个loss,包括3070、3090、A4000、A5000和A6000,总共大约30块GPU,花了一个月左右的时间,几乎在圣诞和元旦都没有休息。第一次调试成功的是alpha=1的loss,我们发现它收敛特别快且稳定。随后我们继续对这个loss的细节进行了修改,然后从多个不同角度尝试解释它为何如此有效。我个人认为,SiD效果好的最核心的原因可能是某种意义上引入了$\psi^*(\theta)$关于$\theta$的梯度,而这一点在以前的算法中是不存在的。
个人认为其他都还好,毕竟有公式支撑,主要是负权重这一点,比较反直觉,可能大家主要困惑的是这一点的灵感来源。
看起来是先有的实验结果,后补充的理论是针对$L^2$的,但$L^2$效果不是很好。我个人是从内外圈优化梯度对齐的角度理解的,Meta Pseudo Labels中有类似的做法。另外源码的复现结果虽然很优秀,但我在所有的实验中都没达到论文里超越teacher的效果。
精度好像是因为论文的训练周期比源码默认的多了几倍(看起来极限精度要跑接近一个月,哈哈)
FID的趋势大致呈现为exponential decay before reaching a plateau,这意味着在linear scale上观察,越到后期,下降速度越慢。如果目标是超过previous SOTA,或者visually看起来表现良好,需要的时间并不长。具体参考page 9关于convergence speed的描述,以及图1、图7和图10的展示。
“内外圈优化梯度对齐”这个听起来挺有意思,也许能从另外一个角度提供解释。
你复线代码的时候的duration设的多少?超过以前的SOTA这个数值不需要大,但超过Teacher得大到一定程度:
1. 要达到文章中报告的效果,需要使用`--duration=500`(在CIFAR-cond上增加duration到800,在ImageNet上增加到1000)
- [run_sid.sh, Line 9]
- [sid_train.py, Line 137]
这些设置在对实验的文字描述和图上都可以找到。
2. 你也可以使用SiD release的checkpoints来验证文章报告的结果。
October 31st, 2024
arxiv.org/abs/2410.16794 完美!
前几天已经刷到,还在学习