In the definition of the VAE, we make the ansatz that we can factor the joint density as $P^\theta_{x,z} = P^\theta_{x|z} P_z$. The the prior $P_z$ is any distribution specified by the experimenter. The hope is that the training procedure adjusts $\theta$ so that the decoder $P^\theta_{x|z}$ uses the entropy in $P_z$ to model $P_x\approx P^\theta_x$.
The training procedure of VAEs should find parameters $\hat\theta$ that cause the marginal of $z$ to approximate the specified prior, i.e. $\int_{dx} P^{\hat\theta}_{x,z} \approx P_z$ ...
- If training converges to a meaningful optimum;
- If the network has sufficient capacity, such that the approximation loss $D[ Q^\phi_{z|x} \| P_{z|x}]$ vanishes and the evidence lower-bound becomes tight;
- If the encoding model used for amortized inference is expressive enough to predict the correct distribution parameters $\phi$ for each $x$ that achieve this minimal $D[ Q^\phi_{z|x} \| P_{z|x}]$.
If these conditions hold, one expects the average of $Q_{z|x}$ over the training data $\mathcal X = \{x_1,..,x_L\}$ to approximate the specified prior $P_z$, i.e. $\langle Q_{z|x} \rangle_{\mathcal X}\approx P_z$. (And I suppose this can be interpreted as reflecting the true marginal $P^\theta_{z}$ if $\mathcal X$ are a good sample of the true $P_x$).
In my experience, if $\langle Q_{z|x} \rangle_{\mathcal X}$ and $P_z$ differ substantially, the following might be true:
- Decoder or encoder training failed (numerical problem, vanishing gradients, mode collapse, finding a poor local optimum)
- The encoder/decoder lack capacity to map $z$ ↔ $x$
- The latent space lacks capacity to describe all variability in $P_x$
I'm not sure whether adjusting the strength of the prior, or adding other penalties, could influence $Q_{z|x}$, to cause the above to be false. In the usual VAE the latents are Gaussian $z\sim\mathcal N(0,\boldsymbol I)$, so the decoding weights should be able to adjust to make any Gaussian distribution on the latents align with this (in principle and with the aforementioned caveats).
For Gaussian VAEs you can get the marginal $\langle Q_{z|x}\rangle_{\mathcal X}$ using the law of total variance. The mean for each component is the average over the training data $\mu = \langle \mu(x) \rangle_{x\in\mathcal X}$. The variance for each component is $\sigma^2 = \langle \sigma^2(x) \rangle_{x\in\mathcal X} + \operatorname{var}_{x\in\mathcal X}[\mu(x)]$.
I've seen this differ substantially from $P_z$ early in training. Since everything is a diagonal Gaussian, you could adjust the adjacent encoder/decoder weights/biases so that this matches $P_z$. I tried this and it may slightly accelerate the early phases of training.
All of this is more subtle in models where the latent $P_z$ is itself trainable. For example, The restricted Boltzmann machine uses Bernoulli latents with a parameter $p$. In this case the "prior" should adapt during training, since you can't guarantee that your latents will have the right distribution of feature frequencies to match the data a priori. There isn't much point is explicitly parameterize $P_z$ in the Gaussian VAE, since these parameters can be absorbed into the adjacent encoding/decoding layers.
TLDR:
- When deriving the loss function, it is valid to assume $P^\theta_z = P_z$, where $P_z$ is fixed, inasmuch as we are searching for a model where this is true by construction. Minimizing this loss, when successful, will make $P^\theta_z \approx P_z$.
- If your prior needs to adjust parameters that can't be handled by backprop + reparameterization trick, there may be good reason to update $P_z \gets P^\theta_z \approx \langle P^\theta_{z|x}\rangle_{\mathcal X}$ (or some approximation) during training. This, at least, was my understanding of the wake-sleep algorithm.
- Sometimes, we don't find a model where $P^\theta_z \approx P_z$. Free energy is not minimized and we should be suspicious of the model. Maybe then, as you say, one should use $\langle P_{z|x}\rangle_{\mathcal X}$ to sample $z$.
I have more experience with restricted Boltzmann machines, where the generative $P_z$ does change during training. As far as I can tell, this shouldn't need to happen in VAEs with Gaussian latent spaces trained via backpropagation. Experts might want to weight in—I may have made some errors.