BAYESFLOW: LEARNING COMPLEX STOCHASTIC MODELS WITH INVERTIBLE NEURAL NETWORKS BAYESFLOW:使用可逆神经网络学习复杂随机模型 https://arxiv.org/pdf/2003.06281

在一天的训练之后,预训练的BayesFlow网络在执行500个数据集的推断时,即使在最大试验次数N=1000的情况下,也只需要不到5秒钟。使用基于案例的SMC-MMD算法,500次推断运行将需要超过半年的时间才能完成

打开网易新闻 查看精彩图片
打开网易新闻 查看精彩图片
打开网易新闻 查看精彩图片

摘要

估计数学模型的参数是几乎所有科学分支中的一个常见问题。然而,当过程和模型描述变得日益复杂,且不存在显式似然函数时,这个问题可能会变得相当困难。在这项工作中,我们提出了一种基于可逆神经网络的全新全局摊销贝叶斯推断方法,我们称之为BayesFlow。该方法使用模拟来学习一个全局估计器用于从观测数据到潜在模型参数的概率映射。通过这种方式预训练的神经网络随后可以在不进行额外训练或优化的情况下,对涉及同一模型家族的任意数量的真实数据集推断出完整后验。此外,我们的方法包含一个摘要网络,该网络被训练用于将观测数据嵌入到信息量最大的摘要统计量中从数据中学习摘要统计量使得该方法适用于标准推断技术无法处理的建模场景,这些技术依赖于手工制作的摘要统计量。我们展示了BayesFlow在来自种群动力学、流行病学、认知科学和生态学的难以处理的模型上的实用性。我们认为,BayesFlow提供了一个通用框架,用于构建适用于任何可以从其模拟数据的正向模型的摊销贝叶斯参数估计机器。

1 引言

贝叶斯分析的目标是根据可观测的表现形式x推断出感兴趣的自然过程的潜在特征。在贝叶斯设置中,我们假设我们已经对正向问题有了足够的理解,即,从给定的隐藏参数θ的配置生成观测的适当模型。这个正向模型可以以两种形式提供:在基于似然的方法中,似然函数p(x | θ)是显式已知的,并且可以对任何一对(x, θ)进行分析或数值评估。相比之下,似然自由的方法仅需要能够从似然中采样。后一种方法通常通过模拟程序实现,该程序通过参数θ和独立噪声(即随机数)ξ的确定性函数g生成合成观测

在这种情况下,似然p(x | θ)仅通过模拟程序g的动作隐式定义,但对于模拟观测xi计算其确切的数值是不可能的。这反过来又阻止了标准的统计推断。

似然自由问题出现在例如p(x | θ)没有封闭形式的情况下,或者当正向模型由随机微分方程、蒙特卡罗模拟或复杂算法定义时[27, 49, 47, 51]。

在本文中,我们提出了一种基于可逆神经网络的似然自由设置的新贝叶斯解决方案。贝叶斯建模利用关于正向模型的可用知识来获得逆模型的后验分布的最佳可能估计

在贝叶斯推断中,后验编码了从一组观测数据中可以获得的所有关于θ的信息。假设观测数据来自正向模型的N次运行,具有固定的但未知的真实参数θ∗。贝叶斯逆建模面临三个挑战:

1. 在似然自由的情况下,贝叶斯公式的右侧总是难以处理的,必须进行近似。

2. 正向模型通常是非确定性的,因此关于θ的真实值存在固有的不确定性

3. 正向模型通常不是信息保持的,因此对于θ的可能值存在歧义。

这些问题的标准解决方案由近似贝叶斯计算(ABC)方法提供[45, 10, 39, 47]。ABC方法通过从提议(先验)分布重复采样参数,然后通过运行正向模型模拟多个数据集来近似后验。如果生成的数据集与实际观测数据集x1:N o足够相似,则相应的θ(l)被保留为所需后验的样本,否则被拒绝。更严格的相似性标准导致所需后验的更准确近似,但代价是更高且通常是禁止性的拒绝率。

更有效的近似推断方法,如序列蒙特卡罗(ABC-SMC)、马尔可夫链蒙特卡罗变体[44]或最近的神经密度估计方法[16, 38, 30],优化从提议分布采样,以平衡普通ABC方法的速度-准确性权衡。更多细节可以在相关工作部分和[9]的优秀综述中找到。

上述所有采样方法都在单个数据集的层面上操作,即对于每个观测序列,整个后验估计过程必须从头开始运行。因此,我们将这种方法称为案例推理。与摊销推理相比,单独为每个单独的数据集运行估计是有区别的,在摊销推理中,估计被分为一个可能昂贵的预先训练阶段,随后是一个便宜得多的推理阶段。预先训练阶段的目标是学习一个适用于任何观测序列的近似后验。然后,评估该模型对于特定的观测非常快,因此训练努力在重复评估中摊销(参见图1的图形说明)。案例推理和摊销推理之间的盈亏平衡取决于应用和模型类型,我们将在实验部分报告比较。然而,本文的主要目的是介绍一种通用的摊销贝叶斯推理方法,并展示其在流行正向模型的后验估计中出色的准确性。

