Append_simulations and multiround

[pardocelsa]

Hi,

I was digging into the SBI code, and I realized that internally, in a multi-round SBI context, each round reuses simulations from all previous rounds by default. However, I thought this was not possible in SNPE, only in TSNPE. My reasoning is that in SNPE the proposal prior changes in every round, so the posterior needs to be corrected with the proposal prior of that round. In contrast, in TSNPE the proposal prior is simply c × prior restricted to the support of the posterior, so this correction step is not needed.

Or is it in fact the case that, when computing the loss for SNPE, you also include samples from previous rounds, applying the appropriate loss correction for each round’s proposal? For instance, if we are in round 2, would we compute the loss using the standard log-likelihood for the samples pooled from the first round, while for the samples from the second-round proposal we would use the corrected posterior with that proposal?

Thanks in advance for your help,
Celsa

[michaeldeistler]

Hi! All SNPE methods except SNPE-A allow re-using simulations.

Or is it in fact the case that, when computing the loss for SNPE, you also include samples from previous rounds, applying the appropriate loss correction for each round’s proposal?

Yes, that's exactly what we do for SNPE-B and for non-atomic SNPE-C.

Notably, atomic SNPE-C does not even require knowledge of the proposal distribution.

All the best
Michael

[pardocelsa]

ah ok, thanks a lot for all the info!

[pardocelsa]

Sorry @michaeldeistler, but another question related to this just came to my mind: how does sbi know whether I want to use SNPE-A, SNPE-B, or SNPE-C? Also, when using a truncated prior, I assume that means I’m using SNPE-A trained with maximum likelihood but with a truncated prior. If that’s the case, would I need to append the simulations myself, or does sbi handle it automatically?

This is the piece of code I am using:

posterior_nn_args = {
        "model": 'mdn',
        "hidden_features": config["density_estimator"]["args"][
            "hidden_features"
        ],
        "num_components": config["density_estimator"]["args"][
            "num_components"
        ],
        "device": device,
      }

embedding_net = config.init_object("arch", learning_models)

weight_initializer = config.init_object(
  "weights_initializer", learning_initializers
)
embedding_net.apply(weight_initializer)
posterior_nn_args["embedding_net"] = embedding_net
neural_posterior = utils.posterior_nn(**posterior_nn_args)

inference = SNPE(
  density_estimator=neural_posterior,
  device=f"{device}",
  prior=prior,
)

train_args = {
    "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,
}

if (
    config["trainer"]["truncated_prior"]
    and model_type == "snpe"
):
    train_args["force_first_round_loss"] = True

kwargs = {}
if (
    model_type == "snpe"
    and not config["trainer"]["truncated_prior"]
):
    kwargs["proposal"] = proposal

density_estimator = inference.append_simulations(
    parameter_round.to(device),
    matrix_round.to(device),
    **kwargs,
).train(**train_args)

[michaeldeistler]

If you do SNPE (or NPE), then it defaults to SNPE_C. You can also make it more explicit:

from sbi.inference import SNPE_A, SNPE_B, SNPE_C

inference = SNPE_C(prior)
# Same for other methods

However, in the first round, all SNPE_... methods use the same loss. They only differ in later rounds.

For truncated proposals, please see the documentation on implemented methods. You have to set the .train(..., force_first_round_loss =True) argument. Here is a full code snippet:

from sbi.inference import NPE
from sbi.utils import RestrictedPrior, get_density_thresholder

inference = NPE(prior)
proposal = prior
for _ in range(num_rounds):
    theta = proposal.sample((num_sims,))
    x = simulator(theta)
    _ = inference.append_simulations(theta, x).train(force_first_round_loss=True)
    posterior = inference.build_posterior().set_default_x(x_o)

    accept_reject_fn = get_density_thresholder(posterior, quantile=1e-4)
    proposal = RestrictedPrior(prior, accept_reject_fn, sample_with="rejection")

Hope that helps!
Michael

[pardocelsa]

Yes, I have used this documentation and added force_first_round_loss=True to the training function. Thanks for the clarification on SNPE-B and SNPE-C.

[michaeldeistler]

Great!