📊 statistics

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

PriorGuide introduces a test-time adaptation technique for diffusion-based amortized simulation-based inference that enables flexible alignment of pre-trained models with new prior distributions using a novel guidance approximation, thereby eliminating the need for costly retraining while incorporating updated expert knowledge.

Original authors: Yang Yang, Severi Rissanen, Paul E. Chang, Nasrulloh Loka, Daolang Huang, Arno Solin, Markus Heinonen, Luigi Acerbi

Published 2026-04-22

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yang Yang, Severi Rissanen, Paul E. Chang, Nasrulloh Loka, Daolang Huang, Arno Solin, Markus Heinonen, Luigi Acerbi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a master chef who has spent years perfecting a recipe for a specific type of soup. You trained your palate on a massive pot of soup made with a very broad, generic set of ingredients (a "uniform prior"). You can now taste any new batch of this soup and instantly tell you exactly what spices were used, even if you've never tasted that specific batch before. This is what modern AI "amortized inference" does: it learns a general rule to guess the hidden causes (ingredients) behind observed data (the taste).

The Problem: The "One-Size-Fits-All" Trap
Here's the catch: Your chef training was based on that generic soup. But what if, tomorrow, you are asked to taste a soup made with a very specific set of ingredients? Maybe it's a soup where the chef only uses organic tomatoes and never uses salt.

If you try to guess the ingredients using your old, generic training, you might get it wrong because your brain expects salt and generic tomatoes. To fix this, the old way of doing things would be to fire your chef, hire a new one, and spend months retraining them specifically on "organic, no-salt" soup. That is slow, expensive, and wasteful.

The Solution: PriorGuide
The paper introduces PriorGuide, a clever technique that acts like a "real-time taste adjuster" for your existing master chef.

Instead of retraining the chef, PriorGuide gives them a pair of smart glasses at test time (when they are actually tasting the soup). These glasses show the chef a note saying: "Hey, remember? For this specific bowl, the chef only uses organic tomatoes and no salt. Ignore your generic training for a second and focus on these specific rules."

The chef then uses their existing skills but adjusts their guess on the fly to match these new rules. They don't need to relearn how to taste; they just need to apply a little bit of extra "guidance" based on the new information.

How It Works (The Metaphor)

The Diffusion Model (The Chef):
Think of the AI model as a chef who learns to reverse-engineer a soup. They start with a "noisy" version of the soup (where all the ingredients are mixed up) and slowly "denoise" it to figure out the original recipe. This process is called diffusion.
The Training Prior (The Generic Menu):
When the chef was trained, they were fed every possible soup recipe from a huge, broad menu. They learned the general rules of cooking.
The New Prior (The Special Request):
Now, a customer orders a soup with very specific constraints (e.g., "No garlic, lots of basil"). In the old days, you'd have to retrain the chef. With PriorGuide, you just hand the chef a guidance card.
The Guidance Approximation (The Smart Glasses):
The paper's magic trick is figuring out how to translate that "No garlic, lots of basil" rule into a mathematical instruction the chef can understand while they are cooking.
- The paper uses a Gaussian Mixture Model (a fancy way of saying "a flexible shape made of overlapping circles") to approximate the difference between the generic menu and the special request.
- It calculates a "correction score" that nudges the chef's guess away from the generic path and toward the specific path.
Langevin Dynamics (The Taste Test):
Sometimes, the chef might still be a little unsure. PriorGuide allows the chef to take a few extra "taste tests" (called Langevin steps) at the end of the process to refine the guess. It's like the chef saying, "Hmm, this still feels a bit off, let me adjust the salt one more time." This costs a tiny bit more time but makes the result much more accurate.

Why Is This a Big Deal?

No Retraining: You don't need to spend weeks retraining the AI. You just change the "guidance" at the moment of use.
Flexibility: It works even if the new rules are complex (like a soup with two distinct flavor profiles, or a "mixture" prior).
Speed vs. Accuracy Trade-off: If you are in a rush, you can use a quick, rough guidance. If you need perfect accuracy (like for climate science or medical diagnosis), you can let the AI take those extra "taste tests" to get it right.

Real-World Analogy: The Weather Forecaster

Imagine a weather forecaster (the AI) trained on 100 years of average global weather. They are great at predicting rain or sun.

