As pyro currently stands, if we use the standard Trace_ELBO for SVI, the ELBO (loss) is estimated using n=num_particles samples from the guide. These are drawn by calling the guide function n times with the same parameters.
But what if the guide is doing significant/slow computation based on its parameters, that could be reused? It can’t simply cache the results, because that would break their gradient connection to the parameters. So, I believe that there should be a way for a guide to (optionally) signal to pyro that it can draw multiple samples in one call, and then a way for pyro to ask it to do so.
For example, I’m currently working on demonstrating the concept of “Laplace family guides”; multivariate normal guides which are parametrized by their mode, but whose precision matrix is calculated by taking the Hessian (observed information) of the model itself. Torch is slow at computing Hessians, because it’s optimized for backwards-mode, not forwards-mode, automatic differentiation. But once you have a Hessian (and have done the extra tricks necessary to ensure it’s positive definite, and to get a Cholesky decomposition), it would be trivial to draw multiple samples from the resulting multivariate normal.
To be clear, I’m talking about drawing multiple samples from the full guide, not just multiple samples from an individual distribution. Each “sample” would be a full set of model parameters, suitable for conditioning the model on, separately from other “samples”. I suppose that as a temporary workaround, I could rewrite the model to be a model over multiple samples of different parameters… but that’s pretty much a hack. It would definitely better if this were possible within pyro.
I’ve already had to write a patch for pyro, for a different reason relating to my problem (to get it to add a gradient from a detached copy of a variable back onto the original variable’s gradient, just before optimization). So if there’s a quick-and-dirty place I could get this working for my own purposes by digging into pyro, I’m not afraid to do so, and I’d appreciate pointers. But more generally, I’d like to give back to pyro, and help this make it into the main code base — after I’ve finished defending my dissertation.