Prior in variational autoencoders

Question

I am currently dealing with variational autoencoders where I've read the original paper "An introduction to variational Bayes" from Kingma and Welling. I am currently still a little confused about the choice of the prior $p(z)$. For my understanding the variational autoencoder has the three possible distribution families we can choose:

• The family of distributions for the encoder $q_\phi(z|x)$ parametrised by $\phi$.
• The family of distribtuions for the decoder $p_\theta(x|z)$ parametrised by $\theta$.
• The family of prior distributions $p_\theta(z)$.

Now for my understanding and also according to this notation used, the parameters for the decoder and the prior have to be the same, since overall we want to compute $p_\theta(x,z)$ with a neural network right?

Then, how is it possible to choose $p_\theta(z)=\mathcal{N}(0,I)$ while at the same time optimising the parameters $\theta$ computed by the decoder network?

As a third question: if not restricting the prior to be a standard normal distribution, how do we learn the parameters of the prior $p_\theta(z)$?

---

Edit: I am aware, that in practice we choose the prior distribution to be standard normal. But why are we allowed to do so? For my understanding, the whole model (looking at the variational autoencoder as one model) we want to fit $p_\theta(x,z)=p_\theta(x|z) p_\theta(z)$ meaning the parameters $\theta$ are coming from the same parameter space for the decoder and the prior distribution. But when on the one hand fixing the prior distribution and on the other hand computing the parameters $\theta$ for the decoder, then we are kind of violating this assumption aren't we?

Answer 1

Now for my understanding and also according to this notation used, the parameters for the decoder and the prior have to be the same, since overall we want to compute $p_\theta(x,z)$ with a neural network right?

Prior is the distribution for the parameters. Parameters come “from” the prior.

Then, how is it possible to choose $p_\theta(z)=\mathcal{N}(0,I)$ while at the same time optimising the parameters $\theta$ computed by the decoder network?

What you are optimizing is the posterior probability, so likelihood times prior. You choose the prior a priori, plug-in to the Bayes theorem and optimize. When optimizing you consider the prior.

As a third question: if not restricting the prior to be a standard normal distribution, how do we learn the parameters of the prior $p_\theta(z)$?

You assume them based on your prior knowledge, on what you can expect the distributions of the parameters to be before seeing the data. There's nothing special about Gaussians here. Moreover, the above prior choice is dictated rather by convenience than any deep mathematical argument.

It may be easier to understand on practical example, Keras’s blog has a nice tutorial on autoencoders.

Answer 2

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.