Improve sampling speed for SNLE/SNRE

[Saheli2001]

I am using SNLE() to estimate the posterior with a 2-d uniform as a prior density. Used 2000 simulations to generate (theta, x ). But while generating sample, it's taking too much time. I know that SNLE uses MCMC to generate the theta's which is why probably it's taking too long but is there a efficient way to generate sample from the posterior? only using SNLE?

[michaeldeistler]

Hi you can either use other sampling methods (variational inference or rejection sampling) or you customize MCMC to improve sampling speed. See also here.

1. Use VI:

inference = SNLE(prior)
_ = inference.append_simulations(theta, x).train()

sampling_algorithm = "vi"
vi_method = "rKL"  # or fKL
posterior = inference.build_posterior(sample_with=sampling_algorithm, vi_method=vi_method)

# Unlike other methods, vi needs a training step for every observation.
posterior = posterior.set_default_x(x_o).train()

2. Use rejection sampling

# Sampling with rejection sampling
sampling_algorithm = "rejection"
posterior = inference.build_posterior(sample_with=sampling_algorithm)

3. Customize MCMC. Use more chains and a vectorized sampler.

from sbi.inference import likelihood_estimator_based_potential, MCMCPosterior

inference = SNLE(prior)
likelihood_estimator = inference.append_simulations(theta, x).train()

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

sampler = 'slice_np_vectorized' # default is 'slice_np'
num_chains = 100 # default is 1
posterior = MCMCPosterior(
    potential_fn, proposal=prior, theta_transform=parameter_transform, method=sampler, num_chains=num_chains,
)