AutoBNN：利用组合贝叶斯神经网络进行概率时间序列预测

时间序列问题无处不在，从预测天气和交通模式到了解经济趋势。贝叶斯方法从对数据模式的假设（先验概率）开始，收集证据（例如新的时间序列数据），并持续更新该假设以形成后验概率分布。传统贝叶斯方法如高斯过程（GPs）和结构时间序列被广泛用于建模时间序列数据，例如常用的莫纳罗亚二氧化碳数据集。然而，这些方法通常依赖领域专家费力选择适当的模型组件，并且计算成本可能很高。神经网络等替代方法缺乏可解释性，使得难以理解它们如何生成预测，并且无法产生可靠的置信区间。

为此，我们推出了AutoBNN，这是一个用JAX编写的新开源软件包。AutoBNN自动化了可解释时间序列预测模型的发现，提供高质量的不确定性估计，并有效扩展以用于大型数据集。我们将描述AutoBNN如何将传统概率方法的可解释性与神经网络的可扩展性和灵活性相结合。

AutoBNN

AutoBNN基于过去十年的一系列研究，这些研究通过使用具有学习核结构的高斯过程对时间序列进行建模，提高了预测准确性。高斯过程的核函数编码了关于被建模函数的假设，例如趋势、周期性或噪声的存在。对于学习过的高斯过程核，核函数是组合定义的：它要么是基核（如线性、二次、周期、Matérn或指数二次），要么是使用加法、乘法或变点等运算符组合两个或多个核函数的复合核。这种组合核结构有两个相关目的。首先，它足够简单，使得对其数据是专家但不一定对高斯过程是专家的用户可以为其时间序列构建合理的先验。其次，像顺序蒙特卡洛这样的技术可以用于对小型结构进行离散搜索，并输出可解释的结果。

AutoBNN在这些思想的基础上进行了改进，用贝叶斯神经网络（BNNs）取代了高斯过程，同时保留了组合核结构。贝叶斯神经网络是一种在权重上具有概率分布而非固定权重集的神经网络。这导致输出的分布，捕捉了预测中的不确定性。贝叶斯神经网络相对于高斯过程带来以下优势：首先，训练大型高斯过程计算成本高昂，传统训练算法的时间复杂度是时间序列数据点数量的立方。相反，对于固定宽度，训练贝叶斯神经网络的时间复杂度通常与数据点数量近似线性。其次，贝叶斯神经网络比高斯过程训练操作更适合GPU和TPU硬件加速。第三，组合贝叶斯神经网络可以轻松与传统深度贝叶斯神经网络结合，后者具有特征发现的能力。可以想象“混合”架构，其中用户指定顶层结构为Add(Linear, Periodic, Deep)，而深度贝叶斯神经网络则负责学习来自潜在高维协变量信息的贡献。

那么，如何将具有组合核的高斯过程转化为贝叶斯神经网络呢？单层神经网络通常会在神经元数量（或“宽度”）趋于无穷大时收敛到高斯过程。最近，研究人员发现了相反方向的对应关系——许多流行的高斯过程核（如Matern、指数二次、多项式或周期）可以通过具有适当选择的激活函数和权重分布的无限宽度贝叶斯神经网络获得。此外，即使宽度远小于无限，这些贝叶斯神经网络仍然保持与相应高斯过程接近。例如，下图显示了真实高斯过程与其对应的宽度为10的神经网络版本在观测对之间的协方差以及回归结果方面的差异。

真实高斯过程核（上行）与其宽度10神经网络近似（下行）的Gram矩阵比较

真实高斯过程核（上行）与其宽度10神经网络近似（下行）的回归结果比较

最后，通过贝叶斯神经网络对应的高斯过程的加法和乘法运算符，以及产生周期核的输入扭曲，完成了这种转换。贝叶斯神经网络的加法直接通过添加组件贝叶斯神经网络的输出给出。贝叶斯神经网络的乘法是通过将贝叶斯神经网络的隐藏层激活相乘，然后应用共享的密集层来实现的。因此，我们仅限于只能相乘具有相同隐藏宽度的贝叶斯神经网络。

使用AutoBNN

AutoBNN包在TensorFlow Probability中可用。它用JAX实现，并使用flax.linen神经网络库。它实现了到目前为止讨论的所有基核和运算符（线性、二次、Matern、指数二次、周期、加法、乘法），外加一个新核和三个新运算符：

