Reading

DDPM

💡 扩散模型:通过加噪的方式去学习原始数据的分布, 从学到的分布中去生成样本
DDPM 关键点:
1. 正向加噪是离散时间马尔可夫链:从 \(x_0\) 逐步加噪得到 \(x_1,x_2,...,x_T\);在合适的噪声调度与足够大的 \(T\) 下,\(x_T\) 近似服从\( N(0,I) \)的各向同性高斯。
2. 每一步噪声方差 \(β_t\) 满足 \(0<β_t<1\),通常随 \(t\) 增大;因此 \(q(x_t|x_{t-1}) \)的均值缩放系数 \(\sqrt{1-β_t} \)逐渐减小。
3. 训练通过最大化对数似然的变分下界(ELBO)来学习反向过程\( p_θ(x_{t-1}|x_t)\),并将其参数化为高斯分布(神经网络预测均值/噪声或 score)。
4. 将目标写成 score/DSM 形式时,loss 的权重与对应噪声层的方差尺度(如 \(1-\bar{α}_t\) 或相关量)有关;采样通常是按学习到的反向转移逐步生成(祖先采样),与经典 Langevin MCMC 更新形式不同,但可在 SDE 视角下统一理解。

“拆楼建楼”角度理解

很多文章在介绍DDPM时,上来就引入转移分布,接着就是变分推断,一堆数学记号下来,先吓跑了一群人(当然,从这种介绍我们可以再次看出,DDPM实际上是VAE而不是扩散模型),再加之人们对传统扩散模型的固有印象,所以就形成了“需要很高深的数学知识”的错觉。事实上,DDPM也可以有一种很“大白话”的理解,它并不比有着“造假-鉴别”通俗类比的GAN更难。

首先,我们想要做一个像GAN那样的生成模型,它实际上是将一个随机噪声\(\boldsymbol{z}\)变换成一个数据样本\(\boldsymbol{x}\)的过程:

\[\begin{CD} \text{随机噪声}\boldsymbol{z}\quad @>\quad\text{变换}\quad>> \quad\text{样本数据}\boldsymbol{x}\\ @V \text{类比} VV @VV \text{类比} V\\ \text{砖瓦水泥}\quad @>\quad\text{建设}\quad>> \quad\text{高楼大厦}\\ \end{CD}\]

我们可以将这个过程想象为“建设”,其中随机噪声\(\boldsymbol{z}\)是砖瓦水泥等原材料,样本数据\(\boldsymbol{x}\)是高楼大厦,所以生成模型就是一支用原材料建设高楼大厦的施工队。

这个过程肯定很难的,所以才有了那么多关于生成模型的研究。但俗话说“破坏容易建设难”,建楼你不会,拆楼你总会了吧?我们考虑将高楼大厦一步步地拆为砖瓦水泥的过程:设\(\boldsymbol{x}_0\)为建好的高楼大厦(数据样本),\(\boldsymbol{x}_T\) 为拆好的砖瓦水泥(随机噪声),假设“拆楼”需要T步,整个过程可以表示为

\[\boldsymbol{x} = \boldsymbol{x}_0 \to \boldsymbol{x}_1 \to \boldsymbol{x}_2 \to \cdots \to \boldsymbol{x}_{T-1} \to \boldsymbol{x}_T = \boldsymbol{z}\]

建高楼大厦的难度在于,从原材料\(\boldsymbol{x}_T\)到最终高楼大厦\(\boldsymbol{x}_0\)的跨度过大,普通人很难理\(解\boldsymbol{x}_T\)是怎么一下子变成\(\boldsymbol{x}_0\)的。但是,当我们有了“拆楼”的中间过程\(\boldsymbol{x}_1,\boldsymbol{x}_2,\cdots,\boldsymbol{x}_T\)后,我们知道\(\boldsymbol{x}_{t-1} \to \boldsymbol{x}_t\)代表着拆楼的一步,那么反过来\(\boldsymbol{x}_t\to \boldsymbol{x}_{t-1}\)不就是建楼的一步?如果我们能学会两者之间的变换关系\(\boldsymbol{x}_{t-1}=\boldsymbol{\mu}(\boldsymbol{x}_t)\),那么从\(\boldsymbol{x}_T\)出发,反复地执行\(\boldsymbol{x}_{T-1}=\boldsymbol{\mu}(\boldsymbol{x}_T)、\boldsymbol{x}_{T-2}=\boldsymbol{\mu}(\boldsymbol{x}_{T-1})、...,\)最终不就能造出高楼大厦\(\boldsymbol{x}_0\)出来?

该怎么拆

正所谓“饭要一口一口地吃”,楼也要一步一步地建,DDPM做生成模型的过程,其实跟上述“拆楼-建楼”的类比是完全一致的,它也是先反过来构建一个从数据样本渐变到随机噪声的过程,然后再考虑其逆变换,通过反复执行逆变换来完成数据样本的生成,所以本文前面才说DDPM这种做法其实应该更准确地称为“渐变模型”而不是“扩散模型”。

具体来说,DDPM将“拆楼”的过程建模为

\[\boldsymbol{x}_t = \alpha_t \boldsymbol{x}_{t-1} + \beta_t \boldsymbol{\varepsilon}_t,\quad \boldsymbol{\varepsilon}_t\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})\tag{1}\]

其中有\(\alpha_t,\beta_t > 0\)\(\alpha_t^2 + \beta_t^2=1\)\(\beta_t\)通常很接近于0,代表着单步“拆楼”中对原来楼体的破坏程度,噪声 \(\boldsymbol{\varepsilon}_t\) 的引入代表着对原始信号的一种破坏,我们也可以将它理解为“原材料”,即每一步“拆楼”中我们都将\(\boldsymbol{x}_{t-1}\)拆解为“\(\alpha_t \boldsymbol{x}_{t-1}\)的楼体 + \(\beta_t \boldsymbol{\varepsilon}_t\)的原料”。(提示:本文\(\alpha_t,\beta_t\)的定义跟原论文不一样。)

反复执行这个拆楼的步骤,我们可以得到:

\[\begin{aligned} \boldsymbol{x}_t =&\, \alpha_t \boldsymbol{x}_{t-1} + \beta_t \boldsymbol{\varepsilon}_t \\ =&\, \alpha_t \big(\alpha_{t-1} \boldsymbol{x}_{t-2} + \beta_{t-1} \boldsymbol{\varepsilon}_{t-1}\big) + \beta_t \boldsymbol{\varepsilon}_t \\ =&\,\cdots\\ =&\,(\alpha_t\cdots\alpha_1) \boldsymbol{x}_0 + \underbrace{(\alpha_t\cdots\alpha_2)\beta_1 \boldsymbol{\varepsilon}_1 + (\alpha_t\cdots\alpha_3)\beta_2 \boldsymbol{\varepsilon}_2 + \cdots + \alpha_t\beta_{t-1} \boldsymbol{\varepsilon}_{t-1} + \beta_t \boldsymbol{\varepsilon}_t}_{\text{多个相互独立的正态噪声之和}} \end{aligned}\tag{2}\]

