NO.54.tip： 变分自编码模型

背景

潜变量模型的似然函数由下式给出：

$$p(\mathbf{x}\mid \mathbf{w})=\int p(\mathbf{x}\mid \mathbf{z},\mathbf{w})p(\mathbf{z})d\mathbf{z}\text{\hspace{1em}(1)}$$

其中 $p(\mathbf{x}\mid \mathbf{z},\mathbf{w})$ 由一个深度神经网络定义。这个似然函数是难以处理的，因为关于 $\mathbf{z}$ 的积分无法解析计算。变分自编码器（Variational Autoencoder, VAE）(Kingma and Welling, 2013; Rezende, Mohamed, and Wierstra, 2014; Doersch, 2016) 在训练模型时使用了这个似然函数的近似表示来解决这一难题。VAE 背后的关键思想如下：

1. 使用证据下界来近似似然函数，从而与 EM 算法形成紧密联系；
2. 借助摊销推断（即使用编码器网络）来近似 E 步骤中潜变量上的后验分布，而不是精确评估每个数据点的后验分布；
3. 使用重参数化（reparameterization trick）技巧实现编码器模型的可微分训练。

考虑一个生成式模型，其条件分布 $p(\mathbf{x}\mid \mathbf{z},\mathbf{w})$ 覆盖了由深度神经网络输出控制的 $D$ 维数据变量 $\mathbf{x}$。例如，$\mathbf{g}(\mathbf{z},\mathbf{w})$ 可能表示高斯条件分布的均值。同时考虑 $M$ 维潜变量 $\mathbf{z}$ 的分布，它由零均值、单位方差的高斯分布给出：

$$p(\mathbf{z})=\mathcal{N}(\mathbf{z}\mid 0,\mathbf{I})\text{\hspace{1em}(2)}$$

为了推导出 VAE 的近似，首先回想一下，对于由潜变量 $\mathbf{z}$ 描述的空间上的任意概率分布 $q(\mathbf{z})$，以下关系成立：

$$\ln p(\mathbf{x}\mid \mathbf{w})=\mathcal{L}(\mathbf{w})+\mathrm{KL}\left(q(\mathbf{z})\Vert p(\mathbf{z}\mid \mathbf{x},\mathbf{w})\right)\text{\hspace{1em}(3)}$$

其中 $\mathcal{L}$ 是证据下界（evidence lower bound）或 ELBO，也称变分下界（variational lower bound），由下式给出：

$$\mathcal{L}(\mathbf{w})=\int q(\mathbf{z})\ln \left\{\frac{p(\mathbf{x}\mid \mathbf{z},\mathbf{w})p(\mathbf{z})}{q(\mathbf{z})}\right\}d\mathbf{z}\text{\hspace{1em}(4)}$$

Kullback-Leibler 散度（简称 KL 散度）$\mathrm{KL}(q\Vert p)$ 定义如下：

$$\mathrm{KL}\left(q(\mathbf{z})\Vert p(\mathbf{z}\mid \mathbf{x},\mathbf{w})\right)=-\int q(\mathbf{z})\ln \left\{\frac{p(\mathbf{z}\mid \mathbf{x},\mathbf{w})}{q(\mathbf{z})}\right\}d\mathbf{z}\text{\hspace{1em}(5)}$$

由于 KL 散度满足 $\mathrm{KL}(q\Vert p)⩾0$，可以得到

$$\ln p(\mathbf{x}\mid \mathbf{w})⩾\mathcal{L}\text{\hspace{1em}(6)}$$

所以 $\mathcal{L}$ 是 $\ln p(\mathbf{x}\mid \mathbf{w})$ 的证据下界。虽然对数似然函数 $\ln p(\mathbf{x}\mid \mathbf{w})$ 难以直接计算，但我们可以使用蒙特卡洛估计方法来评估这个下界，从而使它近似于真实对数似然函数的值。

考虑从模型分布 $p(\mathbf{x})$ 中独立采样的训练数据集 $\mathcal{D}=\{{\mathbf{x}}_1,\dots ,{\mathbf{x}}_N\}$。这个数据集的对数似然函数由下式给出：

