作者：Emily Fertig、Joshua V. Dillon、Wynn Vonnegut、Dave Moore 和 TensorFlow Probability 团队

在本文中，我们介绍了在 TensorFlow Probability 中使用联合分布进行变分推理的新工具，并展示了如何使用它们来估计回归模型中权重的贝叶斯可信区间。

概述

变分推理 (VI) 将近似贝叶斯推理视为一个优化问题，并寻求一个“替代”后验分布，该分布将 KL 散度与真实后验分布最小化。基于梯度的 VI 通常比 MCMC 方法更快，自然地与模型参数的优化相结合，并提供模型证据的下限，该下限可直接用于模型比较、收敛诊断和可组合推理。

TensorFlow Probability (TFP) 提供了用于快速、灵活和可扩展的 VI 的工具，这些工具自然地融入 TFP 堆栈。这些工具能够构建由线性变换或归一化流引起的协方差结构的替代后验分布。

VI 可用于估计回归模型参数的贝叶斯 可信区间，以估计各种处理或观察到的特征对感兴趣结果的影响。可信区间根据参数的后验分布（以观察到的数据为条件，并假设参数的先验分布）以一定的概率对未观察到的参数值进行约束。

在本文中，我们展示了如何使用 VI 获取贝叶斯线性回归模型参数的可信区间，该模型用于测量房屋中的氡含量（使用 Gelman 等人（2007 年）的氡数据集；请参阅 Stan 中的类似示例）。我们展示了 TFP JointDistributions 如何与双射器结合，构建和拟合两种类型的表达性替代后验分布

• 由块矩阵变换的标准正态分布。该矩阵可以反映后验某些分量之间的独立性和其他分量之间的依赖性，从而放宽对平均场或全协方差后验的假设。
• 更复杂、更高容量的 逆自回归流。

替代后验分布经过训练，并与来自平均场替代后验基线的結果进行比较。这些图显示了使用三种 VI 替代后验分布获得的四个模型参数的可信区间，如下所述，以及用于比较的哈密顿蒙特卡罗 (HMC)。

您可以在此 Google Colab 中详细了解并查看所有细节。

示例：氡测量值的贝叶斯分层线性回归

氡是一种放射性气体，它通过与地面的接触点进入房屋。它是一种致癌物质，是非吸烟者肺癌的主要原因。氡含量因房屋而异。

美国环保署对 80,000 户房屋的氡含量进行了研究。两个重要预测因子是

• 进行测量所在的楼层（地下室的氡含量更高）
• 县铀含量（与氡含量呈正相关）

预测按县分组的房屋的氡含量是贝叶斯分层建模中的一个经典问题，由 Gelman 和 Hill（2006 年） 提出。我们对明尼苏达州房屋氡含量与位置（县）关系的可信区间感兴趣。为了隔离这种影响，模型中还包括楼层和铀含量的影响。此外，我们将纳入一个对应于每个县的测量平均楼层的上下文效应，以便如果测量所在的楼层在各县之间存在差异，则不会将其归因于县效应。

回归模型指定如下

其中 i 是观测值的索引，county_i 是进行第 i 次观测的县。

我们使用县级随机效应来捕捉地理差异。参数 uranium_weight 和 county_floor_weight 以概率方式建模，floor_weight 和常数 bias 是确定性的。这些 建模 选择在很大程度上是任意的，并且是为了演示在合理复杂性的概率模型上进行 VI 的目的。有关使用氡数据集在 TFP 中使用固定效应和随机效应进行多级建模的更深入讨论，请参阅 多级建模入门 和 使用变分推理拟合广义线性混合效应模型。

此示例的完整代码可在 Github 上获得。

为确定性参数和正态分布尺度参数定义变量，后者约束为正值。

import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions
tfb = tfp.bijectors

floor_weight = tf.Variable(0.)
bias = tf.Variable(0.)
log_radon_scale = tfp.util.TransformedVariable(1., tfb.Exp())
county_effect_scale = tfp.util.TransformedVariable(1., tfb.Exp())

我们将回归的概率图形模型指定为 TFP JointDistribution。

@tfd.JointDistributionCoroutineAutoBatched
def model():
 uranium_weight = yield tfd.Normal(0., scale=1., name='uranium_weight')
 county_floor_weight = yield tfd.Normal(
     0., scale=1., name='county_floor_weight')
 county_effect = yield tfd.Sample(
     tfd.Normal(0., scale=county_effect_scale),
     sample_shape=[num_counties], name='county_effect')
 yield tfd.Normal(
     loc=(log_uranium * uranium_weight
          + floor_of_house * floor_weight
          + floor_by_county * county_floor_weight
          + tf.gather(county_effect, county, axis=-1)
          + bias),
     scale=log_radon_scale[..., tf.newaxis],
     name='log_radon')