可能刚才读者就想问为什么叠加的系数要满足\(\alpha_t^2 + \beta_t^2 = 1\)了,现在我们就可以回答这个问题。首先,式中花括号所指出的部分,正好是多个独立的正态噪声之和,其均值为0,方差则分别为\((\alpha_t\cdots\alpha_2)^2\beta_1^2、(\alpha_t\cdots\alpha_3)^2\beta_2^2、...、\alpha_t^2\beta_{t-1}^2、\beta_t^2\);然后,我们利用一个概率论的知识——正态分布的叠加性,即上述多个独立的正态噪声之和的分布,实际上是均值为0、方差为\((\alpha_t\cdots\alpha_2)^2\beta_1^2 + (\alpha_t\cdots\alpha_3)^2\beta_2^2 + \cdots + \alpha_t^2\beta_{t-1}^2 + \beta_t^2\) 的正态分布;最后,在\(\alpha_t^2 + \beta_t^2 = 1\)恒成立之下,我们可以得到式(2)的各项系数平方和依旧为1,即

\[(\alpha_t\cdots\alpha_1)^2 + (\alpha_t\cdots\alpha_2)^2\beta_1^2 + (\alpha_t\cdots\alpha_3)^2\beta_2^2 + \cdots + \alpha_t^2\beta_{t-1}^2 + \beta_t^2 = 1\]

所以实际上相当于有

\[\boldsymbol{x}_t = \underbrace{(\alpha_t\cdots\alpha_1)}_{\text{记为}\bar{\alpha}_t} \boldsymbol{x}_0 + \underbrace{\sqrt{1 - (\alpha_t\cdots\alpha_1)^2}}_{\text{记为}\bar{\beta}_t} \bar{\boldsymbol{\varepsilon}}_t,\quad \bar{\boldsymbol{\varepsilon}}_t\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})\tag{3}\]

这就为计算\(\boldsymbol{x}_t\)提供了极大的便利。另一方面,DDPM会选择适当的\(\alpha_t\)形式,使得有\(\bar{\alpha}_T\approx 0\),这意味着经过\(T\)步的拆楼后,所剩的楼体几乎可以忽略了,已经全部转化为原材料\(\boldsymbol{\varepsilon}\)。(提示:本文\(\bar{\alpha}_t\)的定义跟原论文不一样。)

又如何建

“拆楼”是\(\boldsymbol{x}_{t-1}\to \boldsymbol{x}_t\)的过程,这个过程我们得到很多的数据对\((\boldsymbol{x}_{t-1},\boldsymbol{x}_t)\),那么“建楼”自然就是从这些数据对中学习一个\(\boldsymbol{x}_t\to \boldsymbol{x}_{t-1}\) 的模型。设该模型为\(\boldsymbol{\mu}(\boldsymbol{x}_t)\),那么容易想到学习方案就是最小化两者的欧氏距离:

\[\left\Vert\boldsymbol{x}_{t-1} - \boldsymbol{\mu}(\boldsymbol{x}_t)\right\Vert^2\]

其实这已经非常接近最终的DDPM模型了,接下来让我们将这个过程做得更精细一些。首先“拆楼”的(1)可以改写为\(\boldsymbol{x}_{t-1} = \frac{1}{\alpha_t}\left(\boldsymbol{x}_t - \beta_t \boldsymbol{\varepsilon}_t\right)\) ,这启发我们或许可以将“建楼”模型\(\boldsymbol{\mu}(\boldsymbol{x}_t)\)设计成

\[\boldsymbol{\mu}(\boldsymbol{x}_t) = \frac{1}{\alpha_t}\left(\boldsymbol{x}_t - \beta_t \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\right)\tag{4}\]

的形式,其中\(\boldsymbol{\theta}\) 是训练参数,将其代入到损失函数,得到

\[\left\Vert\boldsymbol{x}_{t-1} - \boldsymbol{\mu}(\boldsymbol{x}_t)\right\Vert^2 = \frac{\beta_t^2}{\alpha_t^2}\left\Vert \boldsymbol{\varepsilon}_t - \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\right\Vert^2\]

前面的因子 \(\frac{\beta_t^2}{\alpha_t^2}\) 代表loss的权重,这个我们可以暂时忽略,最后代入结合(1)(3)所给出\(\boldsymbol{x}_t\) 的表达式

\[\boldsymbol{x}_t = \alpha_t\boldsymbol{x}_{t-1} + \beta_t \boldsymbol{\varepsilon}_t = \alpha_t\left(\bar{\alpha}_{t-1}\boldsymbol{x}_0 + \bar{\beta}_{t-1}\bar{\boldsymbol{\varepsilon}}_{t-1}\right) + \beta_t \boldsymbol{\varepsilon}_t = \bar{\alpha}_t\boldsymbol{x}_0 + \alpha_t\bar{\beta}_{t-1}\bar{\boldsymbol{\varepsilon}}_{t-1} + \beta_t \boldsymbol{\varepsilon}_t\]

得到损失函数的形式为:

\[\left\Vert \boldsymbol{\varepsilon}_t - \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\bar{\alpha}_t\boldsymbol{x}_0 + \alpha_t\bar{\beta}_{t-1}\bar{\boldsymbol{\varepsilon}}_{t-1} + \beta_t \boldsymbol{\varepsilon}_t, t)\right\Vert^2\tag{5}\]

可能读者想问为什么要回退一步来给出 \(\boldsymbol{x}_t\)直接根据(3)来给出 \(\boldsymbol{x}_t\) 可以吗?答案是不行,因为我们已经事先采样了\(\boldsymbol{\varepsilon}_t\),而\(\boldsymbol{\varepsilon}_t\)\(\bar{\boldsymbol{\varepsilon}}_t\) 不是相互独立的,所以给定 \(\boldsymbol{\varepsilon}_t\) 的情况下,我们不能完全独立地采样 \(\bar{\boldsymbol{\varepsilon}}_t\)

降低方差

原则上来说,损失函数(5)就可以完成DDPM的训练,但它在实践中可能有方差过大的风险,从而导致收敛过慢等问题。要理解这一点并不困难,只需要观察到(5)实际上包含了4个需要采样的随机变量:

1、从所有训练样本中采样一个\(\boldsymbol{x}_0\)

2、从正态分布\(\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})\)中采样\(\bar{\boldsymbol{\varepsilon}}_{t-1}, \boldsymbol{\varepsilon}_t\)(两个不同的采样结果);

3、从\(1\sim T\)中采样一个\(t\)

要采样的随机变量越多,就越难对损失函数做准确的估计,反过来说就是每次对损失函数进行估计的波动(方差)过大了。很幸运的是,我们可以通过一个积分技巧来将\(\bar{\boldsymbol{\varepsilon}}_{t-1}, \boldsymbol{\varepsilon}_t\)合并成单个正态随机变量,从而缓解一下方差大的问题。

