The paper Ma et al. 2018 states the following statement about a VAE model:

Usually, the prior $p_{\theta}(z)$ is standard normal, but we find that parameterizing it with a trainable mean vector $m$ and variance vector $s^2$ sometimes improves inference.

A VAE has the objective function

$$ \mathcal{L}^{B}=-K L[\overbrace{q_{\phi}\left(\mathbf{z} \mid \mathbf{x}^{(i)}\right)}^{\text {Encoder }} \| \overbrace{p_{\theta}(\mathbf{z})}^{\text {Fixed }}]+\frac{1}{L} \sum_{l=1}^{L} \log \overbrace{p_{\boldsymbol{\theta}}\left(\mathbf{x}^{(i)} \mid \mathbf{z}^{(l)}\right)}^{\text {Decoder }} $$

I was wondering how would the objective for such a VAE with learned $p_{\theta}(z)$ looks like. Also, how would the diagram of the model look like.


1 Answer 1


The objective would look exactly the same as above -- the KL divergence is still closed form (albeit the form now includes $m$ and $s$), and $p_\theta(z)$ doesn't really show up elsewhere.

A really easy way to think about this is to imagine you add a "scaling layer" at the very end of your encoder which adds $m$ to your mean and $\log s^2$ to your log variance. Then at the very start of the decoder, you add an "inverse scaling layer" $z \mapsto \frac{z - m}{s}$. And then everything else just works as before with a standard normal prior for $z$.

It's not too hard to convince yourself that this VAE is equivalent to a VAE without the scalings, but with the fancier prior.

  • $\begingroup$ Thanks. Do you have reference to this or a diagram of what you explained? $\endgroup$
    – Blade
    Jan 24, 2021 at 19:24

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.