2021年以来,生成式模型扛把子diffusion model红红火火,已成为当代人工智能研究者不得不面对的话题。本人作为初学者,最近学习了一些经典扩散模型如DDPM、DDIM、Classifier Guidance等,以及一些较新的模型如Stable Diffusion 3和Flow Matching等。趁着记忆犹新赶紧记录下来。本文梳理了diffusion model的发展脉络,并展望了未来diffusion model的发展方向。
一、开山之作:DDPM
2020年以前,在生成模型(Generative Model)领域占统治地位的一直是GANs (Generative Adversarial Networks)。随后,一种去噪扩散概率模型DDPM (Denoising Diffusion Probabilistic Model)横空出世,凭借较小的参数量就能取得很好的图像生成效果,而且由于是概率模型,生成图像的多样性也能得到保证。因此,从2021年至今,扩散模型的研究就如雨后春笋般涌现。一条大致的发展路线如图1所示。从最开始随机生成图像的DDPM,到最后能根据自然语言prompt生成各种风格图像的Stable Diffusion,发展是显著且迅速的。在接下来的介绍中,会看到这样的成就并非一蹴而就,每一步工作都严格建立在前人的基础之上,最后逐渐累积到今天的成果。
flowchart LR A[DDPM] B[DDIM] C[Classifier Guidance] D[Classifier-free Guidance] E[Latent Diffusion Model] F[Dalle2] G[Stable Diffusion] H[Flow Matching] I[Rectified Flow] A--加速采样过程-->B--分类器指导-->C--隐式条件指导-->D--隐空间推理-->E E-->F E-->G A--不同理解视角-->H--流的改进-->I
图1 扩散模型大致发展脉络
1. 扩散模型基本思想
首先,要从开山之作DDPM说起[2],它把图像生成建模成去噪过程。假设有一张图像$x_0$,对它逐渐施加噪声得到一系列图像$x_1,x_2,…,x_T$。图像变得越来越模糊,直到施加$T$次噪声后,图像彻底变成标准正态分布$p(x_T)=\mathcal{N}(x_T;0,I)$。图像生成则是逆过程,从一张噪声图像开始,逐渐去噪得到原始图像。以上过程如图2所示。
图2 DDPM过程示意(从左到右是加噪,从右到左是去噪)
加噪的过程可以表示为
这是一个随机过程。DDPM基于以下假设:
- 它是马尔科夫随机过程,所以每个$x_t(t\neq 0)$仅由$x_{t-1}$决定;
- 每个中间图像都符合正态分布,它的均值取决于上一步的输出图像。
因此,加噪的过程又可以表示为
这里的$\alpha_t(t=1,2,…,T)$是预先定义好的超参数。
实际上,如果要这样一步一步采样$T$步得到$x_T$就太慢了,所以我们可以通过重参数化的技巧直接推导得到$x_0$一步到$x_T$的结果。采样$x_t\sim q(x_t|x_{t-1})$可以写作
同样,采样$x_{t-1}\sim q(x_{t-1}|x_{t-2})$可以写作
将它代入公式(3),就能得到
其中第一行到第二行推导用到独立正态分布相加还是正态分布的性质,所以只需要计算新的正态分布的期望和方差。这个技巧对于扩散模型的理解非常重要,故在此记录下来。后面将看到,加噪对应了模型的训练过程,去噪(也叫做采样)对应了模型的推理过程。
采样是从$x_T$逐步去噪得到$x_0$的过程。当给定图像$x_t$,我们需要知道$x_{t-1}$的分布,从而采样得到$x_{t-1}$。根据DDPM的假设,我们有
由于$\Sigma_q(t)$由超参数组成,数值固定,故模型只需要预测$\mu_q(x_t,x_0)$就能得到$x_{t-1}$的分布。从构成上看,它由$x_t$和$x_0$共同决定,而$x_0$正是我们要求得的目标。我们可以训练一个神经网络,在很长一段时间里,扩散模型都是用的U-Net,用来预测一个$\hat{x}_0=\hat{x}_\theta(x_t,t)$,从而计算一个$\hat{\mu}_{t-1}=\mu_\theta(x_t,t)$($\theta$是神经网络参数),再采样出$x_{t-1}$。从这里可以看出:
- 神经网络并不是直接预测$\mu_q(x_t,x_0)$,而是通过预测它的参数来间接计算;
- 神经网络并不是预测出$x_0$后就直接作为生成的图像,而是先用它采样出$x_{t-1}$,再预测新的$x_0$,采样出$x_{t-2}$,循环下去直到得到$x_0$。
这些操作都是为了让结果更可靠,例如迭代多次自然比一次出结果要靠谱。事实上,我们还可以通过预测$x_0$到$x_t$加噪时重参数化得到的$\epsilon$来计算$\mu_\theta(x_t,t)$。
2. 预测噪声的DDPM
由公式(5)知
那么
代入$\mu_q(x_t,x_0)$得到
用神经网络预测出$\hat{\epsilon}_\theta(x_t,t)$,得到
从而也能得到$x_{t-1}$的分布。有研究发现预测噪声比预测$x_0$有更好的表现[1]。
3. 预测score function的DDPM
概率分布的对数导数被称为得分函数(Score Function)。用得分函数来理解DDPM的好处是,我们能看到扩散模型的几何意义。准确地说,是跟预测噪声的DPPM结合起来,理解去噪的几何意义。
要引入这种形式,要用到Tweedie公式。对于高斯随机变量$z\sim \mathcal{N}(z;\mu_z;\Sigma_z)$,Tweedie公式表示
根据这个公式,我们能得到
因为$\mu_{x_t}=\sqrt{\bar{\alpha}_t}x_0$,所以有
代入$\mu_q(x_t,x_0)$并整理得到
从而可以用神经网络去预测得分函数$\nabla\log{p(x_t)}$,记为$s_\theta(x_t,t)$,则有
4. DDPM的几何意义
可以观察到公式(9)和公式(13)在结构上的相似性,我们建立一个等式
可以看到得分函数和噪声,在常数倍数的缩放下,分别以一种相反的方向移动。从直觉上讲,这是非常自然的,所谓去噪,就是向着噪声相反的方向移动从而恢复原图。DDPM中得分函数的几何意义非常重要,有利于深入对扩散模型本身改进的研究。例如2022年的Flow Matching和2024年的Rectified Flow,就是把加噪和去噪的过程看成一个“流”[11, 12]。(虽然目前我还没整明白,感觉这是一个比较新鲜的方向)
5. 总结
到此为止,就介绍完了DDPM的基本思想和一些等价的形式。它是一个加噪-去噪模型,加噪时的噪声$\epsilon$作为标签来训练神经网络预测噪声的精度;去噪则在给定$x_T$的前提下,不断采样前一步时刻的图像,最终得到$x_0$。总结一下它的优缺点:
- 相比传统的GAN模型,采用了新的生成范式,实现了更高质量的图像生成;
- 模型需要的参数量很小,即U-Net的参数量加上$\alpha_t$;
- 模型需要较多采样步骤才能得到较好的生成结果,需要在采样速度上优化;
- 模型是随机生成图像的,无法主动控制生成的方向。
讲述DDPM用了不少篇幅,且涉及很多公式细节。我认为这是有必要的,因为这些细节正是扩散模型能取得成功的重要来源,后面所有扩散模型都是在它们的基础上演化的。接下来的讲述就主要提及核心思想,减少过于繁琐的细节推导了。
二、加速采样:DDIM
DDIM (Denosing Diffusion Implicit Model)相比DDPM,训练过程相同,将采样效率提高了10-50倍。DDPM采样慢的原因是它是基于马尔科夫假设采样,需要较多步长构建原图。因此,DDIM对此的做法是:构建一个不依赖于一阶马尔科夫假设的采样分布$p(x_{t-1}|x_t,x_0)$,同时不影响DDPM中前向加噪的过程,即$q(x_t|x_0)=\mathcal{N}(x_t;\sqrt{\bar{\alpha}_t}x_0,(1-\bar{\alpha}_t)I)$。
假设采样分布依然是一个高斯分布,且均值也是关于$x_0,x_t$的线性函数,方差是时间步$t$的函数,那么
现在的问题是,怎样选取合适的$\lambda_t,k_t,\sigma_t$,使得$q(x_t|x_0)=\mathcal{N}(x_t;\sqrt{\bar{\alpha}_t}x_0,(1-\bar{\alpha}_t)I)$对$t=1,2,…,T$皆成立(注意$\alpha_t$依然是固定的超参数)。在$t=T$的时候,这个式子天然成立,因为$x_T$正是通过这个分布得到的。
假设在某个$t$的时候,该式子成立。如果找到一组合适的,,,使得$t-1$时刻也成立,即$q(x_{t-1}|x_0)=\mathcal{N}(x_{t-1};\sqrt{\bar{\alpha}_{t-1}}x_0,(1-\bar{\alpha}_{t-1})I)$,自然根据数学归纳的思想,$t-2$到$1$的时刻的解的通项公式也就找到了。我们可以根据待定系数法进行求解。
根据公式(16)对$x_{t-1}$采样,得到
根据$q(x_t|x_0)$的分布对$x_t$采样,得到
代入(17)有
正向采样时,根据$q(x_{t-1}|x_0)$的分布,得到
为了让两个方向采样的$x_{t-1}$分布一致,必须满足
其中$\sigma_t$作为自由变量,方程有无数组解。将其代入(16)并整理得到DDIM的采样分布
改变$\sigma_t$就得到了不同的采样过程,同时前向加噪过程不变。采样中,深度神经网络发挥的作用,依然是预测加噪的噪声$\epsilon$。根据$x_t,x_0,\epsilon$的关系,可以求得上述正态分布的均值,从而能够采样$x_{t-1}$。值得注意的是:
- 当$\sigma_t$跟DDPM中的一致时,这时的采样过程就是DDPM的采样过程;
- 当$\sigma_t=0$时,就将正态分布的均值作为采样的结果,采样过程完全确定,此时的生成模型是一个隐概率模型(Implicit Probabilistic Model)。DDIM也就是在这种情况下被命名的。
现在知道了DDIM的采样分布,但并不知道DDIM是如何加速采样的。假设DDPM中需要从$t=T$依次采样到$t=1$,这个$T$往往较大,例如1000,那么就要采样1000次。DDIM的思路是,从原本的序列$L=[T,T-1,…,1]$中抽取一个长度远小于$T$的子序列$\tau=[\tau_s,\tau_{s-1},…\tau_{1}]$,在这个子序列中采样。例如$\tau_s=1000,\tau_{s-1}=950,\tau_{s-2}=900,…$。DDIM的论文[3]中显示,在CIFAR10上步长为50就达到了DDPM中步长为1000的效果(评价指标为FID, Fréchet Inception Distance)。
三、类别控制:Classifier Guidance
到目前为止,扩散模型生成的图像都是随机的。类别控制(Classifier Guidance)[4]的引入,让它有了指定类别的生成能力,这是一切应用的始端。
1. DDPM中基于条件的去噪过程
它的大致思路是,前向加噪过程和DDPM保持一致,在逆向采样过程中引入类别控制。假设类别标签为$y$,则带类别信号的采样过程定义为$\hat{p}(x_{t-1}|x_{t},y)$
式中$\hat{p}(y|x_{t-1},x_t)=\hat{p}(y|x_{t-1})$可以理解为$x_{t-1}$比$x_t$更接近原图$x_0$,包含了$x_t$以及更多的信息量,可以只靠它决定$y$。$\hat{p}(x_{t-1}|x_t)=p(x_{t-1}|x_t)$意味着带类别信号的采样过程在不用类别信号的时候,跟DDPM的采样过程保持一致。当然,实际上没有这么理所当然,原文[4]的补充材料中有严谨的定义和推导。
式中容易处理的有$p(x_{t-1}|x_t)$和$\hat{p}(y|x_t)$,前者就是DDPM的去噪过程,后者可以训练一个关于$(x_t,y)$的分类器。但$\hat{p}(y|x_{t-1})$是不得而知的。作者将$\hat{p}(x_{t-1}|x_{t},y)$近似为扰动高斯分布(Perturbed Gaussian Distribution),接下来将看到怎么操作的。
假设$p(x_{t-1}|x_{t})=\mathcal{N}(\mu,\Sigma)$,取其对数(默认底数为$e$)得到
对于$\log{\hat{p}(y|x_{t-1})}$,作者假设其曲率比$\Sigma^{-1}$低(这里没搞懂是什么意思),认为假设是合理的,因为当加噪步骤无限多的时候,$|\Sigma|\rightarrow 0$。在这种情况下,对$\log{\hat{p}(y|x_{t-1})}$在$x_{t-1}=\mu$处进行泰勒展开
从而计算出
对比这个式子和公式(24)的形式,它们在形式上是一致的。不用担心常数,因为常数也会刚好匹配的。通过上述推导,得到了带类别条件的采样过程,还是用高斯分布来近似,只是均值要加上$\Sigma g$。$g$就是$t-1$时刻分类器输出的对数梯度,输入来自于神经网络预测的$\hat{\mu}_{t-1}=\mu_\theta({x_t},t)$和类别标签$y$。
2. DDIM中基于条件的去噪过程
由于DDIM等确定性采样方法中设定了方差为0,上述推导不再适用。作者在研究中采用score-based的思路,利用了扩散模型和score matching之间的联系。首先根据贝叶斯公式
然后需要定义条件下噪声的概念,把(15)中的公式写在下面
这是无条件信号下的噪声,据此将有条件噪声$\hat\epsilon(x_t|y)$定义为
只要将DDIM中预测的$\epsilon_\theta(x_t)$替换为$\hat{\epsilon}(x_t|y)$就得到了基于条件的去噪过程。还可以引入gradient score $s$来控制条件的强度,此时就将(28)改写为
3. 总结
在扩散模型中引入类别控制是一大不错的创新,也是走向实用性的必然趋势。但这样扩散模型的效果好坏就很大程度上取决于分类器的好坏了,而且额外使用一个分类器也使得模型变得复杂。由此设想,能不能不用显式的分类器,而是让扩散模型在采样过程中就自然融入了类别信息,不仅简化了模型结构,还能提升标签信息的识别能力?这就是Classifier-free Guidance要解决的问题了。
四、无分类器的控制:Classifier-free Guidance
2021年分类器控制的扩散模型提出,2022年就出现了无分类控制(Classifier-free Guidance)的模型了[5]。其实它的思路很简单,在训练噪声预测模型的时候,直接在输入中引入类别信息,即$\hat\epsilon_\theta(x_t,t)\rightarrow\hat{\epsilon}_\theta(x_t,y,t)$。
1. 基本思路
融入类别信息有很多种方法,现在比较常用的方法有直接拼接(Concatenation)、交叉注意力(Cross-Attention)、通道级注意力(Channel-wise Attention)。无论哪种方法,都要先将类别信息转化成向量,即需要一个文本编码器(Text Encoder)。在自然语言处理(NLP, Natural Language Processing)领域中,已经有很多成熟的方法可以使用了,例如Bert、T5、CLIP等,将文本转化成向量的过程叫做嵌入(Embedding)。而且不局限于类别信息,任何语义信息都可以转化成嵌入向量,所以Classifier-free的可用性大大提升,这项机制也成为了后面Dalle2、Stable Diffusion等模型的基础。
以交叉注意力为例,文本向量作为注意力中的$K$和$V$,图片表征(Image Representation)作为$Q$,注意力的输出为
2. 训练和采样
如果只用图文对去训练$\hat\epsilon_\theta(x_t,y,t)$,一是受限于图文对的数量,二是会束缚模型生成的多样性。所以,投入一定量的无条件信息的图片训练是恰当的,这时模型的标签输入为空,记为$\hat\epsilon_\theta(x_t,y=\varnothing,t)$。$y=\varnothing$的嵌入向量可以随机初始化,也可以人为给定。
采样公式可由(15)、(27)和(29)进一步推导,
最后一步计算用到了(15)中得分函数和噪声$\epsilon$的关系。于是,采样时跟DDPM、DDIM一样,只是需要将原本的$\hat\epsilon_\theta(x_t,t)$用$\hat\epsilon(x_t,y,t)$替代。
五、隐空间加速推理:LDM
此前扩散模型的扩散和采样过程都是在像素空间进行的,这样会造成很大的训练成本和推理开销。LDM (Latent Diffusion Model)提供了一种思路[6],将图像转换到隐空间进行训练和推理,由于隐空间的参数量更小,从而减小了计算开销。像素空间和隐空间之间的转换用一个变分自编码器(VAE, Variational AutoEncoder),将原始图像$x,x\in R^{H\times W\times 3}$通过VAE的编码器$\mathcal{E}$获得图片的隐空间表示$z,z\in R^{\frac{H}{f}\times\frac{W}{f}\times c}$,$f$为下采样率。扩散阶段就是从将原始图像的隐空间表示作为$z_0$,逐渐加噪得到$z_t$;去噪过程类似Classifier-free,在条件指导下逐步采样得到$z_0$,最后通过VAE的解码器$\mathcal{D}$重建像素空间的表征$x_0$。本质上LDM是一种二阶段的图像生成方法。后来知名的Stable Diffusion就是在LDM的基础上构建的。LDM的结构如图3所示。
图3 LDM结构示意
到此为止,主流的扩散模型的基本结构就已经构建完毕了,后面的大多数工作无非是在它的基础上对某些模块进行优化、加速或者应用。最繁琐的数学推导也都集中在了前期阶段(DDPM、DDIM、Classifier Guidance),这是理论孕育的阶段。到后面,就是应用百花齐放的阶段了。
六、扩散模型的应用
2022年Dalle2诞生,掀起了一片互联网的浪花。它的功能很简单,正如Classifier-free描述的那样,用语义信息控制图像生成,因为效果非常好,受到大家关注。它利用CLIP[8]将图像和文本信息编码到同一表征空间进行训练和推理[7]。如今已经有Dalle3,但是OpenAI并没有将它开源,只是提出了一种re-caption的方法[9],对数据进行预处理。(OpenAI真是越来越往CloseAI的方向靠拢了)
Stable Diffusion是基于LDM基础上的直接应用,如今已经发展到了Stable Diffusion 3[12]。它提出了一种修正流(Rectified Flow),这就不得不提到一篇叫流匹配的工作(Flow Matching)[11],将扩散模型建模成一个流的过程(这个还在学习中,不太懂)。不得不感叹数学和物理在扩散模型发展中的重要性,我还看到一个叫薛定谔的桥的方法[13],突然就看不明白了。不过Stable Diffusion 3还有个重要的价值,就是证明了DiT (Diffusion Transformer)的价值[10],这是用transformer替代扩散模型中长期使用的U-Net。transformer真是有一统三界的趋势。
此外,我还看到了扩散模型在图像编辑方面的应用,即给定一张图片和指令,按指令修改图片[14, 15]。
七、可能的研究前景
从前面的介绍来看,在这里直接列出一些可以研究的方向:
- 加速采样,就像DDIM那样对DDPM的采样进行优化,加快图像生成速度;
- 优化压缩模型,即原始图像空间和潜在空间的VAE模型,以更小的潜在空间维度生成更高质量的图像;
- 语义信息编码,或者说如何让模型充分理解语义信息(这里或许可以结合LLM);
- 网络架构,例如从U-Net到DiT;
- 应用扩展,如视频生成和图像编辑等。
参考文献
[1] Luo, C. (2022). Understanding diffusion models: A unified perspective. arXiv preprint arXiv:2208.11970.
[2] Ho, J., Jain, A., & Abbeel, P. (2020). Denoising diffusion probabilistic models. Advances in neural information processing systems, 33, 6840-6851.
[3] Song, J., Meng, C., & Ermon, S. (2020). Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502.
[4] Dhariwal, P., & Nichol, A. (2021). Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34, 8780-8794.
[5] Ho, J., & Salimans, T. (2022). Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598.
[6] Rombach, R., Blattmann, A., Lorenz, D., Esser, P., & Ommer, B. (2022). High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition (pp. 10684-10695).
[7] Ramesh, A., Dhariwal, P., Nichol, A., Chu, C., & Chen, M. (2022). Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125, 1(2), 3.
[8] Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., … & Sutskever, I. (2021, July). Learning transferable visual models from natural language supervision. In International conference on machine learning (pp. 8748-8763). PMLR.
[9] Betker, J., Goh, G., Jing, L., Brooks, T., Wang, J., Li, L., … & Ramesh, A. (2023). Improving image generation with better captions. Computer Science. https://cdn.openai.com/papers/dall-e-3.pdf, 2(3), 8.
[10] Peebles, W., & Xie, S. (2023). Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision (pp. 4195-4205).
[11] Lipman, Y., Chen, R. T., Ben-Hamu, H., Nickel, M., & Le, M. (2022). Flow matching for generative modeling. arXiv preprint arXiv:2210.02747.
[12] Esser, P., Kulal, S., Blattmann, A., Entezari, R., Müller, J., Saini, H., … & Rombach, R. (2024, March). Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first International Conference on Machine Learning.
[13] Shi, Y., De Bortoli, V., Campbell, A., & Doucet, A. (2024). Diffusion Schrödinger bridge matching. Advances in Neural Information Processing Systems, 36.
[14] Brooks, T., Holynski, A., & Efros, A. A. (2023). Instructpix2pix: Learning to follow image editing instructions. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 18392-18402).
[15] Sheynin, S., Polyak, A., Singer, U., Kirstain, Y., Zohar, A., Ashual, O., … & Taigman, Y. (2024). Emu edit: Precise image editing via recognition and generation tasks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 8871-8879).