这个积分确实有点技巧性,但也不算复杂。由于正态分布的叠加性,我们知道\(\alpha_t\bar{\beta}_{t-1}\bar{\boldsymbol{\varepsilon}}_{t-1} + \beta_t \boldsymbol{\varepsilon}_t\)实际上相当于单个随机变量 \(\bar{\beta}_t\boldsymbol{\varepsilon}|\boldsymbol{\varepsilon}\sim \mathcal{N}(\boldsymbol{0}, \boldsymbol{I})\),同理\(\beta_t \bar{\boldsymbol{\varepsilon}}_{t-1} - \alpha_t\bar{\beta}_{t-1} \boldsymbol{\varepsilon}_t\) 实际上相当于单个随机变量\(\bar{\beta}_t\boldsymbol{\omega}|\boldsymbol{\omega}\sim \mathcal{N}(\boldsymbol{0}, \boldsymbol{I})\),并且可以验证\(\mathbb{E}[\boldsymbol{\varepsilon}\boldsymbol{\omega}^{\top}]=\boldsymbol{0}\),所以这是两个相互独立的正态随机变量。

接下来,我们反过来将\(\boldsymbol{\varepsilon}_t\)\(\boldsymbol{\varepsilon},\boldsymbol{\omega}\)重新表示出来

\[\boldsymbol{\varepsilon}_t = \frac{(\beta_t \boldsymbol{\varepsilon} - \alpha_t\bar{\beta}_{t-1} \boldsymbol{\omega})\bar{\beta}_t}{\beta_t^2 + \alpha_t^2\bar{\beta}_{t-1}^2} = \frac{\beta_t \boldsymbol{\varepsilon} - \alpha_t\bar{\beta}_{t-1} \boldsymbol{\omega}}{\bar{\beta}_t}\]

代入到式(5)得到

\[\begin{aligned} &\,\mathbb{E}_{\bar{\boldsymbol{\varepsilon}}_{t-1}, \boldsymbol{\varepsilon}_t\sim \mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\left\Vert \boldsymbol{\varepsilon}_t - \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\bar{\alpha}_t\boldsymbol{x}_0 + \alpha_t\bar{\beta}_{t-1}\bar{\boldsymbol{\varepsilon}}_{t-1} + \beta_t \boldsymbol{\varepsilon}_t, t)\right\Vert^2\right] \\ =&\,\mathbb{E}_{\boldsymbol{\omega}, \boldsymbol{\varepsilon}\sim \mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\left\Vert \frac{\beta_t \boldsymbol{\varepsilon} - \alpha_t\bar{\beta}_{t-1} \boldsymbol{\omega}}{\bar{\beta}_t} - \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\bar{\alpha}_t\boldsymbol{x}_0 + \bar{\beta}_t\boldsymbol{\varepsilon}, t)\right\Vert^2\right] \end{aligned}\]

注意到,现在损失函数关于 \(\boldsymbol{\omega}\) 只是二次的,所以我们可以展开然后将它的期望直接算出来,结果是

\[\frac{\beta_t^2}{\bar{\beta}_t^2}\mathbb{E}_{\boldsymbol{\varepsilon}\sim \mathcal{N}(\boldsymbol{0}, \boldsymbol{I})}\left[\left\Vert\boldsymbol{\varepsilon} - \frac{\bar{\beta}_t}{\beta_t}\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\bar{\alpha}_t\boldsymbol{x}_0 + \bar{\beta}_t\boldsymbol{\varepsilon}, t)\right\Vert^2\right]+\text{常数}\]

再次省掉常数和损失函数的权重,我们得到DDPM最终所用的损失函数:

\[\left\Vert\boldsymbol{\varepsilon} - \frac{\bar{\beta}_t}{\beta_t}\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\bar{\alpha}_t\boldsymbol{x}_0 + \bar{\beta}_t\boldsymbol{\varepsilon}, t)\right\Vert^2\tag{6}\]

提示:原论文中的\(\boldsymbol{\epsilon}_{\boldsymbol{\theta}}\)实际上就是本文的\(\frac{\bar{\beta}_t}{\beta_t}\boldsymbol{\epsilon}_{\boldsymbol{\theta}}\),所以大家的结果是完全一样的。)

递归生成

至此,我们算是把DDPM的整个训练流程捋清楚了。内容写了不少,你要说它很容易,那肯定说不上,但真要说非常困难的地方也几乎没有——没有用到传统的能量函数、得分匹配等工具,甚至连变分推断的知识都没有用到,只是借助“拆楼-建楼”的类比和一些基本的概率论知识,就能得到完全一样的结果。所以说,以DDPM为代表的新兴起的生成扩散模型,实际上没有很多读者想象的复杂,它可以说是我们从“拆解-重组”的过程中学习新知识的形象建模。

训练完之后,我们就可以从一个随机噪声\(\boldsymbol{x}_T\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})\)出发执行\(T\)(4)来进行生成:

\[\boldsymbol{x}_{t-1} = \frac{1}{\alpha_t}\left(\boldsymbol{x}_t - \beta_t \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\right)\]

这对应于自回归解码中的Greedy Search。如果要进行Random Sample,那么需要补上噪声项:

\[\boldsymbol{x}_{t-1} = \frac{1}{\alpha_t}\left(\boldsymbol{x}_t - \beta_t \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\right) + \sigma_t \boldsymbol{z},\quad \boldsymbol{z}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})\]

一般来说,我们可以让\(\sigma_t=\beta_t\),即正向和反向的方差保持同步。这个采样过程跟传统扩散模型的朗之万采样不一样的地方在于:DDPM的采样每次都从一个随机噪声出发,需要重复迭代\(T\)步来得到一个样本输出;朗之万采样则是从任意一个点出发,反复迭代无限步,理论上这个迭代无限步的过程中,就把所有数据样本都被生成过了。所以两者除了形式相似外,实质上是两个截然不同的模型。

从这个生成过程中,我们也可以感觉到它其实跟Seq2Seq的解码过程是一样的,都是串联式的自回归生成,所以生成速度是一个瓶颈,DDPM设了\(T=1000\),意味着每生成一个图片,需要将\(\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\)反复执行1000次,因此DDPM的一大缺点就是采样速度慢,后面有很多工作都致力于提升DDPM的采样速度。而说到“图片生成 + 自回归模型 + 很慢”,有些读者可能会联想到早期的PixelRNNPixelCNN等模型,它们将图片生成转换成语言模型任务,所以同样也是递归地进行采样生成以及同样地慢。那么DDPM的这种自回归生成,跟PixelRNN/PixelCNN的自回归生成,又有什么实质区别呢?为什么PixelRNN/PixelCNN没大火起来,反而轮到了DDPM?

了解PixelRNN/PixelCNN的读者都知道,这类生成模型是逐个像素逐个像素地生成图片的,而自回归生成是有序的,这就意味着我们要提前给图片的每个像素排好顺序,最终的生成效果跟这个顺序紧密相关。然而,目前这个顺序只能是人为地凭着经验来设计(这类经验的设计都统称为“Inductive Bias”),暂时找不到理论最优解。换句话说,PixelRNN/PixelCNN的生成效果很受Inductive Bias的影响。但DDPM不一样,它通过“拆楼”的方式重新定义了一个自回归方向,而对于所有的像素来说则都是平权的、无偏的,所以减少了Inductive Bias的影响,从而提升了效果。此外,DDPM生成的迭代步数是固定的\(T\),而PixelRNN/PixelCNN则是等于图像分辨率(\(\text{宽}\times\text{高}\times{通道数}\)),所以DDPM生成高分辨率图像的速度要比PixelRNN/PixelCNN快得多。

