Sampling z in VAE

Question

How many times do we sample from $Q(z|x)$ in a Variational Autoencoder?

Let’s say that the autoencoder input $x$ is a single image 28x28 pixels - and $Z$ is is a one dimensional distribution. Then, to reconstruct the output $X'$ - I read (I could be wrong) that we can sample $10 000$ times.

Why do we sample so many? And what do we do to reduce it to 28x28?

EDIT

I am still confused how we can sample $10000$ times from $Q(z|x)$. For example let us consider the example above and assume that $Q(z|x)$ ~ $N(0,1)$. We have a 28x28 pixel image as the VAE input denoted $X$ with $28$ nodes each corresponding to one pixel.

Also, just for simplicity lets calculate the expectation of the likelihood part of ELBO instead of the whole thing.

Since, the posterior is Gaussian and we sample only once

$$ E_{z\sim Q}[\log p(x|z)] = \frac{1}{28} \sum_{1}^{28} (x_i-x'_i)^2 $$

Note that $Z_1$ in this case is a single value, and we are summing the error for each node

If we were to sample twice instead - would the expectation be

$$ E_{z\sim Q}[\log p(x|z)] = \frac{1}{2*28} \sum_{1}^{2*28} (x_i-\overline{x}_i)^2 $$

where $\overline{x}_i$ is the combined output of two samples for example: Node $1, 2,...28, 29$(Node $1$ of the second sample), $30$(Node $2$ of the second sample)...$56$(Node $28$ of the second sample)$.

...and so on for 10000 times. Is this correct?

Now if this is correct then it deserves a follow up question. As mentioned in the answer below, we do not really need to sample so much. So, say we sample twice independently - but it just so happens that the second sample $Z_2$ is much closer to the mean of the distribution.

Do we expect that the reconstruction $X'$ and $X"$ look very similar to each other?

Answer 1

Because VAEs are really a graphical model more than they are autoencoders, it can be confusing what exactly "reconstruction" means in context.

Recall that we have an lower bound on the log probability given by the RHS of:

$$\log p(x) - D_{KL}(Q(z|x)||P(z|x)) = E_{z\sim Q}[\log p(x|z)] - D_{KL}(Q(z|x)||P(z))$$

To compute this lower bound -- a necessary prerequisite for doing a backprop pass to maximize it -- corresponds to passing $x$ through the encoder to arrive at $Q(z|x)$, computing the KL-divergence, then estimating $E_{z \sim Q}[\log p(x|z)]$ by sampling once or more (but usually just once) from $Q$ and running the decoder.

This process of estimating the posterior with the encoder and then sampling to approximate the expectation in the RHS so closely mimics the computation of an autoencoder would do that we call it "reconstruction". However, it's really just a side effect of trying to maximize the log probability of the inputs.

What happens when you sample multiple times from $Q$? The immediate consequence is that you get a better approximation of the expectation, and hence a better approximation of the lower bound on the log probability. You also need to run the decoder multiple times, which can be expensive, so it is usually not done. Of course if you do this, then you end up with many reconstructions rather than just one. Note that it is definitely not possible to average the reconstructions and have a meaningful output.

So you probably just want to sample once.

---

In response to your edit, the correct way to write it would be

$$\begin{align*} E_{z \sim Q}[\log p(x|z)] &\approx \frac{1}{n}\sum_{i=1}^N \log p(x|z_i) \\ &\propto -\frac{1}{n}\sum_i ||x-\text{decode}(z_i)||_2^2 \\ &= -\frac{1}{nm} \sum_i \sum_j (x_j - \text{decode}(z_i)_j)^2\end{align*}$$

We would expect that the reconstructions $\text{decode}(z_i)$ look quite similar to each other, but not exactly the same. Exactly how well depends on the nature of the data and how well the model is fitted.