【论文笔记 - DDPM】Deep Unsupervised Learning using Nonequilibrium Thermodynamics
Read-through
Abstract
We present high quality image synthesis results using diffusion probabilistic models, a class of latent variable models inspired by considerations from nonequilibrium thermodynamics. Our best results are obtained by training on a weighted variational bound designed according to a novel connection between diffusion probabilistic models and denoising score matching with Langevin dynamics, and our models naturally admit a progressive lossy decompression scheme that can be interpreted as a generalization of autoregressive decoding. On the unconditional CIFAR10 dataset, we obtain an Inception score of 9.46 and a state-of-the-art FID score of 3.17. On 256x256 LSUN, we obtain sample quality similar to ProgressiveGAN. Our implementation is available at https://github.com/hojonathanho/diffusion.
数学推导【转载】
数学推导过程来自苏剑林大神的《生成扩散模型漫谈》系列,感谢苏神的无私奉献,深入浅出的讲解让我这样数学功底不好的人也能领略这个当下最为火爆的模型的精髓。
系列中有部分步骤,一眼看过去可能有些费解,所以这里稍微做了展开,作为自己的笔记用。
通俗解释:DDPM=拆楼+建楼
生成模型实际上就是:随机噪声 \(\boldsymbol{z}\ \xrightarrow{变换}\) 样本数据 \(\boldsymbol{x}\)
我们把“拆楼”分为 \(T\) 步:
如果能学会 \(\boldsymbol{x}_{t-1}=\boldsymbol{\mu}(\boldsymbol{x}_t)\),那么反复执行 \(\boldsymbol{x}_{T-1}=\boldsymbol{\mu}(\boldsymbol{x}_T),\,\boldsymbol{x}_{T-2}=\boldsymbol{\mu}(\boldsymbol{x}_{T-1}),\,\cdots,\,\boldsymbol{x}_1=\boldsymbol{\mu}(\boldsymbol{x}_0)\) 即可还原 \(\boldsymbol{x}_0\)
该怎么拆
DDPM将“拆楼”建模为:
\[\begin{equation}\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})\end{equation} \]
其中 \(\alpha_t,\beta_t > 0\) 且 \(\alpha_t^2 + \beta_t^2=1\),通常 \(\beta_t\rightarrow 0\);\(\boldsymbol{\varepsilon}_t\) 为噪声。
反复执行这个拆楼的步骤,可以得到:
式中花括号指出的部分可以看成一个整体的噪声。利用正态分布的叠加性:
显然这些噪声的均值为 \(0\),我们来算它们的方差之和:
所以实际上相当于有:
\[\begin{equation}\begin{aligned}\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})\\=&\bar{\alpha}_t\boldsymbol{x}_0+\bar{\beta}_t\bar{\boldsymbol{\varepsilon}}_t,\quad \bar{\boldsymbol{\varepsilon}}_t\sim\mathcal{N}(\boldsymbol{0},\, \boldsymbol{I})\end{aligned}\end{equation} \]
此外DDPM还会选择适当的 \(\alpha_t\),使得 \(\bar{\alpha}_T\approx 0\),这意味着经过 \(T\) 步的拆楼后,所剩的楼体几乎可以忽略了,已经全部转化为原材料 \(\boldsymbol{\varepsilon}\)。
又如何建
\(\boldsymbol{x}_{t-1}\to \boldsymbol{x}_t\) 有了,现在我们要学习 \(\boldsymbol{x}_t\to \boldsymbol{x}_{t-1}\)。设该模型为 \(\boldsymbol{\mu}(\boldsymbol{x}_t)\),那么学习方案就是最小化欧氏距离:
\[\begin{equation}\|\boldsymbol{x}_{t-1} - \boldsymbol{\mu}(\boldsymbol{x}_t)\|^2\end{equation} \]
首先根据 \((2)\),反解 \(\boldsymbol{x}_{t-1}\) 就是 \(\boldsymbol{x}_{t-1} = \dfrac{1}{\alpha_t}\left(\boldsymbol{x}_t - \beta_t \boldsymbol{\varepsilon}_t\right)\)。所以我们就可以将 \(\boldsymbol{\mu}(\boldsymbol{x}_t)\) 设成:
\[\begin{equation}\boldsymbol{\mu}(\boldsymbol{x}_t) = \dfrac{1}{\alpha_t}\left(\boldsymbol{x}_t - \beta_t \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,\,t)\right)\end{equation} \]
其中 \(\boldsymbol{\theta}\) 是训练参数。代入到 \((7)\) 中,损失函数即:
忽略掉权重 \(\dfrac{\beta_t^2}{\alpha_t^2}\),另外结合 \((2)\ (6)\) 可以将 \(\boldsymbol{x}_t\) 化为:
最终损失函数的形式为:
\[\begin{equation}\| \boldsymbol{\varepsilon}_t - \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,\, t)\|^2=\| \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)\|^2\end{equation} \]
为什么需要 \(\boldsymbol{x}_t = \alpha_t\boldsymbol{x}_{t-1} + \beta_t \boldsymbol{\varepsilon}_t\) 这一步呢?这是因为 \(\bar{\boldsymbol{\varepsilon}}_t\) 和 \(\boldsymbol{\varepsilon}_t\) 不是相互独立的,所以只能用 \(\bar{\boldsymbol{\varepsilon}}_{t-1}\) 和 \(\boldsymbol{\varepsilon}_t\)。
降低方差
原则上 \((11)\) 就可以完成DDPM的训练,但由于这个式子中需要对 \(\boldsymbol{x}_0,\,\bar{\boldsymbol{\varepsilon}}_{t-1},\, \boldsymbol{\varepsilon}_t,\,t\) 四个随机变量分别采样,在实践中可能有方差过大的风险,从而导致收敛过慢等问题。我们可以将 \(\bar{\boldsymbol{\varepsilon}}_{t-1},\,\boldsymbol{\varepsilon}_t\) 合并成单个随机变量,从而缓解方差大的问题。
首先推一下 \(\bar{\beta}_{t-1}^2\) 和 \(\beta_t^2,\,\bar{\beta}_t^2\) 的关系:
然后和上面做过的事情一样,利用一下正态分布的叠加性:
\(\alpha_t\bar{\beta}_{t-1}\bar{\boldsymbol{\varepsilon}}_{t-1} + \beta_t \boldsymbol{\varepsilon}_t\) 均值为 \(0\),方差为 \(\alpha_t^2\bar{\beta}_{t-1}^2 + \beta_t^2=\alpha_t^2\dfrac{\bar{\beta}_{t}^2-\beta_t^2}{1-\beta_t^2} + \beta_t^2=\beta_t^2\),实际相当于 \(\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\) 均值为 \(0\),方差为 \(\alpha_t^2\bar{\beta}_{t-1}^2 + \beta_t^2=\alpha_t^2\dfrac{\bar{\beta}_{t}^2-\beta_t^2}{1-\beta_t^2} + \beta_t^2=\beta_t^2\),实际相当于 \(\bar{\beta}_t\boldsymbol{\omega}|\boldsymbol{\omega}\sim \mathcal{N}(\boldsymbol{0},\, \boldsymbol{I})\)
然后我们来验证一下 \(\boldsymbol{\varepsilon}\) 和 \(\boldsymbol{\omega}\) 是两个相互独立的正态随机变量。这可以用协方差为零证明,不过我们也可以通过 \(\mathbb{E}[\boldsymbol{\varepsilon}\boldsymbol{\omega}^{\top}]=\boldsymbol{0}\) 来说明。我们先算 \(\mathbb{E}[(\bar{\beta}_t\boldsymbol{\varepsilon})(\bar{\beta}_t\boldsymbol{\omega})^{\top}]\):
于是我们也就证明了 \(\mathbb{E}[\boldsymbol{\varepsilon}\boldsymbol{\omega}^{\top}]=\boldsymbol{0}\)。这里用到了结论:\(\boldsymbol{\varepsilon}_1,\,\boldsymbol{\varepsilon}_2\sim \mathcal{N}(\boldsymbol{0},\, \boldsymbol{I})\) ,且 \(\boldsymbol{\varepsilon}_1,\,\boldsymbol{\varepsilon}_2\) 独立,则有 \(\mathbb{E}[\boldsymbol{\varepsilon}_1\boldsymbol{\varepsilon}_1^{\top}]=\mathbb{E}[\boldsymbol{\varepsilon}_2\boldsymbol{\varepsilon}_2^{\top}]=\boldsymbol{I}\),\(\mathbb{E}[\boldsymbol{\varepsilon}_1\boldsymbol{\varepsilon}_2^{\top}]=\mathbb{E}[\boldsymbol{\varepsilon}_1^{\top}\boldsymbol{\varepsilon}_2]=\boldsymbol{0}\)。
接下来我们反过来解 \(\boldsymbol{\varepsilon}_t\):
解得:
代入 \((11)\) 式得到:
我们先来处理 \(\boldsymbol{\omega}\):
直接打开,注意到 \(\mathbb{E}\boldsymbol{\omega}\) 和 \(\mathbb{E}\boldsymbol{\omega}^2\) 都是常数,所以损失函数就相当于:
再次忽略常数和权重,我们得到DDPM最终所用的损失函数:
\[\begin{equation}\mathbb{E}_{\boldsymbol{\varepsilon}\sim \mathcal{N}(\boldsymbol{0},\, \boldsymbol{I})}\left\|\boldsymbol{\varepsilon}-\dfrac{\bar{\beta}_t}{\beta_t}\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\bar{\alpha}_t\boldsymbol{x}_0 + \bar{\beta}_t\boldsymbol{\varepsilon},\, t)\right\|^2\end{equation} \]
这个形式和DDPM原论文中的 \(L_{\mathrm{simple}}(\theta)\) 是完全一致的:
递归生成
训练完之后,我们就可以从一个随机噪声 \(\boldsymbol{x}_T\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})\) 出发执行 \(T\) 步 \((8)\) 来进行生成:
这对应于自回归解码中的Greedy Search。如果要进行Random Sample,那么需要补上噪声项:
一般来说,我们可以让 \(\sigma_t=\beta_t\),即正向和反向的方差保持同步。
超参设置
在DDPM中,\(T=1000\),\(\alpha_t = \sqrt{1 - \dfrac{0.02t}{T}}\)。
在重构的时候我们用了欧氏距离 \((7)\) 作为损失函数,而一般我们用DDPM做图片生成,以往做过图片生成的读者都知道,欧氏距离并不是图片真实程度的一个好的度量,VAE用欧氏距离来重构时,往往会得到模糊的结果,除非是输入输出的两张图片非常接近,用欧氏距离才能得到比较清晰的结果,所以选择尽可能大的 \(T\),正是为了使得输入输出尽可能相近,减少欧氏距离带来的模糊问题。
为什么要选择单调递减的 \(\alpha_t\) 呢?当 \(t\) 比较小时,\(x_t\) 还比较接近真实图片,所以我们要缩小 \(x_{t−1}\) 与 \(x_t\) 的差距,以便更适用欧氏距离 \((7)\),因此要用较大的 \(\alpha_t\);当 \(t\) 比较大时,\(x_t\) 已经比较接近纯噪声了,噪声用欧式距离无妨,所以可以稍微增大 \(x_{t−1}\) 与 \(x_t\) 的差距,即可以用较小的 \(\alpha_t\)。那么可不可以一直用较大的 \(α_t\) 呢?可以是可以,但是要增大 \(T\)。
我们之前说过,应该有 \(\bar{\alpha}_T\approx 0\),我们利用 \(\alpha_t\) 的表达式来计算 \(\bar{\alpha}_T\):
由此可以看出 \(T\) 要足够大,才能达到 \(\approx0\) 的标准。当 \(T=1000\) 时,\(\bar{\alpha}_T\approx \mathrm{e}^{-5}\)。
最后我们留意到,“建楼”模型中的 \(\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\bar{\alpha}_t\boldsymbol{x}_0 + \bar{\beta}_t\boldsymbol{\varepsilon}, t)\) 中,我们在输入中显式地写出了 \(t\),这是因为原则上不同的 \(t\) 处理的是不同层次的对象,所以应该用不同的重构模型,即应该有 \(T\) 个不同的重构模型才对,于是我们共享了所有重构模型的参数,将 \(t\) 作为条件传入。按照论文附录的说法,\(t\) 是转换成位置编码后,直接加到残差模块上去的。
VAE角度
多步突破
在传统VAE中,编码过程和生成过程都是一步到位的:
DDPM将编码过程和生成过程分解为 \(T\) 步:
联合散度
每一步编码过程被建模成 \(q(\boldsymbol{x}_t|\boldsymbol{x}_{t-1})\),每一步生成过程被建模成 \(p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t)\),相应的联合分布就是:
其中 \(\boldsymbol{x}_0\) 代表真实样本,所以 \(\tilde{q}(\boldsymbol{x}_0)\) 就是数据分布;而 \(\boldsymbol{x}_T\) 代表着最终的编码,所以 \(p(\boldsymbol{x}_T)\) 就是先验分布;剩下的 \(q(\boldsymbol{x}_t|\boldsymbol{x}_{t-1})\)、\(p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t)\) 就代表着编码、生成的一小步。
VAE可以理解为在最小化联合分布的KL散度,对于DDPM也是如此,上面我们已经写出了两个联合分布,所以DDPM的目的就是最小化
\[\displaystyle\begin{equation}KL(q\Vert p) = \int q(\boldsymbol{x}_T|\boldsymbol{x}_{T-1})\cdots q(\boldsymbol{x}_1|\boldsymbol{x}_0) \tilde{q}(\boldsymbol{x}_0) \log \frac{q(\boldsymbol{x}_T|\boldsymbol{x}_{T-1})\cdots q(\boldsymbol{x}_1|\boldsymbol{x}_0) \tilde{q}(\boldsymbol{x}_0)}{p(\boldsymbol{x}_0|\boldsymbol{x}_1)\cdots p(\boldsymbol{x}_{T-1}|\boldsymbol{x}_T) p(\boldsymbol{x}_T)} d\boldsymbol{x}_0 d\boldsymbol{x}_1\cdots d\boldsymbol{x}_T\end{equation} \]
接下来,我们要将 \(q(\boldsymbol{x}_t|\boldsymbol{x}_{t-1})\)、\(p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t)\) 的具体形式定下来,然后简化DDPM的优化目标。
分而治之
DDPM将每一步的编码建立为正态分布:\(q(\boldsymbol{x}_t|\boldsymbol{x}_{t-1})=\mathcal{N}(\boldsymbol{x}_t;\,\alpha_t \boldsymbol{x}_{t-1},\, \beta_t^2 \boldsymbol{I})\),其主要的特点是均值向量仅仅由输入 \(\boldsymbol{x}_{t-1}\) 乘以一个标量 \(\alpha_t\) 得到,相比之下传统VAE的均值方差都是用神经网络学习出来的,因此DDPM是放弃了模型的编码能力,最终只得到一个纯粹的生成模型;至于 \(p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t)\),则被建模成均值向量可学习的正态分布 \(\mathcal{N}(\boldsymbol{x}_{t-1};\,\boldsymbol{\mu}(\boldsymbol{x}_t),\, \sigma_t^2 \boldsymbol{I})\)。其中 \(\alpha_t,\,\beta_t,\,\sigma_t\) 都不是可训练参数,而是事先设定好的值,整个模型拥有可训练参数的就只有\(\boldsymbol{\mu}(\boldsymbol{x}_t)\)。
由于目前分布 \(q\) 不含任何可训练参数,因此目标 \((27)\) 中关于 \(q\) 的积分就只是贡献一个可以忽略的常数,目标 \((27)\) 等价于:
\[\begin{equation}\begin{aligned}&\,-\int q(\boldsymbol{x}_T|\boldsymbol{x}_{T-1})\cdots q(\boldsymbol{x}_1|\boldsymbol{x}_0) \tilde{q}(\boldsymbol{x}_0) \log p(\boldsymbol{x}_0|\boldsymbol{x}_1)\cdots p(\boldsymbol{x}_{T-1}|\boldsymbol{x}_T) p(\boldsymbol{x}_T) d\boldsymbol{x}_0 d\boldsymbol{x}_1\cdots d\boldsymbol{x}_T \\ =&\,-\int q(\boldsymbol{x}_T|\boldsymbol{x}_{T-1})\cdots q(\boldsymbol{x}_1|\boldsymbol{x}_0) \tilde{q}(\boldsymbol{x}_0) \left[\log p(\boldsymbol{x}_T) + \sum_{t=1}^T\log p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t)\right] d\boldsymbol{x}_0 d\boldsymbol{x}_1\cdots d\boldsymbol{x}_T \end{aligned}\end{equation} \]
由于先验分布 \(p(\boldsymbol{x}_T)\) 一般都取标准正态分布,也是没有参数的,所以这一项也只是贡献一个常数。因此需要计算的就是每一项
\[\begin{equation}\begin{aligned}&\,-\int q(\boldsymbol{x}_T|\boldsymbol{x}_{T-1})\cdots q(\boldsymbol{x}_1|\boldsymbol{x}_0) \tilde{q}(\boldsymbol{x}_0) \log p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t) d\boldsymbol{x}_0 d\boldsymbol{x}_1\cdots d\boldsymbol{x}_T\\ =&\,-\int q(\boldsymbol{x}_t|\boldsymbol{x}_{t-1})\cdots q(\boldsymbol{x}_1|\boldsymbol{x}_0) \tilde{q}(\boldsymbol{x}_0) \log p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t) d\boldsymbol{x}_0 d\boldsymbol{x}_1\cdots d\boldsymbol{x}_t\\ =&\,-\int q(\boldsymbol{x}_t|\boldsymbol{x}_{t-1})q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_0) \tilde{q}(\boldsymbol{x}_0) \log p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t) d\boldsymbol{x}_0 d\boldsymbol{x}_{t-1}d\boldsymbol{x}_t \end{aligned}\end{equation} \]
这两个等号分别是因为:
场景再现
除去优化无关的常数:
这和第一种推导过程中的 \((7)\) 式是一样的,同样的处理方法可以得到与 \((19)\) 式相同的目标函数:
\[\begin{equation}\mathbb{E}_{\boldsymbol{\varepsilon}\sim \mathcal{N}(\boldsymbol{0},\boldsymbol{I}),\,\boldsymbol{x}_0\sim \tilde{p}(\boldsymbol{x}_0)}\left[\left\Vert\boldsymbol{\varepsilon} - \dfrac{\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]\end{equation} \]
当然这里的系数已经去掉了(原论文中通过实验发现,去掉这个系数后的实际效果更好些)。
超参设置
对于 \(q(\boldsymbol{x}_t|\boldsymbol{x}_{t-1})\) 来说,习惯上约定 \(\alpha_t^2 + \beta_t^2=1\)。在第一种推导 \((6)\) 式中已经证明,由于正态分布的叠加性,在此约束之下我们有 \({x}_t =\bar{\alpha}_t\boldsymbol{x}_0+\bar{\beta}_t\bar{\boldsymbol{\varepsilon}}_t,\, \bar{\boldsymbol{\varepsilon}}_t\sim\mathcal{N}(\boldsymbol{0},\, \boldsymbol{I})\),对应在这里的表示就是:
另一方面,\(p(\boldsymbol{x}_T)\) 一般都取标准正态分布 \(\mathcal{N}(\boldsymbol{x}_T;\,\boldsymbol{0},\, \boldsymbol{I})\)。而我们的学习目标是最小化两个联合分布的KL散度,即希望 \(p=q\),那么它们的边缘分布自然也相等,所以我们也希望
由于数据分布 \(\tilde{q}(\boldsymbol{x}_0)\) 是任意的,所以要使上式恒成立,只能让 \(q(\boldsymbol{x}_T|\boldsymbol{x}_0)=p(\boldsymbol{x}_T)\),即退化为与 \(\boldsymbol{x}_0\) 无关的标准正态分布,这意味着我们要设计适当的 \(\alpha_t\),使得 \(\bar{\alpha}_T\approx 0\)。同时这再次告诉我们,DDPM是没有编码能力了,最终的 \(p(\boldsymbol{x}_T|\boldsymbol{x}_0)\) 可以说跟输入 \(\boldsymbol{x}_0\) 无关的,生成出来的图像也无法回到原图。
至于 \(\sigma_t\),理论上不同的数据分布 \(\tilde{q}(\boldsymbol{x}_0)\) 来说对应不同的最优 \(\sigma_t\),但我们又不想将 \(\sigma_t\) 设为可训练参数,所以只好选一些特殊的 \(\tilde{q}(\boldsymbol{x}_0)\) 来推导相应的最优 \(\sigma_t\),并认为由特例推导出来的 \(\sigma_t\) 可以泛化到一般的数据分布。我们可以考虑两个简单的例子:
- 假设训练集只有一个样本\(\boldsymbol{x}_{*}\),即 \(\tilde{q}(\boldsymbol{x}_0)\) 是狄拉克分布 \(\delta(\boldsymbol{x}_0 - \boldsymbol{x}_*)\),可以推出最优的 \(\sigma_t = \dfrac{\bar{\beta}_{t-1}}{\bar{\beta}_t}\beta_t\);
- 假设数据分布 \(\tilde{q}(\boldsymbol{x}_0)\) 服从标准正态分布,这时候可以推出最优的 \(\sigma_t = \beta_t\)。
具体的推导在下面一节的“遗留问题”部分给出。实验结果显示两个选择的表现是相似的,因此可以选择任意一个进行采样。
贝叶斯角度
请贝叶斯
根据贝叶斯定理:
然而,我们并不知道 \(p(\boldsymbol{x}_{t-1}),\,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)\) 都是已知的(再次复习一下 \((2)\,(6)\) 两式):
所以上式是可计算的,代入各自的表达式得到:
系数部分 \(\dfrac{\bar{\beta}_t}{\sqrt{2\pi}\bar{\beta}_{t-1}\beta_t}\) 可知协方差矩阵是 \(\dfrac{\bar{\beta}_{t-1}^2 \beta_t^2}{\bar{\beta}_t^2}\boldsymbol{I}\),在此基础上化简指数部分求出均值 \(\tilde{\boldsymbol{\mu}}(\boldsymbol{x}_t, \boldsymbol{x}_0)\),可以得到:
\[\begin{equation}p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\, \boldsymbol{x}_0) = \mathcal{N}\left(\boldsymbol{x}_{t-1};\,\dfrac{\alpha_t\bar{\beta}_{t-1}^2}{\bar{\beta}_t^2}\boldsymbol{x}_t + \dfrac{\bar{\alpha}_{t-1}\beta_t^2}{\bar{\beta}_t^2}\boldsymbol{x}_0,\,\dfrac{\bar{\beta}_{t-1}^2\beta_t^2}{\bar{\beta}_t^2} \boldsymbol{I}\right)\end{equation} \]
去噪过程
现在我们得到了 \(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\)。训练完成后,我们就认为
\[\begin{equation}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};\, \dfrac{\alpha_t\bar{\beta}_{t-1}^2}{\bar{\beta}_t^2}\boldsymbol{x}_t + \dfrac{\bar{\alpha}_{t-1}\beta_t^2}{\bar{\beta}_t^2}\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t),\,\dfrac{\bar{\beta}_{t-1}^2\beta_t^2}{\bar{\beta}_t^2} \boldsymbol{I}\right)\end{equation} \]
在 \(\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 = \dfrac{1}{\bar{\alpha}_t}\left(\boldsymbol{x}_t - \bar{\beta}_t \boldsymbol{\varepsilon}\right)\),这启发我们将 \(\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)\)参数化为
\[\begin{equation}\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t) = \dfrac{1}{\bar{\alpha}_t}\left(\boldsymbol{x}_t - \bar{\beta}_t \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,\, t)\right)\end{equation} \]
此时损失函数变为
\[\begin{equation}\Vert \boldsymbol{x}_0 - \bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)\Vert^2 = \dfrac{\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\end{equation} \]
省去前面的系数,就得到DDPM原论文所用的损失函数了。可以发现,本文是直接得出了从 \(\boldsymbol{x}_t\) 到 \(\boldsymbol{x}_0\) 的去噪过程,而不是像之前两篇文章那样,通过 \(\boldsymbol{x}_t\) 到 \(\boldsymbol{x}_{t-1}\) 的去噪过程再加上积分变换来推导,相比之下本文的推导可谓更加一步到位了。
另一边,我们将式 \((42)\) 代入到式 \((41)\) 中,均值部分:
最后一步是因为 \(\dfrac{\alpha_t^2\bar{\beta}_{t-1}^2+\beta_t^2}{\bar{\beta}_t^2\alpha_t}=\dfrac{\alpha_t^2(1-\bar{\alpha}_{t-1}^2)+\beta_t^2}{\bar{\beta}_t^2\alpha_t}=\dfrac{1-\bar{\alpha}^2_t}{\bar{\beta}_t^2\alpha_t}=\dfrac{1}{\alpha_t}\),化简得到:
\[\begin{equation} 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};\, \dfrac{1}{\alpha_t}\left(\boldsymbol{x}_t - \dfrac{\beta_t^2}{\bar{\beta}_t}\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,\, t)\right),\,\dfrac{\bar{\beta}_{t-1}^2\beta_t^2}{\bar{\beta}_t^2} \boldsymbol{I}\right)\end{equation} \]
这就是反向的采样过程所用的分布,连同采样过程所用的方差也一并确定下来了,即上一节中由狄拉克分布推导出来的 \(\sigma_t = \dfrac{\bar{\beta}_{t-1}}{\bar{\beta}_t}\beta_t\)。这里的 \(\boldsymbol{\epsilon}_{\boldsymbol{\theta}}\) 和前两种推导方式不同,反而跟DDPM原论文一致。
预估修正
不知道读者有没有留意到一个有趣的地方:我们要做的事情,就是想将 \(\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在2019年提出的《Lookahead Optimizer: k steps forward, 1 step back》,它同样也包含了预估(k steps forward)和修正(1 step back)两部分,原论文将其诠释为“快(Fast)-慢(Slow)”权重的相互结合,快权重就是预估得到的结果,慢权重则是基于预估所做的修正结果。如果愿意,我们也可以用同样的方式去诠释DDPM的“预估-修正”过程~
遗留问题
\((36)\) 式没法直接算的原因是 \(p(\boldsymbol{x}_{t-1}),\,p(\boldsymbol{x}_t)\) 未知。根据定义:
其中 \(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})\)。继而代入 \((39)\) 式,可以发现结果正好是 \(p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\,\boldsymbol{x}_0)\) 取 \(\boldsymbol{x}_0=\boldsymbol{0}\) 的特例,直接利用 \((40)\) 式的结论,即
\[\begin{equation}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};\,\dfrac{\alpha_t\bar{\beta}_{t-1}^2}{\bar{\beta}_t^2}\boldsymbol{x}_t,\,\dfrac{\bar{\beta}_{t-1}^2\beta_t^2}{\bar{\beta}_t^2} \boldsymbol{I}\right)\end{equation} \]我们主要关心其方差为 \(\dfrac{\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\) 正好也服从标准正态分布。将标准正态分布的概率密度代入 \((36)\) 式:
\[\begin{equation}p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t) = \dfrac{\frac1{\sqrt{2\pi}\beta_t}\mathrm{e}^{-\tfrac{\|\boldsymbol{x}_t-\alpha_t\boldsymbol{x}_{t-1}\|^2}{2\beta_t^2}}\frac1{\sqrt{2\pi}}\mathrm{e}^{-\tfrac{\|\boldsymbol{x}_{t-1}\|^2}{2}}}{\frac1{\sqrt{2\pi}}\mathrm{e}^{-\tfrac{\|\boldsymbol{x}_t\|^2}{2}}}=\dfrac1{\sqrt{2\pi}\beta_t}\mathrm{e}^{\tfrac{\|\boldsymbol{x}_{t-1}-\alpha_t\boldsymbol{x}_t\|^2}{2\beta_t^2}}\end{equation} = \mathcal{N}\left(\boldsymbol{x}_{t-1};\,\alpha_t\boldsymbol{x}_t,\,\beta_t^2 \boldsymbol{I}\right) \]我们同样主要关心其方差为 \(\beta_t^2\),这便是采样方差的另一个选择。
参考文献
苏剑林. (Jun. 13, 2022). 《生成扩散模型漫谈(一):DDPM = 拆楼 + 建楼 》[Blog post]. Retrieved from https://spaces.ac.cn/archives/9119
苏剑林. (Jul. 06, 2022). 《生成扩散模型漫谈(二):DDPM = 自回归式VAE 》[Blog post]. Retrieved from https://spaces.ac.cn/archives/9152
苏剑林. (Jul. 19, 2022). 《生成扩散模型漫谈(三):DDPM = 贝叶斯 + 去噪 》[Blog post]. Retrieved from https://spaces.ac.cn/archives/9164