超参设置

这一节我们讨论一下超参的设置问题。

在DDPM中,\(T=1000\),可能比很多读者的想象数值要大,那为什么要设置这么大的T呢?另一边,对于\(\alpha_t\)的选择,将原论文的设置翻译到本博客的记号上,大致上是

\[\alpha_t = \sqrt{1 - \frac{0.02t}{T}}\]

这是一个单调递减的函数,那为什么要选择单调递减的\(\alpha_t\)呢?

其实这两个问题有着相近的答案,跟具体的数据背景有关。简单起见,在重构的时候我们用了欧氏距离公式5作为损失函数,而一般我们用DDPM做图片生成,以往做过图片生成的读者都知道,欧氏距离并不是图片真实程度的一个好的度量,VAE用欧氏距离来重构时,往往会得到模糊的结果,除非是输入输出的两张图片非常接近,用欧氏距离才能得到比较清晰的结果,所以选择尽可能大的T,正是为了使得输入输出尽可能相近,减少欧氏距离带来的模糊问题。

选择单调递减的\(\alpha_t\)也有类似考虑。当t比较小时,\(\boldsymbol{x}_t\)还比较接近真实图片,所以我们要缩小\(\boldsymbol{x}_{t-1}\)\(\boldsymbol{x}_t\) 的差距,以便更适用欧氏距离公式5,因此要用较大的\(\alpha_t\);当t比较大时,\(\boldsymbol{x}_t\)已经比较接近纯噪声了,噪声用欧式距离无妨,所以可以稍微增大\(\boldsymbol{x}_{t-1}\)\(\boldsymbol{x}_t\)的差距,即可以用较小的\(\alpha_t\)。那么可不可以一直用较大的\(\alpha_t\)呢?可以是可以,但是要增大T。注意在推导公式4时,我们说过应该有\(\bar{\alpha}_T\approx 0\),而我们可以直接估算

\[\log \bar{\alpha}_T = \sum_{t=1}^T \log\alpha_t = \frac{1}{2} \sum_{t=1}^T \log\left(1 - \frac{0.02t}{T}\right) < \frac{1}{2} \sum_{t=1}^T \left(- \frac{0.02t}{T}\right) = -0.005(T+1)\]

代入T=1000大致是\(\bar{\alpha}_T\approx e^{-5}\),这个其实就刚好达到\(\approx 0\)的标准。所以如果从头到尾都用较大的\(\alpha_t\),那么必然要更大的T才能使得\(\bar{\alpha}_T\approx 0\)了。

最后我们留意到,“建楼”模型中的\(\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\bar{\alpha}_t\boldsymbol{x}_0 + \bar{\beta}_t\boldsymbol{\varepsilon}, t)\)中,我们在输入中显式地写出了t,这是因为原则上不同的t处理的是不同层次的对象,所以应该用不同的重构模型,即应该有T个不同的重构模型才对,于是我们共享了所有重构模型的参数,将t作为条件传入。按照论文附录的说法,t是转换成位置编码后,直接加到残差模块上去的。

小结

本文从“拆楼-建楼”的通俗类比中介绍了最新的生成扩散模型DDPM,在这个视角中,我们可以通过较为“大白话”的描述以及比较少的数学推导,来得到跟原始论文一模一样的结果。总的来说,本文说明了DDPM也可以像GAN一样找到一个形象类比,它既可以不用到VAE中的“变分”,也可以不用到GAN中的“概率散度”、“最优传输”,从这个意义上来看,DDPM甚至算得上比VAE、GAN还要简单。

贝叶斯角度下的DDPM

image.png

DDPM建模的是如下变换流程:

\[\boldsymbol{x} = \boldsymbol{x}_0 \rightleftharpoons \boldsymbol{x}_1 \rightleftharpoons \boldsymbol{x}_2 \rightleftharpoons \cdots \rightleftharpoons \boldsymbol{x}_{T-1} \rightleftharpoons \boldsymbol{x}_T = \boldsymbol{z}\]

其中,正向就是将样本数据\(\boldsymbol{x}\)逐渐变为随机噪声\(\boldsymbol{z}\)的过程,反向就是将随机噪声\(\boldsymbol{z}\)逐渐变为样本数据\(\boldsymbol{x}\)的过程,反向过程就是我们希望得到的“生成模型”。

(1)(2)可以求出\(p(\boldsymbol{x}_t|\boldsymbol{x}_0)=\mathcal{N}(\boldsymbol{x}_t;\bar{\alpha}_t \boldsymbol{x}_0,\bar{\beta}_t^2 \boldsymbol{I})\),其中\(\bar{\alpha}_t = \alpha_1\cdots\alpha_t\),而\(\bar{\beta}_t = \sqrt{1-\bar{\alpha}_t^2}\)

DDPM要做的事情,就是从上述信息中求出反向过程所需要的\(p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t)\),这样我们就能实现从任意一个\(\boldsymbol{x}_T=\boldsymbol{z}\)出发,逐步采样出\(\boldsymbol{x}_{T-1},\boldsymbol{x}_{T-2},\cdots,\boldsymbol{x}_1\),最后得到随机生成的样本数据\(\boldsymbol{x}_0=\boldsymbol{x}\)

请贝叶斯

下面我们请出伟大的贝叶斯定理。事实上,直接根据贝叶斯定理我们有

\[p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t) = \frac{p(\boldsymbol{x}_t|\boldsymbol{x}_{t-1})p(\boldsymbol{x}_{t-1})}{p(\boldsymbol{x}_t)}\tag{7}\]

然而,我们并不知道\(p(\boldsymbol{x}_{t-1}),p(\boldsymbol{x}_t)\)的表达式,所以此路不通。但我们可以退而求其次,在给定\(\boldsymbol{x}_0\) 的条件下使用贝叶斯定理:

\[p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t, \boldsymbol{x}_0) = \frac{p(\boldsymbol{x}_t|\boldsymbol{x}_{t-1})p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_0)}{p(\boldsymbol{x}_t|\boldsymbol{x}_0)}\]

这样修改自然是因为\(p(\boldsymbol{x}_t|\boldsymbol{x}_{t-1}),p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_0),p(\boldsymbol{x}_t|\boldsymbol{x}_0)\)都是已知的,所以上式是可计算的,代入各自的表达式得到:

