理解变分自编码器（VAE）

转载翻译自：https://towardsdatascience.com/understanding-variational-autoencoders-vaes-f70510919f73

介绍

在过去几年中，基于深度学习的生成模型因为在该领域中取得了一些惊人的进展而越来越受到关注。依靠大量的数据、精心设计的网络架构和智能训练技术，深度生成模型展现出了令人难以置信的能力，能够生成高度逼真的各种内容，如图像、文本和声音。在这些深度生成模型中，有两个主要的方向引人注目，值得特别关注：生成式对抗网络（GANs, Generative Adversarial Networks）和变分自编码器（VAEs, Variational Autoencoders）。

VAEs是一种自编码器（autoencoder），其编码分布在训练过程中受到正则化，以确保其潜在空间（latent space）中具有良好的属性，使我们能够生成一些新的数据。此外，“变分”一词来自于正则化和统计学中的变分推断方法之间的密切关系。

如果最后两句话很好地总结了VAEs的概念，那么它也可能会引发一些思考，比如什么是自编码器？什么是潜在空间，为什么要对其进行正则化？如何从VAEs生成新数据？VAEs和变分推断之间有什么联系？为了尽可能详细地描述VAEs，我们将尝试回答所有这些问题（以及许多其他问题！），并为读者提供尽可能多的见解（从基本直觉到更高级的数学细节）。因此，本文的目的不仅是讨论变分自编码器所依赖的基本概念，还要从一开始逐步构建推导这些概念的思路。 废话不多说，让我们一起（重新）探索VAEs吧！

大纲

在第一部分中，我们将回顾一些关于降维和自编码器的重要概念，这些概念对于理解变分自编码器（VAEs）非常有用。然后，在第二部分中，我们将展示为什么自编码器不能用于生成新数据，并介绍变分自编码器，它是自编码器的正则化版本，并且使生成过程（generative process）成为可能。在最后一部分中，我们将基于变分推断对VAEs进行更数学化的介绍。

降维（Dimensionality reduction）、主成分分析（PCA）和自编码器（autoencoders）

在这个第一部分中，我们将首先讨论与降维相关的一些概念。具体来说，我们将简要回顾主成分分析（PCA）和自编码器，展示这两个想法如何相互关联。

1. 什么是降维

在机器学习中，降维是减少描述数据的特征数量的过程。这种减少可以通过选择（只保留一些现有特征）或提取（基于旧特征创建少量新特征）来实现，这在许多需要低维数据的情况下非常有用（数据可视化、数据存储、海量计算等）。虽然存在许多不同的降维方法，但我们可以建立一个全局框架，大多这些方法都符合。

首先，让我们称编码器为从“旧特征”表示生成“新特征”表示的过程（通过选择或提取），解码器则是相反的过程。然后，可以将降维解释为数据压缩，其中编码器将数据压缩（从初始空间到编码空间，也称为潜在空间），而解码器将其解压。当然，根据初始数据分布、潜在空间维度和编码器定义，这种压缩可能是有损的，意味着在编码过程中丢失了部分信息，无法在解码时恢复。

$$(e^{\ast },d^{\ast })=\underset{(e,d)\in E\times D}{\mathrm{argmin}}ϵ(x,d(e(x)))$$

其中 $ϵ(x,d(e(x)))$ 定义了定义了输入数据 $x$ 和编解码数据 $d(e(x))$ 之间的重构误差度量（Reconstruction error measure）。最后请注意，接下来我们将用 N 表示数据的数量，用 n_d 表示初始空间和解码空间的维度，用 n_e 表示降维空间和编码空间的维度。

2. 主成分分析（PCA，Principal components analysis）

提到降维方法时，我们首先想到的是主成分分析（PCA）。为了展示它如何适应我们刚刚描述的框架，并与自编码器建立联系，我们先简要介绍一下PCA的工作原理。 PCA的思想是构建 n_e 个新的独立特征，它们是 n_d 个旧特征的线性组合，使得数据在由这些新特征定义的子空间上的投影尽可能接近初始数据（以欧几里德距离衡量）。换句话说，PCA在初始空间中寻找最佳线性子空间（由新特征的正交基描述），使得用它们在该子空间上的投影来近似数据的误差尽可能小。