The Old Way: If a specific region suddenly has a unique micro-climate due to a new dam, the forecaster would need to retrain their entire model on data from that dam.
The PriorGuide Way: The forecaster keeps their 100-year knowledge but gets a real-time update: "Remember, there is a dam here now. The air is 5 degrees cooler and 20% more humid than your average." The forecaster instantly adjusts their prediction to account for the dam without forgetting everything they learned about the rest of the world.

Summary

PriorGuide is a tool that lets powerful AI models adapt to new, specific rules or beliefs on the fly, without needing to be retrained from scratch. It turns a rigid, one-size-fits-all AI into a flexible, context-aware expert that can incorporate new information instantly, saving time and computing power while making better predictions.

1. Problem Statement

Simulation-Based Inference (SBI) is a critical tool for Bayesian inference in fields where the likelihood function $p(x|\theta)$ is intractable, but data can be simulated from a forward model $x \sim p(x|\theta)$ . Modern amortized SBI methods (e.g., using Diffusion Models or Transformers) train a neural network once on a vast dataset of simulated parameter-data pairs $(\theta, x)$ to learn the inverse mapping $x \to p(\theta|x)$ .

The Core Limitation:
Standard amortized methods are tied to the training prior $p_{train}(\theta)$ used to generate the training data. In practice, domain experts often possess specific knowledge or updated beliefs that require a different target prior $q(\theta)$ .

Retraining is prohibitive: Changing the prior requires re-simulating massive datasets and retraining the model, which is computationally expensive, especially for complex simulators (e.g., in neuroscience or engineering).
Existing adaptations are limited: Some recent methods allow prior adaptation but rely on pre-training over a "meta-prior" (a fixed set of predefined prior families, e.g., specific histograms or Gaussian mixtures). This restricts the user to only those pre-defined prior shapes and cannot handle arbitrary, complex, or non-factorized priors defined at test time.

The Goal: Develop a method to adapt a pre-trained diffusion-based SBI model to a new, arbitrary target prior $q(\theta)$ at test time without retraining the base model.

2. Methodology: PriorGuide

PriorGuide leverages the guidance mechanisms inherent in diffusion models to steer the generative process toward a new posterior distribution defined by a new prior.

2.1 Theoretical Foundation

The method is based on the observation that the target posterior $q(\theta|x)$ under a new prior $q(\theta)$ can be related to the original posterior $p(\theta|x)$ (trained under $p_{train}(\theta)$ ) via a prior ratio $r(\theta)$ :
$q(\theta|x) \propto q(\theta)p(x|\theta) = \underbrace{\frac{q(\theta)}{p_{train}(\theta)}}_{r(\theta)} \cdot \underbrace{p_{train}(\theta)p(x|\theta)}_{p(\theta|x)} \propto r(\theta)p(\theta|x)$
Thus, sampling from the new posterior is equivalent to sampling from the original posterior weighted by the ratio $r(\theta)$ .

2.2 Guidance Term Derivation

In diffusion models, the sampling process is guided by a score function $\nabla \log p(\theta_t|x)$ . To incorporate the new prior, PriorGuide modifies the score function during the reverse diffusion process:
$\nabla_{\theta_t} \log q(\theta_t|x) = \underbrace{\nabla_{\theta_t} \log p(\theta_t|x)}_{\text{Original Score } s(\theta_t, t, x)} + \nabla_{\theta_t} \log \mathbb{E}_{p(\theta_0|\theta_t, x)} [r(\theta_0)]$
The second term is the guidance term, representing the influence of the new prior ratio propagated to the noisy state $\theta_t$ .

2.3 Analytical Approximations

Directly computing the expectation $\mathbb{E}[r(\theta_0)]$ is intractable. PriorGuide introduces two key approximations to make this tractable and differentiable:

Reverse Transition Kernel Approximation: The conditional distribution $p(\theta_0|\theta_t, x)$ is approximated as a multivariate Gaussian, where the mean is estimated via Tweedie's formula using the trained score model, and the covariance is approximated analytically.
Prior Ratio Approximation (GMM): The ratio function $r(\theta) = q(\theta)/p_{train}(\theta)$ is approximated as a Gaussian Mixture Model (GMM). Since both $q(\theta)$ and $p_{train}(\theta)$ are known analytically, the parameters of this GMM (weights, means, covariances) are fitted via gradient descent to minimize the $L_2$ error between the GMM and the true ratio.