\[p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t, \boldsymbol{x}_0) = \mathcal{N}\left(\boldsymbol{x}_{t-1};\frac{\alpha_t\bar{\beta}_{t-1}^2}{\bar{\beta}_t^2}\boldsymbol{x}_t + \frac{\bar{\alpha}_{t-1}\beta_t^2}{\bar{\beta}_t^2}\boldsymbol{x}_0,\frac{\bar{\beta}_{t-1}^2\beta_t^2}{\bar{\beta}_t^2} \boldsymbol{I}\right)\tag{8}\]
推导:上式的推导过程并不难,就是常规的展开整理而已,当然我们也可以找点技巧加快计算。首先,代入各自的表达式,可以发现指数部分除掉\(-1/2\)因子外,结果是:
\[\frac{\Vert \boldsymbol{x}_t - \alpha_t \boldsymbol{x}_{t-1}\Vert^2}{\beta_t^2} + \frac{\Vert \boldsymbol{x}_{t-1} - \bar{\alpha}_{t-1}\boldsymbol{x}_0\Vert^2}{\bar{\beta}_{t-1}^2} - \frac{\Vert \boldsymbol{x}_t - \bar{\alpha}_t \boldsymbol{x}_0\Vert^2}{\bar{\beta}_t^2}\]
它关于\(\boldsymbol{x}_{t-1}\)是二次的,因此最终的分布必然也是正态分布,我们只需要求出其均值和协方差。不难看出,展开式中\(\Vert \boldsymbol{x}_{t-1}\Vert^2\)项的系数是
\[\frac{\alpha_t^2}{\beta_t^2} + \frac{1}{\bar{\beta}_{t-1}^2} = \frac{\alpha_t^2\bar{\beta}_{t-1}^2 + \beta_t^2}{\bar{\beta}_{t-1}^2 \beta_t^2} = \frac{\alpha_t^2(1-\bar{\alpha}_{t-1}^2) + \beta_t^2}{\bar{\beta}_{t-1}^2 \beta_t^2} = \frac{1-\bar{\alpha}_t^2}{\bar{\beta}_{t-1}^2 \beta_t^2} = \frac{\bar{\beta}_t^2}{\bar{\beta}_{t-1}^2 \beta_t^2}\]
所以整理好的结果必然是\(\frac{\bar{\beta}_t^2}{\bar{\beta}_{t-1}^2 \beta_t^2}\Vert \boldsymbol{x}_{t-1} - \tilde{\boldsymbol{\mu}}(\boldsymbol{x}_t, \boldsymbol{x}_0)\Vert^2\)的形式,这意味着协方差矩阵是\(\frac{\bar{\beta}_{t-1}^2 \beta_t^2}{\bar{\beta}_t^2}\boldsymbol{I}\)。另一边,把一次项系数拿出来是\(-2\left(\frac{\alpha_t}{\beta_t^2}\boldsymbol{x}_t + \frac{\bar{\alpha}_{t-1}}{\bar{\beta}_{t-1}^2}\boldsymbol{x}_0 \right)\),除以\(\frac{-2\bar{\beta}_t^2}{\bar{\beta}_{t-1}^2 \beta_t^2}\)后便可以得到
\[\tilde{\boldsymbol{\mu}}(\boldsymbol{x}_t, \boldsymbol{x}_0)=\frac{\alpha_t\bar{\beta}_{t-1}^2}{\bar{\beta}_t^2}\boldsymbol{x}_t + \frac{\bar{\alpha}_{t-1}\beta_t^2}{\bar{\beta}_t^2}\boldsymbol{x}_0\]
这就得到了\(p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t, \boldsymbol{x}_0)\)的所有信息了,结果正是式(8)

去噪过程

现在我们得到了\(p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t, \boldsymbol{x}_0)\),它有显式的解,但并非我们想要的最终答案,因为我们只想通过\(\boldsymbol{x}_t\)来预测\(\boldsymbol{x}_{t-1}\),而不能依赖\(\boldsymbol{x}_0\)\(\boldsymbol{x}_0\)是我们最终想要生成的结果。接下来,一个“异想天开”的想法是

如果我们能够通过\(\boldsymbol{x}_t\)来预测\(\boldsymbol{x}_0\),那么不就可以消去\(p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t, \boldsymbol{x}_0)\)中的\(\boldsymbol{x}_0\),使得它只依赖于\(\boldsymbol{x}_t\)了吗?

说干就干,我们用\(\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)\)来预估\(\boldsymbol{x}_0\),损失函数为\(\Vert \boldsymbol{x}_0 - \bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)\Vert^2\)。训练完成后,我们就认为

\[p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t) \approx p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t, \boldsymbol{x}_0=\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)) = \mathcal{N}\left(\boldsymbol{x}_{t-1}; \frac{\alpha_t\bar{\beta}_{t-1}^2}{\bar{\beta}_t^2}\boldsymbol{x}_t + \frac{\bar{\alpha}_{t-1}\beta_t^2}{\bar{\beta}_t^2}\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t),\frac{\bar{\beta}_{t-1}^2\beta_t^2}{\bar{\beta}_t^2} \boldsymbol{I}\right)\tag{9}\]

\(\Vert \boldsymbol{x}_0 - \bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)\Vert^2\)中,\(\boldsymbol{x}_0\)代表原始数据,\(\boldsymbol{x}_t\)代表带噪数据,所以这实际上在训练一个去噪模型,这也就是DDPM的第一个“D”的含义(Denoising)。

具体来说,\(p(\boldsymbol{x}_t|\boldsymbol{x}_0)=\mathcal{N}(\boldsymbol{x}_t;\bar{\alpha}_t \boldsymbol{x}_0,\bar{\beta}_t^2 \boldsymbol{I})\)意味着\(\boldsymbol{x}_t = \bar{\alpha}_t \boldsymbol{x}_0 + \bar{\beta}_t \boldsymbol{\varepsilon},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})\),或者写成\(\boldsymbol{x}_0 = \frac{1}{\bar{\alpha}_t}\left(\boldsymbol{x}_t - \bar{\beta}_t \boldsymbol{\varepsilon}\right)\),这启发我们将\(\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)\)参数化为

\[\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t) = \frac{1}{\bar{\alpha}_t}\left(\boldsymbol{x}_t - \bar{\beta}_t \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\right)\tag{10}\]

此时损失函数变为

\[\Vert \boldsymbol{x}_0 - \bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)\Vert^2 = \frac{\bar{\beta}_t^2}{\bar{\alpha}_t^2}\left\Vert\boldsymbol{\varepsilon} - \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\bar{\alpha}_t \boldsymbol{x}_0 + \bar{\beta}_t \boldsymbol{\varepsilon}, t)\right\Vert^2\]

省去前面的系数,就得到DDPM原论文所用的损失函数了。可以发现,本文是直接得出了从\(\boldsymbol{x}_t\)\(\boldsymbol{x}_0\)的去噪过程,而不是像之前两篇文章那样,通过\(\boldsymbol{x}_t\)\(\boldsymbol{x}_{t-1}\)的去噪过程再加上积分变换来推导,相比之下本文的推导可谓更加一步到位了。

另一边,我们将(10)代入到(9)中,化简得到

\[p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t) \approx p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t, \boldsymbol{x}_0=\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)) = \mathcal{N}\left(\boldsymbol{x}_{t-1}; \frac{1}{\alpha_t}\left(\boldsymbol{x}_t - \frac{\beta_t^2}{\bar{\beta}_t}\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\right),\frac{\bar{\beta}_{t-1}^2\beta_t^2}{\bar{\beta}_t^2} \boldsymbol{I}\right)\]