OneLayer核，一个单隐藏层ReLU贝叶斯神经网络
ChangePoint运算符，允许在两个核之间平滑切换
LearnableChangePoint运算符，与ChangePoint相同，但位置和斜率被赋予先验分布并可以从数据中学习
WeightedSum运算符

WeightedSum使用可学习的混合权重组合两个或多个贝叶斯神经网络，其中可学习的权重遵循Dirichlet先验。默认情况下，使用浓度为1.0的平坦Dirichlet分布。

WeightedSums允许一种“软”版本的结构发现，即同时训练许多可能模型的线性组合。与使用离散结构的结构发现（如AutoGP中）相比，这允许我们使用标准梯度方法来学习结构，而不是使用昂贵的离散优化。WeightedSum允许我们并行评估潜在组合结构，而不是串行评估。

为了便于探索，AutoBNN定义了许多模型结构，这些结构包含顶层或内部的WeightedSums。这些模型的名称可以用作任何估计器构造函数中的第一个参数，包括诸如sum_of_stumps（所有基核的WeightedSum）和sum_of_shallow（将所有基核与所有运算符的所有可能组合相加）等内容。

sum_of_stumps模型图示。上行条形图显示每个基核的贡献量，下行显示基核表示的函数。得到的加权和显示在右侧。

下图展示了在M3数据集中的N374（从1949开始的年度财务时间序列）上进行结构发现的技术。六个基本结构是指数二次（与径向基函数核相同，简称RBF）、Matern、线性、二次、OneLayer和周期核。该图显示了在32个粒子集合上它们的权重的MAP估计。所有高似然粒子都给予周期组件大的权重，给予线性、二次和OneLayer低的权重，并给予RBF或Matern大的权重。

在32个粒子上基核权重的MAP估计的并行坐标图。sum_of_stumps模型在M3数据集的N374系列上训练（蓝色插入）。较暗的线对应较高似然的粒子。

通过使用WeightedSums作为其他运算符的输入，可以表达丰富的组合结构，同时保持模型紧凑且可学习权重的数量少。例如，我们包括sum_of_products模型（如下图所示），它首先创建两个WeightedSums的成对乘积，然后是这两个乘积的和。通过将一些权重设置为零，我们可以创建许多不同的离散结构。此模型中可能的结构总数为2^16，因为有16个基核可以打开或关闭。通过仅训练这一个模型，所有这些结构都被隐式探索。

“sum_of_products"模型图示。四个WeightedSums中的每一个都具有与"sum_of_stumps"模型相同的结构。

然而，我们发现某些核的组合（例如，周期与Matern或指数二次的乘积）在许多数据集上会导致过拟合。为了防止这种情况，我们定义了像sum_of_safe_shallow这样的模型类，在使用WeightedSums执行结构发现时排除此类乘积。

对于训练，AutoBNN提供AutoBnnMapEstimator和AutoBnnMCMCEstimator来分别执行MAP和MCMC推断。任一估计器都可以与六个似然函数中的任何一个结合使用，包括四个基于具有不同噪声特征的正态分布（用于连续数据）和两个基于负二项分布（用于计数数据）。

在我们的示例colab中在莫纳罗亚二氧化碳数据集上运行AutoBNN的结果。模型捕捉了数据中的趋势和季节性成分。外推至未来，均值预测略微低估了实际趋势，而95%置信区间逐渐增加。

要拟合如上图所示的模型，只需使用以下10行代码，使用scikit-learn风格的估计器接口：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


import autobnn as ab

model = ab.operators.Add(
    bnns=(ab.kernels.PeriodicBNN(width=50),
          ab.kernels.LinearBNN(width=50),
          ab.kernels.MaternBNN(width=50)))

estimator = ab.estimators.AutoBnnMapEstimator(
    model, 'normal_likelihood_logistic_noise', jax.random.PRNGKey(42),
    periods=[12])

estimator.fit(my_training_data_xs, my_training_data_ys)
low, mid, high = estimator.predict_quantiles(my_training_data_xs)

结论

AutoBNN为构建复杂的时间序列预测模型提供了一个强大而灵活的框架。通过将贝叶斯神经网络和高斯过程与组合核的优势相结合，AutoBNN为理解和预测复杂数据开辟了无限可能。我们邀请社区尝试colab，并利用这个库来创新和解决现实世界的挑战。

致谢

AutoBNN由Colin Carroll、Thomas Colthurst、Urs Köster和Srinivas Vasudevan编写。我们要感谢Kevin Murphy、Brian Patton和Feras Saad的建议和反馈。