Speeding up training times for SNPE.

[paarth-dudani]

Hello community!

So I have been running SNPE on very high dimensional models and have been getting successful results like this:

Red shows MCMC ground truth and blue is the SNPE estimate using MADE.

While this is all good, I want to now address the problem of speeding up the inference. In particular, my aim is to infer the joint posterior using SNPE faster than that through MCMC (currently MCMC is faster).

Can someone suggest ways in which the neural network training and convergence can sped up as a part of the sbi pipline?

I use the following function to train to wrap up the sbi training pipeline:

def SNE_training_basic_multi_round(neural_net:str, embedding_net, proposal, simulator, observed_data, 
                                    learning_rate, num_rounds, num_sim, initial_conditions, 
                                    hidden_features, num_transforms, sigma, Simulator_Repressilator_noisy, hidden_size):

    input_size = 300
    output_size = 60
    hidden_size = hidden_size
    embedding_net = RNN(input_size, hidden_size, output_size)

    neural_posterior = posterior_nn(
            model=neural_net, embedding_net = embedding_net, hidden_features=hidden_features, 
            num_transforms = num_transforms, num_components = num_transforms)

    inference = NPE(prior=proposal, density_estimator=neural_posterior)
    proposal = proposal
    if num_rounds > 1:
      for _ in range(num_rounds):
          print('yo')
          theta = proposal.sample((num_sim,))
          x = simulator(theta, t, initial_conditions, sigma, Simulator_Repressilator_noisy)
          _ = inference.append_simulations(theta, x, proposal=proposal, exclude_invalid_x=True).train(learning_rate = learning_rate, training_batch_size = 5000)
          posterior = inference.build_posterior().set_default_x(observed_data)
	  proposal = posterior
    else: 
      theta = proposal.sample((num_sim,))
      x = simulator(theta, t, initial_conditions, sigma, Simulator_Repressilator_noisy)
      _ = inference.append_simulations(theta, x, proposal=proposal).train(learning_rate = learning_rate, training_batch_size = 5000)
      posterior = inference.build_posterior()
      proposal = posterior
    return proposal

For reference, the estimate shown above took ~6.5 hours to train (2 rounds, 35k simulations) on a high performance cluster with 100CPUs and and 100GB of memory. This currently precludes the accessibility of these algorithms.

[paarth-dudani]

Any thoughts or ideas anyone?

[janfb]

This looks promising, but the SNPE posterior is still a bit too broad compared to the MCMC reference.

I am wondering why you are using MADE? Have you tried MAF or NSF?

What kind of embedding network are you using for the high-dimensional data?

Also, I am wondering why you are using sequential NPE in the first place. When you want to beat MCMC in terms of speed, doing amortized NPE would be the ideal solution. This might need more simulations, but given that have a cluster with 100 CPU cores, here is what I would do:

• pre-simulate as many simulations as possible, parallelized over 100 CPUs, e.g., using simulate_for_sbi. Even if one simulation takes 10 seconds, this would allow you to generate around 200k simulations in 6h. If you let it run for one day, you will have plenty of simulations from the prior that you can just save on disk and later use for training and testing.
• train single-round NPE on those (200k) simulations, using nsf
• if GPU memory is a problem, use the custom training loop to do the data handling yourself.
• run SBC to check the calibration
• compare the posterior against MCMC

If the resulting posterior comes close MCMC, then you will be able to do inference instantly, i.e., way faster than with MCMC.

[paarth-dudani]

MAF has proved to be very slow to train. I have also been doing hyperparameter optimization using grid search to tune some of the hyperparameters and MAF is not amenable to that since most of the proposals are not accepted. Could you tell me how to sample using MCMC from a trained direct posterior object? I did it earlier, but can't seem to find it again elsewhere.

I have used an RNN to embed the data so far. But I am keen on exploring other alternatives too. I would appreciate any suggestions.

You are correct in your speculation that amortization would be more advantageous to improve the inference time. However, I had gotten some arbitrary estimates from NPE, which is why I wanted to explore whether SNPE gets faster with higher dimensions (which it does). That said, NPE does show promising results and that's one of my main streams of exploration currently.

As for your points:

1. I did try this. Not just hours, I am able to obtain 100k simulations within 30 minutes. The rate limiting step is the training of the network on these simulations, which is where my current obstacle is.
2. I will try amortization with nsf, surely.
3. The custom training loop dramatically slows down the inference. I can revisit it to see if there were some errors at my end.
4. SBC checks have been somewhat arbitrary in the past, which is why I didn't use them. Again, I can revisit this and see if there is a consistent way for me to apply it.

Thank you for your suggestions! If you have a specific suggestion to parallelize or speed up the neural network training of MDN/MADE/MDN/NSF itself, given the simulations, that would be of additional help as I am soon wrapping up my project.

Do you have any further suggestions for benchmarking, given that I have access to the ground truth over the parameters?

Thank you for your reply Jan!

[paarth-dudani]

Hello!

Could someone tell me how to sample using MCMC from a trained direct posterior object? I can't seem to find a way to do it.

[michaeldeistler]

Something like this should work:

from sbi.inference.posteriors import MCMCPosterior
from sbi.inference.potentials import posterior_estimator_based_potential

inference = NPE(prior)
posterior_estimator = inference.append_simulations(theta, x).train()
potential, tf = posterior_estimator_based_potential(posterior_estimator, prior, x_o)
posterior = MCMCPosterior(
    potential,
    proposal=prior,
    theta_transform=tf,
    num_chains=50,
    thin=1,
    method="slice_np_vectorized"
)