PriorGuide: Test-Time Prior Adaptation for Simulation-Based Inference

TL;DR

PriorGuide enables test-time adaptation of diffusion-based amortized SBI models to new priors $q({\bm{\theta}})$ without retraining the score network trained under $p_{\text{train}}({\bm{\theta}})$. It achieves this by modeling the prior ratio $r({\bm{\theta}})=q({\bm{\theta}})/p_{\text{train}}({\bm{\theta}})$ as a Gaussian mixture and deriving a guidance term that augments the learned score during the diffusion sampling process; the reverse transition kernel is approximated as a Gaussian to yield tractable closed-form corrections. The approach supports both posterior and posterior predictive sampling and can be augmented with corrective Langevin steps to trade compute for fidelity, forming an annealed MCMC scheme. Empirically, PriorGuide improves inference accuracy across diverse SBI benchmarks under mild, strong, and mixture priors, while reducing the need for costly retraining and enabling post-hoc prior sensitivity analyses in scientific workflows.

Abstract

Amortized simulator-based inference offers a powerful framework for tackling Bayesian inference in computational fields such as engineering or neuroscience, increasingly leveraging modern generative methods like diffusion models to map observed data to model parameters or future predictions. These approaches yield posterior or posterior-predictive samples for new datasets without requiring further simulator calls after training on simulated parameter-data pairs. However, their applicability is often limited by the prior distribution(s) used to generate model parameters during this training phase. To overcome this constraint, we introduce PriorGuide, a technique specifically designed for diffusion-based amortized inference methods. PriorGuide leverages a novel guidance approximation that enables flexible adaptation of the trained diffusion model to new priors at test time, crucially without costly retraining. This allows users to readily incorporate updated information or expert knowledge post-training, enhancing the versatility of pre-trained inference models.

Figures (7)

• Figure 1: PriorGuide adapts a diffusion model to new prior information at test time for simulator-based inference---here for a time-series model. Left: Original broad training prior and a new, more specific target prior. Middle: Corresponding true posterior distributions over parameters ($\uparrow$) and predictive data ($\downarrow$) for each prior, given observed data. Right: Posterior ($\uparrow$) and posterior-predictive ($\downarrow$) samples from the diffusion model trained on the old prior vs. those from PriorGuide at low or high test-time cost, illustrating PriorGuide's ability to match the new posterior from the middle column.
• Figure 2: Example posterior predictive distributions for OUP and Turin models (strong priors).
• Figure 3: Pareto frontiers with respect to number of function evaluations (NFEs) and MMTV on posterior inference for OUP and Turin, with varying number of diffusion and Langevin steps.
• Figure :
• Figure A1: Left: Original Two Moons posterior and posterior samples from the base diffusion model trained on a uniform prior. Middle: New prior information about the parameters becomes available. Right: PriorGuide steers the diffusion process to match the Bayesian posterior under the new prior.

Theorems & Definitions (7)

• Proposition 1
• proof
• Proposition 2
• Proposition : Proposition \ref{['eq:prop1']}
• proof
• Proposition : Proposition \ref{['eq:prop2']}
• proof