Low-Budget Simulation-Based Inference with Bayesian Neural Networks

## Abstract

Abstract Simulation-based inference methods have been shown to be inaccurate in the data-poor regime, when training simulations are limited or expensive.
Under these circumstances, the inference network is particularly prone to overfitting, and using it without accounting for the computational uncertainty arising from the lack of identifiability of the network weights can lead to unreliable results.
To address this issue, we propose using Bayesian neural networks in low-budget simulation-based inference, thereby explicitly accounting for the computational uncertainty of the posterior approximation.
We design a family of Bayesian neural network priors that are tailored for inference and show that they lead to well-calibrated posteriors on tested benchmarks, even when as few as O ⁢ ( 10 ) 𝑂 10 O(10) italic_O ( 10 ) simulations are available.
This opens up the possibility of performing reliable simulation-based inference using very expensive simulators, as we demonstrate on a problem from the field of cosmology where single simulations are computationally expensive. We show that Bayesian neural networks produce informative and well-calibrated posterior estimates with only a few hundred simulations.

## 1 Introduction

Simulation-based inference aims at identifying the parameters of a stochastic simulator that best explain an observation.
In its Bayesian formulation, simulation-based inference approximates the posterior distribution of the model parameters given an observation. This approximation usually takes the form of a neural network trained on synthetic data generated from the simulator.
In the context of scientific discovery, Hermans et al., () stressed the need for posterior approximations that are conservative – not overconfident – in order to make reliable downstream claims. They also showed that common simulation-based inference algorithms can produce overconfident approximations that may lead to erroneous conclusions.

In the data-poor regime (Villaescusa-Navarro et al.,, ; Zhang and Mikelsons,, ; Zeng et al.,, ), where the simulator is expensive to run and only a small number of simulations are available, training a neural network to approximate the posterior can easily lead to overfitting.
With small amounts of training data, the neural network weights are only loosely constrained, leading to high computational uncertainty (Wenger et al.,, ).
That is, many neural networks can fit the training data equally well, yet they may have very different predictions on test data.
For this reason, the posterior approximation is uncertain and, in the absence of a proper quantification of this uncertainty, potentially overconfident.
Fortunately, computational uncertainty in a neural network can be quantified using Bayesian neural networks (BNNs) (Gal et al.,, ), which account for the uncertainty in the neural network weights. Therefore, in the context of simulation-based inference, BNNs can provide a principled way to quantify the computational uncertainty of the posterior approximation.

Hermans et al., () showed empirically that using ensembles of neural networks, a crude approximation of BNNs (Lakshminarayanan et al.,, ), does improve the calibration of the posterior approximation. A few studies have also used BNNs as density estimators in simulation-based inference (Cobb et al.,, ; Walmsley et al.,, ; Lemos et al.,, ). However, these studies have remained empirical and limited in their evaluation.
This lack of theoretical grounding motivates the need for a more principled understanding of BNNs for simulation-based inference.
In particular, the choice of prior on the neural network weights happens to be crucial in this context, as it can strongly influence the resulting posterior approximation. Yet, arbitrary priors that convey little or undesired information about the posterior density have been used so far.

