深度学习-5-生成模型3-变分自编码

变分自编码(VAE)

  GMM诞生于上世纪，其本身受制于当时的计算机性能以及数据量的大小，伴随着计算机技术的发展，在2010年之后，深度学习领域崛起，生成模型的研究也迎来了大的变革。很多学者研究了如何将利用神经网络模型拟合非线性的强大性能来改善生成模型的效果。2014年，Kingma.D.P 在ICLR上发表了著名的论文《Auto-encoding variational bayes》[2]，提出了变分自编码模型(Variational Auto-encoding)，其采用基于神经网络的编码器-解码器架构，利用编码器对数据的隐变量进行建模，再通过一个解码器将隐变量映射到数据空间来生成数据。同时VAE也提出了变分推断的思想，即对似然函数的变分下界来进行优化。
  VAE的出现对现代生成模型具有重大影响，其建立的基本框架，直到今天仍然被Diffusion等模型采用，就在今年5月份，ICLR颁发了首届时间检验奖，用于表彰具有长期影响力的论文，VAE的原始论文[3]获得了该奖项，由此便可说明其在生成模型领域的巨大影响。

概率模型

  与GMM相同，VAE同样采用了隐变量假设，即认为观测数据 \(x\) 受到隐变量 \(z\) 的影响，但与GMM不同的是，其采用了标准高斯分布作为隐变量的先验分布，同时引入了贝叶斯思想，即后验分布的概念，其概率图如下图3所示：

Image3: VAE 概率图

  对于图中的未知的条件概率 \(p(x|z),q(z|x)\)，VAE采用了参数化的神经网络模型对其进行逼近，其分别被称为解码器 \(p_{\theta}(x|z)\)，表示将隐变量解码到数据空间中，即生成新样本；编码器 \(q_{\phi}(z|x)\)，表示将观测数据编码到隐空间中，编码器设置为高斯分布，其参数由神经网络模型给出。VAE的概率模型可以总结为：

\[\begin{align} & Latend \ prior: p(z) = N(z; \boldsymbol{0, I}) \tag{17} \\ & Encoder: q_{\phi}(z|x) = N(z; \boldsymbol{\mu}_{\phi}(x), \sigma_{\phi}^{2}(x)\boldsymbol{I}) \rightarrow q(z|x) \tag{18} \\ & Decoder: p_{\theta}(x|z) \rightarrow p(x|z) \tag{19} \\ \end{align}\]

变分下界(ELBO)

  在估计概率模型的参数时，我们通常会使用“最大似然法”，然而尽管有EM算法等迭代算法，但在面对复杂的概率模型时，求解似然参数仍然是一件困难且耗时的工作，对此VAE提出变分下界的概念，也叫做证据下界(Evidence Lower Bound)，通过最大化变分下界来代替直接对于似然函数的优化。
  (ELBO) 对于存在隐变量 \(z\) 的模型，其似然函数具有以下不等式成立：

\[\begin{align} \log{p(x)} \ge \mathbb{E}_{q_{\phi}(z|x)} \left[ \log\frac{p(x,z)}{q_{\phi}(z|x)} \right] \tag{20} \end{align}\]

\(proof\)

\[\begin{align} \log{p(x)} &= \log{p(x)} \int_{z} q_{\phi}(z|x) dz = \int_{z} q_{\phi}(z|x)\log{p(x)} dz \notag \\ &= \mathbb{E}_{q_{\phi}(z|x)}[\log{p(x)}] = \mathbb{E}_{q_{\phi}(z|x)}\left[ \log{\frac{p(x,z)q_{\phi}(z|x)}{q(z|x)q_{\phi}(z|x)}} \right] \notag \\ &= \mathbb{E}_{q_{\phi}(z|x)} \left[ \log{\frac{p(x,z)}{q_{\phi}(z|x)}} \right] + \mathbb{E}_{q_{\phi}(z|x)} \left[ \log{\frac{q_{\phi}(z|x)}{q(z|x)}} \right] \notag \\ &= \mathbb{E}_{q_{\phi}(z|x)} \left[ \log{\frac{p(x,z)}{q_{\phi}(z|x)}} \right] + D_{KL}(q_{\phi}(z|x) || q(z|x)) \tag{21} \\ & \ge \mathbb{E}_{q_{\phi}(z|x)} \left[ \log{\frac{p(x,z)}{q_{\phi}(z|x)}} \right] \tag{22} \end{align}\]

  从式(21)可以看出，似然函数 \(\log{p(x)}\) 实际上等于ELBO加上真实后验 \(q(z|x)\) 与近似后验 \(q_{\phi}(z|x)\) 之间的 KL Divergence，由于似然函数 \(\log{p(x)}\)并不依赖于参数 \(\phi\)，其在优化过程中可看作一个固定的常数，因此使用 ELBO 代替似然函数作为最大化的目标函数是有意义的，因为其实际上等价于最小化真实后验 \(q(z|x)\) 与近似后验 \(q_{\phi}(z|x)\) 之间的 KL Divergence.   通过(17)式与(19)式，我们可以对 ELBO 进一步推导：

