使用 TensorFlow Probability 的概率层执行回归
文 / Pavel Sountsov、Chris Suter、Jacob Burnim、Joshua V. Dillon 和 TensorFlow Probability 团队
背景
在 2019 年的 TensorFlow 开发者峰会上，我们推出 TensorFlow Probability (TFP) 的概率层。在本文中，我们会更详细地演示如何使用 TFP 层管理回归预测中固有的不确定性。
回归和概率
回归是机器学习从业者用于解决预测问题的最基本技术之一。然而，许多基于回归的分析会省略适当量化预测中不确定性的步骤，部分原因在于所需的复杂程度。若要开始量化不确定性，提出该问题的一个特别简单的方法是将回归模型编写为 P(y | x, w)，其中标签的概率分布为 (y)，而提供的输入为 (x)，参数为 (w)。我们可以调整此模型以适合数据，方法是最大化标签的概率，或者等效地最小化负对数似然损失：-log P(y | x)。在 Python 中：
negloglik = lambda y, p_y: -p_y.log_prob(y)
我们可以将各种标准连续函数、分类函数和损失函数与此回归模型搭配使用。例如，用于连续标签的均方误差损失函数表示 P(y | x, w) 是具有固定尺度（标准差）的正态分布。用于分类的交叉熵损失函数表示 P(y | x, w) 是指分类分布。
在本文中，我们将展示如何将 TensorFlow Probability (TFP) 中的概率层与 Keras 搭配使用，以在此简单的基础上进行构建，从而逐步推理出现有任务中逐渐增多的不确定性。您可根据 Google Colab 中的说明执行操作（https://colab.research.google.com/github/tensorflow/probability/blob/master/tensorflow_probability/examples/jupyter_notebooks/Probabilistic_Layers_Regression.ipynb）。
案例 1：简单的线性回归
我们首先将介绍适合某些数据的简单的线性回归模型：
import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions
# Build model.
model = tf.keras.Sequential([
tf.keras.layers.Dense(1),
tfp.layers.DistributionLambda(lambda t: tfd.Normal(loc=t, scale=1)),
])
# Do inference.
model.compile(optimizer=tf.optimizers.Adam(learning_rate=0.05), loss=negloglik)
model.fit(x, y, epochs=500, verbose=False)
# Make predictions.
yhat = model(x_tst)
任何使用过 Keras 的人应该都熟悉其推理和预测部分，但模型构建看起来会有所不同。我们来明确说明这一点：我们使用正态分布对标签进行建模，其中尺度 1 的中心位置（平均值）取决于输入值。实际上，tfp.layers.DistributionLambda 层会返回一个 tfd.Distribution 特殊实例（关于这一点的更多详细信息，请参阅附录 A），因此我们可以自由获取其平均值，并在数据旁边进行绘制：
mean = yhat.mean()
如此一来，我们便可捕获数据的整体趋势（蓝色圆圈），以及标签分布的预测平均值。不过，我们可以看出数据的结构较多：似乎 随着 x 的增加， y 的可变性也会增强。到目前为止，我们编写的模型无法捕捉到这个细节，但在下一部分中，我们将展示如何修改模型，以便赋予其这项能力。
案例 2：已知的未知数
在上一部分中，我们已经看出，无论 x 具体 为何值， y 都存在可变性。我们可以将这种可变性视作问题的固有特性。这意味着即使拥有无限的训练集，我们也仍然无法完美地预测标签。这种不确定性的常见示例是公平抛硬币游戏的结果（假设您没有详细的物理模型等）。无论过去看到过多少次硬币的正反面，我们都无法预测将来会看到哪一面。
我们将假设这种可变性与 x 的值之间具有已知的函数关系。我们采用求 y 的平均值时所用的相同线性函数对这种关系进行建模。
# Build model.
model = tfk.Sequential([
tf.keras.layers.Dense(1 + 1),
tfp.layers.DistributionLambda(
lambda t: tfd.Normal(loc=t[..., :1],
scale=1e-3 + tf.math.softplus(0.05 * t[..., 1:]))),
])
# Do inference.
model.compile(optimizer=tf.optimizers.Adam(learning_rate=0.05), loss=negloglik)
model.fit(x, y, epochs=500, verbose=False)
# Make predictions.
yhat = model(x_tst)
现在，除了预测标签分布的平均值外，我们还要预测其尺度（标准差）。在完成训练，并以相同方式形成预测后，我们可以有效预测 y 随 x 变化的可变性。同前：
mean = yhat.mean()
stddev = yhat.stddev()
mean_plus_2_stddev = mean - 2. * stddev
mean_minus_2_stddev = mean + 2. * stddev
效果好多了！我们的模型现在不太确定当 x 变大时，y 应该如何变化。我们将这种不确定性称为偶然 不确定性，因为它代表基础过程的固有变化。虽然我们取得进展，但偶然不确定性并非此问题中唯一的不确定性来源。在进一步探讨之前，让我们思考我们至今忽略的另一个不确定性来源。
案例 3：未知的未知数
数据中存在噪声意味着我们无法完全确定 x 和 y 之间线性关系的参数。例如，我们在上一部分中发现的倾斜似乎很合理，但其实我们并不确定，或许倾斜角度略微平缓或更陡也很合理。我们将这种不确定性称为认知 不确定性；不同于偶然不确定性，如果我们获得更多数据，便可以减少认知不确定性。为了解这种不确定性，我们将标准的 Keras Dense 层替换为 TFP 的 DenseVariational 层。
DenseVariational 层针对权重使用变分后验 Q(w)，以表示其值的不确定性。该层对 Q(w) 进行调整，使其接近 先验 分布 P(w)，而 P(w) 会在我们查看数据前，对权重中的不确定性进行建模。
对于 Q(w)，我们会针对变分后验使用多元正态分布，其中包含以可训练位置为中心的可训练对角协方差矩阵。对于 P(w)，我们将针对先验使用标准多元正态分布，其中包含可训练位置和固定尺度。关于此层工作原理的更多详细信息，请参阅附录 B。
我们将所有内容整合到一起：
# Build model.
model = tf.keras.Sequential([
tfp.layers.DenseVariational(1, posterior_mean_field, prior_trainable),
tfp.layers.DistributionLambda(lambda t: tfd.Normal(loc=t, scale=1)),
])
# Do inference.
model.compile(optimizer=tf.optimizers.Adam(learning_rate=0.05), loss=negloglik)
model.fit(x, y, epochs=500, verbose=False)
# Make predictions.
yhats = [model(x_tst) for i in range(100)]
尽管涉及的算法较为复杂，但使用 DenseVariational 层却很简单。上述代码的一个有趣之处在于，使用具有此类层的模型进行预测时，我们每次都会得到不同的答案。这是因为 DenseVariational 本来就定义一组模型。让我们看看，我们可以从该集合中了解哪些有关模型参数的信息。
每条线代表对后验分布中的模型参数执行不同的随机抽取。正如我们所见，线性关系实际上存在相当多的不确定性。即使我们不关心 y 针对任何特定 x 值的可变性，但如果我们对 x的预测与 0 相差太多，倾斜中的不确定性也会让我们暂停预测。
请注意，在此示例中，我们同时训练 P(w) 和 Q(w)。该训练与使用 经验贝叶斯方法 或 第二类最大似然之间线性关系的参数。通过使用这种方法，我们不需要为倾斜和截距参数指定先验的位置，如果我们不了解有关此问题的先验知识，可能很难得出正确的结果。此外，如果您将先验设置为与其真实值相差太多，则后验可能会受到此选择带来的不利影响。使用第二类最大似然时，请注意，您会失去某些针对权重的正则化优势。如果想对不确定性进行适当的贝叶斯处理（如果您了解一些或较复杂的先验知识），您可以使用不可训练的先验（参见附录 B）。
案例 4：已知和未知的未知数
单独了解偶然不确定性和认知不确定性后，现在我们可以使用 TFP 层的可组合 API，创建同时报告两种不确定性的模型：
# Build model.
model = tf.keras.Sequential([
tfp.layers.DenseVariational(1 + 1, posterior_mean_field, prior_trainable),
tfp.layers.DistributionLambda(
lambda t: tfd.Normal(loc=t[..., :1],
scale=1e-3 + tf.math.softplus(0.01 * t[..., 1:]))),
])
# Do inference.
model.compile(optimizer=tf.optimizers.Adam(learning_rate=0.05), loss=negloglik)
model.fit(x, y, epochs=500, verbose=False);
# Make predictions.
yhats = [model(x_tst) for _ in range(100)]
我们对之前模型所做的唯一改变是在 DenseVariational 层中添加一个额外的输出，以便对标签分布的尺度也进行建模。正如我们在之前的解决方案中所述，我们获得一组模型，但这次所有模型还会报告 y 随 x 变化的可变性。我们来绘制这一集合：
请注意，此模型的预测相比仅考虑偶然不确定性的模型的预测，存在着质的区别：该模型预测可变性更高，因为除了更倾向于正值外， x 也会更倾向于负值，而简单的偶然不确定性线性模型根本无法做到这一点。
结论
TFP 层的指导原则是从业者应该集中精力编写模型而不是损失函数。在整篇博文中，我们确保用户指定的损失函数（即实现负对数似然的 negloglik 函数）始终相同，同时对模型进行局部修改，以处理越来越多类型的不确定性。此外，您还可使用 API 在最大似然学习、第二类最大似然和完整贝叶斯处理之间自由切换。我们认为，此 API 能够显著简化概率模型的构建过程，并且很高兴与全世界分享。
此 API 将在下一个稳定版本（即 TensorFlow Probability 0.7.0）中投入使用，并且现可在 Nightly 版本中使用。请加入我们的 tfprobability@tensorflow.org 论坛，查看最新的 TensorFlow Probability 公告和其他 TFP 讨论。
注：tfprobability@tensorflow.org 链接
https://groups.google.com/a/tensorflow.org/forum/#!forum/tfprobability
彩蛋：不设限制
到目前为止，我们一直在假设数据会呈线性分布。如果我们不知道输入和标签之间的函数关系，结果会如何？假设我们有一种模糊的感觉，即只有当对应的输入接近我们观察到的输入时，预测标签才应与已经看到的标签相似？换言之，我们想要做的唯一假设是我们要与数据进行拟合的函数是平滑函数。
做出这类假设时，执行回归的标准工具是高斯过程。此强大的模型使用核函数对有关以下内容的平滑假设（以及其他全局函数属性）进行编码：输入和标签之间的关系应采用的形式。该模型以数据为条件，针对与这些假设和数据一致的 函数 形成概率分布。
TFP 提供 VariationalGaussianProcess 层，该层对完整的高斯过程使用变分近似法（与我们在上面案例 3 和案例 4 中的做法理念类似），以获得有效且灵活的回归模型。为简单起见，我们将只考虑有关输入和标签之间关系形式的认知不确定性。至于我们要做的假设，我们将简单地假设我们拟合的函数是局部平滑函数：该函数可以在整个数据集中尽可能大地变化，但如果两个输入接近，它将会返回相似的值。
num_inducing_points = 40
model = tf.keras.Sequential([
tf.keras.layers.InputLayer(input_shape=[1], dtype=x.dtype),
tf.keras.layers.Dense(1, kernel_initializer='ones', use_bias=False),
tfp.layers.VariationalGaussianProcess(
num_inducing_points=num_inducing_points,
kernel_provider=RBFKernelFn(dtype=x.dtype),
event_shape=[1],
inducing_index_points_initializer=tf.constant_initializer(
np.linspace(*x_range, num=num_inducing_points,
dtype=x.dtype)[..., np.newaxis]),
unconstrained_observation_noise_variance_initializer=(
tf.constant_initializer(
np.log(np.expm1(1.)).astype(x.dtype))),
),
])
# Do inference.
batch_size = 32
loss = lambda y, rv_y: rv_y.variational_loss(
y, kl_weight=np.array(batch_size, x.dtype) / x.shape[0])
model.compile(optimizer=tf.optimizers.Adam(learning_rate=0.01), loss=loss)
model.fit(x, y, batch_size=batch_size, epochs=1000, verbose=False)
# Make predictions.
yhats = [model(x_tst) for _ in range(100)]
由于此模型具有强大的功能，因此其定义要复杂得多：我们需要定义新的损失函数，而且要指定更多的参数。在不久的将来，TFP 团队将致力于进一步简化此模型。不过，增加这种复杂性非常值得，以下结果证明了这一点：
VariationalGaussianProcess 在训练集中发现一个周期性结构！事实上，在整篇博文中，我们从头到尾使用的数据中就有这种结构。您是否先于模型注意到这一点？值得注意的是，虽然我们并未告知模型数据中存在任何此类周期性，但模型仍然发现了这种结构。而且，正如宣传中所说，该模型仍然为我们带来一定程度的不确定性。例如，接近 0 时，周期性结构不是很明显，因此模型并未确定该区域中存在任何此类关系。
附录 A：DistributionLambda 如何运作？
DistributionLambda 是特殊的 Keras 层，使用 Python lambda 来构建以层输入为条件的分布：
layer = tfp.layers.DistributionLambda(lambda t: tfd.Normal(t, 1.))
distribution = layer(2.)
assert isinstance(distribution, tfd.Normal)
distribution.loc
# ==> 2.
distribution.stddev()
# ==> 1.
该层使我们能够用之前的方法编写 NegLogLik 损失函数，因为 Keras 将模型最后一层的输出传递给损失函数，而对于本篇博文中的模型，所有这些层都会返回分布。关于使用这些层的更多方法，请参阅文章使用 Tensorflow Probability Layers 的变分自编码器。
附录 B：DenseVariational 如何运作？
DenseVariational 层支持使用变分推理来学习权重分布。此操作的实现方式是最大化 ELBO（证据下界）目标：
ELBO 使用三种分布：
P(w) 是针对权重的先验，是我们在训练模型前假设权重要遵循的分布
Q(w; θ) 是由参数 θ 参数化的变分后验。这是我们完成模型训练后权重分布的近似值
P(Y | X, w) 是与所有输入 X、所有标签 Y 以及权重相关的似然函数。当用作针对 Y的概率分布时，它会指定变化 Y ，并提供 X 和权重
ELBO 是对数 P(Y | X) 的下界，即在对权重的不确定性进行边缘化处理后，根据输入确定的标签对数似然性。ELBO 的工作原理是将 Q 的 KL 散度与权重的先验（第二项）进行权衡，并能够根据输入（第一项）预测标签。当数据很少时，由第二项主导，而我们的权重仍然接近先验分布，作为连带作用，这有助于防止过度拟合。
DenseVariational 分别计算 ELBO 的这两个项。第一项的计算方法是使用 Q 中的单个随机样本求取近似值。如果我们仔细查看该项，则会发现，对于任何特定的 w 值，该项的结果恰好是我们在本篇博文中用于回归的负对数似然损失。因此，通过简单地绘制一组来自 Q 的随机权重，然后计算常规损失，我们会自动获得 ELBO 第一项的近似值。
第二项由分析计算得出，然后作为正则化损失添加到层中，类似于我们指定 L2 正则化等内容时的做法。Keras 将此损失添加到我们的第一项中。
用于计算第一项的抽样方法解释了我们如何通过多次调用具有相同输入的模型来生成多个模型：我们每次都会根据 Q 分布，抽样一组新的权重。
我们如何指定先验和变分后验？正如我们在案例 2 中所见，先验和变分后验都是可训练分布。例如，我们在案例 3 中使用的可训练先验定义如下：
def prior_trainable(kernel_size, bias_size=0, dtype=None):
n = kernel_size + bias_size
return tf.keras.Sequential([
tfp.layers.VariableLayer(n, dtype=dtype),
tfp.layers.DistributionLambda(lambda t: tfd.Independent(
tfd.Normal(loc=t, scale=1),
reinterpreted_batch_ndims=1)),
])
这只是一个可调用先验，会返回含有 DistributionLambda 层的常规 Keras 模型！这里唯一的新成分是 VariableLayer，其只返回可训练变量的值，并忽略任何输入（因为先验不以任何输入为条件）。请注意，如果我们想将此转换为不可训练的先验，我们会将 trainable=False 传递给 VariableLayer 构造函数。
更多 AI 相关阅读：