我们将 log_radon 固定到观察到的氡数据以对非归一化后验进行建模。

target_model = model.experimental_pin(log_radon=log_radon)

贝叶斯变分推理的快速总结

假设我们有以下生成过程，其中 𝜃 代表随机参数（回归模型中的 uranium_weight、county_floor_weight 和 county_effect），ω 代表确定性参数（floor_weight、log_radon_scale、county_effect_scale 和 bias）。x_𝑖 是特征（log_uranium、floor_of_house 和 floor_by_county），而 𝑦_𝑖 是 𝑖 = 1...n 个观察数据点的目标值（log_radon）

然后，VI 的特征在于

（从技术上讲，我们假设 q 相对于 r 是 绝对连续 的。另请参阅 詹森不等式。）

由于该边界对所有 q 成立，因此对于

关于术语，我们称

• q* 为“替代后验”，而
• Q 为“替代族”。

ω* 表示 VI 损失上确定性参数的最大似然值。有关变分推理的更多信息，请参阅 这篇综述。

表达性替代后验

接下来，我们使用具有两种不同类型替代后验的 VI 估计参数的后验分布

• 具有由分块矩阵变换引起的协方差结构的约束多元正态分布。
• 由 逆自回归流 变换的多元标准正态分布，然后进行拆分和重构以匹配后验的支持。

多元正态替代后验

为了构建这种替代后验分布，使用可训练的线性算子来诱导后验分量之间的相关性。

我们首先构建一个具有向量值标准正态分量的基本分布，其大小等于相应先验分量的大小。这些分量是向量值的，因此可以由线性算子进行变换。

flat_event_size = tf.nest.map_structure(
     tf.reduce_prod,
     tf.nest.flatten(target_model.event_shape_tensor()))
base_standard_dist = tfd.JointDistributionSequential(
     [tfd.Sample(tfd.Normal(loc=0., scale=1.), s)
      for s in flat_event_size])

对于此分布，我们应用可训练的分块下三角线性算子来诱导后验中的相关性。在线性算子内，可训练的全矩阵块表示后验两个分量之间的完全协方差，而零块（或 None）表示独立性。对角线上的块是下三角矩阵或对角矩阵，因此整个块结构表示下三角矩阵。

将此双射器应用于基本分布会导致具有均值 0 和（Cholesky 因式分解）协方差等于下三角块矩阵的多元正态分布。

operators = (
  (tf.linalg.LinearOperatorDiag,),  # Variance of uranium weight (scalar).
  (tf.linalg.LinearOperatorFullMatrix,  # Covariance between uranium and floor-by-county weights.
   tf.linalg.LinearOperatorDiag),  # Variance of floor-by-county weight (scalar).
  (None,  # Independence between uranium weight and county effects.
   None,  # Independence between floor-by-county and county effects.
   tf.linalg.LinearOperatorDiag)  # Independence among the 85 county effects.
)
block_tril_linop = (
  tfp.experimental.vi.util.build_trainable_linear_operator_block(
      operators, flat_event_size))
scale_bijector = tfb.ScaleMatvecLinearOperatorBlock(block_tril_linop)

最后，我们允许均值取非零值，方法是应用可训练的移位双射器。

loc_bijector = tfb.JointMap(
   tf.nest.map_structure(
       lambda s: tfb.Shift(
           tf.Variable(tf.random.uniform(
               (s,), minval=-2., maxval=2., dtype=tf.float32))),
       flat_event_size))

通过使用尺度和位置双射器变换标准正态分布获得的多元正态分布必须被重新整形和重构以匹配先验，最后被约束到先验的支持。

reshape_bijector = tfb.JointMap(
   tf.nest.map_structure(tfb.Reshape, flat_event_shape))
unflatten_bijector = tfb.Restructure(
       tf.nest.pack_sequence_as(
           event_shape, range(len(flat_event_shape))))
event_space_bijector = target_model.experimental_default_event_space_bijector()

现在，将它们全部组合在一起 - 将可训练的双射器链接在一起并将其应用于基本标准正态分布以构建替代后验。

surrogate_posterior = tfd.TransformedDistribution(
   base_standard_dist,
   bijector = tfb.Chain([  # Chained bijectors are applied in reverse order.
        event_space_bijector,  # Constrain to the support of the prior.
        unflatten_bijector,  # Pack components into the event_shape structure.
        reshape_bijector,  # Reshape the vector-valued components.
        loc_bijector,  # Allow for nonzero mean.
        scale_bijector  # Apply the block matrix transformation.
      ]))

