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

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

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

AutoBNN

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

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

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

最后，通过 BNN 对 GPs 的加法和乘法运算符的类比，以及用于生成周期性核的输入变换，完成了这种转换。BNN 加法直接由组件 BNNs 的输出相加给出。BNN 乘法通过将 BNNs 隐藏层的激活值相乘，然后应用一个共享的密集层来实现。因此，我们仅限于乘具有相同隐藏宽度的 BNNs。

使用 AutoBNN

AutoBNN 包在 Tensorflow Probability 中提供。它使用 JAX 实现，并使用了 flax.linen 神经网络库。它实现了迄今为止讨论的所有基础核和运算符（Linear、Quadratic、Matern、ExponentiatedQuadratic、Periodic、Addition、Multiplication），外加一个新的核和三个新的运算符：

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

WeightedSum 以可学习的混合权重组合两个或多个 BNNs，其中可学习的权重遵循狄利克雷先验。默认情况下，使用浓度为 1.0 的平坦狄利克雷分布。WeightedSums 允许一种“软”版本的结构发现，即同时训练许多可能模型的线性组合。与使用离散结构的结构发现（如 AutoGP 中）相比，这允许我们使用标准的梯度方法来学习结构，而不是使用昂贵的离散优化。WeightedSum 允许我们并行评估潜在组合结构，而不是串行评估。

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

下图在 M3 数据集中的 N374 序列（从 1949 年开始的年度金融时间序列）上展示了结构发现技术。六个基础结构是 ExponentiatedQuadratic（与径向基函数核 RBF 相同）、Matern、Linear、Quadratic、OneLayer 和 Periodic 核。该图显示了在 32 个粒子上它们的权重的最大后验概率估计。所有高似然粒子都给 Periodic 分量赋予了较大的权重，给 Linear、Quadratic 和 OneLayer 赋予较低的权重，并且给 RBF 或 Matern 赋予了较大的权重。

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

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

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

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

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

致谢

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