在我们的全局框架中，我们正在寻找一个编码器，它属于 **n_e * n_d ** 矩阵（线性变换）的 E 族，并且其行是正交的（特征独立），以及与之相关的解码器，它属于 n_d * n_e 矩阵的 D 族。可以证明，与协方差特征矩阵的 n_e 个最大特征值对应的酉特征向量是正交的（或者可以选择为正交的），并且定义了在最小近似误差下将数据投影到的最佳 n_e 维子空间。因此，这些 n_e 个特征向量可以被选择为我们的新特征，因此，降维问题可以被表达为一个特征值/特征向量问题。此外，还可以证明，在这种情况下，解码器矩阵是编码器矩阵的转置。

3. 自编码器（Autoencoders）

现在让我们讨论一下自编码器，并看看如何使用神经网络进行降维。自编码器的基本思想非常简单，就是将编码器和解码器构建为神经网络，并使用迭代优化过程来学习最佳的编码-解码方案。因此，在每次迭代中，我们在自编码器架构（编码器+解码器）中输入一些数据，将编码-解码过程输出与初始数据进行比较，计算误差并通过反向传播来更新网络的权重。

因此，直观地看，整体的自编码器架构（编码器+解码器）为数据创建了一个瓶颈，就是只有部分主要的结构化信息可以通过并进行重建。从总体框架来看，被考虑的编码器集合 E 由编码器网络架构定义，被考虑的解码器集合 D 由解码器网络架构定义，并通过梯度下降在这些网络的参数上进行搜索，以最小化重建误差。

首先假设我们的编码器和解码器架构只有一层并且不是非线性结构（线性自编码器）。这样的编码器和解码器只是简单的线性变换，可以用矩阵表示。在这种情况下，我们可以发现这与主成分分析（PCA）有一定的相似性，我们正在寻找最佳的线性子空间，以尽可能少的信息损失来得到降维后的数据。通过PCA获得的编码和解码矩阵定义了我们通过梯度下降能够得到的其中一个满意的解，但我们应该意识到这并不是唯一的解。实际上，我们可以选择几个基来描述相同的最佳子空间，因此可以有几对编码器/解码器结构能够给出相似的重构误差。此外，对于线性自编码器，与PCA相反，我们得到的新特征之间不一定是独立的（神经网络中没有正交性约束）。

现在，让我们假设编码器和解码器都是深层神经网络结构且是非线性的。在这种情况下，架构越复杂，自编码器能够在保持低重构损失的同时还可以进行更多维度的降维。直观地说，如果我们的编码器和解码器有足够多的自由度，我们可以将任意初始维度降到1。实际上，具有“无限能力”的编码器在理论上可以将N个初始数据点编码为在某一个实轴上的N个数（ E(1)、E(2)、E(3)、…、E(N) ），相关的解码器可以进行相反的转换，并且过程中没有任何损失。

不过我们应该记住两件事。首先，没有重构损失的降维通常会付出代价：潜在空间中缺乏可解释和可利用的结构（缺乏规律性）。其次，在大多数情况下，降维的最终目的不仅仅是减少数据的维度，而是在保持数据结构信息的同时减少维度。出于这两个原因，潜在空间的维度和自编码器的“深度”（定义了压缩的程度和质量）必须根据降维的最终目的进行精准控制和调整。

变分自编码器（Variational Autoencoders）

到目前为止，我们已经讨论了降维问题，并介绍了可以通过梯度下降进行训练的编码器-解码器架构——自编码器。现在让我们将其与内容生成联系起来，看看自编码器在这个问题上的局限性，并介绍变分自编码器。

1. 自动编码器在内容生成方面的局限性

