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

AutoBNN：概率时间序列预测与组合式贝叶斯神经网络

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

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

AutoBNN

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

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

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

最后，通过BNN对GP的加法和乘法运算符的模拟，以及输入扭曲以产生周期性内核，完成了转换。BNN加法直接通过添加组件BNN的输出给出。BNN乘法通过将BNN的隐藏层激活相乘，然后应用共享的密集层来实现。因此，我们仅限于只能相乘具有相同隐藏宽度的BNN。

使用AutoBNN

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

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

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

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

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

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

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

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

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

要拟合如上图所示的模型，只需以下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为构建复杂的时间序列预测模型提供了一个强大而灵活的框架。通过将BNN和具有组合式内核的GP的优势相结合，AutoBNN为理解和预测复杂数据开辟了一个充满可能性的世界。我们邀请社区尝试colab，并利用这个库来创新和解决现实世界的挑战。

致谢

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