Our contributions are twofold. Firstly, we provide an improved understanding of BNNs in the context of simulation-based inference by empirically analyzing their effect on the resulting posterior approximations. Secondly, we introduce a principled way of using BNNs in simulation-based inference by designing meaningful priors. These priors are constructed to produce calibrated posteriors even in the absence of training data. We show that they are conservative in the small-data regime, for very low simulation budgets.
The code is available at [https://github.com/ADelau/low_budget_sbi_with_bnn](https://github.com/ADelau/low_budget_sbi_with_bnn).

## 2 Background

#### Simulation-based inference

We consider a stochastic simulator that takes parameters $\boldsymbol{\theta}$ as input and produces synthetic observations $\boldsymbol{x}$ as ouput.
The simulator implicitly defines the likelihood $p(\boldsymbol{x}|\boldsymbol{\theta})$ in the form of a forward stochastic generative model but does not allow for direct evaluation of its density due to the intractability of the marginalization over its latent variables. In this setup, Bayesian simulation-based inference aims at approximating the posterior distribution $p(\boldsymbol{\theta}|\boldsymbol{x})$ using the simulator. Among possible approaches, *neural* simulation-based inference methods train a neural network to approximate key quantities from simulated data, such as the posterior, the likelihood, the likelihood-to-evidence ratio, or a score function (Cranmer et al.,, ).

Recently, concerns have been raised regarding the calibration of the approximate posteriors obtained with neural simulation-based inference. Hermans et al., () showed that, unless special care is taken, common inference algorithms can produce overconfident posterior approximations. They quantify the calibration using the expected coverage

$$ $\text{EC}(\hat{p},\alpha)=\mathbb{E}_{p(\boldsymbol{\theta},\boldsymbol{x})}[ \mathds{1}(\boldsymbol{\theta}\in\boldsymbol{\Theta}_{\hat{p}}(\alpha))]$ (1) $$

where $\boldsymbol{\Theta}_{\hat{p}}(\alpha)$ denotes the highest posterior credible region at level $\alpha$ computed using the posterior approximate $\hat{p}(\boldsymbol{\theta}|\boldsymbol{x})$. The expected coverage is equal to $\alpha$ when the posterior approximate is calibrated, lower than $\alpha$ when it is overconfident and higher than $\alpha$ when it is underconfident or conservative.

The calibration of posterior approximations has been improved in recent years in various ways.
Delaunoy et al., (, ) regularize the posterior approximations to be balanced, which biases them towards conservative approximations. Similarly, Falkiewicz et al., () regularize directly the posterior approximation by penalizing miscalibration or overconfidence. Masserano et al., () use Neyman constructions to produce confidence regions with approximate Frequentist coverage. Patel et al., () combine simulation-based inference and conformal predictions. Schmitt et al., () enforce the self-consistency of likelihood and posterior approximations to improve the quality of approximate inference in low-data regimes.

#### Bayesian deep learning

Bayesian deep learning aims to account for both the aleatoric and epistemic uncertainty in neural networks.
The aleatoric uncertainty refers to the intrinsic randomness of the variable being modeled, typically taken into account by switching from a point predictor to a density estimator.
The epistemic uncertainty, on the other hand, refers to the uncertainty associated with the neural network itself and is typically high in small-data regimes.
Failing to account for this uncertainty can lead to high miscalibration as many neural networks can fit the training data equally well, yet they may have very different predictions on test data.

Bayesian deep learning accounts for epistemic uncertainty by treating the neural network weights as random variables and considering the full posterior over possible neural networks instead of only the most probable neural network (Papamarkou et al.,, ). Formally, let us consider a supervised learning setting in all generality, where $\boldsymbol{x}$ denotes inputs, $\boldsymbol{y}$ outputs, $\boldsymbol{D}$ a dataset of $N$ pairs $(\boldsymbol{x},\boldsymbol{y})$, and $\mathbf{w}$ the weights of the neural network. The likelihood of a given set of weights is

$$ $p(\boldsymbol{D}|\mathbf{w})\propto\prod_{i=1}^{N}p(\boldsymbol{y}_{i}| \boldsymbol{x}_{i},\mathbf{w}),$ (2) $$

where $p(\boldsymbol{y}_{i}|\boldsymbol{x}_{i},\mathbf{w})$ is the output of the neural network with weights $\mathbf{w}$ and inputs $\boldsymbol{x}_{i}$.
The resulting posterior over the weights is

$$ $p(\mathbf{w}|\boldsymbol{D})=\frac{p(\boldsymbol{D}|\mathbf{w})p(\mathbf{w})}{ p(\boldsymbol{D})},$ (3) $$

where $p(\mathbf{w})$ is the prior. Once estimated, the posterior over the neural network’s weights can be used for predictions through the Bayesian model average

$$ $p(\boldsymbol{y}|\boldsymbol{x},\boldsymbol{D})=\int p(\boldsymbol{y}| \boldsymbol{x},\mathbf{w})p(\mathbf{w}|\boldsymbol{D})d\mathbf{w}\simeq\frac{1 }{M}\sum_{i=1}^{M}p(\boldsymbol{y}|\boldsymbol{x},\mathbf{w}_{i}),\mathbf{w}_{ i}\sim p(\mathbf{w}|\boldsymbol{D}).$ (4) $$

In practice, the Bayesian model average can be approximated by Monte Carlo sampling, with $M$ samples from the posterior over the weights. The quality of the approximation depends on the number of samples $M$, which should be chosen high enough to obtain a good enough approximation but small enough to keep reasonable the computational costs of predictions.

Estimating the posterior over the neural network weights is a challenging problem due to the high dimensionality of the weights.
Variational inference (Blundell et al.,, ) optimizes a variational family to match the true posterior, which is typically fast but requires specifying a variational family that may restrict the functions that can be modeled. Markov chain Monte Carlo methods (Welling and Teh,, ; Chen et al.,, ), on the other hand, are less restrictive in the functions that can be modeled but require careful tuning of the hyper-parameters and are more computationally demanding. The Bayesian posterior can also be approximated by an ensemble of neural networks (Lakshminarayanan et al.,, ; Pearce et al.,, ; He et al.,, ). Laplace methods leverage geometric information about the loss to construct an approximation of the posterior around the maximum a posteriori (MacKay,, ). Similarly, Maddox et al., () use the training trajectory of stochastic gradient descent to build an approximation of the posterior.

## 3 Bayesian neural networks for simulation-based inference

In the context of simulation-based inference, treating the weights of the inference network as random variables enables the quantification of the computational uncertainty of the posterior approximation.
Considering neural networks taking observations $\boldsymbol{x}$ as input and producing parameters $\theta$ as output, the posterior approximation $\hat{p}(\boldsymbol{\theta}|\boldsymbol{x})$ can be modeled as the Bayesian model average

$$ $\hat{p}(\boldsymbol{\theta}|\boldsymbol{x})=\int p(\boldsymbol{\theta}| \boldsymbol{x},\mathbf{w})p(\mathbf{w}|\boldsymbol{D})d\mathbf{w},$ (5) $$

where $p(\boldsymbol{\theta}|\boldsymbol{x},\mathbf{w})$ is the posterior approximation parameterized by the weights $\mathbf{w}$ and evaluated at $(\boldsymbol{\theta},\boldsymbol{x})$, and $p(\mathbf{w}|\boldsymbol{D})$ is the posterior over the weights given the training set $\boldsymbol{D}$.

Remaining is the choice of prior $p(\mathbf{w})$. While progress has been made in designing better priors in Bayesian deep learning (Fortuin,, ), we argue that none of those are suitable in the context of simulation-based inference. To illustrate our point, let us consider the case of a normal prior $p(\mathbf{w})=\mathcal{N}(\boldsymbol{0},\sigma^{2}\boldsymbol{I})$ on the weights, in which case

$$ $\hat{p}_{\text{normal prior}}(\boldsymbol{\theta}|\boldsymbol{x})=\int p( \boldsymbol{\theta}|\boldsymbol{x},\mathbf{w})\ \mathcal{N}(\mathbf{w}| \boldsymbol{\mu}=\boldsymbol{0},\boldsymbol{\Sigma}=\sigma^{2}\boldsymbol{I})d \mathbf{w}.$ (6) $$

As mentioned in Section 2, a desirable property for a posterior approximation is to be calibrated. Therefore we want $\text{EC}(\hat{p}_{\text{normal prior}},\alpha)=\alpha,\forall\alpha$. Although it might be possible for this property to be satisfied in particular settings, this is obviously not the case for all values of $\sigma$ and all neural network architectures. Therefore, and as illustrated in Figure 1, the Bayesian model average is not even calibrated a priori when using a normal prior on the weights. As the Bayesian model average is not calibrated a priori, it cannot be expected that updating the posterior over weights $p(\mathbf{w}|\boldsymbol{D})$ with a small amount of data would lead to a calibrated a posteriori Bayesian model average.

### 3.1 Functional priors for simulation-based inference

We design a prior that induces an a priori-calibrated Bayesian model average. To achieve this, we work in the space of posterior functions instead of the space of weights. We consider the space of functions taking a pair ($\boldsymbol{\theta}$, $\boldsymbol{x}$) as input and producing a posterior density value $f(\boldsymbol{\theta},\boldsymbol{x})$ as output. Each function $f$ is defined by the joint outputs it associates with any arbitrary set of inputs, such that a posterior over functions can be viewed as a distribution over joint outputs for arbitrary inputs. Formally, let us consider $M$ arbitrary pairs $(\boldsymbol{\theta},\boldsymbol{x})$ represented by the matrices $\boldsymbol{\Theta}=[\boldsymbol{\theta}_{1},...,\boldsymbol{\theta}_{M}]$ and $\boldsymbol{X}=[\boldsymbol{x}_{1},...,\boldsymbol{x}_{M}]$ and let $\boldsymbol{f}=[f_{1},...,f_{M}]$ be the joint outputs associated with those inputs. The distribution $p(\boldsymbol{f}|\boldsymbol{\Theta},\boldsymbol{X})$ then represents a distribution over posteriors $\boldsymbol{f}=[\tilde{p}(\boldsymbol{\theta}_{1}|\boldsymbol{x}_{1}),...,
\tilde{p}(\boldsymbol{\theta}_{M}|\boldsymbol{x}_{M})]$. The functional posterior distribution over posteriors for parameters $\boldsymbol{\Theta}$ and observations $\boldsymbol{X}$ given a training dataset $\boldsymbol{D}$ is then $p(\boldsymbol{f}|\boldsymbol{\Theta},\boldsymbol{X},\boldsymbol{D})$ and the Bayesian model average is obtained through marginalization, that is

$$ $p(\boldsymbol{\theta}_{i}|\boldsymbol{x}_{i},\boldsymbol{D})=\int f_{i}\ p( \boldsymbol{f}|\boldsymbol{\Theta},\boldsymbol{X},\boldsymbol{D})d\boldsymbol{ f},\quad\forall i.$ (7) $$

Computing the posterior over functions requires the specification of a prior over functions. We first observe that the prior over the simulator’s parameters is a calibrated approximation of the posterior. That is, for the prior function $p_{\text{prior}}:(\boldsymbol{\theta},\boldsymbol{x})\to p(\boldsymbol{\theta})$, we have that $\text{EC}(p_{\text{prior}},\alpha)=\alpha,\forall\alpha$ (Delaunoy et al.,, ). It naturally follows that the a priori Bayesian model average with a Dirac delta prior around the prior on the simulator’s parameters is calibrated

$$ $\displaystyle\hat{p}(\boldsymbol{\theta}_{i}|\boldsymbol{x}_{i})$ $\displaystyle=\int f_{i}\ \delta([f_{j}=p_{\text{prior}}(\boldsymbol{\theta}_{ j},\boldsymbol{x}_{j})])\ d\boldsymbol{f},\forall i$ (8) $\displaystyle=\int f_{i}\ \delta(f_{i}=p(\boldsymbol{\theta}_{i}))\ df_{i}, \forall i\Rightarrow\text{EC}(\hat{p},\alpha)=\alpha,\forall\alpha.$ $$

However, this prior has limited support, and the Bayesian model average will not converge to the posterior $p(\boldsymbol{\theta}|\boldsymbol{x})$ as the dataset size increases. We extend this Dirac prior to include more functions in its support while retaining the calibration property, which we propose defining as a Gaussian process centered at $p_{\text{prior}}$.

A Gaussian process (GP) defines a joint multivariate normal distribution over all the outputs $\boldsymbol{f}$ given the inputs ($\boldsymbol{\Theta}$, $\boldsymbol{X}$). It is parametrized by a mean function $\mu$ that defines the mean value for the outputs given the inputs and a kernel function $K$ that models the covariance between the outputs. If we have access to no data, the mean and the kernel jointly define a prior over functions as they define a joint prior over outputs for an arbitrary set of inputs. In order for this prior over functions to be centered around the prior $p_{\text{prior}}$, we define the mean function as $\mu(\boldsymbol{\theta},\boldsymbol{x})=p(\boldsymbol{\theta})$.
The kernel $K$, on the other hand, defines the spread around the mean function and the correlation between the outputs $\boldsymbol{f}$. Its specification is application-dependent and constitutes a hyper-parameter of our method that can be exploited to incorporate domain knowledge on the structure of the posterior. For example, periodic kernels could be used if periodicity is observed. Kernel’s hyperparameters can also be chosen such as to incorporate what would be a reasonable deviation of the approximated posterior from the prior. We denote the Gaussian process prior over function outputs as $p_{\text{GP}}(\boldsymbol{f}|\mu(\boldsymbol{\Theta},\boldsymbol{X}),K(
\boldsymbol{\Theta},\boldsymbol{X}))$. Proposition 1 shows that a functional prior defined in this way leads to a calibrated Bayesian model average.

###### Proposition 1 .

The Bayesian model average of a Gaussian process centered around the prior on the simulator’s parameters is calibrated. Formally, let $p_{\text{GP}}$ be the density probability function defined by a Gaussian process, $\mu$ its mean function, and $K$ the kernel. Let us consider $M$ arbitrary pairs $(\boldsymbol{\theta},\boldsymbol{x})$ represented by the matrices $\boldsymbol{\Theta}=[\boldsymbol{\theta}_{1},...,\boldsymbol{\theta}_{M}]$ and $\boldsymbol{X}=[\boldsymbol{x}_{1},...,\boldsymbol{x}_{M}]$ and represent by the vector $\boldsymbol{f}=[f_{1},...,f_{M}]$ the joint outputs associated with those inputs. The Bayesian model average on the $i^{\text{th}}$ pair is expressed

$$ $\hat{p}(\boldsymbol{\theta}_{i}|\boldsymbol{x}_{i})=\int f_{i}\ p_{\text{GP}}( \boldsymbol{f}|\mu(\boldsymbol{\Theta},\boldsymbol{X}),K(\boldsymbol{\Theta}, \boldsymbol{X}))\ d\boldsymbol{f}$ $$

If $\mu(\boldsymbol{\theta},\boldsymbol{x})=p(\boldsymbol{\theta}),\forall
\boldsymbol{\theta},\boldsymbol{x}$, then,

$$ $\mathrm{EC}(\hat{p},\alpha)=\alpha,\forall\alpha,$ $$

for all kernel $K$.

###### Proof.

As $p_{\text{GP}}$ is, by definition of a Gaussian process, a multivariate normal, the expectations of the marginals are equal to the mean parameters

$$ $\hat{p}(\boldsymbol{\theta}_{i}|\boldsymbol{x}_{i})=\mu(\boldsymbol{\theta}_{i },\boldsymbol{x}_{i})=p(\boldsymbol{\theta}_{i}).$ $$

The joint evaluation of the Bayesian model average of the Gaussian process is hence equivalent to the joint evaluation of the prior for any matrices $\boldsymbol{\Theta}$ and $\boldsymbol{X}$. We can therefore conclude that $\hat{p}$ is equivalent to $p_{\text{prior}}:(\boldsymbol{\theta},\boldsymbol{x})\rightarrow p(\boldsymbol
{\theta})$. Since $\text{EC}(p_{\text{prior}},\alpha)=\alpha,\forall\alpha$ (Delaunoy et al.,, ), then, $\text{EC}(\hat{p},\alpha)=\alpha,\forall\alpha$.
∎

###### Proposition 1 .

###### Proof.

### 3.2 From functional to parametric priors

In this section, we now discuss how the functional GP prior over posterior density functions proposed in Section 3.1 can be used in the simulation-based inference setting. We follow Flam-Shepherd et al., () and Sun et al., () for mapping the functional prior to a distribution over neural network weights. We note that other methods for functional Bayesian deep learning have been proposed in the literature, such as (Tran et al.,, ; Rudner et al.,, ; Kozyrskiy et al.,, ; Ma and Hernández-Lobato,, ). These methods could also be used in our setting and are discussed in Appendix A.
Let us first observe that a neural network architecture and a prior on weights jointly define a prior over functions. We parameterize the prior on weights by $\boldsymbol{\phi}$ and denote this probability density over function outputs by

$$ $\displaystyle p_{\text{BNN}}(\boldsymbol{f}\,|\,\boldsymbol{\phi},\boldsymbol{ \Theta},\boldsymbol{X})$ $\displaystyle=\int p(\boldsymbol{f}\,|\,\mathbf{w},\boldsymbol{\Theta}, \boldsymbol{X})\ p(\mathbf{w}|\boldsymbol{\phi})\ d\mathbf{w}$ (9) $\displaystyle=\int\delta([f_{i}=p(\boldsymbol{\theta}_{i}|\boldsymbol{x}_{i}, \mathbf{w})])\ p(\mathbf{w}|\boldsymbol{\phi})\ d\mathbf{w}.$ $$

To obtain a prior on weights that matches the target GP prior, we optimize $\boldsymbol{\phi}$ such that $p_{\text{BNN}}(\boldsymbol{f}\,|\,\boldsymbol{\phi},\boldsymbol{\Theta},
\boldsymbol{X})$ matches $p_{\text{GP}}(\boldsymbol{f}|\mu(\boldsymbol{\Theta},\boldsymbol{X}),K(
\boldsymbol{\Theta},\boldsymbol{X}))$. Following Flam-Shepherd et al., (), given a measurement set $\mathcal{M}=\{\boldsymbol{\theta}_{i},\boldsymbol{x}_{i}\}_{i=1}^{M}$ at which we want the distributions to match, the KL divergence between the two priors can be expressed as

$$ $\displaystyle\text{KL}\left[p_{\text{BNN}}(\boldsymbol{f}\,|\,\boldsymbol{\phi },\mathcal{M})\,||\,p_{\text{GP}}(\boldsymbol{f}\,|\,\mu(\mathcal{M}),K( \mathcal{M}))\right]$ (10) $\displaystyle=$ $\displaystyle\int p_{\text{BNN}}(\boldsymbol{f}\,|\,\boldsymbol{\phi},\mathcal {M})\log\frac{p_{\text{BNN}}(\boldsymbol{f}\,|\,\boldsymbol{\phi},\mathcal{M}) }{p_{\text{GP}}(\boldsymbol{f}\,|\,\mu(\mathcal{M}),K(\mathcal{M}))}d \boldsymbol{y}$ $\displaystyle=$ $\displaystyle-\mathbb{H}\left[p_{\text{BNN}}(\boldsymbol{f}\,|\,\boldsymbol{ \phi},\mathcal{M})\right]-\mathbb{E}_{p_{\text{BNN}}(\boldsymbol{f}\,|\, \boldsymbol{\phi},\mathcal{M})}\left[\log p_{\text{GP}}(\boldsymbol{f}\,|\,\mu (\mathcal{M}),K(\mathcal{M}))\right],$ $$

where the second term $\mathbb{E}_{p_{\text{BNN}}(\boldsymbol{f}\,|\,\boldsymbol{\phi},\mathcal{M})}
\left[\log p_{\text{GP}}(\boldsymbol{f}\,|\,\mu(\mathcal{M}),K(\mathcal{M}))\right]$ can be estimated using Monte-Carlo. The first term $\mathbb{H}\left[p_{\text{BNN}}(\boldsymbol{f}\,|\,\boldsymbol{\phi},\mathcal{M
})\right]$, however, is harder to estimate as it requires computing $\log p_{\text{BNN}}(\boldsymbol{f}\,|\,\boldsymbol{\phi},\mathcal{M})$, which involves the integration of the output over all possible weights combinations. To bypass this issue, Sun et al., () propose to use Spectral Stein Gradient Estimation (SSGE) (Shi et al.,, ) to approximate the gradient of the entropy as

$$ $\nabla\mathbb{H}\left[p_{\text{BNN}}(\boldsymbol{f}\,|\,\boldsymbol{\phi}, \mathcal{M})\right]\simeq\text{SSGE}\left(\boldsymbol{f}_{1},...,\boldsymbol{f }_{N}\sim p_{\text{BNN}}(\boldsymbol{f}\,|\,\boldsymbol{\phi},\mathcal{M}) \right).$ (11) $$

We note that the measurement set $\mathcal{M}$ can be chosen arbitrarily but should cover most of the support of the joint distribution $p(\boldsymbol{\theta},\boldsymbol{x})$. If data from this joint distribution are available, those can be leveraged to build the measurement set. To showcase the ability to create a prior with limited data, in this work, we derive boundaries of the support of each marginal distribution and draw parameters and observations independently and uniformly over this support. If the support is known a priori, this procedure can be performed without (expensive) simulations. We draw a new measurement set at each iteration of the optimization procedure. If a fixed measurement set is available, a subsample of this measurement set should be drawn at each iteration.

As an illustrative example, we chose independent normal distributions as a variational family $p(\mathbf{w}|\boldsymbol{\phi})$ over the weights and minimize (10) w.r.t. $\mathbf{w}$. In Figure 1, we show the coverage of the resulting a priori Bayesian model average using the tuned prior, $p(\mathbf{w}\,|\,\boldsymbol{\phi})$, and normal priors for increasing standard deviations $\sigma$, for the SLCP benchmark. We observe that while none of the normal priors are calibrated, the trained prior achieves near-perfect calibration. This prior hence guides the obtained posterior approximation towards more calibrated solutions, even in low simulation-budget settings.

The attentive reader might have noticed that $p_{\text{BNN}}(\boldsymbol{f}\,|\,\boldsymbol{\phi},\boldsymbol{\Theta},
\boldsymbol{X})$ and $p_{\text{GP}}(\boldsymbol{f}|\mu(\boldsymbol{\Theta},\boldsymbol{X}),K(
\boldsymbol{\Theta},\boldsymbol{X}))$ do not share the same support, as the former distribution is limited to functions that represent valid densities by construction, while the latter includes arbitrarily shaped functions. This is not an issue here as the support of the first distribution is included in the support of the second distribution, and functions outside the support of the first distribution are ignored in the computation of the divergence.

Figure: Figure 1: Visualization of the prior tuned to match the GP prior on the SLCP benchmark. Left: examples of posterior functions sampled from the tuned prior over neural network’s weights. Right: expected coverage of the prior Bayesian model average with the tuned prior and normal priors for varying standard deviations.
Refer to caption: x1.png

## 4 Experiments

In this section, we empirically demonstrate the benefits of replacing a regular neural network with a BNN equipped with the proposed prior for simulation-based inference.
We consider both Neural Posterior Estimation (NPE) with neural spline flows (Durkan et al.,, ) and Neural Ratio Estimation (NRE) (Hermans et al.,, ), along with their balanced versions (BNRE and BNPE) (Delaunoy et al.,, , ) and ensembles (Lakshminarayanan et al.,, ; Hermans et al.,, ). BNNs-based methods are trained using mean-field variational inference (Blundell et al.,, ).
As advocated by Wenzel et al., (), we also consider cold posteriors to achieve good predictive performance. More specifically, the variational objective function is modified to give less weight to the prior by introducing a temperature parameter $T$,

$$ $\mathbb{E}_{\mathbf{w}\sim p(\mathbf{w}|\boldsymbol{\tau})}\left[\sum_{i}\log p (\boldsymbol{\theta}_{i}|\boldsymbol{x}_{i},\mathbf{w})\right]-T\ \text{KL}[p( \mathbf{w}|\boldsymbol{\tau})||p(\mathbf{w}|\boldsymbol{\phi})],$ (12) $$

where $\boldsymbol{\tau}$ are the parameters of the posterior variational family and $T$ is a parameter called the temperature that weights the prior term. In the following, we call BNN-NPE, a Bayesian Neural Network posterior estimator trained without temperature ($T=1$), and BNN-NPE ($T=0.01$), an estimator trained with a temperature of $0.01$, assigning a lower weight to the prior.

A detailed description of the Gaussian process used can be found in Appendix A. For simplicity, in this analysis, we use an RBF kernel in the GP prior. If more information on the structure of the target posterior is available, more informed kernels may be used to leverage this prior knowledge. A description of the benchmarks can be found in Appendix B, and the hyperparameters are described in Appendix C. For clarity, only NPE-based methods are shown in this section; results using NRE can be found in Appendix D.

Following Delaunoy et al., (), we evaluate the quality of the posterior approximations based on the expected nominal log posterior density and the expected coverage area under the curve (coverage AUC). The expected nominal log posterior density $\mathbb{E}_{\boldsymbol{\theta},\boldsymbol{x}\sim p(\boldsymbol{\theta},
\boldsymbol{x})}\left[\log\hat{p}(\boldsymbol{\theta}|\boldsymbol{x})\right]$ quantifies the amount of density allocated to the nominal parameter that was used to generate the observation. The coverage AUC $\int_{0}^{1}(\text{EC}(\hat{p},\alpha)-\alpha)\ d\alpha$ quantifies the calibration of the expected posterior. A calibrated posterior approximation exhibits a coverage AUC of 0 0. A positive coverage AUC indicates conservativeness, and a negative coverage AUC indicates overconfidence.

#### BNN-based simulation-based inference

Figure 2 compares simulation-based inference methods with and without accounting for computational uncertainty We observe that BNNs equipped with our prior and without temperature show positive coverage AUC even for simulation budgets as low as $O(10)$. The coverage curves are reported in Appendix D and show that this corresponds to conservative posterior approximations. We conclude that BNNs can hence be reliably used when the simulator is expensive and few simulations are available. We also observe that the nominal log posterior density is on par with other methods for very high simulation budgets but that more samples are required to achieve high values. Cold posteriors can help achieve high nominal log posterior values with fewer samples at the cost of sometimes producing overconfident posterior approximations. From these observations, general guidelines to set the temperature include increasing $T$ if overconfidence is observed and decreasing it if low predictive performance is observed.

Examples of posterior approximations obtained with and without using a Bayesian neural network are shown in Figure 3. Wide posteriors are observed for low budgets for BNN-NPE, while NPE produces an overconfident approximation and excludes most of the relevant parts of the posterior. As the simulation budget increases, BNN-NPE converges slowly towards the same posterior as NPE. BNN-NPE ($T=0.01$) converges faster than BNN-NPE but, for low simulation budgets, excludes parts of the region that should be accepted according to high budget posteriors. Yet, the posterior approximate is still less overconfident than NPE’s. Finally, further results in Appendix D show that BNN-NRE is more conservative than BNN-NPE. This comes at the cost of lower nominal log posterior density for a given simulation budget while still reaching comparable values as the simulation budgets grow.

Figure: Figure 2: Comparison of different simulation-based inference methods through the nominal log probability and coverage area under the curve. The higher the nominal log probability, the more performant the method is. A calibrated posterior approximation exhibits a coverage AUC of 0 0. A positive coverage AUC indicates conservativeness, and a negative coverage AUC indicates overconfidence. 3 runs are performed, and the median is reported.
Refer to caption: x2.png

Figure: Figure 3: Examples of $95\%$ highest posterior density regions obtained with various algorithms and simulation budgets on the SLCP benchmark for a single observation. The black star represents the ground truth used to generate the observation.
Refer to caption: x3.png

#### Comparison of different priors on weights

We analyze the effect of the prior on the neural network’s weights on the resulting posterior approximation. The posterior approximations obtained using our GP prior are compared to the ones obtained using independent normal priors on weights with zero means and increasing standard deviations. In Figure 4, we observe that when using a normal prior, careful tuning of the standard deviation is needed to achieve results close to the prior designed for simulation-based inference. The usage of an inappropriate prior can lead to bad calibration for low simulation budgets or can prevent learning if it is too restrictive.

Figure: Figure 4: Comparison of posterior approximations obtained using a prior tuned to match the Gaussian process-based prior and using independent normal priors on weights with zero means and various standard deviations on the SLCP benchmark. 3 runs are performed, and the median is reported.
Refer to caption: x4.png

#### Uncertainty decomposition

We decompose the uncertainty quantified by the different methods. Following Depeweg et al., (), the uncertainty can be decomposed as

$$ $\mathbb{H}\left[\hat{p}(\boldsymbol{\theta}|\boldsymbol{x})\right]=\mathbb{E}_ {q(\mathbf{w})}\left[\mathbb{H}\left[\hat{p}(\boldsymbol{\theta}|\boldsymbol{x },\mathbf{w})\right]\right]+\mathbb{I}(\boldsymbol{\theta},\mathbf{w}),$ (13) $$

where $\mathbb{E}_{q(\mathbf{w})}\left[\mathbb{H}\left[\hat{p}(\boldsymbol{\theta}|
\boldsymbol{x},\mathbf{w})\right]\right]$ quantifies the aleatoric uncertainty, $\mathbb{I}(\boldsymbol{\theta},\mathbf{w})$ quantifies the epistemic uncertainty, and the sum of those terms is the predictive uncertainty.
Figure 5 shows the decomposition of the two sources of uncertainty, in expectation, on the SLCP benchmark. Other benchmarks can be found in Appendix D. We observe that BNN-NPE and NPE ensemble methods account for the epistemic uncertainty while other methods do not. BNPE artificially increases the aleatoric uncertainty to be better calibrated. The epistemic uncertainty of BNN-NPE is initially low because most of the models are slight variations of $p_{\boldsymbol{\Theta}}$. The epistemic uncertainty then increases as it starts to deviate from the prior and decreases as the training set size increases. BNN-NPE ($T=0.01$) exhibits a higher epistemic uncertainty for low budgets as the effect of the prior is lowered.

Figure: Figure 5: Quantification of the different forms of uncertainties captured by the different NPE-based methods on the SLCP benchmark. 3 runs are performed, and the median is reported.
Refer to caption: x5.png

Figure: Figure 6: Comparison of the posterior approximations obtained with and without a Bayesian neural network on the cosmological application. First plot: An example observation: particles representing galaxies in a synthetic universe. Second plot: example of $95\%$ highest posterior density regions for increasing simulation budgets. The black star represents the ground truth used to generate the observation. Third plot: Expected coverage with and without using a Bayesian neural network for increasing simulation budgets. Fourth plot: The nominal log posterior.
Refer to caption: x6.png

#### Infering cosmological parameters from N 𝑁 N italic_N -body simulations

To showcase the utility of Bayesian deep learning for simulation-based inference in a practical setting, we consider a challenging inference problem from the field of cosmology. We consider *Quijote* $N$-body simulations (Villaescusa-Navarro et al.,, ) tracing the spatial distribution of matter in the Universe for different underlying cosmological models. The resulting observations are particles with different masses, corresponding to dark matter clumps, which host galaxies.
We consider the canonical task of inferring the matter density (denoted $\Omega_{m}$) and the root-mean-square matter fluctuation averaged over a sphere of radius $8h^{-1}$ Mpc (denoted $\sigma_{8}$) from an observed galaxy field. Robustly inferring the values of these parameters is one of the scientific goals of flagship cosmological surveys. These simulations are very computationally expensive to run, with over $35$ million CPU hours required to generate $44100$ simulations at a relatively low resolution. Generating samples at higher resolutions, or a significantly larger number of samples, is challenging due to computational constraints. These constraints necessitate methods that can be used to produce reliable scientific conclusions from a limited set of simulations – when few simulations are available, not only is the amount of training data low, but so is the amount of test data that is available to assess the calibration of the trained model.

In this experiment, we use $2000$ simulations processed as described in Cuesta-Lazaro and Mishra-Sharma, (). These simulations form a subset of the full simulation suite run with a uniform prior over the parameters of interest. $1800$ simulations are used for training and $200$ are kept for testing. We use the two-point correlation function evaluated at $24$ distance bins as a summary statistic. The observable is, hence, a vector of $24$ features. We observed that setting a temperature lower than $1$ was needed to reach reasonable predictive performance with Bayesian neural networks in this setting. Figure 6 compares the posterior approximations obtained with a single neural network against those obtained with a BNN trained with a temperature of $0.01$. We observe from the coverage plots that while a single neural network can lead to overconfident approximations in the data-poor regime, the BNN leads to conservative approximations. BNN-NPE also exhibits higher nominal log posterior probability. Additionally, we observe that it provides posterior approximations that are calibrated and have a high nominal log probability with only a few hundred samples.

## 5 Conclusion

In this work, we use Bayesian deep learning to account for the computational uncertainty associated with posterior approximations in simulation-based inference. We show that the prior on neural network’s weights should be carefully chosen to obtain calibrated posterior approximations and develop a prior family with this objective in mind. The prior family is defined in function space as a Gaussian process and mapped to a prior on weights. Empirical results on benchmarks show that incorporating Bayesian neural networks in simulation-based inference methods consistently yields conservative posterior approximations, even with limited simulation budgets of $\mathcal{O}(10)$. As Bayesian deep learning continues to rapidly advance (Papamarkou et al.,, ), we anticipate that future developments will strengthen its applicability in simulation-based inference, ultimately enabling more efficient and reliable scientific applications in domains with computationally expensive simulators.

Using BNNs for simulation-based inference also comes with limitations. The first observed limitation is that the Bayesian neural network based methods might need more simulated data in order to reach a predictive power similar to methods that do not use BNNs, such as NPE. While we showed that this limitation can be mitigated by tuning the temperature appropriately, this is something that might require trials and errors. A second limitation is the computational cost of predictions. When training a BNN using variational inference, the training cost remains on a similar scale as standard neural networks. At prediction time, however, the Bayesian model average described in Equation 4 must be approximated, and this requires a neural network evaluation for each Monte Carlo sample in the approximation. The computational cost of predictions then scales linearly with the number $M$ of Monte Carlo samples.

## Acknowledgments and Disclosure of Funding

The authors would like to thank Adrien Bolland for the valuable discussions we had during this project. Arnaud Delaunoy would also like to thank the National Fund for Scientific Research (F.R.S.-FNRS) for his scholarship.
This work is supported by the National Science Foundation under Cooperative Agreement PHY-2019786 (The NSF AI Institute for Artificial Intelligence and Fundamental Interactions, [http://iaifi.org/](http://iaifi.org/)).
SM is partly supported by the U.S. Department of Energy, Office of Science, Office of High Energy Physics of U.S. Department of Energy under grant Contract Number DE-SC0012567.

## Appendix A Prior tuning details

We tune the parameters $\boldsymbol{\phi}$ of a variational distribution over neural network weights $p(\mathbf{w}|\boldsymbol{\phi})$. The variational distribution is chosen to be independent normal distributions, with parameters $\boldsymbol{\phi}$ representing the means and standard deviations of each parameter of $\mathbf{w}$. This variational family defines a prior over function outputs

$$ $p_{\text{BNN}}(\boldsymbol{f}\,|\,\boldsymbol{\phi},\boldsymbol{\Theta}, \boldsymbol{X})=\int p(\boldsymbol{f}\,|\,\mathbf{w},\boldsymbol{\Theta}, \boldsymbol{X})p(\mathbf{w}|\boldsymbol{\phi})d\mathbf{w}.$ (14) $$

The parameters $\boldsymbol{\phi}$ are optimized to obtain a prior on weights that matches the target Gaussian process functional prior $p_{\text{GP}}(\boldsymbol{f}|\mu(\boldsymbol{\Theta},\boldsymbol{X}),K(
\boldsymbol{\Theta},\boldsymbol{X}))$. To achieve this, we repeatedly sample a measurement set $\mathcal{M}=\{\boldsymbol{\theta}_{i},\boldsymbol{x}_{i}\}_{i=1}^{M}$ and $N$ function outputs from the BNN prior $\boldsymbol{f}_{1},...,\boldsymbol{f}_{N}\sim p_{\text{BNN}}(\boldsymbol{f}\,|
\,\boldsymbol{\phi},\mathcal{M})$ and perform a step of gradient descend to minimize the divergence

$$ $\text{KL}\left[p_{\text{BNN}}(\boldsymbol{f}\,|\,\boldsymbol{\phi},\mathcal{M} )\,||\,p_{\text{GP}}(\boldsymbol{f}\,|\,\mu(\mathcal{M}),K(\mathcal{M}))\right].$ (15) $$

The mean function $\mu$ of the Gaussian process is selected as:

$$ $\mu(\boldsymbol{\theta},\boldsymbol{x})=p(\boldsymbol{\theta}).$ (16) $$

The kernel $K$ is a combination of two Radial Basis Function (RBF) kernels

$$ $K(\boldsymbol{\theta}_{1},\boldsymbol{\theta}_{2},\boldsymbol{x}_{1}, \boldsymbol{x}_{2})=\sqrt{\text{RBF}(\boldsymbol{\theta}_{1},\boldsymbol{ \theta}_{2})}*\sqrt{\text{RBF}(\boldsymbol{x}_{1},\boldsymbol{x}_{2})}.$ (17) $$

such that the correlation between outputs is high only if $\boldsymbol{\theta}_{1}$ and $\boldsymbol{\theta}_{2}$ as well as $\boldsymbol{x}_{1}$ and $\boldsymbol{x}_{2}$ are close. The RBF kernel is defined as

$$ $\text{RBF}(\boldsymbol{x}_{1},\boldsymbol{x}_{2})=\sigma^{2}\exp\left(-\frac{1 }{N}\sum_{i}^{N}\frac{(x_{1,i}-x_{2,i})^{2}}{2l_{i}^{2}}\right),$ (18) $$

where $\sigma$ is the standard deviation and $l_{i}$ is the lengthscale associated to the $i^{\text{th}}$ feature. The lengthscale is derived from the measurement set. To determine $l_{i}$, we query observations $\boldsymbol{x}$ from the measurement set and compute the $0.1$ quantile of the squared distance between different observations for each feature. We then set $l_{i}$ such that $2l_{i}^{2}$ equals this quantile. All the benchmarks have a uniform prior over the simulator’s parameters. The mean function is then equal to a constant $C$ for all input values. The standard deviation is chosen to be $C/2$. To ensure stability during the inference procedure, we enforce all standard deviations defined in $\boldsymbol{\phi}$ to be at least $0.001$ by setting any parameters below this threshold to this value.

Note that there are various methods that can be used to perform inference on the neural network’s weights with our GP prior. Instead of minimizing the KL-divergence, the parameters $\boldsymbol{\phi}$ can be optimized using an adversarial training procedure by treating both priors as function generators and training a discriminator between the two (Tran et al.,, ). Another approach to performing inference using a functional prior is to directly use it during inference by modifying the inference algorithm to work in function space. Variational inference can be performed in the space of function (Sun et al.,, ; Rudner et al.,, ). The stochastic gradient Hamiltonian Monte Carlo algorithm (Chen et al.,, ) could also be modified to include a functional prior Kozyrskiy et al., (). Alternatively, a variational implicit process can be learned to express the posterior in function space (Ma and Hernández-Lobato,, ).

## Appendix B Benchmarks description

#### SLCP

The SLCP (Simple Likelihood Complex Posterior) benchmark (Papamakarios et al.,, ) is a fictive benchmark that takes $5$ parameters as input and produces an 8-dimensional synthetic observable. The observation corresponds to the 2D coordinates of 4 points that are sampled from the same multivariate normal distribution. We consider the task of inferring the marginal over $2$ of the $5$ parameters.

#### Two Moons

The Two Moons simulator (Greenberg et al.,, ) models a fictive problem with 2 parameters. The observable x is composed of 2 scalars, which represent the 2D coordinates of a random point sampled from a crescent-shaped distribution shifted and rotated around the origin depending on the parameters’ values. Those transformations involve the absolute value of the sum of the parameters leading to a second crescent in the posterior and, hence
making it multi-modal.

#### Lotka Volterra

The Lotka-Volterra population model (Lotka,, ; Volterra,, ) describes a process of interactions between a predator and a prey species. The model is conditioned on 4 parameters that influence the reproduction and mortality rate of the predator and prey species. We infer the marginal posterior of the predator parameters from a time series of 2001 steps representing the evolution of both populations over time. The specific implementation is based on a Markov Jump Process, as in Papamakarios et al., ().

#### SpatialSIR

The Spatial SIR model (Hermans et al.,, ) involves a grid world of susceptible, infected, and recovered individuals. Based on initial conditions and the infection and recovery rate, the model describes the spatial evolution of an infection. The observable is a snapshot of the grid world after some fixed amount of time. The grid used is of size 50 by 50.

## Appendix C Hyperparameters

All the NPE-based methods use a Neural Spline Flow (NSF) (Durkan et al.,, ) with 3 transforms of 6 layers, each containing 256 neurons. Meanwhile, all the NRE-based methods employ a classifier consisting of 6 layers of 256 neurons. For the spatialSIR and Lotka Volterra benchmarks, the observable is initially processed by an embedding network. Lotka Volterra’s embedding network is a 10 layers 1D convolutional neural network that leads to an embedding of size $512$. On the other hand, SpatialSIR’s embedding network is an 8 layers 2D convolutional neural network resulting in an embedding of size $256$. All the models are trained for $500$ epochs which we observed to be enough to reach convergence.

Bayesian neural network-based methods use independent normal distributions as a variational family. During inference, $100$ neural networks are sampled to approximate the Bayesian model average. Ensemble methods involve training $5$ neural networks independently. The experiments were conducted on a private GPU cluster, and the estimated computational cost is around $25,000$ GPU hours.

## Appendix D Additional experiments

In this section, we provide complementary results. Figure 7 illustrates the performance of the various NRE variants. Figures 8 and 9 display the coverage curves, demonstrating that a higher positive coverage AUC corresponds to coverage curves above the diagonal line. Figures 10 and 11 present the uncertainty decomposition of all methods on all the benchmarks. Figures 12 and 13 illustrate how the performance of the different algorithms varies across different runs.

Figure: Figure 7: Comparison of different NRE simulation-based inference methods through the nominal log probability and coverage area under the curve. The higher the nominal log probability, the more performant the method is. A calibrated posterior approximation exhibits a coverage AUC of 0 0. A positive coverage AUC indicates conservativeness, and a negative coverage AUC indicates overconfidence. 3 runs are performed, and the median is reported
Refer to caption: x7.png

Figure: Figure 8: Coverage of different NPE simulation-based inference methods. A calibrated posterior approximation exhibits a coverage AUC of 0 0. A coverage curve above the diagonal indicates conservativeness and a curve below the diagonal indicates overconfidence. 3 runs are performed, and the median is reported.
Refer to caption: x8.png

Figure: Figure 9: Coverage of different NRE simulation-based inference methods. A calibrated posterior approximation exhibits a coverage AUC of 0 0. A coverage curve above the diagonal indicates conservativeness and a curve below the diagonal indicates overconfidence. 3 runs are performed, and the median is reported.
Refer to caption: x9.png

Figure: Figure 10: Quantification of the different forms of uncertainties captured by the different NPE-based methods. 3 runs are performed, and the median is reported.
Refer to caption: x10.png

Figure: Figure 11: Quantification of the different forms of uncertainties captured by the different NRE-based methods. 3 runs are performed, and the median is reported.
Refer to caption: x11.png

Figure: Figure 12: Comparison of different NPE simulation-based inference methods through the nominal log probability and coverage area under the curve. The higher the nominal log probability, the more performant the method is. A calibrated posterior approximation exhibits a coverage AUC of 0 0. A positive coverage AUC indicates conservativeness, and a negative coverage AUC indicates overconfidence. 3 runs are performed. The median run is reported in plain, and the shaded lines represent the other two runs.
Refer to caption: x12.png

Figure: Figure 13: Comparison of different NRE simulation-based inference methods through the nominal log probability and coverage area under the curve. The higher the nominal log probability, the more performant the method is. A calibrated posterior approximation exhibits a coverage AUC of 0 0. A positive coverage AUC indicates conservativeness, and a negative coverage AUC indicates overconfidence. 3 runs are performed. The median run is reported in plain, and the shaded lines represent the other two runs.
Refer to caption: x13.png

## Peer Reviews

confidence: 4
contribution: 4
initial_rating: 5
presentation: 3
questions: 1) Do these simulations for prior tuning have to be run in addition to the training data?
2) Are they accounted for the in the budgets in the benchmarks?
3) NPE / NRE ensembles seems to perform similarly well compared to the BNN approach. How
  does it perform on the cosmology example? How does it compare to the BNN approach in
  terms of computational budget? Do you think ensembles could be a good alternative when
  computational budgets are limited? What are the disadvantages of ensembles compared to
  BNN for SBI?