打开网易新闻 查看精彩图片

为了使摊销推理在实践中可行,它必须适用于任意大小的数据集N。根据数据采集情况,对于固定的模型参数设置,可用的观测数量可能从N = 1到数百甚至更多不等。这不仅对我们的密度近似器的所需架构有影响,而且对其行为也有影响:它们必须表现出正确的后验收缩。因此,估计的后验随着可用观测数量N的增加而变得更尖锐(即,更尖峰)。在最简单的情况下,后验方差应以1/N的速率减小,但对于困难的(例如,多模态)真实后验可能会出现更复杂的行为

我们将这些考虑纳入我们的方法中,通过整合两个独立的深度神经网络模块(在方法部分详细说明;另见图1),这些模块在正向模型的模拟数据上联合训练:一个摘要网络和一个推理网络。

摘要网络负责将大小可变的观测数据集x1:N减少到一个固定大小的学习到的摘要统计量向量。在传统的似然自由方法中,方法设计者负责为每个应用预先选择合适的统计量[33, 32, 43, 45]。相比之下,我们的摘要网络直接从数据中学习最有信息的统计量,我们将在实验中展示(参见实验3.8),这些统计量优于手工构建的统计量。摘要网络与标准前馈网络不同,因为它们应该独立于输入大小N,并尊重数据固有的功能和概率对称性。例如,对于独立同分布的观测[6],需要排列不变网络,对于具有时间或空间依赖性的数据,需要循环网络[15]或卷积网络[29]。

推理网络负责根据观测数据的摘要统计量学习模型参数的真实后验。由于它只通过摘要网络的视角看到数据,因此后者捕获的所有对称性自动被后验继承。我们将推理网络实现为可逆神经网络。可逆神经网络基于最近关于归一化流[3, 25, 18, 13, 26]的理论和应用。基于流的方法在完美收敛下可以执行精确推断,并且从简单的低维问题到具有复杂依赖性的高维分布(例如,图像的像素[25])都具有良好的扩展性。对于每个感兴趣的应用/模型,我们使用从各自已知的正向模型和合理的先验中模拟的数据,联合训练一个可逆网络和一个相应的摘要网络。这种正向训练收敛后,网络的可逆性确保了对于逆模型的一个模型是免费获得的,只需通过模型反向运行推断即可。因此,我们的网络可以在给定应用领域中对任意多个数据集执行快速摊销贝叶斯推断,而无需昂贵的案例优化。我们将我们的方法称为BayesFlow,因为它结合了贝叶斯推断和基于流的深度学习的思想。

BayesFlow借鉴了现代深度概率建模的重大进展,也被称为深度生成建模[6, 25, 2, 24]。深度概率建模中的一个标志性思想是将复杂的目标分布表示为某种较简单潜在分布(例如,高斯分布或均匀分布)的非线性双射变换,即所谓的推前。因此,目标分布的密度估计,一个非常复杂的问题,被简化为学习非线性变换,这是通过标准反向传播进行梯度训练的神经网络的理想任务。在推理阶段,可以通过从较简单的潜在分布中采样并应用在训练阶段学习到的逆变换来获得来自目标分布的样本(请参见图1b,进行高级概述)。利用这种方法,深度概率模型的最近应用在迄今为止棘手的高维问题上取得了前所未有的性能[6, 25, 18]。

在贝叶斯推理的背景下,目标分布是给定观察数据的模型参数的后验概率。我们利用这样一个事实,即我们可以从前向模型模拟任意大量的训练数据,以确保摘要和可逆网络尽可能准确地近似真实后验分布。在推理阶段,我们的模型可以对任何候选参数θ的后验概率进行数值评估,或者可以为观察数据xo1:N 生成可能参数

打开网易新闻 查看精彩图片
的后验样本。在方法部分,我们展示了我们的网络确实在完美收敛时从正确的后验中采样。

总之,我们BayesFlow方法的贡献如下:

• 具有可逆神经网络的全局分摊近似贝叶斯推理

• 从原始数据集中学习最大信息摘要统计,而不是依赖于限制性的手工摘要统计;

• 对于从任意先验和后验中采样的真实后验分布提供理论保证;

• 并行计算可应用于前向模拟和神经网络优化;

为了说明BayesFlow的实用性,我们首先将其应用于两个具有解析可追踪后验的玩具模型。第一个是具有完整协方差矩阵和单峰后验的多变量高斯模型。第二个是具有多峰后验的高斯混合模型。然后,我们展示了BayesFlow在来自人口动态学、认知科学、流行病学和生态学的难以处理的似然模型中的应用,并展示了其在速度、恢复准确性和概率校准方面的实用性。与此同时,我们还介绍了几种性能验证工具。

1.1 相关工作

