从统一的视角理解扩散模型

说在最前面:最近扩散模型突然很火,甚至出圈到了那些搞艺术搞绘画的圈子里去了。而且这类模型的表现确实是很惊人,不说那些打着open名义却不开源的公司,单说现在最多人用的stable diffusion,输出效果有目共睹。但是这些模型背后的原理和之前的深度学习思路都不太一样,并且原论文的数学内容又非常浓厚,令人难以下咽。所以,本篇正文内容是来自于非常详尽的推导与解释的论文“Understanding Diffusion Models: A Unified Perspective”的翻译,以及一些原文理解起来很怪的地方的自我思考。

另一篇优秀的解读博客:什么是Diffusion模型? - wrong.wang

通篇的通读翻译已经做完,但是细节部分我还在研究,内容随着继续解读和了解会有补充与更正,业内同学请随意提出订正意见以及想法交流。

引言:生成模型

给定一个样本x\boldsymbol{x},这个样本采样于要研究的分布,生成模型的目标是学习如何建立数学模型,用这个模型描述真实的数据分布p(x)p(\boldsymbol{x})。一旦学到了这个模型,我们就可以利用这个近似的模型生成新的想要的样本。进一步,利用一些式子,可以用这个模型评估观测到的或者采样到的数据的似然(likelihood)。【有关“似然”一词的通俗理解,我的类比是,在已知一个事情的先决条件的情况下,这件事发生的概率】【贝叶斯公式:Posterior = Likelihood × Prior ÷ Evidence】

目前有几个比较常见的研究方向,本文只做一些高度概括的介绍。生成对抗网络(GAN)通过对抗学习的方法建模一个采样过程,这个采样是从一个复杂的分布中进行的。另一种“基于似然的(likelihood-based)”生成模型尝试学一个模型,让观测数据在这个模型上获得更高的似然值,包括自回归模型(autoregressive models),标准化流(normalizing flows),变分自编码器(Variational Autoencoders,VAEs)。另一种类似的方法是基于能量的(energy-based)模型,用任意的能量函数学一个分布,之后对他归一化。

基于score的生成模型和本篇主题非常相关,相比于学习如何建模能量函数,这种方法通过神经网络学到基于能量模型的得分(score)。本篇的内容研究了扩散模型,在后文中会看到,这种方法可以通过“基于似然的”也可以通过“基于得分”的视角进行解读。本文尽可能详尽地展开了扩散模型背后的数学细节,即使二年级都可以理解扩散模型以及他的工作原理。

背景知识:ELBO,VAE,Hierarchical VAE

对于许多模态,我们可以把观测到的数据当作是从一个没见过的相关的隐变量(latent variable)中生成或者表示出来的,这个隐变量可以记作随机变量z\boldsymbol{z}。一个很直观的例子是柏拉图的洞穴(Allegory of the Cave),讲的是,有一群人一辈子呆在一个洞穴里,他们只能看到面前洞壁上二维的影子,而这些影子是在一堆火前面的三维物体形成的。对于洞里的人而言,他们能观测到的所有东西实际上是由他们看不到的高维抽象概念决定的。

图:柏拉图的洞穴,来源Wiki百科

类似地,我们在现实世界中遇到的东西也可能是从一些高维函数产生的东西。比如说,有一些表示颜色、大小、形状,等等抽象的表示,然后这些抽象高维表示投影到了三维空间,形成了我们看到的东西,就像洞里的那些人观测到的东西实际上只是三维物体的影子。不过,洞里那些人虽说从没见过或者能理解三维物体,他们还是能思考以及推理一些阴影背后的规律。与此类似的是,我们也可以估计一些隐表示(latent representations)来描述我们观察到的数据。

尽管柏拉图的洞穴描述了这种隐变量背后的思想:隐含的观测不到的东西决定了能观测到的东西,不过这个比喻提醒了我们,在生成模型中我们通常希望学到一种低维的隐表示,而非高维的。因为在缺少很强的先验知识时,很难学习比观测维度还高维的表示。另一方面,学一个低维的隐表示同时可以看作是一种压缩,也许能发现一些对描述观测结果有语义意义的结构。

证据下界(Evidence Lower Bound,ELBO)

从数学角度出发,我们可以想象隐变量和我们观测到的数据服从一个联合分布p(x,z)p(\boldsymbol{x},\boldsymbol{z})。前面说到“基于似然的”生成模型,他能够最大化所有观测到的x\boldsymbol{x}的似然p(x)p(\boldsymbol{x})。有两种方法可以处理这个联合分布,使似然p(x)p(\boldsymbol{x})完全和观测数据相关。

一种方法是边缘化隐变量p(z)p(\boldsymbol{z})(求x\boldsymbol{x}的边缘概率密度):

(1)p(x)=p(x,z)dzp(\boldsymbol{x})=\int p(\boldsymbol{x}, \boldsymbol{z}) d \boldsymbol{z} \tag {1}

或者利用概率链式法则(根据条件概率公式推导得来):

(2)p(x)=p(x,z)p(zx)p(\boldsymbol{x})=\frac{p(\boldsymbol{x}, \boldsymbol{z})}{p(\boldsymbol{z}|\boldsymbol{x})} \tag {2}

不过直接计算并且最大化似然p(x)p(\boldsymbol{x})很难,因为这个过程涉及到公式1中对所有隐变量z\boldsymbol{z}的积分,在复杂模型上这很难做。又或者需要知道公式2里面的实际(ground-truth)隐编码器p(zx)p(\boldsymbol{z}|\boldsymbol{x})【这显然更不可能,以VAE为例,相当于我们要明确知道一个样本要编码成什么隐编码,但是我们又希望模型自己学到这个编码】。不过,通过这两个式子可以导出一个叫做ELBO的项,正如其名,他代表证据(evidence)的下界。这个例子里,证据被量化为观测数据的对数似然。之后,最大化ELBO就成为了优化隐变量模型的一个间接的替代目标。在最佳的情况下,当ELBO被充分的参数化和优化到极致的时候,他就完全等同于证据。形式上,ELBO的公式写作:

(3)Eqϕ(zx)[logp(x,z)qϕ(zx)]\mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})}\left[\log \frac{p(\boldsymbol{x}, \boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})}\right] \tag {3}

为了更清晰的展示ELBO和证据p(x)p(\boldsymbol{x})之间的关系,可以这样写:

(4)logp(x)Eqϕ(zx)[logp(x,z)qϕ(zx)]\log p(\boldsymbol{x}) \geq \mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})}\left[\log \frac{p(\boldsymbol{x}, \boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})}\right] \tag {4}

这里的qϕ(zx)q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})是一个灵活的近似的变分(variational)分布,其参数为ϕ\boldsymbol{\phi},这是我们想要优化的东西。直观地讲,他是一个可参数化的模型,他在观测值为x\boldsymbol{x}时,估计对应隐变量的真实分布。换言之,他希望能逼近一个真实的后验概率p(zx)p(\boldsymbol{z}|\boldsymbol{x})。在探究VAE的时候我们就会看到,通过调整参数ϕ\boldsymbol{\phi}增加下界从而最大化ELBO的时候,我们就能建模真实数据分布并且可以从里面抽取样本,也就是学到了一个生成模型。不过目前我们先深入探究一下为什么要最大化ELBO。

通过公式1可以推出ELBO,过程如下:

上面的推导直接从Jensen不等式得到了一个下界,但是为什么ELBO就是证据的下界?并且只知道ELBO是数据的下界也还是无法理解为什么要把他当作目标函数进行最大化。为了更好的理解证据和ELBO之间的关系,继续用公式2来做另一个推导:

从上面推导的公式15可以很清晰的看到证据等于ELBO+KL散度(近似后验qϕ(zx)q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})和真实后验p(zx)p(\boldsymbol{z} | \boldsymbol{x})的KL散度),而这个KL散度项正是公式8里面被Jensen不等式消掉的东西。理解这一点是理解ELBO与证据的关系以及将ELBO当作目标函数优化的关键。

第一,现在应该能理解为何ELBO是一个下界了:证据比ELBO多出来的就是一个非负的KL散度项,所以ELBO的值永远不会超过证据的值。

