从联合分布角度解析 VAE

A VAE

我们先从 VAE 的角度理解生成模型的优化目标如何获得

A.1 背景知识：数值计算vs采样计算

给定已知概率密度函数 $p (x)$ ,那么 $x$ 的期望定义为

\begin{matrix} (0) & E [x] = \int x p (x) d x \end{matrix}

如果求解基于概率密度函数的期望的数据计算，也就是数据积分，可以选择若干有代表性的点 $x_{0} < x_{1} < . . . < x_{n}$ ，得到数值解：

\begin{matrix} (2) & E [x] \approx \sum_{i = 1}^{n} x_{i} p (x_{i}) (x_{i} - x_{i - 1}) \end{matrix}

这里本质就是借助离散化区间计算积分数值。
此外，我们可以利用采样方法计算期望的近似数值解：从 $p (x)$ 中采样若干个点 $x_{1}, x_{2}, . . ., x_{n}$ ，我们有：

\begin{matrix} (3) & E [x] \approx \frac{1}{n} \sum_{i = 1}^{n} x_{i}, x \sim p (x) \end{matrix}

要理解 $E q .2$ 和 $E q .3$ 的区别，重点在于理解概率密度函数的含义：描述样本 x 出现的可能性。
$E q .2$ 中包含了概率的计算而 $E q .3$ 中仅有 $x$ 的计算，这是因为 $E q .3$ 中的采样过程已经包含了概率过程，概率大的 $x_{i}$ 出现的次数也多。
更一般的，我们有

\begin{matrix} (4) & E [f (x)] = \int f (x) p (x) d x \approx \frac{1}{n} \sum_{i = 1}^{n} f (x_{i}), x_{i} \sim p (x_{i}) \end{matrix}

上式为蒙特卡洛模拟的基础。

此外，关于概率期望，我们后面会重点用到如下视角的变化：

\begin{matrix} (5) & K L (p (x) | | q (x)) = \int p (x) \ln \frac{p (x)}{q (x)} d x = E_{x \sim p (x)} [\ln \frac{p (x)}{q (x)}] \end{matrix}

VAE的优化目标

VAE 作为生成模型的经典任务：给定若干数据样本: ${x_{1}, . . ., x_{n}}$ ，整体用 $x$ 描述，我们希望借助隐变量 $z$ 描述 $x$ 的分布 $\tilde{p} (x)$ :

\begin{matrix} (6) & q (x) = \int q (x | z) q (z) d z \end{matrix}

最终我们希望构造的分布 $q (x)$ 能逼近目标数据分布 $\tilde{p} (x)$ ，即两分布的 $K L$ 散度足够小。

由于引入了隐变量 $z$ ，推导原始 $K L$ 散度是复杂的，需要引入各种变分推理以及 ELBO。因此，我们考虑优化两个概率的联合分布距离：

\begin{matrix} (7) & K L (p (x, z) | | q (x, z)) = \iint p (x, z) \ln \frac{p (x, z)}{q (x, z)} d z d x \end{matrix}

对于联合分布，借助贝叶斯公式有

\begin{matrix} (8) & q (x, z) = q (x | z) q (z) \end{matrix}

\begin{matrix} (9) & p (x, z) = \tilde{p} (x) p (z | x) \end{matrix}

注意 $E q .8$ 和 $E q .9$ 中我们分别从不同从面引入了条件概率。具体的， $q (x | z)$ 可视为解码器，而 $p (z | x)$ 则为编码器，因为 $p (x, z) = \tilde{p} (x) p (z | x)$ 来源与真实数据样本 $\tilde{x}$ 。

\begin{matrix} (10) & \begin{aligned} K L (p (x, y) | | q (x, y)) & = \int \tilde{p} (x) [\int p (z | x) \ln \frac{\tilde{p} (x) p (z | x)}{q (x, z)} d z] d x \\ = E_{x \sim \tilde{p} (x)} [\int p (z | x) \ln \tilde{p} (x) + p (z | x) \ln \frac{p (z | x)}{q (x, z)} d z] \end{aligned} \end{matrix}

对于积分的前一项，有

