使用 TensorFlow Probability Layers 的变分自编码器

文 / Ian Fischer、Alex Alemi、Joshua V. Dillon 以及 TFP 团队

在 2019 年的 TensorFlow 开发者峰会上，我们推出 TensorFlow Probability (TFP) Layers。在当时的演示中，我们展示了如何使用极少行的代码构建强大的回归模型。在本文中，我们将介绍如何使用 TFP Layers 轻松构建变分自编码器 (VAE)。

TensorFlow Probability Layers

TFP Layers 为使用 Keras 构建包含深度网络的分布提供高级 API。您可使用此 API，轻松构建融合深度学习和概率编程的模型。例如，我们可以利用神经网络输出将概率分布参数化。我们会在这里使用此方法。

注：Keras 链接

https://www.tensorflow.org/guide/keras

变分自编码器和 ELBO

变分自编码器 (VAE) 是非常普及的生成模型，应用于多个不同领域，包括协同过滤、图像压缩、强化学习，以及音乐和草图生成。

在传统的 VAE 推导中，我们会设想生成数据的某个过程，例如隐变量生成模型。思考绘制数字的过程，就像在 MNIST 中一样。假设在绘制数字前，您先决定好要绘制的数字，并在脑海中想象模糊的画面。然后，您在纸上落笔，试着真正画出脑海中的画面。我们可以将这两个步骤正式化：

您从某先验分布 z ~ p(z) 中 提取一些隐藏表征 z 的样本。这就是您脑海中的模糊画面，假设是 “3”。

根据样本，您绘制出实际图画的表征 x，并将表征本身建模为随机过程 x ~ p(x|z)。您每次写 “3” 时，该模型就会捕捉您的想法，而结果看起来至少有点不同。

因此，在创建手写数字时，我们设想某些差异的存在是由过程固有的某类 信号 所造成，例如 MNIST 数字的类标识，而有一些差异是由 噪声所造成，例如在同一个人绘制同一个数字的不同样本中，线条角度的精确度有所不同。概括地说，VAE 试图通过两个过程的显式模型将信号和噪声分开。

为了针对该目标进行训练，我们将 ELBO（证据下界）目标最大化：

其中三个概率密度函数为：

• p(z)， 隐藏 表征 z 的先验

• q(z|x)， 变分编码器

• p(x|z)， 解码器 ，鉴于隐藏表征 x ，绘制出图像 z 的可能性有多大

ELBO 是对数 p(x) 的下界，即观测数据点的对数概率。ELBO 方程的第一个积分是重构项，其询问我们从图像 x 开始、将其编码为 z、解码，然后返回原始 x 的可能性有多大。第二项是 KL 散度项。该项衡量编码器和先验的接近程度；您可以认为该项只是想让编码器保持 “诚实”。如果根据先验，不太可能绘制出编码器生成的 z 样本，则目标比编码器生成比先验更典型的 z 样本更糟。因此，如果要让编码器与先验保持不同，除非重构项的益处大于这样做的成本。

代码

从上面的介绍中，我们可以看出，对以下三个不同成分单独建模是很自然的做法： 先验 p(z)、 变分编码器 q(z|x)，以及 解码器 p(x|z)。您可以通过运行 此 Colab 遵循这一做法。该 Colab 在几分钟内于云 GPU 上使用 MNIST 训练 VAE。

注：此 Colab 链接

https://colab.research.google.com/github/tensorflow/probability/blob/master/tensorflow_probability/examples/jupyter_notebooks/Probabilistic_Layers_VAE.ipynb

先验

VAE 惯用的最简单的先验是各向同性高斯：

tfd = tfp.distributions
encoded_size = 16
prior = tfd.Independent(tfd.Normal(loc=tf.zeros(encoded_size), scale=1),
                       reinterpreted_batch_ndims=1)

我们刚刚创建了不含习得参数且无关 TFP 的高斯分布，并且指定隐变量 z 将拥有 16 个维度。

编码器

对于编码器分布，我们将使用全协方差高斯分布，其平均值和协方差矩阵通过神经网络输出参数化。这听起来可能很复杂，但使用 TFP Layers 表达则非常简单：

tfpl = tfp.layers
encoder = tfk.Sequential([
   tfkl.InputLayer(input_shape=input_shape),
   tfkl.Lambda(lambda x: tf.cast(x, tf.float32) - 0.5),
   tfkl.Conv2D(base_depth, 5, strides=1,
               padding='same', activation=tf.nn.leaky_relu),
   tfkl.Conv2D(base_depth, 5, strides=2,
               padding='same', activation=tf.nn.leaky_relu),
   tfkl.Conv2D(2 * base_depth, 5, strides=1,
               padding='same', activation=tf.nn.leaky_relu),
   tfkl.Conv2D(2 * base_depth, 5, strides=2,
               padding='same', activation=tf.nn.leaky_relu),
   tfkl.Conv2D(4 * encoded_size, 7, strides=1,
               padding='valid', activation=tf.nn.leaky_relu),
   tfkl.Flatten(),
   tfkl.Dense(tfpl.MultivariateNormalTriL.params_size(encoded_size),
              activation=None),
   tfpl.MultivariateNormalTriL(
       encoded_size,
       activity_regularizer=tfpl.KLDivergenceRegularizer(prior, weight=1.0)),
])