4) More generally, what is the computational complexity of the BNN approach compared to
  NPE (ensembles)?
soundness: 3
strengths: ### Originality

The idea of using BNN for SBI is not new (as discussed in the paper). However, showing
why previous approaches did not work so well and the proposal of a new type of BNN prior
more suitable for the SBI setting is a valuable contribution. Most of the previous work
on better uncertainty quantification in SBI is discussed adequately in the introduction.
The only paper with a related approach that seems to be missing in the discussion is [Lueckmann et al.
2017, section
2.2](https://proceedings.neurips.cc/paper/2017/hash/addfa9b7e234254d26e9c7f2af1005cb-Abstract.html),
who propose performing SVI on the neural network weights for continual learning across
SNPE rounds. Also, it seems a bit odd to me that the benchmark used are not cited
(except for the spatial SIR by Hermans et al.). The SLCP benchmark was introduced by
Papamakarios et al. (SNLE paper), two moons by Greenberg et al.; and the overall set of
SBI benchmarking tasks was introduced in Lueckmann et al. 2021.

### Quality

The technical contributions of the paper appear as sound and the selected methods for
applying BNN to SBI are appropriate. The experimental results tend to support the
initial claim of obtaining well-calibrated posteriors in low-data regime. However, the
number of performed experiments is quite low (three) and the results appear quite noisy.

In general, it should be made clearer how strong the trade-off is between posterior
calibration and posterior accuracy. At the moment it seems that the BNN approach tends
to be quite underconfident in some scenarios (e.g., BNN-NRE on all benchmarks). I
therefore suggest that more experiments should be run and results should be reported
with error bars (e.g., 5-10 runs with error bars showing the standard error of the mean
performance). Additionally, I think it would be good to also check the posterior
predictive distribution for the cases where the BNN has low nominal log posterior values
even in high-data regimes.

### Clarity

Overall, the paper is well-written and clearly-structured. I have a couple of remarks
that should be addressed to improve the clarity.

- In section 3, especially in section 3.2, it should be explained more clearly why and
  how the prior has to be tuned to be used for the BNN approach. In the classical SBI
  setting, the prior is usually set a priori using expert knowledge. This step usually
  does not involve using simulated data for an optimization procedure. In the BNN
  approach, it seems that the prior has to be tuned with actual simulations (e.g., lines
  254-261). Do these simulations have to be run in addition to the training data? Are
  they accounted for the in the budgets in the benchmarks?
- Figure 1 needs a couple of clarifications. The caption says it's a visualization of
  the tuned prior. Then, the next sentence says the left panels show posterior functions
  sampled from the *tuned prior prior over the neural network's weights*. How are these
  samples obtained and what are thet supposed to show? Are they obtained before or after
  SBI training? Similarly question for the last panel: is the calibration calculated
  after the full SBI training or only after the prior tuning?

In general, it seems that a better explanation of the steps involved in setting up and
training the BNN-NPE (NRE) approach. I suggest adding more details in the next and
adding an algorithm scheme in the text or appendix.

### Significance

Uncertainty quantification in neural SBI methods on the posterior approximation itself
is an important and timely problem. Especially in low-data regimes, common neural SBI
methods like NPE struggle. The proposed BNN approach with a Gaussian process prior
as proposed here appears as an important contribution for addressing this issue.
summary: The paper intodruces a new method for simulation-based inference (SBI) that uses
Bayesian neural networks (BNN) to better account for uncertainty in posterior estimates
due to limited training data. The authors demonstrate that previous attempts of
combining BNN with SBI showed limited success due to the choice of prior on the BNN
weights and they propose an alternative prior more suitable for the SBI setting. To
evaluate the proposed method, the authors show on four benchmarking tasks and on a SBI
use-case from cosmology that the resulting SBI posterior estimates are well-calibrated
even in the low-data regime.
weaknesses: I outlined several concerns and questions above. To summarize to most important points:

- the experimental results are difficult to interpret because they show the media of
  only three repetitions. More repetitions and error bars would be better here.
- the high underconfidence of the BNN approach in how-data regimes is concerning.
  Additional evaluation of the posterior predictive distributions would be appropriate.
- the choice and construction of the prior from simulated should be explained more
  clearly. A better explanation ideally will resolve the questions on the general
  procedure and on Figure 1 above.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 2
questions: - The results in Figure 3 seem to indicate that NPE is indeed outperforming BNN-NPE here. Would the authors argue that BNN-NPE is preferable over NPE for this task?
- What do the numbers in the legend in Figure 3 mean?
soundness: 3
strengths: - The paper proposes a novel contribution for simulation-based inference. The method is in my opinion very relevant to the community and interesting as a low-budget solution for inferential problems.
- The motivation and derivation of the tailored prior is in my opinion particularly well done and convincing.
- Empirically, the method seems to achieve the motivated goal: having calibrated posteriors for low simulation budgets.
- The paper is well written and easy to follow.
summary: The authors propose a novel approach for simulation-based inference for low-budget problems using Bayesian neural networks. They additionally propose to use tailored priors for the neural network weights which generally leads to more conservative posterior estimates. They evaluate their method on several experimental models against multiple baselines demonstrating promising results for SBI when the computational budget is limited.
weaknesses: The evaluations and presentations of the results could in my opinion be improved.

- The authors evaluate their method with the expected coverage (EC). In Figure 2, it is in my opinion very difficult to draw conclusions which methods works best. This is likely due to the fact that only 3 runs have been evaluated. Given the complexity of SBI and the high variance of the inferential results, I think they should do at least 10 evaluations and report these.
- The authors use as a second evaluation metric the expected posterior log density (EPLD). It is not clear what a good value should look like here or even if a high log density should be desired. When the goal is to have a conservative posterior, is a extremely high EPLD a good measure? The authors should, despite possible drawbacks, have used some divergence to the true posterior, such as MMD, in addition to the other two and report this.
- It is not clear how crucial hyperparameters such as the temperature $T$ or the covariance functions of the reference GP are chosen in practice. 
- The authors propose to use cold posteriors, e.g., using a temperature of $T=0.01$. The motivation of this is not clear, given that the entire idea is to use a prior that enforces calibration. Why would it be desirable to in fact reduce the impact of the prior? 

#### Minor
- Figures 2 should show the standard errors.
- Figure 3 should show the reference posterior and not only the true parameter value.

---

confidence: 5
contribution: 3
initial_rating: 3
presentation: 2
questions: No questions.
soundness: 3
strengths: The idea to generate neural network weight priors which lead to a-priori well-calibrated posteriors is novel, interesting, and potentially impactful. The methodology which the authors employ to achieve this is novel, elegant, and rigorous. I thoroughly enjoyed reading these parts of the paper. In addition, the authors demonstrate that the method can be applied across methods (NPE & NRE) and they evaluate the method on a series of useful tasks.
summary: The authors aim to produce well-calibrated posterior in simulation-based inference, in particular in cases where the simulator is very expensive and, therefore, few simulations can be used as training data. To achieve this, the authors propose to use Bayesian neural networks. They develop a novel prior which leads to a-priori well-calibrated posteriors. They apply their method to several toy simulators and to a physics simulator.
weaknesses: (1) The paper overstates its claims.
My main issue with this paper is that it overstates its claims and does not acknowledge the weakness of empirical results. Looking at Figure 2, and in particular for NRE: BNN-NRE _never_ reaches the log posterior of the other methods (even for 1M simulations). Yet the authors state that `the nominal log posterior density is on par with other methods for very high simulation budgets`. Where does this claim come from. Similarly, in Figure 2 (NRE), the authors do not acknowledge that the AUC is not particularly much higher for BNN-NRE as compared to NRE, even for 10 simulations. Indeed, for three of the four tasks NRE has a higher coverage AUC than BNN-NRE for 10 simulations. On top of this, for many tasks and simulation budgets, BNN-NRE does have a negative AUC, but the authors simply claim that they `show positive coverage AUC`. Finally, 

(2) The empirical results are weak.
Second, to me, the empirical results are difficult to interpret. (A) Many of the curves shown in Figure 2 are very noisy and irregular. It might be beneficial to average across more seeds to observe clear trends. (B) As a potential user, I would find it very worrisome that BNN-NRE has _significantly_ lower nominal log-posterior than standard methods. Indeed, for some tasks, it seems to require 5 orders of magnitude (for many other tasks 3 orders of magnitude) more simulations than standard methods.

---

confidence: 2
contribution: 3
initial_rating: 5
presentation: 2
questions: How can we be sure that any benefits are from the BNN modelling epistemic uncertainty, and not the altered initialization? For example, if we "pretrained" NPE/NRE to prior samples (and $x$ simulated from some noise distribution), such that at "initialization" the posterior estimate would be approximately equal to the prior for any $x$, it would be interesting to see if there is still be any benefit to the introduced method.

Most the benefits seem to be in the very small data regime (<100 simulations). I am not convinced that this scenario is of particular interest the scientific community, could you give an example of a practical case where this is useful?

Why is the convergence to the NPE solution so slow in terms of the number of simulations? Can the prior be altered to avoid this problem, rather than introducing a temperature parameter?
soundness: 2
strengths: Being able to perform reliable simulation-based inference with fewer simulations would be broadly useful to the scientific community. The idea to use a prior over the weights of a BNN that is compatible with the simulators prior distribution is sensible, and using a Gaussian process to achieve this is to the best of my knowledge novel and clever. The paper is well structured overall (except for where noted below) and includes some interesting experiments, using a good choice of metrics assessing reliability of the posteriors.
summary: The authors wish to provide a method for simulation-based inference for expensive simulators that accounts for the epistemic uncertainty, which tends to be high when available data (simulations) is low. They wish to do so by using a Bayesian neural network in combination with the SBI methods NRE and NPE. They note issue with current priors for Bayesian neural networks in simulation based inference, stating they are not calibrated a priori. They note that using a Gaussian process, with the prior distribution as its mean, can be used to define a prior which is well-calibrated a priori. By optimizing the Bayesian networks prior on the weights to approximate the Gaussian process prior, the Bayesian networks prior becomes approximately well-calibrated. They perform experiments, and show this often leads to better calibrated posteriors in the low-data (low simulation budget) regime.
weaknesses: Structure/Presentation:
- Lack of method summary: There's no clear summary of the method, e.g. in the contributions section, an algorithm, or the start of section 3. A brief overview of 3-4 sentences would be very beneficial for readers. It's only towards the end of section 3.2, that the overarching method begins to come together on a first read. In my opinion, the abstract also does not include enough information on the method.
- Lack of clear definitions: The concept of an "a priori-calibrated Bayesian model" should be defined more clearly when it is first introduced.  I presume that an example of a well calibrated a priori model, would be one for which the Bayesian model average at initialization is equal to the simulator parameters prior $p(\theta)$ for any $x \in \mathcal{X}$? Presumably this includes a large variety of useful and non-useful models for modelling epistemic uncertainty (as suggested in equation 8).
- Figures: The credibility figure in Figure 1 is confusing, either the standard deviations are not in order, or there is a typo one of the standard deviations. Figure 3 should have a labeled legend. Many of the font sizes are too small.

Experiments:
- Limited experimental runs: Only 3 runs are performed for each experiment, and only the median is reported. This makes it hard to assess if differences between methods are significant. Repeating with more runs would be beneficial, although I am sympathetic to limitations in computational budget.
- The BNN-NRE method appears to perform worse than NRE in terms of mass placed on true parameters, even for low simulation budgets (100-1000 in Figure 2), which is the domain where the method is proposed to be beneficial. Similarly, the results for the NPE case are not particularly convincing in terms of the mass placed on the true parameters. The coverage properties do appear to have improved, but coverage alone is not indicative of a good posterior estimate.
- The results don't align with an intuitive understanding of epistemic uncertainty. For example, in figure 3, in NPE, we can see 4000 simulations produces a reasonable posterior estimate, whereas BNN-NPE is still massively conservative, even with 65536 simulations, suggesting overestimation of epistemic uncertainty. This suggests simple alternative approaches such as training posterior estimates on subsets of the data with standard methods, such as NPE, and mixing the resulting posteriors, would likely be more effective. I understand the temperature parameter is introduced to limit this problem, but this then introduces another hard to choose hyperparameter (along with the Gaussian process parameters introduced).

Meta-review: The authors propose a low-budget Bayesian Neural Network (BNN) designed to provide accurate posterior estimates when limited data is available due to the high cost of simulations. The reviewers commended the algorithm and noted that the paper was generally easy to follow. However, they found the results insufficiently robust to support the claims made in the paper. Both the reviewers and the Area Chair acknowledged the challenge of obtaining more samples, which is a key motivation for this work. Nevertheless, the bar for acceptance as a new, reliable method is higher than simply demonstrating utility for a specific problem. While the results with 10 samples show slight improvements, the performance with only three samples could be a fluke. Moreover, as the number of samples increases, the proposed method fails to match the quality of state-of-the-art results.

Decision: Reject