现在，一个自然而然的问题是：自动编码器和内容生成（Content generation）之间有什么联系？ 一旦自动编码器训练完成，我们有了编码器和解码器，但仍然没有真正的方法来生成任何新内容。乍一看，我们可能会认为，如果潜在空间（Latent space）足够规则（在训练过程中由编码器“组织”得够好），我们就可以随机选择潜在空间中的一个数据点进行解码，以获得新的内容。解码器将会像生成式对抗网络的生成器一样发挥作用。

然而掌握自动编码器生成的潜在空间的规则是一个困难的问题，它取决于数据在初始空间中的分布、潜在空间的维度和编码器的架构。因此，我们很难事先确保编码器会按照我们需求的生成过程那样构造潜在空间。

为了说明这一点，让我们考虑之前提到的例子，其中我们描述了一个足够强大的编码器和解码器，可以将任何N个初始训练数据放置在实轴上（每个数据点被编码为实值），并在解码时没有任何重构损失。在这种情况下，自动编码器的高自由度使得编码和解码可以在没有信息损失的情况下进行（尽管潜在空间的维度较低），但这也会导致严重的过拟合问题，即一些潜在空间的点在解码后会产生无意义的内容。

实际上，自动编码器所训练的任务并没有强制要求获得上述那样的结构，自动编码器的训练仅仅是为了尽可能地进行编码和解码，不管潜在空间的组织方式如何。因此，如果我们在定义自编码器架构时没有多做考虑，自编码器在训练过程中自然会尽可能地降低重建误差，使得神经网络最大化拟合了训练数据，容易导致过拟合的情况，除非我们明确地对其进行正则化！

2. 变分自编码器的定义

所以，为了能够使用我们自编码器的解码器进行内容生成，我们必须确保潜在空间足够规则。在训练过程中引入显式正则化是获得这种规则性的一种可能的解决方案。因此，正如我们在本文介绍中简要提到的那样，变分自编码器可以被定义为一种经过正则化的自编码器，以避免过拟合并确保潜在空间具有良好的属性，从而实现生成过程。

和标准自编码器一样，变分自编码器是由编码器和解码器组成的架构，并且训练目标是最小化编码-解码数据与初始数据之间的重构误差。然而，为了引入潜在空间的一些正则化，我们对编码-解码过程进行了部分修改：我们将输入数据编码为在潜在空间上的分布，而不是单个数据点。模型的训练如下进行：

• 首先，将输入数据编码为潜在空间上的分布。
• 其次，从该分布中随机采样一个点。
• 然后，对采样点进行解码，并计算重构误差。
• 最后，将重构误差通过网络进行反向传播。

实际上，编码分布被选择为正态分布，这样编码器就可以被训练成描述这些高斯分布的均值和协方差矩阵的结构。将输入数据编码为具有一定方差的分布而不是单个点的原因是，这样可以很自然地表达潜在空间的正则化，编码器返回的分布被强制接近标准正态分布。我们将在下一小节中看到，通过这种方式，我们确保了潜在空间的局部和全局正则化（局部是因为方差控制，全局是因为均值控制）。

因此，训练VAE（Variational Autoencoders）时要最小化的损失函数由一个“重构项”（最后一层）和一个“正则化项”（潜在层）组成。重构项旨在使编码-解码方案能够正常运行，而正则化项旨在通过使编码器返回的分布接近标准正态分布来规范潜在空间的结构。这个正则化项被表达为返回分布与标准高斯分布之间的Kullback-Leibler散度，这将在下一节中进一步解释其合理性。我们可以注意到，两个高斯分布之间的 KL散度 (Kullback-Leibler divergenc) 具有一个可以直接用均值和协方差矩阵表示的闭式形式（有限次运算形成的解析表达式）。

Kullback-Leibler散度（Kullback-Leibler divergence），又叫相对熵（relative entropy），是用来度量两个概率分布之间差异的一种方法。一般地，设P和Q是两个概率分布，我们希望知道用Q来拟合P时，产生的信息损失是多少，这时就可以使用Kullback-Leibler散度来衡量这种差异。公式如下所示：

