# Learning prior p(z) in VAEs

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.

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$$.