Training Calibrated Conditional Generative Models with a Sample-Only Coverage Loss

Introduction

Inverse problems are ubiquitous across scientific domains, including medical imaging ([28]), astrophysics ([1, 36]), engineering ([30]), and economics ([20]). Measurements of the signal of interest are typically acquired indirectly, only partially, and corrupted by noise, making the recovery of the underlying parameters ill-posed. Many different parameters can explain the same observations, so we seek to sample from the full range of admissible signals; such samples enable uncertainty quantification.

Bayesian inference provides a principled framework for this, modeling the set of all plausible solutions through the posterior distribution of the parameters given the observations. Deep generative models have made it possible to learn this posterior from pairs of parameters and corresponding measurements. Normalizing Flows (NF, [37, 35]), generative adversarial networks (GANs, [15]), and more recently diffusion ([39, 41] (Ho et al. 2020)) and flow ([27, 2]) models have all been successful at conditional sampling. However, each of these models has known failure modes that produce unreliable posteriors. GANs are prone to mode collapse ([31]), and NFs have been shown to be frequently overconfident for posteriors ([17]). Diffusion models, while achieving remarkable sample quality, have also been found to under-represent the uncertainty of the true posterior ([44]). We seek models with valid conditional coverage: for every measurement $y$ and every level $\alpha$, whenever the model assigns probability $\alpha$ to a region of signals, the true signal should fall in that region with probability $\alpha$ under the posterior $p(\theta\mid y)$, that is, $\alpha\cdot100\%$ of the time over hypothetical draws of the true signal conditioned on that measurement. Requiring conditional coverage for all regions is equivalent to the learned posterior estimate $q_{\phi}(\cdot\mid y)$ equaling the true posterior $p(\cdot\mid y)$, but it cannot be checked directly, since a dataset typically contains a single ground-truth signal per measurement.

Two families of training objectives address this problem, each with a complementary weakness. Coverage objectives ([10, 13]) add to the training loss a differentiable measure of the estimator’s coverage, typically computed on its highest probability density (HPD) regions, or a constraint favoring conservative estimators. They target coverage directly, but they have two limitations. Computing HPD regions requires evaluating the estimator density, which restricts these methods to architectures with tractable likelihoods, such as NFs, and excludes diffusion models and GANs. Moreover, they enforce marginal coverage, which does not identify the posterior: many distributions, including the prior, have valid marginal coverage. Strictly proper scoring rules, such as the continuous ranked probability score (CRPS) or the energy score, require only samples and are uniquely minimized at the posterior, but they do not target coverage, and in practice they can yield poorly calibrated models due to finite-sample effects, model misspecification and the geometry of the loss landscape (§ 2).

The basic idea is to assess a conditional generative model by asking how much of its probability mass lies closer to a reference point than the true signal does. For a given measurement $y$, choose a reference point $\psi(y)$ and consider balls centered at this point. The model assigns a probability mass to each ball, and we can check how often the true signal falls inside it. If the model is correct, the two agree at every probability level, and we measure departures from this agreement with a Cramér–von Mises (CvM) discrepancy. This construction builds on TARP (Tests of Accuracy with Random Points, [26]), and we call the resulting loss TARP-CvM. The converse depends on how reference points are chosen. A reference rule assigns a reference point to every measurement, and a single rule is not enough, even one that depends on the measurement. Its coverage is averaged over measurements, so over-coverage at some measurements can balance under-coverage at others. We therefore compute the discrepancy separately for many rules and average the discrepancies, rather than pooling their coverage. Our main theoretical result shows that if the rules can place reference points anywhere and can differ across regions of measurement space, this average vanishes only when the model equals the true posterior. Our contributions are:

• Theory: a sufficient condition on the distribution of measurement-dependent reference rules under which our TARP-CvM loss vanishes if and only if the model equals the true posterior.

• Algorithm: a sample-only, differentiable estimator of TARP-CvM that applies to any conditional generative model, together with a practical construction of reference rules that satisfies this condition.

• Experiments: on simulation-based inference, TARP-CvM matches a density-based coverage loss without evaluating the density; on high-dimensional imaging inverse problems, it is the only regularizer we tested that brings the posterior into the Mira null band on all three problems, with little change in the quality of point estimates.

Problem formulation. We seek to recover parameters $\theta\in\Theta\subseteq\mathbb{R}^{d_{\theta}}$ from indirect measurements $y \in\mathcal{Y} \subseteq\mathbb{R}^{d_y}$. Knowledge of the data likelihood $p(y \mid\theta)$ may be used to generate training samples but is not otherwise required. There may exist a forward model $\mathcal{A}$ such that $y = \mathcal{A}(\theta) + \epsilon$, where $\epsilon$ is a noise vector, and together $\mathcal{A}$ and the distribution of $\epsilon$ control the data likelihood $p(y \mid\theta)$. This setting encompasses a wide range of problems: in inverse problems, $\theta$ may be a high-dimensional medical image of a patient’s anatomy, $y$ the raw measurements collected by an MRI scanner, and $\mathcal{A}$ a model of the scanner’s measurement process; in simulation-based inference, $\theta$ may be the parameters of a physics simulation, with $\mathcal{A}$ the simulator itself; in which case we seek to infer $\theta$ from observations $y$. We seek to sample from the posterior distribution $p(\theta\mid y)$, but cannot evaluate the prior density $p(\theta)$. We have only samples, a training set $\{(\theta_1,y_1),\ldots,(\theta_n,y_n)\}$ of $n$ pairs drawn from the joint distribution $p(\theta,y)$, so that the $\theta_i$ are samples from the prior. Our strategy is to train a conditional generative model, denoted $q_{\phi}(\cdot\mid y)$, which takes as input a measurement vector $y$ and a random seed and generates a sample $\theta$, where $\phi\in\mathbb{R}^{p}$ is the vector of generative model parameters (e.g., neural network weights) to be learned.

Related work

Generative models for inverse problems Model-agnostic methods train a conditional generative model on paired samples $(\theta_i, y_i)$, without using the forward model $\mathcal{A}$. Model-adaptive methods instead combine a generative prior over $\theta$, pre-trained without measurements, with knowledge of $\mathcal{A}$ at sampling time, as in Diffusion Posterior Sampling [6]. We focus on the model-agnostic setting, but the core ideas translate naturally to model-adaptive samplers.

Coverage tests Two standard approaches for empirically testing coverage of $q_\phi$ are Highest Posterior Density (HPD, Hermans et al., [17]) testing and Tests of Accuracy with Random Points (TARP, Lemos et al., [26]). HPD testing constructs, for each test point $y$, the smallest region $C \subset\mathbb{R}^{d_\theta}$ under $q_\phi(\cdot\mid y)$ that contains $\alpha$ probability mass, and then checks whether the true $\theta$ falls inside that region at the expected rate. However, HPD requires evaluating the density of $q_\phi$, which is often unavailable for generative models, and valid HPD coverage is not a sufficient condition for accuracy: a model can pass the HPD test while still differing from the posterior. Figure 1 illustrates the limitations of using HPD to train a conditional generative model vis-à-vis our proposed measure, TARP-CvM; details on this experiment are in § 5.1.

Vector (PDF)

TARP instead constructs regions as balls centered at randomly chosen reference points, measuring coverage by the fraction of samples from $q_\phi(\cdot\mid y)$ that fall closer to the reference point than the true $\theta$ does. Lemos et al. [26] show that correct TARP coverage for every assignment of reference points to measurements is necessary and sufficient for the estimator to equal the posterior, so TARP can detect errors that HPD-based tests miss, while requiring only samples from $q_\phi$ rather than density evaluations. The key distinction is that HPD regions are implicitly centered by the geometry of the estimated density, whereas TARP regions can be centered at arbitrary reference points. This flexibility lets TARP probe aspects of the posterior that HPD coverage cannot capture. Which assignments are needed in practice is left open; we make this precise in § 3. The output of TARP is a coverage curve that plots the empirical coverage rate against the nominal credibility level $\alpha$, with a perfectly calibrated model lying on the diagonal. Although correct TARP coverage for every possible assignment of reference points characterizes the posterior, this does not immediately provide a practical training objective. In particular, averaging coverage over randomly chosen reference points can conceal errors, because overcoverage and undercoverage may cancel. Our contribution is to identify a tractable family of reference rules whose individual coverage discrepancies can be averaged without losing the ability to identify the posterior.

Because TARP produces a coverage curve rather than a scalar, it does not directly support quantitative comparisons between models. Mira [38] provides such a scalar. Over many randomly drawn regions and joint samples $(\theta, y) \sim p(\theta, y)$, it estimates the probability that the true $\theta$ falls in a region given the fraction of samples from $q_\phi(\cdot\mid y)$ that fall in it. Under calibration, the Mira score equals 2/3, and values below this level indicate overconfidence. Since our loss is built on TARP, evaluating it with TARP coverage alone would be partly circular. We therefore use Mira as an independent metric: it relies on its own random regions, is not optimized by any of the methods we compare, and provides a calibrated reference value with a null band that accounts for sampling noise.

Energy Score and CRPS Using MSE as a loss function assesses a model’s ability to correctly estimate the conditional mean, but it does not reward modeling the spread of possibilities correctly and can prefer overconfident generators. The Energy Score (ES, Gneiting and Raftery, [14])

addresses this by penalizing samples far from the true $\theta$ while rewarding spread. It is strictly proper on distributions with finite first moment, with expectation under $p(\cdot\mid y)$ uniquely minimized at $q_\phi(\cdot\mid y) = p(\cdot\mid y)$, and it can be estimated from samples alone. Its univariate case is the continuous ranked probability score (CRPS). In high dimension, the CRPS is often summed over coordinates (pCRPS), which parallelizes across pixels, but only identifies the marginals of the posterior. Strict propriety says that the population objective has the true posterior as its unique global minimizer. It does not guarantee that training finds this minimizer. The loss is estimated on finite batches and minimized by gradient descent, which can stop at local minima, and if the model class does not contain the posterior, the guarantee says nothing about which approximation is found. Training with a strictly proper scoring rule therefore need not produce a calibrated model.

Two objectives can share the same global minimizer but have very different landscapes. In a Gaussian toy model, the energy score is much flatter than TARP-CvM in the direction of the posterior spread when the mean is biased, so that noisy gradients are less likely to correct the spread, whereas TARP-CvM penalizes miscoverage directly (Appendix § 12.§ 3). Our experiments reflect this difference; on most of our imaging inverse problems, ES and pCRPS improve coverage but remain overconfident, whereas TARP-CvM reaches calibration (§ 5.2).

Calibration objectives Several works have proposed training objectives that encourage calibration directly, including the HPD objective of Falkiewicz et al. [13], the rank-based statistics of Yao et al. [45], and the balanced NRE approach of Delaunoy et al. [10], which favors conservative rather than exactly calibrated estimators. These methods target marginal calibration, which does not ensure calibration conditionally on the measurement (Kuleshov et al., [25]). Conformal prediction (Angelopoulos and Bates, [3]) instead acts after training: using a held-out calibration set, it turns the predictions of any trained model into sets that contain the true value with a chosen probability, without assumptions on the data distribution and for any calibration set size. This guarantee is also marginal, since conditional coverage is unattainable in this distribution-free setting, and the procedure does not provide a loss or score that can be used to train the model. Conditioning the calibration statistic on $y$ gives a stronger notion of calibration, explored both as a diagnostic (Zhao et al., [47]) and for post-hoc recalibration (Dey et al., [11]). Our loss enforces it during training, at the population level.

The TARP-CvM loss: A strict divergence