$D_{KL}(P||Q)={\sum }_{i}P(i)\log \frac{P(i)}{Q(i)}$

其中，P 和 Q 是两个概率分布，i是任意一个事件。这个公式的含义是，如果我们用 Q 来拟合 P，那么产生的信息损失就是 $D_{KL}(P||Q)$。也就是说，P 和 Q 之间的差异越大，$D_{KL}(P||Q)$ 就越大，二者之间的差异就越明显。

需要注意的是，Kullback-Leibler散度不具有对称性。也就是说，$D_{KL}(P||Q)\ne D_{KL}(Q||P)$，因为在这个公式中，我们用 Q 来拟合 P 时得到的结果与用 P 来拟合 Q 时得到的结果是不同的。

除了Kullback-Leibler散度之外，还有其他一些用于量化概率分布差异的方法，比如Earth Mover’s Distance、Jensen-Shannon散度等。在不同的场合下，我们可以根据具体的需求选择不同的方法来量化概率分布的差异。

3. 关于正则化（regularisation）的直觉

为了使生成过程成为可能，潜在空间中所期望的规则性可以表现为两个主要属性：连续性（潜在空间中的两个接近的点在解码后不应该给出完全不同的内容）和完整性（潜在空间中采样的点在解码后应该给出有意义的内容）。

VAEs将输入数据编码为分布而不是简单的数据点并不足以确保潜在空间的连续性和完整性。如果没有明确定义的正则化项，为了最小化重建误差，模型甚至可以“忽略”需要返回分布这一条件，并且表现的像经典的自动编码器一样（导致过拟合）。为了达到这一点，编码器只需要返回方差很小的分布（趋向于点分布），或者返回均值非常不同的分布（在潜在空间中相互远离）。在这两种情况下，分布的方式都是不理想的，并且连续性和/或完整性无法满足。

因此，为了避免这些效果，我们必须对编码器返回的分布的协方差矩阵和均值进行正则化。在实践中，通过强制分布接近标准正态分布（居中且缩小）来实现这种正则化。这样，我们要求协方差矩阵接近单位矩阵，防止点分布，并且要求均值接近0，防止编码分布相互之间相距太远。

通过这个正则化项，我们可以防止模型将数据编码后在潜在空间中相距较远，并尽可能让返回的数据分布“重叠”，从而满足预期的连续性和完整性条件。当然，这和其他任何正则化项一样，也会在训练数据时有更高的的重构误差。然而，重构误差和KL散度之间的权衡（tradeoff）可以进行调整，我们将在下一节中看到表达式是如何自然地推导出来。

我们可以发现通过正则化获得的连续性和完整性倾向于在潜在空间中创建一个“梯度”。例如，潜在空间中的一个点如果恰好位于来自不同训练数据的两个编码分布的均值之间，那么它应该被解码成介于两个分布之间的某个结果，因为自编码器在两个分布下都可能对其进行采样。

VAE 中的数学细节

在之前的内容中，我们给出了以下直观的概述：VAE是自动编码器，将输入数据编码为分布而不是点，并通过约束编码器返回的分布接近标准高斯分布来规范其潜在空间的结构。在本节中，我们将给出更加数学化的 VAE 解释，这将使我们能够更严格地证明正则化项的合理性。为此，我们将建立一个清晰的概率框架，并使用特定的变分推断技术（variational inference technique）。

1. 概率框架和假设 （Probabilistic framework and assumptions）

让我们先定义一个概率图模型来描述我们的数据。我们用 $x$ 表示代表我们的数据的变量，并假设 $x$ 是从一个未直接观察到的潜在变量 $z$（编码表示）生成的。因此，对于每个数据点，假设以下两个步骤：

• 首先，从先验分布 $p(z)$ 中采样一个潜在表示变量 $z$。
• 其次，从条件似然分布 $p(x|z)$ 中采样数据 $x$。

