引言
本文通过对一个最简单的线性拟合的例子，做包括点估计，蒙卡估计和变分推断在内的不同分析，来展示不同统计分析方案如何统一在相同的贝叶斯推断的形式体系之内，并且以此展示 TensorFlow Probability 库的一些基本用法。一但了解了统计推断领域的基本概念，类似的库的用法都是相通的，包括但不限于 Stan，PyMC3，Pyro，Edward2，Probtorch 等。这些库的使用，困难主要是来自理论上的，而不是 API 上的，如果对于统计推断的整体概念没有清晰的认知，那么看上半天这些库的 API 文档，也一样不会用。
本文的灵感，研究问题，和下述方法二与方法三的实现，部分参考了1。本文对其方法二和方法三的模型进行了调整，错误进行了修改，并给出了详细的解释，使得贝叶斯推断的形式体系可以和最终的程序实现无缝衔接。本文额外增添并独立给出了方法一和方法四的实现。
模型与数据
我们考虑一个最简单的线性拟合模型，给出一系列的点 \((\mathbf{x}, y)\)，找到权重向量 \(\mathbf{w}\) 和偏置 \(b\)，使得对应的超平面 \(\hat{y}=\mathbf{w}\cdot\mathbf{x}+b\) 最接近数据 \(y\)。用更统计的说法，这个模型是这样的。\(\mathbf{x}\sim N(\mathbf{0},\mathbf{I})\) （稍后我们会看到，x 变量的边缘分布怎样取均可，和要研究的问题解耦），\(y\sim N(\mathbf{w}\cdot\mathbf{x}+b, \sigma)\)，我们则需要寻找 \(\mathbf{w},b,\sigma\) 这些参数的最优估计甚至其分布情况。
首先我们导入相应的库并生成待使用的数据。本文的全部实现，都基于 TF 2.0。
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability import edward2 as ed
tfd = tfp.distributions
tf.__version__, tfp.__version__
#('2.0.0-dev20190731', '0.8.0')
D = 4 #number of dimensions
N = 2000 #number of datapoints to generate
noise_std_true = 0.5
# Generate data
b_true = np.random.randn(1).astype(np.float32) #bias
w_true = np.random.randn(D, 1).astype(np.float32) #weights
x = np.random.randn(N, D).astype(np.float32)
noise = noise_std_true * np.random.randn(N,1).astype(np.float32)
y = np.matmul(x, w_true) + b_true + noise
以上产生数据的代码很简单，这里不多做解释了，我们可以可视化了解下数据的大体情况。
fig, axes = plt.subplots(int(np.ceil(D/2)), 2, sharex=True)
fig.set_size_inches(6.4, 6)
plt.subplots_adjust(wspace=0.3, hspace=0.2)
for i in range(D):
t_ax = axes[int(i/2), i%2] #this axis
sns.regplot(x[:,i], y[:,0], ax=t_ax)
t_ax.set_ylabel('y')
t_ax.set_xlabel('x[%d]'%(i+1))
plt.show()
可以看到在每个轴上的斜率和w_true的各个分量一致。
打印出生成数据真实的斜率和截距与标准差。下文我们将通过不同的方案，从数据点来推断出以下数值。
w_true, b_true, noise_std_true
# (array([[-0.5026519 ],[ 0.15551816],[ 0.76468086],[-0.07409738]], dtype=float32),
# array([0.43074715], dtype=float32),
# 0.5)
为了之后训练方便，我们将 np.array 的数据点，转化为 Dataset 对象。tf.data 的 API 提供了方便的一键 shuffle，一键 batch，并且 data_train 可以直接迭代。
data_train = tf.data.Dataset.from_tensor_slices(
(x, y)).shuffle(10000).batch(100)
这样当我们之后的代码使用 x，y 时，对应的是全体数据，而通过 for x_train, y_train in data_train 拿到的则是 batch size 为 100 的数据。
方法一：点估计
最基本的对于斜率和截距的估计，并不需要用到统计的信息。所谓最小二乘法，则是直接最小化 \(L(\mathbf{k}, b)=\sum_i(\hat{y}_i(\mathbf{x}_i, \mathbf{k}, b)-y_i)^2\)。也即使得所有真实点离预测点的距离平方和最小，用一句机器学习领域的“黑话”，就是损失函数是 MSE (Mean Square Error)。这一问题可以有解析解，不过来都来了，也顺手用机器学习的范式解决了。其对应的实现如下。
class LinearRegression(tf.keras.Model):
def __init__(self, d, name=None):
super().__init__(name=name)
self.w_loc = tf.Variable(tf.random.normal([d, 1]), name='w_loc')
self.b_loc = tf.Variable(tf.random.normal([1]), name='b_loc')
def call(self, x):
y = x@self.w_loc+self.b_loc
return y
mse = tf.keras.losses.MeanSquaredError()
lrmodel = LinearRegression(D)
adamopt = tf.keras.optimizers.Adam(learning_rate=1e-3)
@tf.function
def train_step(x_data, y_data):
with tf.GradientTape() as tape:
loss = mse(y_data, lrmodel(x_data))
gradients = tape.gradient(loss, lrmodel.trainable_variables)
adamopt.apply_gradients(zip(gradients, lrmodel.trainable_variables))
return loss
训练过程为：
for epoch in range(2000):
for x_data, y_data in data_train:
train_step(x_data, y_data)
adamopt = tf.keras.optimizers.Adam(learning_rate=1e-4)
l = []
for epoch in range(500):
for x_data, y_data in data_train:
l.append(train_step(x_data, y_data))
注意上面训练过程，待参数基本稳定后，我们又继续跑了几百次，用来收集每轮损失函数的值，马上我们会看到其作用。
训练收敛之后可以查看找到的参数。
lrmodel.weights
"""
[<tf.Variable 'w_loc:0' shape=(4, 1) dtype=float32, numpy=
array([[-0.49829516],
[ 0.14661531],
[ 0.75829834],
[-0.08185098]], dtype=float32)>,
<tf.Variable 'b_loc:0' shape=(1,) dtype=float32, numpy=array([0.40996137], dtype=float32)>]
"""
可以看到相应的参数与真值基本吻合。
同样我们可以计算训练中最低损失函数的值，也即
np.sqrt(np.mean(l))
# 0.48789
因为我们用的损失函数是 MSE，这一值开方恰好刻画了从 \(y\) 到 \(\hat{y}\) 的标准差，此数值恰好对应了noise_std_true。
之所以这种方法叫做点估计，是由于我们只给出了几个待定参数的估计值，而对于这些估计值可能的误差或者分布一无所知。要想能够给出每个参数的分布，也即计算 \(P(\mathbf{k}\vert \mathbf{x}, y)\) 等，我们则需要引入贝叶斯推断。下一部分，是关于贝叶斯推断的数学，只会讲到理解代码需要的部分。
一些贝叶斯数学
计算参数也即隐变量的分布，实际上我们需要计算的就是后验分布
\[P(\mathbf{k}, b, \sigma\vert \mathbf{x},y)=\frac{P(\mathbf{x},y\vert \mathbf{k},b, \sigma)P(\mathbf{k})P(b)P(\sigma)}{\sum_{\mathbf{x},y} P(\mathbf{x},y\vert \mathbf{k},b, \sigma)P(\mathbf{k})P(b)P(\sigma)}\\= \frac{P(y\vert \mathbf{x},\mathbf{k},b, \sigma)P(\mathbf{x}\vert \mathbf{k},b, \sigma)P(\mathbf{k})P(b)P(\sigma)}{\sum_{\mathbf{x},y} P(\mathbf{x},y\vert \mathbf{k},b, \sigma)P(\mathbf{k})P(b)P(\sigma)}. \label{Bayes}\]后验一般来讲，分母当然是 intractable。分子中 \(\mathbf{x}\vert \mathbf{k},b,\sigma=\mathbf{x}\) 与这些参数无关是一个常数。注意到反正分子也不是归一化的，那么任何常数可以直接在分析问题时直接扔掉。
需要在此强调的是，\(\eqref{Bayes}\) 中的 \(\mathbf{x}, y\) 都指数据整体，也即 \(y=(y_1…y_N)\)。如果假设各个数据点的独立性，我们对于分子进一步有：
\[P(y\vert \mathbf{x},\mathbf{k},b, \sigma)P(\mathbf{x}\vert \mathbf{k},b, \sigma)P(\mathbf{k})P(b)P(\sigma)\propto \prod_i^N P(y_i\vert \mathbf{x_i},\mathbf{k},b, \sigma)P(\mathbf{k})P(b)P(\sigma)\\ =\prod_i^D P_{N(\mathbf{k}\mathbf{x_i}+b)}(y_i)P_{N(0,1)}(\mathbf{k})P_{N(0,1)}(b)P_{N(0.5,0.01)}(\sigma). \label{mco}\]第二行我们已经带入了对于模型和先验的假设。这里要注意将标准差参数的先验假设为高斯分布是非常不好的实践，首先，高斯分布可以给出负数，但标准差只能是正的。一般来讲，标准差的先验大多选取 HalfCauchy，HalfNormal，InverseGamma 等只在正数上的分布。我们这里选取了一个标准差极小的高斯分布做先验，其意思就是说，我们在之前就对该模型的标准差很了解了。这里这么做主要是为了后面模型处理的方便。tfp 中没有上述其他分布和高斯分布的 KL 散度解析表达式，会复杂化方法三的实现。鉴于本文尽量将实现平行于理论形式体系，因此不去展示更复杂的分布时 VI 怎么调整。因此本文仅仅是个 demo，千万别在任何严肃的场合，用这样的高斯分布做标准差系数的先验。
对于 MCMC，由于 Metropolis-Hasting 算法只需要知道两点概率之比即可模拟分布，而不需归一化，因为 \(\eqref{Bayes}\) 分子各项均已知，且 \(\mathbf{x}, y\) 整体作为数据给出，那么可以通过 MCMC 来得到按一定概率（后验）分布的各参数的 sample。由这些 sample 可以进行参数值和方差的估计，并给出分布。
另一种近似后验的方案，则是变分推断，假设一个后验分布的函数形式 \(q_{\theta}(\mathbf{k},b,\sigma)\)，我们只需最小化 ELBO 即可。也即，我们用 z 指代全部隐变量 \(\mathbf{k}, b,\sigma\)。
\[ELBO = E_q[\ln q(z)-\ln P(x,z)]\\=E_{q}[\ln q(z) -\ln P(x\vert z)-\ln P(z)]. \label{ELBO}\]这一式子根据需要可以看成前两项的 KL 减去第三项，也可以看成第一项和第三项的 KL 减去第二项。再次强调， \(\eqref{ELBO}\) 中提到的 x 也代表了数据全体，而不是指某个数据点，实际上这里 x 表示了全部的 \((\mathbf{x}, y)\)。考虑到后验假设的和先验模型中的统计独立性和 \(\mathbf{x}\) 边缘分布是个常数，我们有：
\[ELBO=E_{q(\mathbf{k},b,\sigma)} [\ln q(\mathbf{k})+\ln q(b)+\ln q(\sigma)-\ln p(\mathbf{k})+\ln p(b)+\ln p(\sigma)-\sum_i^N P(y_i\vert \mathbf{x_i},\mathbf{k},b,\sigma)]. \label{vio}\]其中前三项是隐变量的后验假设，后三项是隐变量的先验分布，而最后一项拆成了独立的各个可能性条件概率之和。请确保真的看懂了上面的式子，比如请回答以下问题，既然 \(\mathbf{x}\) 是数据而不是隐变量，为什么最后一项不是原始的 \(P(\mathbf{x_i},y_i\vert \mathbf{k},b, \sigma)\)？
小结：\(\eqref{mco}\) 是对该模型做蒙卡估计的基础， \(\eqref{vio}\) 是对该模型做变分推断估计的基础。
方法二： MCMC
这里我们选择 Hamiltonian Monte Carlo 方法进行蒙卡取样。该方法在连续变量空间的取样效果远远优于普通的 Metropolis-Hasting 方法。（当然 HMC 本身也是 M-H 算法的一种，只不过是 propose 分布考虑了本来分布的一些性质，关于 HMC 的具体介绍，可以参考2）。
根据 \(\eqref{mco}\) 可知，为了实现蒙卡模拟，只需要告诉取样器分布的联合概率即可。但我们这里进一步偷一个懒，先定义生成 \(y\) 的概率模型，再通过 Edward 自动生成联合概率函数（函数变量自动就是概率模型函数的自变量和函数内部出现的全部随机变量）。需要注意在概率编程的框架里，对于写下的每个分布和取样，需要始终对其指的是先验，后验还是可能性保持高度的警觉，绝不可混淆。
def generate_model(features):
D = features.shape[1] #feature is x
coeffs = ed.Normal( #normal prior on w
loc=tf.zeros([D,1]),
scale=tf.ones([D,1]),
name="coeffs")
bias = ed.Normal( #normal prior on b
loc=tf.zeros([1]),
scale=tf.ones([1]),
name="bias")
noise_std = ed.Normal( #normal prior on sigma
loc = tf.constant([0.5]),
scale = tf.constant([0.01]),
name = "noise_std")
predictions = ed.Normal( #prediction is y, the distribution is P(y|x,k,b,sigma)
loc=features@coeffs+bias,
scale=noise_std,
name="predictions")
return predictions
log_joint = ed.make_log_joint_fn(generate_model)
mc_log_prob_fn = lambda coeffs, bias, noise_std: log_joint(
features=x,
coeffs=coeffs,
bias=bias,
noise_std=noise_std,
predictions=y)
我们看到通过 ed.make_log_joint_fn，edward 可以将描述生成概率模型的函数的联合概率的对数函数自动生成出来。我们蒙卡需要的拟合分布是关于隐变量的，因此 features 和 predictions 是固定的，于是我们构造了蒙卡专用的偏函数。
蒙卡模拟本身，tfp 封装的已经很好了，懂了原理，直接套 API 就好了。
num_results = int(8e3) #number of hmc iterations
n_burnin = int(2e3) #number of burn-in steps
step_size = 0.002
num_leapfrog_steps = 5
kernel = tfp.mcmc.HamiltonianMonteCarlo(
target_log_prob_fn=mc_log_prob_fn,
step_size=step_size,
num_leapfrog_steps=num_leapfrog_steps)
states, kernel_results = tfp.mcmc.sample_chain(
num_results=num_results,
num_burnin_steps=n_burnin,
kernel=kernel,
current_state=[
tf.zeros([D, 1]),
tf.ones([1]),
tf.constant([0.5])
]
)
这里直接用了 HMC 的采样 kernel，你也可以尝试 tfp 自带的其他蒙卡取样器，比如更高级的 NUTS NoUTurnSampler 等，其好处是需要调节的参数更少。像这里的 HMC，调节参数 step_size 是必须的。每次跑完取样，需要看一下 kernel_results.is_accepted 来判断下接受率，如果接受率太低是有问题的，说明步长太大了。纯理论上，由于能量守恒，其接受率应该严格逼近 1。如果步长过大，则取样的接受率会很低甚至是0，由此造成取样完全失败。有报道说 HMC 的理想接受率应该在 60% 到 80% 左右3，不过我有点怀疑这样的接受率也有点低了。蒙卡的真实模拟发生在 mcmc.sample_chain 函数，不过这一过程慢到惊人，而且提供的 parallel_iteration 的参数也并没有用，还是占一个 CPU core 在算。
计算完成之后，每一次的参数值（包括 burnin）都存在 states list 中，而 kernel_results 则记录一些模拟过程的其他信息。让我们观察一下参数模拟的跑动情况。
coeffs, bias, noise_std = states
plt.figure()
for i in range(D):
plt.plot(coeffs[:,i], label='W[{}]'.format(i))
plt.hlines(w_true[i], 0, num_results, linestyle='--')
plt.plot(bias, label='b')
plt.hlines(b_true, 0, 3000, linestyle='--')
plt.plot(noise_std, label='noise_std')
plt.hlines(noise_std_true, 0, 3000, linestyle='--')
plt.title('Weight parameters over HMC')
plt.legend()
plt.show()
plt.close()
我们可以看到，参数几乎从第一次更新开始，就马上来到了基本上正确的位置，这充分显示了 HMC 取样的高效性。我们也可以看一下蒙卡模拟出的这些参数在已知 x，y 情况下的后验概率分布。
fig, axes = plt.subplots(6,1)
fig.set_size_inches(3.5,9)
plt.subplots_adjust(wspace=0.3, hspace=0.5)
for i in range(D):
sns.kdeplot(np.array(coeffs[n_burnin:,i]).reshape(-1), ax=axes[i], shade=True)
axes[i].axvline(w_true[i], linestyle="--")
axes[i].set_ylabel(f"W_{i+1}")
sns.kdeplot(np.array(bias[n_burnin:]).reshape(-1), ax=axes[4], shade=True)
axes[4].axvline(b_true, linestyle="--")
axes[4].set_ylabel("b")
sns.kdeplot(np.array(noise_std[n_burnin:]).reshape(-1), ax=axes[5], shade=True)
axes[5].axvline(noise_std_true, linestyle="--")
axes[5].set_ylabel("noise_std")
plt.show()
plt.close()
方法三：变分推断微观视角
上面我们看到，MCMC 模拟的方式来找后验实在是太慢了（大部分也可能是纯 python 实现的锅）。这里我们改用变分推断估计。在方法三，我们首先按照 \(\eqref{vio}\) 来手动构造模型，使其损失函数为该 ELBO，从而实现参数优化。这里我们假设的所有隐变量的后验分布都是高斯分布，控制隐变量分布的参数组自然是期望和标准差（但是 noise_std 我们假设后验标准差也为 0.01 固定，想法上这样当然不好，但是实践上会避免很多问题，因为本文不是讲解 VI 严格的最佳实践，因此暂时用这样的方式使得噪声参数比较容易优化，只需要记住本文对于噪声项的先验和后验处理的都有点畸形）。我们先放实现，再做分析。
class BayesianLinearRegression(tf.keras.Model):
def __init__(self, d, name=None):
super(BayesianLinearRegression, self).__init__(name=name)
self.w_loc = tf.Variable(tf.zeros([d,1]), name='w_loc')
self.w_std = tf.Variable(tf.zeros([d,1]), name='w_std')
self.b_loc = tf.Variable([0.0], name='b_loc')
self.b_std = tf.Variable([0.0], name='b_std')
self.noise_std = tf.Variable([-0.7], name='noise_std')
@property
def weight(self):
"""Variational posterior for the weight"""
return tfd.Normal(self.w_loc, tf.exp(self.w_std))
@property
def bias(self):
"""Variational posterior for the bias"""
return tfd.Normal(self.b_loc, tf.exp(self.b_std))
@property
def std(self):
"""Variational posterior for the noise standard deviation"""
return tfd.Normal(tf.exp(self.noise_std), scale=[0.01])
def call(self, x, sampling=True):
"""Predict p(y|x,k,b,sigma) with k,b,sigma determined by posterior q"""
sample = lambda x: x.sample() if sampling else x.mean()
loc = x @ sample(self.weight) + sample(self.bias)
std = tf.sqrt(sample(self.std))
return tfd.Normal(loc, std)
@property
def losses(self):
"""Sum of KL divergences between posteriors and priors"""
prior = tfd.Normal(0, 1) #prior for w and b
prior2 = tfd.Normal(0.5, 0.01) #prior for sigma
return (tf.reduce_sum(tfd.kl_divergence(self.weight, prior)) +
tf.reduce_sum(tfd.kl_divergence(self.bias, prior)) +
tf.reduce_sum(tfd.kl_divergence(self.std, prior2))
)
blrmodel = BayesianLinearRegression(D)
adamopt = tf.keras.optimizers.Adam(lr=10e-3)
@tf.function
def vi_train_step(x_data, y_data):
with tf.GradientTape() as tape:
log_prob = tf.reduce_mean(blrmodel(x_data).log_prob(y_data))
kl_loss = blrmodel.losses/N
elbo_loss = kl_loss - log_prob
gradients = tape.gradient(elbo_loss, blrmodel.trainable_variables)
adamopt.apply_gradients(zip(gradients, blrmodel.trainable_variables))
for epoch in range(3000):
for x_data, y_data in data_train:
vi_train_step(x_data, y_data)
训练部分和模型构造没什么说的，但一定要注意区分好每个分布是先验还是后验还是可能性，这真的太重要了，再强调也不为过。参考1省略了 \(\sigma\) 项先验和后验之间的 KL 散度，导致损失函数缺了一项，我们补回来了。这也是没办法的，因为按照传统套路把 \(\sigma\) 的先验后验假设成 InverseGamma 分布之流，tfp 根本无法解析计算两分布间的 KL 散度，因为没有实现，这也是本文强行将 \(\sigma\) 的先验后验都假设称 Gaussian 的原因。当然，也可以进一步微操，把 KL 散度的损失项，也用取样近似而不是完整的分布来解析的算，这就不在本文的简介范围内。
继续看损失函数，KL 散度项除以了 N，这是因为 KL 散度这里是解析给出的，直接是总样本的值，而我们分析 \(E_q(\ln P(y\vert \mathbf{x},\mathbf{k},b,\sigma))\) 时，用了 mini-batch 和 tf.reduce_mean 相当于单个数据点的值。我们也可以对 log_prob 用 tf_reduce_sum，对应 kl_loss要引入*batch_size会使得损失函数依赖更多参数。kl_loss 项比较好理解，因为是直接带的公式。log_prob项对应的是通过 \(q(\mathbf{k}, b, \sigma)\) 取样产生的参数下的可能性，这里边为什么隐变量是从后验近似取样，为什么数据 x 出现在了可能性的条件位置成了隐变量的一部分，这里实现只是做的是 MC sample size 为 1 的蒙卡的巧合等问题，都要了解清楚，才能之后进一步魔改模型。比如，如何实现每个数据点用 10 个后验生成的隐变量来平均？
训练之后的分布，可以直接通过 blrmodel.bias 等等来获得。参数也可以通过 blrmodel.weights 来查看。训练结束参数后验分布的可视化，我们在讲完方法四后一起比较。
方法四：变分推断宏观视角
方法三虽然实现的是变分推断，但是完全绕回了神经网络的语言，只不过损失函数构造的是 ELBO 而已。有时我们想进一步平行变分推断自身的逻辑。也就是我们需要这样一种接口，告诉程序生成模型 model（间接告诉了先验和可能性）以及假设的带参数后验分布 guide，可以直接调用某个 API 来进行模拟，从而给出最优化的后验近似。这也是 Pyro 的核心语言构造，即 pyro.infer.SVI。tfp 中也有类似的 API，也即 tfp.vi.fit_surrogate_posterior。这种 API 只需要告诉其模型的联合分布（类似蒙卡 API 的变量）和假设的后验分布族即可无脑用。而其中的损失函数构造，蒙卡近似与优化等细节都被封装了起来，用起来更省心，使用思维上也更吻合变分推断的理论形式体系。
我们首先需要构造含参的后验概率分布，正如名字暗示的，就是构造一个分布。
defer = tfp.util.DeferredTensor
#posterior for k,b,sigma
q = tfd.JointDistributionNamed({
"coeffs": tfd.Normal(tf.Variable(np.zeros([D,1]), name="w_loc", dtype=tf.float32),
defer(tf.exp,tf.Variable(np.zeros([D,1]),name="w_std", dtype=tf.float32))),
"bias": tfd.Normal(tf.Variable([0.], name="b_loc", dtype=tf.float32),
defer(tf.exp,tf.Variable(0., name="b_std", dtype=tf.float32))),
"noise_std": tfd.Normal(loc=defer(tf.exp, tf.Variable([-0.7], name="noise_std", dtype=tf.float32)),
scale=0.01)
})
注意我们这里用到了 tfp.util.DeferredTensor 防止 exp 在运行之前就作用，使得该部分的导数无法被记录（exp 本身的作用是防止标准差参数变成负数）。
我们还需要在包装一下联合分布函数，因为 JointDistributionNamed 做 sample 方法，默认返回字典。而在fit_surrogate_posterior中默认将 sample 后的变量作为自变量传入联合分布函数，字典时按照 **dict 传入，因此有：
def vi_log_prob_fn(**param):
coeffs, bias, noise_std = param["coeffs"], param["bias"], param["noise_std"]
coeffs=tf.reshape(coeffs,[-1,1])
bias=tf.reshape(bias,[-1])
noise_std=tf.reshape(noise_std,[-1])
return log_joint(
features=x,
coeffs=coeffs,
bias=bias,
noise_std=noise_std,
predictions=y)
这里还有一个 reshape，这里涉及的问题时，vi 引擎默认是对 q 做 sample(1) 方法而不是 sample()，两者返回的随机变量值的 shape 是不一样的，前者多出了一维，这会造成和 log_joint 的自变量 shape 不匹配，因此需要把第一维 num_sample 压掉。当然 vi 引擎也可以每次 sample 多个隐变量取样值算概率，本文就暂不扩展了。还有个问题，就是在方法四的框架下边，添加 mini-batch 支持，似乎是比较复杂的。
最后可以开始开心的训练了，一个函数搞定，比微操来自己实现 VI 轻松多了。
hist_loss = tfp.vi.fit_surrogate_posterior(
target_log_prob_fn=vi_log_prob_fn,
surrogate_posterior=q,
optimizer=tf.optimizers.Adam(learning_rate=0.0001),
num_steps=20000,
)
该函数的返回值是每次迭代的 loss function。这边优化过的分布和参数情况可以通过 q.model['noise_std'] 等查看。
最后我们可视化比较一下 MCMC 和两种 VI 得出的后验分布的区别 (竖直虚线代表真实值)。
fig, axes = plt.subplots(6,1)
fig.set_size_inches(4,16)
plt.subplots_adjust(wspace=0.3, hspace=0.4)
vi_mic = blrmodel.weight.sample(6000)
vi_mac = q.model['coeffs'].sample(6000)
for i in range(D):
sns.kdeplot(np.array(vi_mic[:,i]).reshape(-1), ax=axes[i], label="VI-micro", shade=True)
sns.kdeplot(np.array(vi_mac[:,i]).reshape(-1),ax=axes[i], label="VI-macro", shade=True)
sns.kdeplot(np.array(coeffs[n_burnin:,i]).reshape(-1), ax=axes[i], label="MCMC", shade=True)
axes[i].axvline(w_true[i],linestyle="--")
axes[i].set_ylabel(f"W_{i+1}")
axes[i].legend()
sns.kdeplot(np.array(blrmodel.bias.sample(6000)).reshape(-1), ax=axes[4], label="VI-micro", shade=True)
sns.kdeplot(np.array(q.model['bias'].sample(6000)).reshape(-1),ax=axes[4], label="VI-macro", shade=True)
sns.kdeplot(np.array(bias[n_burnin:]).reshape(-1), ax=axes[4], label="MCMC", shade=True)
axes[4].axvline(b_true,linestyle="--")
axes[4].legend()
axes[4].set_ylabel("b")
sns.kdeplot(np.array(blrmodel.std.sample(6000)).reshape(-1), ax=axes[5], label="VI-micro", shade=True)
sns.kdeplot(np.array(q.model['noise_std'].sample(6000)).reshape(-1),ax=axes[5], label="VI-macro", shade=True)
sns.kdeplot(np.array(noise_std[n_burnin:]).reshape(-1), ax=axes[5], label="MCMC", shade=True)
axes[5].axvline(noise_std_true,linestyle="--")
axes[5].legend()
axes[5].set_ylabel("b")
plt.show()
plt.close()
以上，我们以 tfp 为例展示了概率编程的一些基本范式及其和理论形式体系的一一映射。本文只是为了建立起这种框架，更多实现上的细节，比如增添 VI 中每个样本蒙卡取样的次数，对 VI API 增加 mini-batch 支持，使用更真实的先验后验来模拟标准差参数的分布等等，本文都暂且不讨论。