Learning prior p(z) in VAEs

Question

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.

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.