Making the custom density estimator NPE/NLE/NRE procedure sequential

[paarth-dudani]

Following is my code for implementing a custom density estimator for NPE using a non-sequential approach:

from sbi.neural_nets.net_builders import build_maf

## Density estimator
# unifrom_beta_sigma = BoxUniform(low=torch.tensor([0.0,0.0]), high=torch.tensor([4.0,0.5]))
unifrom_beta_alpha = BoxUniform(low=torch.tensor([0.0, 0.0]), high=torch.tensor([4.0, 3.0]))

def SNE_training_made(num_epochs):
    simulator = partial(simulator_test, num_sim = 15000, input_array=t, initial_condition=initial_condition)
    train_data = NPEData(num_samples = 15000, prior = unifrom_beta_alpha, simulator = simulator) #using custom data loader
    train_loader = torch.utils.data.DataLoader(train_data, batch_size = 250)
    dummy_theta, dummy_x = next(iter(train_loader))
    maf_estimator = build_maf(dummy_theta, dummy_x, hidden_features = 50, num_transforms = 5, device = 'cpu')
    optw = AdamW(list(maf_estimator.parameters()), lr = 1e-4)
    loss_list = []
    for ep in range(num_epochs):
        for idx, (theta_batch, x_batch) in enumerate(train_loader):
            optw.zero_grad()
            losses = maf_estimator.loss(theta_batch, condition=x_batch)
            loss = torch.mean(losses)
            loss.backward()
            optw.step()
        if ep % 10 == 0:
            loss_list.append(loss.item())
            print("last loss", loss.item())
    return loss_list, maf_estimator

result = SNE_training_made(2500)

Of course, I would like to make this inference more specific for a dataset, x_o, as this uses a lot of simulations. However the primary sequential implementation documentation on the website does not allow for as much control, in terms of getting the per epoch loss, using custom data loaders, specifying the number of epochs etc.:

# For SNLE: likelihood_nn(). For SNRE: classifier_nn()
from sbi.neural_nets import posterior_nn

density_estimator_build_fun = posterior_nn(
    model="nsf", hidden_features=60, num_transforms=3
)
inference = NPE(prior=prior, density_estimator=density_estimator_build_fun)

posteriors = []
proposal = prior

for _ in range(num_rounds):
    theta, x = simulate_for_sbi(simulator, proposal, num_simulations=500)

    # In `SNLE` and `SNRE`, you should not pass the `proposal` to
    # `.append_simulations()`
    density_estimator = inference.append_simulations(
        theta, x, proposal=proposal
    ).train()
    posterior = inference.build_posterior(density_estimator)
    posteriors.append(posterior)
    proposal = posterior.set_default_x(x_o)

Can someone assist me in implementing the first algorithm but in a format that is sequential and is conditioned on the observed data set x_o?

[paarth-dudani]

Can someone provide an answer for this? I am approaching a deadline for my project soon.

[janfb]

Dear @paarth-dudani

please be aware that many of use are doing this work on the sbi package on the side as a community project. Of course, it is fine to write reminders when a question or issue remains unanswered for a while. But it can easily take a couple of days or a week or so, until someone in the community finds the time to have a look, and I would kindly ask you keep a friendly and respectful tone.

Regarding your question: For multi-round SNPE, it is less straight-forward to implement a custom training-loop because SNPE requires different loss functions or post-hoc correction to obtain an unbiased posterior estimate (because we are using a proposal different from the prior). For NRE and NLE, you could just wrap your training loop over epochs into a loop over rounds and run the training with the newly sampled data in every new round (because the proposal correction is not needed here). However, overall I would recommend just using our implementations to be on the save side. What is the reason you need low-level access to the training loop?

Cheers,
Jan

[paarth-dudani]

Dear Jan,

Thanks for your reply!
I apologize if it came across as unfriendly. I was under the impression that the queries are being more actively resolved and not just on the side. Regardless, I will be more mindful about my tone next time onwards.

As for the question: Here the link to the earlier discussion I had, which inspired my current approach:
#1306 (comment)

Could you elaborate a little on 'SNPE requires different loss functions or post-hoc correction to obtain an unbiased posterior estimate'? I believe this may be relevant to the following loss curves that I get from multi-round SNPE:

My main reason has to do with the fact that I find independent estimates from multi-round SNE (original implementation) a little unreliable in terms of the variability of the posterior when compared with the ground truth. A custom implementation lets me see more patterns in terms of convergence issues, and has enabled me to get pretty consistent independent estimates.

An idea I had in this regard was to first, train a network with the custom implementation (maybe with lesser number of simulations) and then wrap it as a direct posterior and use it as a proposal for multi-round inference.

I'd appreciate any suggestions in this regard!

Regards,
Paarth

[michaeldeistler]

I agree with Jan, implementing a custom training loop for SNPE-C is tricky (because of the different loss functions). However, it is quite easy to implement a custom training loop for TSNPE. You will have to do something like this:

from sbi.inference.posteriors import DirectPosterior
from sbi.utils import RestrictedPrior, get_density_thresholder

proposal = prior
maf_estimator = build_maf(dummy_theta, dummy_x, hidden_features = 50, num_transforms = 5, device = 'cpu')
optw = AdamW(list(maf_estimator.parameters()), lr = 1e-4)

for _ in range(rounds):
    # In the next round, sample theta from the proposal instead of from the prior.
    theta = proposal.sample((num_sims_per_round,))
    x = run_your_simulator(theta)

    # ...retrain the `maf_estimator`
    for epoch in range(epochs):
        optw.zero_grad()
        losses = maf_estimator.loss(theta_batch, condition=x_batch)
        loss = torch.mean(losses)
        loss.backward()
        optw.step()

    # After training, build the truncated proposal:
    posterior = DirectPosterior(maf_estimator, prior)
    accept_reject_fn = get_density_thresholder(posterior, quantile=1e-4)
    proposal = RestrictedPrior(prior, accept_reject_fn, sample_with="rejection")

With TSNPE, there is no need to modify the loss function or the rest of the training loop at each round.

Hope that helps
Michael

[paarth-dudani]

I see, this looks promising!

Thank you for your reply! I will try it out and revert back.

Could you elaborate a little on the changing loss function or point me to the appropriate text which highlights this in SNPE?

Regards,
Paarth

[michaeldeistler]

Here is the SNPE-C paper which describes the loss function.