$$\ln p(\mathcal{D}\mid \mathbf{w})=\sum _{n=1}^N{\mathcal{L}}_n+\sum _{n=1}^N\mathrm{KL}\left(q_n({\mathbf{z}}_n)\Vert p({\mathbf{z}}_n\mid {\mathbf{x}}_n,\mathbf{w})\right)\text{\hspace{1em}(7)}$$

请注意，这引入了与每个数据向量 ${\mathbf{x}}_n$ 相对应的单独潜变量 ${\mathbf{z}}_n$，就像我们在混合模型和概率 PCA 模型中看到的那样。因此，每个潜变量都有自己的独立分布 $q_n({\mathbf{z}}_n)$，且都可以单独优化。

由于式 (7) 对任意选取的分布 $q_n(\mathbf{z})$ 都成立，我们可以选择能使下界 ${\mathcal{L}}_n$ 最大化的分布，或等价地选择最小化 KL 散度 $\mathrm{KL}(q_n({\mathbf{z}}_n)\Vert p({\mathbf{z}}_n\mid {\mathbf{x}}_n,\mathbf{w}))$ 的分布。在先前讨论的简单高斯混合模型和概率 PCA 模型中，我们可以在 EM 算法的 E 步骤中精确计算这些后验分布，即将每个 $q_n({\mathbf{z}}_n)$ 设置为相应的后验分布 $p({\mathbf{z}}_n\mid {\mathbf{x}}_n,\mathbf{w})$。此时 KL 散度降为 0，因此下界等于真实的对数似然。

${\mathbf{z}}_n$ 的精确后验分布由贝叶斯定理给出：

$$p({\mathbf{z}}_n\mid {\mathbf{x}}_n,\mathbf{w})=\frac{p({\mathbf{x}}_n\mid {\mathbf{z}}_n,\mathbf{w})p({\mathbf{z}}_n)}{p({\mathbf{x}}_n\mid \mathbf{w})}\text{\hspace{1em}(8)}$$

对于深度生成式模型，式 (8) 中的分子是可以直接计算的。然而分母是由似然函数给出的，难以解析计算。因此，我们需要找到一个对后验分布的近似。原则上，我们可以考虑为每个分布 $q_n({\mathbf{z}}_n)$ 单独建立参数化模型，并进行数值优化，但计算成本非常高，特别是对于大数据集来说。此外还需要在每次更新 $\mathbf{w}$ 后重新计算分布。因此，我们决定换一个方向，引入第二个神经网络来探索一个更有效的近似框架。

源与流

摊销推理

在变分自编码器中，我们不再单独为每个数据点 ${\mathbf{x}}_n$ 计算一个后验分布 $p({\mathbf{z}}_n\mid {\mathbf{x}}_n,\mathbf{w})$，而是通过训练一个称为编码器网络的神经网络来近似所有这些分布。这种技术称为摊销推理（amortized inference），要求编码器生成单一的分布 $q(\mathbf{z}\mid \mathbf{x},\varphi )$，该分布以 $\mathbf{x}$ 为条件，其中 $\varphi$ 代表网络参数。由证据下界给出的目标函数现在同时依赖于 $\varphi$ 和 $\mathbf{w}$，我们可以使用基于梯度的优化方法来联合优化这两组参数。

因此，VAE 包含两个参数独立但能联合训练的神经网络：一个编码器网络，用于将数据向量映射到潜空间；一个解码器网络，用于将潜空间向量映射回数据空间。这类似于简单的神经网络自编码器模型，所不同的是，我们在潜空间中定义了一个概率分布。我们将看到，编码器根据贝叶斯定理计算解码器的近似概率逆过程——即根据给定观测数据推断潜变量的分布。

在变分自编码器中，编码器的一个典型选择是具有对角协方差矩阵的高斯分布：

