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

2024年3月28日 | Urs Köster，谷歌研究软件工程师

AutoBNN将传统概率方法的可解释性与神经网络的可扩展性和灵活性相结合，利用复杂数据构建复杂的时间序列预测模型。

快速链接

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

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

AutoBNN

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

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

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

下图展示了在M3数据集的N374（从1949开始的年度财务时间序列）上结构发现技术的演示。六个基结构是指数二次（与径向基函数核相同，简称RBF）、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提供了一个强大而灵活的框架，用于构建复杂的时间序列预测模型。通过将贝叶斯神经网络和高斯过程的优势与组合核结合，AutoBNN为理解和预测复杂数据开辟了可能性世界。我们邀请社区尝试colab，并利用这个库创新和解决现实世界的挑战。

致谢

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

标签： 算法与理论，机器智能，开源模型与数据集

快速链接

AutoBNN包