Embedding Net for Theta (Simulation Inputs) in SNPE or Embedding Net for X (Simulation outputs) in SNLE

[rmehta1987]

🚀 Is it possible to added embedding_nets for both theta and X in SNLE/SNPE

Is your request related to a problem?

Sometimes the simulation inputs are much larger that the simulation outputs, and I would like to have summary statistics for the inputs (theta) rather than the outputs (x). Currently I can do this separately by just training summary statistics of the simulation inputs separately and a prior on the these statistics rather than the simulation inputs themselves. But when training that means we optimizes the loss on learning the simulation input embeddings rather than the simulation inputs themselves.
...

Describe the solution you'd like

Currently I just modified the build_fn methods to include embedding nets for both theta and x, but it is kind of hacky, so if possible a more formal way to do this.

...

Describe alternatives you've considered

Just using the summary statistics of the simulation inputs as the input to SBI or modifying the posterior_nn function to include embedding nets for x and theta.

...

[janfb]

Hi @rmehta1987 thanks for raising this question. In general, when performing conditional density estimation, we can only use embeddings for the conditioning variable, i.e., we can use embeddings for x if we do NPE, and we can use embeddings for theta if we do NLE. We generally cannot embed the variable over which we learn the density.

However, NRE is an exception. Because in NRE we perform classification to estimate the ratio, we can actually use embeddings for both theta and x (this should also be reflected in the build function classifier_nn). Thus, for your use-case I would recommend using NRE.

[rmehta1987]

Hi @janfb thank you for the response. While I agree NRE can be used in this situation, it's performance in inference seems to be lower than NPE. In another way you can think of using the NPE with allowing embeddings for both theta and X is that it is a form of meta-learning where the outer network learns the embedding for the inputs and then is frozen after which the NPE is trained. When the loss of the NPE is created, the loss is not only the logprobs of the NPE but also the loss of the outer network.

[janfb]

I am not sure I understand how this would work in the NPE setting. If you apply an embedding to the input variable, how would your posterior be defined in that case? Would it be a conditional density over the "raw" input theta? how would be corresponding density estimator be defined? How would you sample from it?

Internally, we are constructing a normalizing flow on theta with x as the conditioning variable. Are you suggestion to learn the flow over theta entirely in an embedded space and learn the theta-embedding somehow in parallel and then apply it to the flow to sample from the posterior?

[rmehta1987]

Currently the way NPE is constructed theta or X have to be 1-D vectors. When theta or X are not 1-D vectors the suggestion is to either use NRE if both need to be embedded, NLE if theta needs to be embedded, or NPE if X to be embedded. In the case of NLE or NPE it is possible to learn the embedding in parallel as it is more of a context variable rather than a part of the normalizing flow so the noise is transformed to theta or X conditioned on the embedding.

I would like to do as you suggested it so learn the flow of theta in embedded space and learn the embedding in parallel. Currently I learn the embedding separately, then sample from the posterior P(embedding_theta | X), then reconstruct the embedding to a P(theta_hat). However this means that flow architecture must represent the transformation of the embedding to noise very well to retrieve an accurate reconstruction P(theta_hat) given a "ground truth" X_o Therefore it seems more plausible to train them in parallel.

Ideally it would be nice to learn on the "raw" theta, but since theta is multidimensional (i.e. theta is a multidimensional time series), I do not think that is possible as of yet.

Another idea is to use your recent method of Simformer since it models the joint distribution P(theta, X) akin to NRE where the transformer would learn the embedding of the joint distribution, but I am having a hard time understanding how it (and the code) exactly works.

[janfb]

I don't think I fully understand how it would be possible to learn theta in embedded space and learn the embedding separately. When the using the conditional flow architecture, you need to perform the conditioning during training in order to introduce the connection between the input and the condition. One way I can imagine would be training an encoding decoding architecture on theta separately, then do standard NPE in the embedded space. But then you need to perform a decoding step from the posterior over the encoded theta to the original theta space, which might be difficult to interpret conceptually and mathematically.

Using the SIMFORMER architecture instead makes more sense (btw, I don't think it is technically related to NRE). So I would recommend having a closer look at the paper, and the code here: https://github.com/mackelab/simformer

[rmehta1987]

One way I can imagine would be training an encoding decoding architecture on theta separately, then do standard NPE in the embedded space. But then you need to perform a decoding step from the posterior over the encoded theta to the original theta space,

That is what I am currently doing, and you are correct it is hard to interpret what the posterior over the embedded space represents. The inference performs well on simulated test data, as it the decoder is able to reconstruct the correct simulation parameters, but it kind of seems hand-wavy in a way of interpreting results.

I will take a look at simfomer again, the codes is just hard to follow in it's current state.

Thanks for the discussion!