\[\begin{align} \mathbb{E}_{q_{\phi}(z|x)} \left[ \log{\frac{p(x,z)}{q_{\phi}(z|x)}} \right] &= \mathbb{E}_{q_{\phi}(z|x)} \left[ \log{\frac{p(z)p_{\theta}(x|z)}{q_{\phi}(z|x)}} \right] \notag \\ &= \mathbb{E}_{q_{\phi}(z|x)} \left[ \log{p_{\theta}(x|z)} \right] - \mathbb{E}_{q_{\phi}(z|x)} \left[ \log{\frac{q_{\phi}(z|x)}{p(z)}} \right] \notag \\ &= \underbrace{\mathbb{E}_{q_{\phi}(z|x)} \left[ \log{p_{\theta}(x|z)} \right]}_{reconstruction \ term} - \underbrace{D_{KL}(q_{\phi}(z|x) || p(z))}_{prior \ matching \ term} \tag{23} \\ \end{align}\]

  通过推导，我们发现 VAE 的 ELBO 包括两项，其意义如下：

• reconstruction term: 该项代表了解码器基于隐变量 \(z\) 重构输入数据 \(x\) 的分布的对数似然。通过最大化这个对数似然，模型被训练以更好地从潜在变量重构输入数据，确保潜在变量 \(z\) 能够有效地捕捉到输入数据的关键特征。
  
• prior matching term: 该项了学习到的变分分布与对隐变量的先验分布的相似程度。最小化这个项鼓励编码器学习一个实际的分布，而不是崩溃成狄拉克delta函数。

  要最大化 ELBO 就是要使第一项尽可能大，同时第二项尽可能小，但这两个目标是想抗衡的，若第二项很小，则说明 \(q_{\phi}(z|x)\) 近似于标准高斯分布，则意味着 \(z\) 没有任何辨识度，并没有学习到 \(x\) 的分布特征，此时第一项就不可能非常小；反过来，如果第一项较大，即说明 \(z\) 具有很高的辨识度，\(q_{\phi}(z|x)\) 会大幅度偏离标准高斯分布，则第二项不会很小。因此类似于 GAN ，这两项的优化实际上是一个对抗过程。在变分自编码器中，ELBO是优化目标，它同时包括了对数据重构的质量（第一项）和对模型复杂性的控制（第二项）。最大化ELBO意味着我们尝试达到这两个目标的最佳平衡：即最大化数据的重构质量，同时保持潜在表示的多样性和广泛性。

模型训练

  综上所述，VAE的优化函数可以写成：

\[\begin{align} (\hat{\phi},\hat{\theta}) = \argmax_{\phi,\theta} \mathbb{E}_{q_{\phi}(z|x)} \left[ \log{p_{\theta}(x|z)} \right] - D_{KL}(q_{\phi}(z|x) || p(z)) \tag{24} \end{align}\]

  我们来将目标函数进行细化，对于第一项，我们可以使用Monte Carlo estimate 来进行估计，即：

\[\begin{align} \mathbb{E}_{q_{\phi}(z|x)} \left[ \log{p_{\theta}(x|z)} \right] \approx \frac{1}{N}\sum_{i=1}^{N} \log{p_{\theta}(x|z_i)} \tag{25} \end{align}\]

  其中，\(N\) 为训练集的样本容量，\(z_i\) 为从隐变量的后验分布 \(q_{\phi}(z|x)\) 采样出的样本。同时，由于 \(q_{\phi}(z|x)\) 与 \(p(z)\) 均为高斯分布，故第二项很容易就能得到：

\[\begin{align} D_{KL}(q_{\phi}(z|x) || p(z)) = \frac{1}{2}\sum_{k=1}^{d}\left( \mu_{\phi,k}^{2}(x) + \sigma_{\phi,k}^{2}(x) - \log{\sigma_{\phi,k}^{2}(x)}-1 \right) \tag{26} \end{align}\]

  因此(24)式可以写为：

\[\begin{align} \argmax_{\phi,\theta} \frac{1}{N}\sum_{i=1}^{N} \log{p_{\theta}(x|z_i)} - \frac{1}{2}\sum_{k=1}^{d}\left( \mu_{\phi,k}^{2}(x) + \sigma_{\phi,k}^{2}(x) - \log{\sigma_{\phi,k}^{2}(x)}-1 \right) \tag{27} \end{align}\]

  然而，这里有一个问题：随机采样过程的不可导性。在第一项中，我们需要从 \(q_{\phi}(z|x)\) 中采样出 \(z\) 用于计算损失，但在神经网络中，我们需要计算损失函数关于网络参数的梯度以进行反向传播和参数更新。然而，由于采样步骤是一个随机过程，并且具有“跳跃”的性质(从概率分布中直接抽取样本)，它本身不是一个可导的操作，我们无法对\(\phi\) 进行更新。
  为了解决这个问题，VAE提出了一种重要技术：重参数化技巧(reparameterization trick)。重新参数化技巧将随机变量改写为噪声变量的确定性函数；这允许通过梯度下降对非随机项进行优化。具体而言，对于从 \(N(z; \boldsymbol{\mu}_{\phi}(x), \sigma_{\phi}^{2}(x)\boldsymbol{I})\) 采样出的样本 \(z\)，我们可以将其重写为：

