Is it possible to use SNLE with multi-dimensional/image data?

[msainsburydale]

I'd like to implement neural variational inference (NVI) via the SNLE approach listed here. However, my data is gridded (an image), and so I receive the error message:

AssertionError: SNLE cannot handle multi-dimensional simulator output

The error is clear, but I wanted to double check that I'm not missing something. Are there any workarounds (like embedding_net_x when using the SNRE method) that would allow one to use SNLE with multi-dimensional data?

[janfb]

Hi @msainsburydale, SNLE does not allow the use of embedding nets for x because it learns a density in x. It can only embed the parameters theta.

Is there a specific reason you want to use NVI? If not you could just use NPE which directly learns the posterior and allows embeddings in x. NPE is also fully amortized, i.e., after training you can obtain posterior estimates instantly without having to run VI or MCMC.

Alternatively, you could use NRE which allows to embed both, theta and x. With NRE you could then use VI or MCMC to obtain posterior samples.

I hope this helps.
Best,
Jan

[msainsburydale]

Hi Jan,
Many thanks for you answer.

I had thought that SNLE() might allow for x to be summarised into statistics S(x) by an embedding net, so that the density is learned on S(x), like in the "synthetic-likelihood" approach. But perhaps this isn't possible. For my specific problem with image data, I ended up using hand-crafted summary statistics.

I'd like to use amortised NVI as part of a comparison study with several other amortised likelihood-free methods. I've implemented amortised NRE (using SBI, which works very well by the way) and NPE, but not yet amortised NVI. I saw in the SBI docs that there is variational inference, but that a second training stage needs to be done for every new observed data x_o (i.e., it is not amortised). Is there any reason for this? I had thought it would be possible to:

1. Approximate the likelihood function using an amortised likelihood estimator (e.g., using SNLE()).
2. Construct an amortised posterior (e.g., using SNPE()) based on minimising the expected reverse KL rather than the expected forward KL, and with the amortised likelihood estimator above used in place of the true likelihood when computing the reverse KL.

Is it possible to do this in SBI?

Many thanks again for your help.
Best,
Matt

[michaeldeistler]

Amortized variational inference is not implemented, see #909. There is also no paper that did this (to my knowledge) so I am not certain how well it would work. So, as of now, if you want to use NVI, you will have to retrain the variational posterior for every observation.

[msainsburydale]

Great that there are plans to implement amortised VI. A paper that does something similar is Glöckler et al. (2022), which is also referenced in the SBI package documentation. Figure 1 and Section 3.1 describe how an amortised likelihood function is learnt and then subsequently used when optimising a variational posterior parameterised by a normalising flow (the likelihood is needed since it appears in the objective function when training a variational posterior).

As far as I understand, there's no reason that a variational posterior can't also be trained in an amortised fashion, thereby ensuring that the entire workflow remains amortised. I think this is motivation for #909, which is great.

Thanks for your help and the excellent software.

Cheers,
Matt

[michaeldeistler]

Hi Matt,

indeed. Just to clarify: as in Glöckler et al., sbi allows to train an amortized likelihood. But also like Glöckler et al., sbi does not allow to train the variational posterior in an amortized fashion (this is meant by #909)

Michael