从嘈杂数据中推断复杂模型的参数：CMPE

Consistency Models for Scalable and Fast Simulation-Based Inference

用于可扩展和快速基于仿真的推理的一致性模型

https://arxiv.org/html/2312.05440v2

CMPE 在速度和性能方面都可以同时优于所有其他方法

CMPE不仅仅是FMPE的更快版本。相反，它展示了与FMPE的定性差异，超越了更快的采样速度：(i) 在实验4中，我们表明CMPE比FMPE对神经网络架构的依赖性更低，这使得当任务的最佳架构未知时，CMPE成为一个有前景的替代方案；(ii) 在我们的实验中，我们观察到CMPE在低数据环境中表现良好，使其在训练数据稀缺时成为一种有吸引力的方法。

摘要

基于仿真的推断（SBI Simulation-based inference）不断寻求更具表现力的算法，以准确地从嘈杂数据中推断复杂模型的参数。我们提出了一种用于神经后验估计的一致性模型（CMPE consistency models for neural posterior estimation），这是一种新的自由形式条件采样器，用于可扩展、快速和摊销的SBI，采用生成神经网络。CMPE将归一化流和流匹配 normalizing flows and flow matching 方法的优势结合到一个单一的生成架构中：它本质上提炼了一个连续的概率流，并且能够使用无约束的架构进行快速的少样本推断，可以根据估计问题的结构进行定制。我们的实证评估表明，CMPE不仅在三个困难的低维问题上胜过了当前最先进的算法，而且在高维贝叶斯去噪实验和估计计算要求高的肿瘤球体生长多尺度模型中取得了竞争性表现。

1. 引言

计算机仿真在无数科学学科中发挥着基础性作用，从物理学到生物学，从气候科学到经济学（Lavin等，2021）。

Simulation programs can generate observables 𝐱∈𝒳 as a function of unknown parameters 𝜽∈Θ and latent program states 𝝃∈Ξ (Cranmer et al., 2020).仿真程序可以将观测值x ∈ X 生成为未知参数θ ∈ Θ 和潜在程序状态ξ ∈ Ξ 的函数（Cranmer等，2020）。这种正向问题通常通过科学理论实例化为生成模型而被很好地理解。

然而，逆问题要困难得多，并且构成了贝叶斯（概率）推断的关键：根据观测值x 推理出未知的θ。贝叶斯定理捕获了在观测数据x 条件下的可信参数θ 的完整后验分布，如p(θ | x) ∝ p(θ) p(x | θ)，给定先验p(θ)。虽然从仿真程序中生成合成数据是可能的（尽管可能很慢），但似然密度p(x | θ) 通常不是显式可用的。相反，它只是通过对所有可能的程序执行路径进行高维积分而隐含地定义的（Cranmer等，2020）。

一种新时代的仿真智能（Lavin等，2021）正在解决这个具有挑战性的环境，并在各种应用领域取得了显着的成功（Butter等，2022；Radev等，2021b；Shiono，2021；Bieringer等，2021；von Krause等，2022；Ghaderi-Kangavari等，2023）。

这些在基于仿真的推断（SBI）领域的进展与条件生成建模的快速进展相辅相成。目前，生成神经网络的进展主要表现在生成图像（Goodfellow等，2014）、视频（Vondrick等，2016）和音频（Engel等，2020）方面的显著性能上。在生成模型的类别中，基于分数的扩散模型（Song等，2021；Rombach等，2022；Ho等，2020；Batzolis等，2021）最近因其作为逼真生成器的卓越性能而受到重视。扩散模型具有非常高的灵活性，但需要多步采样阶段来去噪样本（Song等，2021）。为了解决这一缺点，Song等人（2023年）提出了一致性模型，这些模型经过设计可以进行少步生成。在本文中，我们将一致性模型移植到基于仿真的推断中（见图1），以实现可扩展和快速的神经后验估计的前所未有的组合。我们的贡献包括：

1. 我们将一致性模型调整为基于仿真的贝叶斯推断，并提出了一致性模型后验估计（CMPE）；

2. 我们阐明了一致性模型在基于仿真的推断中的基本优势：表达自由形式的架构和快速推断；