这就是反向的采样过程所用的分布,连同采样过程所用的方差也一并确定下来了。至此,DDPM推导完毕~(注:出于推导的流畅性考虑,本文的\(\boldsymbol{\epsilon}_{\boldsymbol{\theta}}\)跟上面”建楼-拆楼“时的推导不一样,反而跟DDPM原论文一致。)

推导:将将(10)代入到(9)中的主要化简难度就是计算
\[\begin{aligned}\frac{\alpha_t\bar{\beta}_{t-1}^2}{\bar{\beta}_t^2} + \frac{\bar{\alpha}_{t-1}\beta_t^2}{\bar{\alpha}_t\bar{\beta}_t^2} =&\, \frac{\alpha_t\bar{\beta}_{t-1}^2 + \beta_t^2/\alpha_t}{\bar{\beta}_t^2} = \frac{\alpha_t^2(1-\bar{\alpha}_{t-1}^2) + \beta_t^2}{\alpha_t\bar{\beta}_t^2} = \frac{1-\bar{\alpha}_t^2}{\alpha_t\bar{\beta}_t^2} = \frac{1}{\alpha_t} \end{aligned}\]

预估修正

不知道读者有没有留意到一个有趣的地方:我们要做的事情,就是想将\(\boldsymbol{x}_T\)慢慢地变为\(\boldsymbol{x}_0\),而我们在借用\(p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t, \boldsymbol{x}_0)\)近似\(p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t)\)时,却包含了“用\(\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)\)来预估\(\boldsymbol{x}_0\)”这一步,要是能预估准的话,那就直接一步到位了,还需要逐步采样吗?

真实情况是,“用\(\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)\)来预估\(\boldsymbol{x}_0\)”当然不会太准的,至少开始的相当多步内不会太准。它仅仅起到了一个前瞻性的预估作用,然后我们只用\(p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t)\)来推进一小步,这就是很多数值算法中的“预估-修正”思想,即我们用一个粗糙的解往前推很多步,然后利用这个粗糙的结果将最终结果推进一小步,以此来逐步获得更为精细的解。

由此我们还可以联想到Hinton三年前提出的《Lookahead Optimizer: k steps forward, 1 step back》,它同样也包含了预估(k steps forward)和修正(1 step back)两部分,原论文将其诠释为“快(Fast)-慢(Slow)”权重的相互结合,快权重就是预估得到的结果,慢权重则是基于预估所做的修正结果。如果愿意,我们也可以用同样的方式去诠释DDPM的“预估-修正”过程~

遗留问题

最后,在使用贝叶斯定理一节中,我们说式(7)没法直接用的原因是\(p(\boldsymbol{x}_{t-1})\)\(p(\boldsymbol{x}_t)\)均不知道。因为根据定义,我们有

\[p(\boldsymbol{x}_t) = \int p(\boldsymbol{x}_t|\boldsymbol{x}_0)\tilde{p}(\boldsymbol{x}_0)d\boldsymbol{x}_0\]

其中\(p(\boldsymbol{x}_t|\boldsymbol{x}_0)\)是知道的,而数据分布\(\tilde{p}(\boldsymbol{x}_0)\)无法提前预知,所以不能进行计算。不过,有两个特殊的例子,是可以直接将两者算出来的,这里我们也补充计算一下,其结果也正好是上一篇文章遗留的方差选取问题的答案。

第一个例子是整个数据集只有一个样本,不失一般性,假设该样本为\(\boldsymbol{0}\),此时\(\tilde{p}(\boldsymbol{x}_0)\)为狄拉克分布\(\delta(\boldsymbol{x}_0)\),可以直接算出\(p(\boldsymbol{x}_t)=p(\boldsymbol{x}_t|\boldsymbol{0})\)。继而代入式(7),可以发现结果正好是\(p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\boldsymbol{x}_0)\)\(\boldsymbol{x}_0=\boldsymbol{0}\)的特例,即

\[p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t) = p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t, \boldsymbol{x}_0=\boldsymbol{0}) = \mathcal{N}\left(\boldsymbol{x}_{t-1};\frac{\alpha_t\bar{\beta}_{t-1}^2}{\bar{\beta}_t^2}\boldsymbol{x}_t,\frac{\bar{\beta}_{t-1}^2\beta_t^2}{\bar{\beta}_t^2} \boldsymbol{I}\right)\]

我们主要关心其方差为\(\frac{\bar{\beta}_{t-1}^2\beta_t^2}{\bar{\beta}_t^2}\),这便是采样方差的选择之一。

第二个例子是数据集服从标准正态分布,即\(\tilde{p}(\boldsymbol{x}_0)=\mathcal{N}(\boldsymbol{x}_0;\boldsymbol{0},\boldsymbol{I})\)。前面我们说了\(p(\boldsymbol{x}_t|\boldsymbol{x}_0)=\mathcal{N}(\boldsymbol{x}_t;\bar{\alpha}_t \boldsymbol{x}_0,\bar{\beta}_t^2 \boldsymbol{I})\)意味着\(\boldsymbol{x}_t = \bar{\alpha}_t \boldsymbol{x}_0 + \bar{\beta}_t \boldsymbol{\varepsilon},\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})\),而此时根据假设还有\(\boldsymbol{x}_0\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})\),所以由正态分布的叠加性,\(\boldsymbol{x}_t\)正好也服从标准正态分布。将标准正态分布的概率密度代入式(7)后,结果的指数部分除掉\(-1/2\)因子外,结果是:

\[\frac{\Vert \boldsymbol{x}_t - \alpha_t \boldsymbol{x}_{t-1}\Vert^2}{\beta_t^2} + \Vert \boldsymbol{x}_{t-1}\Vert^2 - \Vert \boldsymbol{x}_t\Vert^2\]

跟推导\(p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\boldsymbol{x}_0)\)的过程类似,可以得到上述指数对应于

\[p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t) = \mathcal{N}\left(\boldsymbol{x}_{t-1};\alpha_t\boldsymbol{x}_t,\beta_t^2 \boldsymbol{I}\right)\]

我们同样主要关心其方差为\(\beta_t^2\),这便是采样方差的另一个选择。

通过ELBO推导

对数似然下界(Evidence Lower Bound, 简称 ELBO)是现代生成模型(尤其是 VAE 和扩散模型)的基石。简单来说,它就像是我们在攀登一座险峻的高山(最大化数据似然)时,找到的一条虽不完美但可行的登山路径。

在扩散模型中,ELBO 不仅仅是一个数学推导的中间产物,它直接指导了我们如何将复杂的图像生成问题转化为简单的噪声预测问题

什么是 ELBO?

在生成模型中,我们的终极目标是最大化模型生成真实数据 \(x\) 的概率,即最大化 对数似然(Log-Likelihood)\(\log p(x)\)

对于大多数复杂的模型,直接计算 \(\log p(x)\) 是不可能的(Intractable)。因为这需要积分掉所有的潜在变量 \(z\)

