User defined ELBO in AdvancedVI or Turing

Specific Domains Statistics
1

Hi all, I am fairly new to ADVI and Turing (and sub-packages in general). I get the general idea of ADVI in that we make a set of parameter transformations to obtain an estimator to the intractable expectation

\mathbb{E}_{q\left(\theta \middle \vert \phi\right)}\left[\ln{P\left(\mathcal{D} \middle \vert \theta \right)P(\theta)}\right]

If my understanding is correct, AdvancedVI and Turing obtains an estimator over the entire unnormalized posterior in the real space, i.e. does not obtain analytical expression where apllicable? If this is this case, would it possible to specify some user defined function which computes the ELBO and just uses the optimization algorithms and other functionality of AdvancedVI/Turing? An example would be if the expectation of the log prior is tractable, hence only requiring a stochatic estimate for the expectation of the log likelihood.

\mathbb{E}_{q\left(\theta \middle \vert \phi\right)}\left[\ln{P\left(\mathcal{D} \middle \vert \theta \right)}\right] \approx \frac{1}{N} \sum_{i=1}^{N} {\ln{P\left(\mathcal{D} \middle \vert \theta_{i} \right)}} \\ \mathbb{E}_{q\left(\theta \middle \vert \phi\right)}\left[\ln{P(\theta)}\right] = f(\phi)

Where f(\phi) is the a tractable solution for the expectation of the log prior. Here we would estimate the expectation only for the log likelihood, making the parameter transformation and not include the jacobian adjustement. I would expect that the estimate we obtain in this case to be more accurate and computationally efficient. Is there a way in which one can do this using the current functionality of AdvancedVI or Turing? I.e. specifying a user defined ELBO directly?

Any help will be much appreciated!

2

Hi @Garren_Hermanus ! There are few things to consider here. While it is possible to use your own variational objective with closed form expectations as you described, it might not be the best of ideas. Firstly, the standard way of implementing ADVI is to rely on the reparametrization gradient. This restricts the types of variational families that you can use. Typically, Gaussians with appropriate support transformations like exp for positive random variables. These unfortunately rarely have closed for expectations against distributions typically used as priors. If you insist on using variational families that do have closed form expectations, you can use the score gradient estimator instead, but this has much higher variance. So the benefits of closed form expectations might end up being counter-balanced.