有了这样一个概率模型，我们可以重新定义编码器和解码器的概念。与简单的自编码器不同，它考虑了确定性的编码器和解码器，我们现在要考虑这两个对象的概率版本。"概率解码器"自然地由 $p(x|z) $ 定义，它描述了给定编码变量的解码分布，而"概率编码器"由 $p(z|x) $ 定义，它描述了给定解码变量的编码分布。

在这一点上，我们可以注意到，在简单的自编码器中缺乏的潜在空间的正则化表示出现在数据生成过程的定义中：潜在空间中的编码表示 $z$ 被假设遵循先验分布 $p(z)$ 。通过贝叶斯定理，我们可以将先验分布 $p(z)$、似然函数 $p(x|z)$ 和后验分布 $p(z|x)$ 联系起来：
$$p(z|x)=\frac{P(x|z)p(z)}{p(x)}=\frac{p(x|z)p(z)}{\int p(x|u)p(u)du}$$
现在假设 $p(z)$ 是一个标准的高斯分布，而 $p(x|z)$ 是一个高斯分布，其均值由变量 $z$ 的确定性函数 $f$ 定义，协方差矩阵的形式是一个正数 $c$ 乘以单位矩阵 $I$。函数 $f$ 被假设属于一个被标记为 F 的函数族，暂时没有具体指定，稍后会选择。因此，我们有：
$$\begin{array}{rlrlrl}p(z)& \equiv N(0,I)& & & & \\ P(x|z)& \equiv N(f(z),cI)& & f\in F& & c>0\end{array}$$
假设 $f$ 是明确定义且固定的。理论上，我们已知 $p(z)$（潜在空间分布） 和 $p(x|z)$（解码器），可以使用贝叶斯定理来计算 $p(z|x)$。然而因为分母上的积分较难处理，所以需要使用近似技术，比如变分推断。

需要注意的是，在这里我们所提到的 $p(z)$ 和 $p(x|z)$ 都是高斯分布。所以，如果我们有 $E(x|z)=f(z)=z$，那么它意味着 $p(z|x) $ 也应该遵循高斯分布，在理论上，我们可以尝试用 $p(z)$ 和 $p(x|z)$ 的均值和协方差矩阵来表示 $p(z|x)$ 的均值和协方差矩阵。然而，在实践中，这个条件并不满足，我们需要使用一种近似技术，比如变分推断，使得这种方法更加通用，并且模型更具鲁棒性。

2. 变分推断公示（Variational inference formulation）

在统计学中，变分推断（VI，Variational inference）是一种近似复杂分布的技术。其思想是设定一个参数化的分布族（例如高斯分布族，其参数为均值和协方差），并在该族中寻找最佳近似目标分布的元素，该最佳元素则是对描述该族的参数使用梯度下降，通过最小化给定的近似误差度量（大多数情况下是近似分布和目标分布之间的Kullback-Leibler散度）来得到的。更多细节可以参考另一篇关于变分推断的文章（后续也将翻译并转载）。