BayesFlow 结合了以前机器学习和深度学习方法中关于无似然推断的思想【31, 41, 33, 43, 22】。最常见的方法是将参数估计问题转化为监督学习任务。在这种设置中,通过从 p(θ) 中反复采样并通过模拟器运行带有采样参数的人造数据集 来创建形式为的大数据集。通常,通过使用固定的摘要函数 计算摘要统计量来减少模拟数据的维度。然后,训练一个监督学习算法(例如,随机森林【43】或神经网络【41】)来基于模拟数据的摘要统计量输出对真实数据生成参数的估计。因此,试图近似难以处理的逆模型。监督方法的一个主要缺点是它们只提供关于后验的有限信息(例如,点估计、分位数或方差估计)或对后验形状施加过于严格的分布假设(例如,高斯分布)。

我们的思想也与最优传输映射及其在贝叶斯推断中的应用【12, 40, 8, 5】密切相关。传输映射定义了一个(概率)测度之间的转换,可以构造为将简单的概率分布扭曲为更复杂的分布。在贝叶斯推断的背景下,传输映射已被应用于加速 MCMC 采样【40】、执行序列推断【12】以及通过直接优化解决推断问题【5】。事实上,BayesFlow 可以看作是通过可逆神经网络对可逆传输映射的参数化。一个重要的区别是,BayesFlow 不需要显式似然函数来近似目标后验,并且能够进行摊销推断。

类似的无似然推断思想也被最近的自动后验变换(APT)【16】和序列神经似然(SNL)【38】方法所采用。APT 通过掩码自回归流(MAF)网络迭代细化提议分布,以生成与特定观测数据集紧密匹配的参数样本。SNL 则在 MCMC 循环中训练一个掩码自动编码器密度估计器(MADE)神经网络,以加速收敛到真实后验。尽管这些方法也涉及相对昂贵的学习阶段和廉价的推断阶段,但后验推断仅针对单个数据集进行摊销。因此,每个单独的数据集都需要重新运行学习阶段(见图1a)。相比之下,我们提议通过使用条件可逆神经网络(cINN)估计器在整个可能参数和数据集范围内全局学习后验(见图1b)。以前,INNs 已成功地应用于建模天体物理学和医学数据【2】。我们调整了模型以适应数学建模背景下的参数估计任务,并开发了一种概率架构,用于对复杂数学模型进行完全贝叶斯和全局摊销推断。

2 方法

2.1 符号表示

在下文中,数学模型的参数数量将表示为D,数据集中的观测数量表示为N。我们将从感兴趣的数学模型模拟的数据表示为,其中每个个体xi可以表示为标量或向量。观测或测试数据将用上标o表示。数学模型的参数表示为向量,可逆神经网络和摘要神经网络的所有可训练参数分别表示为φ和ψ。当数据集由一段时间内的观测组成时,观测数量将表示为T。

2.2 学习后验分布

打开网易新闻 查看精彩图片
打开网易新闻 查看精彩图片
打开网易新闻 查看精彩图片
打开网易新闻 查看精彩图片
打开网易新闻 查看精彩图片

总结来说,如果摘要网络和可逆网络完美收敛,那么通过BayesFlow方法获得的近似后验是正确的。然而,在实际操作中,完美收敛是不现实的,导致不正确后验的误差主要来自三个方面。第一个误差是通过使用来自的模拟数据来近似方程17中的期望时引入的蒙特卡罗误差。第二个误差是由于摘要网络可能无法完全捕捉数据中的相关信息,或者当充分的摘要统计量不存在时。第三个误差是由于可逆网络不能准确地将真实后验转换为规定的高斯潜在空间。尽管我们可以通过更频繁地运行模拟器来减轻蒙特卡罗误差,后两者的误差在原则上更难检测和缓解。然而,最近关于概率对称性【6】和算法对齐【52】的研究可以为如何为特定问题选择合适的摘要网络提供一些指导。此外,可逆链的深度以及构建模块(稍后解释)可以进行调整,以提高从θ空间到z空间的变换表达能力。神经网络深度的益处已经在理论和实践中得到确认【28, 4】,因此我们期望在复杂环境中,随着网络深度的增加,性能会更好。

2.3 组成可逆网络

打开网易新闻 查看精彩图片
打开网易新闻 查看精彩图片

2.4 摘要网络

由于实际场景中观察到的样本数量通常不同(例如,不同的测量数量或时间点),以及数据集可能表现出各种冗余,因此条件不变神经网络(cINN)可以从某种形式的降维中受益。如前所述,我们希望避免通过限制性的手工摘要统计数据而丢失信息,相反,我们直接从数据中学习最有信息量的摘要统计数据。因此,我们不是将原始的模拟或观察到的数据输入到每个ACB,而是将数据通过一个额外的摘要网络传递,以获得一个固定大小的学习摘要统计向量。