3. 我们证明了CMPE在三个基准实验（见图2）、高维贝叶斯去噪和肿瘤球体模型中均优于归一化流和流匹配。

2.预备工作和相关工作

本节概述了基于模拟推断的方法，包括基于正规化流、流匹配和基于得分的扩散模型。

2.1.基于模拟推断（SBI）

SBI方法的定义属性是它们仅依赖于从数据生成过程p(x,θ)中采样的能力，而不像依赖于评估似然函数p(x|θ)的基于似然的方法。与此密切相关的，用于SBI的概率神经方法可以分为顺序方法（Papamakarios＆Murray，2016; Lueckmann等，2017; Greenberg等，2019; Durkan等，2020; Deistler等，2022; Sharrock等，2022）与摊销方法（Ardizzone等，2019; Gonc¸alves等，2020; Radev等，2020; Pacchiardi＆Dutta，2022; Avecilla等，2022; Dax等，2023b; Geffner等，2022; Radev等，2021a; Elsem¨uller等，2023b）。

顺序推断算法通过迭代改进先验p(θ)来生成在目标观测附近的模拟。因此，它们不是摊销的，因为每个新的观察数据集都需要昂贵的重新训练神经逼近器，以适应特定的目标观察。

相反，摊销方法训练神经逼近器以泛化整个模型的先验预测空间。这使我们能够查询逼近器，以获取假定来自模型范围的任何新数据集。事实上，摊销可以跨模型的任何组件进行，包括多个数据集（Gonc¸alves等，2020）和上下文因素，例如数据集中的观测数量（Radev等，2020），异构数据源（Schmitt等，2023b）甚至不同的概率模型和数据配置（Elsem¨uller等，2023a）。我们当前的工作处于摊销设置中。

2.2. 用于神经后验估计的正规化流

传统上，用于SBI Neural posterior estimation 的神经后验估计（NPE）方法依赖于条件正规化流，用于从模拟的参数和数据对（θ，x）中学习神经密度估计器pϕ(θ | x)（Rezende＆Mohamed，2015; Papamakarios等，2021）。可学习的神经网络权重被视为ϕ。正规化流定义了参数θ和具有简单基本分布（例如，高斯分布）的潜变量z之间的一次性变换。通过最大似然训练，正规化流可以学习闭式近似密度：

在SBI中，两个著名的正规化流架构是仿射耦合流（ACF；Dinh等，2016）和神经样条流（NSF；Durkan等，2019）。这些架构的一个主要缺点是反演要求：它将神经网络的设计空间限制为具有廉价雅可比计算的可逆函数。接下来，我们将简要回顾最近的一些工作，这些工作首创了使用多步、自由形式架构进行SBI的方法，灵感来自这些模型在各种生成任务中的成功应用（Batzolis等，2021; Rombach等，2022）。

2.3. 用于后验估计的流匹配

Dax等（2023b）应用了流匹配技术（Lipman等，2023; Liu等，2022）在SBI中，这种方法被称为流匹配后验估计（FMPE flow matching posterior estimation ）。流匹配消除了反演限制，并使得可以使用自由形式的神经网络。FMPE基于最优输运，在此过程中，基本分布和目标分布之间的映射由每个时间步t∈[0,1]上的样本空间上的矢量场驱动的连续过程参数化。这里，t = 1时的分布是单位高斯分布N（0，I），而t = 0时的分布是目标后验分布。FMPE损失将NPE目标（Eq. 1）中的最大似然项替换为条件流匹配目标。

其中，µϕ表示由可训练参数ϕ参数化的自由形式神经网络的条件矢量场，而ut表示边际矢量场，可以简单地表示为ut(θt | θ) = θ1 − θ，其中t ∈ [0, 1]（Liu等，2022）。我们通过在t ∈ [0, 1]上反向解决dθt = −µ(θt, t; x)，从初始噪声样本θ1 ∼ N (0, I)开始，获得后验抽样θ0 ∼ p(θ | x)。我们可以使用任何现成的ODE求解器将噪声θ1转换为从近似后验中抽取的θ0。原则上，ODE求解器中的步数K可以通过设置步长dt = 1/K 来调整。这会增加采样速度，但FMPE并非旨在优化少步采样性能。我们在使用鼓励直线路径解的矫正流进行实验时证实了这一点（Liu等，2022）。

