run_sbc won't sample from inferred Direct Posterior

[paarth-dudani]

Describe the bug
I am trying to establish an SBC pipeline. I inferred a few independent SNPE estimates using concurrent futures. However, the run_sbc sampler won't sample from my inferred Direct Posterior.

To Reproduce

1. Versions
   Python version: 3.9.13
   SBI version: 0.23.1

2. Minimal code:

num_sbc_samples = 200  # choose a number of sbc runs, should be ~100s

t = np.linspace(0,4,n)
input_array = t
initial_condition = 2.0

# 2D SBC
thetas = unifrom_beta_sigma.sample((num_sbc_samples,))
xs = simulator_2D(thetas, num_sbc_samples, t, initial_condition)

# results[0] is one of the 3 direct posteriors
num_posterior_samples = 1_000
num_workers = 1
ranks, dap_samples = run_sbc(
    thetas, xs, results[0], num_posterior_samples=num_posterior_samples, num_workers=num_workers

Code for context:

if __name__ == "__main__":
  with concurrent.futures.ProcessPoolExecutor(max_workers = 3) as executor:
    futures = [executor.submit(SNE_training_basic_multi_round, neural_net, proposal, simulator, observed_data, learning_rate, num_rounds, num_sim, initial_condition, hidden_features) for i in array_indices]

results = [futures[i].result() for i in range(len(futures))]

3. Error message
   No error message but the sampler won't sample:

Expected behavior
Fast sampling

[janfb]

Hi @paarth-dudani thanks for reporting this.

The run_sbc does nothing to the posterior than just sampling from the posterior for the sbc ranking. Two questions:

1. Have you checked that the posterior you obtained from your parallelization routine is "healthy", i.e., it's fast to sample / has no leakage etc?
2. Is the simulator_2d simulator the same one you used for training SNPE?

Cheers,
Jan

[paarth-dudani]

Hi @janfb

1. The posterior is indeed quick to sample from. I don't exactly know how to check for leakage. If you are referring to the warning given in the following link on the sbi website: FAQ, then I get no such warnings while training. The training is smooth.
2. Yes, it is the same simulator that I used for both training and the SBC procedure.

I get this for other parallel estimates as well. What should I do?

[janfb]

OK, good to know that posterior itself is fast to sample.

Then it's indeed something within the run_sbc method. Can you please try it with setting use_batched_sampling=False (True by default). I can imagen that because you have 2D data and are using an embedding network, passing the entire batch of 200 xs at once to the density estimator might be too much.

Does this help?

[paarth-dudani]

No, unfortunately it does not help.

Just as a clarification. My data is still 1D, it's just that I am trying to infer the standard deviation of the noise as well. This is the reason why my prior is 2D although my data is still 1D.

As a result, I have also not used an embedding network in this example.

As you can see, it takes too long and I don't see any output for minutes.

[patriglesias]

I am experiencing the same problem with the new version: No error message but the sampler won't sample

[janfb]

OK, interesting! Can one of you please share a minimal code example? Or at least some details on the use-case, the sbi method, the hyperparameters etc?

[michaeldeistler]

I have also just heard from another user that is experiencing this issue: slow sampling, but no warning.

[rdgao]

this was somebody else's issue that I tried to help with but will try to describe it here:

• NPE with MAF, 50 parameters (tried NSF as well)
• tried various dimensionality for summary statistics (1 up to ~400)

they report that at some dimensionality of parameter space (around 35), this problem arises where NPE trains, but is very slow to sample and no warnings. It does however produce samples, and is faster on higher end GPU, but still on the order of 10s per sample.

manually the sampling the density estimator via:

net = posterior.density_estimator
net.sample(...)

samples much faster, basically "normal" speed. But we did not check extensively the quality of the samples, i.e., if some subset of the 50 parameters are often out of prior bounds.

Not sure if its the same problem as the OP, since they report that sampling from the posterior is normal but SBC is slow, whereas here posterior sampling is slow.