\[p(x) = \int p(x|z)p(z) dz \]

在高维空间中,这个积分根本算不出来。

既然算不出确切的高度(似然),数学家们想出了一个办法:找一个下界(Lower Bound)。如果我们能不断抬高这个地板(Maximize ELBO),那么天花板(Log-Likelihood)通常也会跟着升高。

根据 Jensen 不等式,我们可以推导出:

\[ \log p(x) \ge \mathbb{E}_{q(z|x)}[\log p(x|z)] - D_{KL}(q(z|x) || p(z)) = \text{ELBO} \]

或者更直观的写法:

\[ \log p(x) = \text{ELBO} + D_{KL}(\text{近似后验} || \text{真实后验}) \]

由于 KL 散度永远非负(\(\ge 0\)),所以 ELBO 永远小于或等于对数似然。

  • ELBO 的作用:它把一个不可算的积分问题,转化成了一个可优化的期望问题。

ELBO 在扩散模型中的角色

扩散模型(Diffusion Models)本质上是一类 特殊的、层级很深的变分自编码器(Hierarchical VAE)。理解了这一点,就能明白 ELBO 在其中的统治地位。

扩散模型的特殊结构

与传统 VAE 不同,扩散模型的结构非常独特:

  • 编码过程(Forward Process, \(q\)):是固定的,没有参数!它只是不断地加高斯噪声,直到图片变成纯噪声 \(x_T\)
  • 解码过程(Reverse Process, \(p_\theta\)):是我们需要学习的,它负责一步步去噪,从 \(x_T\) 恢复到 \(x_0\)

扩散模型的 ELBO 推导

在扩散模型中,潜在变量 \(z\)就是中间所有的噪声图像 \(x_{1:T}\)。我们要最大化 \(x_0\) 的似然,于是写出 ELBO:

\[\begin{aligned}-\log p_{\theta}(\mathbf{x}_{0})&\leq-\log p_{\theta}(\mathbf{x}_{0})+D_{\mathrm{KL}}(q(\mathbf{x}_{1:T}|\mathbf{x}_{0})\|p_{\theta}(\mathbf{x}_{1:T}|\mathbf{x}_{0}))\\&=-\log p_\theta(\mathbf{x}_0)+\mathbb{E}_{\mathbf{x}_{1:T\sim q}(\mathbf{x}_{1:T}|\mathbf{x}_0)}\Big[\log\frac{q(\mathbf{x}_{1:T}|\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})/p_\theta(\mathbf{x}_0)}\Big]\\&=-\log p_\theta(\mathbf{x}_0)+\mathbb{E}_q\Big[\log\frac{q(\mathbf{x}_{1:T}|\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})}+\log p_\theta(\mathbf{x}_0)\Big]\\&=\mathbb{E}_q\Big[\log\frac{q(\mathbf{x}_{1:T}|\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})}\Big]\\\mathrm{Let~}L_{\mathrm{VLB}}&=\mathbb{E}_{q(\mathbf{x}_{0:T})}\Big[\log\frac{q(\mathbf{x}_{1:T}|\mathbf{x}_{0})}{p_{\theta}(\mathbf{x}_{0:T})}\Big]\geq-\mathbb{E}_{q(\mathbf{x}_{0})}\log p_{\theta}(\mathbf{x}_{0})\end{aligned}\]

经过数学展开(利用马尔可夫链性质),这个复杂的式子会神奇地分解为三个部分:

\[\begin{aligned} L_{\mathrm{VLB}}& =\mathbb{E}_{q(\mathbf{x}_{0:T})}\bigg[\log\frac{q(\mathbf{x}_{1:T}|\mathbf{x}_{0})}{p_{\theta}(\mathbf{x}_{0:T})}\bigg] \\ &=\mathbb{E}_q\Big[\log\frac{\prod_{t=1}^Tq(\mathbf{x}_t|\mathbf{x}_{t-1})}{p_\theta(\mathbf{x}_T)\prod_{t=1}^Tp_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)}\Big] \\ &=\mathbb{E}_q\Big[-\log p_\theta(\mathbf{x}_T)+\sum_{t=1}^T\log\frac{q(\mathbf{x}_t|\mathbf{x}_{t-1})}{p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)}\Big] \\ &=\mathbb{E}_q\Big[-\log p_\theta(\mathbf{x}_T)+\sum_{t=2}^T\log\frac{q(\mathbf{x}_t|\mathbf{x}_{t-1})}{p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)}+\log\frac{q(\mathbf{x}_1|\mathbf{x}_0)}{p_\theta(\mathbf{x}_0|\mathbf{x}_1)}\Big] \\ &=\mathbb{E}_q\Big[-\log p_\theta(\mathbf{x}_T)+\sum_{t=2}^T\log\Big(\frac{q(\mathbf{x}_{t-1}|\mathbf{x}_t,\mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)}\cdot\frac{q(\mathbf{x}_t|\mathbf{x}_0)}{q(\mathbf{x}_{t-1}|\mathbf{x}_0)}\Big)+\log\frac{q(\mathbf{x}_1|\mathbf{x}_0)}{p_\theta(\mathbf{x}_0|\mathbf{x}_1)}\Big]\\ &=\mathbb{E}_q\Big[-\log p_\theta(\mathbf{x}_T)+\sum_{t=2}^T\log\frac{q(\mathbf{x}_{t-1}|\mathbf{x}_t,\mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)}+\sum_{t=2}^T\log\frac{q(\mathbf{x}_t|\mathbf{x}_0)}{q(\mathbf{x}_{t-1}|\mathbf{x}_0)}+\log\frac{q(\mathbf{x}_1|\mathbf{x}_0)}{p_\theta(\mathbf{x}_0|\mathbf{x}_1)}\Big] \\ &=\mathbb{E}_q\Big[-\log p_\theta(\mathbf{x}_T)+\sum_{t=2}^T\log\frac{q(\mathbf{x}_{t-1}|\mathbf{x}_t,\mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)}+\log\frac{q(\mathbf{x}_T|\mathbf{x}_0)}{q(\mathbf{x}_1|\mathbf{x}_0)}+\log\frac{q(\mathbf{x}_1|\mathbf{x}_0)}{p_\theta(\mathbf{x}_0|\mathbf{x}_1)}\Big] \\ &=\mathbb{E}_{q}\Big[\log\frac{q(\mathbf{x}_{T}|\mathbf{x}_{0})}{p_{\theta}(\mathbf{x}_{T})}+\sum_{t=2}^{T}\log\frac{q(\mathbf{x}_{t-1}|\mathbf{x}_{t},\mathbf{x}_{0})}{p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_{t})}-\log p_{\theta}(\mathbf{x}_{0}|\mathbf{x}_{1})\Big] \\ &=\mathbb{E}_{q}[\underbrace{D_{\mathrm{KL}}(q(\mathbf{x}_{T}|\mathbf{x}_{0})\parallel p_{\theta}(\mathbf{x}_{T}))}_{L_{T}}+\sum_{t=2}^{T}\underbrace{D_{\mathrm{KL}}(q(\mathbf{x}_{t-1}|\mathbf{x}_{t},\mathbf{x}_{0})\parallel p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_{t}))}_{L_{t-1}}\underbrace{-\log p_{\theta}(\mathbf{x}_{0}|\mathbf{x}_{1})}_{L_{0}}] \end{aligned}\]

