系统理解扩散模型（Diffusion Models）：从柏拉图洞穴之喻开始（上）

本文参考"Understanding Diffusion Models: A Unified Perspective"

柏拉图洞穴之喻

洞穴之喻（Allegory of the Cave）是柏拉图在《理想国》514-521中描述的对人类知识的基本想象。

设想在洞穴中有一批囚徒，他们终身呆在那里，被铁链束缚、不能转头，只能看到面前洞壁上的影子。他们身后有一条横贯洞穴的小道，沿小道有一堵矮墙，临着矮墙上方有一堆篝火。有一些特定的人，扛着各种根据现实中的实物所做的模型走过墙后的小道，而火光则把透出墙的器具投影到囚徒面前的洞壁上。对这些囚徒而言，他们能观测到的就是这些从不曾见识过的高维实物的低维投影。

生成模型

类似地，我们在现实世界中观测到的物体也可能是从一个抽象的表示方法（例如结合了形状、颜色、大小等多种特征组合的表示）生成的。然后，正如洞穴之喻中的囚徒观测到的是三维物体的二维投影，我们观测到的物体也是这些抽象概念的某种“投影”。尽管洞穴中的人从未看到过隐藏在身后的物体实体，他们也依然能够这些实体进行推断。同样的道理，我们也能够估计表示我们观测到的数据的潜在变量（latent variable）。一旦掌握了这种表示方法，我们就可以利用其生成新的数据，这就是生成模型的底层思想。

不同的是，洞穴之喻中的潜在变量的维度比实际观测的高；而对于生成模型，我们一般是希望寻找一个更低维度的潜在表示。这大致出于两个原因：一方面，在没有很强先验知识的前提下找到一个比实际观测维度更高的表示十分困难；另一方面，采用更低维度的表示可以看作是一个压缩过程，有助于挖掘有语意意义的结构（semantically meaningful structure）。

证据下界（Evidence Lower Bound）

数学上，我们通过联合分布$p(x,z)$对观测值$x$和潜在变量$z$进行建模。那么，基于似然（likelihood-based）的生成旨在学习模型，最大化所有观测值$x$的似然函数$p(x)$。根据联合分布$p(x,z)$，我们有两种方式可以推出$p(x)$。第一种是求对于$x$的边缘概率分布：
$$\begin{array}{c}p(x)=\int p(x,z)dz\end{array}$$
第二种则是利用贝叶斯公式进行推导：
$$\begin{array}{c}p(x)=\frac{p(x,z)}{p(z|x)}\end{array}$$
直接计算并最大化似然函数$p(x)$十分困难：从(1)出发，我们需要对所有的潜在变量$z$进行积分，遇到复杂的模型更是难上加难；从(2)出发，我们则需要一个真实的潜在编码器（ground truth latent encoder）$p(x,z)$。

但是，我们可以利用这两个式子推出证据下界（Evidence Lower Bound, ELBO）。所谓的“证据”即是观测值的对数似然函数（log likelihood），而ELBO正是其下界：
$$\begin{array}{c}\log p(x)\ge {\mathbb{E}}_{q_{\varphi }(z|x)}[\log \frac{p(x,z)}{q_{\varphi }(z|x)}]\end{array}$$
$q_{\varphi }(z|x)$是一个近似变分分布，$\varphi$是要优化的参数，其旨在近似真实的后验（posterior）分布$p(z|x)$。这里，我们利用(2)推导(3)：
$$\begin{array}{rl}\log p(x)& =\log p(x)\int q_{\varphi }(z|x)dz\\ & =\int q_{\varphi }(z|x)(\log p(x))dz\\ & ={\mathbb{E}}_{q_{\varphi }(z|x)}\log p(x)\\ & ={\mathbb{E}}_{q_{\varphi }(z|x)}\left[\log \frac{p(x,z)}{p(z|x)}\right]\\ & ={\mathbb{E}}_{q_{\varphi }(z|x)}\left[\log \frac{p(x,z)q_{\varphi }(z|x)}{p(z|x)q_{\varphi }(z|x)}\right]\\ & ={\mathbb{E}}_{q_{\varphi }(z|x)}\left[\log \frac{p(x,z)}{q_{\varphi }(z|x)}\right]+{\mathbb{E}}_{q_{\varphi }(z|x)}\left[\log \frac{q_{\varphi }(z|x)}{p(z|x)}\right]\\ & ={\mathbb{E}}_{q_{\varphi }(z|x)}\left[\log \frac{p(x,z)}{q_{\varphi }(z|x)}\right]+D_{KL}(q_{\varphi }(z|x)\Vert p(z|x))\\ & \ge {\mathbb{E}}_{q_{\varphi }(z|x)}\left[\log \frac{p(x,z)}{q_{\varphi }(z|x)}\right]\end{array}$$

