PaperlayerView exact revision on arXiv

Source-linked AI summary

Guidance for Prior Change via Density Ratio Estimation

Yichen Zang, Song Liu, Jiun-Yi Lin

arXiv:2608.21729v1stat.MLcs.LGstat.ME

TL;DR

Amortized SBI is efficient but tied to its training prior, making adaptation difficult as prior knowledge changes. The paper introduces DRE-based test-time guidance that directly learns the guidance term, and reports gains on most tasks, especially under OOD prior shifts, without retraining the base model.

  • Problem

    Amortized SBI models are tied to the fixed training prior, motivating prior adaptation without retraining the base model.

  • Method

    The method uses Density Ratio Estimation across diffusion steps to directly learn an unbiased, robust guidance term and supports different ratio estimators.

  • Results

    The method gains on the majority of tasks, particularly under out-of-distribution prior shifts where PriorGuide degrades due to biased guidance.

  • Takeaways & Limitations

    The framework enables flexible test-time prior adaptation without retraining the base model, supporting scientific workflows with evolving prior knowledge.

  • Takeaways & Limitations

    Simformer+SIR can collapse posterior samples when the posterior is poorly covered, producing metric-dependent comparisons because duplication affects C2ST more than MMD or W2.

Abstract

from arXiv · show

Simulation-Based Inference (SBI) serves as a vital framework for parameter inference in scientific fields where simulators involve intractable likelihoods, yet while amortized generative models offer rapid posterior estimation, they are often restricted by the specific priors used during training, thereby limiting their flexibility as prior knowledge evolves. To address this prior dependency, PriorGuide was introduced as an inference-time guidance method, but due to its intractable formulation, it relies on Gaussian approximations of the reverse transition kernel and Gaussian mixture model fitting for the prior ratio, both of which introduce systematic bias. Motivated by these limitations, we propose an unbiased test-time guidance framework that leverages Density Ratio Estimation (DRE) to learn a score guidance term, effectively decoupling the inference process from the prior training. Moreover, our framework remains agnostic to the specific density ratio estimators, making it a general and flexible framework for handling prior changes. Experimental results across multiple tasks demonstrate that our method matches or outperforms PriorGuide on C2ST and MMD in most tasks while maintaining robustness even under limited overlap between the training and target priors. Furthermore, we apply our method to Bayesian updating for parameter inference from planetary light-curve data, where it also demonstrates strong effectiveness and robustness. Code is available at https://github.com/a-chenchen/dre-based-prior-guidance .

1 Introduction

SBI enables inference when simulator likelihoods are intractable, but amortized models remain tied to their training priors. The proposed DRE-based test-time guidance adapts inference without PriorGuide’s structural approximations.

  • SBI uses simulator-generated data to learn posterior estimates when likelihoods are analytically intractable.
  • Amortized posterior models generalize across observations but remain dependent on the fixed training prior.
  • Changing the prior otherwise requires costly resimulation and retraining, while prior-amortization methods constrain the allowable prior family or space.
  • PriorGuide adapts a pretrained diffusion model at test time but uses GMM prior-ratio fitting and Gaussian reverse-kernel approximations that can bias guidance.
  • In the Ornstein-Uhlenbeck example, the method closely tracks the true posterior, whereas PriorGuide deviates.
  • The proposed method trains a lightweight DRE across diffusion steps to learn an unbiased guidance term for flexible prior adaptation.

2 Background

Bayesian inference depends on the prior, while SBI addresses intractable likelihoods through learned parameter–observation relationships. Diffusion-based posterior models support amortized sampling, and DRE-based guidance provides a less restrictive route to prior adaptation.

  • 2.1 Simulation-based Inference: The posterior combines likelihood and prior, making prior specification important for parameter inference.
  • 2.1 Simulation-based Inference: SBI learns statistical relationships from simulator-generated parameter–observation pairs when likelihood evaluation is difficult or impossible.
  • 2.2 Diffusion Models: Diffusion models generate posterior samples by reversing a noise-adding process, with a neural network approximating the intractable score.
  • 2.4 Density Ratio Estimation: DRE estimates density ratios from samples of two distributions without requiring access to their density functions, encompassing methods such as KLIEP and NRE.
  • 2.5 Prior Adaptation in Amortized SBI: Prior adaptation research includes meta-amortization and inference-time guidance, with PriorGuide modifying diffusion scores using the target-to-training prior ratio.
  • 2.5 Prior Adaptation in Amortized SBI: PriorGuide’s GMM ratio and Gaussian reverse-kernel assumptions can cause systematic bias and instability under low prior overlap, whereas DRE directly approximates guidance.