2.4. 神经后验分数估计

另一种利用自由形式网络进行基于模拟推断的方法是神经后验分数估计（NPSE neural posterior score estimation ）。这个方向使用条件得分驱动的扩散模型（Song＆Ermon，2019; Song等，2021）学习后验分布。Sharrock等（2022）提出了顺序NPSE，以顺序（非摊销）方式训练基于分数的扩散模型。在略有不同的重点方面，（Geffner等，2022）将后验分布分解，并为较大数据集的子集（甚至是单个观察结果）学习（扩散的）后验分数。随后，通过组合学习到的分数，将子集的信息聚合以近似整个数据集的后验分布。

关键是，这两种方法都依赖于基于分数的扩散模型的基本公式：它们通过随机微分方程逐渐通过扩散过程

扩散目标分布。

具有漂移系数 𝜇 、扩散系数 𝜎 、时间 𝑡∈[0,𝑇] 和布朗运动 {𝐰𝑡}𝑡∈[0,𝑇] 。在每次 𝑡 时，以 𝐱 为条件的 𝜽 的当前（扩散）分布表示为 𝑝𝑡⁢(𝜽|𝐱) 。至关重要的是， 𝑡=0 处的分布等于目标后验分布 𝑝0⁢(𝜽|𝐱)≡𝑝⁢(𝜽|𝐱) ，而 𝑡=𝑇 处的分布是纯噪声（见下文）。宋等人。(2021) 证明存在一个常微分方程（“概率流 ODE”），其在时间 𝑡 时的解轨迹按照 𝑝𝑡⁢(𝜽|𝐱) 分布：

其中，

是

的得分函数。这个微分方程通常设计为在扩散过程后产生一个球形高斯噪声分布

。由于我们无法访问目标后验分布p(θ | x)，基于分数的扩散模型通过分数匹配训练一个时间相关的分数网络sϕ(θt, t, x) ≈ ∇ log pt(θt | x)，并将其插入到方程4中。设定µ(θt, t; x) = 0和σt = √(2t)，概率流ODE的估计变为dθt = −tsϕ(θt, x, t)dt (Karras等，2022)。最后，我们可以从噪声分布生成一个随机抽样θT ∼ N (0, T^2 I)，并向后解概率流ODE，得到一个轨迹

。轨迹的终点θ0表示从近似后验p0(θ0 | x) ≈ p(θ | x)中抽取的一个样本。

数值稳定性 通常在固定的小正数t = ε处停止求解器，以防止数值不稳定性（Song等，2023），因此我们使用θε表示从近似后验中抽取的样本。为简单起见，我们也将θε称为轨迹的起点。

3. 一致性模型后验估计

扩散模型有一个关键缺点：在推理时，它们需要求解许多微分方程，这会减慢采样速度。这在SBI（模拟贝叶斯推理）应用中特别麻烦，因为这些应用通常需要为成千上万个数据集生成成千上万个样本。一致性模型（Song等，2023）通过一种新的生成模型来解决这个问题，该模型支持单步和多步采样。下面，我们将Song等人（2023）的一致性模型公式扩展，以适应条件信息，如数据x或联合学习的嵌入h(x)（Radev等，2020）。

3.1. 条件一致性模型

一致性函数

将解轨迹

上的点映射到给定固定条件变量x和概率流ODE（微分方程）在公式4中的轨迹起点θε。为了通过已建立的基于得分的扩散模型架构实现这一点，我们可以使用一个通过跳跃连接参数化的自由形式神经网络Fϕ(θ, t; x)

其中

是可微的，并且满足边界条件

和

。一致性模型最初是作为扩散模型的蒸馏技术而提出的。然而，song等人。（2023）表明单独训练一致性模型是可能的，我们的方法基于它们的直接方法。