在这里，我们将通过一个高斯分布 $q_x(z)$ 来近似 $p(z|x)$，其均值和协方差由参数 $x$ 的两个函数 $g$ 和 $h$ 定义。这两个函数分别被假设属于函数族 G 和 H，这些函数族将在稍后进行具体说明，但假设它们是参数化的。因此，我们可以表示为：
$$\begin{array}{rlrlr}{q}_x(z)\equiv N(g(x),h(x))& & g\in G& & h\in H\end{array}$$
所以，我们已经定义了一组候选变分推断方法，并且现在需要通过优化函数 $g$ 和 $h$（实际上是它们的参数）来找到这组方法中最佳的近似，以最小化近似与目标分布 $p(z|x)$ 之间的Kullback-Leibler散度。换句话说，我们正在寻找最优的 $g$ 和 $h$，使得：
$$\begin{array}{rl}(g^{\ast },h^{\ast })& =\underset{(g,h)\in G\times H}{\mathrm{argmin}}KL(q_x(z),p(z|x))\\ & =\underset{(g,h)\in G\times H}{\mathrm{argmin}}(E_{z\sim q_x}(\log q_x(z))-E_{z\sim q_x}(\log \frac{p(x|z)p(z)}{p(x)}))\\ & =\underset{(g,h)\in G\times H}{\mathrm{argmin}}(E_{z\sim q_x}(\log q_x(z))-E_{z\sim q_x}(\log p(z))-E_{z\sim qx}(\log p(x|z))+E_{z\sim qx}(\log p(x)))\\ & =\underset{(g,h)\in G\times H}{\mathrm{argmax}}(E_{z\sim q_x}(\log p(x|z))-KL(q_x(z),p(z)))\\ & =\underset{(g,h)\in G\times H}{\mathrm{argmax}}(E_{z\sim q_x}(-\frac{||x-f(z)|{|}^2}{2c})-KL(q_x(z),p(z)))\end{array}$$
在上述公式的倒数第二个方程中可以观察到，在近似后验概率 $p(z|x)$ 时，我们在最大化“观测”似然 $E_{z\sim qx}(\log q_x(z))$ 和保持接近先验分布 $KL(q_x(z),p(z))$ 之间有个权衡。这种权衡对于贝叶斯推断问题是自然的，并且表达了我们需要在数据和先验之间找到的平衡。 其中，$q_x(z)$ 表示在给定数据 $x$ 的条件下,对隐变量 $z$ 的一个假设后验分布。$E_{z\sim qx}(\log p(x))$ 表示从分布 $q(x)$ 中采样 $z$ 并对 $\log p(x)$ 求期望，因为 p(x) 与变量 z 无关，所以实际上 $E_{z\sim qx}(\log p(x))$ 的值就是一个常数 $\log p(x)$，而常数对于上述方程来说并不影响，于是可以直接忽略。

到目前为止，我们假设函数 $f$ 是已知且固定的，并且我们已经展示了在这些假设下，我们可以使用变分推断技术来近似后验概率 $p(z|x)$。然而在实践中，定义解码器的函数 $f$ 是未知的，同样需要选择。还记得我们的初始目标是找到一个性能良好的编码-解码方案，让其潜在空间足够规则，可以用于生成。如果规则性主要由潜在空间上假设的先验分布 $p(x|z)$ 决定，那么整体编码-解码方案的性能高度依赖于函数 $f$ 的选择。实际上，由于 $p(z|x)$ 可以从 $p(z)$ 和 $p(x|z)$ 近似（通过变分推断），而 $p(z)$ 是一个简单的标准高斯分布，我们在模型中可以进行优化的唯一两个变量是参数 c（定义似然的方差）和函数 f（定义似然的均值）。

正如我们之前讨论的那样，对于 F 中的任何函数 f（每个函数定义了不同的概率解码器$p(x|z)$ ），我们可以得到 $p(z|x) $ 的最佳近似，记为 $q_x^{\ast }(z)$。尽管它具有概率性质，但我们尽可能寻找高效的编码-解码方案，因此我们希望能够找到最佳的函数 $f^{\ast }$ 使得从分布 $q_x^{\ast }(z)$ 中采样 $z$ 时，能够最大化对数似然函数 $p(x|z)$，并且从 $p(x|z)$ 采样得到 $x^{\ast }$：
$$\begin{array}{rl}{f}^{\ast }& =\underset{f\in F}{\mathrm{argmax}}{E}_{z\sim q_x^{\ast }}(\log p(x|z))\\ & =\underset{f\in F}{\mathrm{argmax}}{E}_{z\sim q_x^{\ast }}(-\frac{||x-f(z)|{|}^2}{2c})\end{array}$$
我们需要找到最佳的 $f^{\ast }$, $g^{\ast }$ 和 $h^{\ast }$，其中 $q_x^{\ast }(z)$ 取决于函数 $f$，使得：