$$q(\mathbf{z}\mid \mathbf{x},\varphi )=\prod _{j=1}^M\mathcal{N}\left(z_j\mid {\mu }_j(\mathbf{x},\varphi ),{\sigma }_j^2(\mathbf{x},\varphi )\right)\text{\hspace{1em}(10)}$$

其均值 ${\mu }_j$ 和方差 ${\sigma }_j^2$ 由一个以 $\mathbf{x}$ 作为输入的神经网络给出。

请注意，均值 ${\mu }_j(\mathbf{x},\varphi )$ 位于区间 $(-\infty ,\infty )$，因此相关的输出单元通常使用线性激活函数；而方差 ${\sigma }_j^2(\mathbf{x},\varphi )$ 必须是非负的，所以相关的输出单元通常使用指数函数 $\exp (\cdot )$ 作为激活函数。

我们的优化目标是使用基于梯度的优化方法，针对参数 $\varphi$ 和 $\mathbf{w}$ 最大化这个下界，这通常使用基于小批量数据的随机梯度下降法来实现。尽管是联合优化参数，但从概念上我们可以交替地优化 $\varphi$ 和 $\mathbf{w}$。

（ELBO）优化与 EM 算法的一个关键区别是，对于给定的 $\mathbf{w}$ 值，相对于编码器参数 $\varphi$ 的优化通常不会将 KL 散度降低到零，因为编码器网络并不能完美地预测后验潜在分布。所以下界和真实对数似然之间存在残差。尽管基于深度神经网络的编码器非常灵活，但它并不能精确地模拟真实的后验分布，原因如下：

1. 真实的条件后验分布未必服从因子分解的高斯分布；
2. 即使是一个大型神经网络，也存在灵活性方面的限制；
3. 训练过程只是一个近似的优化。

重参数化技巧

遗憾的是，到目前为止，下界［式 (9)］仍然难以直接计算，因为在潜变量 ${\mathbf{z}}_n$ 的积分过程中，解码器网络会导致被积函数对潜变量存在复杂的依赖关系。对于数据点 ${\mathbf{x}}_n$，其对下界的贡献可以表示为：

$${\mathcal{L}}_n(\mathbf{w},\varphi )=\int q({\mathbf{z}}_n\mid {\mathbf{x}}_n,\varphi )\ln \left\{\frac{p({\mathbf{x}}_n\mid {\mathbf{z}}_n,\mathbf{w})p({\mathbf{z}}_n)}{q({\mathbf{z}}_n\mid {\mathbf{x}}_n,\varphi )}\right\}d{\mathbf{z}}_n=\int q({\mathbf{z}}_n\mid {\mathbf{x}}_n,\varphi )\ln p({\mathbf{x}}_n\mid {\mathbf{z}}_n,\mathbf{w})d{\mathbf{z}}_n-\mathrm{KL}\left(q({\mathbf{z}}_n\mid {\mathbf{x}}_n,\varphi )\Vert p({\mathbf{z}}_n)\right)\text{\hspace{1em}(11)}$$

式 (11) 右侧的第二项是两个高斯分布之间的 KL 散度，并且可以解析评估（见习题 2.27）：

$$\mathrm{KL}\left(q({\mathbf{z}}_n\mid {\mathbf{x}}_n,\varphi )\Vert p({\mathbf{z}}_n)\right)=\frac{1}{2}\sum _{j=1}^M\left\{1+\ln {\sigma }_j^2({\mathbf{x}}_n)-{\mu }_j^2({\mathbf{x}}_n)-{\sigma }_j^2({\mathbf{x}}_n)\right\}\text{\hspace{1em}(12)}$$

对于式 (11)，我们可以尝试用一个简单的蒙特卡洛估计来近似潜变量 ${\mathbf{z}}_n$ 上的积分：

$$\int q({\mathbf{z}}_n\mid {\mathbf{x}}_n,\varphi )\ln p({\mathbf{x}}_n\mid {\mathbf{z}}_n,\mathbf{w})d{\mathbf{z}}_n\approx \frac{1}{L}\sum _{l=1}^L\ln p\left({\mathbf{x}}_n\mid {\mathbf{z}}_n^{(l)},\mathbf{w}\right)\text{\hspace{1em}(13)}$$