This section builds the TARP-CvM loss in three stages. We first describe how the TARP coverage test of Lemos et al. [26] assesses a conditional generative model: it centers a ball at a reference point chosen from the measurement, measures how much of the estimator’s mass lies closer to the reference point than the true parameter does, and checks whether these masses are uniformly distributed (§ 3.1). We then relate this test to credible regions and to the expected coverage probability, which places TARP alongside HPD-based tests (§ 3.2). Finally, we turn the test into a population objective and give conditions under which it vanishes only at the true posterior (§ 3.3).

The TARP construction

TARP replaces the HPD level sets of the density (§ 2) by balls, whose mass can be estimated from samples alone. Given a pair $(\theta,y)$ drawn from the joint distribution $p(\theta,y)$, for instance a training pair $(\theta_i,y_i)$, the construction proceeds in four steps.

Step 1: choose a reference point. For the measurement $y$, choose a reference point $\psi(y) \in\Theta$. Since this choice must be made for every measurement, we formalize it as a rule.

Step 2: construct a ball. Center a ball at $\psi(y)$, with radius chosen so that the ball just contains the true $\theta$, that is, the ball $B(\psi(y),d(\theta,\psi(y)))$, where $d$ is a distance on $\Theta$ and $B(c,r)=\{\tilde{\theta}:d(\tilde{\theta},c)\le r\}$ is the closed ball.

Step 3: evaluate its mass under the estimator. Define the credibility value as the probability mass of this ball, where the probability is with respect to the distribution $q_\phi(\cdot\mid y)$:

In practice, it is the fraction of draws $\tilde{\theta}\sim q_\phi(\cdot\mid y)$ that fall inside the ball (see Figure 2b).

Vector (PDF)

Step 4: evaluate coverage. Repeating Steps 1–3 over pairs $(\theta,y)\sim p(\theta,y)$, with the same rule $\psi$, gives a distribution of credibility values. The conditional coverage of rule $\psi$ at measurement $y$ and the marginal coverage, its average over measurements, are

(conditional coverage), (1)

(marginal coverage). (2)

Since $F_\psi(\alpha\mid y)$ depends on $\psi$ only through the reference point $\psi(y)$, we also write $F_c(\alpha\mid y)$ for the coverage at $y$ of balls centered at a point $c\in\Theta$, so that $F_\psi(\alpha\mid y)=F_{\psi(y)}(\alpha\mid y)$. If the estimated posterior is correct, the credibility values of the true parameters are uniformly distributed on $[0,1]$ at every measurement, $F_\psi(\alpha\mid y)=\alpha$ (Lemma 1 below). The original TARP test pools the credibility values over randomly drawn reference points. We instead compare them to the uniform distribution separately for each rule and average the resulting discrepancies (§ 3.3).

Credible regions and expected coverage

To relate TARP to HPD-based tests, we view the balls of § 3.1 as credible regions of the estimator and their marginal coverage $F_\psi(\alpha)$ as expected coverage probability.

In this paper, we compare two credible region maps. The HPD region is the level set of $q_\phi(\cdot\mid y)$ (also the region with the smallest volume) with mass $\alpha$ ([5, 21]). Given a rule $\psi$, the TARP region is a ball centered at reference point $\psi(y)\in\Theta$, with radius ensuring the ball has mass $\alpha$:

The credibility value $s_\psi(\theta,y)$ is the level at which $\theta$ enters the TARP region: $\theta\in R_{\alpha,\psi}^{\mathrm{TARP}}(y)\Longleftrightarrow s_\psi(\theta,y)\le\alpha$.

By construction, $F_{\psi}(\alpha)=\operatorname{ECP}(q_{\phi},\alpha,R_{\alpha,\psi}^{\mathrm{TARP}})$. So $q_{\phi}$ is calibrated for the TARP region given by $\psi$, when $F_{\psi}$ is the identity, i.e., when the credibility values are uniformly distributed.

In other words, a correct estimator is always calibrated, no matter how the regions are chosen (Proof in Appendix A). However the converse, calibration for a map implies $q_{\phi}=p$, depends on the choice of the map, and is the central question of § 3.3.

The population objective

$D_{\rho}=0$ iff the estimator is calibrated for every rule $\psi$ in the support of $\rho$. For a single rule, the integrand compares each credibility level $\alpha$ with the frequency $F_{\psi}(\alpha)$ at which the true parameter falls in the corresponding region. The inner integral is thus a continuous analogue of the expected calibration error used in classification ([33]), and closely resembles the calibration error of Kuleshov et al. [25] for regression. The outer expectation over rules is what distinguishes it from these per-model calibration scores, and the choice of $\rho$ determines the properties of $D_{\rho}$. We motivate the requirements on $\rho$ below and formalize them in Definition 6.

For a fixed rule $\psi$, the integral is the Cramér–von Mises distance ([7]) $D_{\mathrm{CvM}}(X_{1},X_{2})=\int_{0}^{1}\left[\mathbb{P}[X_{1}\leq\alpha]-\mathbb{P}[X_{2}\leq\alpha]\right]^{2}\,\mathrm{d}\alpha$, where $X_{1}=s_{\psi}(\theta,y)$, the credibility of the true parameters, whose CDF is $F_{\psi}$, and $X_{2}\sim\mathcal{U}[0,1]$ is the uniform distribution. $D_{\rho}$ is non-negative and, under Assumption 1 and for any admissible $\rho$ (Definition 6), vanishes if and only if $q_{\phi}=p$ (Theorem 1); in this sense it is a divergence between conditional distributions. It is not a metric, since it is not symmetric in $(p,q_{\phi})$ and does not satisfy the triangle inequality, and it is not a scoring rule, since it is quadratic in the distribution of $(\theta,y)$ rather than an average of per-example scores.

Constant reference points cannot identify the posterior. A rule that ignores the measurement lets the prior pass the coverage test. Let us consider a constant rule $\psi(y)=c$. If our estimator returns the prior, $q_{\phi}(\theta\mid y)=p(\theta)$, the TARP region $R_{\alpha,\psi}^{\mathrm{TARP}}(y)=R_{\alpha,c}^{\mathrm{TARP}}$ no longer depends on $y$, since both its center and the distribution defining its radius are fixed. Then $\operatorname{ECP}(q_{\phi},\alpha,R_{\alpha,\psi}^{\mathrm{TARP}})=\mathbb{E}_{p(\theta,y)}\left[\mathbf{1}_{\theta\in R_{\alpha,c}^{\mathrm{TARP}}}\right]=\int p(y)\int_{R_{\alpha,c}^{\mathrm{TARP}}}p(\theta\mid y)\,\mathrm{d}\theta\,\mathrm{d}y=\int_{R_{\alpha,c}^{\mathrm{TARP}}}\left(\int p(\theta\mid y)p(y)\,\mathrm{d}y\right)\mathrm{d}\theta=\int_{R_{\alpha,c}^{\mathrm{TARP}}}p(\theta)\,\mathrm{d}\theta$. Since $q_{\phi}(\theta\mid y)=p(\theta)$ and by definition of a credible region, we get $\operatorname{ECP}(q_{\phi},\alpha,R_{\alpha,\psi}^{\mathrm{TARP}})=\alpha$, which evaluates the estimator as calibrated.

A single measurement-dependent rule is still insufficient. Reference points must depend on the measurement, but this alone does not guarantee that the loss identifies the posterior. A single rule gives each measurement its own reference point, but its coverage $F_{\psi}(\alpha)=\int F_{\psi}(\alpha\mid y)p(y)\,\mathrm{d}y$ is an average over measurements, and errors at different measurements can cancel. Suppose there are two measurements $y_{0}$ and $y_{1}$, each with probability $\frac{1}{2}$, and an estimator that is too wide at $y_{0}$ and too narrow at $y_{1}$, with $F_{\psi}(\alpha\mid y_{l})=\alpha+(-1)^{l}\delta\,\alpha(1-\alpha)$ for $l\in\{0,1\}$ and $\delta\in(0,1)$. The coverage curves at $y_{0}$ and $y_{1}$ deviate from the diagonal in opposite directions and average to it, so $F_{\psi}(\alpha)=\alpha$ at every level: the estimator is calibrated under $\psi$ yet wrong at both measurements. Choosing a better single rule does not help, since the counterexample holds for any $\psi$.

Comparing rules requires local differences. For a wrong estimator, the conditional coverage $F_{\psi}(\alpha\mid y)$ (Equation 1) depends on where the ball is centered, so errors can be revealed by comparing different reference points at the same measurement. The loss, however, only sees marginal coverages $F_{\psi}(\alpha)$, so this comparison must be made across several rules. It isolates a subset of measurements only if the collection contains pairs of rules that agree on most measurements and differ on that subset: the difference between their marginal coverages then depends only on the subset. If every pair of rules differs at all measurements, errors can still cancel within each marginal average.

For instance, a simpler reference would add noise to a point estimate $b(y)$ at scale $\varsigma$ (see § 4.2 for definitions), $\psi(y_i) = b(y_i) + \varsigma\xi_i$, with $\xi_i$ drawn either independently for each example or once for the whole batch. Neither is sufficient. With independent draws, the batch yields a single coverage curve, $\bar{F}(\alpha) = \int\mathbb{E}_{\xi} [F_{b(y)+\varsigma\xi}(\alpha\mid y)]p(y)\,dy$, and a single statistic $\int_0^1 (\bar{F}(\alpha)-\alpha)^2\,d\alpha$. The expectation over reference points then sits inside the square of Equation 4, and the examples where $F_{b(y_i)+\varsigma\xi_i}(\alpha\mid y_i) > \alpha$ can balance those where $F_{b(y_i)+\varsigma\xi_i}(\alpha\mid y_i) < \alpha$, so that $\bar{F}(\alpha) = \alpha$ even if $F_{b(y_i)+\varsigma\xi_i}(\alpha\mid y_i) \ne\alpha$ for every example. If instead a single displacement $\xi_r$ is shared by the batch, each draw is a rule $\psi_r(y) = b(y)+\varsigma\xi_r$ giving a statistic $I_{\psi_r}$, and the rules satisfy (R1) below. But each $I_{\psi_r}$ constrains $F_{\psi_r}(\alpha\mid y)$ only through an average over all measurements, so nothing prevents over-coverage at some measurements from balancing under-coverage at others for every $\xi_r$ at once, and Theorem 1 does not apply to these rules.

Patching rules resolves the problem. Given two reference rules, a patched rule uses the first rule for measurements on one side of a hyperplane and the second rule for measurements on the other side. Thus, we can change the reference points for one subset of measurements while leaving the rest unchanged. In the example above, with $\delta= 0.8$ and $\alpha= \frac{1}{2}$, the coverages are 0.7 and 0.3 under $\psi$; take a second calibrated rule $\psi'$, centered elsewhere, with coverages 0.6 and 0.4. Switching rules across measurements, $\psi$ at $y_0$ and $\psi'$ at $y_1$, gives $\frac{1}{2}(0.7+0.4)=0.55$, and the miscalibration appears.