一旦一致性模型训练完毕，从近似后验生成样本变得简单，只需从噪声分布中抽取样本 θT ∼ N (0, T²I)，然后将其转换为目标分布的样本，就像在标准扩散模型中一样。然而，与扩散模型不同的是，我们不需要为这种转换求解一系列的微分方程。相反，我们可以使用学习到的一致性函数 fϕ 来获得一步的目标样本 θε = fϕ(θT, T; x)。此外，用一致性模型进行推断实际上并不限于一步采样。事实上，通过迭代采样过程，可以实现多步生成，我们将在下文中描述这一过程。对于时间点序列

和初始噪声

，我们计算

对于 k = K − 1, K − 2, ... , 1，其中 zk ∼ N(0, I)，并且 K − 1 是采样步骤的数量（Song & Dhariwal, 2023；Song 等人，2023）。生成的样本通常比一步采样的样本更好：调整步骤数量可以在计算量和样本质量之间进行权衡（Song 等人，2023）。

3.2 一致性模型用于SBI

一致性模型最初是为图像生成而开发的，但可以应用于学习任意分布。自由形式的架构使得可以集成针对数据 x 和参数 θ 的专门化架构。由于采样所需的遍历次数较少（与流匹配和扩散模型相比），可以在保持低推理时间的同时使用更复杂的网络。理论上，一致性模型结合了两者的优点（见表1）：无约束的网络可最佳地适应参数结构和数据模式，同时通过少量的网络遍历实现快速推理速度。

这是以显式可逆性为代价的，这导致了计算后验密度的困难，如第3.4节所述。根据Cranmer等人（2020）的分类法，我们将我们的方法称为一致性模型后验估计（CMPE）。虽然我们专注于后验估计，但使用一致性模型进行似然模拟是我们工作的自然扩展。

由于其根本不同的训练目标，CMPE不仅仅是FMPE的更快版本。相反，它展示了与FMPE的定性差异，超越了更快的采样速度：(i) 在实验4中，我们表明CMPE比FMPE对神经网络架构的依赖性更低，这使得当任务的最佳架构未知时，CMPE成为一个有前景的替代方案；(ii) 在我们的实验中，我们观察到CMPE在低数据环境中表现良好，使其在训练数据稀缺时成为一种有吸引力的方法。事实上，有限的数据可用性是科学（例如，分子动力学；Kadupitiya等人，2020）和工程（Heringhaus等人，2022）中的复杂模拟程序的常见限制因素。

3.3. 优化目标

3.4. 密度估计

使用变量变化公式，我们可以将单步样本的后验密度表示为

3.5 选择采样步骤数

一致性模型的设计允许一步采样（见方程6）。然而，在实际应用中，使用两步采样显著提高了图像生成任务的样本质量（Song & Dhariwal, 2023）。我们观察到，使用大约 K = 5 到 10 步的少步骤采样在样本质量和计算量之间提供了最佳平衡，特别是在低维问题中。这大致相当于实验1-3中一步估计器（如仿射耦合流或神经样条流）的速度（见图3）。采样步骤数可以在推理时选择，因此实践者可以轻松调整这种权衡以适应特定情况。我们注意到，接近训练中使用的最大离散化步骤数时，可能会导致过于自信的后验分布。当训练期间超过最大离散化步骤数时，这种效应变得更加明显，因此我们基于实验证据不建议这样做。

4. 实证评估

我们的实验涵盖了SBI的三个基本方面。首先，我们在三个具有双峰后验分布的低维实验中进行了广泛评估，这些实验来自逆问题的基准套件（Lueckmann et al., 2021；Kruse et al., 2021）。基于模拟的训练阶段基于一个固定的训练集

，该训练集包含M个数据集

及其对应的数据生成参数（即，真实值）

。其次，我们专注于一个图像去噪示例，该示例在SBI背景下作为一个足够高维的案例研究（Ramesh et al., 2022；Pacchiardi & Dutta, 2022）。第三，我们将我们的CMPE方法应用于一个计算上具有挑战性的肿瘤球体生长科学模型，并展示了其在一个非平凡科学模拟器中的卓越性能（Jagiella et al., 2017）。我们使用BayesFlow Python库实现所有实验，该库用于摊销贝叶斯工作流（Radev et al., 2023b）。