第二,来研究一下为什么要最大化ELBO。引入想建模的隐变量z\boldsymbol{z}以后,我们的目标是学到描述观测到的数据背后的隐式结构。换言之,我们想优化这个变分后验qϕ(zx)q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})的参数,使他能完美贴近真实后验p(zx)p(\boldsymbol{z} | \boldsymbol{x}),这个过程可以通过最小化他们的KL散度实现(理想情况下等于0)。只可惜,因为没法获取到真实情况下的p(zx)p(\boldsymbol{z} | \boldsymbol{x}),所以很难直接最小化这个KL散度。不过,注意看公式15等号左边,数据的似然(也就是证据项logp(x)\log p(\boldsymbol{x}))对于ϕ\boldsymbol{\phi}来说一直是个常数项,原因是这个似然是通过边缘化联合分布p(xz)p(\boldsymbol{x} | \boldsymbol{z})里面的z\boldsymbol{z}得来的,因此和ϕ\boldsymbol{\phi}没有关系。既然ELBO+KL散度一直是个常数,那么针对ϕ\boldsymbol{\phi}最大化ELBO就等同于最小化KL散度。因此ELBO可以当作一个替代的目标函数,优化模型学习建模真实后验。越优化ELBO,近似后验就越贴近真实后验。并且因为ELBO学习去近似模型的证据logp(x)\log p(\boldsymbol{x}),所以一旦训练完成,ELBO也可以用来估计观测数据或者生成数据的似然。

变分自编码器(Variational Autoencoders,VAE)

在VAE原始论文里面[1]直接最大化ELBO,这个方法是变分的(variational),因为我们要在一群参数为ϕ\boldsymbol{\phi}的可能的后验分布里面优化出最佳的qϕ(zx)q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})自编码器(autoencoder)这个名字的由来和传统的自动编码模型有关系,这种模型输入数据后,经过一些中间层,再输出这个数据本身。为了清楚的看到一些关联,下面再仔细研究一下ELBO项:

在这里,需要学习一个可以被视作是编码器(encoder)的中间分布qϕ(zx)q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x}),他能将输入转换到一个可能的隐式分布上。同时,再学一个确定(deterministic)函数pθ(xz)p_{\boldsymbol{\theta}}(\boldsymbol{x} | \boldsymbol{z})来把隐向量z\boldsymbol{z}转换成观测数据x\boldsymbol{x},这个称为解码器(decoder)

公式19里面的两项都有直观解释:第一项从变分分布衡量解码器的重建似然,确保学到的分布能够有效建模隐分布,从而确保原始数据能重新生成。第二项衡量了学到的变分分布与隐变量的先验信念(prior belief)之间有多像,最小化这一项能够促使编码器真正学到分布,而非塌缩到Dirac delta函数上【防止模式崩溃】。最大化ELBO等同于最大化第一项,最小化第二项。

VAE的一个关键特征是ELBO是如何在参数ϕ\boldsymbol{\phi}θ\boldsymbol{\theta}上被联合优化的。VAE的编码器通常建模一个具有对角协方差的多元高斯模型,而先验通常是一个标准多元高斯模型:

(20)qϕ(zx)=N(z;μϕ(x),σϕ2(x)I)q_\boldsymbol{\phi}(\boldsymbol{z} | \boldsymbol{x}) =\mathcal{N}\left(\boldsymbol{z} ; \boldsymbol{\mu}_\boldsymbol{\phi}(\boldsymbol{x}), \boldsymbol{\sigma}_\boldsymbol{\phi}^2(\boldsymbol{x}) \mathbf{I}\right) \tag{20}

(21)p(z)=N(z;0,I)p(\boldsymbol{z}) = \mathcal{N}(\boldsymbol{z} ; \boldsymbol{0} , \mathbf{I})\tag{21}

然后ELBO的KL散度项可以解析求得,公式19里面那个reconstruction项可以用蒙特卡洛估计法近似出来。由此目标函数可以写成:

对数据集内任意的观测值x\boldsymbol{x},其中的隐变量{z(l)}l=1L\left\{\boldsymbol{z}^{(l)}\right\}_{l=1}^L是从qϕ(zx)q_\boldsymbol{\phi}(\boldsymbol{z}|\boldsymbol{x})中抽样的。但是这种方法会产生一个问题:对于计算loss时的每一个z(l)\boldsymbol{z}^{(l)}都是从随机采样中产生的,这个过程不可微。不过当qϕ(zx)q_\boldsymbol{\phi}(\boldsymbol{z} | \boldsymbol{x})设计用来建模一个特定的分布(包括多元高斯分布)时,这个问题可以通过重参数化reparameterization)的技巧来解决。

重参数化将一个随机变量表达成一个噪声变量的确定的函数,这样可以对非随机项使用梯度下降法进行优化。比如说,从一个均值为μ\mu,方差为σ2\sigma^2的正态分布xN(x;μ,σ2)x \sim \mathcal{N}(x ; \mu, \sigma^2)中采样的过程可以写作:

x=μ+σϵ with ϵN(ϵ;0,I)x=\mu+\sigma \epsilon \quad \text { with } \epsilon \sim \mathcal{N}(\epsilon ; 0, \mathrm{I})

也就是说,任意的高斯分布都可以转换成均值从0用加法移到μ\mu、方差从1拉伸到σ2\sigma^2的标准正态分布(这里的ϵ\epsilon是一个抽样)。由此,通过重采样的办法,从任意的高斯分布抽样都可以从标准正态分布中抽样实现,只是将抽样结果缩放一些方差和平移一下均值即可。

在VAE里,每个z\boldsymbol{z}通过输入x\boldsymbol{x}和辅助的噪声变量ϵ\boldsymbol{\epsilon}经过一个确定的函数算得:

z=μϕ(x)+σϕ(x)ϵ with ϵN(ϵ;0,I)\boldsymbol{z}=\boldsymbol{\mu}_\boldsymbol{\phi}(\boldsymbol{x})+\boldsymbol{\sigma}_\boldsymbol{\phi}(\boldsymbol{x}) \odot \boldsymbol{\epsilon} \quad \text { with } \boldsymbol{\epsilon} \sim \mathcal{N}(\boldsymbol{\epsilon} ; \mathbf{0}, \mathbf{I})

其中\odot代表逐元素乘积。对于这个重参数化的z\boldsymbol{z},对于ϕ\boldsymbol{\phi}的梯度就可以计算出来,并且去优化μϕ\boldsymbol{\mu}_\boldsymbol{\phi}σϕ\boldsymbol{\sigma}_\boldsymbol{\phi}了。因此VAE实现了在ϕ\boldsymbol{\phi}θ\boldsymbol{\theta}上使用重参数化和蒙特卡洛估计优化ELBO。

训练完VAE之后,生成新的数据就可以直接从隐空间p(z)p(\boldsymbol{z})抽样并且用解码器算出来。VAE很有意思,尤其在z\boldsymbol{z}的维度小于输入x\boldsymbol{x}的时候,这时候他可能学到紧密且有用的表示。并且当模型学到一个有语义意义的隐空间之后,可以在解码之前编辑隐向量,更精确的控制数据的生成。

图1:VAE的图示。其中,编码器q(zx)q(\boldsymbol{z}|\boldsymbol{x})定义了对于观测数据x\boldsymbol{x},隐变量z\boldsymbol{z}的分布,解码器p(xz)p(\boldsymbol{x}|\boldsymbol{z})将隐变量解码成为观测数据。

层次化的变分自编码器(Hierarchical Variational Autoencoders,HVAE)

HVAE[2, 3]是VAE的一个推广形式,他将隐变量扩展到了多个层级,模型里的隐变量本身又是从更高级更抽象的隐变量中生成的。直觉上说,柏拉图洞穴里那些人把三维物体当作是生成二维影子的隐变量,就像是我们把更高级的抽象东西当作生成我们观测到的三维物体的隐变量。所以,在柏拉图洞穴里面的人可以把他们的观测情况建模成一个有两层(或者更多层)隐层次的过程(理解为,高级抽象->产生三维物体->产生二维影子)。

图2:有TT个层级的MHVAE。生成过程用一个马尔科夫链建模表示,其中每个隐变量zt\boldsymbol{z}_t仅从前一个zt+1\boldsymbol{z}_{t+1}生成而来。

对于有TT个层级的一般HVAE,每个隐变量都取决于前面所有隐变量,在这篇里,我们关注一个特殊情况:马尔可夫HVAE(Markovian HVAE,MHVAE)。在MHVAE中,生成过程是一个马尔科夫链,即,每一次往下个层级的转移过程是马尔可夫过程,即,解码每个隐变量zt\boldsymbol{z}_t仅取决于前一个zt+1\boldsymbol{z}_{t+1}。直观上讲,可以看作是把VAE堆叠起来,如图2所示。另一个符合这个模型描述的是Recursive VAE。这个联合分布以及这个MHVAE的后验的数学表达为:

之后,ELBO很容易扩展一下:

把联合分布(公式23)和后验(公式24)替换到公式28里面:

下文会阐述,当我们研究变分扩散模型时,这个目标函数会进一步分解为可解释的成分。

变分扩散模型(Variational Diffusion Model,VDM)