Definition 5. For rules $\psi$, $\psi'$ and a measurable set $H \subseteq\mathcal{Y}$, the patched rule uses $\psi$ on $H$ and $\psi'$ elsewhere, $(\psi\oplus_H \psi')(y) = \psi(y)\,1_{y\in H} + \psi'(y)\,1_{y\notin H}$.

Patching along all half-spaces is enough (Lemma 5), which is condition (R2) below. The reference rules must have two properties. First, at every measurement, the available reference points must be sufficiently rich to probe the entire parameter space. Second, the collection of rules must allow reference points to be changed on arbitrary half-spaces of measurement space while remaining unchanged elsewhere. These two properties ensure that coverage errors cannot be hidden by averaging, and that the resulting coverage tests distinguish different conditional distributions.

Definition 6. A distribution $\rho$ over rules is admissible if there exists a countable family of rules $\Psi\subseteq\operatorname{supp}\rho$ (Appendix § 9.§ 1) such that

(R1) For every $y$, the centers $\{\psi(y):\psi\in\Psi\}$ are dense in $\Theta$.

(R2) For all $\psi,\psi' \in\Psi$ and every half-space $H = \{y:\langle u,y\rangle\le t\}$, with $u$ a unit vector of $\mathbb{R}^{d_y}$ and $t \in\mathbb{R}$, the patched rule $\psi\oplus_H \psi'$ is in $\operatorname{supp}\rho$.

With (R1), the loss can place reference points anywhere near the posterior mass at each measurement. With (R2), they may be used on the two sides of a hyperplane in measurement space, which is enough to tell measurements apart.

The switch rule of Equation 8 avoids both problems of the noisy point estimate above. Because each side of the hyperplane shares one displacement, each draw is a rule with its own statistic. Using one displacement on $H$ and another on its complement changes the reference points of the measurements in $H$ only, so errors inside $H$ can no longer be balanced by errors outside it. Calibration of every such rule, for every $H$, forces the coverage at almost every measurement to be the same for every displacement (condition (R2) and Lemma 5).

Definition 7. The TARP-CvM loss is (1) for an admissible $\rho$.

Proof in Appendix § 9. The direction $q_{\phi}=p \Rightarrow D_{\rho}=0$ follows from Lemma 1. The converse proceeds in three steps. First, zero loss implies calibration for every rule in the support of $\rho$. Second, patching (R2) implies that the conditional coverage at almost every measurement does not depend on the reference point. Third, since reference points are dense (R1), this invariance holds only if the estimated and true posteriors agree.

Lemos et al. [26] show that a posterior estimator is exact if its expected coverage is correct for every function assigning reference points to measurements. Their argument relies on a functional derivative with respect to this function. They do not specify how reference points should be selected from the measurement, and in their experiments reference points are drawn independently of it, a setting in which the prior also passes the test. Moreover, the original TARP pools coverage over randomly drawn reference points, and the pooled coverage can match the nominal level even when a distribution is miscalibrated with respect to an individual rule, because their errors can cancel in the average. Theorem 1 shows how our TARP-CvM addresses both issues. It gives an explicit sufficient condition on how reference points depend on the measurement (Definition 6). Our loss averages the miscalibration of a distribution with respect to each rule rather than pooling, so that zero loss implies calibration for every rule in the support.

Making the objective trainable

At each training iteration, we draw a batch of paired observations and ground-truth signals $(\theta_i, y_i)$, generate samples for each observation $\tilde{\theta}_{i,m}$, and construct several reference rules. For each rule, we compute the credibility of every ground-truth signal $\theta_i$ relative to the corresponding generated samples. We then calculate a Cramér–von Mises discrepancy between the empirical distribution of these credibility values and the uniform distribution. The average discrepancy over reference rules is added to the base generative-model training loss.

The closed form estimator

The objective term in Equation (1) is the squared area $I_{\psi} = \int_{0}^{1} (F_{\psi}(\alpha)-\alpha)^2\,\mathrm{d}\alpha$ between the cumulative distribution function of the credibility values and the diagonal. $F_{\psi}$ is empirically estimated with $N$ credibility values $s_i=s_{\psi}(\theta_i,y_i)$ from pairs $(\theta_i,y_i)$ of the batch, $\widehat{F}_{\psi}(\alpha) \approx\frac{1}{N}\sum_i \mathbf{1}_{s_i\leq\alpha}$, computed with the same rule $\psi$. $I_{\psi}$ is a Cramér–von Mises statistic, and its derivative admits a closed-form expression ([45]). The statistic measures the squared discrepancy between the empirical distribution of the credibility values and the uniform distribution. The usual integral representation involves indicator functions and is not convenient for gradient-based optimization. By integrating analytically over the credibility levels, we obtain the following expression in terms of the sorted credibility values. See Appendix § 11 for details on deriving this empirical estimator:

where $s_{(1)} \leq\cdots\leq s_{(N)}$ are the sorted credibility values. Writing $\operatorname{rank}(i)$ for the position of $s_i$ in this ordering, so that $s_i=s_{(\operatorname{rank}(i))}$, the gradient is

The closed form expression is what makes the statistic differentiable, by turning evaluation on a grid of levels $\alpha$, where most indicator counts would have zero gradient, into a polynomial of the sorted values. Integrating over $\alpha$ first removes the indicators, and the resulting expression depends on $s_i$ values and rank. It was used by Yao et al. [45] to turn a rank-uniformity test into a differentiable objective for simulation-based inference. We differ in the choice of the statistic. Coordinate-wise ranks constrain only the one-dimensional marginals of the estimator, so the prior estimator also passes the test, whereas uniformity of the credibility value across admissible rules forces the estimator to equal the posterior as stated in Theorem 1.

Unbiased estimator. Equation Equation (2) is the plug-in estimate of the inner integral of Equation (1). Written as an average over all pairs $(i,j)$ of credibility values, it includes the diagonal pairs $i=j$, so its expectation exceeds it by $\frac{1}{N}\int_0^1 F_\psi(1-F_\psi)\,d\alpha$, a term that does not vanish at the posterior. Excluding the diagonal pairs gives an unbiased U-statistic (Hoeffding, [19]),

Equation Equation (4) is therefore preferable when an unbiased estimate of the population discrepancy is desired (Proposition 2, see Appendix § 11 for derivation). In our experiments, however, we use the simpler plug-in statistic in Equation (2). Its finite-sample bias is not removed, but the same estimator is used consistently throughout training.

Reference points

The objective Equation (1) is an average over reference rules. Each optimization step draws $R$ rules and applies each one to every pair in the batch. It then computes the CvM statistic of Equation (2) separately for each rule and averages the $R$ statistics. This keeps the expectation over rules outside the squared deviation, as in Equation (1). The rules are drawn from a mixture of three groups: switch rules, prior rules and posterior rules.

Switch rules. Figure 2 illustrates the construction. A switch rule needs a point estimate $b(y)$ with nontrivial posterior likelihood. For a linear forward model $y=A\theta+\epsilon$, we take the adjoint reconstruction $b(y)=A^{\mathsf{T}}y$, or the pseudo-inverse $A^+y$. Without an adjoint, $b$ can be a regressor of the posterior mean, trained beforehand on the same pairs (§ 5.3). A switch rule splits the measurements of the batch with a hyperplane and displaces the point estimate by one displacement on one side and by another on the other side, so that it is a single rule patched across a half-space (Definition 5). It is determined by two displacements $\xi,\xi'$, a unit vector $u$ in measurement space that orients the hyperplane, and an example $j$ of the batch through which it passes, and assigns to every example $y_i$ the reference point

where $\varsigma\in\mathbb{R}^{d_b}_{>0}$ is a fixed reference scale, applied component-wise and treated as a constant. Drawing $\xi,\xi' \sim\mathcal{N}(0,I)$, a random unit vector $u$ and a random example $j$ chooses a switch rule at random; once drawn, it is a fixed rule in the sense of Definition 1. In our experiments, we set $\varsigma$ to the root-mean-square residual of $b$ on a held-out tenth of the training simulations,

with the square and root taken element-wise, so that the reference displacements match the average posterior scale. $b$ and $\varsigma$ are both frozen before the posterior estimator is trained.

The hyperplane splits the batch into two groups that receive different displacements, so within a rule each measurement receives one reference point. Across the $R$ rules of a step, each measurement receives $R$ reference points. The switch rules satisfy both admissibility conditions in Definition 6: their reference points are dense in parameter space, and they can implement arbitrary half-space patches. The displacements have positive densities, so reference points come arbitrarily close to every point of $\Theta$ (R1). Every rule is a patch of two rules across a half-space, and every patch can be drawn, since the displacements, the direction and the position of the hyperplane are all drawn from distributions with full support (R2).

Prior and posterior rules. Admissibility guarantees that the true posterior is the unique minimizer of the loss, but we observe that adding other rules helps training. A prior rule uses a single reference point drawn from the prior and shared by the batch. A posterior rule uses the $M$ posterior draws of each pair as reference points, detached from the computational graph. For the reference $\tilde{\theta}_m$, the credibility of $\theta_i$ is computed from the $M-1$ other draws. When a generated posterior sample is used as the reference point, that same sample must be excluded from the set of samples used to estimate the credibility value. Otherwise, the reference point would contribute a guaranteed zero-distance sample, changing the distribution of credibility values even when the estimator is correct.

Theorem 1 still applies. Every term of the loss is nonnegative, and the switch rules keep a fixed share of the mixture, so the loss can only vanish if the loss restricted to switch rules vanishes, which forces $q_\phi=p$. The switch rules are essential for the theoretical guarantee, and the other two are supplementary choices intended to help optimization.

Relaxation of the indicator

To differentiate the loss with respect to $\phi$, we backpropagate through the sampler. For a diffusion model, we draw all the noises of the sampler in advance and hold them fixed, so that the sampler becomes a deterministic function of $\phi$. This is the reparameterization trick [24]. The exact credibility value is the fraction of generated samples whose distance from the reference point is no greater than the distance of the true signal. The indicator in this definition is nondifferentiable, so we replace it with a sigmoid approximation $\sigma$ at temperature $\tau>0$, for $\tilde{\theta}_m \sim q_\phi(\cdot\mid y)$,

Algorithm 1 summarizes one training step, and Appendix B details its complexity.

Experiments

TARP-CvM as training signal

We first consider a setting in which the true posterior is known analytically. This experiment isolates the distinction between marginal coverage and posterior accuracy, allowing us to directly compare the minimizers of the HPD and TARP-CvM objectives. We illustrate Theorem 1 on a Gaussian linear inverse problem $y=A\theta+\epsilon$, with a Gaussian prior and fewer measurements than parameters (Appendix E.1). The posterior is Gaussian, $p(\theta\mid y)=\mathcal{N}(W^{*}y,S^{*})$, and we fit a Gaussian estimator $q_{\phi}(\theta\mid y)=\mathcal{N}(Wy,S)$ by gradient Algorithm 1 One training step with the TARP-CvM regularizer

Require: batch $\{(\theta_i,y_i)\}_{i=1}^{N}$, posterior point estimator $b:\mathcal{Y}\to\Theta$, reference scale $\varsigma$ ((6)), base loss $\mathcal{L}_{\mathrm{base}}$ of the generative model, weight $\lambda$, number of draws $M$, numbers of switch, prior and posterior rules $R_s,R_p,R_{\mathrm{post}}$ with $R=R_s+R_p+R_{\mathrm{post}}$, temperature $\tau$

1. for $i=1,\ldots,N$ do

2. Draw $\tilde{\theta}_{i,1},\ldots,\tilde{\theta}_{i,M}\sim q_\phi(\cdot\mid y_i)$ with the reparameterization trick

3. end for

4. for $r=1,\ldots,R$ do

5. Let $k\in\{\text{switch},\text{prior},\text{posterior}\}$ be the group of rule $r$

6. Draw a rule $\psi_r\sim\rho_k$:

7. if $k=\text{switch}$ then

8. Draw $\xi,\xi'\sim\mathcal{N}(0,I)$, a unit vector $u$ and an index $j$

9. $\psi_r(y)=b(y)+\varsigma\left(\xi\,\mathbf{1}_{\langle u,y\rangle\le\langle u,y_j\rangle}+\xi'\,\mathbf{1}_{\langle u,y\rangle>\langle u,y_j\rangle}\right)$

10. $\mathcal{M}_r\gets\{1,\ldots,M\}$

11. else if $k=\text{prior}$ then

12. Draw $c\sim p(\theta)$ and set $\psi_r(y)=c$

13. $\mathcal{M}_r\gets\{1,\ldots,M\}$

14. else

15. Draw an index $m'$ and set $\psi_r(y_i)=\operatorname{sg}(\tilde{\theta}_{i,m'})$

