Choice of approximate posterior in variational inference with positive support

Question

I have a simple probabilistic graphical model: $z \longrightarrow x$ where $z_i \sim Exp\left(\lambda_i\right)$ where subscript $i$ denotes the $i$th dimension and $x|z \sim \mathcal{N}\left(f\left(z\right), \sigma^2\right)$ where $f$ is some real valued function (usually non-linear). I'm using variational inference to approximate the posterior $p\left(z|x\right)$. Let's say $q_{\phi}\left(z|x\right)$ is the approximate posterior with $\theta$ being the parameters.

What are some well-informed choices for $q$?

Since $z$ is a positive random variable, I cannot use the prototypical normal distribution assumption for $q$. Furthermore I also want the posterior to have some correlation structure like in a multivariate normal distribution.

My idea is to indeed use a multivariate normal distribution but only as the base distribution and add a positive, bijective transformation, $T$ on top such that my target distribution has positive support, retains a correlation structure and the bijective transformation renders a full likelihood computation. Basically: $q_{\phi}\left(z|x\right) \equiv \mathcal{N}\left(z|x\right)\left|\operatorname{det} \frac{\partial T}{\partial \mathbf{z}}\right|^{-1}$.

What are some good choices for $T$? And what are some other distributions I can use in general in that allow dependencies between dimensions, and have positive support? I believe I can also just choose univariate distributions as long as I also use copulas but I lack the knowledge on that and some advice will be much appreciated.

So far I have tried the exponential function as a transformation on top of multivariate normal, which is essentially a multivariate lognormal.

Finally, disclaimer: sorry that I'm cross posting this from mathoverflow where I posted the same question (https://mathoverflow.net/questions/431744/choice-of-approximate-posterior-in-variational-inference-with-positive-support), but I believe Cross Validated is a better forum for discussion for this topic. Please feel free to let me know otherwise if you strongly think so.

Answer 1

Very interested to hear others' thoughts on this.

One approach to choosing a prior distribution goes something like "If we're not sure which distribution to use and think the conjugate prior is reasonable and sufficiently flexible, then let's use that."

Going somewhat along this logic, it may make sense to use a $\Gamma$ distribution, a distribution which has a well known KL divergence with itself (keeping in mind that the exponential distribution can also be thought of as a gamma and using this CV post).

But we could just as well motivate a Lognormal distribution like that, as this CV post shows, which uses Mathematica to solve the expecations.

One advantage to doing this is keeping things simple. In numerical algorithms, the fewer moving parts the better if you're clumsy like me. A second advantage is that you'll have an analytic formula for what your penalty is, and understanding it may help you choose between the two based on your needs. (After all, KL-based Variational Bayes is just finding a distribution with maximum expected likelihood added with to a KL divergence between our variational and prior distributions). Finally, you'll have one less source of variance for the KL divergence/ELBO gradient estimation.