研究VDM[4, 5, 6]的一个最简单的方法是把他看作一个MHVAE,但是有3个关键的限定条件:

  1. 隐变量维度严格等于数据维度
  2. 每个时间步的隐编码器的结构不是学到的,是一个预定义好的线性高斯模型。换句话说,是一个关于前一个时间步的输出的高斯分布
  3. 隐编码器的高斯分布参数随着时间步而改变,让最后一个时间步TT的隐变量分布是一个标准正态分布

而且还要保持着像MHVAE那种层级之间转移的马尔可夫性。

接下来进一步阐述这些假设的含义。

从第一项约束看,我们可以把真实样本和隐变量一起表示为xt\boldsymbol{x}_t,当t=0t=0时代表真实样本,t[1,T]t \in [1, T]时代表对应tt层的隐变量。VDM的后验和MHVAE的后验一样(公式24),但这里可以表示成为:

(30)q(x1:Tx0)=t=1Tq(xtxt1)q\left(\boldsymbol{x}_{1: T} | \boldsymbol{x}_0\right)=\prod_{t=1}^T q\left(\boldsymbol{x}_t | \boldsymbol{x}_{t-1}\right) \tag{30}

第二项假设,我们已知编码器里的每个隐变量的分布是有关前一层隐变量的高斯分布。与MHVAE不同,每个时间步tt的编码器结构不是学到的,而是一个固定的线性高斯模型,其中的均值和标准差是由超参数设置的[5],或者作为参数学到的[6]。用均值为μt(xt)=αtxt1\boldsymbol{\mu}_t(\boldsymbol{x}_t)=\sqrt{\alpha_t} \boldsymbol{x}_{t-1},方差为Σt(xt)=(1αt)I\boldsymbol{\Sigma}_t(\boldsymbol{x}_t)=(1-\alpha_t) \mathbf{I}对这个高斯编码器做参数化,要选择系数的形式,以确保隐变量的方差保持在相似的尺度。换言之,编码过程是保持方差variance-preserving)的。注意:交替进行高斯参数化也是可以的,并且能得到类似的推导。主要的结论是,αt\alpha_t是一个随着层深tt变化的参数(可以是可学习的)。编码器的数学表述为:

(31)q(xtxt1)=N(xt;αtxt1,(1αt)I)q(\boldsymbol{x}_t | \boldsymbol{x}_{t-1})=\mathcal{N}\left(\boldsymbol{x}_t ; \sqrt{\alpha_t} \boldsymbol{x}_{t-1},(1-\alpha_t) \mathbf{I}\right) \tag{31}

第三项假设,可以看到αt\alpha_t是一个随时间演化的固定或可学习的设计好的结构,以使得最终的隐分布p(xT)p(\boldsymbol{x}_T)是标准高斯分布。之后就可以升级MHVAE的联合分布(公式23)来获得VDM的联合分布:

总的来说,这组假设描述了输入图像随时间的稳定加噪过程,通过添加高斯噪声来逐渐破坏图像,直到完全变成纯高斯噪声。该过程如图3所示。

图3:VDM的示意图。x0\boldsymbol{x}_0代表真实观测到的数据,例如普通图像。xT\boldsymbol{x}_T代表纯高斯噪声,xt\boldsymbol{x}_t是中间的带噪声的x0\boldsymbol{x}_0。每一个q(xtxt1)q(\boldsymbol{x}_t|\boldsymbol{x}_{t-1})是一个高斯分布模型,他使用前一个状态的输出作为均值。

需要注意这里的编码器分布q(xtxt1)q(\boldsymbol{x}_t|\boldsymbol{x}_{t-1})不再是带参数ϕ\boldsymbol{\phi}的了,因为模型是完全由高斯模型定义的,而且模型每一步的均值方差也是预定好的。因此VDM中我们只关心pθ(xt1xt)p_\boldsymbol{\theta}(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t),这样就能生成新数据了。VDM优化好以后,采样就是简单的从高斯噪声p(xT)p(\boldsymbol{x}_T)采样,然后迭代运行去噪pθ(xt1xt)p_\boldsymbol{\theta}(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t)的过程,TT步之后就可以获得x0\boldsymbol{x}_0

就像众多HVAE一样,VDM也可以通过最大化ELBO优化,具体推导如下:

这个形式的ELBO的每个部分有如下的解释:

  1. Eq(x1x0)[logpθ(x0x1)]\mathbb{E}_{q\left(\boldsymbol{x}_1 | \boldsymbol{x}_0\right)}\left[\log p_{\boldsymbol{\theta}}(\boldsymbol{x}_0 | \boldsymbol{x}_1)\right]可以解释为重建项reconstruction term),用来预测给定第一步的隐变量下原始样本的对数概率。这一项也在原版VAE里面出现过,都可以用类似的办法训练。
  2. Eq(xT1x0)[DKL(q(xTxT1)p(xT))]\mathbb{E}_{q(\boldsymbol{x}_{T-1} | \boldsymbol{x}_0)}\left[D_{\mathrm{KL}}\left(q(\boldsymbol{x}_T | \boldsymbol{x}_{T-1}) \| p(\boldsymbol{x}_T)\right)\right]先验匹配项prior matching term),在最终的隐分布和高斯先验匹配的时候最小。这一项没有可学习的参数,不需要优化。并且我们假设了一个足够大的TT来使最终的分布是高斯分布,所以这一项最终会化为0。
  3. Eq(xt1,xt+1x0)[DKL(q(xtxt1)pθ(xtxt+1))]\mathbb{E}_{q(\boldsymbol{x}_{t-1}, \boldsymbol{x}_{t+1} | \boldsymbol{x}_0)}\left[D_{\mathrm{KL}}\left(q(\boldsymbol{x}_t | \boldsymbol{x}_{t-1}) \| p_{\boldsymbol{\theta}}(\boldsymbol{x}_t | \boldsymbol{x}_{t+1})\right)\right]一致项consistency term),让xt\boldsymbol{x}_t在前向和后向时的分布保持一致。即,对于每个中间步骤,对于带噪声图像的去噪步骤应当与干净图像的加噪步骤相匹配。这个过程反应在KL散度上面。这一项在训练pθ(xtxt+1)p_{\boldsymbol{\theta}}(\boldsymbol{x}_t | \boldsymbol{x}_{t+1})与高斯分布q(xtxt1)q(\boldsymbol{x}_t | \boldsymbol{x}_{t-1})(公式31)匹配时被最小化。

直观来看,这个关于ELBO的解释如图4所示。优化VDM的成本开销主要在第三项上,因为必须在所有时间步tt上面做优化。

图4:在第一个推导里面,对于每一个中间的xt\boldsymbol{x}_t,通过确保隐变量的后验pθ(xtxt+1)p_{\boldsymbol{\theta}}(\boldsymbol{x}_t | \boldsymbol{x}_{t+1})和隐变量上的高斯加噪q(xtxt1)q(\boldsymbol{x}_t | \boldsymbol{x}_{t-1})之间匹配,VDM可以被优化。反映在图中就是对于每一个中间的xt\boldsymbol{x}_t,最小化粉色和绿色箭头表示的分布之间的差距。

在这个推导下,ELBO的所有项都用期望计算,因此可以使用蒙特卡洛估计进行近似。但是直接用我们刚刚推导出来的东西去优化ELBO并不是最优的,因为一致项在每个时间步下是两个随机变量的期望{xt1,xt+1}\{\boldsymbol{x}_{t-1}, \boldsymbol{x}_{t+1}\},他的蒙特卡洛估计的方差可能高于只有一个随机变量的情况。当TT很大时,在T1T-1个一致项上面累积,ELBO的最终估计值可能方差很大。

接下来尝试让每个期望都只和一个随机变量相关,再推导一下ELBO。一个重要的想法是,编码器的转移可以重写成q(xtxt1)=q(xtxt1,x0)q(\boldsymbol{x}_t | \boldsymbol{x}_{t-1})=q(\boldsymbol{x}_t | \boldsymbol{x}_{t-1}, \boldsymbol{x}_0),根据马尔可夫性质,其中的额外条件项是多余的。之后,利用贝叶斯法则,可以重写为

(46)q(xtxt1,x0)=q(xt1xt,x0)q(xtx0)q(xt1x0)q(\boldsymbol{x}_t | \boldsymbol{x}_{t-1}, \boldsymbol{x}_0)=\frac{q(\boldsymbol{x}_{t-1} | \boldsymbol{x}_t, \boldsymbol{x}_0) q(\boldsymbol{x}_t | \boldsymbol{x}_0)}{q(\boldsymbol{x}_{t-1} | \boldsymbol{x}_0)} \tag{46}

有了这个新式子以后,可以尝试重新从公式37开始,写一下ELBO的推导过程:

