Generalized Bayesian Inference for Scientific Simulators via Amortized Cost Estimation

Richard Gao,Michael Deistler,Jakob H. Macke
DOI: https://doi.org/10.48550/arXiv.2305.15208
2023-11-02
Abstract:Simulation-based inference (SBI) enables amortized Bayesian inference for simulators with implicit likelihoods. But when we are primarily interested in the quality of predictive simulations, or when the model cannot exactly reproduce the observed data (i.e., is misspecified), targeting the Bayesian posterior may be overly restrictive. Generalized Bayesian Inference (GBI) aims to robustify inference for (misspecified) simulator models, replacing the likelihood-function with a cost function that evaluates the goodness of parameters relative to data. However, GBI methods generally require running multiple simulations to estimate the cost function at each parameter value during inference, making the approach computationally infeasible for even moderately complex simulators. Here, we propose amortized cost estimation (ACE) for GBI to address this challenge: We train a neural network to approximate the cost function, which we define as the expected distance between simulations produced by a parameter and observed data. The trained network can then be used with MCMC to infer GBI posteriors for any observation without running additional simulations. We show that, on several benchmark tasks, ACE accurately predicts cost and provides predictive simulations that are closer to synthetic observations than other SBI methods, especially for misspecified simulators. Finally, we apply ACE to infer parameters of the Hodgkin-Huxley model given real intracellular recordings from the Allen Cell Types Database. ACE identifies better data-matching parameters while being an order of magnitude more simulation-efficient than a standard SBI method. In summary, ACE combines the strengths of SBI methods and GBI to perform robust and simulation-amortized inference for scientific simulators.
Machine Learning
What problem does this paper attempt to address?
The problem that this paper attempts to solve is how to conduct more robust inference when using scientific simulators for inference, when the model cannot accurately reproduce the observed data (i.e., the model is mis - specified) or when researchers are more concerned about the quality of the predictive simulation rather than the exact Bayesian posterior distribution. Specifically, traditional simulation - based inference (SBI) methods usually target the Bayesian posterior, which is effective when the model can accurately describe the data - generation process. However, in many practical applications, the model is often a rough approximation and may even be unable to generate simulation results that match the observed data. In this case, directly targeting the Bayesian posterior may be too strict and limit the exploration range of the parameter space. To address this challenge, the paper introduces a new method - conducting generalized Bayesian inference (GBI) through amortized cost estimation (ACE). GBI improves the robustness to model mis - specification by using a cost function to evaluate the goodness of parameters relative to data, rather than the traditional likelihood function. However, existing GBI methods usually need to run multiple simulations in each inference process to estimate the cost function, which is computationally very expensive, especially when dealing with complex simulators. The core idea of ACE is to train a neural network to approximate the cost function, which is defined as the expected distance between the simulation results generated by the parameters and the observed data. Once trained, this neural network can be used for GBI posterior inference of any observed data without running additional simulations. This method not only improves computational efficiency but also can provide more reasonable parameter estimates in the case of model mis - specification. The paper demonstrates the effectiveness of ACE through multiple benchmark tasks and practical applications (such as parameter inference in the Hodgkin - Huxley neuron model), proving that ACE can significantly reduce the number of required simulations while ensuring the quality of the predictive simulation.