[论文解读] All-in-one simulation-based inference

Manuel Gloeckler, Michael Deistler+3 more|arXiv (Cornell University)|Apr 15, 2024

Simulation Techniques and Applications被引用 5

一句话总结

Simformer 是一种基于扩散模型的摊销SBI方法，使用一个对联合参数与数据输入的变换器来采样任意条件分布（后验、似然或其他），并处理非结构化数据、函数值参数和观测区间。

ABSTRACT

Amortized Bayesian inference trains neural networks to solve stochastic inference problems using model simulations, thereby making it possible to rapidly perform Bayesian inference for any newly observed data. However, current simulation-based amortized inference methods are simulation-hungry and inflexible: They require the specification of a fixed parametric prior, simulator, and inference tasks ahead of time. Here, we present a new amortized inference method -- the Simformer -- which overcomes these limitations. By training a probabilistic diffusion model with transformer architectures, the Simformer outperforms current state-of-the-art amortized inference approaches on benchmark tasks and is substantially more flexible: It can be applied to models with function-valued parameters, it can handle inference scenarios with missing or unstructured data, and it can sample arbitrary conditionals of the joint distribution of parameters and data, including both posterior and likelihood. We showcase the performance and flexibility of the Simformer on simulators from ecology, epidemiology, and neuroscience, and demonstrate that it opens up new possibilities and application domains for amortized Bayesian inference on simulation-based models.

研究动机与目标

• 激发对更灵活、高效的摊销SBI的需求，该方法能处理不规则/结构化数据和函数值参数。
• 开发一个单一模型，能够对参数与数据的联合分布的任意条件进行采样。
• 利用扩散模型和注意力屏蔽的变换器来编码依赖关系，并允许有条件/推断输出。
• 在生态学、流行病学和神经科学等多样化的仿真器上展示该方法，以体现广泛的适用性。

提出的方法

• 为所有变量（参数和数据）构建基于令牌的表示，作为输入给变换器。
• 使用去噪分数匹配，在联合分布 p(theta, x) 上训练基于分数的扩散模型。
• 引入注意力掩码 M_E 以编码依赖结构并允许采样任意条件。
• 通过引导扩散实现对区间和约束的条件化，以从受约束的分布中采样。
• 通过对条件掩码的随机化来训练模型，使其输出联合分布的任意条件。
• 从单个训练好的网络中实现后验、似然以及其他条件的采样。

实验结果

研究问题

• RQ1 单一的扩散-变换器模型是否能准确地对联合分布 p(theta, x) 的任意条件进行采样？
• RQ2 引入依赖结构掩码是否提高SBI任务的准确性和效率？
• RQ3 该方法是否能处理非结构化数据、缺失值以及函数值（无限维）参数？
• RQ4 引导扩散如何实现对观测区间或附加约束的条件化？
• RQ5 该方法在生态学、流行病学和神经科学等多样化的仿真器上是否有效？

主要发现

• 在给定仿真预算下，Simformer 在基准SBI任务中的准确性优于神经后验估计。
• 使用领域知情的注意力掩码提高了准确性，尤其是在依赖结构较稀疏时。
• 平均而言，Simformer 在任务和观测方面所需的仿真次数比NPE少约10倍。
• 该模型能够采样任意条件，包括后验和似然，并处理非结构化数据和函数值参数。
• 引导扩散使对观测区间及其他约束的条件化成为可能，影响后验和后验预测样本。
• 该方法在 Lotka-Volterra 的非结构化观测、SIRD 的无限维参数空间以及带能量约束的 Hodgkin-Huxley 中实现了准确推断。