评估指标 我们基于公认的指标评估实验结果，以衡量结果的准确性和校准性。所有指标均在一个测试集上计算，该测试集包含 J 个未见过的实例

。在以下内容中，S 表示我们为每个实例 J 抽取的（近似）后验样本的数量。首先，易于理解的均方根误差（RMSE）指标量化了测试集中近似后验样本的偏差和方差，如下所示：

其次，我们估计近似样本与参考后验样本之间的平方最大均值差异（MMD；Gretton et al., 2012），这是一种基于核的分布距离。第三，作为SBI中广泛应用的指标，C2ST分数使用一个MLP分类器来区分近似后验和参考后验的样本。分类器的测试准确度即为C2ST分数，范围从0.5（最差；样本无法区分）到1（最好；样本完全可分）。最后，通过基于模拟的校准（SBC；Talts et al., 2018）来评估不确定性校准：对于每个分位数 q ∈ (0, 1)，真实后验 p(θ | x) 的所有不确定性区间

都被很好地校准。

其中

是指示函数。偏离方程(12)的情况表明近似后验的校准不足。期望校准误差（ECE）汇总了20个线性间隔分位数 q 的中心可信区间的中位SBC误差，并在整个测试集中取平均值。

竞争方法 我们比较了仿射耦合流（ACF；Dinh et al., 2016）、神经样条流（NSF；Durkan et al., 2019）、流匹配后验估计（FMPE；Dax et al., 2023b）和一致性模型后验估计（CMPE；我们的）。我们为FMPE和CMPE使用相同的自由形式神经网络以确保可比性。

4.1 实验1：高斯混合模型

我们在二维高斯混合模型上展示CMPE，该模型具有两个对称分量（Geffner et al., 2022；Schmitt et al., 2023a）。对称分量权重相等，方差相同。

结果 如图2所示，ACF和NSF均未能拟合具有分离模式的双峰后验分布。FMPE成功形成了分离的后验模式，但由于过窄的后验分布表现出明显的过度自信。CMPE的视觉采样性能优于所有其他方法。在这项任务中，我们观察到，与其他近似方法相比，CMPE不需要在速度或性能之间进行选择。相反，CMPE可以在速度和性能方面同时优于所有其他方法，这从J = 100个测试实例中的较低C2ST分数可以看出（见图3）。如果我们能容忍较慢的采样，CMPE在K = 10个推理步骤时能达到最佳性能。最值得注意的是，CMPE在性能上大幅超过了1000步的FMPE，尽管后者的速度大约慢了75倍。

4.2 实验2：两个月亮

结果 在M = 1024个样本的模拟预算下训练，CMPE始终能够探索两个新月形，并成功捕捉到后验分布的局部模式（见图2，中排）。仿射耦合流和神经样条流都未能完全分离模式。特别是，如果我们试图通过减少推理期间的采样步骤来实现FMPE的快速采样速度，FMPE在30个采样步骤下表现出明显的过度自信，并且在10个采样步骤下后验近似显著不足。相反，CMPE在少步采样情况下表现出色。图4显示了所有架构在更大的训练预算下受益。CMPE在10步采样的情况下，在最多M = 4096个训练实例的低数据和中等数据情况下表现出最佳性能。在最大的训练预算M = 8192实例下，多步FMPE表现优于其他近似方法。然而，请记住，在这个任务中，多步FMPE的速度比ACF、NSF和CMPE大约慢30-70倍。

4.3. 实验3：逆运动学

4.4 实验4：贝叶斯去噪

本实验展示了 CMPE 在高维逆问题上的可行性，即在 Fashion MNIST 数据集上进行贝叶斯去噪。未知参数

是原始图像展平后的向量，观测

是模糊且展平后的清晰图像，来自模拟嘈杂相机 (Ramesh et al., 2022；Pacchiardi & Dutta, 2022；Radev et al., 2023a)。我们在标准和小数据制度下比较了 CMPE 和 FMPE。由于这两种方法都允许自由形式的架构，因此我们可以评估架构选择对结果的影响。因此，我们使用次优的简单架构和已建立的 U-Net (Ronneberger et al., 2015) 架构评估这两种方法。

