发布时间:2025-12-09 16:48:26 浏览次数:4
继OpenAI在2021提出的文本转图像模型DALLE之后,越来越多的大公司卷入这个方向,例如谷歌相继推出了Imagen和Parti。一些主流的文本转图像模型,例如DALL·E 2,stable-diffusion和Imagen采用了扩散模型(Diffusion Model)作为图像生成模型,这也引发了对扩散模型的研究热潮。
与GAN相比,扩散模型训练更稳定,而且能够生成更多样的样本,OpenAI的论文Diffusion Models Beat GANs on Image Synthesis也证明了<font color='red>扩散模型能够超越GAN。
论文地址:https://arxiv.org/abs/2105.05233
简单来说,扩散模型包含两个过程:前向扩散过程和反向生成过程。前向扩散过程是对一张图像逐渐添加高斯噪音直至变成随机噪音,而反向生成过程是去噪音过程,我们将从一个随机噪音开始逐渐去噪音直至生成一张图像,这也是我们要求解或者训练的部分。扩散模型与其它主流生成模型的对比如下所示:
目前所采用的扩散模型大都是来自于2020年的工作DDPM: Denoising Diffusion Probabilistic Models,DDPM对之前的扩散模型(具体见Deep Unsupervised Learning using Nonequilibrium Thermodynamics)进行了简化,并通过变分推断(variational inference)来进行建模,这主要是因为扩散模型也是一个隐变量模型(latent variable model),相比VAE这样的隐变量模型,扩散模型的隐变量是和原始数据是同维度的,而且推理过程(即扩散过程)往往是固定的。本文将基于DDPM详细介绍扩散模型的原理,并给出具体的代码实现和分析。
论文地址:https://arxiv.org/abs/2006.11239
扩散模型包括两个过程:前向过程(forward process)和反向过程(reverse process),其中前向过程又称为为扩散过程(diffusion process),如下图所示。无论是前向过程还是反向过程都是一个参数化的马尔可夫链(Markov chain ),其中反向过程可以用来生成数据,这里我们将通过变分推断来进行建模和求解。
扩散过程是指的对数据逐渐增加高斯噪声直至数据完全变成随机噪声的过程。对于原始数据x0∼q(x0)x_0\sim q(x_0)x0∼q(x0),总共包含TTT步的扩散过程的每一步都是对上一步得到的数据xt−1x_{t-1}xt−1按如下方式增加高斯噪声:
q(xt∣xt−1)=N(xt;1−βtxt−1,βtI)q(x_t|x_{t-1})=\mathcal{N}(x_t;\sqrt{1-\beta_t}x_{t-1},\beta_t I)q(xt∣xt−1)=N(xt;1−βtxt−1,βtI)
这里{βt}t=1T\{\beta_t\}_{t=1}^T{βt}t=1T为每一步所采用的方差,它介于010~10 1之间。对于扩散模型,往往称不同step的方差设定为variance schedule或者noise schedule,通常情况下,越后面的step会采用更大的方差,即满足β1<β2<...<βT\beta_1<\beta_2<...<\beta_Tβ1<β2<...<βT。在一个设计好的variance schedule下,如果扩散步数TTT足够大,那么最终得到的xTx_TxT就完全丢失了原始数据而变成了一个随机噪音。 扩散过程的每一步都生成一个带噪音的数据xtx_txt,整个扩散过程也就是一个马尔卡夫链:
q(x1:T∣x0)=∏t=1Tq(xt∣xt−1)q(x_{1:T}|x_0)= {\textstyle \prod_{t=1}^{T}}q(x_t|x_{t-1})q(x1:T∣x0)=∏t=1Tq(xt∣xt−1)
另外,要指出的是,扩散过程往往是固定的,即采用一个预先定义好的variance schedule,比如DDPM就是一个线性的variance schedule。扩散过程的一个重要特性是我们可以直接基于原始数据x0x_0x0来对任意ttt步的xtx_txt进行采样:xtq(xt∣x0)x_t~ q(x_t|x_0)xt q(xt∣x0)。这里定义αt=1−βt\alpha_t=1-\beta_tαt=1−βt和αˉt=∏i=1tαi\bar{\alpha}_t= {\textstyle \prod_{i=1}^{t}\alpha _i}αˉt=∏i=1tαi,通过重参数技巧(和VAE类似),那么有:
上述推导过程利用了两个方差不同的高斯分布N(0,σ12I)\mathcal{N}(0,\sigma _1^2\mathtt{I})N(0,σ12I)和N(0,σ22I)\mathcal{N}(0,\sigma _2^2\mathtt{I})N(0,σ22I)相加等于一个新的高斯分布N(0,(σ12+σ22)I)\mathcal{N}(0,(\sigma _1^2+\sigma _2^2)\mathtt{I})N(0,(σ12+σ22)I)。反重参数化后,我们得到:
q(xt∣x0)=N(xt;αˉtx0,(1−αˉt)I)q(\mathrm{x}_t|\mathrm{x}_0)=\mathcal{N}(\mathrm{x}_t;\sqrt{\bar{\alpha}_t}\mathrm{x}_0,(1-\bar{\alpha}_t)\mathtt{I})q(xt∣x0)=N(xt;αˉtx0,(1−αˉt)I)
扩散过程的这个特性很重要。首先,可以看到xt\mathrm{x}_txt其实可以看成是原始数据x0\mathrm{x}_0x0和随机噪音ϵ\epsilonϵ的线性组合,其中αˉt\sqrt{\bar{\alpha}_t}αˉt和1−αˉt\sqrt{1-\bar{\alpha}_t}1−αˉt为组合系数,它们的平均和等于1,可以称两者分别为signal_rate和noise_rate
参考:https://keras.io/examples/generative/ddim/#diffusion-schedule和Variational Diffusion Models
更进一步,可以基于αˉt\bar{\alpha}_tαˉt而不是βt\beta_tβt来定义noise schedule(见Improved Denoising Diffusion Probabilistic Models所设计的cosine schedule),因为这样处理更直接,比如我们直接将αˉT\bar{\alpha}_TαˉT设定为一个接近0的值,那么就可以保证最终得到的xT\mathrm{x}_TxT近似为一个随机噪声。其次,后面的建模和分析过程将使用这个特性。
扩散过程是将数据噪声化,那么反向过程就是一个去噪的过程。如果我们知道反向过程的每一步的真实分布q(xt−1∣xt)q(\mathrm{x}_{t-1}|\mathrm{x}_t)q(xt−1∣xt),那么从一个随机噪音xt∼N(0,I)\mathrm{x}_t\sim \mathcal{N}(0,\mathtt{I})xt∼N(0,I)开始,逐渐去噪,就能生成一个真实的样本,所以反向过程也就是生成数据的过程。
估计分布q(xt−1∣xt)q(\mathrm{x}_{t-1}|\mathrm{x}_t)q(xt−1∣xt)需要用到整个训练样本,可以用神经网络来估计这些分布。这里,将反向过程也定义为一个马尔可夫链,只不过它是由一系列用神经网络参数化的高斯分布来组成:pθ(x0;T)=p(xT)∏t=1Tp(xt−1∣xt)p_\theta(\mathrm{x}_{0;T})=p(\mathrm{x}_T)\prod_{t=1}^{T}p(\mathrm{x}_{t-1}|\mathrm{x}_t) pθ(x0;T)=p(xT)t=1∏Tp(xt−1∣xt)pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),∑θ(xt,t))p_\theta(\mathrm{x}_{t-1}|\mathrm{x}_t)=\mathcal{N}(\mathrm{x}_{t-1};\mu_\theta(\mathrm{x}_t,t), {\textstyle \sum_{\theta}(\mathrm{x}_t,t)})pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),∑θ(xt,t))
这里p(xT)=N(xT;0,I)p(\mathrm{x}_T)=\mathcal{N}(\mathrm{x}_T;0,\mathtt{I})p(xT)=N(xT;0,I),而pθ(xt−1∣xt)p_\theta(\mathrm{x}_{t-1}|\mathrm{x}_t)pθ(xt−1∣xt)为参数化的高斯分布,它们的均值和方差由训练的网络μθ(xt,t)\mu_\theta(\mathrm{x}_t,t)μθ(xt,t)和∑θ(xt,t){\textstyle \sum_{\theta}}(\mathrm{x}_t,t)∑θ(xt,t)给出。实际上,扩散模型就是要得到这些训练好的网络,因为它们构成了最终的生成模型。
虽然分布q(xt−1∣xt)q(\mathrm{x}_{t-1}|\mathrm{x}_t)q(xt−1∣xt)是不可直接处理的,但是加上条件x0\mathrm{x}_0x0的后验分布q(xt−1∣xt,x0)q(\mathrm{x}_{t-1}|\mathrm{x}_t, \mathrm{x}_0)q(xt−1∣xt,x0)却是可处理的,这里有:q(xt−1∣xt,x0)=N(xt−1;μ~(xt,x0,β~tI))q(\mathrm{x}_{t-1}|\mathrm{x}_t, \mathrm{x}_0)=\mathcal{N}(\mathrm{x}_{t-1};\tilde{\mu}(\mathrm{x}_t,\mathrm{x}_0,\tilde{\beta}_t\mathtt{I}))q(xt−1∣xt,x0)=N(xt−1;μ~(xt,x0,β~tI))
下面来具体推导这个分布。首先根据贝叶斯公式,有:q(xt−1∣xt,x0)=q(xt∣xt−1,x0)q(xt−1∣x0)q(xt∣x0)q(\mathrm{x}_{t-1}|\mathrm{x}_t,\mathrm{x}_0)=q(\mathrm{x}_t|\mathrm{x}_{t-1}, \mathrm{x}_0)\frac{q(\mathrm{x}_{t-1}|\mathrm{x}_0)}{q(\mathrm{x}_t|\mathrm{x}_0)}q(xt−1∣xt,x0)=q(xt∣xt−1,x0)q(xt∣x0)q(xt−1∣x0)
由于扩散过程的马尔可夫链特性,知道分布q(xt−1∣xt,x0)=q(xt∣xt−1)=N(xt;1−βtxt−1,βtI)q(\mathrm{x}_{t-1}|\mathrm{x}_t,\mathrm{x}_0)=q(\mathrm{x}_t|\mathrm{x}_{t-1})=\mathcal{N}(\mathrm{x}_t;\sqrt{1-\beta_t}\mathrm{x}_{t-1},\beta_t\mathtt{I})q(xt−1∣xt,x0)=q(xt∣xt−1)=N(xt;1−βtxt−1,βtI)(这里条件x0\mathrm{x}_0x0是多余的),而由前面得到的扩散过程特性可知:q(xt−1∣x0)=N(xt−1;αˉt−1x0,(1−αˉt−1)I)q(\mathrm{x}_{t-1}|\mathrm{x}_{0})=\mathcal{N}(\mathrm{x}_{t-1};\sqrt{\bar{\alpha}_{t-1}}\mathrm{x}_0,(1-\bar{\alpha}_{t-1})\mathtt{I})q(xt−1∣x0)=N(xt−1;αˉt−1x0,(1−αˉt−1)I)q(xt∣x0)=N(xt;αˉtx0,(1−αˉt)I)q(\mathrm{x}_{t}|\mathrm{x}_{0})=\mathcal{N}(\mathrm{x}_{t};\sqrt{\bar{\alpha}_{t}}\mathrm{x}_0,(1-\bar{\alpha}_{t})\mathtt{I})q(xt∣x0)=N(xt;αˉtx0,(1−αˉt)I)
所以,有:
这里的C(xt,x0)C(\mathrm{x}_t, \mathrm{x}_0)C(xt,x0)是一个和xt−1\mathrm{x}_{t-1}xt−1无关的部分,所以省略。根据高斯分布的概率密度函数定义和上述结果(配平方),可以得到分布q(xt∣xt−1,x0)q(\mathrm{x}_{t}|\mathrm{x}_{t-1}, \mathrm{x}_{0})q(xt∣xt−1,x0)的均值和方差:
可以看到方差是一个定量(扩散过程参数固定),而均值是一个依赖x0\mathrm{x}_0x0和xt\mathrm{x}_txt的函数。这个分布将会被用于推导扩散模型的优化目标。
上面介绍了扩散模型的扩散过程和反向过程,现在我们来从另外一个角度来看扩散模型:如果把中间产生的变量看成隐变量的话,那么扩散模型其实是包含T个隐变量的隐变量模型(latent variable model),它可以看成是一个特殊的Hierarchical VAEs(见Understanding Diffusion Models: A Unified Perspective)
相比VAE来说,扩散模型的隐变量是和原始数据同维度的,而且encoder(即扩散过程)是固定的。既然扩散模型是隐变量模型,那么我们就可以基于变分推断来得到variational lower bound(VLB,又称ELBO)作为最大化优化目标,这里有:
这里最后一步是利用了Jensen's inequality(不采用这个不等式的推导见博客What are Diffusion Models?),对于网络训练来说,其训练目标为VLB取负:
进一步对训练目标进行分解可得:
可以看到最终的优化目标共包含T+1T+1T+1项,其中L0L_0L0可以看成是原始数据重建,优化的是负对数似然,L0L_0L0可以用估计的N(x0;μθ(x1,1),∑θ(x1,1)\mathcal{N}(\mathrm{x}_0;\mu_\theta (\mathrm{x}_1,1), {\textstyle \sum_{\theta }}(\mathrm{x}_1,1)N(x0;μθ(x1,1),∑θ(x1,1)来构建一个离散化的decoder来计算(见DDPM论文3.3部分);
而LTL_TLT计算的是最后得到的噪音的分布和先验分布的KL散度,这个KL散度没有训练参数,近似为0,因为先验p(xT)=N(0,I)p(\mathrm{x}_T)=\mathcal{N}(0,\mathtt{I})p(xT)=N(0,I),而扩散过程最后得到的随机噪音p(xT∣x0)p(\mathrm{x}_T|\mathrm{x}_0)p(xT∣x0)也近似为N(0,I)\mathcal{N}(0,\mathtt{I})N(0,I);
而Lt−1L_{t-1}Lt−1则是计算的是估计分布pθ(xt−1∣xt)p_\theta(\mathrm{x}_{t-1}|\mathrm{x}_t)pθ(xt−1∣xt)和真实后验分布q(xt−1∣xt,x0)q(\mathrm{x}_{t-1}|\mathrm{x}_t,\mathrm{x}_0)q(xt−1∣xt,x0)的KL散度,这里希望我们估计的去噪过程和依赖真实数据的去噪过程近似一致:
之所以在前面将pθ(xt−1∣xt)p_\theta(\mathrm{x}_{t-1}|\mathrm{x}_t)pθ(xt−1∣xt)定义为一个用网络参数化的高斯分布N(xt−1;μθ(xt,t),∑θ(xt,t))\mathcal{N}(\mathrm{x}_{t-1};\mu_\theta(\mathrm{x}_t,t), {\textstyle \sum_{\theta}(\mathrm{x}_t,t)})N(xt−1;μθ(xt,t),∑θ(xt,t)),是因为要匹配的后验分布q(xt−1∣xt,x0)q(\mathrm{x}_{t-1}|\mathrm{x}_t,\mathrm{x}_0)q(xt−1∣xt,x0)也是一个高斯分布。对于训练目标L0L_0L0和Lt−1L_{t-1}Lt−1来说,都是希望得到训练好的网络μθ(xt,t)\mu_\theta(\mathrm{x}_t,t)μθ(xt,t)和∑θ(xt,t){\textstyle \sum_{\theta}(\mathrm{x}_t,t)}∑θ(xt,t)(对于L0L_0L0,t=1t=1t=1)。
DDPM对pθ(xt−1∣xt)p_\theta(\mathrm{x}_{t-1}|\mathrm{x}_t)pθ(xt−1∣xt)做了进一步简化,采用固定的方差:∑θ(xt,t)=σt2I{\textstyle \sum_{\theta}(\mathrm{x}_t,t)}=\sigma_t^2\mathtt{I}∑θ(xt,t)=σt2I,这里的σt2\sigma_t^2σt2可以设定为βt\beta_tβt或者β~t\tilde{\beta}_tβ~t(这其实是两个极端,分别是上限和下限,也可以采用可训练的方差,见论文Improved Denoising Diffusion Probabilistic Models和Analytic-DPM: an Analytic Estimate of the Optimal Reverse Variance in Diffusion Probabilistic Models)。
这里假定σt2=β~t\sigma_t^2=\tilde{\beta}_tσt2=β~t,那么:
q(xt−1∣xt,x0)=N(xt−1;μ~(xt,x0),σt2I)pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),σt2I)q(\mathrm{x}_{t-1}|\mathrm{x}_t,\mathrm{x}_0)=\mathcal{N}(\mathrm{x}_{t-1};\tilde{\mu}(\mathrm{x}_t,\mathrm{x}_0),\sigma_t^2\mathtt{I})p_\theta (\mathrm{x}_{t-1}|\mathrm{x}_t)=\mathcal{N}(\mathrm{x}_{t-1};\mu_\theta (\mathrm{x}_t,t),\sigma_t^2\mathtt{I})q(xt−1∣xt,x0)=N(xt−1;μ~(xt,x0),σt2I)pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),σt2I)
对于两个高斯分布的KL散度,其计算公式为:
那么就有:
那么优化目标Lt−1L_{t-1}Lt−1即为:Lt−1=Eq(xt∣x0)[12σt2∣∣μ~t(xt,x0)−μθ(xt,t)∣∣2]L_{t-1}=\mathbb{E}_{q(\mathrm{x}_t|\mathrm{x}_0)}[\frac{1}{2\sigma _t^2} ||\tilde{\mu}_t(\mathrm{x}_t,\mathrm{x}_0)-\mu_\theta (\mathrm{x}_t,t)||^2]Lt−1=Eq(xt∣x0)[2σt21∣∣μ~t(xt,x0)−μθ(xt,t)∣∣2]
从上述公式来看,我们希望网络学习到的均值μθ(xt,t)\mu_\theta (\mathrm{x}_t,t)μθ(xt,t)和后验分布的均值μ~t(xt,x0)\tilde{\mu}_t(\mathrm{x}_t,\mathrm{x}_0)μ~t(xt,x0)一致。不过DDPM发现预测均值并不是最好的选择。根据前面得到的扩散过程的特性,有:xt(x0,ϵ)=αˉtx0+1−αˉtϵwhere ϵ∼N(0,I)\mathrm{x}_{t}(\mathrm{x}_0,\epsilon)=\sqrt{\bar{\alpha }_t}\mathrm{x}_0+\sqrt{1-\bar{\alpha}_t\epsilon} \text{ where }\epsilon \sim\mathcal{N}(0,\mathtt{I} )xt(x0,ϵ)=αˉtx0+1−αˉtϵ where ϵ∼N(0,I)
将这个公式带入上述优化目标,可以得到:
进一步地,对μθ(xt(x0,ϵ),t)\mu_\theta(\mathrm{x}_t (\mathrm{x}_0, \epsilon), t)μθ(xt(x0,ϵ),t)也进行重参数化,变成:μθ(xt(x0,ϵ),t)=1αt(xt(x0,ϵ)−βt1−αˉtϵθ(xt(x0,ϵ),t))\mu_\theta(\mathrm{x}_t (\mathrm{x}_0, \epsilon), t)=\frac{1}{\sqrt{\alpha}_t}(\mathrm{x}_t(\mathrm{x}_0,\epsilon)-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon_\theta(\mathrm{x}_t(\mathrm{x}_0,\epsilon), t))μθ(xt(x0,ϵ),t)=αt1(xt(x0,ϵ)−1−αˉtβtϵθ(xt(x0,ϵ),t))
这里的ϵθ\epsilon_\thetaϵθ是一个基于神经网络的拟合函数,这意味着由原来的预测均值而换成预测噪音ϵ\epsilonϵ。我们将上述等式代入优化目标,可以得到:
DDPM进一步对上述目标进行了简化,即去掉了权重参数,变成了Lt−1simple=Ex0,ϵ∼N(0,I)[∣∣ϵ−ϵθ(αˉtx0+1−αˉtϵ,t)∣∣2]L_{t-1}^{simple}=\mathbb{E}_{\mathrm{x}_0,\epsilon \sim\mathcal{N}(0,\mathtt{I})}[||\epsilon -\epsilon _\theta (\sqrt{\bar{\alpha}_t}\mathrm{x}_0+\sqrt{1-\bar{\alpha}_t}\epsilon ,t )||^2]Lt−1simple=Ex0,ϵ∼N(0,I)[∣∣ϵ−ϵθ(αˉtx0+1−αˉtϵ,t)∣∣2]。这里的ttt在[1,T][1,T][1,T]范围内取值(如前所述,其中取1时对应L0L_0L0)。由于去掉了不同ttt的权重系数,所以这个简化的目标其实是VLB优化目标进行了reweight。从DDPM的对比实验结果来看,预测噪音比预测均值效果要好,采用简化版本的优化目标比VLB目标效果要好:
虽然扩散模型背后的推导比较复杂,但是最终得到的优化目标非常简单,就是让网络预测的噪音和真实的噪音一致。DDPM的训练过程也非常简单,如下图所示:随机选择一个训练样本->从1-T随机抽样一个t->随机产生噪音-计算当前所产生的带噪音数据(红色框所示)->输入网络预测噪音->计算产生的噪音和预测的噪音的L2损失->计算梯度并更新网络。
一旦训练完成,其采样过程也非常简单,如上所示:我们从一个随机噪音开始,并用训练好的网络预测噪音,然后计算条件分布的均值(红色框部分),然后用均值加标准差乘以一个随机噪音,直至t=0完成新样本的生成(最后一步不加噪音)。
不过实际的代码实现和上述过程略有区别(见https://github.com/hojonathanho/diffusion/issues/5:先基于预测的噪音生成x0\mathrm{x}_0x0,并进行了clip处理(范围[-1, 1],原始数据归一化到这个范围),然后再计算均值。这应该算是一种约束,既然模型预测的是噪音,那么我们也希望用预测噪音重构处理的原始数据也应该满足范围要求。
前面我们介绍了扩散模型的原理以及优化目标,那么扩散模型的核心就在于训练噪音预测模型,由于噪音和原始数据是同维度的,所以我们可以选择采用AutoEncoder架构来作为噪音预测模型。DDPM所采用的模型是一个基于residual block和attention block的U-Net模型。如下所示:
U-Net属于encoder-decoder架构,其中encoder分成不同的stages,每个stage都包含下采样模块来降低特征的空间大小(H和W),然后decoder和encoder相反,是将encoder压缩的特征逐渐恢复。U-Net在decoder模块中还引入了skip connection,即concat了encoder中间得到的同维度特征,这有利于网络优化。DDPM所采用的U-Net每个stage包含2个residual block,而且部分stage还加入了self-attention模块增加网络的全局建模能力。
另外,扩散模型其实需要的是个噪音预测模型,实际处理时,我们可以增加一个time embedding(类似transformer中的position embedding)来将timestep编码到网络中,从而只需要训练一个共享的U-Net模型。具体地,DDPM在各个residual block都引入了time embedding,如上图所示。
最后,我们基于PyTorch框架给出DDPM的具体实现,这里主要参考了三套代码实现:
Framework
Component组件:
Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding
https://arxiv.org/abs/2205.11487
Decoder can be trained without labelled data.
Learning Transferable Visual Models From Natural Language Supervision
https://arxiv.org/abs/2103.00020
The easier way to use a Diffusion Model in PyTorch is to use the denoising-diffusion-pytorch package.
pip install denoising_diffusion_pytorch一个示例代码:
import torchfrom denoising_diffusion_pytorch import Unet, GaussianDiffusionmodel = Unet(dim = 64,dim_mults = (1, 2, 4, 8))model = Unet(dim = 64,dim_mults = (1, 2, 4, 8))diffusion = GaussianDiffusion(model,image_size = 128,timesteps = 1000, # number of stepsloss_type = 'l1' # L1 or L2)training_images = torch.randn(8, 3, 128, 128)loss = diffusion(training_images)loss.backward()sampled_images = diffusion.sample(batch_size = 4)