\[\begin{align} z = \boldsymbol{\mu}_{\phi}(x) + \sigma_{\phi}^{2}(x) \odot \epsilon,\quad \epsilon \sim N(\epsilon; \boldsymbol{0,I}) \tag{28} \end{align}\]

  其中 \(\boldsymbol{\mu}_{\phi}(x)\) 与 $ _{}^{2}(x)$ 均由神经网络模型计算，通以上处理，我们便可以对参数 \(\phi,\theta\) 进行求导，进而利用SGD等算法进行参数更新。
  对于解码器 \(p_{\theta}(x|z)\) 的分布，在VAE的原始论文给出了两种候选方案：伯努利分布或高斯分布。综上所述，VAE的训练过程可以用下图4表示：

Image4: VAE 训练示意图

采样

  在完成对VAE的训练后，我们便可以使用VAE来进行采样，即生成新的样本 \(x'\)。首先重隐变量的先验分布 \(N(z; \boldsymbol{0,I})\) 中采样出隐变量 \(z\)，再利用解码器 \(p_{\theta}(x|z)\) 将隐变量映射到数据空间得到新样本 \(x'\)。

分层 VAE

  在对VAE的研究中，一些学者提出了其更具有拓展性的框架，即分层变分自变量(Hierarchical Variational Auto-encoders, HVAE)，将 VAE 由单个隐变量拓展到多个隐变量的层次，下图5即为 HVAE 的概率图：

Image5: HVAE 概率图

  通过设置多个隐变量，模型可以逐步学习到更高层次、更抽象的特征。在具有 \(T\) 个隐变量的 HVAE中，我们来考虑一种特殊的情况，即具有马尔可夫性质的 HVAE (MHVAE)，此时，某个隐变量的分布仅取决于与之相邻的隐变量。这样，我们可以对重写概率模型：

\[\begin{align} & p(x, z_{1:T}) = p(z_{T})p_{\theta}(x|z_1)\prod_{t=2}^{T}p_{\theta}(z_{t-1}|z_{t}) \tag{29}\\ & q_{\phi}(z_{1:T}|x) = q_{\phi}(z_{1}|x) \prod_{t=2}^{T}q_{\phi}(z_{t}|z_{t-1}) \tag{30} \\ \end{align}\]

  同时我们可以将 MHVAE 的 ELBO 写为：

\[\begin{align} \log{p(x)} &= \log{\int_{z_{1:T}} p(x, z_{1:T})} dz_{1:T} \notag \\ &= \log{\int_{z_{1:T}} \frac{p(x, z_{1:T})q_{\phi}(z_{1:T}|x)}{q_{\phi}(z_{1:T}|x)}dz_{1:T}} \notag \\ &= \log{\mathbb{E_{q_{\phi}(z_{1:T}|x)}} \left[ \frac{p(x, z_{1:T})}{q_{\phi}(z_{1:T}|x)} \right]} \notag \\ &\ge \mathbb{E_{q_{\phi}(z_{1:T}|x)}} \left[ \log\frac{p(x, z_{1:T})}{q_{\phi}(z_{1:T}|x)} \right] \tag{31} \end{align}\]

  利用(29)式与(30)式，ELBO可以进一步写为：

\[\begin{align} \mathbb{E_{q_{\phi}(z_{1:T}|x)}} \left[ \log\frac{p(x, z_{1:T})}{q_{\phi}(z_{1:T}|x)} \right] = \mathbb{E_{q_{\phi}(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_{\phi}(z_{1}|x) \prod_{t=2}^{T}q_{\phi}(z_{t}|z_{t-1})} \right] \tag{32} \end{align}\]

  想必读者们已经注意到，MHVAE的基本架构和 Diffusion Model 已经非常相似，接下来我们将介绍 Diffusion Model 在 MHVAE 上又做了哪些工作。

References

[1] Paper: Luo C. Understanding diffusion models: A unified perspective[J]. arXiv preprint arXiv:2208.11970, 2022.
[2] Paper: Kingma, D. P. and Welling, M. (2014). Auto-encoding variational bayes. In Proceedings of the International Conference on Learning Representations (ICLR).
[4] Blog: 苏剑林 变分自编码(1-3), 科学空间