神经架构 简单架构由一个卷积神经网络 (CNN；LeCun et al., 2015) 组成，将观测转换为潜在摘要统计的向量。我们将输入向量、摘要统计和时间嵌入串联，并将它们馈送到由每层包含 2048 个单元的四层隐藏层的多层感知器 (MLP) 中。这对应于一种情况，即观测的结构（即图像数据）是已知的，但参数的结构是未知的或不适用于专门的网络架构。

然而，在这个示例中，我们可以利用我们的参数是图像的先验知识。具体来说，我们可以通过选择适用于图像处理的 U-Net 架构（即 Nain，2022 的调整版本）将归纳偏差纳入我们的网络架构中。再次，CNN 学习了嘈杂观测的摘要向量，然后将其与时间嵌入串联以用作神经密度估计器的条件向量。

结果 我们报告了两种方法和架构的聚合 RMSE、MMD 以及每个样本的时间（见表2）。FMPE 无法为简单架构生成良好的样本，而 CMPE 即使在这种次优设置下也能产生可接受的样本。这种对次优架构的降低敏感性可能在高维问题中成为一个有价值的特性，其中参数的结构无法利用。U-Net 架构为两种方法都提供了良好的样本质量，突显了自由形式架构的优势。MMD 值与对样本质量的视觉评估很好地吻合，因此我们认为它是比较结果的一个信息丰富的指标（有关可视检查，请参见第 C.4 节）。配合大型训练集，U-Net 架构为 CMPE 和 FMPE 提供了详细且多样的相似质量的样本。在所有情况下，CMPE 的推理时间明显较低，因为仅需要两次网络传递进行采样，同时实现更好或具有竞争力的质量。

4.5 实验5：肿瘤球体生长

我们以一个复杂的多尺度模型结束我们的实证评估，这个模型用于描述二维肿瘤球体生长（Jagiella等人，2017），其中有7个未知参数（取自pyABC，2017）。关键是，使用标准硬件运行该模型的仿真相当昂贵（在消费级计算机上单次运行约需1分钟），因此人们希望能够使用有限的离线训练预算提供合理的估计。在这里，我们比较了本文中四种竞争方法在一个固定的训练集上的性能，该训练集包含M = 19,600个仿真，J = 400个仿真作为测试集来计算性能指标。在所有方法中，我们使用混合的LSTM-Transformer架构将可变长度的高维专家数据转换为固定长度的摘要向量 h(x)。

附录 C.5 提供了有关神经网络架构和训练超参数的更多细节。

结果 通过在400个未见测试实例上的更低的 RMSE 和 ECE 所反映的更好的准确性和校准指标，CMPE 在性能上优于其他替代神经方法（见表3）。简单的 ACF 的速度无法与其他方法匹敌。与自由形式竞争者 FMPE 的直接比较中，CMPE 同时表现出 (i) 稍微更高的准确性；(ii) 大幅改善的校准；以及 (iii) 更快的采样速度（请参见附录中的图6）。对于这个模型，FMPE 在高达 K = 100 推理步骤时没有达到令人满意的校准性能，因此表3报告了 K = 1000 步时的最佳 FMPE 结果。

5. 结论

我们提出了一致性模型后验估计（CMPE），这是一种先进的方法，能够在大型模型上执行准确的基于仿真的贝叶斯推断，并实现快速的推断速度。CMPE通过将扩散模型和流匹配的自由形式神经架构特性与高性能的少步采样相结合，增强了神经后验估计的能力。为了评估CMPE的有效性，我们将其应用于一组允许直观视觉检查的三个低维基准任务，以及一个高维贝叶斯去噪实验和一个科学肿瘤生长模型。在所有实验中，通过对后验准确性、校准和推断速度进行整体评估，CMPE始终表现出色。

未来的工作可能会致力于进一步减少采样步骤，以实现一步推断，例如通过广泛的自动化超参数优化或为CMPE设计的定制训练方案。总的来说，我们的结果证明了CMPE作为一种新的基于仿真的推断工具的潜力，使其成为科学和工程领域基于仿真的推断工作流的新竞争者。