I am learning about variational inference and am implementing a couple of things from scratch. I am trying to build a Gaussian mixture model where the prior on the mixture component selection is a Dirichlet. That's fine.
However, I need a variational distribution that will produce values that are reasonable for the Dirichlet prior, namely, they need to be positive. My current implementation is just a Dirichlet distribution without any constraints on its parameters. This works sometimes when the optimization produces variational parameter values that are positive. However, there is no constraint on these variational parameters, so when they become negative, nothing works anymore.
My initial idea was to just exponentiate the variational parameters since this is what we can do when using a Gaussian variational distribution. To get the variance, we just exp(.) the variance. However, things stopped working after that. I suspect the reason behind this is that exp will explode one value and the other can't seem to keep up, and all the probability mass will go on a single mixture component. Obviously, this is not okay.
One approach I have found is from ADVI where the variational distribution is just a bunch of univariate Gaussians. We sample from these Gaussians, and softmax them when comparing them to the Dirichlet prior. But here, we have to include a Jacobian term, which is cumbersome.
There must be a better way to solve this. Maybe through some better parameterization of the Dirichlet as the variational distribution? Any help would be appreciated!