KL散度：$D_{KL}(P\Vert Q)={\sum }_{i}P(i)\log \frac{P(i)}{Q(i)}\ge 0$。

根据以上分析，我们引入$q_{\varphi }(z|x)$来近似真实的后验分布$p(z|x)$，这个目标可以通过最小化两者间的KL散度实现。但是，由于我们不知道$p(z|x)$，我们无法直接最小化该KL散度。不过，我们可以通过(10)发现，对数似然$\log p(x)$与参数$\varphi$无关，并且它等于ELBO和上述KL散度的和。因此，我们可以通过最大化ELBO来实现缩小$q_{\varphi }(z|x)$和$p(z|x)$间的差距。此外，优化后的ELBO是对数似然的近似，可以用来估计似然函数$p(x)$。

变分自编码器（Variational Autoencoders）

在原始的变分自编码器（Variational Autoencoders，VAE）模型中，我们就是直接最大化ELBO。为了看清它们之间的联系，我们对ELBO作进一步变形：
$$\begin{array}{rl}{\mathbb{E}}_{q_{\varphi }(z|x)}\left[\log \frac{p(x,z)}{q_{\varphi }(z|x)}\right]& ={\mathbb{E}}_{q_{\varphi }(z|x)}\left[\log \frac{p_{\theta }(x|z)p(z)}{q_{\varphi }(z|x)}\right]\\ & ={\mathbb{E}}_{q_{\varphi }(z|x)}\left[\log p_{\theta }(x|z)\right]+{\mathbb{E}}_{q_{\varphi }(z|x)}\left[\log \frac{p(z)}{q_{\varphi }(z|x)}\right]\\ & ={\mathbb{E}}_{q_{\varphi }(z|x)}\left[\log p_{\theta }(x|z)\right]-D_{KL}(q_{\varphi }(z|x)\Vert p(z))\end{array}$$
对于VAE，我们学习一个可以看作是编码器（encoder）的分布$q_{\varphi }(z|x)$，它将输入转化成可能的潜在表示；同时，我们学习一个可以解释为解码器（decoder）的函数$p_{\theta }(x|z)$，它将给定的潜在变量$z$转化成观测值$x$。那么，(14)中，第一项衡量了恢复的似然；第二项衡量了学习到的变分分布和潜在变量先验的差距。因此，最大化ELBO相当于最大化第一项并且最小化第二项，即最大化“似然”和最小化“差距”。

VAE的一个特征是通过变量$\theta$和$\varphi$对ELBO进行联合优化的方式。具体而言，VAE的编码器通常对多维正态分布进行建模，而先验通常选择为标准多维正态分布：
$$\begin{array}{rl}{q}_{\varphi }(z|x)& =\mathcal{N}(z;{\mu }_{\varphi }(x),{\sigma }_{\varphi }^2(x)I)\\ p(z)& =\mathcal{N}(z;0,I)\end{array}$$
因此，(14)式中的KL散度可以写出具体形式，而第一项可以通过蒙特卡洛估计（Monte Caro Estimate）来近似。此时，我们的目标可以写为：
$$\begin{array}{rl}& \underset{\varphi ,\theta }{\mathrm{argmax}}{\mathbb{E}}_{q_{\varphi }(z|x)}\left[\log p_{\theta }(x|z)\right]-D_{KL}(q_{\varphi }(z|x)\Vert p(z))\\ \approx & \underset{\varphi ,\theta }{\mathrm{argmax}}\sum _l^L\log p_{\theta }(x|z^{(l)})-D_{KL}(q_{\varphi }(z|x)\Vert p(z))\end{array}$$
其中，潜在变量 { z ( l ) } l = 1 L \{z^{(l)}\}^{L}_{l=1} {z(l)}l=1L​是由数据集中的所有观测值$x$通过$q_{\varphi }(z|x)$采样得到的。但是，这一随机采样过程通常是不可导的（non-differentiable）。不过，由于$q_{\varphi }(z|x)$被设计成一个特定的概率分布，我们可以通过重参数化（reparameterization）来解决这个问题。

