变分推断：贝叶斯神经网络#

PyMC中的贝叶斯神经网络#

生成数据#

首先，让我们生成一些示例数据——一个不是线性可分的简单二元分类问题。

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pymc as pm
import pytensor
import seaborn as sns

from sklearn.datasets import make_moons
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import scale

%config InlineBackend.figure_format = 'retina'
floatX = pytensor.config.floatX
RANDOM_SEED = 9927
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-darkgrid")

X, Y = make_moons(noise=0.2, random_state=0, n_samples=1000)
X = scale(X)
X = X.astype(floatX)
Y = Y.astype(floatX)
X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.5)

fig, ax = plt.subplots()
ax.scatter(X[Y == 0, 0], X[Y == 0, 1], color="C0", label="Class 0")
ax.scatter(X[Y == 1, 0], X[Y == 1, 1], color="C1", label="Class 1")
sns.despine()
ax.legend()
ax.set(xlabel="X1", ylabel="X2", title="Toy binary classification data set");
../../../_images/cad6382102331ff90c2c382f63af66a7275b32896cbbb27e8ff1d4210c4118a6.png

模型规范#

神经网络其实相当简单。基本单元是一个感知器，它只不过是逻辑回归。我们并行使用许多这样的感知器，然后将它们堆叠起来以获得隐藏层。这里我们将使用2个隐藏层，每层5个神经元，这对于这样一个简单的问题已经足够了。

def construct_nn(ann_input, ann_output):
    n_hidden = 5

    # Initialize random weights between each layer
    init_1 = rng.standard_normal(size=(X_train.shape[1], n_hidden)).astype(floatX)
    init_2 = rng.standard_normal(size=(n_hidden, n_hidden)).astype(floatX)
    init_out = rng.standard_normal(size=n_hidden).astype(floatX)

    coords = {
        "hidden_layer_1": np.arange(n_hidden),
        "hidden_layer_2": np.arange(n_hidden),
        "train_cols": np.arange(X_train.shape[1]),
        "obs_id": np.arange(X_train.shape[0]),
    }
    with pm.Model(coords=coords) as neural_network:
        ann_input = pm.Data("ann_input", X_train, dims=("obs_id", "train_cols"))
        ann_output = pm.Data("ann_output", Y_train, dims="obs_id")

        # Weights from input to hidden layer
        weights_in_1 = pm.Normal(
            "w_in_1", 0, sigma=1, initval=init_1, dims=("train_cols", "hidden_layer_1")
        )

        # Weights from 1st to 2nd layer
        weights_1_2 = pm.Normal(
            "w_1_2", 0, sigma=1, initval=init_2, dims=("hidden_layer_1", "hidden_layer_2")
        )

        # Weights from hidden layer to output
        weights_2_out = pm.Normal("w_2_out", 0, sigma=1, initval=init_out, dims="hidden_layer_2")

        # Build neural-network using tanh activation function
        act_1 = pm.math.tanh(pm.math.dot(ann_input, weights_in_1))
        act_2 = pm.math.tanh(pm.math.dot(act_1, weights_1_2))
        act_out = pm.math.sigmoid(pm.math.dot(act_2, weights_2_out))

        # Binary classification -> Bernoulli likelihood
        out = pm.Bernoulli(
            "out",
            act_out,
            observed=ann_output,
            total_size=Y_train.shape[0],  # IMPORTANT for minibatches
            dims="obs_id",
        )
    return neural_network


neural_network = construct_nn(X_train, Y_train)

这并不算太糟。Normal 先验有助于正则化权重。通常我们会向输入中添加一个常数 b，但为了保持代码更简洁，我在这里省略了它。

变分推断：扩展模型复杂度#

我们现在可以运行一个MCMC采样器，比如pymc.NUTS，在这种情况下效果很好，但正如已经提到的，当我们把模型扩展到更深的架构和更多层时，这会变得非常慢。

相反，我们将使用pymc.ADVI变分推断算法。这会快得多，并且扩展性更好。请注意，这是一个均值场近似，因此我们忽略了后验中的相关性。

%%time

with neural_network:
    approx = pm.fit(n=30_000)