摘要网络的架构应与观察数据的概率对称性相一致。对于时间序列数据,一个明显的选择是LSTM网络[15],因为递归网络可以自然地处理可变长度的长序列。另一个选择可能是1D全卷积网络[29],它已经在概率无似然推断的背景下被应用[41]。当处理具有可变大小的独立同分布(i.i.d.)样本时,可能需要不同的架构。这些数据通常被称为可交换的,或排列不变的,因为改变个体元素的顺序不会改变相关的似然或后验。换句话说,如果SN(·)是N个元素的任意排列,则后验应满足以下条件:

按照[6]的方法,我们通过一个等变非线性变换后跟一个池化操作符(例如,求和或平均)和另一个非线性变换来实现概率排列不变性:

打开网易新闻 查看精彩图片

其中是两个不同的全连接神经网络。在实践中,我们将多个等变和不变函数堆叠成一个不变网络,以提高表达能力[6]。

我们通过反向传播联合优化摘要网络的参数ψ和cINN链的参数φ。因此,训练过程完全是端到端的,并且BayesFlow通过在训练期间适当变化N,学习推广到不同大小的数据集,无论是通过排列不变的摘要网络还是通过变化递归/卷积网络的序列长度。

为了结合观察或模拟的数据,每个ACB的每个内部网络都被增强,以接受学习到的摘要向量˜x作为额外的输入。每个ACB的输出然后变为:

因此,通过整个条件不变链的完整传递可以表示为以及逆操作。图2展示了给定真实观察数据时,通过独立样本从学习到的推进分布逼近后验的过程。因此,关于数据和参数(逆模型)的映射知识被紧凑地编码在两个网络的权重中。

打开网易新闻 查看精彩图片

2.5 整合所有内容

算法1描述了使用任意摘要网络并采用在线学习方法的BayesFlow方法的基本步骤。

反向传播算法通过计算损失函数相对于神经网络参数的梯度,然后调整参数,以使损失函数达到最小。在训练过程中,我们使用由方程18给出的损失函数,没有遇到任何不稳定或收敛问题。注意,算法1的步骤3-14和18-22可以在GPU支持下并行执行,以显著加速收敛和推断。此外,步骤18-22可以并行应用于任意数量的观察数据集(见图2以获得完整的图形说明)。

接下来,我们将BayesFlow应用于两个具有单峰和多峰后验的玩具模型,然后使用它对来自人口动力学、认知科学、流行病学和生态学领域的具有不可处理似然的挑战性模型进行贝叶斯推断。我们认为这些模型适合于初步验证,因为它们在实现和解释观察数据的生成机制方面差异很大。因此,在这些不同的例子上表现良好,强调了BayesFlow方法的广泛经验效用。模型设置的详细信息可以在附录C中找到。

打开网易新闻 查看精彩图片

3.1 训练网络

在本文中描述的所有可逆和摘要网络都是通过反向传播联合训练的。对于以下所有实验,我们使用Adam优化器,起始学习率为,指数衰减率为0.95。我们为每个实验执行50,000到100,000次迭代(即,小批量更新步骤),并报告收敛网络得到的结果。注意,我们没有对网络超参数的最优值进行广泛的搜索,而是在本文的所有示例中使用了一个默认的BayesFlow,具有5到10个ACB和大小为128的摘要向量(有关摘要网络架构的更多细节,请参见附录C)。所有网络都是使用TensorFlow库在配备NVIDIA® GTX1060显卡的单GPU机器上实现和训练的。

关于数据生成步骤,我们采取了一种结合了在线学习思想的方法,即按需通过方程1进行数据模拟。相应地,数据集x1:N或M个数据集的批次是即时生成的,然后传递给神经网络。这种训练方法的优点是网络永远不会两次遇到相同的输入数据。此外,只要网络不断改进(即,损失不断减少),训练就可以继续,因为在线学习中的经典过拟合几乎是不可能的。然而,如果模拟计算成本高昂,研究人员需要尝试不同的网络或训练超参数,那么存储和重用模拟可能会有益,因为在线学习中的模拟和训练是紧密相连的。

一旦网络收敛,我们就存储训练好的网络,并使用它们对单独的验证数据集进行摊销推断。预训练的网络也可以在研究社区中共享,以便多个研究人员/实验室可以从推断的摊销中受益。

3.2 性能验证

为了评估以下应用示例中BayesFlow的性能,我们考虑了几种不同的指标:

- 归一化均方根误差(NRMSE) - 评估恢复真实参数值的点估计的准确性;

- 决定系数(R^2) - 评估点估计捕获的真实参数方差的比例;

- 重模拟误差(Errsim) - 评估真实数据分布与使用估计参数生成的数据分布之间的预测不匹配(即,后验预测检验);

- 校准误差(Errcal,[2]) - 评估近似后验的覆盖范围(即,可信区间是否真的可信);

- 基于模拟的校准(SBC,[46]) - 直观检测近似后验中的系统偏差。

计算所有指标的详细信息在附录B中给出。

3.3 概念证明:多元正态分布