重参数化技巧在于将一个随机变量表示成一个噪声变量的确定性函数。例如，一个正态分布中的样本$x\sim \mathcal{N}(x;\mu ,{\sigma }^2)$可以重新表示为：
$$x=\mu +\sigma ϵ,\text{with}ϵ\sim \mathcal{N}(x;0,I)$$
因此，我们可以先从一个标准正态分布分布采样，然后对结果进行相应的变换，以实现在任意一个正态分布上进行采样。在VAE模型中，$z$可以通过以下方式得到：
$$z={\mu }_{\varphi }(x)+{\sigma }_{\varphi }(x)\odot ϵ,\text{with}ϵ\sim \mathcal{N}(x;0,I)$$
$\odot$表示对应元素相乘。此时，我们可以计算$\varphi$的梯度来优化${\mu }_{\varphi }$和${\sigma }_{\varphi }$。

训练好VAE模型后，我们可以从潜在空间$p(z)$采样后输入到解码器中生成新数据。

层级变分自编码器（Hierarchical Variational Autoencoders）

层级变分自编码器（Hierarchical Variational Autoencoders，HVAE）是对VAE模型的自然扩展，顾名思义，此时我们有多层级的潜在变量，也就是说潜在变量本身也来自于一个其它更高层级、更抽象的潜在变量。回到柏拉图的洞穴之喻，正如我们观测到的三维实体是来自于更高层级的抽象潜在变量，洞穴中的囚徒是将我们观测到的三维实体作为他们二维观测值的潜在变量。因此，对于洞穴中的囚徒而言，他们的观测值此时经过了两层“潜在”建模。

对于一般的HVAE来说，某一层级的潜在变量可由之前所有潜在变量决定。在这里，我们仅讨论一类特殊的HAVE——Markovian HVAE （MHVAE）。对于MHAVE，生成过程是一个马尔可夫链（Markov chain），也就是说每层的潜在变量$z_t$仅由前一层的潜在变量$z_{t+1}$决定。

数学上，我们可以把MHVAE的联合分布和后验分布表示为：
$$\begin{array}{rl}p(x,z_{1:T})& =p(z_T)p_{\theta }(x|z_1)\prod _{t=2}^T{p}_{\theta }(z_{t-1}|z_t)\\ q_{\varphi }(z_{1:T}|x)& =q_{\varphi }(z_1|x)\prod _{t=2}^T{q}_{\varphi }(z_t|z_{t-1})\end{array}$$
进而，我们可以推导出ELBO：
$$\begin{array}{rl}\log p(x)& =\log \int p(x,z_{1:T})d{z}_{1:T}\\ & =\log \int \frac{p(x,z_{1:T})q_{\varphi }(z_{1:T}|x)}{q_{\varphi }(z_{1:T}|x)}d{z}_{1:T}\\ & =\log {\mathbb{E}}_{q_{\varphi }(z_{1:T}|x)}\left[\frac{p(x,z_{1:T})}{q_{\varphi }(z_{1:T}|x)}\right]\\ & \ge {\mathbb{E}}_{q_{\varphi }(z_{1:T}|x)}\left[\log \frac{p(x,z_{1:T})}{q_{\varphi }(z_{1:T}|x)}\right]\end{array}$$

琴生不等式（Jensen’s Inequality）: $\log ({\sum }_{i=1}^n\frac{x_i}{n})\ge {\sum }_{i=1}^n\frac{\log x_i}{n}$。

代入(19)和(20)可得：
$$\begin{array}{rl}& {\mathbb{E}}_{q_{\varphi }(z_{1:T}|x)}\left[\log \frac{p(x,z_{1:T})}{q_{\varphi }(z_{1:T}|x)}\right]\\ =& {\mathbb{E}}_{q_{\varphi }(z_{1:T}|x)}\left[\log \frac{p(z_T)p_{\theta }(x|z_1){\prod }_{t=2}^T{p}_{\theta }(z_{t-1}|z_t)}{q_{\varphi }(z_1|x){\prod }_{t=2}^T{q}_{\varphi }(z_t|z_{t-1})}\right]\end{array}$$