新闻  |   论坛  |   博客  |   在线研讨会
大一统视角理解扩散模型Understanding Diffusion Models: A Unified Perspective(2)
计算机视觉工坊 | 2022-10-19 07:08:47    阅读:1337   发布文章

Diffusion Model recap

在扩散模型里,有几个重要的假设。其中一个就是每一步扩散过程的变换,都是对前一步结果的高斯变换(上一节MHVAE的限制条件2):

图片与MHVAE不同,编码器侧的潜在向量分布并不经过学习得到,而是固定为线性高斯模型

这一点和VAE有很大不同。VAE里编码器侧的潜在向量的分布是通过模型训练得到的。而扩散模型里,前向加噪过程里的每一步都是基于上一步结果的高斯变换。其中 alpha_t 一般当作超参设置得到。这点对于我们计算扩散模型的证据下界有很大帮助。因为我们可以基于输入x0确切地知道前向过程里的某一步的具体状态,从而监督我们的预测。

基于式31,我们可以递归式地对x0不断加噪变换,得到最终xt的表达式:

图片xt可以写为关于x0的一个高斯分布的采样结果

所以对于式58里噪音匹配项里的监督信号,我们可以重写成以下形式,其中根据式70,我们可以得到q(xt|x0)和q(xt-1|x0)的表达式,而q(xt|xt-1, x0)因为是前向扩散过程,可以应用马尔可夫性质看做q(xt|xt-1)使用式31得到具体表达式。

图片式58里的监督信号可以通过x0计算具体的值

代入每一项q所代表的高斯函数表达式后,我们最后可以得到一个新的高斯分布表达式,其中每一项都是具体可求的:

图片q(xt-1|xt,x0)的解析形式

参考已经证明了前向加噪过程可以写为一个高斯分布了。在扩散模型的初始论文[2]里提到,对于一个连续的高斯扩散过程,其逆过程与前向过程的方程形式(functional form)一致。所以我们将对去噪匹配项里的p_theta(xt-1|xt)也采用高斯分布的形式(更加具体的一些推导放在了末尾的补充里)。注意式58里,对两个高斯分布求KL散度,其解析解的形式如下:

图片两个高斯分布的KL散度解析解

我们现在已知其中一个高斯分布(左侧)的参数,现在如果我们令右侧的高斯分布和左侧高斯分布的方差保持一致。那么优化该KL散度的解析式将简化为以下形式:

图片式58的噪音匹配项简化为最小化前后向均值的预测误差