其中 ${\mathbf{z}}_n^{(l)}$ 是从编码器分布 $q({\mathbf{z}}_n\mid {\mathbf{x}}_n,\varphi )$ 中采样。关于 $\mathbf{w}$ 求梯度很容易，但关于 $\varphi$ 求梯度却存在问题，因为参数 $\varphi$ 的变化也会改变采样分布，所以我们没有办法获取这些样本相对于 $\varphi$ 的导数。从概念上讲，我们可以认为将潜变量 $\mathbf{z}$ 固定到特定样本值的过程阻止了误差信号向编码器网络的反向传播。

我们可以使用重参数化技巧来解决这个问题。重新设计蒙特卡洛采样过程，以便明确计算相对于 $\varphi$ 的导数。首先请注意，如果 $\epsilon$ 是一个均值为零且方差为单位方差的高斯随机变量，那么：

$$z=\sigma \epsilon +\mu \text{\hspace{1em}(14)}$$

将具有均值为 $\mu$、方差为 ${\sigma }^2$ 的高斯分布的样本。

其中 $\mu$ 和 $\sigma$ 由编码器网络的输出 ${\mu }_j({\mathbf{x}}_n,\varphi )$ 和 ${\sigma }_j^2({\mathbf{x}}_n,\varphi )$ 确定，分别对应于式 (10) 中的各分量。我们不直接从潜变量 ${\mathbf{z}}_n$ 采样，而是对 $\epsilon$ 采样，并使用式 (14) 计算对应的 $z_n$ 样本：

$$z_{nj}^{(l)}={\mu }_j({\mathbf{x}}_n,\varphi )+{\sigma }_j^{(l)}({\mathbf{x}}_n,\varphi ){\epsilon }_j^{(l)}\text{\hspace{1em}(15)}$$

其中 $l=1,\dots ,L$ 是样本的索引。该方法建立 $\varphi$ 与采样值之间的依赖关系，使得计算可行。

重参数化技巧可以扩展到其他分布，但仅限于连续变量上的分布。有的估计器可以直接计算梯度而无须使用重参数化技巧（Williams, 1992），但这些估计器的方差很大，因此重参数化技巧也可以看作一种有效的方差降低技术。

使用我们特定的建模假设，VAE 的完整损失函数变为

$$\mathcal{L}=\sum _n\left\{\frac{1}{2}\sum _{j=1}^M\left\{1+\ln {\sigma }_{nj}^2-{\mu }_{nj}^2-{\sigma }_{nj}^2\right\}+\frac{1}{L}\sum _{l=1}^L\ln p\left({\mathbf{x}}_n\mid {\mathbf{z}}_n^{(l)},\mathbf{w}\right)\right\}\text{\hspace{1em}(16)}$$

其中 $z_{nj}^{(l)}={\mu }_{nj}+{\sigma }_{nj}{\epsilon }_{nj}^{(l)}$，${\mu }_{nj}={\mu }_j({\mathbf{x}}_n,\varphi )$，${\sigma }_{nj}={\sigma }_j({\mathbf{x}}_n,\varphi )$。在式 (16) 中，求和操作遍历当前小批量中的所有数据点。每个数据点 ${\mathbf{x}}_n$ 的样本数 $L$ 通常设置为 1，即仅使用单个样本进行估计。尽管这种设置会导致对证据下界的计算结果存在噪声，但由于随机梯度优化过程本身就具有噪声，该策略在整体上仍能实现更高效的参数优化。

我们可以将 VAE 训练总结如下。对于小批量中的每个数据点，通过编码器网络前向传播来计算近似潜变量分布的均值和方差，并使用重参数化技巧从该分布中采样。然后将这些样本通过解码器网络传播计算证据下界［式 (16)］。最后使用自动微分计算相对于 $\mathbf{w}$ 和 $\varphi$ 的梯度。算法 1 对 VAE 训练做了总结。为了清晰起见，我们省略了使用小批量数据点这一细节。

