Mixed neural posterior estimation for simulators with discrete and continuous parameters

技术摘要：针对具有离散与连续参数的模拟器的混合神经后验估计

问题陈述
神经后验估计（NPE）已成为基于模拟的推断（SBI）的标准工具，能够针对具有不可处理似然函数的复杂模拟器，实现快速、摊销的参数推断。然而，标准的 NPE 方法假设参数空间是连续的，通常使用归一化流或扩散模型将后验建模为连续概率密度。这一假设对于一大类科学模拟器而言并不适用，因为它们的参数空间是混合的，既包含连续维度（例如速率、电导），也包含离散维度（例如分类变量、模型选择索引、变点位置）。

现有的混合推断方法面临显著局限。标准的马尔可夫链蒙特卡洛（MCMC）方法需要仔细调整，无法在不同观测值之间摊销推断成本，且通常难以处理离散变量（例如需要解析边缘化或专门的非连续哈密顿动力学）。最近的基于扩散的 SBI 方法尝试将离散参数嵌入连续空间或将其视为模型索引，但这些方法计算负担沉重，或需要复杂的黎曼嵌入。因此，当前 SBI 工具箱对具有混合参数空间的模拟器支持仍显不足。

方法：混合神经后验估计（MNPE）
作者提出了混合神经神经后验估计（MNPE），这是 NPE 的一种扩展，通过因子化推断网络联合处理离散和连续参数。

1. 因子化策略：
   离散参数 $\theta_d$ 和连续参数 $\theta_c$ 上的联合后验 $p(\theta_d, \theta_c | x)$ 被因子化为：
   $$p(\theta_d, \theta_c | x) = p(\theta_d | x) \cdot p(\theta_c | \theta_d, x)$$
   这将问题分解为两个子任务：离散空间上的分类，以及以离散状态为条件的连续空间中的条件密度估计。

2. 网络架构：

   • 离散组件： 项 $p(\theta_d | x)$ 使用**掩码自回归密度估计器（MADE）**进行建模。该自回归网络输出每个离散维度的类别概率，条件于 preceding 离散维度和观测值 $x$。这使得该方法能够处理任意数量的离散参数，且各参数的类别数量可变。
   • 连续组件： 项 $p(\theta_c | \theta_d, x)$ 使用标准的条件生成模型（例如神经样条流或扩散模型）进行建模，条件于观测值 $x$ 和离散参数 $\theta_d$。
   • 训练目标： 网络联合训练以最小化混合参数的负对数概率：
     $$\mathcal{L} = -\log q(\theta_d, \theta_c | x) = -\log q(\theta_c | \theta_d, x) - \log q(\theta_d | x)$$
     连续部分使用标准的流似然，而离散部分简化为交叉熵损失。

3. 校准诊断：
   为了在没有参考解的情况下评估混合后验的质量，作者提出了一种综合诊断框架：

   • 连续维度： 使用基于秩统计的**基于模拟的校准（SBC）**来检查边缘秩分布的均匀性。
   • 离散维度： 使用**期望校准误差（ECE）**和可靠性图，将离散推断视为分类问题。
   • 基线： 作者推导了这两个指标的实证有限样本基线，以区分实际的校准不良与统计噪声，承认在有限样本下无法实现完美校准。

关键结果
作者在三个复杂度递增的模拟器上验证了 MNPE：

1. 可处理的 Gaussian 玩具示例：

   • 一个包含一个连续和一个离散潜在变量的模拟器，通过共享观测值耦合。
   • MNPE 准确恢复了分析真实后验，包括离散状态不确定时的双峰机制。
   • 分类器双样本检验（C2ST）分数在约 1,000 次模拟后收敛至随机水平（约 0.5），表明其与真实后验不可区分。

2. 串联排队模拟器：

   • 一个包含三个连续速率和两个离散服务器数量（$|D|=25$）的模拟器。
   • 将 MNPE 与 MCMC 参考（PyMC 和 NumPyro）进行比较，其中离散维度被解析边缘化。
   • 在 100,000 次模拟下，MNPE 后验在连续和离散维度上均与 MCMC 参考匹配。
   • 校准诊断（SBC 秩直方图和 ECE 可靠性图）证实后验校准良好，连续维度接近有限样本基线，离散维度匹配完美校准分类器的预期误差。

3. Hodgkin–Huxley 生物物理模拟器：

   • 一个用于神经活动的“黑盒”模拟器，具有不可处理的似然函数，涉及连续电导和离散通道身份（例如特定离子通道的存在）。
   • MNPE 成功推断出连续电导和分类通道类型的联合后验。
   • 该方法在不同训练预算下产生校准良好的后验，使得能够分析离子通道之间的补偿机制（例如 L 型钙通道与钾电导之间的关系）。

意义与主张
本文主张 MNPE 填补了 SBI 的关键空白，为具有混合参数空间的模拟器提供了统一的、摊销的框架。其主要贡献包括：

• 通用性： 支持任意数量的离散和连续参数，并与任何连续生成网络（流、扩散）兼容。
• 效率： 与 MCMC 不同，MNPE 摊销了模拟和训练的成本，允许通过单次前向传递对新观测值进行快速推断。
• 校准： 提供了一种原则性方法来评估混合设置中的校准，结合基于秩和基于分类的指标以及实证基线。
• 实现： 该方法已实现在公开可用的 sbi Python 包中，使领域科学家无需大量调整即可使用。

作者指出了 NPE 家族固有的局限性，例如复杂后验需要大量训练数据集，以及对模型误设的敏感性。他们还澄清，其校准程序评估的是每个维度的边缘校准，而非完整的联合校准，后者在混合空间中仍然具有挑战性。