Skip to content

How to choose sampling algorithms

sbi implements three methods: NPE, NLE, and NRE. When using NPE, the trained neural network directly approximates the posterior. Thus, sampling from the posterior can be done by sampling from the trained neural network. The neural networks trained in NLE and NRE approximate the likelihood(-ratio). Thus, in order to draw samples from the posterior, one has to perform additional sampling steps, e.g. Markov-chain Monte-Carlo (MCMC). In sbi, the implemented samplers are:

  • Markov-chain Monte-Carlo (MCMC)

  • Rejection sampling

  • Variational inference (VI)

  • Importance sampling (IS)

Which sampler should you choose? And how can you do this in the sbi toolbox?

Explicit recommendations

For NPE: - use the DirectPosterior. If you have a likelihood available, you can refine the posterior with importance sampling (see tutorial here)

For NLE or NRE: - If you have very few parameters (<3), use RejectionPosterior or MCMCPosterior - If you have a medium number of parameters (3-10), use MCMCPosterior - If you have many parameters (>10) and the MCMCPosterior is too slow, use the VIPosterior. Optionally combine the VIPosterior with an ImportanceSamplingPosterior to improve its accuracy.

Overview

  • MCMCPosterior: very accurate, but can be slow if you have many parameters.
  • VIPosterior: can be much faster if you have many parameter. May be inaccurate.
  • RejectionPosterior: accurate and fast, but only for very few parameters (typically less than 3).
  • ImportanceSamplingPosterior: typically inaccurate, but can be very useful to improve the accuracy of a VIPosterior (see above).

Main syntax

Below, we show how you can use these different sampling algorithms in the sbi toolbox. We begin with the MCMCPosterior:

from sbi.inference import MCMCPosterior, likelihood_estimator_based_potential

trainer = NLE()
likelihood_estimator = trainer.append_simulations(theta, x).train()

potential_fn, parameter_transform = likelihood_estimator_based_potential(
    likelihood_estimator, prior, x_o
)
posterior = MCMCPosterior(
    potential_fn, proposal=prior, theta_transform=parameter_transform, warmup_steps=10
)

If you want to use variational inference or rejection sampling, you have to replace the last line with VIPosterior or RejectionPosterior:

from sbi.inference import RejectionPosterior, VIPosterior

# For VI, we have to train.
posterior = VIPosterior(
    potential_fn, prior=prior, theta_transform=parameter_transform
).train()

posterior = RejectionPosterior(
    potential_fn, proposal=prior, theta_transform=parameter_transform
)

Finally, it is also possible to improve the accuracy of a VIPosterior with importance sampling:

from sbi.inference import ImportanceSamplingPosterior, VIPosterior

vi_posterior = VIPosterior(
    potential_fn, prior=prior, theta_transform=parameter_transform
).train()
refined_posterior = ImportancerSamplingPosterior(
    potential_fn, vi_posterior, oversampling_factor=32
)
samples = refined_posterior.sample((1000,))

At this point, you could also plug the potential_fn into any sampler of your choice and not rely on any of the in-built sbi-samplers.

Further explanation

The first lines are the same as always: 

inference = NLE()
likelihood_estimator = inference.append_simulations(theta, x).train()

Next, we obtain the potential function. A potential function is a function of the parameter \(f(\theta)\). The posterior is proportional to the product of likelihood and prior: \(p(\theta | x_o) \propto p(x_o | \theta)p(\theta)\). The potential function is the logarithm of the right-hand side of this equation: \(f(\theta) = \log(p(x_o | \theta)p(\theta))\)

potential_fn, parameter_transform = likelihood_estimator_based_potential(
    likelihood_estimator, prior, x_o
)

By calling the potential_fn, you can evaluate the potential:

# Assuming that your parameters are 1D.
potential = potential_fn(
    torch.zeros(1, num_dim)
)  # -> returns f(0) = log( p(x_o|0) p(0) )

The other object that is returned by likelihood_estimator_based_potential is a parameter_transform. The parameter_transform is a pytorch transform. The parameter_transform is a fixed transform that is can be applied to parameter theta. It transforms the parameters into unconstrained space (if the prior is bounded, e.g. BoxUniform), and standardizes the parameters (i.e. zero mean, one std). Using parameter_transform during sampling is optional, but it usually improves the performance of MCMC.

theta_tf = parameter_transform(torch.zeros(1, num_dim))
theta_original = parameter_transform.inv(theta_tf)
print(theta_original)  # -> tensor([[0.0]])

After having obtained the potential_fn, we can sample from the posterior with MCMC or rejection sampling:

posterior = MCMCPosterior(
    potential_fn, proposal=prior, theta_transform=parameter_transform
)
posterior = RejectionPosterior(potential_fn, proposal=prior)

On the usage of NPE

NPE usually does not require MCMC or rejection sampling (if you still need it, you can use the same syntax as above with the posterior_estimator_based_potential function). Instead, NPE samples from the neural network. If the support of the prior is bounded, some samples can lie outside of the support of the prior. The DirectPosterior class automatically rejects these samples:

from sbi.inference import NPE, DirectPosterior

inference = NPE()
posterior_estimator = inference.append_simulations(theta, x).train()

posterior = DirectPosterior(posterior_estimator, prior=prior)