打开网易新闻 查看精彩图片
打开网易新闻 查看精彩图片
打开网易新闻 查看精彩图片

3.4 多模态后验-高斯混合模型

为了测试BayesFlow方法是否能够恢复多模态后验分布,我们将其应用于生成高斯混合模型(GMM)。在实际中,多模态后验分布会出现,例如,当前向模型定义为不同过程的混合时,或当模型在其参数空间中表现出较大的多变量权衡时(例如,存在多个独立的后验密度区域具有合理的参数值)。因此,展示我们的方法能够捕捉这种行为并且不会遭受模态崩溃是很重要的。

根据[2],我们构建了一个场景,其中观测数据x是一个独热编码向量,表示红色、绿色、蓝色或黄色的硬标签之一(即单次观测,因此 \(N = 1\))。参数是从八个高斯簇混合中抽取的点的二维坐标,这些簇的中心以单位方差在原点周围顺时针分布(见图3左上角)。前四个簇被标记为红色,接下来的两个簇标记为绿色,剩下的两个簇分别标记为蓝色和黄色。后验分布由对应标签索引的簇组成。我们通过增加BayesFlow的深度从1个ACB块到5个ACB块多次进行实验。通过这种方式,我们可以研究cINN深度对近似多模态后验质量的影响。我们将每个BayesFlow训练50个周期,并从训练模型获得的近似后验中抽取8000个样本。

打开网易新闻 查看精彩图片

所有BayesFlow的结果如图3所示。我们观察到,近似效果随着cINN链的加深而受益,使用更多的ACB时簇分离变得更清晰。这证实了我们的方法能够恢复多模态后验分布。

3.5 随机时间序列模型 - Ricker模型

在本节中,我们估计一个著名的离散随机人口动态模型的参数[51]。通过这个例子,我们旨在实现几个目标:首先,我们想要证明BayesFlow方法能够通过从原始数据中学习摘要统计量,准确恢复具有不可处理似然函数的实际模型的参数。其次,我们展示BayesFlow能够适当地处理与数据完全无关的参数,通过将估计值减少到相应参数的先验分布。第三,我们比较BayesFlow方法的全局性能与能够进行摊销似然无似然推断的相关方法。最后,我们展示随着观测数据数量的增加,估计的改进和期望的后验收缩。

离散人口动态模型描述了种群中的个体数量如何在离散时间单位内变化[51]。特别是,Ricker模型通过以下非线性方程将第t代的个体数量xt描述为前一代中预期个体数量的函数:

打开网易新闻 查看精彩图片

对于t = 1, ..., T,其中Nt是时间点t的预期个体数量,r是增长率,ρ是缩放参数,ξt是随机高斯噪声。Ricker模型的似然函数没有封闭形式的解,并且该模型以其混沌行为而闻名,因此它是进行无似然推断的一个合适的候选模型。参数估计任务是从观察到的一维时间序列数据中恢复参数θ = (ρ, r, σ),其中每个。

如果数据不包含关于某个特定参数的信息,该怎么办?在这种情况下,任何好的估计方法都应该检测到这一点,并返回该特定参数的先验。为了测试这一点,我们在参数向量θ中附加了一个随机均匀变量u ∼ U(0, 1)并使用这个额外的虚拟参数训练BayesFlow。我们期望网络忽略这个虚拟参数,即我们假设估计的u的后验与均匀先验相似。

我们将BayesFlow的性能与以下能够进行摊销无似然推断的最新方法进行了比较:条件变分自编码器(cVAE)[35]、带有自回归流的cVAE(cVAE-IAF)[26]、带有异方差损失的深度推断(DeepInference)[41]、通过LSTM神经网络学习信息摘要统计量的近似贝叶斯计算(ABC-NN)[22]和分位数随机森林(ABC-RF)[43]。为了训练这些模型,我们模拟了长度不同的Ricker模型时间序列。在每次训练迭代中,时间点数T从均匀分布T ∼ U(100, 500)中抽取。

所有神经网络方法在从Ricker模型模拟的数据上训练了100个周期,每个周期进行了1000次迭代。由于ABC-RF方法不支持在线学习且增加参考表似乎并未提高性能,我们在一个包含200,000个数据集的参考表上拟合了ABC-RF方法。为了避免为ABC-RF方法使用手工制作的摘要统计量,我们输入了通过与cINN联合训练的摘要网络获得的摘要向量。因此,ABC-RF方法利用了最大信息量的统计量作为输入。我们在一个独立的测试集(由T = 500生成的500个数据集)上验证了所有方法的性能。我们在表1中报告了每种方法和每个参数的性能指标。

打开网易新闻 查看精彩图片

参数r和ρ似乎可以通过这里考虑的所有方法很好地恢复。σ参数则较难估计,其中BayesFlow和ABC-NN方法表现最佳。此外,BayesFlow在所有参数和指标上的表现都非常好。重要的是,BayesFlow获得的校准误差Errcal始终较低,这表明近似后验的形状与真实后验的形状非常匹配。变分方法(cVAE,cVAE-IAF)在恢复σ的后验时遇到了一些问题。ABC-NN和ABC-RF方法似乎能够高精度地恢复点估计,但前者的近似后验表现出相对较高的校准误差。ABC-RF方法只能估计后验分位数,因此无法计算可比的校准指标。

