SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Richard Gao\*1,2
r.dg.gao@gmail.com
Michael Deistler\*1,2
michael.deistler@uni-tuebingen.de
Jakob H. Macke $^{1,2,3}$
jakob.macke@uni-tuebingen.de
<sup>1</sup>Machine Learning in Science, Excellence Cluster Machine Learning, University of Tübingen
$^{2}$ Tübingen AI Center
$^{3}$ Department Empirical Inference, Max Planck Institute for Intelligent Systems
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.
Mechanistic models expressed as computer simulators are used in a wide range of scientific domains, from astronomy, geophysics, to neurobiology. The parameters of the simulator,  $\theta$ , encode mechanisms of interest, and simulating different parameter values produces different outputs, i.e.,  $\sin (\theta_{i}) \rightarrow \mathbf{x}_{i}$  where each model-simulation  $\mathbf{x}_i$  can be compared to experimentally observed data,  $\mathbf{x}_o$ . Using such simulators, we can quantitatively reason about the contribution of mechanisms behind experimental
measurements. But to do so, a key objective is often to find all those parameter values that can produce simulations consistent with observed data.
One fruitful approach towards this goal is simulation-based inference (SBI) [1], which makes it possible to perform Bayesian inference on such models by interpreting simulator outputs as samples from an implicit likelihood [2],  $\mathbf{x} \sim p(\mathbf{x}|\boldsymbol{\theta})$ . Standard Bayesian inference targets the parameter posterior distribution given observed data, i.e.,  $p(\boldsymbol{\theta}|\mathbf{x}_o) = \frac{p(\mathbf{x}_o|\boldsymbol{\theta})p(\boldsymbol{\theta})}{p(\mathbf{x}_o)}$ , where  $p(\boldsymbol{\theta})$  captures prior knowledge and constraints over model parameters, and the likelihood function  $p(\mathbf{x}_o|\boldsymbol{\theta})$  is evaluated as a function of  $\boldsymbol{\theta}$  for a fixed  $\mathbf{x}_o$ . SBI methods can differ in whether they aim to approximate the likelihood [3, 4, 5, 6] or the posterior directly [7, 8, 9], and can be amortized, i.e., do not require new simulations and retraining for new data [10, 11]. In the end, each method provides samples from the posterior, which are all, in theory, capable of producing simulations that are identical to the observation we condition on. Furthermore, by definition, the posterior probability of drawing a sample scales as the product of its prior probability and, critically, the likelihood that this sample can produce a simulation that is exactly equal to the observation.
However, targeting the exact posterior may be overly restrictive. In many inference scenarios, modelers are primarily interested in obtaining a diverse collection of parameter values that can explain the observed data. This desire is also reflected in the common usage of posterior predictive checks, where seeing predictive simulations that resemble the data closely (in some specific aspects) is used to gauge the success of the inference process. In particular, it is often clear that the scientific model is only a coarse approximation to the data-generating process, and in some cases even cannot generate data-matching simulations, i.e., is misspecified [12]. For example, in the life-sciences, it is not uncommon to use idealized, theoretically motivated models with few parameters, and it would be unrealistic to expect that they precisely capture observations of highly complex biological systems. In such cases, or in cases where the model is fully deterministic, it is nonsensical to use the probability of exactly reproducing the data. In contrast, it would still be useful to find parameter values that produce simulations which are 'close enough', or as close as possible to the data. Therefore, instead of sampling parameters according to how often they produce simulations that match the data exactly, many use cases call for sampling parameters according to how closely their corresponding simulations reproduce the observed data.
Generalized Bayesian Inference (GBI) [13] offers a principled way to do so by replacing the (log) likelihood function with a cost function that simply scores a parameter given an observation, such as the expected distance between  $\mathbf{x}_o$  and all possible simulations  $\mathbf{x}$  produced by  $\theta_i$  (Fig. 1). Several recent works have leveraged this framework to perform inference for models with implicit or intractable likelihoods, especially to tackle model misspecification: Matsubara et al. [14] use a Stein Discrepancy as the cost function (which requires the evaluation of an unnormalized likelihood and multiple i.i.d. data samples), and Cherief-Abdellatif et al. and Dellaporta et al. [15, 16] use simulator samples to estimate maximum mean discrepancy and directly optimize over this cost function via stochastic gradient descent (which
Figure 1: Estimating cost from simulations. Using the expected distance between simulated and target data as the cost function, GBI assigns high probability to parameter values that, on average, produce simulations that are close—but not necessarily equal—to the observation.
requires a differentiable simulator). More broadly, cost functions such as scoring rule estimators have been used to generalize approximate Bayesian computation (ABC) [17] and synthetic likelihood approaches [3, 18], where the Monte Carlo estimate requires (multiple) simulations from  $p(\mathbf{x}|\boldsymbol{\theta})$ .
Thus, existing GBI approaches for SBI either require many simulations to be run during MCMC sampling of the posterior (similar to classical ABC methods), or are limited to differentiable simulators. Moreover, performing inference for new observations requires re-running simulations, rendering such methods simulation-inefficient and expensive at inference-time, and ultimately impractical for scientific simulators with even moderate computational burden.
We here propose to perform GBI for scientific simulators with amortized cost estimation (ACE), which inherits the flexibility of GBI but amortizes the overhead of simulations by training a neural
network to predict the cost function for any parameter-observation pair. We first outline the GBI formalism in Sec. 2, then introduce ACE in Sec. 3. In Sec. 4, we show that ACE provides GBI posterior predictive simulations that are close to synthetic observations for a variety of benchmark tasks, especially when the simulator is misspecified. We showcase its real-world applicability in Sec. 5: using experimental data from the Allen Cell Types Database, ACE successfully infers parameters of the Hodgkin-Huxley single-neuron simulator with superior predictive performance and an order of magnitude higher simulation efficiency compared to neural posterior estimation [10]. Finally, we discuss benefits and limitations of GBI and ACE, and related work (Sec. 6).
To construct the GBI posterior, the likelihood,  $p(\mathbf{x}_o|\boldsymbol{\theta})$ , is replaced by a 'generalized likelihood function',  $L(\boldsymbol{\theta};\mathbf{x}_o)$ , which does not need to be a probabilistic model of the data-generating process, as long as it can be evaluated for any pair of  $\boldsymbol{\theta}$  and  $\mathbf{x}_o$ . Following the convention in Bissiri et al. [13], we define  $L(\boldsymbol{\theta};\mathbf{x}_o)\equiv e^{-\beta \ell (\boldsymbol{\theta};\mathbf{x}_o)}$ , where  $\ell (\boldsymbol {\theta};\mathbf{x}_o)$  is a cost function that encodes the quality of  $\boldsymbol{\theta}$  relative to an observation  $\mathbf{x}_o$ , and  $\beta$  is a scalar inverse temperature hyperparameter that controls how much the posterior weighs the cost relative to the prior. Thus, the GBI posterior can be written as
$$
p (\boldsymbol {\theta} | \mathbf {x} _ {o}) \propto e ^ {- \beta \ell (\boldsymbol {\theta}; \mathbf {x} _ {o})} p (\boldsymbol {\theta}). \tag {1}
$$
As noted previously [13], if we define  $\ell(\pmb{\theta}; \mathbf{x}_o) \equiv -\log p(\mathbf{x}_o|\pmb{\theta})$  (i.e., self-information) and  $\beta = 1$ , then we recover the standard Bayesian posterior, and "tempered" or "power" posterior for  $\beta \neq 1$  [19, 20]. The advantage of GBI is that, instead of adhering strictly to the (implicit) likelihood, the user is allowed to choose arbitrary cost functions to rate the goodness of  $\pmb{\theta}$  relative to an observation  $\mathbf{x}_o$ , which is particularly useful when the simulator is misspecified. Previous works have referred to  $\ell(\pmb{\theta}; \mathbf{x}_o)$  as risk function [21], loss function [13], or (proper) scoring rules when they satisfy certain properties [22, 18] (further discussed in Section 6.1). Here we adopt 'cost' to avoid overloading the terms 'loss' and 'score' in the deep learning context.
3 Amortized Cost Estimation for GBI
3.1 Estimating cost function with neural networks
In this work, we consider cost functions that can be written as an expectation over the likelihood:
$$
\ell (\boldsymbol {\theta}; \mathbf {x} _ {o}) \equiv \mathbb {E} _ {p (\mathbf {x} | \boldsymbol {\theta})} [ d (\mathbf {x}, \mathbf {x} _ {o}) ] = \int_ {\mathbf {x}} d (\mathbf {x}, \mathbf {x} _ {o}) p (\mathbf {x} | \boldsymbol {\theta}) d \mathbf {x}. \tag {2}
$$
Many popular cost functions and scoring rules can be written in this form, including the average mean-squared error (MSE) [13], the maximum mean discrepancy  $(\mathrm{MMD}^2)$  [23], and the energy score (ES) [22] (details in Appendix A2). While  $\ell(\theta; \mathbf{x}_o)$  can be estimated via Monte Carlo sampling, doing so in an SBI setting is simulation-inefficient and time-intensive, since the inference procedure must be repeated for every observation  $\mathbf{x}_o$ , and simulations must be run in real-time during MCMC sampling of the posterior. Furthermore, this does not take advantage of the structure in parameter-space or data-space around neighboring points that have been simulated.
We propose to overcome these limitations by training a regression neural network (NN) to learn  $\ell (\pmb {\theta};\mathbf{x}_o)$ . Our first insight is that cost functions of the form of Eq. 2 can be estimated from a dataset consisting of pairs of parameters and outputs—in particular, from a single simulation run per  $\pmb{\theta}$  for MSE, and finitely many simulation runs for  $\mathrm{MMD}^2$  and ES. We leverage the well-known property that NN regression converges to the conditional expectation of the data labels given the data: If we compute the distances  $d(\mathbf{x},\mathbf{x}_o)$  between a single observation  $\mathbf{x}_o$  and every  $\mathbf{x}$  in our dataset, then a neural network  $f_{\phi}(\cdot)$  trained to predict the distances given parameters  $\pmb{\theta}$  will denoise the noisy distance labels  $d(\mathbf{x},\mathbf{x}_o)$  and converge onto the desired cost  $f_{\phi}(\pmb {\theta})\rightarrow \ell (\pmb {\theta};\mathbf{x}_o) = \mathbb{E}_{p(\mathbf{x}|\pmb {\theta})}[d(\mathbf{x},\mathbf{x}_o)]$  approximating the cost of any  $\pmb{\theta}$  relative to  $\mathbf{x}_o$  (see Appendix A3.1 for formal statement and proof).
3.2 Amortizing over observations
As outlined above, a regression NN will converge onto the cost function  $\ell(\pmb{\theta}; \mathbf{x}_o)$  for a particular observation  $\mathbf{x}_o$ . However, naively applying this procedure would require retraining of the network
for any new observation  $\mathbf{x}_o$ , which prevents application of this method in time-critical or high-throughput scenarios. We therefore propose to amortize cost estimation over a target distribution  $p(\mathbf{x}_t)$ . Specifically, a NN which receives as input a parameter  $\theta$  and an independently sampled target datapoint  $\mathbf{x}_t$  will converge to  $\ell(\theta; \mathbf{x}_t)$  for all  $\mathbf{x}_t$  on the support of the target distribution (Fig. 2a,b), enabling estimation of the cost function for any pair of  $(\theta, \mathbf{x}_t)$  (see Appendix A3.2 for formal statement and proof).
Naturally, we use the already simulated  $\mathbf{x} \sim p(\mathbf{x})$  as target data during training, and therefore do not require further simulations. In order to have good accuracy on potentially misspecified observations, however, such datapoints should be within the support of the target distribution. Thus, in practice, we augment this target dataset with noisy simulations to broaden the support of  $p(\mathbf{x}_t)$ . Furthermore, if the set of observations (i.e., real data) is known upfront, they can also be appended to the target dataset during training. Lastly, to keep training efficient and avoid quadratic scaling in the number of simulations, we randomly subsample a small number of  $\mathbf{x}_t$  per  $\theta$  in each training epoch (2 in our experiments), thus ensuring linear scaling as a function of simulation budget. Fig. 2a,b summarizes dataset construction and network training for ACE (details in Appendix A4.1).
3.3 Sampling from the generalized posterior
Given a trained cost estimation network  $f_{\phi}(\cdot, \cdot)$ , an observed datapoint  $\mathbf{x}_o$ , and a user-selected inverse temperature  $\beta$ , the generalized posterior probability (Eq. 1) can be computed up to proportionality for any  $\pmb{\theta}$ :  $p(\pmb{\theta} | \mathbf{x}_o) \propto \exp(-\beta \cdot f_{\phi}(\pmb{\theta}, \mathbf{x}_o)) p(\pmb{\theta})$  and, thus, this term can be sampled with MCMC (Fig. 2c). The entire algorithm is summarized in Algorithm 1.
Algorithm 1: Generalized Bayesian Inference with Amortized Cost Estimation (ACE)
Inputs: prior  $p(\theta)$ , simulator with implicit likelihood  $p(\mathbf{x}|\boldsymbol{\theta})$ , number of simulations  $N$ , feedforward NN  $f_{\phi}$  with parameters  $\phi$ , NN learning rate  $\eta$ , distance function  $d(\cdot, \cdot)$ , noise level  $\sigma$ , number of noise-augmented samples  $S$ , inverse temperature  $\beta$ , number of target datapoints per  $\theta$ $N_{\mathrm{target}}$ ,  $K$  observations  $\mathbf{x}_o^{(1,\dots,K)}$ .  
Outputs:  $M$  samples from generalized posteriors given  $K$  observations.
Generate dataset:
sample prior and simulate:  $\pmb{\theta},\mathbf{x}\gets \{\pmb{\theta}_i\sim p(\pmb {\theta}),\mathbf{x}_i\sim p(\mathbf{x}|\pmb {\theta}_i)\}_{i:1\dots N}$  add noise and concatenate:  $\mathbf{x}_{\mathrm{target}} = [\mathbf{x},\mathbf{x}_{1\dots S} + \epsilon ,\mathbf{x}_o],\epsilon \sim \mathcal{N}(\mathbf{0},\sigma^2\mathbf{I})$
while not converged do
for  $(\theta ,\mathbf{x})$  in batch do  $\mathbf{x}_t^{\mathrm{used}}\gets$  sample  $N_{\mathrm{target}}$  datapoints from  $\mathbf{x}_{\mathrm{target}}$  for  $\mathbf{x}_t$  in  $\mathbf{x}_t^{\mathrm{used}}$  do  $\mathcal{L}\leftarrow \mathcal{L} + (f_{\phi}(\pmb {\theta},\mathbf{x}_t) - d(\mathbf{x},\mathbf{x}_t))^2$ $\phi \gets \phi -\eta \cdot \mathrm{Adam}(\nabla_{\phi}\mathcal{L})$  // and reset L to zero   
Sampling:   
for  $k\in [1,\dots,K]$  do Draw M samples, with MCMC, from:  $\exp (-\beta \cdot f_{\phi}(\pmb {\theta},\mathbf{x}_o^{(k)}))p(\pmb {\theta})$
3.4 Considerations for choosing the value of  $\beta$
We note that the choice of value for  $\beta$  is an important decision, though this is an issue not only for ACE but GBI methods in general [13]. A good "baseline" value for  $\beta$  is such that the average distance across a subset of the training data is scaled to be in the same range as the (log) prior probability, both of which can be computed on prior simulations. From there, increasing  $\beta$  sacrifices sample diversity for predictive distance, and as  $\beta$  approaches infinity, posterior samples converge onto the minimizer of the cost function. In practice, we recommend experimenting with a range of (roughly) log-spaced values since, as we show below, predictive sample quality tends to improve with increasing  $\beta$ .
Figure 2: Schematic of dataset construction, network training, and inference. (a-b) The neural network is trained to predict the distance between pairs of  $\mathbf{x}$  (red) and  $\mathbf{x}_{\mathrm{t}}$  (green), as a noisy sample of the cost function (i.e., expected distance) evaluated on  $\theta$  (grey) and  $\mathbf{x}_{\mathrm{t}}$ . (c) At inference time, the trained ACE network predicts the cost for any parameter  $\theta$  given observation  $\mathbf{x}_o$  (top row), which is used to evaluate the GBI posterior under different  $\beta$  (bottom row, darker for larger  $\beta$ ) for MCMC sampling without running additional simulations. The distance is well-defined and can be approximated even when the simulator is misspecified (dashed lines).
4 Benchmark experiments
Tasks We first evaluated ACE on four benchmark tasks (modified from Lueckmann et al. [24]) with a variety of parameter- and data-dimensionality, as well as choice of distance measure: (1) Uniform 1D: 1D  $\theta$  and  $\mathbf{x}$ , the simulator implements an even polynomial with uniform noise likelihood, uniform prior (Fig. 2c); (2) 2 Moons: 2D  $\theta$  and  $\mathbf{x}$ , simulator produces a half-circle with constant mean radius and radially uniform noise of constant width, translated as a function of  $\theta$ , uniform prior; (3) Linear Gaussian: 10D  $\theta$  and  $\mathbf{x}$ , Gaussian model with mean  $\theta$  and fixed covariance, Gaussian prior; (4) Gaussian Mixture: 2D  $\theta$  and  $\mathbf{x}$ , simulator returns five i.i.d. samples from a mixture of two Gaussians, both with mean  $\theta$ , and fixed covariances, one with broader covariance than the other, uniform prior.
For the first three tasks, we use the mean-squared error between simulation and observation as the distance function. For the Gaussian Mixture task, we use maximum mean discrepancy  $(\mathrm{MMD}^2)$  to measure the statistical distance between two sets of five i.i.d. samples. Importantly, for each of the four tasks, we can compute the integral in Eq. 2 either analytically or accurately capture it with quadrature over  $x$ . Hence, we obtain the true cost  $\ell(\theta; x_o)$  and subsequently the 'ground-truth' GBI posterior (with Eq. 1), and use that to draw, for each value of  $\beta$  and  $x_o$ , 5000 samples (GT-GBI, black in Fig. 3). See Appendix A4.2 for more detailed descriptions of tasks and distance functions.
Training data For each task, we simulate 10,000 pairs of  $(\theta, \mathbf{x})$  and construct the target dataset as in Fig. 2a, with 100 additional noise-augmented targets and 20 synthetic observations—10 well-specified and 10 misspecified—for a total of  $10120 \mathbf{x}_t$  data points. Well-specified observations are additional prior predictive samples, while misspecified observations are created by moving prior predictive samples outside the boundaries defined by the minimum and maximum of 100,000 prior predictive simulations (e.g., by successively adding Gaussian noise). See Appendix A4.3 for details.
Test data To evaluate inference performance, we use ACE to sample approximate GBI posteriors conditioned on 40 different synthetic observations, 20 of which were included in the target dataset  $\mathbf{x}_t$ , and 10 additional well-specified and misspecified observations which were not included in the target dataset. We emphasize that including observations in the target data is not a case of test data leakage, but represents a real use case where some experimental data which one wants to perform inference on are already available, while the network should also be amortized for unseen observations measured after training. Nevertheless, we report in Fig. 3 results for 'unseen' observations, i.e., not in the target dataset. Results are almost identical for those that were in the target dataset (Appendix A1). We drew 5000 posterior samples per observation, for 3 different  $\beta$  values for each task.
Metrics We are primarily interested in two aspects of performance: approximate posterior predictive distance and cost estimation accuracy. First, as motivated above, we want to find parameter
Figure 3: Performance on benchmark tasks. ACE obtains posterior samples with low average distance to observations, and accurately estimates cost function. Rows: results for each task. Columns: average predictive distance compared to SBI methods and GT (1st and 2nd), cost estimation accuracy evaluated on ACE posterior samples for different  $\beta$  (lighter blue shades are lower values of  $\beta$ ) (3rd), and C2ST accuracy relative to GT GBI posterior (4th, lower is better).
configurations which produce simulations that are as close as possible to the observation, as measured by the task-specific distance function. Therefore, we simulate using each of the 5000 ACE GBI posterior samples, and compute the average distance between predictive simulations and the observation. Mean and standard deviation are shown for well-specified and misspecified observations separately below (Fig. 3, 1st and 2nd columns). Second, we want to confirm that ACE accurately approximates  $\ell(\theta; \mathbf{x}_o)$ , which is a prerequisite for correctly inferring the GBI posterior. Therefore, we compare the ACE-predicted and true cost across 5000 samples from each GBI posterior, as well as the classifier 2-sample test (C2ST, [25, 24]) score between the ACE approximate and ground-truth GBI posterior (Fig. 3, 3rd and 4th columns). Note that cost estimation accuracy can be evaluated for parameter values sampled in any way (e.g., from the prior), but here we evaluate accuracy as samples become more concentrated around good parameter values, i.e., from GBI posteriors with increasing  $\beta$ . We expect that these tasks become increasingly challenging with higher values of  $\beta$ , since these settings require the cost estimation network to be highly accurate in tiny regions of parameter-space.