$$\begin{array}{r}(f^x,g^{\ast },h^{\ast })=\underset{(f,g,h)\in F\times G\times H}{\mathrm{argmax}}(E_{z\sim q_x}(-\frac{||x-f(z)|{|}^2}{2c})-KL(q_x(z),p(z)))\end{array}$$

在这个目标函数中，我们可以识别出在前一节对 VAE 的直观描述中所介绍的元素：$x$ 和 $f(z)$ 之间的重构误差，以及 $q_x(z)$ 和 $p(z)$ 之间的KL散度构成的正则化项（其中 $p(z)$ 是一个标准高斯分布）。我们还可以注意到常数 $c$，它决定了两项之间的平衡，当 $c$ 的值越大，更倾向于优化正则化项，因为重构误差的变化受 $c$ 的影响变得更小。反之当 $c$ 的值越小，重构误差的变化更大，于是更倾向于重构误差项。

3. 结合神经网络（Bringing neural networks into the model）

到目前为止，我们已经建立了一个依赖于三个函数 $f$、$g$ 和 $h$ 的概率模型，并使用变分推断来表达解决优化问题的方案，以获得给定该模型的最佳编码-解码方案的 $f^{\ast }$、$g^{\ast }$ 和 $h^{\ast }$。由于我们无法对整个函数空间进行优化，于是我们选择限制优化域，并决定将$f$、$g$ 和 $h$ 表示为神经网络。因此，F、G和 H 分别对应于由网络架构定义的函数族，优化是在这些网络的参数上进行的。

在实践中，$g$ 和 $h$ 并不是由两个完全独立的网络定义的，而是共享部分架构和权重，这样我们就有：
$$\begin{array}{rlrlr}g(x)=g_2(g_1(x))& & h(x)=h_2(h_1(x))& & g_1(x)=h_1(x)\end{array}$$
由于 $h$ 定义了 $q_x(z)$ 的协方差矩阵，所以 $h(x)$ 应该是一个方阵。然而，为了简化计算并减少参数数量，我们额外假设对 $p(z|x)$ 的近似 $q_x(z)$ 是一个具有对角协方差矩阵的多维高斯分布（假设变量之间相互独立）。在这个假设下，$h(x)$ 就是协方差矩阵对角元素的向量，并且与 $g(x)$ 具有相同的大小。然而，通过这种方式我们减少了我们在变分推断中所考虑的分布族，因此对 $p(z|x)$ 的近似可能不够准确。

编码器部分建模了$p(z|x)$，它是一个均值和协方差都是 $x$ 的函数（ $g$ 和 $h$ ）所构成的的高斯分布。而解码器假设 $p(x|z)$ 是一个具有固定协方差的高斯分布。变量$z$ 定义的均值由一个神经网络建模，并可以表示如下：

整个架构是通过连接编码器和解码器部分构成的，然而，在训练过程中，我们仍然需要非常小心地对来自编码器返回的分布进行采样。采样过程必须允许误差能够通过神经网络进行反向传播。一个简单的技巧叫重参数化（reparametrisation trick），它使得梯度下降能够进行，尽管在架构的中间发生了随机采样，但它却利用了以下事实：如果 $z$ 是一个随机变量，其遵循均值为 $g(x)$ 和协方差为 $H(x)=h(x).h^T(x)$ 的高斯分布，那么它可以表示为：
$$\begin{array}{rlr}z=h(x)\zeta +g(x)& & \zeta \sim N(0,I)\end{array}$$

最终损失函数由重构项、正则化项和一个常数（$C=\frac{1}{2c}$）组成，这个常数用于定义这两个项的相对权重。
$$\begin{array}{rl}loss& =C|||x-\hat{x}|{|}^2+KL(N({\mu }_x,{\sigma }_x),N(0,I))\\ & =C|||x-f(z)|{|}^2+KL(N(g(x),h(x)),N(0,I))\end{array}$$