由此我们就推出了一个估算方差更小的ELBO,其中每个时间步的每一项都是最多一个随机变量的期望。这个表达中的每一项也有很好的解释:

  1. Eq(x1x0)[logpθ(x0x1)]\mathbb{E}_{q\left(\boldsymbol{x}_1 | \boldsymbol{x}_0\right)}\left[\log p_{\boldsymbol{\theta}}(\boldsymbol{x}_0 | \boldsymbol{x}_1)\right]是重建项,与原版VAE的ELBO类似,该项可以使用蒙特卡洛估计进行近似和优化。
  2. DKL(q(xTx0)p(xT))D_{\mathrm{KL}}(q(\boldsymbol{x}_T | \boldsymbol{x}_{0}) \| p(\boldsymbol{x}_T))表示加噪的输入的最终分布与标准高斯先验的接近程度。它没有可训练的参数,在我们的假设下等于0。
  3. Eq(xtx0)[DKL(q(xt1xt,x0)pθ(xt1xt))]\mathbb{E}_{q(\boldsymbol{x}_{t} | \boldsymbol{x}_0)}\left[D_{\mathrm{KL}}\left(q(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t}, \boldsymbol{x}_{0} ) \| p_{\boldsymbol{\theta}}(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t})\right)\right]去噪匹配项denoising matching term)。我们想要得到的去噪转移pθ(xt1xt)p_{\boldsymbol{\theta}}(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t})作为一个好处理的真实去噪转移q(xt1xt,x0)q(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t}, \boldsymbol{x}_{0})的近似进行学习。q(xt1xt,x0)q(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t}, \boldsymbol{x}_{0})可以作为一个真实的信号,定义了在已知最终完全去掉噪声的x0\boldsymbol{x}_{0}应该是什么样子的条件下,如何给一个带噪图像xt\boldsymbol{x}_{t}去噪。当两个去噪步骤非常匹配时,衡量他们之间的KL散度,这一项得以最小化。

在两个ELBO推导过程(公式45和公式58)中,只用到了马尔可夫假设。因此,这些公式适用于任意的MHVAE。此外,当T=1T=1时,VDM的两种ELBO的解释都准确地再现了原版VAE的ELBO方程(公式19)。

在这个ELBO的推导过程中,优化成本大部分花在了求和项也就是重建项(不应该是去噪匹配项吗?)上面。而由于同时学习编码器导致复杂度增加,每一个KL散度项DKL(q(xt1xt,x0)pθ(xt1xt))D_{\mathrm{KL}}\left(q(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t}, \boldsymbol{x}_{0} ) \| p_{\boldsymbol{\theta}}(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t})\right)对于任意的复杂MHVAE的任意的后验而言很难最小化。在VDM中,我们可以利用高斯转移假设使优化易于处理。利用贝叶斯定律,可以得到:

q(xt1xt,x0)=q(xtxt1,x0)q(xt1x0)q(xtx0)q\left(\boldsymbol{x}_{t-1} | \boldsymbol{x}_t, \boldsymbol{x}_0\right)=\frac{q\left(\boldsymbol{x}_t | \boldsymbol{x}_{t-1}, \boldsymbol{x}_0\right) q\left(\boldsymbol{x}_{t-1} | \boldsymbol{x}_0\right)}{q\left(\boldsymbol{x}_t | \boldsymbol{x}_0\right)}

从编码器转移的假设(公式31)可以知道q(xtxt1,x0)=q(xtxt1)=N(xt;αtxt1,(1αt)I)q\left(\boldsymbol{x}_t | \boldsymbol{x}_{t-1}, \boldsymbol{x}_0\right)=q\left(\boldsymbol{x}_t | \boldsymbol{x}_{t-1}\right)=\mathcal{N}\left(\boldsymbol{x}_t ; \sqrt{\alpha_t} \boldsymbol{x}_{t-1},\left(1-\alpha_t\right) \mathbf{I}\right),接下来要做的就是推导q(xtx0)q(\boldsymbol{x}_t | \boldsymbol{x}_{0})q(xt1x0)q(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{0})的形式。而VDM编码器转移的模型是一个线性高斯模型,这就让事情好办了起来。回忆一下重参数化的技巧,从xtq(xtxt1)\boldsymbol{x}_t \sim q(\boldsymbol{x}_t | \boldsymbol{x}_{t-1})中采样就可以写为:

(59)xt=αtxt1+1αtϵ with ϵN(ϵ;0,I)\boldsymbol{x}_t=\sqrt{\alpha_t} \boldsymbol{x}_{t-1}+\sqrt{1-\alpha_t} \boldsymbol{\epsilon} \quad \text { with } \boldsymbol{\epsilon} \sim \mathcal{N}(\boldsymbol{\epsilon} ; \mathbf{0}, \mathbf{I}) \tag{59}

与之类似,从xt1q(xt1xt2)\boldsymbol{x}_{t-1} \sim q(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t-2})中采样就可以写为:

(60)xt1=αt1xt2+1αt1ϵ with ϵN(ϵ;0,I)\boldsymbol{x}_{t-1}=\sqrt{\alpha_{t-1}} \boldsymbol{x}_{t-2}+\sqrt{1-\alpha_{t-1}} \boldsymbol{\epsilon} \quad \text { with } \boldsymbol{\epsilon} \sim \mathcal{N}(\boldsymbol{\epsilon} ; \mathbf{0}, \mathbf{I}) \tag{60}

之后,q(xtx0)q(\boldsymbol{x}_{t} | \boldsymbol{x}_{0})的形式就可以用重参数化技巧递归着推出来。假设我们现在有2T2T个随机噪声变量{ϵt,ϵt}t=0TiidN(ϵ;0,I)\left\{\boldsymbol{\epsilon}_t^*, \boldsymbol{\epsilon}_t\right\}_{t=0}^T \stackrel{\mathrm{iid}}{\sim} \mathcal{N}(\boldsymbol{\epsilon} ; \mathbf{0}, \mathbf{I}),之后对于任意一个样本xtq(xtx0)\boldsymbol{x}_t \sim q(\boldsymbol{x}_t | \boldsymbol{x}_{0}),可以把上面公式59和公式60对应的过程写为:

在公式64里面应用了独立随机高斯变量求和的方法,即均值相加,方差相加。将1αtϵt1\sqrt{1-\alpha_t} \boldsymbol{\epsilon}_{t-1}^*解释为一个从高斯分布N(0,(1αt)I)\mathcal{N}(\mathbf{0},(1-\alpha_t) \mathbf{I})中抽样的样本,αtαtαt1ϵt2\sqrt{\alpha_t-\alpha_t\alpha_{t-1}} \boldsymbol{\epsilon}_{t-2}^*是一个从高斯分布N(0,(αtαtαt1)I)\mathcal{N}(\mathbf{0},(\alpha_t-\alpha_t\alpha_{t-1}) \mathbf{I})抽样的样本,这样就可以把他们的和当作一个从高斯分布N(0,(1αt+αtαtαt1)I)=N(0,(1αtαt1)I)\mathcal{N}(\mathbf{0},(1-\alpha_t+\alpha_t-\alpha_t \alpha_{t-1}) \mathbf{I})=\mathcal{N}(\mathbf{0},(1-\alpha_t \alpha_{t-1}) \mathbf{I})中采样的随机变量(对应公式63-66的过程)。一个从这个分布抽样的样本可以用重参数化写作1αtαt1ϵt2\sqrt{1-\alpha_t\alpha_{t-1}} \boldsymbol{\epsilon}_{t-2},对应公式66。

图5:一个低方差的优化VDM的替代方案。根据贝叶斯公式计算真实的去噪的形式q(xt1xt,x0)q(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t}, \boldsymbol{x}_{0}),最小化他和我们估计的去噪分布pθ(xt1xt)p_\boldsymbol{\theta}(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t})的KL散度。表现在图中,就是将绿色箭头表示的分布与粉色箭头表示的分布进行匹配,每个粉红色箭头也必须出自条件项x0\boldsymbol{x}_{0}

由此便获得了q(xtx0)q(\boldsymbol{x}_{t} | \boldsymbol{x}_{0})的高斯分布的形式。这个推导亦可以改一下,然后表示q(xt1x0)q(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{0})。目前已知了q(xtx0)q(\boldsymbol{x}_{t} | \boldsymbol{x}_{0})q(xt1x0)q(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{0})的形式,进而就可以代入贝叶斯公式,求出q(xt1xt,x0)q(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t}, \boldsymbol{x}_{0})的形式:

在公式75里,C(xt,x0)C(\boldsymbol{x}_{t}, \boldsymbol{x}_{0})只和xt,x0,α\boldsymbol{x}_{t}, \boldsymbol{x}_{0}, \alpha有关,是一个对于xt1\boldsymbol{x}_{t-1}而言是常数项的东西。这一项隐式地放在了公式83到公式84的过程里,平方项得到补全。