数据集将通过从先验分布 $p(\mathbf{z})$ 中采样并通过解码器网络前向传播来生成数据空间中的样本。

算法 1: VAE 训练
─────────────────────────────────────────────
Input:  Training data set D = {x₁, …, x_N}
        Encoder network {μ_j(x_n, φ), σ_j²(x_n, φ)},  j ∈ {1, …, M}
        Decoder network q(z, w)        
        Initial weight vectors w, φ
        Learning rate η
Output: Final weight vectors w, φ

repeat
    ℒ ← 0
    for j ∈ {1, …, M} do
        ε_{nj} ~ N(0, 1)
        z_{nj} ← μ_j(x_n, φ) + σ_j²(x_n, φ) · ε_{nj}   
        ℒ ← ℒ + ½ { 1 + ln σ_{nj}² − μ_{nj}² − σ_{nj}² }
    end for
    ℒ ← ℒ + ln p(x_n | z_n, w)
    w ← w + η ∇_w ℒ      // 更新解码器权重
    φ ← φ + η ∇_φ ℒ      // 更新编码器权重
until converged
return w, φ

训练成功后，我们想要评估模型表示新测试点 $\hat{\mathbf{x}}$ 的效果如何。由于似然是不可计算的，我们可以使用证据下界 $\mathcal{L}$ 作为近似。对此，我们可以从 $q(\mathbf{z}\mid \hat{\mathbf{x}},\varphi )$ 中采样。

VAE 有许多变种。当应用于图像数据时，编码器通常基于卷积操作，而解码器则基于转置卷积（有时也称反卷积）操作。在条件 VAE 中，编码器和解码器都将一个条件变量 $\mathbf{c}$ 作为额外的输入。例如，我们可能想要生成物体的图像，此时 $\mathbf{c}$ 代表物体的类别。潜空间的先验分布 $p(\mathbf{z}\mid \mathbf{c})$ 既可以是简单的高斯分布，也可以扩展成另一个神经网络给出的条件分布 $p(\mathbf{z}\mid \mathbf{c})$，训练与测试过程与之前类似。

注意，式 (11) 中的第一项鼓励编码器分布 $q(\mathbf{z}\mid \mathbf{x},\varphi )$ 接近先验分布 $p(\mathbf{z})$，从而使解码器在生成模式下，通过从 $p(\mathbf{z})$ 采样，能够产生逼真的输出。在训练 VAE 时，这意味着变分分布可能会出现一个问题，即变分分布 $q(\mathbf{z}\mid \mathbf{x},\varphi )$ 收敛到先验分布 $p(\mathbf{z})$，即发生后验坍塌（posterior collapse）。实际上，潜变量在编码时被忽略了，因为它对输入 $\mathbf{x}$ 没有区分能力，我们对输入 $\mathbf{x}$ 编码然后解码，我们将得到一个看起来模糊不清的重建结果。在这种情况下，KL 散度很小。

另一种问题发生在潜在编码没有被压缩的情况下，具体表现是，即使能得到高度准确的重建结果，但从 $p(\mathbf{z})$ 采样并通过解码器网络生成的输出质量却很差，不像训练数据。在这种情况下，KL 散度相对较大，并且因为训练系统获得了一个与先验分布完全不同的变分分布，从先验分布中采样的样本不会生成逼真的输出。

这两个问题都可以通过式 (11) 的第一项引入一个系数 $\beta$ 来解决（Higgins et al., 2017）。如果重建结果很差，则可以增大 $\beta$；而如果生成的样本很差，通常 $\beta > 1$，则可以减小 $\beta$。$\beta$ 的值也可以根据退火调度来设置，即它从一个小的值开始，并在训练期间逐渐增大。

最后请注意，虽然我们在本章中只讨论了高斯分布，但我们可以扩展 VAE，让它表示其他更复杂的参数分布，包括表示高斯分布方差的输出，或进一步让它表示其他分布。