NaN Samples in SNLE with slice_np_vectorized

[pardocelsa]

Hi,

I am trying to infer the posterior distribution using SNLE, running 3 rounds with 10K simulations in the first round and 1K in subsequent rounds. In each round, the proposal prior is set to the approximated posterior at the observed sample from the previous round.

I successfully inferred the same parameters for the same model using TSNPE without encountering any major issues. However, with SNLE, although the posterior distribution yields results similar to those from TSNPE, some of the samples turn out to be NaN when using both a normalizing flow and an MDN with the slice_np_vectorized MCMC sampler. On the other hand, when I switch the MCMC sampler to slice_np, everything works, though it runs much slower.

Code Snapshot for snle

inference = SNLE(
    device=f"{device}",
    prior=prior,
)

density_estimator = inference.append_simulations(
    [parameter_round.to](http://parameter_round.to/)(device), [matrix_round.to](http://matrix_round.to/)(device)
).train(
    learning_rate=config["trainer"]["lr"],
    training_batch_size=config["trainer"]["batch_size"],
    validation_fraction=config["trainer"]["validation_fraction"],
    show_train_summary=True,
    retrain_from_scratch=retrain_from_scratch,
)

posterior = inference.build_posterior(
    mcmc_method="slice_np_vectorized",
    mcmc_parameters={"num_chains": 20, "thin": 5},
)

posterior_obs = posterior.set_default_x(x_o)
proposal = posterior_obs

observed_samples_posterior = posterior_obs.sample(
    (10000,), show_progress_bars=True
).cpu()

I am running this process for three rounds, but in the second round, some (~2 out of 10.000) posterior samples in observed_samples_posterior are NaN when using slice_np_vectorized.

When trying to understand what is going on I have observed the following:

1. Training Curve Behavior

• When using an MDN, the training and validation curves exhibit an unusual sharp drop to lower loss values at some point:

• When using a normalizing flow, the loss function sometimes drops to negative values:

2. Negative Log-Likelihood for Some Training Data Points

When I inspect the log-likelihood of the training dataset for both cases (mdn and flow) using the following code:

inference_path = path_round_2_exp_2 + '/inference.pickle'

with open(inference_path, "rb") as f:
    inference = pickle.load(f)

training_ln_likelihood = inference._neural_net.log_prob(matrix_round, context=parameter_round)

I observe negative log-likelihood values for some training data points:

tensor([  -9.9781,   -3.8988,   -2.6384,   -7.6448, -290.4689,   -2.7224,
          -0.9067], device='cuda:0', grad_fn=<IndexBackward0>)

My guess is that during MCMC sampling, the likelihood for certain values is extremely close to zero, which might cause the sampled values to become NaN.

Questions

1. Does this explanation seem correct?

   • Could the near-zero likelihood values during MCMC sampling be causing NaN samples?

2. Should I trust my results despite this behavior?

   • Given that I have early stopping, would it be reasonable to simply remove the NaN samples from my posterior distribution and still use the results or should I switch to the slice_np sampler instead of the vectorized version?Additionally, the results remain consistent with those obtained using TSNPE.
   • Or does the observed instability in the training curves and posterior samples indicate that the results are unreliable? If so, do you have any advice on how to mitigate this issue?

Any insights or suggestions would be greatly appreciated!

Thanks in advanced,
Celsa.

[michaeldeistler]

Thanks for creating this!

We remember we had experienced NaN samples with MCMC at some point. In all of these cases, I don't think they ever caused performance issues. As such, I tend to believe that removing the NaN samples still good and reliable results. I would recommend to check though, e.g. with SBC or with MCMC diagnostics.

I am not sure if the "near-zero likelihood" is the cause for this, but it would be possible.

The training curves look good imo. The negative values are likely a result of overfitting in these particular epochs. The neural net seems to recover quite well though, so I think all is good there.

Michael

[pardocelsa]

Thanks, @michaeldeistler, for the quick response. To check the results in this case, where I'm using NLE, the SBC procedure will likely take a long time, right? In my case, sampling the posterior distribution 10,000 times for a given data point (x) takes about 7 minutes when using the vectorized sampler. From my experience, I would need at least 300 test samples to perform rank tests or coverage tests, which would require roughly 35 hours of computation for this experiment. Do you actually perform SBC when using NLE? I did so when using TSNPE, but I thought that with NLE, people generally don’t use it because of computational reason.

I will definitely try the MCMC diagnostic, I wasn’t aware of that approach. I will also perform SBC for the benchmark experiment. Thanks again!

[michaeldeistler]

Yes running SBC will be expensive for NLE. It's not impossible though: the newest sbi versions use sample_batched within SBC to vectorize the MCMC sampler also across observations. This might make it scale better, so it could be worth a try.

[pardocelsa]

Okay, perfect! I didn't know about this new feature of SBI, I will definitely try it. Thanks for all the help!