编码器只是一个标准的 Keras 序列模型，由卷积和密集层组成，但其输出会传送给 TFP 层 MultivariateNormalTril()，而该层将最后一层 Dense() 中的激活透明地分割成指定平均值和（下三角）协方差矩阵（多元正态分布参数）所需的部分。我们使用辅助函数 tfpl.MultivariateNormalTriL.params_size(encoded_size)，确保 Dense() 层输出正确数量的激活（即分布的参数）。最后，我们认为该分布应该为最终损失提供一个 “正则化” 项。具体来说，我们将在编码器和损失的先验 （即上文所述的 ELBO 的 KL 项）之间添加 KL 散度。（小资料：我们可以将此 VAE 转换成 β-VAE，方法非常简单，只需将 weight 参数更改为除 1 之外的数值即可！）

解码器

对于解码器，我们将使用简单的 “平均场解码器”，而在此案例中，我们将使用与像素无关的伯努利分布：

decoder = tfk.Sequential([
   tfkl.InputLayer(input_shape=[encoded_size]),
   tfkl.Reshape([1, 1, encoded_size]),
   tfkl.Conv2DTranspose(2 * base_depth, 7, strides=1,
                        padding='valid', activation=tf.nn.leaky_relu),
   tfkl.Conv2DTranspose(2 * base_depth, 5, strides=1,
                        padding='same', activation=tf.nn.leaky_relu),
   tfkl.Conv2DTranspose(2 * base_depth, 5, strides=2,
                        padding='same', activation=tf.nn.leaky_relu),
   tfkl.Conv2DTranspose(base_depth, 5, strides=1,
                        padding='same', activation=tf.nn.leaky_relu),
   tfkl.Conv2DTranspose(base_depth, 5, strides=2,
                        padding='same', activation=tf.nn.leaky_relu),
   tfkl.Conv2DTranspose(base_depth, 5, strides=1,
                        padding='same', activation=tf.nn.leaky_relu),
   tfkl.Conv2D(filters=1, kernel_size=5, strides=1,
               padding='same', activation=None),
   tfkl.Flatten(),
   tfpl.IndependentBernoulli(input_shape, tfd.Bernoulli.logits),
])

这里采用的方式与编码器基本相同，但现在我们要使用转置的卷积来获取隐藏表征（即一个 16 维向量），并将其转换成 28 x 28 x 1 的张量。最后完成转换的张量会参数化与像素无关的伯努利分布。

损失

至此，我们已经准备好构建完整的模型，并指定损失函数的其余部分。

vae = tfk.Model(inputs=encoder.inputs,
               outputs=decoder(encoder.outputs[0]))

我们的模型只是 Keras 模型，其输出被定义为编码器和解码器的组合。由于编码器已经将 KL 项添加到损失中，我们只需指定重构损失（上文 ELBO 的第一项）即可。

negative_log_likelihood = lambda x, rv_x: -rv_x.log_prob(x)

vae.compile(optimizer=tf.optimizers.Adam(learning_rate=1e-3),
           loss=negative_log_likelihood)

损失函数有两个参数：原始输入 x 以及模型的输出。我们将其称为 rv_x ，因为它是一个随机变量。此示例展示 TFP Layers 的一些核心概念。虽然 Keras 和 TensorFlow 将 TFP Layers 视为输出张量，但 TFP Layers 实际上是 分布 对象。因此，根据模型，我们可以让损失函数成为数据的负对数似然函数：-rv_x.log_prob(x)。

小知识

我们应该花点时间了解，TFP Layers 实际通过执行哪些操作，从而以透明的方式与 Keras 集成。正如我们所言，TFP Layer 的输出实际上是分布对象。我们可以通过以下代码进行确认：

x = eval_dataset.make_one_shot_iterator().get_next()[0][:10]
xhat = vae(x)
assert isinstance(xhat, tfd.Distribution)

但如果 TFP Layer 返回一个分布，那么当我们用编码器的输出：decoder_model(encoder_model.outputs[0])) 组合解码器时会发生什么？为了让 Keras 将编码器分布视作张量，TFP Layers 实际上将分布 “具体化” 为该分布中的样本，而 Keras 将分布对象视作我们将获得的张量（前提是我们调用了 encoder_model.sample()）只是一种好听的说法而已。但是，如果我们需要直接访问分布对象，可以按照调用 rv_x.log_prob(x) 时在损失函数中的做法操作。TFP Layers 自动提供类似分布和张量的行为，因此您无需担心 Keras 会将两者混淆。

训练

训练此模型与训练任何 Keras 模型一样简单：我们只需调用 vae_model.fit()：

vae.fit(train_dataset,
       epochs=15,
       validation_data=eval_dataset)

使用此模型，我们可以获得约 115 奈特的 ELBO（ 奈特是等同于比特的自然对数，115 奈特大约等于 165 比特）。当然，这并不是最佳的表现，但从这个基本设置开始，要让这三个成分中的任何一个变得更强大并非难事。而且，该模型目前生成的数字外观非常美观、清晰！

通过解码 MNIST 测试集中的图像生成的解码器模式

通过从先验中采样生成的解码器模式

总结

在本文中，我们演示了如何将深度学习与概率编程相结合：我们构建了一个变分自动编码器。该编码器使用 TFP Layers 将 Keras 序列模型的输出传递给 TFP 中的概率分布。

我们利用 TFP Layers 中类似张量和分布的语义，相对简化我们的代码。

更多 AI 相关阅读：