由此可以看到,对于每一步,xt1q(xt1xt,x0)\boldsymbol{x}_{t-1} \sim q(\boldsymbol{x}_{t-1} | \boldsymbol{x}_t, \boldsymbol{x}_0)都是均值为μq(xt,x0)\boldsymbol{\mu}_q(\boldsymbol{x}_t, \boldsymbol{x}_0)(为xt\boldsymbol{x}_tx0\boldsymbol{x}_0的函数),方差为Σq(t)\boldsymbol{\Sigma}_q(t)(是一个系数为α\alpha的函数)的正态分布。这些系数α\alpha是已知的,并且在每个时间步中是固定的,他们也是模型设定好然后不变的超参数,或者是一个对其建模的网络模型当前的预测输出。从公式84可以重新写出方差的公式为Σq(t)=σq2(t)I\boldsymbol{\Sigma}_q(t)=\sigma_q^2(t) \mathbf{I},其中

(85)σq2(t)=(1αt)(1αˉt1)1αˉt\sigma_q^2(t)=\frac{\left(1-\alpha_t\right)\left(1-\bar{\alpha}_{t-1}\right)}{1-\bar{\alpha}_t} \tag{85}

为了尽量匹配近似的去噪转移步骤pθ(xt1xt)p_{\boldsymbol{\theta}}(\boldsymbol{x}_{t-1} | \boldsymbol{x}_t)和真实的去噪转移步骤q(xt1xt,x0)q(\boldsymbol{x}_{t-1} | \boldsymbol{x}_t, \boldsymbol{x}_0),我们也可以把他弄成高斯模型。并且,因为所有的α\alpha在每个时间步下已知是定值,所以近似去噪转移步骤的方差可以构建为Σq(t)=σq2(t)I\boldsymbol{\Sigma}_q(t)=\sigma_q^2(t) \mathbf{I}。由于pθ(xt1xt)p_{\boldsymbol{\theta}}(\boldsymbol{x}_{t-1} | \boldsymbol{x}_t)并不以x0\boldsymbol{x}_0为条件,均值μθ(xt,t)\boldsymbol{\mu}_{\boldsymbol{\theta}}\left(\boldsymbol{x}_t, t\right)必须参数化成xt\boldsymbol{x}_t的函数。

回忆一下,两个高斯分布的KL散度

当我们设置两个高斯分布的方差让他俩尽可能匹配,优化KL散度项即为最小化两个分布的均值的差距:

这里μq\boldsymbol{\mu}_qμq(xt,x0)\boldsymbol{\mu}_q(\boldsymbol{x}_t, \boldsymbol{x}_0)的简记,μθ\boldsymbol{\mu}_\boldsymbol{\theta}μθ(xt,t)\boldsymbol{\mu}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)的简记。换言之,我们要优化μθ(xt,t)\boldsymbol{\mu}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)来匹配μq(xt,x0)\boldsymbol{\mu}_q(\boldsymbol{x}_t, \boldsymbol{x}_0),这一项在公式84里,形式如下:

(93)μq(xt,x0)=αt(1αˉt1)xt+αˉt1(1αt)x01αˉt\boldsymbol{\mu}_q\left(\boldsymbol{x}_t, \boldsymbol{x}_0\right)=\frac{\sqrt{\alpha_t}\left(1-\bar{\alpha}_{t-1}\right) \boldsymbol{x}_t+\sqrt{\bar{\alpha}_{t-1}}\left(1-\alpha_t\right) \boldsymbol{x}_0}{1-\bar{\alpha}_t} \tag{93}

因为μθ(xt,t)\boldsymbol{\mu}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)也取决于xt\boldsymbol{x}_t,所以可以写成下面这种形式,以匹配μq(xt,x0)\boldsymbol{\mu}_q(\boldsymbol{x}_t, \boldsymbol{x}_0)

(94)μθ(xt,t)=αt(1αˉt1)xt+αˉt1(1αt)x^θ(xt,t)1αˉt\boldsymbol{\mu}_{\boldsymbol{\theta}}\left(\boldsymbol{x}_t, t\right)=\frac{\sqrt{\alpha_t}\left(1-\bar{\alpha}_{t-1}\right) \boldsymbol{x}_t+\sqrt{\bar{\alpha}_{t-1}}\left(1-\alpha_t\right) \hat{\boldsymbol{x}}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)}{1-\bar{\alpha}_t} \tag{94}

其中x^θ(xt,t)\hat{\boldsymbol{x}}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)由神经网络参数化,用来从带噪图像xt\boldsymbol{x}_t和时刻tt预测x0\boldsymbol{x}_0。之后,优化问题就简化为:

因此,优化VDM成为了训练一个神经网络从任意的带噪图像预测原始的真实图像的问题[5]。并且最小化之前推导出的ELBO(公式58)的各个噪音级别上的求和项也可以通过最小化所有时间步的期望的方法近似,并且可以通过在时间步上随机采样进行优化:

(100)argminθEtU{2,T}[Eq(xtx0)[DKL(q(xt1xt,x0)pθ(xt1xt))]]\underset{\boldsymbol{\theta}}{\arg \min } \mathbb{E}_{t \sim U\{2, T\}}\left[\mathbb{E}_{q\left(\boldsymbol{x}_t | \boldsymbol{x}_0\right)}\left[D_{\mathrm{KL}}\left(q\left(\boldsymbol{x}_{t-1} | \boldsymbol{x}_t, \boldsymbol{x}_0\right) \| p_{\boldsymbol{\theta}}\left(\boldsymbol{x}_{t-1} | \boldsymbol{x}_t\right)\right)\right]\right] \tag{100}

学习扩散噪声参数

下面看一下VDM的噪声参数是如何联合学习的。一种可行的方法是使用一个参数为η\boldsymbol{\eta}的神经网络α^η(t)\hat{\alpha}_{\boldsymbol{\eta}}(t)去建模αt{\alpha}_t。但是这种方法并不高效,因为要连续推理很多次,在每个时间步tt上计算αˉt\bar{\alpha}_t。虽然用缓存可以减少这种计算成本,但我们也可以推导出另一种方法来学习扩散噪声参数。通过将公式85中的方差代入到公式99中的每个时间步,可以得到:

看公式70,q(xtx0)q\left(\boldsymbol{x}_t | \boldsymbol{x}_0\right)是一个形式为N(xt;αˉtx0,(1αˉt)I)\mathcal{N}\left(\boldsymbol{x}_t ; \sqrt{\bar{\alpha}_t} \boldsymbol{x}_0,\left(1-\bar{\alpha}_t\right) \mathbf{I}\right)的高斯分布。根据信噪比的定义,SNR=μ2σ2\mathrm{SNR}=\frac{\mu^2}{\sigma^2},我们可以写出每个时间步tt的信噪比:

(109)SNR(t)=αˉt1αˉt\operatorname{SNR}(t)=\frac{\bar{\alpha}_t}{1-\bar{\alpha}_t} \tag{109}

之后,公式108(和公式99)可以简化为:

正如其名,信噪比表示的是原始信号和所包含噪声的比例,信噪比高代表信号更强,低代表噪声更强。在扩散模型里面,我们需要信噪比随着时间步tt的增长而单调递减,这就是让输入xt\boldsymbol{x}_t随时间逐渐变成噪声,直到最后在t=Tt=T时刻变成标准正态分布。

简化公式110的目标之后,我们可以直接在每个时间步用一个神经网络参数化信噪比,并且和扩散模型一起联合学习。因为信噪比是随时间单调递减的,可以表示为:

(111)SNR(t)=exp(ωη(t))\operatorname{SNR}(t)=\exp \left(-\omega_{\boldsymbol{\eta}}(t)\right) \tag{111}

这里ωη(t)\omega_{\boldsymbol{\eta}}(t)是一个参数为η\boldsymbol{\eta}的单调递增神经网络模型,给他加负号就是单调递减方程,其中指数项使得结果始终为正。公式100里面的目标现在也需要在η\boldsymbol{\eta}上进行优化。结合公式111对信噪比的参数化以及公式109对信噪比的定义,也可以将αˉt\bar{\alpha}_t1αˉt1-\bar{\alpha}_t的更优雅的形式显式表达:

这些项在各种计算中都有用,例如在优化时,他们可以用来通过重参数化从输入x0\boldsymbol{x}_0创建任意的噪声xt\boldsymbol{x}_t,如公式69所示。

三个等价表示

在前面的证明中,VDM可以通过简单的训练一个神经网络来预测原始的自然图像x0\boldsymbol{x}_0,这个输出是通过任意的噪声图像xt\boldsymbol{x}_t和对应的时间步tt获得的。但是x0\boldsymbol{x}_0有另外两个等价的参数化方式,进一步能带来两个对于VDM的解释。