进一步的结果展示在图4中。检查所有方法在一个示例测试数据集上获得的完整后验分布,我们注意到只有BayesFlow和ABC-NN方法能够恢复虚拟噪声变量u的无信息后验分布(图4a)。此外,在观察σ的后验分布时,Ricker模型的贝叶斯处理的重要性变得清晰。在大多数测试数据集中,后验密度分布在整个先验范围内(高后验方差),这表明获得的估计存在很大的不确定性。此外,边际参数后验的形状在验证数据集之间差异很大,这强调了避免对允许的后验形状进行随意限制的重要性(参见图S5中的示例)。我们还观察到,当有更多时间点时,使用BayesFlow的参数估计变得越来越准确(图4b)。在时间点数量最多的情况下,参数恢复尤其出色(见图4c)。最后,图4d揭示了随着提供给摘要网络的时间点数量增加,后验收缩显著。

打开网易新闻 查看精彩图片

3.6 感知决策模型 - Lévy飞行模型

接下来,我们估计一个人类决策的随机微分方程模型的参数。我们首次对最近提出的Lévy飞行模型(LFM)进行了贝叶斯处理,因为其难以处理性迄今为止使得传统的非摊销贝叶斯推断方法极其缓慢【49】。

在这个示例中,我们首先希望通过实验证明BayesFlow能够处理由N次独立运行复杂随机模拟器产生的不同大小的独立同分布(i.i.d.)数据集。为此,我们检查了BayesFlow在各种数据集大小范围内的全局性能。此外,我们还希望展示与基于案例的推断相比,摊销推断在效率和恢复方面的优势。为此,我们将BayesFlow与四种其他最新的无似然推断方法一起应用于单个数据集,并显示在某些情况下,即使只有5个数据集,摊销推断的速度优势也变得显著。重要的是,研究人员经常将相同的模型拟合到不同的数据集上,因此,如果存在预训练模型,在效率和生产力方面将具有巨大的优势。

我们专注于证据累积模型(EAM)家族,这些模型通过一组神经认知动机参数描述人类决策【42】。EAM最常应用于选择反应时间(RT)数据,以估计控制分类和(感知)决策的潜在过程。在其最通用的公式中,EAM的前向模型采用随机常微分方程(ODE)的形式:

其中dx表示累积器激活的变化,v表示信息累积的平均速度(通常称为漂移率),ξ代表随机的附加成分,通常被建模为来自以0为中心的高斯分布:。

EAM对于无似然推断特别适用,因为这个模型家族中大多数有趣的成员的似然性都是难以处理的【34】。这种难以处理性排除了许多有趣的应用和实证驱动的模型改进。在这里,我们将BayesFlow应用于估计最近提出的Lévy-Flight模型(LFM)[49]的参数。LFM假设证据累积过程的α稳定噪声分布,这允许对决策过程中的不连续性进行建模。然而,包含α稳定噪声(而不是通常假定的高斯噪声)导致了一个具有难以处理似然性的模型:

打开网易新闻 查看精彩图片

α控制噪声分布中异常值的概率。LFM还有三个额外的参数:决策过程终止所需的证据量由阈值a确定;相对起始点zr确定在实际决策备选方案呈现之前累加器可用的起始证据量;以及附加的非决策时间t0。

在网络训练期间,我们从两个具有不同漂移率的实验条件模拟响应时间数据,因为这样的设计在心理学研究中经常遇到。因此,参数估计任务是从二维独立同分布的RT数据x1:N中恢复参数,其中每个表示在两个条件下获得的RTs。每次训练迭代时,试验次数从均匀分布N ∼ U(100, 1000)中抽取。采用在线学习方法训练网络花费了不到一天的时间。对于每个参数,对1000个数据集进行推断,每个参数的后验样本约为7.39秒。

为了研究摊销推断对这个模型是否有利,我们还将SMC-ABC算法的一个版本应用于一个单独的数据集,其中N = 500。由于EAM数据没有足够的摘要统计量可用,我们应用了最大均值距离(MMD)度量作为完整原始实验RT分布之间的距离,以防止信息损失。由于MMD计算成本高昂,我们使用了GPU实现,以确保MMD的计算不会成为比较的瓶颈。为了从SMC-MMD近似后验中获得2000个样本的良好近似,我们对20个种群运行了算法,最终的拒绝阈值为ε = 0.04。我们还从通过将预训练的BayesFlow网络应用于相同数据集而获得的近似后验中抽取了2000个样本。

