SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Daolang Huang*
daolang.huang@aalto.fi
Ayush Bharti*
ayush.bharti@aalto.fi
Amauri H. Souza
Federal Institute of Ceará amauri.souza@aalto.fi
Luigi Acerbi
luigi.acerbi@helsinki.fi
Samuel Kaski
University of Manchester samuel.kaski@aalto.fi
Simulation-based inference (SBI) methods such as approximate Bayesian computation (ABC), synthetic likelihood, and neural posterior estimation (NPE) rely on simulating statistics to infer parameters of intractable likelihood models. However, such methods are known to yield untrustworthy and misleading inference outcomes under model misspecification, thus hindering their widespread applicability. In this work, we propose the first general approach to handle model misspecification that works across different classes of SBI methods. Leveraging the fact that the choice of statistics determines the degree of misspecification in SBI, we introduce a regularized loss function that penalizes those statistics that increase the mismatch between the data and the model. Taking NPE and ABC as use cases, we demonstrate the superior performance of our method on high-dimensional time-series models that are artificially misspecified. We also apply our method to real data from the field of radio propagation where the model is known to be misspecified. We show empirically that the method yields robust inference in misspecified scenarios, whilst still being accurate when the model is well-specified.
Bayesian inference traditionally entails characterizing the posterior distribution of parameters assuming that the observed data came from the chosen model family [7]. However, in practice, the true data-generating process rarely lies within the family of distributions defined by the model, leading to model misspecification. This can be caused by, among other things, measurement errors or contamination in the observed data that are not included in the model, or when the model fails to capture the true nature of the physical phenomenon under study. Misspecified scenarios are especially likely for simulator-based models, where the goal is to describe some complex real-world phenomenon. Such simulators, also known as implicit generative models [25], have become prevalent in many domains of science and engineering such as genetics [71], ecology [3], astrophysics [40], economics [29], telecommunications [8], cognitive science [79], and agent-based modeling [87]. The growing field of simulation-based inference (SBI) [21, 51] tackles inference for such intractable likelihood models, where the approach relies on forward simulations from the model, instead of likelihood evaluations, to estimate the posterior distribution.
Traditional SBI methods include approximate Bayesian computation (ABC) [4, 51, 76], synthetic likelihood [67, 85] and minimum distance estimation [15]. More recently, the use of neural networks
Figure 1: Inference on misspecified Ricker model with two parameters, see Section 4 for details. Left: Learned statistics obtained using NPE; each point (in black) corresponds to a parameter value sampled from the prior. The observed statistic (in red) is outside the set of statistics the model can simulate. Middle: Statistics learned using our robust method; the observed statistic is inside the distribution of simulated statistics. Right: The resulting posteriors obtained from NPE and our method. The posterior obtained from our method is close to the true parameter value, while that from NPE is completely off, going outside the prior range (denoted by dashed grey lines).
as conditional density estimators [72] has fuelled the development of new SBI methods that target either the likelihood (neural likelihood estimation) [52, 61, 75, 83], the likelihood-to-evidence ratio (neural ratio estimation) [22, 23, 28, 30, 39, 46, 77], or the posterior (neural posterior estimation) [38, 41, 53, 58, 59, 70, 75, 83]. A common feature of all these SBI methods is the choice of summary statistics of the data, which readily impacts their performance [14]. The summary statistics are either manually handcrafted by domain experts, or learned automatically from simulated data using neural networks [1, 17, 18, 26, 31, 47, 84].
The unreliability of Bayesian posteriors under misspecification is exacerbated for SBI methods as the inference relies on simulations from the misspecified model. Moreover, for SBI methods using conditional density estimators [5, 13, 41, 61, 69], the resulting posteriors can be wildly inaccurate and may even go outside the prior range [9, 16, 35]; see [27] for an instance of this problem in the study of molecules. This behavior is exemplified in Figure 1 for the posterior-targeting neural SBI method called neural posterior estimation (NPE). We see that when the model is misspecified, the NPE posterior goes outside the prior range. This is because the model is unable to match the observed statistic for any value of the parameters, as shown in Figure 1(left), and hence, NPE is forced to generalize outside its training distribution, yielding a posterior that is outside the prior range.
We argue that the choice of statistics is crucial to the notion of misspecification in SBI. Suppose we wish to fit a Gaussian model to data from some skewed distribution. If we choose, for instance, the sample mean and the skewness as statistics, then the Gaussian model would fail to match these statistics for any choice of parameters, and we would end up in the same situation as in Figure 1 where the SBI method is forced to generalize outside its training distribution. However, choosing the sample mean and sample variance as statistics instead may solve this issue, as the model may be able to replicate the chosen statistics for some value of parameters, even though the Gaussian model itself is misspecified for any skewed dataset. This problem of identifying a low-dimensional statistic of a high-dimensional dataset for which the model of interest is correctly specified is termed the data selection problem [82]. We refer to the statistics that solve the data selection problem as being robust.
Contributions. In this paper, we learn robust statistics using neural networks to solve the data selection problem. To that end, we introduce a regularized loss function that balances between learning statistics that are informative about the parameters, and penalizing those choices of statistics or features of the data that the model is unable to replicate. By doing so, the observed statistic does not go outside the set of statistics that the model can simulate, see Figure 1(middle), thereby yielding a posterior around the true parameter even under misspecification (Figure 1(right)). As our method relies on learning appropriate statistics and not on the subsequent inference procedure, it is applicable to all the statistics-based SBI methods, unlike other robust methods in the literature [24, 48, 81]. We substantiate our claim in Section 4 through extensive numerical studies on NPE and ABC — two fundamentally different SBI methods. Additionally in Section 5, we apply our method to a real case of model misspecification in the field of radio propagation [8].
Related works. Model misspecification has been studied in the context of different SBI methods such as ABC [9, 33, 35, 36, 74], synthetic likelihood [32, 34, 57], minimum distance estimation [24], and neural estimators [48, 81]. It was first analyzed for ABC methods such as the regression-adjustment ABC [5] in [35], and a robust approach was proposed in [33]. In [9], the authors proposed to handle model misspecification in ABC by incorporating the domain expert while selecting the statistics. Generalized Bayesian inference [12, 43, 49], which is a popular framework for designing methods that are robust to misspecification, has also been linked to ABC [74] and synthetic likelihood [57]. In [24], the authors propose a robust SBI method based on Bayesian nonparametric learning and the posterior bootstrap. More recently, robustness to model misspecification has been tackled for neural likelihood [48] and posterior estimators [81]. We use the latter method, called robust neural posterior estimation (RNPE) [81], as a baseline for comparing our proposed approach. Our method differs from these previous robust approaches in that it is applicable across distinct SBI methods.
Simulation-based inference (SBI). Consider a simulator  $\{\mathbb{P}_{\theta} : \theta \in \Theta\}$ , which is a family of distributions parameterized by  $\theta$  on data space  $\mathcal{X} \subseteq \mathbb{R}^d$ . We assume that  $\mathbb{P}_{\theta}$  is intractable, but sampling independent and identically distributed (iid) data  $\mathbf{x}_{1:n} = \{\mathbf{x}^{(1)}, \ldots, \mathbf{x}^{(n)}\} \sim \mathbb{P}_{\theta}$  is straightforward for  $\mathbf{x}^{(i)} \in \mathcal{X}, 1 \leq i \leq n$ . Given iid observed data  $\mathbf{y}_{1:n} \sim \mathbb{Q}, \mathbf{y}^{(i)} \in \mathcal{X}$ , from some true data-generating process  $\mathbb{Q}$  and a prior distribution of  $\theta$ , we are interested in estimating a  $\theta^{\star} \in \Theta$  such that  $\mathbb{P}_{\theta^{\star}} = \mathbb{Q}$ , or alternatively, in characterizing the posterior distribution  $p(\theta | \mathbf{y}_{1:n})$ . Since  $\mathcal{X}$  is a high-dimensional space in most practical cases, it is common practice to project both the simulated and the observed data onto a low-dimensional space of summary statistics  $S$  via a mapping  $\eta : \mathcal{X}^n \to S$ , such that  $\mathbf{s} = \eta(\mathbf{x}_{1:n})$  is the vector of simulated statistics, and  $\mathbf{s}_{\mathrm{obs}} = \eta(\mathbf{y}_{1:n})$  the observed statistic. In the following, we assume that this deterministic map  $\eta$  is not known a priori, but needs to be learned using a neural network, henceforth called the summary network  $\eta_{\psi}$ , with trainable parameters  $\psi$ . We now introduce two different paradigms for learning the statistics.
Learning statistics in neural posterior estimation (NPE) framework. NPEs are a class of SBI methods which are becoming popular in astrophysics [40] and nuclear fusion [37, 80], among other fields. They map each statistic vector  $\mathbf{s}$  to an estimate of the posterior  $p(\theta|\mathbf{s})$  using conditional density estimators such as normalizing flows [41, 53, 59, 69]. The posterior is assumed to be a member of a family of distributions  $q_{\nu}$ , parameterized by  $\nu$ . NPEs map the statistics to distribution parameters  $\nu$  via an inference network  $h_{\phi}$ , where  $\phi$  constitutes the weights and biases of  $h$ , such that  $q_{h_{\phi}(\mathbf{s})}(\theta) \approx p(\theta|\mathbf{s})$ . The inference network is trained on dataset  $\{(\theta_i, \mathbf{s}_i)\}_{i=1}^m$ , simulated from the model using prior samples  $\{\theta_i\}_{i=1}^m \sim p(\theta)$ , by minimising the loss  $\mathcal{L}(\phi) = -\mathbb{E}_{p(\theta, \mathbf{s})}[\log q_{h_{\phi}(\mathbf{s})}(\theta)]$ . Once trained, the posterior estimate is then obtained by simply passing the observed statistic  $\mathbf{s}_{\mathrm{obs}}$  through the trained network. Hence, NPEs enjoy the benefit of amortization, wherein inference on new observed dataset is straightforward after a computationally expensive training phase. In cases where the statistics are not known a priori, the NPE framework allows for joint training of both the summary and the inference networks by minimizing the loss function
$$
\mathcal {L} _ {\mathrm {N P E}} (\phi , \psi) = - \mathbb {E} _ {p \left(\theta , \mathbf {x} _ {1: n}\right)} \left[ \log q _ {h _ {\phi} \left(\eta_ {\psi} \left(\mathbf {x} _ {1: n}\right)\right)} (\theta) \right], \tag {1}
$$
on the training dataset  $\{(\theta_{i},\mathbf{x}_{1:n,i})\}_{i = 1}^{m}$  [69]. Even though NPEs are flexible and efficient SBI methods, they are known to perform poorly when the model is misspecified [16, 81]. This is because the observed statistic  $\mathbf{s}_{\mathrm{obs}}$  becomes an out-of-distribution sample under misspecification (see Figure 1 for an example), forcing the inference network in NPE to generalize outside its training distribution.
Learning statistics using autoencoders. Unlike NPEs, statistics for other SBI methods are learned prior to carrying out the inference procedure. This is achieved, for instance, by training an autoencoder with the reconstruction loss [1]
$$
\mathcal {L} _ {\mathrm {A E}} (\psi , \psi_ {d}) = \mathbb {E} _ {p (\theta , \mathbf {x} _ {1: n})} \left[ \left(\mathbf {x} _ {1: n} - \tilde {\eta} _ {\psi_ {d}} \left(\eta_ {\psi} \left(\mathbf {x} _ {1: n}\right)\right)\right) ^ {2} \right], \tag {2}
$$
where  $\psi$  and  $\psi_d$  are the parameters of the encoder  $\eta$  and the decoder  $\tilde{\eta}$ , respectively. The trained encoder  $\eta_{\psi}$  is then taken as the summarizing function in the SBI method. In this paper, we will use the encoder  $\eta_{\psi}$  to perform inference in an ABC framework, which we now recall.
Approximate Bayesian computation (ABC). ABC is arguably the most popular SBI method, and is widely used in many research domains [64]. ABC relies on computing the distance between the simulated and the observed statistics to obtain samples from the approximate posterior of a simulator. Given a discrepancy  $\rho(\cdot, \cdot)$  and a tolerance threshold  $\delta$ , the basic rejection-ABC method involves repeating the following algorithm: (i) sample  $\theta' \sim p(\theta)$ , (ii) simulate  $\mathbf{x}_{1:n}' \sim \mathbb{P}_{\theta'}$ , (iii) compute  $\mathbf{s}' = \eta(\mathbf{x}_{1:n}')$ , and (iv) if  $\rho(\mathbf{s}', \mathbf{s}_{\mathrm{obs}}) < \delta$ , accept  $\theta'$ , until a sample  $\{\theta_i'\}_{i=1}^{n_\delta}$  is obtained from the approximate posterior. Conditional density estimators have also been used in ABC to improve the posterior approximation of the rejection-ABC method [5, 11, 13, 14]. These so-called regression adjustment ABC methods involve fitting a function  $g$  between the accepted parameters and statistics pairs  $\{(\theta_i', \mathbf{s}_i')\}_{i=1}^{n_\delta}$  as  $\theta_i' = g(\mathbf{s}_i') + \omega_i$ ,  $1 \leq i \leq n_\delta$ , where  $\{\omega_i\}_{i=1}^{n_\delta}$  are the residuals. Once fitted, the accepted parameter samples are then adjusted as  $\tilde{\theta}_i' = \hat{g}(\mathbf{s}_{\mathrm{obs}}) + \hat{\omega}_i$ , where  $\hat{g}(\mathbf{s})$  is the estimate of  $\mathbb{E}[\theta|\mathbf{s}]$  and  $\hat{\omega}_i$  is the empirical residual. Despite the popularity of the regression adjustment ABC methods, they have also been shown to behave wildly under misspecification [9, 35], similar to NPEs.
Model misspecification in SBI. Bayesian inference assumes that there exists a  $\theta^{\star} \in \Theta$  such that  $\mathbb{P}_{\theta^{\star}} = \mathbb{Q}$ . Under model misspecification,  $\mathbb{P}_{\theta} \neq \mathbb{Q}$  for any  $\theta \in \Theta$ , which can lead to unreliable predictive distributions. However, this definition does not apply for models whose inference is carried out in light of certain summary statistics. For instance, a Gaussian model is clearly misspecified for a bimodal data sample. However, if inference is based on, say, the sample mean and the sample variance of the data, the model may still match the observed statistics accurately. Hence, even when a simulator is misspecified with respect to the true data-generating process, it may still be well-specified with respect to the observed statistic. The choice of statistics therefore dictates the notion of misspecification in SBI. We propose an analogous definition for misspecification in SBI:
Definition 3.1 (Misspecification of model-summulator pair in SBI). Let  $\eta_{\#} \mathbb{P}_{\theta}^{n}$  be a pushforward of the probability measure  $\mathbb{P}_{\theta}^{n} = \mathbb{P}_{\theta} \times \dots \times \mathbb{P}_{\theta}$  on  $\mathcal{X}^n$  under  $\eta : \mathcal{X}^n \to \mathcal{S}$ , meaning that for  $A \subset \mathcal{S}$ ,  $\eta_{\#} \mathbb{P}_{\theta}^{n}(A) = \mathbb{P}_{\theta}^{n}(\eta^{-1}(A)) = \mathbb{P}_{\theta}(\eta^{-1}(A)_1) \times \dots \times \mathbb{P}_{\theta}(\eta^{-1}(A)_n)$ . Then,  $\{\eta_{\#} \mathbb{P}_{\theta}^{n} : \theta \in \Theta\}$  is a set of distributions on the statistics space  $\mathcal{S}$  induced by the model for a given summulator  $\eta$ , and  $\eta_{\#} \mathbb{Q}^n$  is the distribution of the statistics from the true data-generating process. Then, the model-summulator pair is misspecified if,  $\forall \theta \in \Theta$ ,  $\eta_{\#} \mathbb{Q}^n \neq \eta_{\#} \mathbb{P}_{\theta}^{n}$ .
It is trivial to see that if the model is well-specified in the data space  $\mathcal{X}$ , it will also be well-specified in the statistic space  $S$  for any choice of  $\eta$ . The converse, however, is not true, as previously mentioned. Using this definition, we quantify the level of misspecification of a model-summurator pair as follows:
Definition 3.2 (Misspecification margin in SBI). Let  $\mathcal{D}$  be a distance defined on the set of all Borel probability measures on  $\mathcal{S}$ . Then, for a given  $\eta$ , we define the misspecification margin  $\varepsilon_{\eta}$  in SBI as
$$
\varepsilon_ {\eta} = \inf  _ {\theta \in \Theta} \mathcal {D} \left(\eta_ {\#} \mathbb {P} _ {\theta} ^ {n}, \eta_ {\#} \mathbb {Q} ^ {n}\right). \tag {3}
$$
Note that the margin is equal to zero for the well-specified case. Moreover, the larger the margin, the bigger the mismatch between the simulated and the observed statistics w.r.t  $\eta$ . Our aim is to learn an  $\eta$  such that  $\mathbb{P}_{\theta}$  is represented well whilst the model-summarizer pair is no longer misspecified as per Definition 3.1, or alternatively, has zero margin as per Definition 3.2. We formulate this as a constrained optimization problem in Equation (6).
Learning robust statistics for SBI. Suppose we have the flexibility to choose the summarizer  $\eta$  from a family of functions parameterized by  $\psi$ . Then, a potential approach for tackling model misspecification in SBI is to select the summarizer  $\eta$  that minimizes the misspecification margin  $\varepsilon_{\eta}$ . However, that involves computing the infimum for all possible choices of  $\psi$ . To avoid that, we minimize an upper bound on the margin instead, given as
$$
\varepsilon_ {\eta_ {\psi}} \leq \varepsilon_ {\eta_ {\psi}} ^ {\mathrm {u p p e r}} = \mathbb {E} _ {p (\theta)} \left[ \mathcal {D} \left(\eta_ {\psi \#} \mathbb {P} _ {\theta} ^ {n}, \eta_ {\psi \#} \mathbb {Q} ^ {n}\right) \right]. \tag {4}
$$
Note that Equation (4) is valid for any choice of  $p(\theta)$ , and we call  $\varepsilon_{\eta_{\psi}}^{\mathrm{upper}}$  the margin upper bound.
Minimizing the margin upper bound alone would lead to  $\eta$  being a constant function, which when used for inference, would yield the prior distribution. Hence, there is a trade-off between choosing an  $\eta$  that is informative about the parameters and an  $\eta$  that minimizes the margin upper bound. We
propose to navigate this trade-off by including the margin upper bound as a regularization term in the loss function for learning  $\eta$ . To that end, let  $\mathcal{L}(\omega, \psi)$  be the loss function used to learn  $\eta_{\psi}$  (along with additional parameters  $\omega$ ) which, for instance, can be either  $\mathcal{L}_{\mathrm{NPE}}$  from Equation (1) or  $\mathcal{L}_{\mathrm{AE}}$  from Equation (2). Then, our proposed loss with robust statistics (RS) is written as
$$
\mathcal {L} _ {\mathrm {R S}} (\omega , \psi) = \mathcal {L} (\omega , \psi) + \lambda \underbrace {\mathbb {E} _ {p (\theta)} \left[ \mathcal {D} \left(\eta_ {\psi \#} \mathbb {P} _ {\theta} ^ {n} , \eta_ {\psi \#} \mathbb {Q} ^ {n}\right) \right]} _ {\text {r e g u l a r i z a t i o n}}, \tag {5}
$$
where  $\lambda \geq 0$  is the weight we place on the regularization term relative to the standard loss  $\mathcal{L}$ . Minimizing the loss function in Equation (5) corresponds to solving the Lagrangian relaxation of the optimization problem:
$$
\min  _ {\omega , \psi} \quad \mathcal {L} (\omega , \psi) \tag {6}
$$
$$
\begin{array}{l l}\text {s . t .}&\mathbb {E} _ {p (\theta)} \left[ \right. \mathcal {D} \left( \right.\eta_ {\psi \#} \mathbb {P} _ {\theta} ^ {n}, \eta_ {\psi \#} \left(\mathbb {Q} ^ {n}\right)\left. \right] \leq \xi , \quad \xi > 0,\end{array}
$$
with  $\lambda \geq 0$  being the Lagrangian multiplier fixed to a constant value.
Estimating the margin upper bound. In order to implement the proposed loss from Equation (5), we need to estimate the margin upper bound using the training dataset  $\{(\theta_i, \mathbf{x}_{1:n,i})\}_{i=1}^m \sim p(\theta, \mathbf{x}_{1:n})$  sampled from the prior and the model. However, we only have one sample each from the distributions  $\eta_{\psi\#} \mathbb{P}_{\theta_i}^n$ ,  $i = 1, \dots, m$ , and one sample of the observed statistic from  $\eta_{\psi\#} \mathbb{Q}^n$ . Taking  $\mathcal{D}$  to be the Euclidean distance  $\|\cdot\|$ , we can estimate the margin upper bound as  $\frac{1}{m} \sum_{i=1}^{m} \|\eta_{\psi}(\mathbf{x}_{1:n,i}) - \eta_{\psi}(\mathbf{y}_{1:n})\|$ . Although easy to compute, we found this choice of  $\mathcal{D}$  to yield overly conservative posteriors even in the well-specified case, see Appendix B.3 for the results. This is because the Euclidean distance is large even if only a handful of simulated statistics — corresponding to different parameter values in the training data — are far from the observed statistic. As a result, the regularizer dominates the loss function while the standard loss term  $\mathcal{L}$  related to  $\theta$  is not minimized, leading to underconfident posteriors. Hence, we need a distance  $\mathcal{D}$  that is robust to outliers and can be computed between the set  $\{\eta_{\psi}(\mathbf{x}_{1:n,i})\}_{i=1}^m$  and  $\eta_{\psi}(\mathbf{y}_{1:n})$ , instead of computing the distance point-wise.
To that end, we take  $\mathcal{D}$  to be the maximum mean discrepancy (MMD) [42], which is a notion of distance between probability distributions or datasets. MMD has a number of attractive properties that make it suitable for our method: (i) it can be estimated efficiently using samples [15], (ii) it is robust against a few outliers (unlike the KL divergence), and (iii) it can be computed for unequal sizes of datasets. The MMD has been used in a number of frameworks such as ABC [8, 10, 50, 54, 62], minimum distance estimation [2, 15, 19, 56], generalized Bayesian inference [20, 57], and Bayesian nonparametric learning [24]. Similar to our method, the MMD has previously been used as a regularization term to train Bayesian neural networks [66], and in the NPE framework to detect model misspecification [73]. While we include the mmD-based regularizer to ensure that the observed statistic is not an out-of-distribution sample in the summary space, the regularizer in [73] involves computing the MMD between the simulated statistics and samples from a standard Gaussian, thus ensuring that the learned statistics are jointly Gaussian. They then conduct a goodness-of-fit test [42] to detect if the model is misspecified. Their method is complementary to ours, such that our method can be used once misspecification has been detected using [73].