SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Maciej Falkiewicz $^{1,2}$  Naoya Takeishi $^{3,4}$  Imahn Shekhzadeh $^{1,2}$  Antoine Wehenkel $^{5}$  Arnaud Delaunoy $^{5}$  Gilles Louppe $^{5}$  Alexandros Kalousis $^{2}$ $^{1}$ Computer Science Department, University of Geneva  $^{2}$ HES-SO/HEG Genève  
 $^{3}$ The University of Tokyo  $^{4}$ RIKEN  $^{5}$ University of Liège  
{maciej.falkiewicz, imahn.shekhzadeh, alexandros.kalousis}@hesge.ch ntake@g.ecc.u-tokyo.ac.jp  
{antoine.wehenkel, a.delaunoy, g.louppe}@uliege.be
Bayesian inference allows expressing the uncertainty of posterior belief under a probabilistic model given prior information and the likelihood of the evidence. Predominantly, the likelihood function is only implicitly established by a simulator posing the need for simulation-based inference (SBI). However, the existing algorithms can yield overconfident posteriors (Hermans et al., 2022) defeating the whole purpose of credibility if the uncertainty quantification is inaccurate. We propose to include a calibration term directly into the training objective of the neural model in selected amortized SBI techniques. By introducing a relaxation of the classical formulation of calibration error we enable end-to-end backpropagation. The proposed method is not tied to any particular neural model and brings moderate computational overhead compared to the profits it introduces. It is directly applicable to existing computational pipelines allowing reliable black-box posterior inference. We empirically show on six benchmark problems that the proposed method achieves competitive or better results in terms of coverage and expected posterior density than the previously existing approaches.
Inverse problems [46] appear in a variety of fields of human activity, including engineering [2], geology [8], medical imaging [29], economics [21], and particle physics [1]. In some domains, uncertainty quantification is required for results to be used in the real-world [51]. A particularly well-suited tool for this task is Bayesian inference, where the procedure of updating prior beliefs (knowledge) with observed evidence is formally described in the framework of probability theory. In order to compute the posterior probability according to Bayes' rule, the likelihood function has to be evaluated. Often, it is only implicitly established by a simulator which given input parameters generates (possibly stochastically) outputs. Such a situation requires the use of likelihood-free techniques, also referred to as simulation-based inference (SBI) [9].
There exist two major classical approaches that address the problem of Bayesian inference in the likelihood-free setting. In the first one, simulated data are used to estimate the distribution of expert-crafted low-dimensional summary statistics with a nonparametric method like kernel density estimation [14]. Next, the approximate density is used as a likelihood surrogate turning the problem into likelihood-based inference. In the second approach, termed as Approximate Bayesian Computation (ABC) [40, 38, 48, 42, 3, 44] no data is simulated upfront, but instead for every newly appearing observation (instance of the inverse problem), an independent sampling-rejection procedure with a simulator in-the-loop is run. Given some proposal distribution, candidate
parameters are accepted if the simulator's output is in sufficient agreement with the observation. While ABC effectively sorts out the lack of explicit likelihood, it is generally costly, especially when the underlying simulator is computationally expensive.
Deep Neural Networks have been increasingly utilized in SBI to alleviate these costs, by acting as universal and powerful density estimators or surrogating other functions. The following approaches of neural SBI can be identified: Neural Likelihood Estimation (NLE) [35], where the missing likelihood function is approximated based on the simulator; Neural Ratio Estimation (NRE) [19, 15, 32], where the likelihood-to-evidence ratio (other ratios have also been considered) is approximated; Neural Posterior Estimation (NPE) [34, 30, 18], where the learned density estimator directly solves the inverse problem for a given observation. In the last category the score-based methods [41, 16] are emerging, where the posterior is implicitly expressed by the learned score function. All of the approaches mentioned above amortize the use of a simulator, however, inference with NLE and NRE still requires the use of Markov chain Monte Carlo [7], while NPE can be fully amortized with the use of appropriate generative models [47].
The question of reliability is raised with the use of approximate Bayesian inference methods. Many works have been devoted to inference evaluation techniques in the absence of the ground-truth posterior reference [45, 53, 27, 28]. Lueckmann et al. [31] propose an SBI benchmarking protocol taking into account the effect of the simulation budget, with the observation that the assessment's conclusions will strongly depend on the metric chosen. Hermans et al. [20] suggest relying on coverage analysis of posterior's multi-dimensional credible regions, showing that numerous SBI techniques yield overconfident posteriors. In a follow-up work [11] they introduce a balancing regularizer for training NRE, and show empirically that it leads to the avoidance of overconfident solutions. Very recently, an interesting study by Delaunoy et al. [12] extends the applicability of balancing regularizer to the NPE approach. An alternative perspective is studied in Dey et al. [13] where an additional regression model for posthoc re-calibration based on Probability Integral Transform is learned for the outputs of a pre-trained inference model.
In this work, we introduce and empirically evaluate a new regularizer that directly targets the coverage of the approximate posterior. It is applicable to NRE and NPE based on neural density estimators such as Normalizing Flows [26, 36]. As an additive term, it is easy for practitioners to use as a plug-in to existing pipelines. The additional computational overhead introduced by the regularizer is justified in view of the advantages it brings.
The rest of the work is structured as follows. Section 2 introduces the formal description of the studied problem and presents the coverage analysis that is the motivation for the presented method. Section 3 describes the proposed method and a differentiable formulation of the regularizer that allows gradient-based learning. Section 4 shows the empirical results of applying the method on a wide range of benchmark problems, as well as providing a hyper-parameter sensitivity analysis. In section 5, we provide summarize our work with conclusions and potential future directions.
In this work, we consider the problem of estimating posterior distribution  $p(\theta | x)$  of model parameters  $\theta$ , given the observation  $x$ . We approach the problem from a likelihood-free [43] perspective, where the model's likelihood  $p(x | \theta)$  cannot be evaluated. Instead, there exists a simulator that generates observations  $x$  based on the input parameters  $\theta$ . We can prepare a dataset  $\{ (\theta_i, x_i) \}_{i=1}^N$ , by sampling parameters from a prior  $p(\theta)$  and running them through the simulator. In that way, the problem of finding the approximate posterior marked as  $\hat{p}(\theta | x)$  turns into a supervised learning problem. However, it is important to note that usually there is not only a single generating parameter  $\theta$  for a given  $x$  possible. Instead, for each observation multiple parameter values are plausible [23]. These values form a posterior distribution  $p(\theta | x)$  under a prior distribution  $p(\theta)$ .
Assessing the performance of an SBI technique is challenging due to the non-availability of the ground-truth posterior  $p(\theta | x)$  in the general case. Theoretically, methods such as the rejection-sampling ABC [44] are proven to converge to  $p(\theta | x)$  under an infinite simulation budget but often require a large number of simulation runs to achieve a good approximation. Avoiding this costly procedure is exactly what motivates the introduction of different techniques. Therefore, it is common to resort to performance metrics evaluated on test datasets sampled from the simulator.
If a Bayesian inference method allows  $\hat{p} (\theta |x)$  evaluation, its performance can be assessed by looking at the expected value of the approximate log posterior density of the data-generating parameters (nominal parameters)  $\mathbb{E}_{p(\theta ,x)}[\log \hat{p} (\theta |x)]$ , estimated over a test dataset. In domains such as natural sciences, it is also accurate uncertainty quantification that matters, i.e. reliable posterior density estimation given the available evidence in data. It has been shown in Hermans et al. [20] that approaches yielding high expected posterior density values can give overconfident posteriors, and therefore, verifying the quality of credible regions returned by inference engines as well as designing more reliable ones is essential. In this work, we define a regularizer suitable for training NRE and NPE based on the Simulation-Based Calibration (SBC) diagnostic of Talts et al. [45]. To make the paper self-contained, in the rest of Section 2 we will introduce SBC formally, adapting the definitions and lemmas from Lemos et al. [27], and making the link with the coverage analysis from Hermans et al. [20].
2.1 Coverage analysis
Coverage analysis offers a way to evaluate the quality of an approximate posterior. Given an observation and an estimated distribution, one can assess the probability that a certain credible region [22] of the parameter space contains the nominal parameter value found in the dataset. For a conditional probability distribution, a credible region at level  $1 - \alpha$  is a region of the support to which the parameter used to generate the conditioning observation belongs with a  $1 - \alpha$  probability, called the credibility level. There are infinitely many credible regions at any given level, and coverage analysis requires a particular way of choosing one. In this work, we will consider the Highest Posterior Density Region (HPDR).
Definition 1 (Highest Posterior Density Region).
$$
\Theta_ {\hat {p} (\theta | x)} ^ {H P D R} (1 - \alpha) := \underset {\Theta} {\arg \min } | \Theta | s. t. \int_ {\theta \in \Theta} \hat {p} (\theta | x) d \theta = 1 - \alpha , \tag {1}
$$
where  $|\Theta| \coloneqq \int_{\Theta} \mathrm{d}\mu$  is the volume of a subset  $\Theta$  of the parameters' space, under some measure  $\mu$ .
Next, we introduce the central concept of coverage analysis.
Definition 2 (Expected Coverage Probability). The Expected Coverage Probability (ECP) over a joint distribution  $p(\theta, x)$  for a given  $\Theta_{\hat{p}(\theta|x)}^{HPDR}(1 - \alpha)$  is
$$
E C P \left(\Theta_ {\hat {p} (\theta | x)} ^ {H P D R} (1 - \alpha)\right) := \mathbb {E} _ {p (\theta , x)} \left[ \mathbb {1} \left[ \theta \in \Theta_ {\hat {p} (\theta | x)} ^ {H P D R} (1 - \alpha) \right] \right], \tag {2}
$$
where  $\mathbb{1}[\cdot ]$  is the indicator function.
While the expectation in eq. (2) can be easily estimated by using a set of i.i.d. samples from the joint distribution  $p(\theta, x)$ , the challenge of finding  $\Theta_{\hat{p}(\theta|x)}^{\mathrm{HPDR}}(1 - \alpha)$  remains non-trivial. In general, directly following eq. (1) requires solving an optimization problem. This would be prohibitive to do for every instance during training, thus it is crucial to benefit from the two lemmas we introduce below.
Lemma 1. A pair  $(\theta^{*},x^{*})$  and a distribution  $\hat{p} (\theta |x)$  uniquely define an HPDR:
$$
\Theta_ {\hat {p} (\theta | x)} ^ {H P D R} \left(1 - \alpha_ {H P D R} \left(\hat {p}, \theta^ {*}, x ^ {*}\right)\right) := \left\{\theta \mid \hat {p} \left(\theta | x ^ {*}\right) \geqslant \hat {p} \left(\theta^ {*} \mid x ^ {*}\right) \right\}, \tag {3}
$$
$$
\alpha_ {H P D R} (\hat {p}, \theta^ {*}, x ^ {*}) := \int \hat {p} (\theta | x) \mathbb {1} [ \hat {p} (\theta | x ^ {*}) <   \hat {p} (\theta^ {*} | x ^ {*}) ] d \theta . \tag {4}
$$
See proof in Appendix A.
Lemma 2.
$$
E C P \left(\Theta_ {\hat {p} (\theta | x)} ^ {H P D R} (1 - \alpha)\right) = \mathbb {E} _ {p (\theta , x)} \left[ \mathbb {1} \left[ \alpha_ {H P D R} (\hat {p}, \theta , x) \geqslant \alpha \right] \right] \tag {5}
$$
See proof of Lemma 1 in Lemos et al. [27].
If  $\mathrm{ECP}\left(\Theta_{\hat{p} (\theta |x)}^{\mathrm{HPDR}}(1 - \alpha)\right)$  is equal to  $1 - \alpha$  the inference engine yields posterior distributions that are calibrated at level  $1 - \alpha$ , or simply calibrated if the condition is satisfied for all  $\alpha$ . The ground-truth posterior  $p(\theta |x)$  is calibrated, but another such distribution is the prior  $p(\theta)$  as discussed in Hermans
et al. [20] and Lemos et al. [27]. Therefore, coverage analysis should be applied in conjunction with the standard posterior predictive check. Formally, we aim for:
$$
p ^ {*} (\theta | x) := \underset {\hat {p} (\theta | x)} {\arg \max } \mathbb {E} _ {p (\theta , x)} [ \log \hat {p} (\theta | x) ] \quad \text {s . t .} \forall \alpha \operatorname {E C P} \left(\Theta_ {\hat {p} (\theta | x)} ^ {\mathrm {H P D R}} (1 - \alpha)\right) = 1 - \alpha \tag {6}
$$
In Hermans et al. [20] it has been pointed out that for practical applications often the requirement of the equality sign in the condition in eq. (6) can be relaxed and replaced with  $\mathrm{ECP}\left(\Theta_{\hat{p} (\theta |x)}^{\mathrm{HPDR}}(1 - \alpha)\right)$  above or equal the  $1 - \alpha$  level. Such a posterior distribution is called conservative.
We define the calibration error as
$$
\frac {1}{M} \sum_ {i = 1} ^ {M} | (1 - \alpha_ {i}) - \operatorname {E C P} \left(\Theta_ {\hat {p} (\theta | x)} ^ {\mathrm {H P D R}} (1 - \alpha_ {i})\right) |, \tag {7}
$$
and the conservativeness error as
$$
\frac {1}{M} \sum_ {i = 1} ^ {M} \max  \left((1 - \alpha_ {i}) - \mathbf {E C P} \left(\Theta_ {\hat {p} (\theta | x)} ^ {\mathrm {H P D R}} (1 - \alpha_ {i})\right), 0\right), \tag {8}
$$
where  $\{\alpha_i\}_{i=1}^M$  is a set of credibility levels of arbitrary cardinality such that  $\forall_i 0 < \alpha_i < 1$ ; we can exclude  $\alpha = 0$  and  $\alpha = 1$  because calibration/conservativeness at these levels is achieved by construction. For a motivating graphical example of coverage analysis of an overconfident distribution see appendix D.
We propose a method supporting the training of neural models underlying SBI engines such that the resulting distributions are calibrated or conservative against the real  $p(\theta, x)$  in the data. In addition to the standard loss term in the optimization objective function, we include a regularizer corresponding to differentiable relaxations of the calibration (conservativeness) error as defined in eq. (7) (eq. (8)). In the remainder of this section, we will show how to efficiently estimate the expected coverage probability during training and how to alleviate the arising non-continuities to allow gradient-based optimization. After a few simplifications, the regularizer will become a mean square error between the sorted list of estimated  $\alpha_{\mathrm{HPDR}}(\hat{p}, \theta, x)$ s and a step function on the  $(0,1)$  interval.
The proposed method is based on the Monte Carlo (MC) estimate of the integral in eq. (4):
$$
\hat {\alpha} _ {\mathrm {H P D R}} ^ {L} (\hat {p}, \theta^ {*}, x ^ {*}) := \frac {1}{L} \sum_ {j = 1} ^ {L} \mathbb {1} [ \hat {p} \left(\theta_ {j} \mid x ^ {*}\right) <   \hat {p} \left(\theta^ {*} \mid x ^ {*}\right) ], \tag {9}
$$
where  $\theta_{1},\dots,\theta_{L}$  are samples drawn from  $\hat{p} (\theta |x)$ . The entity defined in eq. (9) is a special case of what is called rank statistic in Talts et al. [45] (with the  $1 / L$  normalization constant skipped there). The non-normalized variant is proven to be uniformly distributed over the integers  $[0,L]$  for  $\hat{p} (\theta |x) = p(\theta |x)\forall (\theta ,x)\sim p(\theta ,x)$ , which is the accurate posterior as defined in Lemos et al. [27]. Taking  $L\to \infty$  and applying the normalization (as done in eq. (9)) one shows that distribution of  $\hat{\alpha}_{\mathrm{HPDR}}^{\infty}(\hat{p},\theta^{*},x^{*})$  converges to a standard continuous uniform distribution (proof in appendix A.2).
We return to coverage analysis with the following remark:
Remark 1. An inference engine yields calibrated posterior distributions if and only if it is characterized by uniformly distributed  $\alpha_{HPDR}(\hat{p},\theta ,x)s$
Following remark 1, we will enforce the uniformity of the normalized rank statistics which coincides with a necessary, though not sufficient, condition for the approximate posterior to be accurate. The one-sample Kolmogorov-Smirnov test provides a way to check whether a particular sample set follows some reference distribution. We will use its test statistic to check, and eventually enforce, the condition. An alternative learning objective is to directly minimize the calibration (conservativeness) error at arbitrary levels with ECP computed following eq. (5).
3.1 One-sample Kolmogorov-Smirnov test
Let  $B$  be a random variable taking values  $\beta \in \mathcal{B}$ . Given a sample set  $\hat{\mathcal{B}} = \{\beta_i | i = 1, \dots, N\}$ , the one-sample Kolmogorov-Smirnov test statistic [5] for r.v.  $B$  is defined as:
$$
D _ {N} = \sup  _ {\beta \in \mathcal {B}} | F _ {N} (\beta) - F (\beta) |, \tag {10}
$$
where  $F_{N}$  is the empirical cumulative distribution function (ECDF) of the r.v.  $B$ , with regard to the sample set  $\hat{B}$  [10].  $F$  is the reference cumulative distribution function (CDF) against which we compare the ECDF.
In the proposed method, we instantiate the KS test statistic for r.v.  $\alpha_{\mathrm{HPDR}}(\hat{p},\theta ,x)$  where  $(\theta ,x)\sim p(\theta ,x)$ , thus  $F_{N}$  is its ECDF, which we estimate from a sample set  $\hat{\mathcal{A}} := \{\hat{\alpha}_{\mathrm{HPDR}}^{L}(\hat{p},\theta_{i},x_{i})|(x_{i},\theta_{i})\sim p(x,\theta),i = 1,\ldots ,N\}$ . The reference CDF,  $F$ , is that of the standard uniform distribution because we want to check and enforce the uniformity of the elements of  $\hat{\mathcal{A}}$ .
By the Glivenko-Cantelli theorem [49] we know that  $D_N \to 0$  almost surely with  $N \to \infty$  if samples follow the  $F$  distribution. The claim in the opposite direction, that is if  $D_N \to 0$  then samples follow  $F$ , is also true and we show that in appendix A.3. As a result, minimizing  $D_N$  will enforce uniformity and thus posterior's calibration. We can upper-bound the supremum in eq. (10) with the sum of absolute errors (SAE), but instead, to stabilize gradient descent, we use the mean of squared errors (MSE), which shares its optimum with SAE. Minimizing the one-sample KS test statistic is just one of the many ways of enforcing uniformity. We can deploy any divergence measure between the sample set and the uniform distribution.
3.2 Approximating  $D_N$
In practice we cannot evaluate the squared errors over the complete support of the r.v., but only over a finite subset. There are two ways of evaluating this approximation which differ by how we instantiate  $F_{N}(\alpha)$  and  $F(\alpha)$ .
Direct computation In the first, we compute  $F_{N}(\alpha) = \sum_{i}\mathbb{1}[\alpha_{i}\leq \alpha ] / N$  summing over all  $\alpha_{i}\in \hat{\mathcal{A}}$  and use the fact that  $F(\alpha) = \alpha$  for the target distribution. To approximate  $D_N$  we aggregate the squared error over an arbitrary set of  $\alpha_{k}$  values:  $\sum_{k}(F_{N}(\alpha_{k}) - \alpha_{k})^{2}$ . Including this term as a regularizer means that we need to backpropagate through  $F_{N}(\alpha_{k})$ . Such an approach can be seen as directly aiming for calibration at arbitrary levels  $\alpha_{k}$ , i.e. applying eq. (5) to evaluate the calibration error.
Sorting-based computation Alternatively, we can sort the values in  $\hat{\mathcal{A}}$  to get the ordered sequence  $(\alpha_{i}|i:=1\dots N)$ . Given the sorted values,  $\alpha_{i}$ , and their indices,  $i$ , the  $F_{N}(\alpha_{i})$  values are given directly by the indices normalized by the sample size  $(i/N)$ , and by the corresponding target distribution  $F(\alpha_{i}) = \alpha_{i}$ , which leads to the  $\sum_{i}(i/N - \alpha_{i})^{2}$  approximation. Under this approach, the gradient will now flow through  $F(\alpha_{i})$ . This is the approach that we use in our experiments.
3.3 Importance Sampling integration
In the definition of  $\alpha_{\mathrm{HPDR}}(\hat{p},\theta^{*},x^{*})$  in eq. (4) as well as in its MC estimate  $\hat{\alpha}_{\mathrm{HPDR}}^L (\hat{p},\theta^*,x^*)$  defined in eq. (9) we sample  $\theta_{j}\mathbf{s}$  from the approximate posterior  $\hat{p} (\theta |x)$ . Optimizing the model constituting the approximate posterior requires backpropagating through the sampling operation which can be done with the reparametrization trick [24]. However, the reparametrization trick is not always applicable; for example, in NRE an MC procedure is required to sample from the posterior [19]. Here we can use instead self-normalized Importance Sampling (IS) to estimate  $\hat{\alpha}_{\mathrm{HPDR}}^L (\hat{p},\theta ,x)$  and lift the requirement of sampling from the approximate posterior. We will be sampling from a known proposal distribution and only evaluate the approximate posterior density. We thus propose to use IS integration [e.g., 17] to estimate the integral in eq. (4), in the following way:
$$
\hat {\alpha} _ {\mathrm {H P D R}} ^ {L, I S} (\hat {p}, \theta^ {*}, x ^ {*}) = \frac {\sum_ {j = 1} ^ {L} \hat {p} \left(\theta_ {j} \mid x ^ {*}\right) / I \left(\theta_ {j}\right) \mathbb {1} \left[ \hat {p} \left(\theta_ {j} \mid x ^ {*}\right) <   \hat {p} \left(\theta^ {*} \mid x ^ {*}\right) \right]}{\sum_ {j = 1} ^ {L} \hat {p} \left(\theta_ {j} \mid x ^ {*}\right) / I \left(\theta_ {j}\right)}. \tag {11}
$$
Algorithm 1 Computing the regularizer loss with calibration objective.
Require: Data batch  $\{(\theta_i, x_i)\}_{i=1}^N$ , model  $\hat{p}(\theta|x)$ , number of samples  $L$ , proposal distribution  $I(\theta)$
Ensure: Regularizer's loss  $R$
1: for  $i\gets 1$  to  $N$  do
2:  $p_i \gets \hat{p}(\theta_i | x_i)$
3: for  $j\gets 1$  to  $L$  do
4:  $\theta_i^j\sim I(\theta)$
5:  $p_i^j \gets \hat{p}(\theta_i^j | x_i)$
6: end for
7:  $\hat{\alpha}_{\mathbf{HPDR}}^{L,IS}(\hat{p},\theta_i,x_i)\gets \frac{\sum_{j = 1}^L p_i^j / I(\theta_i^j)\mathbb{1}[p_i^j < p_i]}{\sum_{j = 1}^L p_i^j / I(\theta_i^j)}$ $\triangleright$  eq. (11)
8: end for
9:  $(\alpha_{i}|i = 1,\dots ,N)\gets \mathrm{sort}(\{\hat{\alpha}_{\mathrm{HPDR}}^{L,IS}(\hat{p},\theta_{i},x_{i})|i = 1,\dots ,N\})$
10:  $R\gets \frac{1}{N}\sum_{i}^{N}\left(i / N - \alpha_{i}\right)^{2}$   the second term in eq. (12)
The  $\theta_{j}$ s are sampled from the proposal distribution  $I(\theta)$ . For a uniform proposal distribution, the  $I(\theta_{j})$  terms are constant and thus cancel out. By default, we suggest using the prior as the proposal distribution. While IS is not needed in the case of NPE, we use it also for this method to keep the regularizer's implementation consistent with the NRE setting.
3.4 Learning objective
So far we have shown how to derive a regularizer that enforces a uniform distribution of the normalized rank statistics that characterizes a calibrated model as shown in remark 1. This is also a necessary condition for the accuracy of the posterior as shown in Talts et al. [45] but not a sufficient one, as exemplified by the approximate posterior equal to the prior for which the regularizer trivially obtains its minimum value. Thus the final objective also contains the standard loss term and is a weighted sum of the two terms:
$$
\frac {1}{N} \sum_ {i} ^ {N} \log \hat {p} \left(\theta_ {i} \mid x _ {i}\right) + \lambda \frac {1}{N} \sum_ {i} ^ {N} \left(\operatorname {i n d e x} \left(\hat {\alpha} _ {\mathrm {H P D R}} ^ {L, I S} (\hat {p}, \theta_ {i}, x _ {i})\right) / N - \hat {\alpha} _ {\mathrm {H P D R}} ^ {L, I S} (\hat {p}, \theta_ {i}, x _ {i})\right) ^ {2}, \tag {12}
$$