SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Maximum mean discrepancy (MMD) as  $\mathcal{D}$ . The MMD is a kernel-based distance between probability distributions, computed by mapping the distributions to a function space called the reproducing kernel Hilbert space (RKHS). Let  $\mathcal{H}_k$  be the RKHS associated with the symmetric and positive definite function  $k: S \times S \to \mathbb{R}$ , called a reproducing kernel [6], defined on the space of statistics  $S$  such that  $f(\mathbf{s}) = \langle f, k(\cdot, \mathbf{s}) \rangle_{\mathcal{H}_k} \forall f \in \mathcal{H}_k$  (called the reproducing property). We denote the norm and inner product of  $\mathcal{H}_k$  by  $\| \cdot \|_{\mathcal{H}_k}$  and  $\langle \cdot, \cdot \rangle_{\mathcal{H}_k}$ , respectively. Any distribution  $\mathbb{P}$  can be mapped to  $\mathcal{H}_k$  via its kernel-mean embedding  $\mu_{k,\mathbb{P}} = \int_S k(\cdot, \mathbf{s})\mathbb{P}(\mathrm{d}s)$ . The MMD between two arbitrary probability measures  $\mathbb{P}$  and  $\mathbb{Q}$  on  $S$  is defined as the distance between their embeddings in  $\mathcal{H}_k$ , i.e.,  $\mathrm{MMD}_k(\mathbb{P}, \mathbb{Q}) = \| \mu_{k,\mathbb{P}} - \mu_{k,\mathbb{Q}} \|_{\mathcal{H}_k}$  [55]. Using the reproducing property, we can express the squared-MMD as
$$
\mathrm {M M D} _ {k} ^ {2} [ \mathbb {P}, \mathbb {Q} ] = \mathbb {E} _ {\mathbf {s}, \mathbf {s} ^ {\prime} \sim \mathbb {P}} [ k (\mathbf {s}, \mathbf {s} ^ {\prime}) ] - 2 \mathbb {E} _ {\mathbf {s} \sim \mathbb {P}, \mathbf {s} ^ {\prime} \sim \mathbb {Q}} [ k (\mathbf {s}, \mathbf {s} ^ {\prime}) ] + \mathbb {E} _ {\mathbf {s}, \mathbf {s} ^ {\prime} \sim \mathbb {Q}} [ k (\mathbf {s}, \mathbf {s} ^ {\prime}) ], \tag {7}
$$
which is computationally convenient as the expectations of the kernel can be estimated using iid samples from  $\mathbb{P}$  and  $\mathbb{Q}$ . Given samples of simulated statistics  $\{\eta_{\psi}(\mathbf{x}_{1:n,i})\}_{i = 1}^{l}$  and the observed
statistic  $\eta_{\psi}(\mathbf{y}_{1:n})$  , we estimate the squared-MMD between them using the V-statistic estimator [42]:
$$
\begin{array}{l} \mathrm {M M D} _ {k} ^ {2} \left[ \eta_ {\psi \#} \mathbb {P} _ {\theta} ^ {n}, \eta_ {\psi \#} \mathbb {Q} ^ {n} \right] \approx \mathrm {M M D} _ {k} ^ {2} \left[ \left\{\eta_ {\psi} \left(\mathbf {x} _ {1: n, i}\right) \right\} _ {i = 1} ^ {l}, \eta_ {\psi} \left(\mathbf {y} _ {1: n}\right) \right] \tag {8} \\ = \frac {1}{l ^ {2}} \sum_ {i, j = 1} ^ {l} k \left(\eta_ {\psi} \left(\mathbf {x} _ {1: n, i}\right), \eta_ {\psi} \left(\mathbf {x} _ {1: n, j}\right)\right) - \frac {2}{l} \sum_ {i = 1} ^ {l} k \left(\eta_ {\psi} \left(\mathbf {x} _ {1: n, i}\right), \eta_ {\psi} \left(\mathbf {y} _ {1: n}\right)\right). \\ \end{array}
$$
Note that the last term in Equation (7) is always constant for the estimator above as we only have one data-point of the observed statistic, hence we disregard it. Equation (8) therefore corresponds to estimating the MMD between the distribution of the simulated statistics for a given  $\psi$  and a Dirac measure on  $\eta_{\psi}(\mathbf{y}_{1:n})$ . The computational cost of estimating the squared-MMD is  $\mathcal{O}(l^2)$  and its rate of convergence is  $\mathcal{O}(l^{-\frac{1}{2}})$ . The NPE loss with robust statistics (NPE-RS) can then be estimated as
$$
\hat {\mathcal {L}} _ {\mathrm {N P E - R S}} (\phi , \psi) = - \frac {1}{m} \sum_ {i = 1} ^ {m} \log q _ {h _ {\phi} \left(\eta_ {\psi} \left(\mathbf {x} _ {1: n, i}\right)\right)} \left(\theta_ {i}\right) + \lambda \mathrm {M M D} _ {k} ^ {2} \left[ \left\{\eta_ {\psi} \left(\mathbf {x} _ {1: n, i}\right) \right\} _ {i = 1} ^ {l}, \eta_ {\psi} \left(\mathbf {y} _ {1: n}\right) \right]. \tag {9}
$$
Similarly, the autoencoder loss with robust statistics (AE-RS) reads
$$
\hat {\mathcal {L}} _ {\mathrm {A E} - \mathrm {R S}} (\psi , \psi_ {d}) = \frac {1}{m} \sum_ {i = 1} ^ {m} \left(\mathbf {x} _ {1: n, i} - \tilde {\eta} _ {\psi_ {d}} \left(\eta_ {\psi} \left(\mathbf {x} _ {1: n, i}\right)\right)\right) ^ {2} + \lambda \mathrm {M M D} _ {k} ^ {2} \left[ \left\{\eta_ {\psi} \left(\mathbf {x} _ {1: n, i}\right) \right\} _ {i = 1} ^ {l}, \eta_ {\psi} \left(\mathbf {y} _ {1: n}\right) \right]. \tag {10}
$$
We use a subset of the training dataset of size  $l < m$  to compute the MMD instead of  $m$  to avoid incurring additional computational cost in case  $m$  is large. The MMD-based regularizer can also be used as a score in a classification task to detect model misspecification, as shown in Appendix A.
Role of the regularizer  $\lambda$ . The regularizer  $\lambda$  penalizes learning those statistics for which the observed statistic  $\eta_{\psi}(\mathbf{y}_{1:n})$  is far from the set of statistics that the model can simulate given a prior distribution. When  $\lambda$  tends to zero, maximization of the likelihood dictates learning of both  $\phi$  and  $\psi$  in Equation (9), and our method converges to the NPE method. On the other hand, the regularization term is minimized when the summary network outputs the same statistics for both simulated and observed data, i.e., when  $\mathcal{D}$  is zero. In this case, the inference network can only rely on information from the prior  $p(\theta)$ . As a result, for large values of  $\lambda$ , we expect the regularization term to dominate the loss and the resulting posterior to converge to the prior distribution. Similar argument holds for the autoencoder loss in Equation (10). We empirically observe this behavior in Section 4. Hence,  $\lambda$  encodes the trade-off between efficiency and robustness of our inference method. Choosing  $\lambda$  can be cast as a hyperparameter selection problem, for which we can leverage additional data (if available) as a validation dataset, or use post-hoc qualitative analysis of the posterior predictive distribution.
We apply our method of learning robust statistics to two different SBI frameworks — NPE [41] and ABC [5] (see Appendix B.6 for results on applying our method to neural likelihood estimator [52]). We also compare the performance of our method against RNPE [81], by using the output of the trained summary network in NPE as statistics for RNPE. Experiments are conducted on synthetic data from two time-series models, namely the Ricker model from population ecology [67, 85] and the Ornstein-Uhlenbeck process [18]. Real data experiment on a radio propagation model is presented in Section 5. Analysis of the computational cost of our method and results on a 10-dimensional Gaussian model is presented in Appendix B.4 and B.5, respectively. The code to reproduce our experiments is available at https://github.com/huangdaolang/robust-sbi.
Ricker model simulates the evolution of population size  $N_{t}$  over the course of time  $t$  as  $N_{t + 1} = \exp (\theta_1)N_t\exp (-N_t + e_t), t = 1,\dots ,T,$  where  $\exp (\theta_1)$  is the growth rate parameter,  $e_t$  are zero-mean iid Gaussian noise terms with variance  $\sigma_e^2$ , and  $N_0 = 1$ . The observations  $x_{t}$  are assumed to be Poisson random variables such that  $x_{t}\sim \mathrm{Poisss}(\theta_{2}N_{t})$ . For simplicity, we fix  $\sigma_e^2 = 0.09$  and estimate  $\theta = [\theta_1,\theta_2]^\top$  using the prior distribution  $\mathcal{U}([2,8]\times [0,20])$ , and  $T = 100$  time-steps.
Ornstein-Uhlenbeck process (OUP) is a stochastic differential equation model widely used in financial mathematics and evolutionary biology. The OU process  $x_{t}$  is defined as:
$$
x _ {t + 1} = x _ {t} + \Delta x _ {t}, \quad t = 1, \dots , T,
$$
$$
\Delta x _ {t} = \theta_ {1} [ \exp (\theta_ {2}) - x _ {t} ] \Delta t + 0. 5 w,
$$
Figure 2: Performance of the SBI methods in terms of RMSE and MMD for both the Ricker model and OUP. Each box represents the median and interquartile range (IQR), while the whiskers extend to the furthest points within 1.5 times the IQR from the edges of the box. For the well-specified case  $(\epsilon = 0\%)$ , the proposed NPE-RS and ABC-RS methods perform similar to their counterpart NPE and ABC, respectively. Under misspecification  $(\epsilon > 0\%)$ , NPE-RS and ABC-RS achieve lower RMSE and MMD values, demonstrating robustness to model misspecification.
where  $T = 25$ ,  $\Delta t = 0.2$ ,  $x_0 = 10$ , and  $w \sim \mathcal{N}(0, \Delta t)$ . A uniform prior  $\mathcal{U}([0, 2] \times [-2, 2])$  is placed on the parameters  $\theta = [\theta_1, \theta_2]^\top$ .
Contamination model. We test for robustness relative to the  $\epsilon$ -contamination model [65], similar to [24]. For both the Ricker model and the OUP, the observed data comes from the distribution  $\mathbb{Q} = (1 - \epsilon)\mathbb{P}_{\theta_{\mathrm{true}}} + \epsilon \mathbb{P}_{\theta_c}$ , where a large proportion of the data comes from the model with  $\theta_{\mathrm{true}}$  while  $\epsilon$  proportion of the data comes from the distribution  $\mathbb{P}_{\theta_c}$ . Hence,  $\epsilon \in [0,1]$  denotes the level of misspecification within both the models, with  $\epsilon = 0$  resulting in the well-specified case. We take  $\theta_{\mathrm{true}} = [0.5,1.0]^{\top}$ ,  $\theta_c = [-0.5,1.0]^{\top}$  for OUP, and  $\theta_{\mathrm{true}} = [4,10]^{\top}$ ,  $\theta_c = [4,100]^{\top}$  for the Ricker model.
Implementation details. We take  $m = 1000$  samples for the training data and  $n = 100$  realizations of both the observed and simulated data for each  $\theta$ . We set  $\lambda$  using an additional dataset simulated from the models using  $\theta_{\mathrm{true}}$ . For the MMD, we take the kernel to be the exponentiated-quadratic kernel  $k(\mathbf{s},\mathbf{s}^{\prime}) = \exp (-\| \mathbf{s} - \mathbf{s}^{\prime}\|_{2}^{2} / \beta^{2})$ , and set its lengthscale  $\beta$  using the median heuristic  $\beta = \sqrt{\mathrm{med}/2}$  [42], where med denotes the median of the set of squared two-norm distances  $\| \mathbf{s}_i - \mathbf{s}_j\| _2^2$  for all pairs of distinct data points in  $\{\mathbf{s}_i\}_{i = 1}^m$ . We use  $l = 200$  samples of the simulated statistics to estimate the MMD. For the Ricker model, the summary network  $\eta_{\psi}$  is composed of 1D convolutional layers, whereas for the OUP,  $\eta_{\psi}$  is a combination of bidirectional long short-term memory (LSTM) recurrent modules and 1D convolutional layers. The dimension of the statistic space is set to four for both the models. We take  $q$  to be a conditional normalizing flow for all the three NPE methods. We take the tolerance  $\delta$  in ABC to be the top  $5\%$  of samples that yield the smallest distance.
Performance metrics. We evaluate the accuracy of the posterior, as well as the posterior predictive distribution of the SBI methods. For the posterior distribution, we compute the root mean squared error (RMSE) as  $(1 / N\sum_{i = 1}^{N}(\theta_{i} - \theta_{\mathrm{true}})^{2})^{\frac{1}{2}}$  where  $\{\theta_i\}_{i = 1}^N$  are posterior samples. For the predictive accuracy, we compute the MMD between the observed data and samples from the posterior predictive distribution of each method. As the models simulate high-dimensional data, we use the kernel specialized for time-series from [8]. The lengthscale of this kernel is set to  $\beta = 1$  for all misspecification levels to facilitate fair comparison.
Results. Figure 2 presents the results for both the Ricker model and the OUP across 100 runs. We observe that both NPE and ABC with our robust statistics (RS) outperform their counterparts, including RNPE, under misspecification ( $\epsilon = 10\%$  and  $20\%$ ). Moreover, our performance is similar
(a) Well-specified  $(\epsilon = 0)$
(b) Misspecified  $(\epsilon = 10\%)$
(c) Misspecified  $(\epsilon = 20\%)$
Figure 3: Posteriors obtained from NPE, RNPE, and our NPE-RS method for the Ricker model. We perform similar to NPE in the well-specified case, unlike RNPE. Under misspecification, NPE and RNPE posteriors drift away from  $\theta_{\mathrm{true}}$ , going even beyond the prior range (denoted by dashed gray lines) in the case of NPE. Our method is robust to model misspecification.
to NPE in the well-specified case ( $\epsilon = 0\%$ ). An instance of the NPE posteriors is shown in Figure 3 for the Ricker model. Our NPE-RS posterior is very similar to the NPE posterior for the well-specified case, whereas the RNPE posterior is underconfident, as noted in [81]. Although RNPE is more robust than NPE under misspecification, its posterior still drifts away from  $\theta_{\mathrm{true}}$ , while NPE-RS posterior stays around  $\theta_{\mathrm{true}}$  even for  $\epsilon = 20\%$ . This is because our method has the flexibility to choose appropriate statistics based on misspecification level, while RNPE is bound to a fixed choice of pre-selected statistics. Posterior plots for ABC and OUP are given in Appendix B.2.
Prior misspecification. So far our discussion has pertained to the case of likelihood misspecification. However, in Bayesian inference, the prior is also a part of the model. Hence, a potential form of misspecification is prior misspecification, in which the prior gives low or even zero probability to the "true" data generating parameters. We investigate the issue of prior misspecification using the Ricker model as an example. To that end, we modify the prior distribution of  $\theta_{2}$  by setting it to  $\mathrm{LogNormal}(0.5,1.0)$  while keeping the prior of  $\theta_{1}$  unchanged. We use  $\theta_{\mathrm{true}} = [4,25]^{\top}$  to generate the observed data, as  $\theta_{2} = 25$  has very low density under the lognormal prior, and the probability that a value of  $\theta_{2} \geq 25$  will be sampled is 0.00372. The result in Figure 4 shows that NPE and RNPE yield posteriors that do not include the true value of  $\theta_{2}$ , whereas our NPE-RS posterior is still around  $\theta_{\mathrm{true}}$ . This highlights the effectiveness of our method to handle cases of prior misspecification as well.
Figure 4: Posteriors from NPE, RNPE, and NPE-RS for Ricker model. NPE-RS is robust to prior misspecification.
Varying regularizer  $\lambda$ . As mentioned in Section 3, our method converges to NPE as  $\lambda$  tends to zero, and to the prior distribution as  $\lambda$  becomes high. To investigate this effect, we vary  $\lambda$  and measure the distance between the NPE-RS and the NPE posteriors, and between the NPE-RS posterior and the prior. We use the MMD from Equation (7) to measure this distance on the Ricker model in the misspecified case ( $\epsilon = 20\%$ ). We see in Figure 5(left) that the MMD values between the prior and the NPE-RS posteriors go close to zero for  $\lambda = 10^3$ , indicating that our method essentially returns the prior for high values of  $\lambda$ . On the other hand, the MMD values between NPE and NPE-RS posteriors in Figure 5(middle) are smallest for  $\lambda = 0.01$  and  $\lambda = 0.1$ . For non-extreme values of  $\lambda$  (i.e.  $\lambda = 1$  or 10), we observe maximum difference between NPE-RS and NPE. This is because the NPE posteriors tend to move out of the prior support under misspecification, whereas the NPE-RS posteriors remain around  $\theta_{\mathrm{true}}$ , as shown in Figure 3. Finally, the NPE-RS posteriors are also close to the NPE posteriors for small values of  $\lambda$  in the well-specified case, as shown in Figure 5(right).
Figure 5: MMD between the NPE-RS posteriors and (left) the prior in the misspecified setting, (middle) the NPE posteriors in the misspecified setting, (right) the NPE posteriors in the well-specified setting, for different values of  $\lambda$ .
5 Application to real data: A radio propagation example
Our final experiment involves a stochastic radio channel model, namely the Turin model [45, 63, 78], which is used to simulate radio propagation phenomena in order to test and design wireless communication systems. The model is driven by an underlying point process which is unobservable, thereby rendering its likelihood function intractable. The model simulates high-dimensional complex-valued time-series data, and has four parameters. The parameters govern the starting point, the slope of the decay, the rate of the point-process, and the noise floor of the time-series, as shown in Figure 6.
We attempt to fit this model to a real dataset from a measurement campaign [44] for which the model is known to be misspecified [8], on account of the data samples being non-iid. The averaged power of the data (square of the absolute value) as a function of time is shown by the black curve in Figure 6. The data dimension is 801 and we have  $n = 100$  realizations. We take the power of the complex data in decibels as input to the summary network, which consists of 1D convolutional layers, and we set uniform priors for all the four parameters. Descriptions of the model, the data, and the prior ranges are provided in Appendix C.
We fit the model to the real data using NPE, RNPE and our NPE-RS methods. As there is no notion of ground truth in this case, we plot the resulting posterior predictive distributions in Figure 6. Note that the multiple peaks present in the data are not replicated by the Turin model
Figure 6: Predictive performance of NPE, RNPE and our NPE-RS method on the Turin model under misspecification. NPE-RS leads to the best fit. The shaded regions denote  $\pm 1$  std. deviation of the posterior predictive distribution of each method.
for any method, which is due to the model being misspecified for this dataset. Despite that, our NPE-RS method appears to fit the data well, while NPE performs the worst. Moreover, our method is better than RNPE at matching the starting point of the data, the slope, and the noise floor, with RNPE underestimating all three aspects on average. The MMD between the observed data and the posterior predictive distribution of NPE, RNPE and NPE-RS is 0.11, 0.09, and 0.03, respectively. Hence, NPE-RS provides a reasonable fit of the model even under misspecification.
We proposed a simple and elegant solution for tackling misspecification of simulator-based models. As our method relies on learning robust statistics and not on the subsequent inference procedure, it is applicable to any SBI method that utilizes summary statistics. Apart from achieving robust inference under misspecified scenarios, the method performs reasonably well even when the model is well-specified. The proposed method only has one hyperparameter that encodes the trade-off between efficiency and robustness, which can be selected like other neural network hyperparameters, for instance, via a validation set.
*In wireless communications, the power of the signal is measured in decibels (dB). A power  $P$  in linear scale corresponds to  $10\log_{10}(P)$  dB.
Limitations and future work. A limitation of our method is the increased computational complexity due to the cost of estimating the MMD, which can be alleviated using the sample-efficient MMD estimator from [10] or quasi-Monte Carlo points [56]. Moreover, as our method utilizes the observed statistic during the training procedure, the corresponding NPE is not amortized anymore — a limitation we share with RNPE. Thus, working on robust NPE methods which are still amortized (to some extent) is an interesting direction for future research. An obvious extension of the work could be to investigate if the ideas translate to likelihood-free model selection methods [68, 76] which can suffer from similar problems as SBI if all (or some) of the candidate models are misspecified.
The authors would like to thank Dr. Carl Gustafson and Prof. Fredrik Tufvesson (Lund University) for providing the measurement data. We also thank Masha Naslidnyk, Dr. François-Xavier Briol, and Dr. Markus Heinonen for their useful comments and discussion. We acknowledge the computational resources provided by the Aalto Science-IT Project from Computer Science IT. This work was supported by the Academy of Finland Flagship programme: Finnish Center for Artificial Intelligence FCAI. SK was supported by the UKRI Turing AI World-Leading Researcher Fellowship, [EP/W002973/1].
[1] Albert, C., Ulzega, S., Ozdemir, F., Perez-Cruz, F., and Mira, A. (2022). Learning Summary Statistics for Bayesian Inference with Autoencoders. SciPost Phys. Core, 5:043.
[2] Alquier, P. and Gerber, M. (2021). Universal robust regression via maximum mean discrepancy. arXiv:2006.00840.
[3] Beaumont, M. A. (2010). Approximate Bayesian computation in evolution and ecology. Annual Review of Ecology, Evolution, and Systematics, 41(1):379-406.
[4] Beaumont, M. A. (2019). Approximate Bayesian computation. Annual Review of Statistics and Its Application, 6(1):379-403.
[5] Beaumont, M. A., Zhang, W., and Balding, D. J. (2002). Approximate Bayesian computation in population genetics. Genetics, 162(4):2025-2035.
[6] Berlinet, A. and Thomas-Agnan, C. (2004). Reproducing Kernel Hilbert Spaces in Probability and Statistics. Springer Science+Business Media, New York.
[7] Bernardo, S. (2000). Bayesian Theory. John Wiley & Sons.
[8] Bharti, A., Briol, F.-X., and Pedersen, T. (2022a). A general method for calibrating stochastic radio channel models with kernels. IEEE Transactions on Antennas and Propagation, 70(6):3986-4001.
[9] Bharti, A., Filstroff, L., and Kaski, S. (2022b). Approximate Bayesian computation with domain expert in the loop. In International Conference on Machine Learning, volume 162, pages 1893-1905.
[10] Bharti, A., Naslidnyk, M., Key, O., Kaski, S., and Briol, F.-X. (2023). Optimally-weighted estimators of the maximum mean discrepancy for likelihood-free inference. arXiv:2301.11674.
[11] Bi, J., Shen, W., and Zhu, W. (2021). Random forest adjustment for approximate Bayesian computation. Journal of Computational and Graphical Statistics, 0(0):1-10.
[12] Bissiri, P. G., Holmes, C. C., and Walker, S. G. (2016). A general framework for updating belief distributions. Journal of the Royal Statistical Society. Series B (Statistical Methodology), 78(5):1103-1130.
[13] Blum, M. G. B. and François, O. (2010). Non-linear regression models for approximate Bayesian computation. Statistics and Computing, 20(1):63-73.