3 Methodology

The method adapts a diffusion-based SBI model from training prior p(θ) to target prior q(θ) at inference time by learning guidance through density-ratio estimation. It uses an importance-weighting reformulation for simulation-free training and supports estimator-agnostic guidance with robustness to mismatched prior supports.

  • Problem formulation: The framework adapts a score-based diffusion SBI model trained on p(θ) to posterior inference under a new target prior q(θ).The target posterior is q(θ | x) ∝ q(θ)p(x | θ), and the corresponding posterior predictive distribution is also targeted.
  • Problem formulation: Under q(θ) = r(θ)p(θ), the target joint satisfies q(θ, x) = r(θ)p(θ, x) because the likelihood remains unchanged.This reweighting relationship underlies the diffusion-space importance-weighting strategy.
  • Guidance derivation: The guided score is decomposed into the original posterior score plus a guidance term involving the target and training diffusion-space joint distributions.The reformulation expresses prior adaptation through estimating the joint log ratio log q(θ_t, x)/p(θ_t, x).
  • Simulation-free training: Importance weighting enables simulation-free DRE training by reusing base-model training pairs and perturbing θ₀ through the diffusion forward kernel.This avoids additional simulator calls while supplying the required target-joint expectations.
  • Estimator framework: NRE and KLIEP estimate the same diffusion-space log density ratio despite using different objectives, motivating an estimator-agnostic guidance framework.Both recover the guidance term through the gradient of the optimal estimator.
  • Estimator framework: The generalized DRE framework uses a convex-generator-dependent activation and recovers different estimators within a shared formulation, while diffusion noise stabilizes ratios under mismatched prior supports.NRE is a special case; KLIEP recovers the same guidance term at optimality through a separate argument.

4 Experiments

The experiments evaluate DRE-based prior adaptation across posterior and posterior predictive SBI tasks under ID and OOD shifts, plus sequential planetary light-curve updating. The method performs strongly across benchmarks and remains stable in settings where PriorGuide exhibits bias or instability.

  • Posterior Inference under Prior Shift: DRE-based guidance is evaluated on diverse SBI benchmarks under both in-distribution and out-of-distribution prior shifts.The benchmarks span low- to high-dimensional spaces and static and time-series data.
  • Posterior Inference under Prior Shift: NRE and KLIEP are compared with vanilla Simformer, PriorGuide, WSM, and importance-resampling references using lightweight test-time adaptation.The ranked comparison includes methods that leave the base model unchanged and require no new simulator calls.
  • Posterior Inference under Prior Shift: ∼0.58 and ∼0.54 C2ST scores for DRE-based guidance on Gaussian Linear 6D and OUP outperform PriorGuide’s ∼0.69 and ∼0.65 under OOD shift.On SLCP (OOD), NRE is the only method in the best-performing group.
  • Posterior Predictive under Prior Shift: PriorGuide’s predictive distribution is systematically biased on OUP (OOD) and misplaces predictive mass in Turin (ID), despite comparable RMSE in Turin.DRE-based guidance and WSM exhibit neither failure mode in these experiments.
  • Planet Light Curve Posterior Update: The sequential light-curve task infers orbital parameters from TESS observations, incorporating five normalized transit-centered segments into posterior updates.The broader task contains eight parameters, while results report k, a, inc, and P.
  • Planet Light Curve Posterior Update: DRE guidance shifts and concentrates posteriors toward reference values as observations accumulate, while PriorGuide becomes unstable after a few updates.Across systems, DRE maintains low and stable RMSE with slight improvements as observations increase; PriorGuide’s error progressively increases.

5 Discussion