\begin{matrix} (11) & \begin{aligned} E_{x \sim \tilde{p} (x)} [\int p (z | x) \ln \tilde{p} (x) d z] & = E_{x \sim \tilde{p} (x)} [\ln \tilde{p} (x) \int p (z | x) d z] \\ = E_{x \sim \tilde{p} (x)} [\ln \tilde{p} (x] \\ = C \end{aligned} \end{matrix}

为常数。所以我们将优化目标设为

\begin{matrix} (12) & \begin{aligned} L & = E_{x \sim \tilde{p} (x)} [\int p (z | x) \ln \frac{p (z | x)}{q (x, z)} d z] \\ = E_{x \sim \tilde{p} (x)} [\int p (z | x) \ln \frac{p (z | x)}{q (x | z) q (z)} d z] \\ = E_{x \sim \tilde{p} (x)} [- \int p (z | x) \ln q (x | z) d z + \int p (z | x) \ln \frac{p (z | x)}{q (z)} d z] \\ = E_{x \sim \tilde{p} (x)} [E_{z \sim p (z | x)} [- \ln (q (x | z))] + K L (p (z | x) | | q (z))] \end{aligned} \end{matrix}

式 $E q .12$ 中的两项，即分别对应了 VAE 的 MSE 损失和 KL 损失。

具体实现

分析得到了上面的优化目标，其实还不能直接实现的。因为要求解最小化损失 $L$ ，我们是需要其中的 $q (z), q (x | z), p (z | x)$ 的具体解析式的。下面我们逐项分析。
首先为了方便采样，我们假设 $z \sim N (0, 1)$ ，如何 $q (z)$ 有了解析式。
对于另外两项，我们需要在基于一定先验假设的情况下，使用神经网络拟合！
首先，假设编码器 $p (z | x)$ 服从正态分布，其均值和方差由 $x$ 决定，由此，我们构建网络 $μ_{θ} (x)$ , $σ_{θ} (x)$ 分别建模均值和方差，那么就可以获得解析式：

\begin{matrix} (13) & p (z | x) = \frac{1}{\prod_{k = 1}^{d} \sqrt{2 π σ_{θ}^{2} (x)}} \exp (- \frac{1}{2} ∥ \frac{z - μ_{θ} (x)}{σ_{θ} (x)} ∥^{2}) \end{matrix}

这样一来，借助两个正态分布的 KL 散度， $E q .12$ 中的 KL 散度就可以详细计算

\begin{matrix} (14) & K L (p (z | x) ∥ q (z)) = μ_{θ}^{2} (x) + σ_{θ}^{2} (x) - \ln σ_{θ}^{2} (x) - 1 \end{matrix}

现在只剩解码器 $q (x | z)$ ，该选择什么分布呢？。在 VAE 的原始论文中，给出了两种候选方案：伯努利分布和正态分布。这里我们只看后者。

我们将解码器建模成高斯分布，是否过于简单？在生成模型领域，其实我们并没有什么其他选择。因为我们要构造一个分布，而不是任意一个函数，既然是分布就得满足归一化的要求，而要满足归一化，又要容易算，候选分布模型并不多

已知 $q (x | z)$ 为高斯分布，那么我们同样我们用神经网络构建均值和方差 ${\tilde{μ}}_{θ} (z)$ ， ${\tilde{σ}}_{θ} (z)$ 。这里由于我们的优化目标中的两项是同时优化，所以我们统一使用 $θ$ 代表网络参数。注意到解码器输入为 $z$ ，因此我们有解码器的解析式

\begin{matrix} (15) & q (x | z) = \frac{1}{\prod_{k = 1}^{D} \sqrt{2 π {\tilde{σ}}_{θ}^{2} (x)}} \exp (- \frac{1}{2} ∥ \frac{z - {\tilde{μ}}_{θ} (x)}{{\tilde{σ}}_{θ} (z)} ∥^{2}) \end{matrix}

通常情况下，对于生成器的方差我们设置为固定值 ${\tilde{σ}}^{2}$ ,此时

\begin{matrix} (16) & \begin{aligned} - \ln q (x | z) & = \frac{1}{2} ∥ \frac{z - {\tilde{μ}}_{θ} (x)}{{\tilde{σ}}_{θ} (z)} ∥^{2} + \frac{D}{2} \ln 2 π + \frac{1}{2} \sum_{k = 1}^{D} \ln {\tilde{σ}}_{θ}^{2} (z) \\ \sim \frac{1}{2 {\tilde{σ}}_{θ}^{2}} ∥ z - {\tilde{μ}}_{θ} (x) ∥^{2} \end{aligned} \end{matrix}

对整个过程来看，从 $\tilde{p} (x)$ 从随机采样样本 $x$ ，经过编码器，我们获得 $p (z | x)$ 的均值和方差，进而得知 $p (z | x)$ 的具体分布解析式，从中采样得到得到隐变量 $z$ ，并可以计算 KL 损失。然后将 $z$ 经过解码器，获得输出 $\hat{x}$ ，从本质上这代表估计的均值，我们将其作为生成目标，此外根据 $E q .16$ 计算 MSE 损失。

分析

从 $E q .12$ 的两项损失中我们可以得知，VAE 中也存在一种对抗过程。假设 KL loss优化为 0，那么编码器的输出为正态分布，隐变量 $z$ 没有任何辨识度，此时的生成器 $q (x | z)$ 是不可能和样本 $\tilde{x}$ 预测准确。而如果前一项 MSE 损失足够小，此时 $q (x | z)$ 大，预测的准确，此时的网络非常接近自编解码器，此时的隐变量 $z$ 不会太随机，KL损失也不会小。所以 VAE 损失中的两项也存在互相对抗的过程，而且相较于 GAN，是共同优化，不需要交替优化过程。

但 VAE 存在的模糊的缺点，其本质也是因为我们将解码器 $q (x | z)$ 预定义成正态分布，显然正态分布的表征能力是不足的，因此 VAE 从理论层面就不具备足够的生成能力。而 GAN 的生成器不需要满足正态分布的先验。对于 vae 来讲，通过迭代多层解码器（HVAE），既使用多个嵌套正态分布提高表征能力，能够大大提升生成效果。从这个角度来看，VAE 和 DDPM 是同根同源的。