首先,利用重参数化技巧,在对于q(xtx0)q(\boldsymbol{x}_t | \boldsymbol{x}_0)的推导中,公式69可以重新排列一下,写成

(115)x0=xt1αˉtϵ0αˉt\boldsymbol{x}_0=\frac{\boldsymbol{x}_t-\sqrt{1-\bar{\alpha}_t} \boldsymbol{\epsilon}_0}{\sqrt{\bar{\alpha}_t}} \tag{115}

代入到前面推导的真实去噪转移均值μq(xt,x0)\boldsymbol{\mu}_q(\boldsymbol{x}_t, \boldsymbol{x}_0)里面(公式93),可以重新这样推导:

由此,近似去噪转移均值μθ(xt,t)\boldsymbol{\mu}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)可以写作:

(125)μθ(xt,t)=1αtxt1αt1αˉtαtϵ^θ(xt,t)\boldsymbol{\mu}_{\boldsymbol{\theta}}\left(\boldsymbol{x}_t, t\right)=\frac{1}{\sqrt{\alpha_t}} \boldsymbol{x}_t-\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t} \sqrt{\alpha_t}} \hat{\boldsymbol{\epsilon}}_{\boldsymbol{\theta}}\left(\boldsymbol{x}_t, t\right) \tag{125}

以及对应的优化问题(公式95)变成:

这里,ϵ^θ(xt,t)\hat{\boldsymbol{\epsilon}}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)是一个神经网络,用来学习预测源噪声ϵ0N(ϵ;0,I)\epsilon_0 \sim \mathcal{N}(\boldsymbol{\epsilon} ; \mathbf{0}, \mathbf{I}),这个决定了从x0\boldsymbol{x}_0过来的xt\boldsymbol{x}_t。由此我们展示了,通过预测原始图像x0\boldsymbol{x}_0学习一个VDM,等价于学习预测噪声。但是从经验上来说,一些工作[5, 7]发现预测噪声的性能更好。

为了推导第三种常见的VDM的表示,要引入Tweedie公式[8],Tweedie公式表示,给定从指数族分布中抽取的样本,这个分布的真实平均值可以通过样本的最大似然估计值(又称经验平均值)加上与估计分数有关的一些修正项来估计。对于只有一个观测样本的情况,经验平均值只是样本本身。如果观测样本都位于某个分布的一端,则负分数变大,并将样本的朴素最大似然估计值往真实平均值修正,用于减轻样本偏差。

从数学上来说,对于一个高斯变量zN(z;μz,Σz)z \sim \mathcal{N}(z ; \boldsymbol{\mu}_z, \boldsymbol{\Sigma}_z),Tweedie公式表示

E[μzz]=z+Σzzlogp(z)\mathbb{E}[\boldsymbol{\mu}_z | \boldsymbol{z}]=\boldsymbol{z}+\boldsymbol{\Sigma}_z \nabla_{\boldsymbol{z}} \log p(\boldsymbol{z})

在这里,我们用它来预测给定样本xt\boldsymbol{x}_t的真实后验均值。从公式70已知:

q(xtx0)=N(xt;αˉtx0,(1αˉt)I)q\left(\boldsymbol{x}_t | \boldsymbol{x}_0\right)=\mathcal{N}\left(\boldsymbol{x}_t ; \sqrt{\bar{\alpha}_t} \boldsymbol{x}_0,\left(1-\bar{\alpha}_t\right) \mathbf{I}\right)

然后使用Tweedie公式:

(131)E[μxtxt]=xt+(1αˉt)xtlogp(xt)\mathbb{E}[\boldsymbol{\mu}_{x_t} | \boldsymbol{x}_t]=\boldsymbol{x}_t+\left(1-\bar{\alpha}_t\right) \nabla_{\boldsymbol{x}_t} \log p\left(\boldsymbol{x}_t\right) \tag{131}

这里xtlogp(xt)\nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{x}_t)简写成了logp(xt)\nabla \log p(\boldsymbol{x}_t)。根据Tweedie公式,生成xt\boldsymbol{x}_t所用的真实均值μxt=αˉtx0\boldsymbol{\mu}_{x_t}=\sqrt{\bar{\alpha}_t} \boldsymbol{x}_0的最佳估计,定义为:

之后可以将公式133代入真实的去噪转移均值μq(xt,x0)\boldsymbol{\mu}_q(\boldsymbol{x}_t, \boldsymbol{x}_0)得到新的形式:

由此可以把近似去噪转移均值μθ(xt,t)\boldsymbol{\mu}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)写作:

(143)μθ(xt,t)=1αtxt+1αtαtsθ(xt,t)\boldsymbol{\mu}_{\boldsymbol{\theta}}\left(\boldsymbol{x}_t, t\right)=\frac{1}{\sqrt{\alpha_t}} \boldsymbol{x}_t+\frac{1-\alpha_t}{\sqrt{\alpha_t}} \boldsymbol{s}_{\boldsymbol{\theta}}\left(\boldsymbol{x}_t, t\right) \tag{143}

以及对应的优化问题:

这里的sθ(xt,t)\boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)是一个神经网络,用来学习预测得分函数(score function)xtlogp(xt)\nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{x}_t),即数据空间上任意噪声级别的xt\boldsymbol{x}_t的梯度。

细心的读者应该已经发现得分函数logp(xt)\nabla \log p(\boldsymbol{x}_t)看起来很像源噪声ϵ0\boldsymbol{\epsilon}_0的形式。这个可以通过结合Tweedie公式(公式133)和重参数化技巧(公式115)明确的看出来:

事实证明,这两项之间仅仅差一个随时间变化的常数因子!得分函数衡量如何在数据空间内移动来最大化对数概率,直观地说,由于源噪声被添加到自然图像中以扰乱图像,因此朝相反方向移动会对图像进行“去噪”,这是提升后续对数概率的最佳更新。我们的数学证明也证实了这个直觉,已经明确表明,学习对得分函数建模等同于对源噪声的负值建模(相差一个比例因子)。

至此我们推出了三种等价的优化VDM的目标:学习一个神经网络来预测(1)原始图像x0\boldsymbol{x}_0,(2)源噪声ϵ0\boldsymbol{\epsilon}_0,(3)一个任意噪声级别下的图像的得分logp(xt)\nabla \log p(\boldsymbol{x}_t),VDM可以通过随机采样时间步长tt和最小化真值目标与预测的范数来进行可缩放的训练。

基于得分的(Score-based)生成模型

前面展示了VDM可以通过简单的优化一个神经网络sθ(xt,t)\boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)来预测得分函数logp(xt)\nabla \log p(\boldsymbol{x}_t),但是,在推导过程中,得分项是应用Tweedie公式获得的。这个过程不能给我们充分的直觉,让我们认识到究竟什么是得分函数,或者说为什么要建模这个东西。刚好延续这种直觉,我们可以看一下另一种生成模型,基于得分的(Score-based)生成模型[9, 10, 11]。事实上,我们可以证明,我们之前推导的VDM和基于得分的生成模型公式之间等价,这样就可以灵活切换这两种解释。

为了开始理解为什么优化得分函数是有意义的,我们绕道而行,重新访问基于能量的模型[12, 13],任意可变的概率分布可以写成

(152)pθ(x)=1Zθefθ(x)p_{\boldsymbol{\theta}}(\boldsymbol{x})=\frac{1}{Z_{\boldsymbol{\theta}}} e^{-f_{\boldsymbol{\theta}}(\boldsymbol{x})} \tag{152}

这里fθ(x)f_{\boldsymbol{\theta}}(\boldsymbol{x})是哟个任意可变的参数化的函数,称作能量函数(energy function),通常由一个神经网络建模。Z_{\boldsymbol{\theta}是一个归一化常数,用来确保pθ(x)dx=1\int p_{\boldsymbol{\theta}}(\boldsymbol{x}) d \boldsymbol{x}=1。学习这种分布的一种方法是极大似然,但是,这需要归一化常数Zθ=efθ(x)dxZ_{\boldsymbol{\theta}}=\int e^{-f_{\boldsymbol{\theta}}(\boldsymbol{x})} d \boldsymbol{x}是容易计算的才行,对于复杂的fθ(x)f_{\boldsymbol{\theta}}(\boldsymbol{x})来说这种方法不可行。

一种避免计算或者建模这个归一化常数的方法是使用神经网络sθ(x)\boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x})来学习得分函数logp(x)\nabla \log p(\boldsymbol{x})或者分布p(x)p(\boldsymbol{x})。这个想法来自于对公式152两侧的对数求导:

这样可以在不涉及任何归一化常数的情况下利用神经网络表示。得分模型可以通过最小化与真实得分函数的Fisher Divergence优化:

