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.

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?

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.

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

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!