训练多元正态替代后验。

optimizer = tf.optimizers.Adam(learning_rate=1e-2)
@tf.function(jit_compile=True)
def run_vi():
 return tfp.vi.fit_surrogate_posterior(
   target_model.unnormalized_log_prob,
   surrogate_posterior,
   optimizer=optimizer,
   num_steps=10**4,
   sample_size=16)
mvn_loss = run_vi()
mvn_samples = surrogate_posterior.sample(1000)

由于训练好的替代后验分布是 TFP 分布，因此我们可以从它中采样并处理这些样本以生成参数的后验 可信区间。

以下箱线图显示了两个最大县的县效应以及对土壤铀测量值和每个县的平均楼层的回归权重的 50% 和 95% 可信区间。县效应的后验可信区间表明，在考虑其他变量后，圣路易斯县的位置与较低的氡含量相关，而亨内平县的位置影响接近中性。

回归权重的后验可信区间表明，土壤铀含量越高，氡含量越高，在较高楼层进行测量（可能是因为房屋没有地下室）的县往往氡含量较高，这可能与土壤性质及其对建筑类型的影响有关。

楼层的（确定性）系数为 -0.7，表明较低楼层氡含量较高，符合预期。

逆自回归流代理后验

逆自回归流（IAF）是利用神经网络捕捉分布分量之间复杂非线性依赖关系的归一化流。接下来，我们构建一个 IAF 代理后验，以查看这种更高容量、更灵活的模型是否优于受限多元正态分布。

我们首先构建一个具有向量事件形状的标准正态分布，其长度等于后验中自由度的总数。

base_distribution = tfd.Sample(
   tfd.Normal(loc=0., scale=1.),
   sample_shape=[tf.reduce_sum(flat_event_size)])

一个可训练的 IAF 转换正态分布。

num_iafs = 2
iaf_bijectors = [
   tfb.Invert(tfb.MaskedAutoregressiveFlow(
       shift_and_log_scale_fn=tfb.AutoregressiveNetwork(
           params=2,
           hidden_units=[256, 256],
           activation='relu')))
   for _ in range(num_iafs)
]

IAF 双射器与其他双射器一起链接，构建一个与先验具有相同事件形状和支撑的代理后验。

iaf_surrogate_posterior = tfd.TransformedDistribution(
   base_distribution,
   bijector=tfb.Chain([
        event_space_bijector,  # Constrain to the support of the prior.
        unflatten_bijector,  # Pack components into the event_shape structure.
        reshape_bijector,  # Reshape the vector-valued components.
        tfb.Split(flat_event_size),  # Split into parts, same size as prior.
   ] + iaf_bijectors))  # Apply a flow model.

与多元正态代理后验类似，IAF 代理后验使用 tfp.vi.fit_surrogate_posterior 进行训练。IAF 代理后验的可信区间与受限多元正态分布的可信区间相似。

均值场代理后验

VI 代理后验通常假设为均值场（独立）正态分布，具有可训练的均值和方差，这些均值和方差通过双射变换约束到先验的支持范围。除了两个更具表现力的代理后验之外，我们还定义了一个均值场代理后验，使用与多元正态代理后验相同的通用公式。我们不使用分块下三角线性算子，而是使用分块对角线性算子，其中每个块都是对角线。

operators = (
  tf.linalg.LinearOperatorDiag,
  tf.linalg.LinearOperatorDiag,
  tf.linalg.LinearOperatorDiag,
)
block_diag_linop = (
   tfp.experimental.vi.util.build_trainable_linear_operator_block(
       operators, flat_event_size))

在这种情况下，均值场代理后验给出了与更具表现力的代理后验类似的结果，表明这种更简单的模型可能足以完成推理任务。作为“真实值”，我们还使用哈密顿蒙特卡罗方法进行采样（有关完整示例，请参见 Colab）。所有三个代理后验都生成了与 HMC 采样在视觉上相似的可信区间，尽管由于 ELBO 损失的影响，有时会发生分散不足，这在 VI 中很常见。

结论

在这篇文章中，我们使用联合分布和多部分双射器构建了 VI 代理后验，并对其进行拟合，以估计氡数据集回归模型中权重的可信区间。对于这个简单的模型，更具表现力的代理后验似乎与均值场代理后验的表现类似。但是，我们演示的工具可用于构建各种灵活的代理后验，适用于更复杂的模型。

请查看我们的代码、文档和更多示例，请访问 TFP 主页。