(157)Ep(x)[sθ(x)logp(x)22]\mathbb{E}_{p(\boldsymbol{x})}\left[\left\|s_{\boldsymbol{\theta}}(\boldsymbol{x})-\nabla \log p(\boldsymbol{x})\right\|_2^2\right] \tag{157}

图6:用朗之万动力学生成的三个随机采样轨迹的可视化,所有这些轨迹都从同一初始点开始,用于高斯混合。左图在三维轮廓上绘制这些采样轨迹,而右图根据真实的得分函数绘制采样轨迹。由于朗之万动力学采样过程中的随机噪声项,我们能够从相同的初始点从不同的模式生成样本。如果没有噪声,每次试验从一个固定的点进行抽样总是会遵循固定的得分,最终到达相同的模式。

得分函数代表了什么?对于任意x\boldsymbol{x},取关于x\boldsymbol{x}的对数似然的梯度,本质上是描述了在数据空间上朝哪个方向移动,以进一步提高其似然。直观上,得分函数定义了一个数据x\boldsymbol{x}所在的整个空间上的向量场,指向某个模式(mode)。关于这个可视化的表述如图6所示。之后通过学习真实数据分布的得分函数,我们可以从相同空间的任意点开始,迭代地跟随得分直到达到一个模式,从而生成样本。这个采样过程称为朗之万动力学(Langevin dynamics),数学描述为:

(158)xi+1xi+clogp(xi)+2cϵ,i=0,1,,K\boldsymbol{x}_{i+1} \leftarrow \boldsymbol{x}_i+c \nabla \log p\left(\boldsymbol{x}_i\right)+\sqrt{2 c} \boldsymbol{\epsilon}, \quad i=0,1, \ldots, K \tag{158}

其中x0\boldsymbol{x}_0是一个从先验分布(例如均匀分布)中的随机采样,$$是一个额外的噪声项,确保生成的样本不要总是崩塌到一种模式上,而要在附近徘徊一下以确保多样性。另外,由于学到的得分函数是确定的,带着一个噪声项采样导致生成过程中增加了随机性,这可以规避确定性的轨迹。在采样从处于多种模式之间的位置初始化时,这尤其有用。一个可视化描述如图6所示,其中展示了朗之万动力学采样以及噪声项的好处。

注意公式157的目标依赖于真实的得分函数,但这在复杂分布(例如对自然图像的建模)中是不可用的。好在其他的技术,称之为得分匹配(score matching)[14, 15, 16, 17]被用于最小化这个Fisher divergence而不必知道真实的得分,同时它可以利用随机梯度下降方法做优化。

总的来说,学习将分布表示为一个得分函数,并使用马尔可夫链蒙特卡罗(MCMC)技术生成样本(如朗之万动力学),被称为基于得分的生成模型[9, 10, 11]。

正如Song和Ermon[9]所详述的那样,原版的得分匹配存在三个主要问题。

  1. 首先,当x\boldsymbol{x}位于高维空间中的低维流形上时,得分函数没有被很好的定义。这在数学上也能看出来,所有不在低维流形上的点的概率值是0,他们的对数是没有定义的。当我们试图学习一个自然图像的生成模型时,这相当不方便,因为自然图像位于整个环境空间的低维流形上。

  2. 第二,在低密度区域,通过原版的得分匹配训练到的得分函数不准确,从我们在公式157中最小化的目标可以明显看出这一点。因为它是对p(x)p(\boldsymbol{x})的一个期望,并通过在其中采样进行了显式的训练,对于少见或没见过的情况,该模型将无法受到准确的训练。这是有问题的,因为我们的采样策略涉及从高维空间中的随机位置开始(很可能是随机噪声),并根据学习的得分函数移动。因为我们正在跟随着一个有噪声或不准确的得分估计,所以最终生成的样本也可能是次优的,或者需要更多轮迭代才能收敛到准确的输出。

  3. 最后,郎之万动力学采样可能不会混合,即使使用真实得分进行采样。假设真实数据分布是两个不相交分布的混合:

    (159)p(x)=c1p1(x)+c2p2(x)p(\boldsymbol{x})=c_1 p_1(\boldsymbol{x})+c_2 p_2(\boldsymbol{x}) \tag{159}

    之后,当计算得分时,这些混合系数会消失,因为对数运算将系数从分布中分离出来,梯度运算将其归零。为了可视化这一点,请注意图6的右图,真实的得分函数不知道三种分布之间的不同权重。尽管右下角的那个模式在实际的高斯混合中具有更高的权重,但从这个初始化点进行的朗之万动力学采样到达每个模式的机会大致相等。

事实证明,通过向数据中添加多级高斯噪声,可以同时解决这三个缺点。首先,由于高斯噪声分布的支持(support)是整个空间,受扰动的数据样本将不再局限于低维流形。其次,加入较大的高斯噪声将增加每个模式在数据分布中覆盖的面积,在低密度区域增加更多的训练信号。最后,添加多级方差增加的高斯噪声将带来与真实的混合系数相关的中间分布。

形式上,我们可以选择正的噪声级别序列{σt}t=1T\left\{\sigma_t\right\}_{t=1}^T,定义一个逐步扰动的数据分布序列:

(160)pσt(xt)=p(x)N(xt;x,σt2I)dxp_{\sigma_t}\left(\boldsymbol{x}_t\right)=\int p(\boldsymbol{x}) \mathcal{N}\left(\boldsymbol{x}_t ; \boldsymbol{x}, \sigma_t^2 \mathbf{I}\right) d \boldsymbol{x} \tag{160}

之后,训练一个神经网络,通过得分匹配来同时在所有噪声级别上学习得分函数:

(161)argminθt=1Tλ(t)Epσt(xt)[sθ(x,t)logpσt(xt)22]\underset{\boldsymbol{\theta}}{\arg \min } \sum_{t=1}^T \lambda(t) \mathbb{E}_{p_{\sigma_t}\left(\boldsymbol{x}_t\right)}\left[\left\|\boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}, t)-\nabla \log p_{\sigma_t}\left(\boldsymbol{x}_t\right)\right\|_2^2\right] \tag{161}

其中λ(t)\lambda(t)是一个正的权重函数,取决于噪声级别tt。注意这个目标几乎与公式148训练VDM推导出的目标一致。此外,作者建议将退火朗之万动力学采样作为一个生成过程,其中对于序列中每个t=T,T1,...,2,1t=T,T-1,...,2,1,每个样本都通过运行朗之万动力学生成。初始化是从一些固定的先验(如均匀分布)中选择的,每个后续采样步骤都从先前模拟的最终样本开始。由于噪声水平随时间步tt而稳定降低,并且我们随时间减小步长,因此样本最终收敛到一个真实模式。这与VDM的MHVAE解释中执行的采样程序相当类似,其中,随机初始化的数据向量在不断降低的噪声水平上进行迭代细化。

至此,我们在VDM和基于得分的生成模型的训练目标和抽样过程中建立了明确的联系。

一个问题是,如何自然地将扩散模型推广到无限的时间步数。在MHVAE的视角下,这可以解释为将层次的数量扩展到无穷TT \rightarrow \infty。从等价的基于得分生成模型的角度来表示这一点更为清晰,在无穷多的噪声尺度下,图像在连续时间内的扰动可以表示为一个随机过程,因此可以用随机微分方程(SDE)来描述。然后通过反转SDE进行采样,这自然需要在每个连续值的噪声级估计得分函数[10]。不同的SDE参数化过程本质上描述了不同的随时间变化的扰动方案,从而实现了加噪过程的灵活建模[6]。

引导

到现在我们一置关注着对数据分布p(x)p(\boldsymbol{x})的建模,但是我们也会想要学习这样一个条件概率分布p(xy)p(\boldsymbol{x}|y),这个分布能够让我们明确的利用条件信息yy控制数据的生成。这种主要的思想应用在图像超分辨率(Cascaded Diffusion Models [18])和目前最前沿的图像文本模型(DALL-E 2 [19], Imagen [7])。

一种自然的想法是,要想加入条件信息,就把他加在每轮的时间步信息上。回想一下公式32的联合分布:

p(x0:T)=p(xT)t=1Tpθ(xt1xt)p\left(\boldsymbol{x}_{0: T}\right)=p\left(\boldsymbol{x}_T\right) \prod_{t=1}^T p_{\boldsymbol{\theta}}\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t\right)

然后,为了把这个转换成条件扩散模型,可以简单的加入任意条件信息yy到每一个转移步上面:

(162)p(x0:Ty)=p(xT)t=1Tpθ(xt1xt,y)p\left(\boldsymbol{x}_{0: T} \mid y\right)=p\left(\boldsymbol{x}_T\right) \prod_{t=1}^T p_{\boldsymbol{\theta}}\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t, y\right) \tag{162}

