CP4SBI: Local Conformal Calibration of Credible Sets in Simulation-Based Inference

Technical Summary: CP4SBI – Local Conformal Calibration of Credible Sets in Simulation-Based Inference

1. Problem Statement

Simulation-Based Inference (SBI) has become a standard framework for likelihood-free Bayesian inference in fields such as astrophysics, neuroscience, and fluid dynamics, where complex stochastic simulators generate data $x$ given parameters $\theta$ but lack tractable likelihood functions. While deep generative models (e.g., normalizing flows, diffusion models) have enabled efficient posterior approximation, these approximations often suffer from statistical miscalibration. Specifically, credible regions constructed from approximate posteriors $\hat{p}(\theta|x)$ frequently fail to achieve their nominal coverage levels, leading to overconfident inferences.

Existing attempts to correct this using Conformal Prediction (CP) typically yield only marginal coverage guarantees ($P(\theta \in C(X)) = 1-\alpha$), which average over all possible data and parameters. This is insufficient for Bayesian inference, which requires conditional coverage ($P(\theta \in C(X) | X = x_{obs}) = 1-\alpha$) for specific observations. Furthermore, many current CP methods for SBI require closed-form posterior densities, a constraint that excludes modern sample-based generative models like diffusion models.

2. Methodology: CP4SBI

The authors propose CP4SBI, a model-agnostic, post-hoc framework that integrates conformal prediction into the SBI pipeline to construct credible sets with finite-sample guarantees. The core insight is to reinterpret the scoring function $s(\theta; x)$ (derived from the approximate posterior) as a nonconformity score within a conformal inference framework.

The framework requires a calibration dataset $D_{cal} = \{(\theta_i, x_i)\}_{i=1}^B$, where pairs are simulated from the prior and the forward model. CP4SBI offers two distinct variants to achieve local and conditional coverage:

A. LoCart CP4SBI (Local Calibration via Regression Trees)

• Mechanism: This method partitions the feature space $\mathcal{X}$ using a regression tree trained to predict the conformity score $s(\theta; x)$ from the observation $x$. The tree is augmented with an estimate of the conditional variance of the score to improve partitioning.
• Calibration: For a new observation $x_{obs}$, the method identifies the leaf node (region $A_j$) containing $x_{obs}$. It then computes a local conformal threshold $t_{1-\alpha}$ using only the calibration scores falling within that specific region.
• Guarantee: This ensures local coverage ($P(\theta \in C(X) | X \in A_j) \geq 1-\alpha$). Under regularity conditions, as the calibration set size $B \to \infty$, the partition becomes finer, and the method achieves asymptotic conditional coverage ($P(\theta \in C(X) | X = x) \to 1-\alpha$).

B. CDF CP4SBI (CDF-based Calibration)

• Mechanism: This variant transforms the raw nonconformity score $s(\theta; x)$ using an estimate of its conditional cumulative distribution function (CDF), $\hat{F}(s | x)$. The transformed score is $s'(\theta; x) = \hat{F}(s(\theta; x) | x)$.
• Implementation: Unlike standard settings where the conditional CDF must be estimated from scratch, CP4SBI leverages the already-trained posterior approximator. It estimates $\hat{F}$ via a Monte Carlo approach: drawing $M$ samples $\{\theta_j\}_{j=1}^M \sim \hat{p}(\theta|x)$ and computing the empirical CDF of the scores.
• Calibration: A single global threshold is computed on the transformed scores $s'$. Because the transformation aims to make the scores uniformly distributed conditional on $x$, the resulting credible set achieves asymptotic conditional coverage as the training size $K$, calibration size $B$, and Monte Carlo samples $M$ all approach infinity.

Flexibility:

• Model Agnostic: Works with density-based estimators (e.g., Neural Posterior Estimation) and sample-based generative models (e.g., Diffusion, Flow Matching). For sample-based models, scores are approximated using kernel density estimation or empirical moments.
• Scoring Functions: Compatible with Highest Posterior Density (HPD), symmetric, and quantile-based regions.
• Transformations: Naturally handles nuisance parameters or transformations $\phi = g(\theta)$ by applying the method to the transformed space.

3. Key Contributions

1. Guaranteed Marginal Coverage: Unlike standard approximate credible sets, CP4SBI provides non-asymptotic marginal coverage guarantees by reinterpreting Bayesian scores as nonconformity scores.
2. Enhanced Local/Conditional Adaptivity:
   • LoCart CP4SBI provides finite-sample local coverage and asymptotic conditional coverage by adapting region sizes to inference difficulty via regression trees.
   • CDF CP4SBI achieves asymptotic conditional coverage by leveraging the posterior approximator to transform scores, improving calibration without requiring closed-form densities.
3. Generality: The framework is compatible with any posterior approximator (density or sample-based) and various scoring functions, addressing a gap where previous CP methods required closed-form posteriors unavailable in modern SBI.

4. Experimental Results

The authors evaluated CP4SBI on ten SBI benchmarks (e.g., Two Moons, Gaussian Mixture, Lotka-Volterra) and a computational neuroscience Neural Mass Model (NMM).

• Metrics: Performance was measured using Mean Absolute Error (MAE) of conditional coverage and Average Marginal Coverage (AMC).
• Baselines: Compared against Self-calibration, Global Conformal Prediction, and HDR (Highest Density Region) recalibration.
• Findings:
  • Conditional Coverage: Both LoCart and CDF variants significantly outperformed baselines in 8 out of 10 benchmarks when using Neural Posterior Estimation (NPE). LoCart CP4SBI was particularly effective for score-diffusion models (NPSE), ranking best in 7 of 10 tasks.
  • Marginal Coverage: The proposed methods consistently maintained coverage close to the nominal level (0.9), whereas baselines like Self-calibration and HDR often failed to maintain marginal validity.
  • Adaptivity: In the NMM case study, LoCart CP4SBI produced adaptive credible regions that varied in size based on the specific parameter combination, whereas global methods produced uniformly large regions and self-calibration produced uniformly small ones.
  • Robustness: The methods performed well across different simulation budgets, though LoCart required sufficient calibration data to populate tree partitions effectively.

5. Significance and Claims

The paper claims that CP4SBI represents a meaningful advance in uncertainty quantification for computational science. By providing finite-sample guarantees and local/conditional coverage, it addresses the critical issue of miscalibration in SBI, which undermines the reliability of downstream decision-making.

The authors emphasize that their approach is modest regarding the tightness of the resulting regions: the quality of the credible sets still depends on the initial quality of the posterior approximation. Poor approximations will yield wider, albeit calibrated, sets. However, the framework ensures that the coverage is statistically valid, preventing the overconfidence common in current SBI practices.

The work positions CP4SBI as a general-purpose tool that can be applied to any SBI method (including those using diffusion models or flow matching) to enhance trustworthiness. The authors note that while they focused on two specific local conformal approaches, the framework is open to other conditional coverage techniques and future extensions to hierarchical or discrete parameter settings.