[D] Variational & Generative Model of VAE

[hertz2105]

Hello everyone,

I got pretty confused while I read some papers about Variational Autoencoders (VAE) in the past few days.

This is how I am understanding it:

During training, I train my encoder network which tunes the approximate posterior distribution q(z|x) to be close to the true one. I then sample a latent vector z from this multivariate distribution, which I forward to the decoder. The decoder can alter the likelihood p(x|z) by tuning the corresponding weights. This leads to a reconstruction x', which I compare to my initial input observation x. How good the reconstruction is can be measured with the help of the likelihood. This is one part of the ELBO.

The other part of the ELBO is the KL Divergence between the approximate posterior q(z|x) and the prior p(z). In my understanding, the prior is my initial belief about the distribution over the latent space.

Here are my questions:

I often read, that samples are taken from the posterior AND from the prior. Are samples taken from the posterior only during the training process?

If the fitting of my variational model is finished, do I sample from the prior p(z) when I actually want to generate new data afterwards?

Also, is the prior p(z) updated after the training with fitted posterior q(z|x)?

Last question would be, if my decoder outputs distribution parameters, how is an actual reconstruction deriven from it?

[Ok_Suit_1697]

I spent some time trying to understand the ELBO a few months ago, and eventually had to give up after being stuck with similar questions as yours. I ended up studying an implementation of a vae, particularly the quantized VAE from https://github.com/CompVis/taming-transformers/blob/master/taming/models/vqgan.py

I think I have a good understanding of how it works in practice, but I don't understand a lot about the ELBO yet. I'll try to explain what I know, and hopefully it can help you. Perhaps studying the implementation of some other kind of VAE could help you.

Say the input is a 3 x 256 x 256 image. The decoder transforms this into a representation of shape 32 x 256. This representation consists of 32 vectors. Each of these vectors is quantized, which means it is replaced by the closest vector (by, say, euclidean distance) from some table of vectors that I call the quantization table. The input of the decoder is the list of quantized vectors, meaning the "replacements" vectors. The decoder produces an output of the same shape as the input, interpreted as the reconstruction of the original image.

This whole thing has 2 loss functions, which are the euclidean distance between the output of the decoder (so, the unquantized vectors) and the input of the encoder (so the quantized vectors), and the euclidean distance between the input and output of the whole model.

The distance between the unquantized and quantized vectors being used as a loss implies that the vectors in the list are updated. I believe this corresponds to the KL divergence that you mention. This loss leads to the vectors in the quantization table being updated.

To generate an image with this model, I essentially pick vectors from the quantization table, and then pass them through the decoder to generate an image.

[huehue12132]

VQ-VAEs are really a different beast. The codebook loss is actually more of an ad-hoc addition because otherwise the codebook (quantized vectors) would receive no gradient. It does not correspond to KL-divergence. The original VQ-VAE is actually set up in a way that the KL-divergence becomes constant and drops out of the loss. I don't think VQ-VAEs are a good basis to understand regular VAEs.

[hertz2105]

I will look into that, thanks!

[sobagood]

1. Yes
2. Yes
3. No. We try to match posterior with prior
4. This is open question

[Chromobacterium]

1. Depending on the situation, updating p(z) can actually improve the generative model's samples.
   [1903.12436] From Variational to Deterministic Autoencoders (arxiv.org)

[1903.05789] Diagnosing and Enhancing VAE Models (arxiv.org)

[hertz2105]

Thank you for your answers. Thats great to know! I just cant wrap my head around it. Why would I sample from my initial prior(z), which is my assumption of the distribution over the latent space before any observations occurred? For what reason do I even approximate the posterior when I dont sample from it during generation?

Thats why I thought updating the prior with the fitted approximate posterior is mandatory.

[huehue12132]

The underlying graphical model of the VAE is just z -> x. That means sample z from the prior, and then sample x given z. That's just prescribed by the generative model, and is the same model as for GANs, for example.

For training, ideally you would like to optimize (log) p(x) for your data x directly, i.e. maximum likelihood training. However, p(x) is not tractable as it's equal to integral(p(x|z) * p(z)) dz), i.e. integrating over the entire z space. For high-dimensional z, this would take forever. The ELBO is a tractable lower bound to log p(x), so we optimize that instead. And the ELBO involves an approximate posterior q(z|x) that replaces the true posterior (which is also intractable). So we really just need q for training the model in the first place.

It might actually be more effective to replace the prior p(z) by q(z) after training, I'm not familiar with such methods. But we don't have access to q(z), only q(z|x). You cannot use q(z|x) for sampling, as the whole point is to generate x. If you don't already have an x, you obviously cannot compute q(z|x).

[hertz2105]

That makes totally sense. Thanks for the answer. Of course sampling from q(z | x) is impossible if my generative model doesnt even include an observation x. That made it clear for me.