How to use permutation-invariant embeddings for iid observations (for NPE)

Important: This how-to guide is only relevant if you have iid observations, are planning to use NPE, and want to use the default sample_with="direct" method. However, standard NPE can natively handle iid observations if you use sample_with="mcmc" or sample_with="vi". NLE and NRE can also naturally deal with iid observations and do not need to be modified. See also this how-to guide on using iid observations with NLE or NRE.

In many cases, you want to estimate a parameter set given multiple observations. To use NPE in these scenarios, we advise that you use a PermutationInvariantEmbedding. In sbi, this can be done as shown below:

Main syntax

from sbi.neural_nets import posterior_nn
from sbi.neural_nets.embedding_nets import FCEmbedding, PermutationInvariantEmbedding

latent_dim = 10
single_trial_net = FCEmbedding(
    input_dim=theta_dim,
    num_hiddens=40,
    num_layers=2,
    output_dim=latent_dim,
)
embedding_net = PermutationInvariantEmbedding(
    single_trial_net,
    trial_net_output_dim=latent_dim,
    num_layers=1,
    num_hiddens=10,
    output_dim=10,
)
density_estimator = posterior_nn("maf", embedding_net=embedding_net)
trainer = NPE(density_estimator=density_estimator)

What to do if you have varying numbers of iid trials?

If you have varying numbers of trials across multiple observations, pad the missing simulations with NaN.

Example

For a detailed explanation and example, see this tutorial.

Citation

@article{radev2020bayesflow,
  title={BayesFlow: Learning complex stochastic models with invertible neural networks},
  author={Radev, Stefan T and Mertens, Ulf K and Voss, Andreas and Ardizzone, Lynton and K{\"o}the, Ullrich},
  journal={IEEE transactions on neural networks and learning systems},
  volume={33},
  number={4},
  pages={1452--1466},
  year={2020},
  publisher={IEEE}
}