Does VAE backprop start from the decoder all the way to encoder?

Question

In neural networks that start with input layer, run through hidden layers, and ultimately reach the output layer, we start back-propagation from weights closer to output layer and go backward towards the input layer.

Now, in case of VAE, I was confused as to whether back-propagation start from the decoder all the way to encoder with an aim to maximize ELBO or is it something different?

We know,

log p(x) = ELBO + KLD between approx posterior and true posterior [Equation 1 of paper]

Since the KLD part of log evidence equation above (bold part) has true posterior which is intractable, in order to minimize it, we instead try to maximize ELBO.

Hence, we can say that maximizing ELBO is basically minimizing the KL divergence between approximate posterior and true posterior (i.e. q(z|x) and p(z|x)). We are maximizing ELBO so as to indirectly obtain the objective of minimizing the KLD between approximate posterior q(z|x) and true posterior p(z|x), which I feel is related to encoder as we are trying to find best approximate distribution of z given x.

So, my confusion and question is that since maximizing ELBO is, as explained above, associated with encoder, how does maximizing it helps to correctly adjust weights for the decoder?

Is it because of the following reason:

Although we are maximizing ELBO so as to indirectly minimize KLD between approximate posterior and true posterior, when ELBO is derived, we can see that it consists of reconstruction term along with the constraining KLD term (i.e. KLD between approximate posterior and prior). Since it consists of the reconstruction term, it is somehow connected to decoder as well. Hence, maximizing ELBO has impact on both encoder and decoder.

Additional Question: Is it ELBO that is taken derivative of with respect to weights when back propagating the VAE neural network?

Answer 1

Yes, I think the explanation you give is quite correct. You always need the decoder which is used to parametrize $p_{\theta}(x|z)$, because what you are really doing is maximizing a lower bound on the likelihood based on the generative model decomposed as $p_{\theta}(x,z)=p_{\theta}(z)p_{\theta}(x|z)$.

More precisely, when training a VAE you need to solve two problems simultaneously. The first is about learning the parameters $\theta$ of the true distribution $p_{\theta}(x,z)=p_{\theta}(z)p_{\theta}(x|z)$, and this part is about maximizing the likelihood, which you would normally do with the EM algorithm (maximizing ${\rm E}_{z\sim p_{\theta}(z|x)}(p_{\theta}(z,x))$). However, you do not know how to compute the true posterior $p_{\theta}(z|x)$. So there is a second problem of learning the parameters $\phi$ of the variational approximation $q_{\phi}(z|x)$ of the true posterior which is parametrized by the encoder.

If you see training a VAE as minimizing ${\rm KL}(q_{\phi}(z|x)||p_{\theta}(z|x))$ with respect to $\phi$ and $\theta$, you also end up using the expression with the ELBO and the loglikelihood (and thus $p_{\theta}(x,z)=p_{\theta}(z)p_{\theta}(x|z)$ and thus making use of the decoder) since you don't know the true posterior.

Note that in other contexts you could just want to fit a variational posterior to a true and known posterior $p(z|x)$ by minimizing ${\rm KL}(q_{\phi}(z|x)||p(z|x))$, but this is not what is done in VAE training.