举例来说,在图像文本生成里,yy可以是一个文本的编码,又或者在超分里是一个低分辨率的图像。因此,对于每个想要的解释和实现,通过预测x^θ(xt,t,y)x0\hat{\boldsymbol{x}}_{\boldsymbol{\theta}}\left(\boldsymbol{x}_t, t, y\right) \approx \boldsymbol{x}_0ϵ^θ(xt,t,y)ϵ0\hat{\boldsymbol{\epsilon}}_{\boldsymbol{\theta}}\left(\boldsymbol{x}_t, t, y\right) \approx \epsilon_0,或者sθ(xt,t,y)logp(xty)\boldsymbol{s}_{\boldsymbol{\theta}}\left(\boldsymbol{x}_t, t, y\right) \approx \nabla \log p\left(\boldsymbol{x}_t | y\right),我们能够像之前一样学习一个VDM的核心神经网络。

这种原始形式公式的一个警告是,以这种方式训练的条件扩散模型可能会学会忽略或淡化任何给定的条件信息。因此,提出了引导(guidance)作为一种更明确地控制模型赋予条件信息权重的方法,需要使用数据多样性作为代价。两种最流行的引导形式被称为分类器引导(Classifier Guidance)[10, 20]和无分类器引导(Classifier-Free Guidance)[21]。

分类器引导

让我们从扩散模型的基于得分的方程入手,这里我们目标是学习任意噪声级别tt时,条件模型的得分logp(xty)\nabla \log p(\boldsymbol{x}_t | y)。回忆一下,\nablaxt\nabla_{\boldsymbol{x}_t}的简写。通过贝叶斯法则,可以推出下面的等价形式:

这里的logp(y)\log p(y)对于xt\boldsymbol{x}_t的梯度为0。

我们的最终推导结果可以解释为,学习一个无条件的得分函数和一个分类器的对抗梯度p(yxt)p(y|\boldsymbol{x}_t)。所以在分类器引导方法[10, 20]中,像之前推导的那样,学一个无条件扩散模型的得分,同时伴随一个分类器,他接受任意噪声xt\boldsymbol{x}_t并且尝试预测他的条件信息yy。然后在采样阶段,退火朗之万动力学所用的整体的条件得分函数,是无条件得分函数和噪声分类器的对抗梯度之和。

为了引入细粒度控制以鼓励或阻止模型考虑条件信息,分类器引导使用超参数项γ\gamma缩放噪声分类器的对抗梯度。分类器引导中学到的得分函数可以总结为:

(166)logp(xty)=logp(xt)+γlogp(yxt)\nabla \log p(\boldsymbol{x}_t | y)=\nabla \log p(\boldsymbol{x}_t)+\gamma \nabla \log p(y | \boldsymbol{x}_t) \tag{166}

直觉上,当γ=0\gamma=0时,条件扩散模型学会完全忽略条件信息,当γ\gamma很大时,模型学会生成与条件信息高度一致的样本。这需要以样本多样性为代价,因为模型只会生成易于重新生成提供的条件信息的数据,即使是在带噪图像的级别。(简要来说就是防止数据多样性不足带来的分类器过拟合噪声的问题,如果数据非常匮乏,分类器会在中间生成的带噪图像上学习类别信息,结果分类器只记住了噪声对应着某个类别,结果就是学不会分类)

分类器引导的一个显著缺点是它依赖于单独学习的分类器。由于分类器必须处理任意噪声输入,而大多数现有的预训练分类模型都没有对此进行优化,因此它必须与扩散模型一起专门做训练。

无分类器引导

在无分类器引导[21]中,作者没有单独训练分类器模型,转而使用无条件扩散模型和条件扩散模型。为了推导无分类器引导的得分函数,首先重新排列公式165:

(167)logp(yxt)=logp(xty)logp(xt)\nabla \log p(y | \boldsymbol{x}_t)=\nabla \log p(\boldsymbol{x}_t | y)-\nabla \log p(\boldsymbol{x}_t) \tag{167}

之后,把这个替换进公式166,可以得到:

同样的,γ\gamma是一个控制条件模型对于条件信息关注程度的参数。当γ=0\gamma=0,条件模型完全忽略条件信息,最终学成一个无条件扩散模型。当γ=1\gamma=1,模型在没有引导的情况下显式地学习原始条件分布。当γ>1\gamma > 1,扩散模型不仅优先考虑条件得分函数,而且会向远离无条件得分函数的方向移动。换言之,它降低了没有使用条件信息的样本的生成概率,有利于明确使用条件信息的样本生成。这也会出现样本多样性降低的效应,因为这个做法导致生成样本与条件信息精确匹配。

由于学习两个单独的扩散模型开销太大,我们可以将条件和无条件扩散模型作为一个单一的条件模型一起学习。无条件扩散模型可以通过将条件信息替换为固定的常量值(例如零),这本质上是对条件信息执行随机的dropout。无分类器引导是优雅的做法,因为只需要单一扩散模型的训练就可以更好地控制条件生成过程。

结语

概括一下我们在探索过程中的发现。首先,我们将VDM作为MHVAE的一个特例进行推导,其中有三个关键假设使得ELBO的计算易于处理和可放缩的优化。然后我们证明,优化VDM归结为学习神经网络来预测三个潜在目标中的其中一个:来自任意噪声级别的源图像、来自任意噪声图像的源噪声、或者任意噪声水平下带噪图像的得分函数。然后深入探讨了学习得分函数意味着什么,并将其与基于得分的生成模型的观点明确联系起来。最后介绍了如何使用扩散模型学习条件分布。

总之,扩散模型显示出了作为生成模型的难以置信的能力。实际上,它们为当前最先进的文本图像生成模型提供了动力,如Imagen和DALL-E 2。此外,这些模型背后的数学过程非常优雅。然而,仍有一些缺点需要考虑:

  • 这不太可能是我们人类自然建模和生成数据的方式,我们不会将迭代去噪一个随机噪声去生成样本。
  • VDM无法产生可解释的隐信息(latents)。尽管VAE非常有希望能够通过优化其编码器来学习结构化的隐空间,但在VDM中,每个时间步长的编码器已经作为线性高斯模型给出,无法可变的优化。因此,中间的隐信息被限制为原始输入的噪音版本。
  • 隐信息被限制与原始输入的维度相同,进一步影响了对有意义、压缩过的隐结构的学习。
  • 采样过程开销巨大,因为必须在两种公式下运行多个去噪步骤。回想一下,其中一个限制是选择足够大的时间步TT,以确保最终的隐噪声完全是高斯噪声。在采样期间必须迭代所有这些时间点以生成样本。

最后,扩散模型的成功突出了HVAE作为生成模型的强大能力。我们已经证明,当推广到无限的隐层次结构时,即使编码器很普通,隐维度是固定的,并且假设了马尔科夫变换,仍然能够学到强大的数据模型。这表明在一般的、深的HVAE情况下,可以进一步提高性能,这种情况下,有可能学到复杂的编码器和有语义意义的隐空间。

参考文献

[1] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114, 2013.
[2] Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improved variational inference with inverse autoregressive flow. Advances in neural information processing systems, 29, 2016.
[3] Casper Kaae Sønderby, Tapani Raiko, Lars Maaløe, Søren Kaae Sønderby, and Ole Winther. Ladder variational autoencoders. Advances in neural information processing systems, 29, 2016.
[4] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages 2256–2265. PMLR, 2015.
[5] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840–6851, 2020.
[6] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems, 34:21696–21707, 2021.
[7] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint arXiv:2205.11487, 2022.
[8] Bradley Efron. Tweedie’s formula and selection bias. Journal of the American Statistical Association, 106(496):1602–1614, 2011.
[9] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 32, 2019.
[10] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Scorebased generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020.
[11] Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. Advances in neural information processing systems, 33:12438–12448, 2020.
[12] Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and F Huang. A tutorial on energy-based learning. Predicting structured data, 1(0), 2006.
[13] Yang Song and Diederik P Kingma. How to train your energy-based models. arXiv preprint arXiv:2101.03288, 2021.
[14] Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
[15] Saeed Saremi, Arash Mehrjou, Bernhard Schölkopf, and Aapo Hyvärinen. Deep energy estimator networks. arXiv preprint arXiv:1805.08306, 2018.
[16] Yang Song, Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: A scalable approach to density and score estimation. In Uncertainty in Artificial Intelligence, pages 574–584. PMLR, 2020.
[17] Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation, 23(7):1661–1674, 2011.
[18] Jonathan Ho, Chitwan Saharia, William Chan, David J Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusion models for high fidelity image generation. J. Mach. Learn. Res., 23:47–1, 2022.
[19] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125, 2022.
[20] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780–8794, 2021.
[21] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. In NeurIPS 2021 Workshop on Deep Generative Models and Downstream Applications, 2021.