The discussion attributes prior-adaptation performance to accurate guidance estimation and contrasts DRE’s accuracy and inference cost with PriorGuide. It also identifies high-dimensional ratio estimation and weak prior-ratio signal as boundaries of the approach.

  • Discussion: Learning the gradient of the joint density ratio via DRE avoids PriorGuide’s structural approximations and provides an unbiased guidance estimate.The paper connects this mechanism with gains on most tasks, particularly under OOD prior shifts.
  • Discussion: DRE-based guidance gains are concentrated particularly under OOD prior shifts, where PriorGuide degrades due to biased guidance.This summarizes the reported majority-task empirical pattern rather than a universal guarantee.
  • Discussion: Both methods avoid additional simulator calls but require lightweight target-prior training, so neither is purely test-time.DRE fits a ratio estimator, whereas PriorGuide fits a GMM for the prior ratio.
  • Discussion: DRE-based sampling avoids PriorGuide’s Langevin correction during inference, eliminating a cost that recurs at every observation.The distinction concerns inference-time sampling rather than target-prior fitting.
  • Discussion: Ratio-estimator quality degrades in high dimensions, and nearly flat prior ratios can leave guidance with little signal over the posterior support.Learned observation embeddings are identified as a standard remedy and future-work direction for the high-dimensional case.
  • Discussion: The framework enables flexible prior adaptation without retraining the base model, supporting scientific workflows in which prior knowledge evolves.This is presented as the paper’s practical test-time-training consequence.

A Meta-Amortization Methods for Prior Change

This section relates the target and original joint distributions through the prior ratio and derives equivalence between the guidance formulations using logarithmic gradients.

  • Meta-Amortization Methods for Prior Change: The target joint distribution is expressed as the original joint distribution multiplied by the prior ratio.The relation is written as q(θ0, x) = r(θ0)p(θ0, x).
  • Meta-Amortization Methods for Prior Change: Because the transition kernel is independent of the prior, the target and original diffused joint distributions can be related through the same ratio.The section uses this independence to simplify the marginal joint distribution.
  • Meta-Amortization Methods for Prior Change: Taking the gradient of the logarithm of the joint-density ratio establishes equivalence between the proposed and prior guidance terms.The derivation explicitly concludes the equivalence after logarithmic differentiation.

C Proof of Theorem 1

The proof characterizes the density-ratio objective pointwise at each diffusion timestep, uses a change of measure for simulation-free training, and recovers the guidance term from the log-ratio gradient.

  • Proof of Theorem 1: The variational density-ratio objective is separable across diffusion timesteps, allowing analysis at a fixed state and timestep.The proof reduces the problem to minimizing a pointwise integrand.
  • Proof of Theorem 1: Strict convexity of the objective component yields a unique minimizer by setting the first-order derivative to zero.The minimizer is derived for the functional at a fixed state (θt, x, t).
  • Proof of Theorem 1: A change-of-measure identity computes expectations under the target distribution using samples from the original joint distribution.The prior ratio r(θ0) = q(θ0)/p(θ0) supplies the reweighting relation.
  • Proof of Theorem 1: Taking the gradient of the log-ratio recovers the guidance term required by the method.This completes the proof after substituting the density-ratio objective.

D.1 Derivation of KLIEP as an Extended Instance

KLIEP, although not pointwise in the generalized DRE template, can be reduced to the same variational form and recover the correct diffusion guidance term. The surrounding framework applies this guidance through a learned score correction during reverse diffusion.

  • Derivation: KLIEP is reduced to the generalized variational DRE form despite not being pointwise in the theorem’s template.The reduction uses an exponential generator and a normalized network output.
  • Derivation: The generator derivative satisfies g′(b) = exp(b −1), so the optimal solution recovers the density ratio q/p.At optimality, g′(b⋆) equals q/p.
  • Derivation: KLIEP’s optimal discriminator recovers the joint log-density ratio up to an additive constant, yielding the correct guidance term.Its gradient therefore matches the gradient of the log-density ratio.
  • Benchmark tasks: The numerical benchmarks include Two Moons, Gaussian Linear, SLCP, OUP, and Turin, using task-specific simulators and posterior-sampling procedures.Ground-truth posteriors are obtained analytically, by surrogate score-based NPE, rejection sampling, or NUTS with symmetry completion, depending on the task.
  • Benchmark tasks: The Turin simulator produces complex-valued frequency-domain observations that are transformed into real-valued outputs through an inverse Fourier transform and logarithmic power scaling.The setup uses 101 frequency points and z-score normalization for simulated outputs.

G.2 Prior Specifications