16. $\mathcal{M}_r\gets\{1,\ldots,M\}\setminus\{m'\}$

17. end if

18. for $i=1,\ldots,N$ do

19. $s_i\gets\frac{1}{|\mathcal{M}_r|}\sum_{m\in\mathcal{M}_r}\sigma\left(\left[d(\theta_i,\psi_r(y_i))-d(\tilde{\theta}_{i,m},\psi_r(y_i))\right]/\tau\right)$

20. end for

21. $\mathcal{L}_r\gets\operatorname{CvM}(s_{1:N})$ ▷ (2)

22. end for

23. Compute $\mathcal{L}_{\mathrm{base}}(\phi)$ on its own batch ▷ e.g. DSM or NLL

24. $\mathcal{L}\gets\mathcal{L}_{\mathrm{base}}(\phi)+\lambda\frac{1}{R}\sum_{r=1}^{R}\mathcal{L}_r$

25. Update $\phi$ with a gradient step on $\mathcal{L}$

descent on either an HPD-loss or TARP-CvM. The estimator is HPD-calibrated for every $W$, provided that $S=S^*+(W^*-W)\Sigma_y(W^*-W)^\top$, where $\Sigma_y$ is the covariance of the measurements. Thus, an estimator with a biased conditional mean can still achieve correct HPD coverage if its covariance is sufficiently inflated. The HPD objective therefore admits a family of minimizers, including the prior, even though most of these distributions do not equal the true posterior. This manifold of minimizers contains the prior, at $W=0$, and every point of it other than the posterior has a biased mean and a covariance wider than the posterior’s. Gradient descent on HPD coverage can therefore stop at any point of this manifold, not only at the posterior. With measurement-dependent references, TARP-CvM instead has the posterior as its unique minimizer (Theorem 1). In Figure 1, TARP-CvM recovers the posterior contours, whereas the HPD fit is biased and too conservative. The HPD coverage error cannot distinguish the two models, although their KL divergence to the posterior differs strongly.

Imaging inverse problems

We study posterior coverage calibration in high-dimensional posteriors on three inverse problems from InverseBench Zheng et al. [48]. Each has the form $y=A(\theta)+\epsilon$. These are linear inverse scattering (IS) and compressed-sensing MRI (MRI), whose forward operators are linear, and Navier–Stokes (NS), whose forward operator is nonlinear and only available as a black-box solver (Appendix § 13.§ 3). For IS, $\epsilon$ is additive Gaussian noise. For MRI and NS, $\epsilon\equiv0$, so the posterior has no density on the full signal space. On MRI, TARP-CvM therefore measures distances in the null space of $A$, where Theorem 1 applies (Appendix § 13.§ 3). Unlike the plug-and-play methods InverseBench was designed for, we train fully amortized conditional models on pairs $(\theta,y)$. The measurement is concatenated to the network input, as its adjoint reconstruction $A^\top y$ for IS and MRI, and at its native $64\times64$ resolution for NS, which matches the resolution of the signal after $2 \times2$ patching. All models are pretrained with denoising score matching (DSM, Song et al. [41]). They are then fine-tuned for 24k steps with one of four objectives, DSM alone ($\varnothing$), a control that isolates the effect of additional training, or DSM plus a calibration regularizer, which is the energy score (ES), the pixel-wise CRPS (pCRPS), or TARP-CvM. All three regularizers use the same calibration batch, the same number of posterior samples per measurement, and the same $T = 20$-step DDIM sampler (Song et al. [40]), through which gradients are backpropagated (see Appendix § 13.§ 3 for details).

We measure coverage with the $L^1$ TARP coverage error, $\int_0^1 \lvert\operatorname{ECP}(q_\phi, \alpha, R^{\mathrm{TARP}}_{\alpha,\psi}) - \alpha\rvert\,d\alpha$. We compute the coverage using standard binary indicators rather than the sigmoid relaxations used during training, and use the $\ell_2$ distance to define the TARP balls. We also use the Mira score (Sharief et al. [38]), which equals 2/3 for a calibrated model; values below the null band indicate overconfidence. For quality, we report in Table 3 the PSNR, SSIM, and relative $L^2$ error of the posterior mean over $K = 32$ samples, the PSNR of a single draw, and the posterior spread (the per-pixel standard deviation across the $K$ samples, in physical units), all computed on 64 held-out observations, together with the DSM loss on the full validation split, averaged over three seeds.

Problem	Stage	Posterior mean (MMSE)			Draw
		PSNR↑	SSIM↑	rel. $L_2$↓	PSNR	spread	DSM↓
NS	DSM	27.83	0.7851	0.2368	24.58	1.1169	0.0595
	+ ∅	28.22 ± 0.02	0.7969 ± 0.0004	0.2280 ± 0.0005	25.01 ± 0.01	1.0606 ± 0.0025	0.0594 ± 0.0000
	+ ES	29.26 ± 0.03	0.8246 ± 0.0009	0.2061 ± 0.0008	25.81 ± 0.04	0.9847 ± 0.0032	0.0590 ± 0.0000
	+ pCRPS	29.23 ± 0.04	0.8257 ± 0.0010	0.2067 ± 0.0008	25.79 ± 0.04	0.9843 ± 0.0037	0.0590 ± 0.0000
	+ TCvM	27.88 ± 0.05	0.7874 ± 0.0011	0.2354 ± 0.0013	24.22 ± 0.05	1.2178 ± 0.0101	0.0597 ± 0.0000
IS	DSM	29.62	0.9221	0.0906	27.93	0.0142	0.1159
	+ ∅	29.79 ± 0.01	0.9229 ± 0.0002	0.0898 ± 0.0000	28.23 ± 0.01	0.0131 ± 0.0000	0.1182 ± 0.0001
	+ ES	29.72 ± 0.02	0.9217 ± 0.0003	0.0908 ± 0.0002	28.21 ± 0.01	0.0128 ± 0.0000	0.1190 ± 0.0001
	+ pCRPS	29.71 ± 0.02	0.9216 ± 0.0004	0.0909 ± 0.0003	28.18 ± 0.01	0.0129 ± 0.0000	0.1194 ± 0.0002
	+ TCvM	29.64 ± 0.04	0.9224 ± 0.0001	0.0905 ± 0.0001	27.88 ± 0.11	0.0179 ± 0.0008	0.1186 ± 0.0000
MRI	DSM	37.11	0.8825	0.1232	35.10	0.0285	0.1916
	+ ∅	37.47 ± 0.00	0.8872 ± 0.0000	0.1186 ± 0.0000	35.41 ± 0.00	0.0280 ± 0.0000	0.1901 ± 0.0000
	+ ES	37.79 ± 0.01	0.8902 ± 0.0003	0.1150 ± 0.0001	34.71 ± 0.06	0.0320 ± 0.0003	0.1902 ± 0.0000
	+ pCRPS	37.75 ± 0.07	0.8905 ± 0.0007	0.1153 ± 0.0007	34.93 ± 0.34	0.0331 ± 0.0030	0.1902 ± 0.0000
	+ TCvM	36.68 ± 0.07	0.8751 ± 0.0017	0.1291 ± 0.0010	33.49 ± 0.37	0.0443 ± 0.0054	0.1910 ± 0.0001

All pretrained models are strongly overconfident (Mira between 0.51 and 0.61, Figure 3). Additional DSM training does not fix this, so the gains below come from the regularizers. On NS and IS, TARP-CvM gives the lowest coverage error (0.11 and 0.28, against 0.17 and 0.35 for the best scoring rule). On NS all three regularizers reach the Mira null band, but on IS only TARP-CvM does (0.68, against 0.53 for ES and 0.57 for pCRPS). On MRI, the two metrics disagree. ES has the lowest coverage error, but remains overconfident according to Mira, while TARP-CvM is the only method inside the band.

Vector (PDF)

Increasing posterior variance can reduce the PSNR of individual posterior draws even when the posterior mean remains unchanged. This is not necessarily a deterioration in reconstruction accuracy: it can reflect the fact that the model is representing uncertainty that was previously missing. Calibration has little effect on the quality of generated samples. The validation DSM loss stays within 1% of the control for every method. Posterior-mean PSNR is unchanged on IS and NS and drops by 0.8 dB on MRI. Per-draw PSNR drops more, by 0.35 to 1.9 dB, but this reflects the metric rather than accuracy. The expected squared error of a draw $\tilde{\theta} \sim q(\cdot\mid y)$ decomposes as $\mathbb{E}[\lVert\tilde{\theta} - \theta^*\rVert^2] = \lVert\mathbb{E}_q[\tilde{\theta}] - \theta^*\rVert^2 + \operatorname{Var}_q(\tilde{\theta})$, where $\operatorname{Var}_q(\tilde{\theta}) = \operatorname{tr}\operatorname{Cov}_q(\tilde{\theta})$ denotes the total variance, so a model that collapses its variance gains per-draw PSNR without improving its mean. Correct coverage requires nonzero variance, and TARP-CvM widens the posterior by 15 to 58% relative to the control, which lowers per-draw PSNR accordingly. On NS, the scoring rules instead improve the posterior mean by about 1 dB while slightly reducing the spread. This suggests that the pretrained model had errors due to both a biased mean, which the scoring rules correct, and insufficient spread, which TARP-CvM corrects. The wider spread of TARP-CvM may also partly compensate for residual error in the posterior mean.

Sampler depth. Because the calibration term backpropagates through the sampler, training fixes the sampler depth, but any depth can be used at inference. We train with $T = 20$ and evaluate at $T \in\{20, 100, 500\}$ (Figure 4). On NS and IS, Mira stays inside the null band at all depths, so calibration learned with a short chain survives at longer ones. On MRI, however, Mira decreases and leaves the null band at higher $T$. At $T = 100$ and 500 it even falls below the DSM-only model, whose Mira improves with $T$. This suggests that on MRI part of the learned calibration compensates for the discretization error of the 20-step sampler rather than correcting the underlying score, so the correction no longer matches a finer discretization. In practice, the model should be sampled at the depth it was calibrated for, or calibrated at the depth used for inference, at a training cost that grows linearly with $T$.

Vector (PDF)

SBI benchmark

In contrast to the imaging experiments, the SBI benchmarks are low-dimensional problems for which standard NPE is already relatively well calibrated. Our objective is to establish that TARP-CvM can match a density-based calibration method while requiring only posterior samples. We evaluate our method on four problems from the SBI benchmark introduced in Hermans et al. [17], namely SLCP, Weinberg, Spatial SIR and Lotka–Volterra, each at three simulation budgets $\{2^{10}, 2^{13}, 2^{16}\}$. We compare conditional normalizing flows trained with Neural Posterior Estimation (NPE, Papamakarios and Murray, [34]; Lueckmann et al., [29]; Deistler et al., [9]; Greenberg et al., [16]), with the HPD-based regularizer of CalNPE (Falkiewicz et al., [13]), and with TARP-CvM. In SBI, we usually only have access to samples from the simulator, and cannot evaluate the adjoint operator of the forward model. In order to get a point estimate $b(y)$ close to the posterior for reference points when only pairs of $(\theta, y)$ are available, we train a multilayer perceptron (MLP) network targeting the posterior mean $y \mapsto E_{p(\theta\mid y)}[\theta]$, with a mean-square-error loss, and freeze it during calibration of the posterior generative model. All methods share the same training protocol, with the same simulation budgets, validation set and checkpoint selection (Appendix § 13.§ 2).

Under this matched protocol, NPE is already nearly calibrated on these low-dimensional problems. On the two TARP coverage maps, log posterior density and joint MMD on $(\theta, b(y))$, a single method is significantly better than the other two in only two of the 48 comparisons: TARP-CvM on the switch map for Weinberg at $2^{10}$, and CalNPE on log density for Spatial SIR at $2^{10}$ (Table 1, Appendix § 13.§ 2).

Problem	$n$	HPD error $\downarrow$ NPE	HPD error $\downarrow$ CalNPE	HPD error $\downarrow$ Ours	TARP error, $y$-indep. $\downarrow$ NPE	TARP error, $y$-indep. $\downarrow$ CalNPE	TARP error, $y$-indep. $\downarrow$ Ours	TARP error, switch $\downarrow$ NPE	TARP error, switch $\downarrow$ CalNPE	TARP error, switch $\downarrow$ Ours
SLCP	$2^{10}$	.096$\pm$.017	.086$\pm$.014	.101$\pm$.018	.015$\pm$.004	.012$\pm$.002	.015$\pm$.005	.013$\pm$.003	.012$\pm$.002	.014$\pm$.004
	$2^{13}$	.047$\pm$.013	.029$\pm$.007	.038$\pm$.013	.010$\pm$.003	.010$\pm$.004	.013$\pm$.004	.009$\pm$.003	.009$\pm$.003	.013$\pm$.004
	$2^{16}$	.013$\pm$.007	.006$\pm$.004	.018$\pm$.007	.007$\pm$.002	.008$\pm$.002	.006$\pm$.002	.007$\pm$.002	.008$\pm$.002	.006$\pm$.001
Weinberg	$2^{10}$	.021$\pm$.005	.010$\pm$.004	.031$\pm$.011	.011$\pm$.003	.010$\pm$.003	.014$\pm$.004	.047$\pm$.003	.050$\pm$.004	.041$\pm$.004
	$2^{13}$	.034$\pm$.005	.007$\pm$.005	.031$\pm$.007	.012$\pm$.002	.012$\pm$.004	.013$\pm$.003	.015$\pm$.002	.019$\pm$.004	.016$\pm$.003
	$2^{16}$	.016$\pm$.006	.008$\pm$.003	.013$\pm$.005	.009$\pm$.003	.008$\pm$.003	.009$\pm$.002	.009$\pm$.003	.011$\pm$.002	.008$\pm$.002
Spatial SIR	$2^{10}$	.081$\pm$.024	.045$\pm$.021	.093$\pm$.012	.029$\pm$.008	.027$\pm$.006	.034$\pm$.003	.041$\pm$.010	.041$\pm$.016	.027$\pm$.003
	$2^{13}$	.053$\pm$.021	.017$\pm$.006	.041$\pm$.016	.025$\pm$.006	.020$\pm$.006	.027$\pm$.004	.021$\pm$.004	.022$\pm$.005	.021$\pm$.005
	$2^{16}$	.031$\pm$.009	.026$\pm$.014	.031$\pm$.013	.020$\pm$.007	.011$\pm$.003	.016$\pm$.006	.015$\pm$.003	.011$\pm$.003	.015$\pm$.005
Lotka–Volterra	$2^{10}$	.082$\pm$.016	.064$\pm$.008	.081$\pm$.008	.016$\pm$.003	.015$\pm$.005	.013$\pm$.006	.067$\pm$.004	.067$\pm$.005	.067$\pm$.003
	$2^{13}$	.024$\pm$.014	.014$\pm$.006	.015$\pm$.008	.017$\pm$.005	.011$\pm$.005	.012$\pm$.005	.030$\pm$.005	.032$\pm$.006	.033$\pm$.004
	$2^{16}$	.014$\pm$.007	.010$\pm$.004	.013$\pm$.010	.014$\pm$.008	.014$\pm$.005	.012$\pm$.004	.014$\pm$.007	.017$\pm$.002	.013$\pm$.003
Problem	$n$	$\mathbb{E}[\log q_{\phi}] \uparrow$ NPE	$\mathbb{E}[\log q_{\phi}] \uparrow$ CalNPE	$\mathbb{E}[\log q_{\phi}] \uparrow$ Ours	MMD $z$, $(\theta,b(y)) \downarrow$ NPE	MMD $z$, $(\theta,b(y)) \downarrow$ CalNPE	MMD $z$, $(\theta,b(y)) \downarrow$ Ours	$|\mathrm{MIRA}-2/3|\downarrow$ NPE	$|\mathrm{MIRA}-2/3|\downarrow$ CalNPE	$|\mathrm{MIRA}-2/3|\downarrow$ Ours
SLCP	$2^{10}$	-2.34$\pm$.05	-2.32$\pm$.08	-2.32$\pm$.03	11.1$\pm$4.1	12.2$\pm$3.8	14.0$\pm$6.2	.0017$\pm$.0009	.0033$\pm$.0010	.0022$\pm$.0013
	$2^{13}$	-1.41$\pm$.08	-1.38$\pm$.02	-1.35$\pm$.06	5.5$\pm$2.7	6.3$\pm$3.0	5.6$\pm$2.4	.0034$\pm$.0010	.0040$\pm$.0014	.0027$\pm$.0016
	$2^{16}$	-0.87$\pm$.03	-0.88$\pm$.01	-0.86$\pm$.02	1.3$\pm$1.0	1.9$\pm$.6	2.1$\pm$1.1	.0045$\pm$.0004	.0051$\pm$.0008	.0048$\pm$.0006
Weinberg	$2^{10}$	-0.11$\pm$.01	-0.10$\pm$.01	-0.11$\pm$.01	56.3$\pm$6.2	59.3$\pm$13.6	58.2$\pm$8.9	.0050$\pm$.0017	.0068$\pm$.0015	.0028$\pm$.0015
	$2^{13}$	0.02$\pm$.01	0.02$\pm$.01	0.01$\pm$.01	8.0$\pm$4.8	14.9$\pm$9.1	8.1$\pm$4.3	.0014$\pm$.0008	.0039$\pm$.0010	.0013$\pm$.0008
	$2^{16}$	0.10$\pm$.00	0.09$\pm$.01	0.09$\pm$.01	3.8$\pm$3.1	7.1$\pm$5.3	5.3$\pm$3.4	.0034$\pm$.0012	.0071$\pm$.0009	.0043$\pm$.0006
Spatial SIR	$2^{10}$	0.82$\pm$.02	0.86$\pm$.03	0.77$\pm$.05	114.2$\pm$11.6	106.4$\pm$26.1	115.5$\pm$14.9	.0023$\pm$.0019	.0055$\pm$.0030	.0059$\pm$.0019
	$2^{13}$	1.31$\pm$.03	1.33$\pm$.02	1.26$\pm$.04	43.8$\pm$15.2	39.3$\pm$15.3	55.0$\pm$18.5	.0033$\pm$.0018	.0031$\pm$.0022	.0023$\pm$.0005
	$2^{16}$	1.63$\pm$.01	1.66$\pm$.03	1.64$\pm$.03	18.4$\pm$12.4	10.7$\pm$6.0	10.0$\pm$4.3	.0011$\pm$.0009	.0063$\pm$.0018	.0016$\pm$.0017
Lotka–Volterra	$2^{10}$	-2.59$\pm$.01	-2.57$\pm$.03	-2.60$\pm$.01	77.1$\pm$4.1	83.7$\pm$15.9	83.1$\pm$24.2	.0060$\pm$.0012	.0064$\pm$.0013	.0054$\pm$.0006
	$2^{13}$	-2.18$\pm$.02	-2.17$\pm$.02	-2.20$\pm$.02	38.1$\pm$13.7	29.9$\pm$2.4	30.4$\pm$6.2	.0065$\pm$.0012	.0071$\pm$.0012	.0065$\pm$.0009
	$2^{16}$	-1.98$\pm$.02	-1.99$\pm$.02	-2.00$\pm$.01	18.3$\pm$16.1	23.0$\pm$15.4	18.3$\pm$11.7	.0053$\pm$.0012	.0068$\pm$.0016	.0050$\pm$.0016

TARP-CvM thus matches CalNPE using only posterior samples, without evaluating the density. Its log density does not differ significantly from that of NPE in any setting, but is lower than that of CalNPE on Spatial SIR at the two smaller budgets. CalNPE achieves a lower HPD coverage error, the statistic it optimizes, with a significant difference in four settings. On the independent Mira score, TARP-CvM is closest to the calibrated value $2/3$ in half of the settings, and CalNPE deviates most from it in 10 of 12 settings, consistent with its conservativeness objective, a one-sided penalty on under-coverage that leaves over-coverage unpenalized, which we used to match the configuration reported by Falkiewicz et al. [13].

None of these differences is significant.

Conclusion

We introduced TARP-CvM, a differentiable coverage loss that vanishes if and only if the learned conditional distribution equals the true posterior. It requires only samples from the model, applies to any generative architecture, and is a strict divergence (Theorem 1). On SBI benchmarks, it matches a density-based coverage loss with the same number of model evaluations, without evaluating the density. On imaging inverse problems, including one with a nonlinear black-box forward model, it is the only regularizer that brings the posterior into the Mira null band on all three problems, with little change in the quality of point estimates. The main practical design choice is the construction of reference points, which are built around a point estimate $b(y)$. For linear forward models, the adjoint $A^{\mathsf{T}}y$ works well. Without an adjoint, an auxiliary regressor of the posterior mean can be trained at extra cost, as in our SBI experiments. Future work includes other nonlinear inverse problems, such as full waveform inversion and black hole imaging, and SBI problems in cosmology, where the forward model is a physics simulator and posteriors are often multimodal.

AI use statement

In this work, we used generative AI tools for formulating mathematical claims, assisting in the writing and checking of proofs, providing feedback on the design of the method and experiments, and assisting in implementing the method and experiments. Additionally, we used generative AI tools to edit and rewrite the manuscript for clarity. We have reviewed all AI-assisted work. All theorems and proofs were checked line by line by the authors. AI-generated code was reviewed and tested by the authors. Results were interpreted and conclusions drawn by the authors, and all edited text was reviewed by the authors. We take responsibility for the final content of this work, including text, claims, or artifacts produced with the aid of generative AI.

Acknowledgments

All authors gratefully acknowledge the support of the NSF-Simons AI-Institute for the Sky (SkAI) via grants NSF AST-2421845 and Simons Foundation MPS-AI-00010513. RW further gratefully acknowledges the support of the NSF-Simons National Institute for Theory and Mathematics in Biology (NITMB), via grants NSF DMS-2235451 and Simons Foundation MP-TMPS-00005320, and AFOSR FA9550-18-1-0166. This research used the DeltaAI advanced computing and data resource, which is supported by the National Science Foundation (award OAC 2320345) and the State of Illinois. DeltaAI is a joint effort of the University of Illinois Urbana-Champaign and its National Center for Supercomputing Applications.

Proofs

This section proves Theorem 1. Section A.1 defines the support of a distribution over rules, Section A.2 recalls the notation, and Section A.3 collects preliminary facts: the regularity of the radial mass function, the calibration of the true posterior for every rule, and a Cramér–Wold argument showing that integrals over half-spaces determine a function of the measurement. The direction “$q_{\phi}=p$ implies $D_{\rho}=0$” holds for any $\rho$ and follows from Lemma 1. The converse is proved in three steps, one per subsection:

1. if $D_{\rho}=0$, every rule in the support of $\rho$ is calibrated (Section A.4)

2. by patching (R2), the coverage at each measurement then does not depend on the reference point (Section A.5)

3. since reference points are dense (R1), this invariance forces $p(\cdot\mid y)=q_{\phi}(\cdot\mid y)$ (Section A.6).

Section A.7 assembles the three steps.

Rules, their topology and the support of $\rho$

A reference rule is a measurable map $\psi:\mathcal{Y}\to\Theta$. We define how far apart two rules $\psi$ and $\psi'$ are by

the metric of convergence in probability. Two rules are close if their reference points are close for most measurements. The support of a distribution $\rho$ over rules is the set of rules that $\rho$ approaches arbitrarily closely with positive probability: $\psi$ is in the support if, for every $\nu>0$, a rule drawn from $\rho$ is within distance $\nu$ of $\psi$ with positive probability,

A rule in $\operatorname{supp}\rho$ need not be drawn with positive probability. For the switch sampler, any particular rule has probability zero, but rules arbitrarily close to it are drawn with positive probability.

Conventions

Throughout the appendix we write $q$ for $q_{\phi}$. The distance $d$ on $\Theta$ is induced by a norm, $B(c,t)=\{\theta:d(\theta,c)\le t\}$ is the closed ball, and $\operatorname{vol}$ is the Lebesgue measure on $\mathbb{R}^{d_{\theta}}$. By Assumption 1, $\Theta$ is open and, for every $y$, $p(\cdot\mid y)$ and $q(\cdot\mid y)$ have continuous densities that are strictly positive on $\Theta$. We recall the objects of the main text. For a reference point $c\in\Theta$, the radial mass function of the estimator is

For a rule $\psi$, the credibility value is $s_{\psi}(\theta,y)=m_{\psi(y)}(d(\theta,\psi(y))\mid y)$, the conditional coverage is $F_{\psi}(\alpha\mid y)=\mathbb{P}_{\theta\sim p(\cdot\mid y)}[s_{\psi}(\theta,y)\le\alpha]$, and the expected coverage is $F_{\psi}(\alpha)=\int F_{\psi}(\alpha\mid y)p(y)\,dy$. For a constant rule $\psi\equiv c$ we write $F_c(\alpha\mid y)$. Finally,

Two remarks are used repeatedly. First, the conditional coverage at $y$ depends on the rule only through its reference point at $y$: $F_{\psi}(\alpha\mid y)=F_{\psi(y)}(\alpha\mid y)$. Second, $\alpha\mapsto F_{\psi}(\alpha\mid y)$ and $\alpha\mapsto F_{\psi}(\alpha)$ are cumulative distribution functions, hence non-decreasing and right-continuous.

Preliminary facts

We restate and prove Lemma 1.

Under Assumption 1, if $q(\cdot\mid y) = p(\cdot\mid y)$ for every $y$, then $F_{\psi}(\alpha\mid y) = \alpha$ for every rule $\psi$, every level $\alpha\in[0,1]$ and every $y$, and hence $F_{\psi}(\alpha) = \alpha$.

Proof. A correct estimator is calibrated, no matter how the regions are chosen. Let $R_{\alpha}$ be any credible region map of $q$. Factoring the joint density as $p(\theta,y) = p(y)\,p(\theta\mid y)$ and exchanging the order of integration, we obtain

where the fourth equality uses $p(\cdot\mid y) = q(\cdot\mid y)$, and the fifth uses the definition of a credible region. The inner integral equals $\alpha$ for every $y$, so the statement holds at each measurement, before averaging over $y$. For the TARP region of a rule $\psi$, the event $\theta\in R_{\alpha,\psi}^{\mathrm{TARP}}(y)$ is the event $s_{\psi}(\theta,y) \leq\alpha$ (§ 3.2): the credibility value of $\theta$ is the mass of the ball where $\theta$ is at the border, so it is at most $\alpha$ exactly when $\theta$ lies in the ball of mass $\alpha$. Hence $F_{\psi}(\alpha\mid y) = \mathbb{P}_{\theta\sim p(\cdot\mid y)}[\theta\in R_{\alpha,\psi}^{\mathrm{TARP}}(y)] = \alpha$ for every $y$, and averaging over $p(y)$ gives $F_{\psi}(\alpha) = \alpha$. □

Letting $n \to\infty$ and then $\kappa\to\infty$ shows that the left-hand side is zero. The same argument applies to sin, so $\int e^{i\lambda\langle u,y\rangle}g(y)\,dy=0$.

Conclusion. Every vector $v \in\mathbb{R}^{d_y}$ is of the form $\lambda u$, so the Fourier transform $\hat{g}(v)=\int e^{i\langle v,y\rangle}g(y)\,dy$ vanishes everywhere. An integrable function whose Fourier transform vanishes is zero almost everywhere, so $f(y)p(y)=0$ for almost every $y$, that is, $f(y)=0$ for $p(y)$-almost every $y$. This is the Cramér–Wold argument [8].

Half-spaces are used only through Proposition 1: any other family of sets for which this proposition holds can replace them in (R2).

From zero loss to calibration on the support

Proof. Let $\psi_n \to\psi$ for $d_0$, that is, in $p(y)$-probability. It suffices to show that every subsequence has a further subsequence along which $I_{\psi_n} \to I_\psi$. Convergence in probability gives a subsequence, still denoted $\psi_n$, such that $\psi_n(y) \to\psi(y)$ for $p(y)$-almost every $y$. For such $y$ and every $\theta$,

by continuity of $d$ and joint continuity of the radial mass function (Lemma 2). So $s_{\psi_n}\to s_\psi$ almost surely under $p(\theta,y)$, hence in distribution, and $F_{\psi_n}(\alpha)\to F_\psi(\alpha)$ at every continuity point of $F_\psi$, that is, for all but countably many $\alpha$. Since $\lvert F_{\psi_n}(\alpha)-\alpha\rvert\le1$, bounded convergence gives $I_{\psi_n}\to I_\psi$.

Proof. Since $I_\psi\ge0$ and $\mathbb{E}_{\psi\sim\rho}[I_\psi]=0$, we have $I_\psi=0$ for $\rho$-almost every $\psi$. The set $Z=\{\psi:I_\psi=0\}$ is closed by Lemma 3, and $\rho(Z)=1$. If some $\psi\in\operatorname{supp}\rho$ were outside $Z$, a small ball around $\psi$ would lie in the open complement of $Z$ and have positive probability under $\rho$, a contradiction. Hence $I_\psi=0$ on $\operatorname{supp}\rho$, so $F_\psi(\alpha)=\alpha$ for almost every $\alpha$, and for every $\alpha$ by right-continuity of both sides.

From calibration to invariance to the reference point

Every rule in $\operatorname{supp}\rho$ is now calibrated, but calibration is an average over measurements, in which errors at different measurements can cancel. This step obtains a statement at each measurement: the coverage curve at $y$ does not depend on the reference point. As in the main-text example, it compares a rule $\psi'$ with the rule patched from $\psi$ on a half-space, which isolates what happens on that half-space.

Step 2: both rules are calibrated. The rule $\psi'$ is in $\operatorname{supp}\rho$ because $\Psi\subseteq\operatorname{supp}\rho$, and the patched rule $\psi\oplus_H\psi'$ is in $\operatorname{supp}\rho$ by (R2). Both are calibrated, so both coverages equal $\alpha$ and the left-hand side of Step 1 is zero:

Step 3: from half-spaces to each measurement. By Proposition 1, applied to $f=\Delta$, we get $\Delta(y)=0$ for $p(y)$-almost every $y$. In other words, $F_{\psi}(\alpha\mid y)=F_{\psi'}(\alpha\mid y)$ for every measurement $y$ except those in a set $N_{\psi,\psi',\alpha}$ that has probability zero under $p(y)$.

Step 4: all pairs and all levels at once. The exceptional set of Step 3 depends on $\psi$, $\psi'$ and $\alpha$. Let $N$ be the union of these sets over all pairs $\psi,\psi'\in\Psi$ and all rational levels $\alpha\in[0,1]\cap\mathbb{Q}$. There are countably many of them, so $N$ still has probability zero. For $y\notin N$, the two coverage curves agree at every rational level. Both are distribution functions, hence right-continuous, and every $\alpha\in[0,1)$ is a limit of rationals from the right, so they agree at every level; at $\alpha=1$ both equal one.

From invariance to the reference point to equal densities

Lemma 6. Fix $y$ such that the reference points $\{\psi(y):\psi\in\Psi\}$ are dense in $\Theta$ and $F_{\psi}(\cdot\mid y)$ is the same function $G_y$ for all $\psi\in\Psi$. Then $p(\cdot\mid y)=q(\cdot\mid y)$.

Proof. Let $c=\psi(y)$ for some $\psi\in\Psi$, write $m=m_c(\cdot\mid y)$, and take $t_c$ from Lemma 2. For $0<t<t_c$, $m$ is strictly increasing on $[0,t_c]$ and non-decreasing beyond, so $m(D)\le m(t)$ if and only if $D\le t$. Applying this to $D=d(\theta,c)$ with $\theta\sim p(\cdot\mid y)$,

Since $m(t)=\mathbb{P}_{\theta\sim q(\cdot\mid y)}[\theta\in B(c,t)]$, dividing gives

As $t\downarrow0$, both densities are continuous and positive at $c$, so each probability equals the density at $c$ times $\operatorname{vol}(B(c,t))$, up to a factor $1+o(1)$. The volumes cancel and the right-hand side tends to $p(c\mid y)/q(c\mid y)$. On the left, $m$ is a continuous increasing bijection from $(0,t_c)$ onto $(0,m(t_c))$, so $\alpha\downarrow0$ exactly when $t\downarrow0$. Therefore $L=\lim_{\alpha\downarrow0}G_y(\alpha)/\alpha$ exists and

The limit $L$ depends on $G_y$ only, not on $c$, so this identity holds at every reference point $c=\psi(y)$, $\psi\in\Psi$, a dense subset of $\Theta$. Both densities are continuous, so it holds on all of $\Theta$. Integrating over $\Theta$ gives $L=1$, hence $p(\cdot\mid y)=q(\cdot\mid y)$.

Proof of Theorem 1

Under Assumption 1 and for any admissible $\rho$,

Proof. $(\Leftarrow)$ By Lemma 1, $F_{\psi}(\alpha)=\alpha$ for every rule $\psi$ and every $\alpha$, so $I_{\psi}=0$ for every $\psi$ and $D_{\rho}=0$. This direction holds for any $\rho$.

$(\Rightarrow)$ Suppose $D_{\rho}=0$. By Lemma 4, every rule in $\operatorname{supp}\rho$ is calibrated. By Lemma 5, there is a null set $N_2$ outside which $F_{\psi}(\cdot\mid y)$ is the same function for all $\psi\in\Psi$. By (R1), there is a null set $N_1$ outside which the reference points $\{\psi(y):\psi\in\Psi\}$ are dense in $\Theta$. For every $y\notin N_1\cup N_2$, Lemma 6 applies and gives $p(\cdot\mid y)=q(\cdot\mid y)$. Since $N_1\cup N_2$ is null, this holds for $p(y)$-almost every $y$.

Training algorithm

Algorithm 1 summarizes one training step with the TARP-CvM regularizer. The base loss $\mathcal{L}_{\mathrm{base}}$ is the training objective of the generative model, such as denoising score matching for diffusion models or the negative log-likelihood for normalizing flows.

Complexity For each rule, computing the $N$ credibility values requires $NM$ distance evaluations, and the closed-form statistic of Equation (2) only requires sorting them, at cost $O(N \log N)$. Its gradient follows analytically from the ranks, so no differentiable sorting operator is needed. One step therefore costs $O(RN(M+\log N))$ operations on top of drawing the $NM$ posterior samples. In practice, the cost is dominated by sampling. For a diffusion model with a $T$-step sampler, a step requires $NMT$ network evaluations, and backpropagating through the chains that carry gradient multiplies their cost, and the memory they require, by a factor proportional to $T$. Training cost thus grows linearly with the sampler depth, which is why we train with $T=20$ (§ 5.2, Appendix § 13.§ 3).

Unbiased estimator of the Cramér–von Mises statistic

For a fixed rule $\psi$, the inner integral of Equation (1) is

Given the credibility values $s_1,\ldots,s_N$ of a batch, Equation (2) replaces $F_\psi$ by the empirical CDF $\widehat{F}_\psi(\alpha)=\frac{1}{N}\sum_i \mathbf{1}_{s_i\leq\alpha}$, so that $\operatorname{CvM}(s_{1:N})=\int_0^1\left(\widehat{F}_\psi(\alpha)-\alpha\right)^2\,\mathrm{d}\alpha$. Writing $g_i(\alpha)=\mathbf{1}_{s_i\leq\alpha}-\alpha$, we have $\widehat{F}_\psi(\alpha)-\alpha=\frac{1}{N}\sum_i g_i(\alpha)$, and

This is an average over all $N^2$ pairs $(i,j)$, including the $N$ diagonal pairs $i=j$, that is, a V-statistic ([32]). The U-statistic ([19]) $\operatorname{CvM}_{U}$ of Equation (4) averages $h$ over the $N(N-1)$ pairs with $i\ne j$ only.

Proof. Fix a level $\alpha$. For $i\ne j$, the values $s_i$ and $s_j$ are independent, so

which is the integrand of $I_\psi$. For $i=j$, we obtain the second moment instead,

By Fubini’s theorem, $\mathbb{E}[h(s_i,s_j)]$ is the integral over $\alpha$ of these expectations. Averaging over the $N(N-1)$ off-diagonal pairs gives $\mathbb{E}[\operatorname{CvM}_{U}]=I_\psi$, and averaging over all $N^2$ pairs, of which $N$ are diagonal, gives the expectation of $\operatorname{CvM}$.

It remains to check that Equation (4) is the off-diagonal average of $h$. Using $\int_0^1\mathbf{1}_{\max(s,s')\leq\alpha}\,\mathrm{d}\alpha=1-\max(s,s')$, $\int_0^1\alpha\mathbf{1}_{s\leq\alpha}\,\mathrm{d}\alpha=(1-s^2)/2$ and $\int_0^1\alpha^2\,\mathrm{d}\alpha=1/3$,

Summing over the ordered pairs $i\ne j$, the first term sums to $(N-1)\sum_i s_i^2$, and the $k$-th smallest value $s_{(k)}$ is the maximum of the $2(k-1)$ ordered pairs it forms with the smaller values. Dividing by $N(N-1)$ gives Equation (4). ▫

Note that the bias term does not vanish at the posterior. When $F_{\psi}(\alpha)=\alpha$, it equals $\frac{1}{N}\int_{0}^{1}\alpha(1-\alpha)\,d\alpha=\frac{1}{6N}$. It also shifts the minimizer of the expected loss.

Scoring rules and loss landscapes

Definitions

The continuous ranked probability score (CRPS) compares the cumulative distribution function (CDF) of the model with the CDF of a point mass at the observation. Part of its appeal is that CDFs are easy to estimate from an ensemble of generated samples. For models that are easy to sample from but have no closed-form CDF, the CRPS admits the energy form

The first term penalizes a model whose samples are far from the true $\theta$ on average, and the second rewards samples that are spread out, so that a model cannot minimize the CRPS by collapsing to a point prediction. Unlike CDFs, this energy form generalizes directly to several dimensions by replacing the absolute value with the Euclidean norm, which gives the energy score of § 2. On the class of distributions with finite first moment, the energy score is strictly proper:

with equality if and only if $q_{\phi}(\cdot\mid y)=p(\cdot\mid y)$. Like CRPS, it can be evaluated from samples of a generative model alone.

In high dimension, the CRPS is often computed for each coordinate separately rather than replaced by the energy score. We call this the pixel-wise CRPS,

where the superscript $j$ denotes restriction to the $j$-th coordinate. Its motivation is often computational, since it parallelizes across pixels. It is however minimized by any model with the correct one-dimensional marginals, and therefore does not constrain the dependence between coordinates.

Strict propriety and training

Strict propriety is a statement about the global minimizer of the population loss, and two gaps separate it from what training achieves. First, the loss is minimized over a finite training set, and the minimizer of this empirical approximation need not be the minimizer of the population loss. Second, gradient descent generally finds local rather than global minima. Two strictly proper scoring rules share the same global minimizer, but they can have very different loss landscapes, with different local minima, so that training with one can yield much better coverage than training with the other. Section D.3 illustrates this on a Gaussian model.

Loss landscapes

We examine the same setting as the inverse problem in Figure 1, setting $\theta$ to have dimension 8 and $Y$ dimension 4.

We identify the true parameters of the posterior distribution $\theta\mid Y$ and perturb them by deviating the mean in the direction of largest variance and the scale of the variance in that same direction. See the equation below for this perturbation model:

where $a,b$ are the mean and scale perturbation parameters, respectively. $h(a)=a^3-3a-4.048$ parameterizes the mean perturbation and vanishes at $a=2.2$. The scale perturbation vanishes at $b=0$. We see in Figure 5 the loss landscapes for Energy Score and TARP-CvM in terms of $a$ and $b$.

Vector (PDF)

In Figure 6, we have slices of the loss functions at the local and global minima of $a$, the mean parameter. We see that TARP-CvM has much more sensitivity to changing the spread parameter, $b$, when $a$ is at the local minimum than energy score does. In other words, we are comparing the steepness of the blue lines in each plot.

Vector (PDF)

Intuitively, when one gets stuck in a local minimum for the mean parameter, TARP-CvM recovers correct coverage by pushing the conditional distribution toward wider spread to compensate for the bias in the mean. In contrast, the energy score’s loss function is relatively flat with respect to the spread parameter in this local optimum. So, when noisy gradients provide imperfect optimization of this population loss, the energy score is less likely to correct coverage when gradient descent is near a local minimum for the conditional mean, since the signal in the loss is too flat to overcome gradient noise.

Experiments

Toy model

Setup. The parameters have a Gaussian prior $\theta\sim\mathcal{N}(0, \Sigma_0)$, and the measurements follow $y = A\theta+ \epsilon$, with $\epsilon\sim\mathcal{N}(0, \sigma^2 I)$, $d_y = 4$, and $d_\theta= 8$, so that $A$ cannot be inverted. The posterior is Gaussian,

where $\Sigma_y = A\Sigma_0 A^{\mathsf{T}} + \sigma^2 I$ is the covariance of the measurements. The estimator is $q_\phi(\theta\mid y) = \mathcal{N}(Wy, S)$, where $\phi= (W, S)$ are parameters to fit.

HPD calibration does not identify the posterior. The density of $q_\phi(\cdot\mid y)$ is a decreasing function of the Mahalanobis norm $m(\theta) = (\theta- Wy)^{\mathsf{T}}S^{-1}(\theta- Wy)$, so its HPD regions are the ellipsoids $\{\theta: m(\theta) \le t_\alpha\}$ with mass $\alpha$ under $q_\phi(\cdot\mid y)$. Under $p(\theta, y)$, the residual $e = \theta- Wy$ decomposes as $(\theta- W^*y) + (W^* - W)y$. The first term has distribution $\mathcal{N}(0, S^*)$ given $y$, hence is independent of $y$, and the second has covariance $(W^* - W)\Sigma_y(W^* - W)^{\mathsf{T}}$, so that marginally

If $S = C_W$, the residual of the true parameter has the same distribution as $\tilde{\theta} - Wy$ for a draw $\tilde{\theta} \sim q_\phi(\cdot\mid y)$. The true parameter then falls in each HPD region with probability $\alpha$, so the estimator is HPD-calibrated, for every $W$.

Points of the manifold. At $W = W^*$, $S = S^*$ is the true posterior. At $W = 0$, $S = S^* + W^*\Sigma_yW^{*\mathsf{T}} = S^* + W^*A\Sigma_0 = \Sigma_0$, so the prior is HPD-calibrated. For every $W \ne W^*$, the mean $Wy$ is biased and $S - S^* = (W^* - W)\Sigma_y(W^* - W)^{\mathsf{T}}$ is positive semi-definite and nonzero, so the estimator is wider than the posterior at every measurement: marginal calibration is reached by trading the bias of the mean for extra spread.

Simulation Based Inference

Problems. We use the simulators of Hermans et al. [17] with the problem settings of Falkiewicz et al. [13]. For each problem we generate up to $2^{16}$ simulations for training, with $10^4$ for validation, and $10^4$ held-out pairs for evaluation.

• SLCP (simple likelihood, complex posterior) is used in the marginal variant of Hermans et al. [17], in which only two parameters are inferred and the likelihood becomes intractable. The parameters are $\theta\in[-3, 3]^2$. The observation $y \in\mathbb{R}^8$ consists of four draws from a bivariate Gaussian with fixed mean $(0.7, -2.9)$, standard deviations $\theta_1$ and $\theta_2$, and correlation $\tanh(\rho)$, where nuisance parameter $\rho\sim\mathcal{U}(-3, 3)$ is marginalized out. Since $\theta$ enters only through its squares, the posterior has four symmetric modes.

• Weinberg infers the Fermi constant, normalized to its nominal value, $\theta\in[0.25, 2]$, from $y \in[-1, 1]^{20}$. The observation contains the cosines of the scattering angle of twenty $e^+e^- \to\mu^+\mu^-$ events drawn by rejection sampling from the differential cross-section at a beam energy of 40 GeV.

• Spatial SIR infers the infection and recovery rates $\theta= (\beta, \gamma) \in[0, 1]^2$ of a stochastic epidemic on a $50 \times50$ lattice. The epidemic is seeded with $1 + \operatorname{Poisson}(3)$ infected cells at random positions and propagated through $3 \times3$ neighborhoods. The observation $y \in\{0, 1\}^{3 \times50 \times50}$ records the susceptible, infected and recovered state of every cell at time $t = 1$.

• Lotka–Volterra is a predator-prey Markov jump process started from 50 predators and 100 prey. The parameters $\theta\in[-4, 1]^2$ are the log-rates of the two predator reactions. The two prey log-rates are drawn from the same prior and marginalized out. The observation $y \in\mathbb{R}^{2 \times1001}$ is the trajectory of both populations.

Architecture and training. Neural posterior estimation (NPE, Papamakarios and Murray, [34]; Lueckmann et al., [29]; Deistler et al., [9]; Greenberg et al., [16]) fits a conditional density estimator $q_\phi(\theta\mid y)$ by maximum likelihood on simulated pairs, minimizing $-(1/N)\sum_i \log q_\phi(\theta_i \mid y_i)$. This is equivalent to minimizing the forward KL divergence to the posterior, averaged over $p(y)$. We use the architecture of Falkiewicz et al. [13], implemented with the SBI library (Tejero-Cantero et al., [42]). It is a conditional neural spline flow (Durkan et al., [12]), with rational-quadratic spline transforms (10 bins, hidden width 64) and a Gaussian base distribution. There are three transforms, each a coupling layer followed by a linear layer, except for the one-dimensional Weinberg problem, which uses a single transform whose spline parameters depend on the conditioning only. The flow is conditioned on an embedding of $y$ given by an MLP with three hidden layers of width 64, SELU activations and a 10-dimensional output. The 7500-dimensional Spatial SIR observations and the 2002-dimensional Lotka-Volterra trajectories are flattened before embedding.

Parameters and observations are standardized per coordinate using the training data. We train with Adam, using a learning rate of $10^{-3}$, a batch size of 128, gradient-norm clipping at 5, for 50 epochs without early stopping. We keep the checkpoint with the lowest validation loss on the $10^4$ validation simulations. All three objectives of Table 1 share this training recipe.

For CalNPE, we use the conservativeness objective of Falkiewicz et al. [13], which penalizes only coverage below the nominal level, with weight $\lambda= 100$, the default of their released code, instead of the published $\lambda= 5$, which gave worse results in our experiments. HPD credibilities are estimated from the density of the estimator at 16 draws from the prior, with straight-through gradients for the indicators and differentiable sorting (Blondel et al., [4]).

TARP-CvM uses weight $\lambda= 30$ and $M = 16$ reparameterized posterior draws per pair, so that both regularizers use 16 flow evaluations per pair: CalNPE evaluates the density at 16 prior draws, TARP-CvM draws 16 posterior samples. Each step uses $R = 48$ references in three groups of 16, described below. The indicators are relaxed with a sigmoid at temperature $\tau= 0.05$ times the median sample-to-reference distance of each reference, treated as a constant, and the loss is the closed-form CvM statistic of Equation (2), computed per reference and averaged over the $R$ references.

Reference points without a forward operator Without access to the forward operator to estimate $b(y)$, we resort to regressing the posterior mean. The point estimate $b(y)$ is an MLP (two hidden layers of width 128, SELU) trained beforehand on the flow’s own training pairs with a mean-squared-error loss, with Adam (learning rate $10^{-3}$, weight decay $10^{-4}$, batch 256). The $R = 48$ references of a step form three groups of 16: (1) switch-rule references, (2) points drawn from the prior distribution and shared by the batch, and (3) the posterior draws of each pair, detached from the computation graph. The credibility of a posterior reference is computed from the 15 other draws, so that $s_i$ remains uniform when $q_\phi= p$, by exchangeability.

Imaging inverse problems

Problems We selected three inverse problems introduced in InverseBench (Zheng et al., [48]). They all have the form $y = \mathcal{A}(\theta) + \varepsilon$, where $y$ are measurements, $\mathcal{A}$ is a forward model, $\theta$ is the signal we seek to reconstruct and $\varepsilon$ additive noise.

• Navier–Stokes (NS). $\theta$ is the initial vorticity of a 2D incompressible flow on $(0, 2\pi)^2$ with periodic boundary conditions (128 × 128 pixels, Reynolds number 200, forcing $f(x) = -4\cos(4x_2)$), obtained by evolving a draw from $\mathcal{N}(0, (-\Delta+ 9I)^{-4})$ for five time units. The forward operator is $\mathcal{A} = P \circ L$, where the solution operator $L$, computed with a black-box pseudo-spectral solver, maps $\theta$ to the vorticity at time $T = 1$ and is nonlinear. The subsampling operator $P$ keeps one grid point out of two in each direction, giving 64 × 64 noise-free observations.

• Linear inverse scattering (IS). $\theta$ is a 128 × 128 permittivity contrast map of HL60 cell nuclei, generated with the simulator of Wiesner et al. [43]. Under the first Born approximation, $\mathcal{A} = A$ is linear, with 20 transmitters and 360 receivers on a circle of radius 1.6 m, and noise level $\sigma_y = 10^{-4}$, resampled at every training batch.

• Compressed sensing MRI (MRI). $\theta$ is a complex 320 × 320 knee image from fastMRI (Zbontar et al., [46]), a coil-combined minimum-variance unbiased estimate (Jalal et al., [22]), stored as two real

Table 1: SBI benchmark results for NPE, CalNPE ($\lambda= 100$, Falkiewicz et al. [13]) and TARP-CvM (Ours, $\lambda= 30$, $M = 16$), all trained with the same protocol, for simulation budgets $n \in\{2^{10}, 2^{13}, 2^{16}\}$. Mean $\pm$ standard deviation over five independently trained estimators per cell. Bold marks the best mean when it is significantly better than both other methods. HPD and TARP errors are the $L^{1}$ coverage errors under each credible-region map, $\mathbb{E}[\log q_{\phi}]$ is the expected log posterior density of the true parameters, MMD $z$ the $z$-score of an MMD test on the pairs $(\theta, b(y))$, and $|\mathrm{Mira} - 2/3|$ the deviation of the Mira score from its calibrated value.

channels. Since these images have no coil sensitivity maps, we use the single-coil operator $\mathcal{A} = PF$, with $F$ the orthonormal 2D Fourier transform and $P$ a $\times4$ subsampling mask that keeps 80 of the 320 vertical k-space lines (25 central lines and 55 random outer lines, redrawn for every example). The measurements are noise-free.

We did not select the Black hole imaging problem because no training set of images is publicly available, and we did not select Full waveform inversion, because it did not fit the shared setup of conditioning architecture with fixed image-shape. Our method is nonetheless applicable to these problems as well.

Noise-free measurements On MRI and NS, measurements are noise-free, so the posterior is supported on the set $\{\theta: \mathcal{A}(\theta) = y\}$ and has no density on $\mathbb{R}^{d_{\theta}}$, which Assumption 1 requires. For the linear MRI operator, writing $\theta= A^{+}y + z$ with $z = (I - A^{+}A)\theta$, the range component $A^{+}y$ is fixed by the measurement, and the posterior is a distribution over the null-space component $z$, which has a density on the null space of $A$. Restricted to this subspace, the null-space distance $d(\theta,\theta') = \lVert(I-A^{+}A)(\theta-\theta')\rVert$ is a norm, so Theorem 1 applies to the posterior of $z$ given $y$, with $\Theta$ replaced by the null space. This is why TARP-CvM uses this distance on MRI. On NS, the solution set of the nonlinear operator has no such linear parametrization, and Theorem 1 does not strictly apply. There, the experiment tests the loss beyond its assumptions, and the Mira score, which requires no density, still measures coverage. On IS, the measurement noise gives the posterior a density, as Assumption 1 requires.

Diffusion models Diffusion models ([39] (Ho et al. 2020)) map the target distribution $p(x)$ to white Gaussian noise by progressively perturbing the data $x$ through a diffusion process. Following Song et al. [41], we express the continuous diffusion process as a stochastic differential equation (SDE):

where $x_t \in\mathbb{R}^{N}$ is the perturbed data at time $t$, $f(x_t,t) \in\mathbb{R}^{N}$ is the drift coefficient, $g(t) \in\mathbb{R}_{+}$ is the diffusion coefficient, and $w_t \in\mathbb{R}^{N}$ is a standard Wiener process. This SDE admits a reverse process, mapping noise to the target distribution:

where $\bar{w}$ is a standard Wiener process when time flows backward. Among the different kinds of diffusion processes, we used the Variance Exploding schedule, which parametrizes the noising kernel $x_t \sim\mathcal{N}(\alpha_t x,\sigma_t^2 I_N)$, with $\alpha_t = 1$, $\log\sigma_t = (1-t)\log\sigma_{\min} + t\log\sigma_{\max}$, with $\sigma_{\min}=10^{-3}$ and $\sigma_{\max}=10^2$. We also adopt the preconditioning strategy of Karras et al. [23], which parametrizes a denoising network $d_\phi$ instead of the score function.

Our target distribution is the conditional distribution of $x = \theta\mid y$, where $y$ are measurements and $\theta$ is the signal we aim to recover. The measurement $y$ enters as a fixed image, concatenated channel-wise to the noisy input of the network. For both MRI and IS, we compute the adjoint solution to transform measurement into images $A^\top y$. NS observations are already $64 \times64$ images, the resolution of the signal after $2 \times2$ patching, so they are concatenated to the patched input without resampling. NS has no linear adjoint, so the switch rules use the nearest-neighbor upsampled observation as the point estimate $b(y)$. We describe all our choices for the network architectures and training strategies in Table 2. All fine-tuning runs share the calibration batch (16 pairs per step on NS, 8 otherwise), the sampler depth ($T = 20$), the optimizer, the number of steps, and the sampler chains that carry gradients ($M_{\mathrm{grad}} = 2$ per pair on NS and MRI, 8 on inverse scattering). $\lambda$ is tuned once on Navier–Stokes and reused for IS and MRI ($\lambda= 10$ for TARP-CvM and $\lambda= 0.003$ for ES and pCRPS).

	Navier–Stokes	Inverse scattering	MRI
Data
Signal $x$ shape	(128, 128, 1)	(128, 128, 1)	(320, 320, 2)
Conditioning	$y$, (64, 64, 1)	$A^{\mathsf{T}}y$, (128, 128, 1)	$[A^{\mathsf{T}}y,\ \mathrm{mask}]$, (320, 320, 3)
Train / validation images	19,456 / 1,024	8,500 / 1,500	22,964 / 2,048
Architecture
Architecture	U-Net	U-Net	U-Net
Patch size	2	2	4
Input shape	(64, 64, 5)	(64, 64, 8)	(80, 80, 80)
Residual blocks per level	(2, 2, 2, 2)	(2, 2, 2, 2)	(3, 3, 3, 3, 3)
Channels per level	(64, 128, 256, 256)	(64, 128, 256, 256)	(128, 256, 384, 512, 512)
Kernel size	3	3	3
Activation	SiLU	SiLU	SiLU
Normalization	LayerNorm	LayerNorm	LayerNorm
Noise embedding	256	256	256
Dropout	0.1	0.1	0.1
Parameters	16.9 M	16.9 M	106.9 M
Diffusion
Noise schedule	VE, $\sigma\in[10^{-3}, 10^{2}]$, EDM preconditioning
Training time distribution	$t \sim\operatorname{Beta}(3, 3)$
Sampler	DDIM ($\eta= 0$), 20 steps
Pretraining (DSM)
Optimizer	Adam	Adam	Adam
Weight decay	0.0	0.0	0.0
Learning rate	$10^{-4}$	$10^{-4}$	$10^{-4}$
Gradient norm clipping	1.0	1.0	1.0
EMA decay	0.999	0.999	0.999
Augmentation	none	none	h-flip, roll $\pm4$ px
Batch size	256	512	256
Epochs (2 passes each)	560	580	204
Optimizer steps	85,120	18,560	36,312
Early-stopping patience	10	10	20
Fine-tuning
Optimizer	AdamW	AdamW	AdamW
Weight decay	0.0	0.0	0.0
Learning rate	$2 \times10^{-5}$	$2 \times10^{-5}$	$2 \times10^{-5}$
Gradient norm clipping	1.0	1.0	1.0
EMA decay	0.999	0.999	0.999
DSM batch size	64	64	64
Optimizer steps	24,000	24,000	24,000
Calibration weight $\lambda$, ES and pCRPS	0.003	0.003	0.003
Calibration weight $\lambda$, TARP-CvM	10	10	10
Calibration batch $N$	16	8	8
Posterior samples $M$ (grad. + value)	2 + 32	8 + 16	2 + 16
References $R$ (mixed)	128	64	64
Temperature $\tau$	cosine, 0.1 $\to$ 0.01$\times$ median pairwise posterior distance
TARP distance	pixel, 2$\times$ avg-pooled	pixel	null space of $A$
Seeds	3	3	3

On one NVIDIA GH200, a fine-tuning step takes 1.9 s with TARP-CvM against 0.7 s with the energy score or pCRPS on NS, 1.0 s against 0.7 s on IS, and 3.7 s against 2.6 s on MRI, a 1.4 to $2.7\times$ increase in training time with respect to the ES and pCRPS, and an 18 to $33\times$ increase with respect to DSM alone. Figure 7 shows example observations and target ground truth signals for the three inverse problems.

Vector (PDF)

The fastMRI data are used under the NYU fastMRI Dataset Sharing Agreement. All other code we use is available under open-source licenses (Apache 2.0, BSD 3-Clause or MIT), and the remaining datasets are publicly available.