Finished [100%]: Average Loss = 138.57





CPU times: user 8.64 s, sys: 229 ms, total: 8.87 s
Wall time: 9.41 s









绘制目标函数（ELBO），我们可以看到优化过程逐步改进了拟合效果。






plt.plot(approx.hist, alpha=0.3)
plt.ylabel("ELBO")
plt.xlabel("iteration");








../../../_images/be1dab9d23db079a95d711fd6539413a604759e799a30fc6598650d00fd989d8.png









trace = approx.sample(draws=5000)









现在我们已经训练了我们的模型，让我们使用后验预测检查（PPC）在保留集上进行预测。我们可以使用sample_posterior_predictive()从后验（从变分估计中采样）生成新数据（在这种情况下是类别预测）。






with neural_network:
    pm.set_data(new_data={"ann_input": X_test})
    ppc = pm.sample_posterior_predictive(trace)
    trace.extend(ppc)









Sampling: [out]













我们可以对每个观测值的预测结果进行平均，以估计类别1的潜在概率。






pred = ppc.posterior_predictive["out"].mean(("chain", "draw")) > 0.5













fig, ax = plt.subplots()
ax.scatter(X_test[pred == 0, 0], X_test[pred == 0, 1], color="C0", label="Predicted 0")
ax.scatter(X_test[pred == 1, 0], X_test[pred == 1, 1], color="C1", label="Predicted 1")
sns.despine()
ax.legend()
ax.set(title="Predicted labels in testing set", xlabel="X1", ylabel="X2");








../../../_images/715d304f781aef165782096c9c5c9186733bda5bef2ac08b0505d75f952532ff.png









print(f"Accuracy = {(Y_test == pred.values).mean() * 100:.2f}%")









Accuracy = 96.20%









嘿，我们的神经网络做得不错！








让我们看看分类器学到了什么#


为此，我们在整个输入空间上对网格上的类概率预测进行评估。






grid = pm.floatX(np.mgrid[-3:3:100j, -3:3:100j])
grid_2d = grid.reshape(2, -1).T
dummy_out = np.ones(grid_2d.shape[0], dtype=np.int8)













coords_eval = {
    "train_cols": np.arange(grid_2d.shape[1]),
    "obs_id": np.arange(grid_2d.shape[0]),
}

with neural_network:
    pm.set_data(new_data={"ann_input": grid_2d, "ann_output": dummy_out}, coords=coords_eval)
    ppc = pm.sample_posterior_predictive(trace)









Sampling: [out]

















y_pred = ppc.posterior_predictive["out"]











概率表面#






cmap = sns.diverging_palette(250, 12, s=85, l=25, as_cmap=True)
fig, ax = plt.subplots(figsize=(16, 9))
contour = ax.contourf(
    grid[0], grid[1], y_pred.mean(("chain", "draw")).values.reshape(100, 100), cmap=cmap
)
ax.scatter(X_test[pred == 0, 0], X_test[pred == 0, 1], color="C0")
ax.scatter(X_test[pred == 1, 0], X_test[pred == 1, 1], color="C1")
cbar = plt.colorbar(contour, ax=ax)
_ = ax.set(xlim=(-3, 3), ylim=(-3, 3), xlabel="X1", ylabel="X2")
cbar.ax.set_ylabel("Posterior predictive mean probability of class label = 0");








../../../_images/59ddf6dffee4b5bc2ff4043d14782f83e6da5062d1a0cf843dad8dc9359e30db.png









预测值的不确定性#


请注意，我们可以使用非贝叶斯神经网络完成上述所有操作。每个类别标签的后验预测均值应与最大似然预测值相同。然而，我们还可以查看后验预测的标准差，以了解我们预测中的不确定性。以下是这种情况的示意图：






cmap = sns.cubehelix_palette(light=1, as_cmap=True)
fig, ax = plt.subplots(figsize=(16, 9))
contour = ax.contourf(
    grid[0], grid[1], y_pred.squeeze().values.std(axis=0).reshape(100, 100), cmap=cmap
)
ax.scatter(X_test[pred == 0, 0], X_test[pred == 0, 1], color="C0")
ax.scatter(X_test[pred == 1, 0], X_test[pred == 1, 1], color="C1")
cbar = plt.colorbar(contour, ax=ax)
_ = ax.set(xlim=(-3, 3), ylim=(-3, 3), xlabel="X1", ylabel="X2")
cbar.ax.set_ylabel("Uncertainty (posterior predictive standard deviation)");