让我们逐一看看这三项在理论中的作用:

  • 第一项:先验匹配 (\(L_T\))
    • 含义:比较最终加噪后的分布 \(q(x_T|x_0)\) 和标准高斯分布 \(p(x_T)\)
    • 作用:因为前向过程是固定的(加噪直到变成纯噪声),这一项是一个常数(接近 0),求导后可以省略。在训练中可以直接忽略。
  • 第二项:去噪匹配 (\(L_{t-1}\))
    • 含义:这是 ELBO 的灵魂。它要求我们的模型 \(p_\theta\)(逆向去噪一步)尽可能接近真实的前向过程的逆运算 \(q(x_{t-1}|x_t, x_0)\)
      • \(q(x_{t-1}|x_t, x_0)\)(已知 \(x_0\) 时的后验)是可以解析计算的,它是一个高斯分布。
      • 我们的模型 \(p_\theta\)也是预测一个高斯分布。
      • 两个高斯分布的 KL 散度,等价于它们均值(Mean)之间的均方误差(MSE)。
    • 作用:这一项告诉我们,训练扩散模型,本质上就是让模型预测的均值去逼近真实的后验均值。
  • 第三项:重建项 (\(L_0\))
    • 含义:最后一步,从微弱噪声 \(x_1\) 恢复出原图 \(x_0\) 的概率。
    • 作用:通常通过像素级的似然来计算,保证最后生成的图片细节准确。

这里论文将\(p_\theta(x_{t-1}|x_t)\)分布的方差设置成一个与\(\beta\)相关的常数,因此可训练的参数只存在于其均值中。对于两个单一变量的高斯分布\(p\)\(q\)而言,它们的KL散度为

\[KL(p,q)=log\frac{\sigma_2}{\sigma_1}+\frac{\sigma^2+(\mu_1-\mu_2)^2}{2\sigma_2^2}-\frac12\]

可得,

\[L_{t-1}=\mathbb{E}_q\bigg[\frac{1}{2\sigma_t^2}\|\tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t,\mathbf{x}0)-\boldsymbol{\mu}\theta(\mathbf{x}_t,t)\|^2\bigg]+C\]

既然这里的loss是从KL divergence出发的,或者说是与分布有关的那我们可以设计一个黑箱子神经网络,把它称之为\(D_\theta\) 网络。对于\(D_\theta\) 网络,输入是\(x_t\)和时间编码\(t\),对于输出是什么,取决于我们的建模目标。

  • 直观的做法是让\(D_\theta\)网络的输出等于前向过程中的后验分布均值\(\tilde{\mu}(x_t,x_0)\),这种建模方法俗称预测后验分布的期望值;
  • 根据\(\tilde{\mu } ( x_t, x_0)\) 表达式 , 它里面的 \(x_0\) 对于\(D_\theta\) 网络是未知的,因此第二种做法是让\(D_\theta\)网络的输出等于\(x_0\), 这种做法即直接预测原始数据。有人问,既然可以通过\(D_\theta\)网络直接预测\(x_0\)了 , 那是不是采样过程就直接计算\(D_\theta(x_T,T)\)的输出即可认为是生成了样本了呢?答案是直接一步到位,质量会比较差,还是需要通过马尔科夫高斯条件迭代而获得最终高质量的生成样本;
  • 当我们把\(\tilde{\mu } ( x_t, x_0)\) 中的\(x_0\)\(x_t\)去表示的时候,\(\tilde{\mu}(x_t,x_0)\)就变成了如下只包含\(x_t\)和随机变量\(\epsilon\)的式子。其中\(x_t\)对于\(D\)网络是已知的,而 \(\epsilon\) 是未知的,因此这个时候,我们可以选择建模目标是让\(D_\theta\) 网络的输出等于\(\epsilon\) 了,这种建模方法俗称预测随机变量(噪音)法:
\[\begin{gathered} L_{t-1}-C =\mathbb{E}_{\mathbf{x}_{0},\boldsymbol{\epsilon}}\left[\frac{1}{2\sigma_{t}^{2}}\left\|\tilde{\boldsymbol{\mu}}_{t}\left(\mathbf{x}_{t}(\mathbf{x}_{0},\boldsymbol{\epsilon}),\frac{1}{\sqrt{\bar{\alpha}_{t}}}(\mathbf{x}_{t}(\mathbf{x}_{0},\boldsymbol{\epsilon})-\sqrt{1-\bar{\alpha}_{t}}\boldsymbol{\epsilon})\right)-\boldsymbol{\mu}_{\theta}(\mathbf{x}_{t}(\mathbf{x}_{0},\boldsymbol{\epsilon}),t)\right\|^{2}\right] \\ =\mathbb{E}_{\mathbf{x}_{0},\boldsymbol{\epsilon}}\left[\frac{1}{2\sigma_{t}^{2}}\left\|\frac{1}{\sqrt{\alpha_{t}}}\left(\mathbf{x}_{t}(\mathbf{x}_{0},\boldsymbol{\epsilon})-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}}\boldsymbol{\epsilon}\right)-\boldsymbol{\mu}_{\theta}(\mathbf{x}_{t}(\mathbf{x}_{0},\boldsymbol{\epsilon}),t)\right\|^{2}\right] \end{gathered}\]

于是\(L_{t-1}\)可以化简成如下表达式

\[\mathbb{E}_{\mathbf{x}_0,\boldsymbol{\epsilon}}\left[\frac{\beta_t^2}{2\sigma_t^2\alpha_t(1-\bar{\alpha}t)}\left\|\boldsymbol{\epsilon}-\boldsymbol{\epsilon}\theta(\sqrt{\bar{\alpha}_t}\mathbf{x}_0+\sqrt{1-\bar{\alpha}_t}\boldsymbol{\epsilon},t)\right\|^2\right]\]

DDPM作者又发现,干脆将系数丢掉,训练更加稳定质量更好,于是有了下面的\(L_{simple}\)

\[L_{\mathrm{simple}}(\theta):=\mathbb{E}_{t,\mathbf{x}0,\boldsymbol{\epsilon}}\Big[\left\|\boldsymbol{\epsilon}-\boldsymbol{\epsilon}_\theta(\sqrt{\bar{\alpha}_t}\mathbf{x}_0+\sqrt{1-\bar{\alpha}_t}\boldsymbol{\epsilon},t)\right\|^2\Big]\]

其中 \(\epsilon_{\theta}\) 可以理解为使用神经网络对噪声预测,此处的噪声可以看成是数据分布的增广,此处的神经网络一般用 Unet

至此,也推导出了完整训练目标,也就得到了原论文中的训练伪代码:

image.png

Reference

生成扩散模型漫谈(一):DDPM = 拆楼 + 建楼

生成扩散模型漫谈(三):DDPM = 贝叶斯 + 去噪