SlowGuess/NeurIPS_2023_OCR huggingface dataset:
For any  $c$ ,  $\mathrm{LPD}^{\mathrm{val}}(\Phi, c) - \sum_{k=0}^{K-1} w_k \log w_k$  becomes a lower bound estimate of the (generalized) divergence  $D^{\mathrm{opt}}(p, q, \Phi)$  in Thm. 1, and an estimate of  $D^{\mathrm{opt}}(p, q, \Phi)$  itself when the classier  $c$  is good enough. In addition to the point estimate of the divergence, we can compute the confidence interval. It is straightforward to obtain the standard error of the sample mean (9). To take into account the potentially heavy tail of the log predictive densities, we can also use Bayesian bootstrap [37] that reweights the sum in (9) by a uniform Dirichlet weight.
Hypothesis testing. Our approach facilitates rigorous frequentist hypothesis testing. The null hypothesis is that the approximate inference matches the exact posterior, i.e.,  $p(\theta | y) = q(\theta | y)$  almost everywhere. We adopt the permutation test: We train the classifier  $c$  once on the training set and keep it fixed, and evaluate the validation set log predictive density  $\mathrm{LPD}^{\mathrm{val}}(\Phi, c)$  in (9). Next, permute the validation set  $B$  times: at time  $b$ , keep the features unchanged and randomly permute the validation labels  $t$  within each batch of examples ( $\Phi$  generates a batch of  $L$  examples each time), and reevaluate the validation set log predictive density (9) on permuted labels, call it  $\mathrm{LPD}_b^{\mathrm{val}}$ . Then we compute the one-sided permutation  $p$ -value as  $p = \sum_{b=1}^{B} \mathbb{1}(\mathrm{LPD}_b^{\mathrm{val}} \geq \mathrm{LPD}^{\mathrm{val}}) / B$ . Given a significance level, say 0.05, we will reject the null if  $p < 0.05$  and conclude a miscalibration.
Theorem 4 (Finite sample frequentist test). For any finite simulation size  $S$  and posterior draw size  $M$ , and any classifier  $c$ , under the null hypothesis  $p(\theta | y) = q(\theta | y)$  almost everywhere, the permutation test is valid as the  $p$ -value computed above is uniformly distributed on  $[0,1]$ .
Our test is exact when the simulation size  $S$  and  $M$  is finite, while the original SBC [4] relied on asymptotic approximation. Further, we learn the test statistic via the classifier, and our test is valid regardless of the dimension of  $\theta$  and there is no need to worry about post-learning testing [3], while the usual SBC rank test will suffer from low power due to multiple testing.
Our test is always valid even if the classifier  $c$  is not optimal. Why does our learning step help? For notional brevity, here we only reason for the binary classification. For any  $p, q$ , we apply the binary classification as described in Example 1,  $t = 0$  or 1 and  $\phi = (\theta, y)$ .
Theorem 5 (Sufficiency). Let  $\hat{c}(\theta, y) = \operatorname*{Pr}(t = 1|\theta, y)$  be the probability of label 1 in the optimal classifier as per (7), and let  $\pi_c^p$  and  $\pi_c^q$  be the one-dimensional distributions of this  $\hat{c}(\theta, y)$  when  $(\theta, y)$  is sampled from  $p(\theta, y)$  or from  $p(y)q(\theta|y)$  respectively, then (i) Conditional on the summary statistic  $\hat{c}$ , the label  $t$  is independent of features  $\phi = (\theta, y)$ . (ii) Under regularity conditions, there is no loss of information in divergence as the joint divergence is the same as the projected divergence in the one-dimensional  $\hat{c}$ -space  $D_1(p, q) = D_1(\pi_c^p, \pi_c^q)$ .
That is, the best prediction  $\hat{c}$  entails the best one-dimensional summary statistics of the high dimensional  $\theta \times y$  space. The enhancement of the test power from using the sufficient statistics is then assured by the Neyman-Pearson lemma [27].
4.1 Feature engineering: use rank, log density and symmetry
Reassuringly, whatever the classifier  $c$ , the prediction ability  $D(p,q,\Phi ,c)$  is always a lower bound to the corresponding divergence  $D^{\mathrm{opt}}(p,q,\Phi)$ , and the permutation test is always exact. But we still wish to train a "good" classifier in terms of its out-of-sample performance, for a tighter bound in divergence, and a higher power in testing. We will typically use a flexible parametric family such as a multilayer perceptron (MLP) network to train the classifier.
The oracle optimal probabilistic classifier is the true label probability, and in the binary and multiclass classification (Example 1 and 4), the oracle has closed-form expressions, although we cannot evaluate:
$$
\Pr_ {\text {b i n a r y}} (t = 0 | \theta , y) = \frac {p (\theta | y)}{p (\theta | y) + q (\theta | y)}, \quad \Pr_ {\text {m u l t i}} (t | \theta_ {0}, \dots , \theta_ {M}, y) = \frac {p \left(\theta_ {t} , y\right) / q \left(\theta_ {t} | y\right)}{\sum_ {k = 0} ^ {M} p \left(\theta_ {k} , y\right) / q \left(\theta_ {k} | y\right)}. \tag {10}
$$
Statistically-meaningful feature. Depending on the inference task, we have more information than just the sample points and should use them in the classifier. In light of the shape and component of the optimal classifiers (10), the following statistically-meaningful features are useful whenever available: (i) The log target density  $\log p(\theta |y)$ . As proxies, the log joint density  $\log p(\theta ,y)$  and the log likelihood  $\log p(y|\theta)$  are often known to traditional Bayesian models. (ii) The log approximate density  $\log q(\theta |y)$ , known to volitional and normalizing flows. (iii) When the log approximate density is unknown, a transformation is to integrate the density and obtain the posterior CDF,  $Q(\theta |y) = \mathbb{E}_{x\sim q(x|y)}\mathbb{1}(x > \theta)$ , and this CDF can be approximated by the rank statistics among the approximate draws up to a rescaling  $r(\theta ,y)\coloneqq \sum_{m = 1,i = 1}^{MS}\mathbb{1}(y_i = y)\mathbb{1}(\theta < \tilde{\theta}_{im})$ . See Appendix Table 1 for the usage and extension of these features in binary and multiclass classifiers.
Linear features. We call the log likelihood, the log approximate density, or log prior (whenever available) linear features, and denote them to be  $l(\theta, y)$ . For example if both likelihood and  $q$  density
is known, then  $l(\theta, y) = (\log p(y|\theta), \log q(\theta|y))$ . Because they appear in the oracle 10, we will keep linear features in the last layer of the network (followed by a softmax). We recommend parameterizing the binary classification in the following form:
$$
\Pr (t = 1 | (\theta , y)) = \operatorname {i n v} _ {-} \operatorname {l o g i t} [ \mathrm {M L P} (\theta , y) + w ^ {T} l (\theta , y) ]. \tag {11}
$$
Symmetry in multiclass classifier. The multi-class classifier is harder to train. Luckily, we can use the symmetry of the oracle classifier (10): the probability of class  $k$  is proportional to a function of  $(\theta_k, y)$  only, hence we recommend parameterizing the multiclass probability as
$$
\Pr (t = k \mid (\theta_ {0}, \theta_ {1}, \dots , \theta_ {M}, y)) = \frac {\exp (g (\theta_ {k} , y))}{\sum_ {k ^ {\prime} = 0} ^ {M} \exp (g (\theta_ {k ^ {\prime}} , y))}, g (\theta , y) = \operatorname {M L P} (\theta , y) + w ^ {T} l (\theta , y), \tag {12}
$$
where  $l(\theta, y)$  is available linear features. We only need to learn a reduced function from  $\theta \times y$  to  $\mathbb{R}$ , instead of from  $\theta^{M + 1} \times y$  to  $\mathbb{R}$ , reducing the complexity while still keeping the oracle (10) attainable.
Zero waste calibration for MCMC. If  $\tilde{\theta}_1, \ldots, \tilde{\theta}_M \sim q(\theta | y)$  are produced by MCMC sampler, typically  $\tilde{\theta}$  has autocorrelations. Although classical SBC originated from MCMC application, the rank test requires independent samples, so SBC can only handle autocorrelation by thinning: to subsample  $\tilde{\theta}_m$  and hope the thinned  $\tilde{\theta}$  draws are independent. Thinning is a waste of draws, inefficient when the simulations are expensive, and the thinned samples are never truly independent. With MCMC draws  $\tilde{\theta}$ , our method could adopt thinning as our Thm. 1 and 4 are valid even when  $M = 1$ .
Yet we can do better by using all draws. The "separable" network architecture (12) is ready to use for MCMC samples. For example, we sample  $(\theta ,y)\sim p(\theta ,y)$ , and sample  $(\tilde{\theta}_1,\dots ,\tilde{\theta}_M)$  from a MCMC sampler whose stationary distribution we believe is  $q(\theta |y)$ , and generate examples from the multiclass permutation (Example 4). Then we run a separable classifier(12) to predict  $t$ . Intuitively, the separable network design (12) avoids the interaction between  $\tilde{\theta}_m$  with  $\tilde{\theta}_{m'}$ , and disallows the network to predict  $t$  based on the autocorrelation or clustering of  $\tilde{\theta}$ . The next theorem states the validity of our method in MCMC settings without thinning.
Theorem 6 (MCMC). Suppose we sample  $(\theta, y) \sim p(\theta, y)$ , and sample  $(\tilde{\theta}_1, \dots, \tilde{\theta}_M)$  from a MCMC sampler whose stationary distribution we believe is  $q(\theta | y)$  (i.e., marginally  $\tilde{\theta}_i$  is from  $q(\theta | y)$ ), and generate examples  $((t_1, \phi_1), \dots, (t_{M+1}, \phi_{M+1}))$  from the multiclass permutation, such that  $\phi = (\theta_0, \theta_1, \dots, \theta_M)$ . Then we run an exchangeable classifier (12) in which  $g$  is any  $\Theta \times \mathcal{V} \to \mathbb{R}$  mapping. Denote  $D_4^{\mathrm{MCMC,sep}}(p, q)$  to be the predictive ability of the optimal classifier among all separable classifiers (12), then  $D_4^{\mathrm{MCMC,sep}}(p, q) = D_4(p, q)$ .
Dimension reduction and nuisance parameter. Sometimes we only care about the sampling quality of one or a few key dimensions of the parameter space, then we only need to restrict the classifier to use these targeted dimensions, as a result of Theorem 1. For example, in binary classification, if we reduce the feature  $\phi = (\theta ,y)$  to  $\phi = (h(\theta),y)$  in the classifier, where  $h(\theta)$  can be a subset of  $\theta$  dimensions, then the resulting classification divergence becomes projected divergence between  $h(\theta)|y,\theta \sim p(\theta |y)$  and  $h(\theta)|y,\theta \sim q(\theta |y)$ , and other nuisance parameters do not impact the diagnostics.
Weighing for imbalanced binary classification. When  $M$  is big, the binary classification (Example 1) can suffer from imbalanced labels, and the divergence in (1) degenerates:  $D_{1}(p,q)\to 0$  as  $M\rightarrow \infty$ . One solution is to use the multiclass classification which has a meaningful limit (Thm. 3). Another solution is to reweigh the loss function or log predictive density by the label  $t$ . If the weights of class 1 and class 0 examples are  $\frac{M + 1}{2M}$  and  $\frac{M + 1}{2}$ , used in both training and validation log prediction density, then regardless of  $M$ , the weighted classification is equivalent to balanced classification, and the resulting divergence is the symmetric Jensen-Shannon (JS) divergence  $\frac{1}{2}\mathrm{KL}[p(\theta |y)||r(\theta |y)] + \frac{1}{2}\mathrm{KL}[q(\theta |y)||r(\theta |y)]$ , where  $r(\theta |y) = \frac{1}{2} [p(\theta |y) + q(\theta |y)]$ . See Appendix A for proofs.
Closed-form example. Consider a multivariate normal parameter prior  $\theta \in \mathbb{R}^d\sim \mathrm{MVN}(\mathbf{0},\mathrm{Id}_d)$  and a normal data model  $y|\theta \sim \mathrm{MVN}(\theta ,\Sigma)$ , so the exact posterior  $p(\theta |y)$  is explicit. In the experiments,
Figure 2: Our test is valid as it yields uniformly distributed  $p$ -values under the null  $p(\theta | y) = q(\theta | y)$ . We check all dimensions at once.
Figure 3: Our test has a uniformly higher power than SBC rank test. We simulate wrong inference  $q$  by multiplying the posterior covariance or adding biases to the posterior mean.
Figure 4: In the Gaussian example with bias  $= 1$  or 2, with a moderately big simulation sample size  $S$ , the divergence estimated from a binary (left panel) or multiclass (right panel) classifier matches its theory quantity (dashed): the Jensen-Shannon or Kullback-Leibler divergence in the big- $M$  limit.
Figure 5: We apply binary classifier calibration to posterior inferences on three models from the SBI benchmark, and check the sampling quality after some iterations. The  $x$ -axis is the number of MCMC iterations. The  $y$ -axis is the estimated JS divergence (± standard error) between the true target and sampled distribution at a given number of iterations, indicating gradual convergence.
we use neural nets to train binary (11) or symmetric multiclass classifiers (12), which we generally recommend. We assume the inference posterior  $\log q(\theta |y)$  is known and added to the classifier.
Legitimacy of testing. First, to validate our hypothesis test under the null, we use true inference  $q(\theta | y) = p(\theta | y)$ . With  $d = 16$ , we simulate  $S = 500$  draws of  $(\theta, y) \sim p(\theta, y)$ , then generate true posterior  $\tilde{\theta}_m \sim p(\theta | y)$ , and run our binary-classifier and obtain a permutation  $p$ -value. We repeat this testing procedure 1000 times to obtain the distribution of the  $p$ -values under the null, which is uniform as shown in Fig. 2, in agreement with Thm. 4. Because  $\theta$  has  $d$  dimensions, a SBC rank test on each margin separately is invalid without adjustment, and would require Bonferroni corrections.
Power. We consider two sampling corruptions: (i) bias, where we add a scalar noise (0.01 to 0.2) to each dimension of the true posterior mean, (ii) variance, where we inflate the true posterior covariance matrix by multiplying a scalar factor (0.8 to 1.2). We compare our discriminative tests to a Bonferroni-corrected SBC chi-squared rank-test. In both settings, we fix a  $5\%$  type-I error and compute the power from 1000 repeated tests. Fig. 3 shows that our method has uniformly higher power than SBC, sometimes as good as SBC with a 10 times bigger simulation sample size.
Divergence estimate. The left panel of Fig. 4 validates Thm. 1: We run a weighted binary classifier on the Gaussian simulations with a scalar bias 1 or 2 added to the posterior mean in  $q$ . As the number of simulations  $S$  grows, the learned classifier quickly approaches the optimal, and the prediction ability matches the theory truth (dashed line): the Jensen-Shannon divergence between  $p(\theta | y)$  and  $q(\theta | y)$ . The right panel validates Thm. 3: With a fixed simulation size  $S$ , we increase  $M$ , the number of posterior draws from  $q$ , and run a multiclass classification. When  $M$  is big, the estimated divergence converges from below to the dashed line, which is the theory limit  $\mathrm{KL}(p(\theta | y), q(\theta | y))$  in Thm. 3.
Benchmark examples. Next, we apply our calibration to three models from the SBI benchmark [23]: the simple likelihood complex posterior (SLCP), the Gaussian linear, and the Gaussian mixture model. In each dataset, we run adaptive No-U-Turn sampler (NUTS) and check the quality of the sampled distribution after a fixed number of iterations, varying from 2 to 2000 (we use equal number of iterations for warm-up and for sampling, and the warm-up samples were thrown away). At each given MCMC iterations, we run our classifier calibration, and estimate the JS divergence, as reported
Figure 6: Scatter plots of the classifier prediction v.s. Figure 7: In the cosmology example we visually check a one-dimensional parameter we need to check can the classifier predicted log odds v.s. cosmology param-reveal marginal over- or under-confidence in inference  $\Omega_{b}$  and conc, and  $\log q(\theta |y)$ . Underconfidence is  $q(\theta |y)$  from a U or inverted-U shape. detected in conc marginally, and more so in the joint.
in Fig. 5. In all three panels, we are able to detect the inference flaws at early iterations and observe a gradual convergence to zero.
Visual check. Our diagnostic outputs rigorous numerical estimates, but it also facilities visual checks. We make a scatter plot of the binary classifier prediction  $\operatorname*{Pr}(t = 1|\phi)$ , a proxy of  $p(\theta |y) / q(\theta |y)$ , against any one-dimensional parameter we need to check: If that parameter is under- or over-confident, this scatter plot will display a U- or inverted-U-shaped trend. Compared with SBC rank histograms, our visualization can further check the magnitude of mismatch (how far away  $\operatorname*{Pr}(t = 1|\phi)$  is from 0.5), tail behavior (small  $q(\theta |y)$ ), and several dimensions jointly. Fig 6 is a visual check in the Gaussian example when we multiply the posterior covariance in inference  $q$  by 0.8, 1, or 1.2.
Galaxy clustering. We would like to model galaxy spatial clustering observations in a Lambda	model 1, dy=38	model 2, dy=151	model 3, dy=1031
1-flow	0.015 (0.004)	0.076 (0.007)
5-ens.	0.013 (0.004)	0.044 (0.006)
cold dark matter framework [34], where observations  $y$  correspond to statistical descriptors of galaxy clustering, and a 14-dimensional parameter  $\theta$  encodes key cosmological information about the Universe. The forward simulation  $y|\theta$  involves expensive  $N$ -body simulations (each single  $y$  simulation takes around 5000 cpu hours). We consider three different physical models  $p$ : they contain the same cosmology parameter  $\theta$  but three types of observations whose dimension  $d_y$  varies from 38 to 1031, reflecting three power-spectrum designs. We apply simulation based inference to sample from each of the models, where we try various normalizing flow architectures and return either the best architecture or the ensemble of five architectures. We now apply our discriminative calibration to these six approximate inferences using a weighted binary classifier trained from  $S = 2000$ ,  $M = 500$  simulations for each inference (that is 1 million examples). We have added the log densities  $\log q(\theta |y)$  as extra features in the classifier, as they are known in the normalizing-flows. The table above shows the estimated Jensen-Shannon distances and standard deviation. Compared this table with the estimates from blackbox MLP classifier (Appendix Fig. 8), using statistical features greatly improves the label prediction and tightens the divergence estimate. From the table above, all inferences  $q$  have a significantly non-zero but small divergence, among which the 5-ensemble that uses a mixture of 5 flows always has a lower divergence so is more trustworthy. Further visualization of the classification log odds (Fig. 7) reveals that  $q$  is under-confident in the parameter that encodes the galaxy concentration rate. To interpret the tail behavior, the predicted log odds are negatively correlated with log joint normalizing-flow density  $q(\theta |y)$ , suggesting  $q(\theta |y) > p(\theta |y)$  in the tail of  $q(\theta |y)$ , another evidence of  $q$  underconfidence.
We share Jax implementation of our binary and multiclass classifier calibration in  $\mathbf{G h i b u^{2}}$
This paper develops a classifier-based diagnostic tool for Bayesian computation, applicable to MCMC, variational and simulation based inference, or even their ensembles. We learn test statistics from data using classification. Through a flexible label mapping, our method yields a (family of) computable divergence metric and an always valid testing  $p$ -value, and the statistical power is typically higher than the traditional rank based SBC. Various statically-meaningful features are available depending on the task and we include them in the classifier. Visual exploration is also supported.
$^{2}$ https://github.com/yao-yl/DiscCalibration
Related diagnostics. The early idea of simulation based calibration dates back to [13] who compared the prior and the data-averaged posterior, i.e., to check  $p(\theta) = \int q(\theta |y)p(y)dy$  using simulations. [42] further developed techniques for the first and second moment estimate of the data-averaged posterior using the law of total variance. Our method includes such moment comparison as special cases: using the no-  $y$  binary labeling (Example 2), then comparing the empirical moments of the  $q$  sample the  $p$  sample can be achieved through a naive Bayes classifier or a linear discriminant analysis.
The rank-based SBC [4, 38] can be recovered by ours using binary labels and taking ranks to be features (Example 3). The rank statistic is central to SBC, which follows the classical frequentist approach—using the tail probability under the posited model quantifies the extremeness of the realized value[11] is only a convenient way to locate the observations in the reference distribution, especially in the past when high dimensional data analysis was not developed—it is the use of modern learning tools that sharpens our diagnostic. Recently, SBC has developed various heuristics for designing (one-dimensional) test statistics  $\phi(\theta, y)$ . Such rank tests are recovered by our method by including the rank of  $\phi$  in our features (Sec. 4.1). For example, [26] proposed to test the rank of the likelihood  $\phi = p(\theta | y)$  in MCMC, [20] looked at the rank of proposal density  $q(\theta | y)$ , [22] used the  $q$ -probability integral transformation in normalizing flows. In light of our Theorem 5, the optimal test statistic is related to the density ratio and hence problem-specific, which is why our method includes all known useful features in the classifier and learns the test statistic from data.
Classifier two-sample test. Using classification to compare two-sample closeness is not a new idea. The classifier two-sample test (C2ST) has inspired the generative adversarial network (GAN) [16] and conditional GAN [25]. [33] has developed GAN-typed inference tools for SBI. In the same vein, [19] used classifiers to diagnose multiple-run MCMC, and [24] used classifiers to learn the likelihood ratio for frequentist calibration, which is further related to using regression to learn the propensity score [36] in observational studies. The theory correspondence between binary classification loss and distribution divergence has been studied in [21, 28, 35]. This present paper not only applies this classifier-for-two-sample-test idea to amortized Bayesian inference to obtain a rigorous test, but also advances the classifier framework by developing the theory of the "label mapping", while the traditional GAN-type approach falls in the one-class-per-group category and deploy binary classifiers as in Example 1. Our extension is particularly useful when samples are overlapped (multiple  $\theta$  per  $y$ ), autocorrelated, and imbalanced.
KL divergence estimate from two samples. As a byproduct, our proposed multiclass-classifier provides a consistent KL divergence estimate (Thm. 3) from two samples. Compared with existing two-sample KL estimate tools from the  $f$ -divergence representation [29] or the Donsker-Varadhan representation [2], our multiclass-classifier estimate appears versatile for it applies to samples with between-sample dependence or within-sample auto-correlation. It is plausible to apply our multiclass-classifier approach to other two-sample divergence estimation tasks such as the  $f$ -GAN [30], which we leave for future investigation.
Limitations and future directions. Like traditional SBC, our method assesses the difference between  $p(\theta |y)$  and  $q(\theta |y)$ , averaged over  $y$ . This is a "global" measure relevant to developing algorithms and statistical software. But sometimes the concern is how close  $p$  and  $q$  are for some particular observation  $y = y^{\mathrm{obs}}$ . "Local" diagnostics have been developed for MCMC [14, 12, 40, 17], variational inference [41, 8] and simulation-based inference [42, 22] that try to assess  $q(\theta |y^{\mathrm{obs}})$  only. It would be valuable to extend our approach to address these cases. Another future direction would be to extend our approach to posterior predictive checks [11, 9, 39] that diagnose how well the statistical model  $p$  fits the observed data.
The authors thank Bruno Regaldo-Saint Blancard for his help on galaxy clustering data and simulations. The authors thank Andreas Buja, Andrew Gelman, and Bertrand Clarke for discussions.