../../../_images/2b74ec6c0dd5e52fbe989d7ec0d6dce9a704dfd43ac4f29b377a9b40d50ea752.png





我们可以看到，在决策边界附近，我们对预测哪个标签的不确定性最高。你可以想象，将预测与不确定性相关联是许多应用（如医疗保健）的关键属性。为了进一步提高准确性，我们可能希望主要在来自高不确定性区域的样本上训练模型。








小批量ADVI#


到目前为止，我们已经一次性地在所有数据上训练了我们的模型。显然，这种方法无法扩展到像ImageNet这样的大规模数据集。此外，在数据的小批量上进行训练（随机梯度下降）可以避免局部最小值，并且可以更快地收敛。


幸运的是，ADVI也可以在小型批次上运行。这只需要一些设置：






minibatch_x, minibatch_y = pm.Minibatch(X_train, Y_train, batch_size=50)
neural_network_minibatch = construct_nn(minibatch_x, minibatch_y)
with neural_network_minibatch:
    approx = pm.fit(40000, method=pm.ADVI())













Finished [100%]: Average Loss = 134.51













plt.plot(approx.hist)
plt.ylabel("ELBO")
plt.xlabel("iteration");








../../../_images/31594df58e1ae32104356e82f5df1ca3293c6c9ad0d699aaef2567efd6dce7fd.png





如您所见，小批量ADVI的运行时间要低得多。它似乎也收敛得更快。


为了好玩，我们也可以看看跟踪。关键是，我们还可以得到神经网络权重的置信度。






az.plot_trace(trace);








../../../_images/c2c68f100683eca15ef605e20ea35ce0dd8fc2da08a1bb2b80045e7dfe887812.png





你可能会认为上述网络并不真正深，但请注意，我们可以很容易地扩展它以拥有更多层，包括卷积层，以训练更具挑战性的数据集。






致谢#


吉冈拓也在PyMC3中对ADVI做了很多工作，包括小批量实现的以及从变分后验采样的部分。我还想感谢Stan团队（特别是Alp Kucukelbir和Daniel Lee）推导了ADVI并为我们讲解了它。同时感谢Chris Fonnesbeck、Andrew Campbell、吉冈拓也和Peadar Coyle对早期草稿的有益评论。






参考资料#







[1]

Alp Kucukelbir, Rajesh Ranganath, Andrew Gelman, 和 David M. Blei. 在Stan中的自动变分推断。2015年。arXiv:1506.03431。





[2]

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, 和 Martin Riedmiller. 使用深度强化学习玩雅达利游戏。2013年。arXiv:1312.5602。





[3]

C. Maddison 等人。D. Silver, A. Huang. 使用深度神经网络和树搜索掌握围棋游戏。自然，529:484–489, 2016. URL: https://doi.org/10.1038/nature16961.





[4]

Diederik P Kingma 和 Max Welling。自动编码变分贝叶斯。2014年。arXiv:1312.6114。





[5]

Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, 和 Andrew Rabinovich. 深入研究卷积。2014年。arXiv:1409.4842。












作者#


    

  • 这个笔记本最初是由Thomas Wiecki在2016年作为博客文章撰写的
    • 

  • 由 Chris Fonnesbeck 更新，适用于 PyMC v4，2022 年
    • 

  • 由 Oriol Abril-Pla 和 Earl Bellinger 于 2023 年更新
    • 







水印#






%load_ext watermark
%watermark -n -u -v -iv -w -p xarray









Last updated: Sat Jun 15 2024

Python implementation: CPython
Python version       : 3.11.9
IPython version      : 8.25.0

xarray: 2024.6.0

matplotlib: 3.8.4
seaborn   : 0.13.2
pymc      : 5.15.1
pytensor  : 2.22.1
numpy     : 1.26.4
arviz     : 0.18.0

Watermark: 2.4.3