如此一来式58的噪音匹配项就被简化为最小化前后向均值的预测误差(式92)读者请注意,以下的大一统的三个角度来看待Diffusion model,实质上都是对式92里mu_q的不同变形所推论出来的。 其中mu_q是关于xt, x0的函数,而mu_theta是关于xt和t的函数。其中通过式84,我们有mu_q的准确计算结果,而因为mu_theta是关于xt的函数。我们可以将其写为类似式84的形式(注意,有关为什么可以忽略方差并且让均值选取这个形式放在了最末尾的补充讨论里。但关于这个形式的选择的深层原因实质上开辟了一个全新的领域来研究,并且关于该领域的研究直接导向了扩散模型之后的一系列加速采样技术的出现

图片将后向预测的均值写为类似前向加噪的形式

比较式84与94可知,x_hat是我们通过噪音数据xt来预测原始数据x0的神经网络。那么我们可以将式58里证据下界的噪音匹配项,最终写为

图片噪声匹配项的最终形式

那么,我们最后得到扩散模型的优化,最终表现为训练一个神经网络,以任意时间步的噪音图像为输入,来预测最初的原始图像!此时优化目标转化为了最小化预测误差。同时式58上的对所有时间步的噪音匹配项求和的优化,可以近似为对每一时间步上的预测误差的期望的最小值,而该优化目标可以通过随机采样近似:

图片该优化目标可以通过随机采样实现

Three Equivalent Perspective

为什么Calvin Luo的这篇论文叫做大一统视角来看待扩散模型?以上我们花了不菲的篇幅论证了扩散模型的优化目标可以最终转化为训练一个神经网络在任意时间步从xt预测原始输入x0。以下我们将论述如何通过对mu_q不同的推导得到类似的角度看待扩散模型

首先,我们已经知道给定每个时间步的噪声系数alpha_t之后,我们可以由初始输入x0递归得到xt。同理,给定xt我们也可以求得x0。那么对式69重置后,我们可以得到式115.

图片将式69里的xt和x0关系重置后可得式115

重新将式115代入式84里,我们所得的关于时间步t的真实均值表达式mu_q后,我们可以得到以下推导:

图片在推导真实均值时替换x0

注意在上一次推导的过程中,mu_q里的xt在计算kl散度的解析式时被抵消掉了,而x0我们采取的是用神经网络直接拟合的策略。而在这一次的推导过程中,x0被替换成了关于xt的表达式(关于alpha_bar和epsilon_0)后,我们可以得到mu_q的新的表达式,依旧关于xt,只是不再与x0相关,而是与epsilon_0相关(式124)。其中,和式94一样,我们忽略方差(将其设为与前向一致)并将希望拟合的mu_theta写成与真实均值mu_q一样的形式,只是将epsilon_0替换为神经网络的拟合项后我们可以得到式125。

图片与上次推导时替换x0为神经网络所拟合项一样,这次换为拟合初始噪声项

将我们新得到的两个均值表达式重新代入KL散度的表达式里,xt再次被抵消掉(因为mu_theta和mu_q选取的形式一致)最终只剩下epsilon_0和epsilon_theta的差值。注意式130和式99的相似性!

图片最终对证据下界里的去噪匹配项的优化可以写成关于初始噪声和其拟合项的差的最小化

至此,我们得到了对扩散模型的第二种直观理解。对于一个变分扩散模型VDM,我们优化该模型的证据下界既等价于优化其在所有时间步上对初始图像的预测误差的期望,也等价于优化在所有时间步上对噪声的预测误差的期望! 事实上DDPM采取的做法就是式130的做法(注意DDPM里的表达式实际上用的是epsilon_t,关于这点在文末也会讨论)。

下面笔者将概括第三种看待VDM的推导方式。这种方式主要来自于SongYang博士的系列论文,非常直观。并且该系列论文将扩散模型这种离散的多步去噪过程统一成了一个连续的随机微分方程(SDE)的特殊形式。SongYang博士因此获得了ICLR2021的最佳论文奖!后续来自清华大学的基于将该SDE转化为常微分方程ODE后的采样提速论文,也获得了2022ICLR的最佳论文奖!关于该论文的一些细节和直观理解,SongYang博士在他自己的博客里给出了非常精彩和直观的讲解。有兴趣的读者可以点开本文初始的第二个链接查看。以下只对大一统视角下的第三种视角做简短的概括。

第三种推导方式主要基于Tweedie's formula.该公式主要阐述了对于一个指数家族的分布的真实均值,在给定了采样样本后,可以通过采样样本的最大似然概率(即经验均值)加上一个关于分数(score)预估的校正项来预估。注意score在这里的定义是真实数据分布的对数似然关于输入xt的梯度。即

图片score的定义

根据Tweedie's formula,对于一个高斯变量z~N(mu_z, sigma_z)来说, 该高斯变量的真实均值的预估是:

图片Tweedie's formula对高斯变量的应用

我们知道在训练时,模型的输入xt关于x0的表达式如下

图片上文里的式70

我们也知道根据Tweedie's formula的高斯变量的真实均值预估我们可以得到下式

图片将式70的方差代入Tweedie's formula

那么联立两式的关于均值的表达式后,我们可以得到x0关于score的表达式133

图片将x0写为关于score的表达式

如上一种推导方式所做的一样,再一次重新将x0的表达式代入式84对真实均值mu_q的表达式里:(注意式135到136的变形主要在分子里最右边的alpha_bar_t到alpha_t, 约去了根号下alpha_bar_t-1)

图片将x0的关于score表达式代入式84

同样,将mu_theta采取和mu_q一样的形式,并用神经网络s_theta来近似score后, 我们得到了新的mu_theta的表达式143。

图片关于score的mu_theta的表达式

再再再同样,和上种推导里的做法一样,我们再将新的mu_theta, mu_q代入证据下界里KL散度的损失项我们可以得到一个最终的优化目标

图片将新的mu的表达式代入证据下界的优化目标里

事实上,比较式148和式130的形式,可以说是非常的接近了。那么我们的score function delta_p(xt)和初始噪声epsilon_0是否有关联呢?联立关于x0的两个表达式133和115我们可以得到

图片score function和初始噪声间的关系

读者如果将式151代入148会发现和式130等价!直观上来讲,score function描述的是如何在数据空间里最大化似然概率的更新向量。而又因为初始噪声是在原输入的基础上加入的,那么往噪声的反方向(也是最佳方向)更新实质上等价于去噪的过程。而数学上讲,对score function的建模也等价于对初始噪声乘上负系数的建模!

至此我们终于将扩散模型的三个形式的所有推导整理完毕!即对变分扩散模型VDM的训练等价于训练一个神经网络来预测原输入x0,也等价于预测噪声epsilon, 也等价于预测初始输入在特定时间步的score delta_logp(xt)。

读到这里,相比读者也已经发现,不同的推导所得出的不同结果,都来自于对证据下界里去噪匹配项的不同推导过程。而不同的变形,基本上都是利用了MHVAE里最开始提到的三点基本假设所得。

Drawbacks to Consider

尽管扩散模型在最近两年成功出圈,引爆了业界,学术界甚至普通人对文本生成图像的AI模型的关注,但扩散模型这个体系本身依旧存在着一些缺陷:

  1. 扩散模型本身尽管理论框架已经比较完善,公式推导也十分优美。但仍然非常不直观。最起码从一个完全噪声的输入不断优化的这个过程和人类的思维过程相去甚远。
  2. 扩散模型和GAN或者VAE相比,所学的潜在向量不具备任何语义和结构的可解释性。上文提到了扩散模型可以看做是特殊的MHVAE,但里面每一层的潜在向量间都是线性高斯的形式,变化有限。
  3. 而扩散模型的潜在向量要求维度与输入一致这一点,则更加死地限制住了潜在向量的表征能力。
  4. 扩散模型的多步迭代导致了扩散模型的生成往往耗时良久。

不过学术界对以上的一些难题其实也提出了不少解决方案。比如扩散模型的可解释性问题。笔者最近就发现了一些工作将score-matching直接应用在了普通VAE的潜在向量的采样上。这是一个非常自然的创新点,就和数年前的flow-based-vae一样。而耗时良久的问题,今年ICLR的最佳论文也将采样这个问题加速和压缩到了几十步内就可以生成非常高质量的结果。

但是对于扩散模型在文本生成领域的应用最近似乎还不多,除了prefix-tuning的作者xiang-lisa-li的一篇论文[3]

之外笔者暂未关注到任何工作。而具体来讲,如果将扩散模型直接用在文本生成上,仍有诸多不便。比如输入的尺寸在整个扩散过程必须保持一致就决定了使用者必须事先决定好想生成的文本的长度。而且做有引导的条件生成还好,要用扩散模型训练出一个开放域的文本生成模型恐怕难度不低。

本篇笔记着重的是在探讨大一统角度下的扩散模型推断。但具体对score matching如何训练,如何引导扩散模型生成我们想要的条件分布还没有写出来。笔者打算在下一篇探讨最近一些将扩散模型应用在受控文本生成领域的方法调研里详细记录和比较一下

补充
  • 关于为什么扩散核是高斯变换的扩散过程的逆过程也是高斯变换的问题,来自清华大神的一篇知乎回答里[4] 给出了比较直观的解释。其中第二行是将p_t-1和p_t近似。第三行是对logpt(x_t-1)使用一阶泰勒展开消去了logpt(xt)。第四行是直接代入了q(xt|xt-1)的表达式。于是我们得到了一个高斯分布的表达式。

图片扩散的逆过程也是高斯分布

  • 在式94和式125,我们都将对真实高斯分布q的均值mu_q的近似mu_theta建模成了与我们所推导出的mu_q一致的形式,并且将方差设置为了与q的方差一致的形式。直观上来讲,这样建模的好处很多,一方面是根据KL散度对两个高斯分布的解析式来说,这样我们可以约掉和抵消掉绝大部分的项,简化了建模。另一方面真实分布和近似分布都依赖于xt。在训练时我们的输入就是xt,采取和真实分布形式一样的表达式没有泄漏任何信息。并且在工程上DDPM也验证了类似的简化是事实上可行的。但实际上可以这样做的原因背后是从2021年以来的一系列论文里复杂的数理证明所在解释的目标。 同样引用清华大佬[4]的回答:

图片DDPM里简化去噪的高斯分布的做法其实蕴含着深刻的道理

  • 在DDPM里,其最终的优化目标是epsilon_t而不是epsilon_0。即预测的误差到底是初始误差还是某个时间步上的初始误差。谁对谁错?实际上这个误解来源于我们对xt关于x0的表达式的求解中的误解。从式63开始的连续几步推导,都应用到了一个高斯性质,即两个独立高斯分布的和的均值与方差等于原分布的均值和与方差和。而实质上我们在应用重参数化技巧求xt的过程中,是递归式的不断引入了新的epsilon来替换递归中的x_n里的epsilon。那么到最后,我们所得到的epsilon无非是一个囊括了所有扩散过程中的epsilon。这个噪声即可以说是t,也可以说是0,甚至最准确来说应该不等于任何一个时间步,就叫做噪声就好!

图片DDPM的优化目标

  • 关于对证据下界的不同简化形式。其中我们提到第二种对噪声的近似是DDPM所采用的建模方式。但是对初始输入的近似其实也有论文采用。也就是上文提及的将扩散模型应用在可控文本生成的论文里[3]所采用的形式。该论文每轮直接预测初始Word-embedding。而第三种score-matching的角度可以参照SongYang博士的系列论文[5]来看。里面的优化函数的形式用的是第三种。
  • 本篇笔记着重于讲述扩散模型的变分下界的公式推导,关于扩散模型与能量模型,朗之万动力学,随机微分方程等一系列名词的关系本篇笔记并无涉及。 笔者将在另外一篇笔记里梳理相关的理解。

参考

  1. ^Improving Variational Inference with Inverse Autoregressive Flow https://arxiv.org/abs/1606.04934
  2. ^Deep Unsupervised Learning using Nonequilibrium Thermodynamics https://arxiv.org/abs/1503.03585
  3. ^abDiffusion-LM Improves Controllable Text Generation https://arxiv.org/abs/2205.14217
  4. ^abdiffusion model最近在图像生成领域大红大紫,如何看待它的风头开始超过GAN?- 我想唱high C的回答 - 知乎 https://www.zhihu.com/question/536012286/answer/2533146567
  5. ^SCORE-BASED GENERATIVE MODELING THROUGH STOCHASTIC DIFFERENTIAL EQUATIONS https://arxiv.org/abs/2011.13456


*博客内容为网友个人发布,仅代表博主个人观点,如有侵权请联系工作人员删除。

参与讨论
登录后参与讨论
推荐文章
最近访客