By substituting these approximations, the guidance term becomes an analytical closed-form solution, allowing for efficient computation without Monte Carlo sampling.

2.4 Refined Sampling with Langevin Dynamics

The Gaussian approximation of the reverse kernel introduces errors, particularly at higher noise levels. To ensure asymptotic correctness, PriorGuide interleaves Langevin Dynamics steps between standard diffusion steps. This transforms the sampling process into an annealed MCMC procedure, allowing users to trade off computational cost (more steps) for higher inference fidelity.

3. Key Contributions

Test-Time Prior Adaptation Framework: PriorGuide is the first method to enable flexible adaptation of pre-trained diffusion-based SBI models to arbitrary target priors at inference time without retraining.
Novel Guidance Approximation: It introduces a principled way to turn a target prior into a tractable guidance term using a GMM approximation of the prior ratio, avoiding the need for pre-training on meta-priors.
Asymptotic Correctness: By integrating Langevin dynamics, the method provides a mechanism to correct approximation errors, offering a tunable trade-off between compute time and accuracy.
Broad Applicability: The method handles complex, non-factorized, and multimodal priors, overcoming the limitations of previous amortized methods restricted to specific prior families.

4. Experimental Results

The authors evaluated PriorGuide on six SBI benchmarks, ranging from 2D toy models (Two Moons) to high-dimensional engineering and neuroscience models (Ornstein-Uhlenbeck Process, Turin Radio Propagation, Bayesian Causal Inference).

Posterior Inference Accuracy:
- PriorGuide consistently outperformed the base Simformer (which ignores the new prior) and the Amortized Conditioning Engine (ACE) (which relies on pre-defined meta-priors).
- In scenarios with strong or mixture priors (complex, specific beliefs), PriorGuide achieved significantly lower Mean Marginal Total Variation (MMTV) and better Classifier Two-Sample Test (C2ST) scores, accurately recovering the true posterior.
- ACE often failed to match the performance of PriorGuide when the target prior was complex or non-factorized, as it was constrained by its pre-training distribution.
Posterior Predictive Inference:
- PriorGuide successfully adapted to new priors for forecasting tasks (e.g., predicting time-series trajectories), generating data distributions that closely matched ground truth under the new prior assumptions.
Compute-Accuracy Trade-off:
- Experiments varying the number of diffusion steps ( $N$ ) and Langevin steps ( $N_L$ ) showed that increasing compute (via $N_L$ ) improves accuracy.
- PriorGuide achieved Pareto-optimal performance, offering high accuracy with moderate compute costs compared to retraining or rejection sampling.
Comparison with Baselines:
- Compared to Rejection Sampling and Sampling-Importance-Resampling (SIR), PriorGuide was significantly more efficient in high-dimensional spaces (e.g., 20D Gaussian Linear), where traditional sampling methods suffered from extremely low acceptance rates or statistical instability.
- It avoided the "mode collapse" issues seen in Neural Likelihood Estimation (NLE) combined with MCMC on multimodal problems.

5. Significance and Impact

Decoupling Training from Inference: PriorGuide allows scientists to train a single, expensive "foundation model" on a broad prior and then instantly adapt it to specific expert knowledge or updated hypotheses without re-simulation. This drastically reduces the computational footprint of scientific workflows.
Enabling Sensitivity Analysis: It facilitates robust prior sensitivity analysis, a crucial step in Bayesian workflows, by allowing rapid re-evaluation of conclusions under different prior assumptions.
Test-Time Compute Paradigm: The paper exemplifies the "test-time compute" paradigm, where models are designed to be flexible at inference time through dedicated computations (guidance + Langevin steps) rather than attempting to amortize every possible scenario during training.
Practical Utility: The method is particularly valuable for domains with expensive simulators (e.g., climate modeling, epidemiology, physics) where retraining is infeasible, and for applications requiring complex, non-standard priors that cannot be pre-defined.

Limitations: The method requires the target prior $q(\theta)$ to have substantial overlap with the training prior $p_{train}(\theta)$ to avoid out-of-distribution (OOD) issues. If the target prior lies in regions where the score model was never trained, the guidance becomes unstable. Additionally, the GMM fitting and matrix operations can pose scalability challenges for extremely high-dimensional parameter spaces ( $D \gg 20$ ).