Help navigating implementation of new inference algorithm

I think you should prioritize implementing it in PyMC v4, since this is new feature back porting it to current major release is very necessary.

Now to the approximation itself, it is a pretty interesting algorithm so I spent some time to draft out some components - this should help you implement it.
Some notes:

  1. Forward sampling is more straightforward in PyMC v4, although conditioning is a feature we are working on (see do operator / conditioning · pymc-devs/pymc · Discussion #5280 · GitHub).
  2. Getting the log likelihood function that can condition on new x_sim takes a bit of effort (thanks @ricardoV94 for guiding me how to do this correctly in v4), tldr is that you need to point to the correct random variable in PyMC
  3. Inference is done in the transformed space (unbounded), so we need a way to link the original space and unbounded space for forward sample.

I use your Neal’s Funnel example, details are in notebook: https://github.com/junpenglao/Planet_Sakaar_Data_Science/blob/main/PyMC3QnA/discourse_8528%20(MUSE).ipynb

(It is not quite work yet, I guess it depends on how to specify and update the Hessian H)

Side question on the approximation: while your paper discuss a bit of the parallel and differences with Laplace Approximation, I am wondering the comparison between MUSE and INLA (which try to do Laplace Approximation to Integrate z). INLA also require the latent field z being Gaussian like, and I find that in MUSE you stated “Even for mildly non-Gaussian latent spaces, one expects the data dependence of the integral to be small, with most of the data-dependence instead captured by the MAP term.” - do you have any comment on wildly non-Gaussian latent space (e.g. mixture)?

3 Likes