How to use variational inference

How to use variational inference#

Variational inference (VI) fits a tractable distribution \(q\) to the posterior by optimizing a divergence against the potential. In sbi it is a sampler: it works with any potential, so it is available for NLE and NRE, and it is the option to reach for when MCMC is too slow or when sampling latency matters. See how to choose sampling algorithms for how it compares to the alternatives.

VIPosterior has two modes, and they are mutually exclusive:

  • Fixed-observation VI (train()) learns one approximation for one \(x_o\).

  • Amortized VI (train_amortized()) learns a conditional approximation \(q(\theta \mid x)\) that serves many observations.

Calling one on a posterior that was already trained in the other mode discards the earlier fit and raises a warning, so use separate VIPosterior objects when both results are needed.

Fixed-observation VI#

Building and training are separate steps. Pass VIPosteriorParameters to choose the objective:

from sbi.inference.posteriors import VIPosteriorParameters

vi_posterior = inference.build_posterior(
    posterior_parameters=VIPosteriorParameters(vi_method="rKL"),
)
vi_posterior.train(x=x_o)
vi_samples = vi_posterior.sample((1000,), x=x_o)

Fixed-observation VI supports vi_method="rKL", "fKL", "IW", or "alpha". Some are mode-seeking ("rKL", "alpha" above 1) and some are mass-covering ("fKL", "IW", "alpha" below 1).

Check the approximation#

VI is approximate, so verify the fit before relying on it. The evaluate() method prints a quality score: by default the shape parameter \(\hat{k}\) of a generalized Pareto distribution fitted to the tail of the importance weights of \(q\) against the potential, where below roughly 0.5 is good and above 1.0 indicates a poor approximation (Yao et al., 2018). Passing quality_control_metric="prop" instead reports an \(R^2\) of \(q\) against the joint, where values above 0.5 are good.

vi_posterior.evaluate()

Where feasible, also compare against MCMC on the same observation. evaluate() and the vi_method choices apply to fixed-observation VI only.

Refine a fixed-observation VI posterior#

Sampling-importance-resampling (SIR) can use the trained VI posterior as a proposal. This workflow is fixed at x_o; it does not imply batched refinement of an amortized posterior.

from sbi.inference import ImportanceSamplingPosterior

refined_posterior = ImportanceSamplingPosterior(
    potential_fn=vi_posterior.potential_fn,
    proposal=vi_posterior,
    method="sir",
).set_default_x(x_o)
refined_samples = refined_posterior.sample((1000,), oversampling_factor=32)

Check the importance weights before relying on the refined result. See the dedicated importance-sampling guide for details.

Sequential NLE with VI#

Fixed-observation VI can also act as the sampler in sequential NLE. In each round, train or update the NLE, fit a VIPosterior to x_o, sample parameters from it for the next simulations, append those simulations, and repeat. This is SNLE with VI instead of MCMC (Glöckler et al., 2022).

proposal = prior
for _ in range(num_rounds):
    theta = proposal.sample((num_sims,))
    x = simulator(theta)
    inference.append_simulations(theta, x).train()
    posterior = inference.build_posterior(
        posterior_parameters=VIPosteriorParameters(vi_method="fKL"),
    ).set_default_x(x_o)
    proposal = posterior.train()

Amortized VI#

Amortized VI learns a conditional approximation \(q(\theta \mid x)\) from simulation pairs, so one fit serves many observations. The variational family is set with q, which accepts the usual estimator strings:

amortized_vi = inference.build_posterior(
    posterior_parameters=VIPosteriorParameters(
        q="nsf", num_transforms=2, hidden_features=32
    ),
)
amortized_vi.train_amortized(theta, x)

Use sample() for one observation, sample_batched() for several observations at once, and log_prob() to evaluate the learned normalized approximation. There is no log_prob_batched() for VI posteriors, so evaluate one observation at a time.

samples = amortized_vi.sample((1000,), x=x_o)
batched_samples = amortized_vi.sample_batched((1000,), x=x_batch)
log_probs = amortized_vi.log_prob(samples, x=x_o)

Amortized VI currently optimizes the evidence lower bound (ELBO), which corresponds to reverse KL; the alternative vi_method choices do not apply. Check accuracy on representative held-out observations, for example with expected coverage or SBC.

Navigation