As I understand it, I'm basically minimizing the KL Divergence between the Prior and Encoder latent distribution, and the log probability of the decoder distribution. I have a model that does generate some useful latent distribution for classification using this loss function, however, it has weird learning dynamics. My decoder distribution is a 300 dimension Gaussian and my latent is a 40 dimension Gaussian, so in all instances, both terms of my ELBO should vary wildly in order of magnitude. It seems from looking around on the internet, most use cases of VAEs look similar - high dimension decoders and low dimensional encoders, so I would imagine some people see a similar a problem, but I haven't mentioned any note of this in any literature or tutorial. Is there any way to correct for this?
1 Answer
There is a nice paper that discusses this problem of named as 'Modelling Bias' in VAE's. ELBO objective is:
$$L_{ELBO}(x) = E_{q_φ(z|x)}[\log p_θ(x|z)] − D_{KL}(q_{φ}(z|x) || p(z))$$
The first term forces the VAE to put all the mass of $q_φ(z|x)$ the mode of $p_θ(x|z)$ while the second term forces it to be as diffuse as possible. But generally for images, the dimensionality of X is often orders of magnitude larger than the dimensionality of Z.
Because the same per dimensional modeling error incurs a much larger loss in X space than Z space, when the two objectives are conflicting (e.g., because of limited modeling capacity), the model will tend to sacrifice divergences on Z and focus on minimizing divergences on X. Hence it may lead to over-fitting rather than generalization.