在SMC-MMD中,我们应用了三种最近的神经密度估计方法,SNPE-A [37]、SNPE-B [30]和SNPE-C ([16],也称为APT)。由于这些方法都依赖于数据的摘要统计量,我们计算了每个实验响应时间分布的前6个时刻以及正确/错误响应的比例。我们对每种方法进行了一轮的训练,每轮训练进行了100个周期和5000个模拟数据集,以保持运行时间最少。此外,当进行多轮训练时,我们没有观察到性能的改善。对于每个模型,我们从近似联合后验中抽取了2000个样本,以使样本数量与通过SMC-MMD获得的样本数量一致。

比较结果如图5所示。我们首先关注与单个数据集上的SMC-MMD的比较。图5a展示了BayesFlow和SMC-MMD获得的边缘和双变量后验。BayesFlow的近似后验似乎更加锐利。观察SCB图(图6b),我们可以得出结论,BayesFlow的近似后验反映了真实后验的锐利性,否则SCB图将显示出与均匀性的明显偏差。此外,图5b展示了应用每种方法获得的边缘后验。值得注意的是,各种方法和参数的性能和锐利度都有所变化,所有方法都通过后验均值在NRMSE和R2指标方面获得了良好的点估计恢复。

打开网易新闻 查看精彩图片
打开网易新闻 查看精彩图片
打开网易新闻 查看精彩图片

我们注意到,在一天的训练之后,预训练的BayesFlow网络在执行500个数据集的推断时,即使在最大试验次数N=1000的情况下,也只需要不到5秒钟。使用基于案例的SMC-MMD算法,500次推断运行将需要超过半年的时间才能完成。我们还注意到,将独立的推断线程并行化到多个核心或(GPU)计算集群的节点上,可以显著提高这里考虑的基于案例方法的墙钟速度。然而,同样适用于BayesFlow训练,因为它最昂贵的部分——正向模型的模拟——将从并行计算中获益最多。

图6展示了BayesFlow在所有验证数据集和所有试验大小N上的全局性能。首先,我们观察到所有LFM参数的恢复效果非常好,NRMSE值在0.008到0.048之间,R2值在最大试验次数时在0.972到0.99之间。重要的是,估计在所有试验次数上都保持非常好,并且随着更多试验的可用性而提高(见图6a)。参数α似乎是最具挑战性的估计对象,需要更多的数据来进行良好的估计,而决策时间参数t0几乎可以完美地恢复所有试验大小。最后,SCB直方图表明边际后验没有系统偏差(见图6b)。

3.7 随机微分方程 - SIR流行病学模型

通过这个例子,我们想要进一步证实在非独立同分布(non-i.i.d.)随机常微分方程(ODE)模型上观察到的LFM模型的卓越全局性能和概率校准。为此,我们研究了一个来自流行病学的隔室模型,其输出包含可变大小的多维且相互依赖的时间序列。因此,研究我们的方法在应用于直接由ODE模拟器输出的数据时的表现是有意义的。

流行病学中的隔室模型描述了传染病在个体群体中传播的随机动态[23, 20]。这些模型的参数编码了疾病的重要特征,如感染率和恢复率。随机SIR模型描述了N个个体在三个离散状态之间 - 易感者(S)、感染者(I)和康复者(R) - 的转换,其动态遵循以下方程:

打开网易新闻 查看精彩图片

其中S + I + R = N表示易感者、感染者和康复者的个体数量。参数β控制着从易感到感染的转换率,而γ控制着从感染到恢复的转换率。上述随机系统没有解析解,因此需要数值模拟方法从数据中恢复参数值。将问题视为参数估计任务,挑战在于从三维时间序列数据中恢复参数θ = {β, γ},其中每个是包含时间t的易感者(S)、感染者(I)和康复者(R)数量的三元组。

在网络训练过程中,我们通过随机SIR模型模拟不同长度的时间序列。在每次训练迭代中,时间点数T从均匀分布T ∼ U(200, 500)中抽取。对于较小的T,系统尚未达到平衡状态(即并非所有个体都已从I状态转变为R状态)。特别有趣的是观察BayesFlow是否能在过程动态仍在展开时恢复出速率参数。使用在线学习方法训练网络大约花费了两个小时。在1000个数据集上进行推断,每个参数有2000个后验样本,大约花费了1.1秒。

打开网易新闻 查看精彩图片

3.8 学习总结与手工总结:Lotka-Volterra 总体模型 参见附录 A

4 讨论

在当前的工作中,我们提出并探讨了一种新方法,该方法使用可逆神经网络来执行全局变动的近似贝叶斯推断。我们将这种方法命名为BayesFlow,仅需使用前向模型的模拟来学习数据与参数之间的高效概率映射。通过将BayesFlow应用于不同研究领域的模型和数据,我们展示了其效用。此外,我们探索了每次迭代中观测数量可变的在线学习方法。结果表明,该方法在当前工作的所有示例中均能实现出色的参数估计。从理论上讲,BayesFlow适用于任何可以实现为计算机模拟的数学前向模型。以下,我们将重点介绍BayesFlow的主要优势。

