生成扩散模型漫谈(二十五):基于恒等式的蒸馏(上)
By 苏剑林 | 2024-05-01 | 69594位读者 |今天我们分享一下论文《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。假设我们有一个在目标数据集训练好的教师扩散模型ϵφ∗(xt,t)
SiD及前作DDE使用了一个看上去很绕但是也很聪明的思路:
如果gθ(z)产生的数据分布跟目标分布很相似,那么拿gθ(z) 生成的数据集去训练一个扩散模型ϵψ∗(xt,t) 的话,它也应该跟ϵφ∗(xt,t) 很相似?
初级形式 #
这个思路的聪明之处在于,它绕开了对教师模型生成样本的需求,也不需要训练教师模型的真实样本,因为“拿gθ(z)
当然,这还只是思路,将其转换为实际可行的训练方案还有一段路要走。首先回顾一下扩散模型,我们采用《生成扩散模型漫谈(三):DDPM = 贝叶斯 + 去噪》的形式,我们使用如下方式对输入x0
xt=ˉαtx0+ˉβtε,ε∼N(0,I)

























换言之p(xt|x0)=N(xt;ˉαtx0,ˉβ2tI)



































φ∗=argminφEx0∼˜p(x0),ε∼N(0,I)[‖ϵφ(ˉαtx0+ˉβtε,t)−ε‖2]























































这里的˜p(x0)











ψ∗=argminψEx(g)0∼pθ(x(g)0),ε∼N(0,I)[‖ϵψ(x(g)t,t)−ε‖2]=argminψEz,ε∼N(0,I)[‖ϵψ(x(g)t,t)−ε‖2]





























































































这里x(g)t=ˉαtx(g)0+ˉβtε=ˉαtgθ(z)+ˉβtε






















































θ∗=argminθEz,ε∼N(0,I)[‖ϵφ∗(x(g)t,t)−ϵψ∗(x(g)t,t)‖2]⏟L1






































































































注意式(3)
















点睛之笔 #
谈到GAN,有读者可能会“闻之色变”,因为它是出了名的容易训崩。很遗憾,上述提出的式(3)
1、理论上要求先求出式(3)
的最优解,然后才去优化式(4) ,但实际上从训练成本考虑,我们并没有将它训练到最优就去优化式(4) 了; 2、理论上ψ∗
随θ 而变,即应该写成ψ∗(θ) ,从而在优化式(4) 时应该多出一项ψ∗(θ) 对θ 的梯度,但实际上在优化式(4) 时我们都只当ψ∗ 是常数。
这两个问题非常本质,它们也是GAN训练不稳定的根本原因,此前论文《Revisiting GANs by Best-Response Constraint: Perspective, Methodology, and Application》也特意从第2点出发改进了GAN的训练。看上去,这两个问题哪一个都无法解决,尤其是第1个,我们几乎不可能总是将ψ
幸运的是,对于上述扩散模型的蒸馏问题,SiD提出了一个有效缓解这两个问题的方案。SiD的想法可谓非常“朴素”:既然ψ∗
这就是SiD的核心贡献,也是让人拍案叫绝的“点睛之笔”。
恒等变换 #
接下来我们具体来看做了什么恒等变换。我们先来看式(2)
Ex0∼˜p(x0),ε∼N(0,I)[‖ϵφ(ˉαtx0+ˉβtε,t)−ε‖2]=Ex0∼˜p(x0),xt∼p(xt|x0)[‖ϵφ(xt,t)−xt−ˉαtx0ˉβt‖2]=Ex0∼˜p(x0),xt∼p(xt|x0)[‖ϵφ(xt,t)+ˉβt∇xtlogp(xt|x0)‖2]

























































































































































根据《生成扩散模型漫谈(五):一般框架之SDE篇》的得分匹配相关结果,上述目标的最优解是ϵφ∗(xt,t)=−ˉβt∇xtlogp(xt)


































































Ez,ε∼N(0,I)[‖ϵφ∗(x(g)t,t)−ϵψ∗(x(g)t,t)‖2]=Ez,ε∼N(0,I)[⟨ϵφ∗(x(g)t,t)−ϵψ∗(x(g)t,t),ϵφ∗(x(g)t,t)+ˉβt∇x(g)tlogpθ(x(g)t)⟩]=Ez,ε∼N(0,I)[⟨ϵφ∗(x(g)t,t)−ϵψ∗(x(g)t,t),ϵφ∗(x(g)t,t)⟩]+Ex(g)t∼pθ(x(g)t)[⟨ϵφ∗(x(g)t,t)−ϵψ∗(x(g)t,t),ˉβt∇x(g)tlogpθ(x(g)t)⟩]


















































































































































































































































接下来要用到在《生成扩散模型漫谈(十八):得分匹配 = 条件得分匹配》证明过的一个恒等式,来化简上式的红色部分:
∇xtlogp(xt)=Ex0∼p(x0|xt)[∇xtlogp(xt|x0)]











































这是由概率密度定义以及贝叶斯公式推出的恒等式,不依赖于p(xt),p(xt|x0),p(x0|xt)























Ex(g)t∼pθ(x(g)t)[⟨ϵφ∗(x(g)t,t)−ϵψ∗(x(g)t,t),ˉβt∇x(g)tlogpθ(x(g)t)⟩]=Ex(g)t∼pθ(x(g)t),x(g)0∼pθ(x(g)0|x(g)t)[⟨ϵφ∗(x(g)t,t)−ϵψ∗(x(g)t,t),ˉβt∇x(g)tlogp(x(g)t|x(g)0)⟩]=−Ex(g)0∼pθ(x(g)0),x(g)t∼pθ(x(g)t|x(g)0)[⟨ϵφ∗(x(g)t,t)−ϵψ∗(x(g)t,t),xt−ˉαtx0ˉβt⟩]=−Ez,ε∼N(0,I)[⟨ϵφ∗(x(g)t,t)−ϵψ∗(x(g)t,t),ε⟩]


































































































































































































































































