The experiments specify training and target priors for each simulator and train lightweight DRE guidance models on augmented data while leaving the base Simformer unchanged. At inference, DRE guidance additively corrects the base score without extra Langevin steps.

  • Prior specifications: The target-prior specifications use explicit vectors, matrices, and mixture covariances, with Σmix = 0.2 16×6 + 0.8 I6 in the stated setup.The notation defines all-ones vectors and matrices, zero vectors, and identity matrices.
  • Training setup: The DRE training set contains 470,000 samples: 50,000 original simulations plus 420,000 posterior-derived pairs from the pretrained Simformer.The added pairs come from 420 observations with 1,000 posterior samples each.
  • Training setup: The base Simformer is trained with a VP-SDE and a transformer using six layers, four attention heads, token dimension 40, and a 128-dimensional time embedding.The diffusion interval is t ∈[10−5, 1] with βmin = 0.1 and βmax = 10.0.
  • Guidance models: The DRE guidance model takes noisy parameters, observations, and diffusion time as input and outputs an estimate of log q(θt, x)/p(θt, x).NRE and KLIEP share a four-layer, 256-hidden-unit MLP architecture with SiLU activations.
  • Inference: At each reverse diffusion step, DRE guidance additively corrects the base score without additional correction steps, while PriorGuide applies four Langevin steps.All methods use the same 25-step reverse diffusion sampler; evaluation draws 5,000 posterior samples per observation.

H Additional Numerical Experimental Results

The additional numerical results evaluate posterior and predictive inference under prior shifts using repeated held-out observations and multiple metrics. Effective sample fraction diagnoses whether the prior update carries enough usable signal, while Turin remains a performance boundary for DRE guidance.

  • Evaluation: Posterior predictive results are summarized by RMSE over 10 observations, with bold entries marking the best or statistically indistinguishable results.The comparison uses means and standard deviations averaged over three runs.
  • Posterior prediction: Figure 5 organizes posterior predictive curves by OUP and Turin task settings, ID versus OOD regimes, and the first five observations.Solid lines show predictive means, dashed lines and shading show 95% confidence intervals, and the vertical line marks the prediction boundary.
  • Prior-overlap diagnostics: ESS/N is 0.990 for Two Moons, indicating nearly constant prior-ratio weights over the base posterior support and an exact guidance term close to zero.This makes approximation error more prominent than genuine guidance signal and provides an inexpensive pre-adaptation diagnostic.
  • Limitations: PriorGuide outperforms DRE guidance on Turin, whose combined parameter-and-observation dimension is dθ + dx = 105.The paper identifies high-dimensional DRE degradation as a possible explanation and leaves learned observation compression for future work.

J.1 Data Acquisition and Preprocessing

The planetary application preprocesses TESS light curves into aligned transit windows and performs sequential posterior updating across exoplanet observations. Synthetic training, DRE-based adaptation, and a Monte Carlo PriorGuide baseline are specified for this setting.

  • Data acquisition: TESS light curves for WASP-126 b, WASP-39 b, and GJ 436 b are downloaded, sigma-clipped, normalized, detrended, stitched, and period-estimated with BLS.The systems span different transit depths and orbital configurations.
  • Preprocessing: Each transit window is recentered on the detected flux minimum using smoothing and a sub-cadence parabolic fit, then cropped to ±0.5 days.This aligns transit events to t = 0, matching simulator training conventions.
  • Preprocessing: The final datasets contain five windows for WASP-39 b and seven windows each for GJ 436 b and WASP-126 b, with each window represented in R128.The windows are incorporated sequentially in posterior updates.
  • Synthetic training: Synthetic training uses a quadratic limb-darkening simulator with an 8-dimensional diagonal Gaussian reference prior.The base score model is trained on 3 × 105 synthetic samples.
  • Sequential adaptation: Sequential DRE adaptation trains a prior-ratio estimator from previous-window posterior samples and then trains an NRE guidance model at each transit window.The guidance model uses a four-layer MLP with 256 hidden units and a 16-dimensional Fourier time embedding.
  • Baseline and inference: The Monte Carlo PriorGuide baseline draws 1,000 Tweedie-posterior samples at each of 25 reverse-diffusion steps and applies eight Langevin corrections.Guided inference draws 105 posterior samples per transit window and excludes the independently recentered midpoint offset.

K Additional Planet Light Curve Inference Results

Figures 7–9 examine sequential posterior evolution and normalized posterior RMSE for exoplanet light-curve inference. Across sequential updates, the method maintains stable error while PriorGuide’s error progressively increases.

  • Figure 7 shows sequential posterior evolution for the GJ 436 b system.The figure uses the reference value from [16], with layout and notation following Figure 4.
  • Figure 8 shows sequential posterior evolution for the WASP-126 b system.The figure uses the reference value from [24], with layout and notation following Figure 4.
  • Across three exoplanet systems, normalized posterior RMSE remains stable for our method but progressively increases for PriorGuide during sequential updates.RMSE is computed from transit parameters k, a, inc, and P, normalized per parameter and then averaged.
Loading 2608.21729v1…