Learned harmonic mean estimation of the marginal likelihood for multimodal posteriors with flow matching

技术摘要：基于流匹配（Flow Matching）的习得性调和平均值估计用于多峰后验分布

问题陈述
边缘似然（或称贝叶斯证据），记为 $z = p(d|M)$，是贝叶斯模型比较中的一个基本量，通过贝叶斯因子实现定量评估。然而，计算该量需要对整个参数空间进行高维积分，这对于复杂的、高维的模型而言在计算上是难以实现的。虽然马尔可夫链蒙特卡洛（MCMC）和哈密顿蒙特卡洛（HMC）方法擅长生成高维后验样本，但它们本身并不直接提供边缘似然。相反，嵌套采样（Nested Sampling）虽然能提供证据，但其采样策略受限于特定的嵌套层次结构。

“习得性调和平均值估计器”（Learned harmonic mean estimator）此前被引入，用于仅利用后验样本来估计边缘似然，而无需关心所使用的采样方法。该方法依赖于学习一个内部目标密度 $\phi(\theta)$，该密度包含在后验分布内，且具有更薄的尾部。然而，以往使用离散归一化流（Discrete Normalizing Flows）或特定密度估计器的实现，在处理高度多峰的后验分布时表现不佳。这些架构往往无法捕捉复杂的拓扑结构，除非进行启发式修改（例如手动构建高斯混合作为基分布），或者会导致估计器分配非零概率质量到密度桥接区域，从而导致估计器的方差变为无穷大。

方法论
本研究通过将**流匹配（Flow Matching）**集成到习得性调和平均值框架中，解决了先前密度估计器的局限性。其方法论如下：

1. 习得性调和平均值估计器： 该估计器利用后验样本 $\theta_i \sim p(\theta|d)$ 和学习到的目标密度 $\phi(\theta)$ 来近似倒数边缘似然 $\rho = z^{-1}$：
   $$\hat{\rho} = \frac{1}{N} \sum_{i=1}^{N} \frac{\phi(\theta_i)}{L(\theta_i)\pi(\theta_i)}$$
   为了使估计器具有有限方差，$\phi(\theta)$ 必须集中在后验分布的高概率区域，并且在尾部衰减得更快。

2. 流匹配架构： 作者没有使用离散归一化流，而是采用了通过流匹配训练的连续归一化流（Continuous Normalizing Flows, CNFs）。

   • 定义： CNF 通过一个随时间变化的速度场 $v_t(\theta)$，参数化了简单基分布（标准高斯分布）与目标后验分布之间的变换。
   • 训练目标： 网络通过无模拟回归目标进行训练。定义一个条件概率路径，即噪声样本 $\theta_0$ 与后验样本 $\theta_1$ 之间的直线：$\theta_t = (1-t)\theta_0 + t\theta_1$。网络最小化预测速度场与真实条件向量场 $(\theta_1 - \theta_0)$ 之间的均方误差。
   • 优势： 这种方法允许使用架构不受限的高容量神经网络（带有 Swish 激活函数的多层感知机 MLP），并避免了早期 CNF 方法在训练期间所需的昂贵的 ODE 模拟。

3. 温度缩放（Temperature Scaling）： 为了确保目标密度 $\phi(\theta)$ 的尾部比后验分布更薄，应用了一个温度参数 $T \in (0, 1)$ 来缩放流的基分布的方差。这使得流分布更加集中，从而满足调和平均值估计器的包含性要求。

核心贡献

• 新颖架构： 引入了基于流匹配训练的连续归一化流作为内部密度估计器。
• 处理多峰性： 证明了该架构可以准确捕捉复杂且高度多峰的后验结构，而无需对基分布进行启发式修改或微调。
• 可扩展性： 提供了一种与采样策略无关的方法（兼容 MCMC、HMC、NUTS、SBI 等），并能在中高维空间中运行。

实验结果
作者在两个具有挑战性的基准测试上验证了该方法：

1. Rastrigin 函数（2 维）：

   • 这是一个具有许多局部模态的标准基准。
   • 流匹配模型（6 层，256 个隐藏单元）在温度 $T=0.98$ 时成功捕捉到了所有模态。
   • 来自 100 次实验运行的结果表明，估计的边缘似然值是准确的，并且与通过数值积分获得的地面真值密切匹配。误差估计是可靠的。

2. 20 维高斯混合模型：

   • 由五个高斯分量组成的 20 维混合模型，这些分量是非重叠的且引入了相关性。
   • 模型（10 层，256 个隐藏单元）在 40,000 个样本上进行了训练。
   • 学习到的目标分布被证明正确地集中在后验分布内，捕捉到了具有挑战性的拓扑结构。
   • 在 100 次运行中的估计边缘似然值与解析地面真值一致，证实了该方法即使在这种复杂的、中等维度的情况下也具有准确性和鲁棒的误差估计。

意义与主张
论文声称，通过利用流匹配训练的 CNF 的表达能力，习得性调和平均值估计器现在可以稳健地估计高度复杂的多峰后验分布的证据。其主要意义在于消除了以往处理多峰性时对启发式调整（如手工挑选的高斯混合）的依赖。作者断言，这种方法能够利用来自任何加速或可扩展采样方法的样本来计算证据，从而促进了像天体物理学这样数据丰富的领域中的贝叶斯模型比较，同时避免了传统嵌套采样的计算瓶颈或原始调和平均值估计器的不稳定性。

正如作者所指出的，未来的工作包括将该方法应用于现实世界的问题，如引力波天体物理学，并探索针对成像及场级宇宙学推断等极高维任务的进一步架构改进。