跟绿色部分合并,就得到学生模型新的损失函数
L2=Ez,ε∼N(0,I)[⟨ϵφ∗(x(g)t,t)−ϵψ∗(x(g)t,t),ϵφ∗(x(g)t,t)−ε⟩]





























































这就是SiD的核心结果,原论文的实验结果显示它能够高效地实现蒸馏,而式(4)



相比式(4)
其他细节 #
到目前为止,本文的推导基本上是原论文推导的重复,但出了个别记号上的不一致外,还有一些细节上的不同,下面简单澄清一下,以免读者混淆。
首先,论文的推导默认了ˉαt=1
最后,原论文发现损失函数L1
相比其他蒸馏方法,SiD的缺点是对显存的需求比较大,因为它同时要维护三个模型ϵφ(xt,t)
延伸思考 #
笔者相信,对于一开始的“初级形式”,即式(3)
同时,SiD也留下了不少值得进一步思考和探索的问题。比如,学生模型的损失(9)
Ez,ε∼N(0,I)[‖ϵφ∗(x(g)t,t)−ϵψ∗(x(g)t,t)‖2]=Ez,ε∼N(0,I)[‖ϵφ∗(x(g)t,t)‖2−2⟨ϵφ∗(x(g)t,t),ϵψ∗(x(g)t,t)⟩+‖ϵψ∗(x(g)t,t)‖2]=Ex(g)t∼pθ(x(g)t)[‖ϵφ∗(x(g)t,t)‖2−2⟨ϵφ∗(x(g)t,t),−ˉβt∇x(g)tlogpθ(x(g)t)⟩+⟨ϵψ∗(x(g)t,t),−ˉβt∇x(g)tlogpθ(x(g)t)⟩]






































































































































































































































这里的每一个−ˉβt∇x(g)tlogpθ(x(g)t)








































































Ez,ε∼N(0,I)[‖ϵφ∗(x(g)t,t)−ϵψ∗(x(g)t,t)‖2]=Ez,ε∼N(0,I)[‖ϵφ∗(x(g)t,t)‖2−2⟨ϵφ∗(x(g)t,t),ε⟩+‖ϵψ∗(x(g)t,t)‖2]=Ez,ε∼N(0,I)[‖ϵφ∗(x(g)t,t)−ε‖2+‖ϵψ∗(x(g)t,t)‖2]+常数




























































































































































这是学生模型,也就是生成器的损失,然后我们再对比学生数据去噪模型的损失(3)



ψ∗=argminψEz,ε∼N(0,I)[‖ϵψ(x(g)t,t)−ε‖2]











































这两个式子联合起来看,我们可以发现学生模型实则在向教师模型看齐,并且试图远离学生数据所训练的去噪模型,形式上很像LSGAN,ϵψ(x(g)t,t)











1、GAN:一开始造假者(生成器)和鉴别者(判别器)都是小白,鉴别者不断对比真品和赝品来提供自己的鉴宝水平,造假者则通过鉴别者的反馈不断提高自己的造假水平;
2、SiD:完全没有真品,但有一个绝对权威的鉴宝大师(教师模型),造假者(学生模型)不断制作赝品,同时培养自己的鉴别者(学生数据训练的去噪模型),然后通过自家鉴别者跟大师的交流来提高自己造假水平。
可能有读者会问:为什么SiD中的造假者不直接向大师请教,而是要通过培养自己的鉴别者来间接获得反馈呢?这是因为直接跟大师交流的话,可能会出现的问题就是长期都只交流同一个作品的技术,最终只制造出了一种能够以假乱真的赝品(模式坍缩),而通过培养自己的鉴别者一定程度上就可以避免这个问题,因为造假者的学习策略是“多得到大师的好评,同时尽量减少自家人的好评”,如果造假者还是只制造一种赝品,那么大师和自家的好评都会越来越多,这不符合造假者的学习策略,从而迫使造假者不断开发新的产品而不是固步自封。
此外,读者可以发现,SiD整个训练并没有利用到扩散模型的递归采样的任何信息,换句话说它纯粹是利用了去噪这一训练方式所训练出来的去噪模型,那么一个自然的问题是:如果单纯为了训练一个单步的生成模型,而不是作为已有扩散模型的蒸馏,那么我们训练一个只具有单一噪声强度的去噪模型会不会更好?比如像DDE一样,固定ˉαt=1某个常数
文章小结 #
在这篇文章中,我们介绍了一种新的将扩散模型蒸馏为单步生成模型的方案,其思想可以追溯到前两年的利用去噪自编码器训练生成模型的工作,它不需要获得教师模型的真实训练集,也不需要迭代教师模型来生成样本对,而引入了类似GAN的交替训练,同时提出了关键的恒等变换来稳定训练过程,整个方法有颇多值得学习之处。
转载到请包括本文地址:https://spaces.ac.cn/archives/10085
更详细的转载事宜请参考:《科学空间FAQ》
如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。
如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!
如果您需要引用本文,请参考:
苏剑林. (May. 01, 2024). 《生成扩散模型漫谈(二十五):基于恒等式的蒸馏(上) 》[Blog post]. Retrieved from https://spaces.ac.cn/archives/10085
@online{kexuefm-10085,
title={生成扩散模型漫谈(二十五):基于恒等式的蒸馏(上)},
author={苏剑林},
year={2024},
month={May},
url={\url{https://spaces.ac.cn/archives/10085}},
}
November 28th, 2024
实在不敢恭维原论文的写作水平,明明很简单的想法写的这么故弄玄虚