首先,引入独立的摘要网络和推断网络,使得该方法与观察数据的形状或大小无关。摘要网络以自动且数据驱动的方式学习数据的固定大小向量表示。由于摘要网络与推断网络共同优化,因此所学到的数据表示在推断参数后验时会被鼓励为最大信息量。这在适当的摘要统计量未知的情况下尤为有用,因为选择次优摘要函数会导致相关信息丢失。然而,如果在特定领域中存在充分的统计量,可以完全省略摘要网络,并将这些统计量直接输入可逆网络。

其次,我们展示了在完美收敛的情况下,BayesFlow可以生成来自正确后验的样本,而不对后验的形状做出分布假设。这与变分方法形成对比,后者优化的是后验的下界【26, 24】,且通常假设高斯近似后验。此外,我们还在所有示例中展示了BayesFlow方法生成的后验均值大多是对真实值的优秀估计。更重要的是,BayesFlow方法恢复了参数的完整后验,不需要使用点估计或后验的摘要统计量。此外,我们观察到理想的后验收缩(随着观测数量的增加后验方差减小)和随观测数量增加的更好恢复。这些都是任何贝叶斯参数估计方法所必备的性质,因为它们反映了认识不确定性的减少和由于数据增多而信息量的同时增加。

第三,BayesFlow的最大计算成本发生在训练阶段。一旦训练完成,网络可以高效地计算来自前向模型的任何观察数据集的后验。这类似于最近引入的预付方法【33】。然而,该方法通过记忆大量预先计算的摘要统计量进行快速的最近邻推断,而BayesFlow的网络权重定义了数据与参数之间关系的抽象表示,覆盖了整个隐参数空间。传统上,这样的抽象表示仅存在于解析可逆的模型族中,而更复杂的前向模型则需要基于个案的推断,即每个观察数据集都需进行昂贵的重新训练。BayesFlow实现的变动推断因此对于在需要适应多个独立数据集的研究领域中探索、测试和比较竞争性的科学假设尤其有利。

最后,BayesFlow方法的所有计算都受益于高度的并行性,因此可以利用现代GPU加速的优势。

尽管如此,提出的方法也有一些局限性需要提及。尽管我们可以理论上保证在完美收敛的情况下BayesFlow从真实联合后验中采样,但在实践中可能无法实现。因此,必须对方法的每次应用进行点估计和联合后验估计的适当校准。幸运的是,由于变动推断,验证已训练的BayesFlow架构相对容易。下面,我们讨论该方法的潜在挑战和局限性。

首先,摘要网络和推断网络的设计是实现方法最佳性能的关键选择。如前所述,摘要网络应能够在不丢失重要信息的情况下表示观测数据,而可逆网络应足够强大以捕捉前向模型的行为。然而,在某些现实世界场景中,可能很难找到合适的摘要网络设计。最近关于概率对称性【6】和算法对齐【52】的研究以及我们当前的实验确实提供了一些关于摘要网络设计的见解。例如,独立同分布的数据会引发一个置换不变分布,使用深度不变网络【6】可以很好地建模。具有时间或空间依赖性的数据最好使用循环网络【21】或卷积网络【41】进行建模。当成对或多向关系特别有信息时,注意力机制【48】或图网络【52】是合理的选择。另一方面,可逆网络的深度应根据感兴趣的数学模型的复杂性进行调整。更多的ACB(可逆块)将使网络能够编码更复杂的分布,但也会增加训练时间。非常高维的问题可能还需要非常大的网络,拥有数百万参数,直到估计变得实际不可行。然而,大多数生命科学中的数学模型优先考虑简洁性和可解释性,因此不包含数百或数千个潜在参数。无论如何,未来的应用可能需要超出我们初步建议的新网络架构和解决方案。

另一个潜在问题是大量的神经网络和优化超参数可能需要用户进行微调,以在特定任务上实现最佳性能。我们观察到,通常使用默认设置就能实现优秀的性能。使用由5到10个ACB组成的较大网络似乎不会损害性能或使训练不稳定,即使所需学习的模型相对简单。根据我们的结果,我们预计单个架构应能在给定领域的模型上表现良好。未来的研究应通过将方法应用于不同或甚至竞争的模型来研究这一普遍性问题。未来的研究还应探讨现代超参数优化方法(如贝叶斯优化【14】)的影响。

最后,尽管现代深度学习库允许快速且相对简单地开发各种神经网络架构,但与该方法相关的实现负担仍然不可忽视。因此,我们目前正在开发一个通用的用户友好软件,以便将我们方法的复杂性抽象化,使用户可以轻松使用。

我们希望新的BayesFlow方法能够帮助各领域的研究人员加速基于模型的推断,并能在本文考虑的示例之外进一步证明其实用性。

相关论文:

Consistency Models for Scalable and Fast Simulation-Based Inference https://arxiv.org/abs/2312.05440v2

Flow Matching for Scalable Simulation-Based Inference https://arxiv.org/abs/2305.17161