Amortized Simulation-Based Inference in Generalized Bayes via Neural Posterior Estimation

Technical Summary: Amortized Simulation-Based Inference in Generalized Bayes via Neural Posterior Estimation

Problem Statement
Generalized Bayesian Inference (GBI) offers a robust framework for statistical inference under model misspecification by tempering the likelihood (or a loss function) with a temperature parameter $\beta > 0$. This results in a family of power posteriors defined as $p_\beta(\theta | x) \propto \pi(\theta)p(x|\theta)^\beta$. While GBI mitigates overconfidence and improves robustness, existing methods for simulators with implicit likelihoods typically rely on costly Markov Chain Monte Carlo (MCMC) or Stochastic Differential Equation (SDE) samplers. Crucially, these methods must be re-run for every new dataset and every specific $\beta$-value, creating a significant computational bottleneck for sensitivity analysis and robustness checks.

Methodology
The authors propose a fully amortized variational approximation for the tempered posterior family. Instead of sampling per dataset, they train a single neural network, $q_\phi(\theta | x, \beta)$, that conditions on both the observation $x$ and the temperature $\beta$. Once trained, sampling from $p_\beta(\theta | x)$ requires only a single forward pass, eliminating the need for simulator calls or inference-time MCMC.

The paper introduces two complementary training routes to generate the necessary training data and optimize the estimator:

1. Route A: Score-Assisted Tempered Synthesis

   • Mechanism: This route constructs an offline dataset of triples $(\theta, x, \beta)$ by approximately sampling from the unnormalized tempered joint density $\tilde{p}_\beta(\theta, x) \propto \pi(\theta)p(x|\theta)^\beta$.
   • Process: A joint score network $s_\psi(\theta, x)$ is first learned via Denoising Score Matching (DSM) on the base joint distribution. Short-run Langevin dynamics are then run across an annealed noise schedule, utilizing the learned score and the prior to synthesize pairs $(\theta, x)$ that target the tempered joint for various $\beta$.
   • Training: The amortized posterior estimator $q_\phi$ is trained via conditional Maximum Likelihood Estimation (MLE) on these synthesized triples.
   • Characteristics: This method can generate "off-manifold" samples (samples not strictly adhering to the simulator's forward mapping), potentially improving coverage in misspecified regimes, but relies on the stability of the score network and Langevin discretization.

2. Route B: SNIS-Weighted NPE (Self-Normalized Importance Sampling)

   • Mechanism: This route reuses a single base dataset $(\theta_i, x_i)_{i=1}^N \sim \pi(\theta)p(x|\theta)$ across all temperatures.
   • Process: It employs self-normalized importance sampling (SNIS) to reweight the base samples. The importance weights are defined as $w_\beta(\theta, x) \propto p(x|\theta)^{\beta-1}m(x)$, where $m(x)$ is a normalizing function.
   • Instantiations:
     • NLE-based: Uses a learned neural likelihood $\hat{p}_\eta(x|\theta)$ with $m(x)=1$.
     • NRE-based: Uses a learned likelihood ratio $\hat{r}_\psi(x|\theta) \approx p(x|\theta)/p(x)$ with $m(x)=p(x)^{1-\beta}$.
   • Training: The objective minimizes a weighted negative log-likelihood: $L_{SNIS}(\phi) = -\sum \tilde{w}_{\beta,i} \log q_\phi(\theta_i | x_i, \beta)$.
   • Theoretical Guarantee: The authors prove that for the NRE-based construction, the variance of the importance weights is finite for $\beta \in [1/2, 1]$, ensuring a consistent forward-KL fit to the tempered posterior.

Key Contributions

• First Fully Amortized GBI: The paper presents the first fully amortized variational approximation for the specific case of the tempered posterior family $p_\beta(\theta | x)$, enabling single-pass sampling for any $(x, \beta)$ without retraining.
• Dual Training Routes: It introduces two distinct strategies (Score-assisted synthesis vs. SNIS reweighting) to handle the challenge of learning a distribution that varies with $\beta$, offering trade-offs between offline compute cost, support coverage, and weight stability.
• Theoretical Analysis: The work provides proofs regarding the consistency of the SNIS-weighted objective (Proposition 4.1) and establishes finite variance conditions for the importance weights in the NRE-based approach (Proposition 4.2). It also proposes a "defensive" SNIS proposal to bound weights and control variance.

Results
The method was evaluated on four standard Simulation-Based Inference (SBI) benchmarks: Gaussian Mixture, Two Moons, SLCP, and the chaotic Lorenz–96 system, as well as a scientific application involving the Hodgkin–Huxley (HH) neuronal model.

• Performance: The $\beta$-amortized estimator achieves competitive posterior approximations (measured by MMD and C2ST metrics) compared to non-amortized MCMC-based power-posterior samplers across a wide range of temperatures.
• Trade-offs:
  • Route B performs efficiently when $\beta$ is close to 1 but can suffer from weight concentration and Effective Sample Size (ESS) collapse as $\beta$ moves away from 1 (particularly $\beta < 1$).
  • Route A demonstrates greater robustness at extreme $\beta$ values (e.g., small $\beta$) where Route B's weights become unstable, though it is sensitive to score errors and requires careful tuning of Langevin step sizes.
• Application: In the HH model experiment, the method produced stable posterior marginals and qualitatively accurate posterior predictive checks for real electrophysiological data across $\beta \in \{0.1, 1.0, 2.0\}$.

Significance and Claims
The authors claim that their approach significantly reduces the cost of robustness and sensitivity analyses in simulation-based inference. By enabling fast posterior and predictive queries after a single training phase, the method lowers the computational overhead associated with repeated sampling for different temperatures.

The paper positions this work as complementary to, rather than a replacement for, traditional adaptive inference. It acknowledges limitations, noting that amortization trades substantial offline simulation costs for test-time speed and that performance may degrade if the test data distribution shifts significantly or if $\beta$ falls outside the trained range. The authors emphasize that while the method eliminates inference-time MCMC, users must still monitor importance-weight stability (ESS) and validate against higher-fidelity samplers to ensure calibration, particularly in misspecified settings.