Ok, I had to ponder this for ages. Since when I first looked, I sort of came to the same conclusions as yourself :) So, I first wrote:
I'm not sure I entirely agree that 'the choice of this output distribution is often Gaussian', though it can be: I think we generally will have some kind of highly non-linear neural network, eg to take some latent values in z, and transform them into image of the digits 0,1,2, ... ; or some other images, etc.
And I concur with you: in a VAE, we are putting a prior on $z$, and sampling $z$, typically from an axis-aligned Gaussian, based on the output of the input, encoder, network.
However, looking at the tutorial you linked, the text you have quoted is from the section '1.1 Preliminaries: latent variable models', and describes the graphical model shown in figure 1:
Later on, in section '2. Variational Autoencoders', it relates this model to the VAE model that we are typically used to seeing, with the input encoder neural net, followed by axis-aligned $z$-sampling, followed by a deterministic neural network, that reconstructs $X$. ie this bit:
... but, having written all this, I still couldn't quite see the answer to your question, i.e. why/how is $P(X \mid z)$ stochastic, and what does it mean?
Since normally in a VAE we essentially have the following layers, if we write it as a net:
- input inference network, models $q(z \mid X)$. Might have a bunch of layers, like convolutions or whatever. deterministic, given the weights, and the inputs
- then we take the outputs of the inference layer and transform them linearly into the mean and variance of an axis-aligned Gaussian. this bit is still deterministic
- we sample from this. stochastic and we use the reparameterization trick, so we can backprop through this
- at this point we have: $z$
- then we pass $z$ through the generation network, which is 100% entirely deterministic, though is probably another deep neural net, with deconvolution layers and so on
- the output of this network is then $X$
So, ... where is the stochastic bit of the generation network?
I think that the answer is hinted at by the likelihood expression written at the bottom of page 10:
$$
\def\Exp{\mathbb{E}}
\Exp_{z \sim Q}[\log P(X \mid z)]
$$
I think the point is that, whilst the $X$ is generated deterministically from $z$ in our network, but the error/loss term is calculated on the basis that these $X$ outputs are normally distributed, so that the loss is not almost surely zero. Otherwise if the outputs are not assumed to be normally distributed, and have a Dirac delta distribution, the probability of our output $X$ given our input $X$ is almost surely $0$.
There's a little bit of discussion in 2.4.3 'VAEs and the reparameterization parameter', that makes this clearer I think? They point out the effect of modifying the $\sigma$ in the Gaussian $P(X \mid z) \sim \mathcal{N}(f(z), \sigma^2 I)$, ie that as $\sigma \rightarrow 0$, then the output distribution tends to Dirac delta, and the likelihood of our reconstructed $X$ tends to $0$.
Edit: You ask:
The encoder gives us mean and std of the latent variables and we are forcing them to be a unit normal. So eventually, the encoder should only give us a zero and a one for all images ?
The z values will form a continuum around the origin, out to a radius of ~3 or so, at which radius the density of the Gaussians will be very tiny. The model will learn to place inputs that are similar to each other into nearby positions in the latent space. The z values will be reals, not binary 0/1, but eg maybe images of cars will end up near [0.3, 0.2] (if z is 2-dimensional), and maybe ostriches would be near [0.7, -0.3], and so on. The exact layout would depend on the distribution of the input data, and is learned in an entirely unsupervised manner. There are some great images around on the web, for example at http://blog.fastforwardlabs.com/2016/08/12/introducing-variational-autoencoders-in-prose-and.html :
This is showing how images learned from mnist look at different positions in the 2-dimensional z space that a model has learned.
You ask:
how can we expect the decoder to generate the exact image in the input when we randomly sample from a normal distribution?
So, the idea is that points in the latent space which are nearby should be similar. So, by encoding to a particular part of the latent space, the generator will generator an image similar-ish (but not identical) to the input. We want it to sample so that it can adequately explore the latent space, make the meaning of adjacent points in the latent space similar-ish to each other. Otherwise we might as well just throw the data into a lookup table :)
On the output, the Gaussian blurs somewhat the output images, but it's not like the latent space cannot influence the generated image at all. The extent to which the output is blurred is controlled by the $\sigma$ hyperparameter. Too large, and the output will be very blurry. Too small and learning will be slow.
Edit 2: You ask:
can you dumb this down for me a little - "I think the point is that, whilst the XX is generated deterministically from zz in our network, but the error/loss term is calculated on the basis that these XX outputs are normally distributed, so that the loss is not almost surely zero. Otherwise if the outputs are not assumed to be normally distributed, and have a Dirac delta distribution, the probability of our output XX given our input XX is almost surely 0." the loss is Variational lower bound + L2 norm between generated and input images?
whilst we can think of what we are calculating as a loss, and it works as a loss, what we are actually calculating is a probability distribution over output images, and the likelihood of our original image, given that probability distribution. It may or may not end up as an expression similar or identical to MSE loss, but the actual underlying approach is different.
MSE Loss:
- difference between target and output
- square it
- take the mean across all the elements in the vector (if there's more than one)
In our case, we have a probability distribution for $X$ given $z$:
$$
p(X \mid z)
$$
and its a Gaussian:
$$
p(X \mid z) = \mathcal{N}(X; \cdot, \cdot)
$$
The variance of the Gaussian is fixed, to be $\sigma^2$, a hyperparameter:
$$
p(X \mid z) = \mathcal{N}(X; \cdot, \sigma^2)
$$
The mean of the Gaussian is the output of our decoder neural net:
$$
p(X \mid z) = \mathcal{N}(X; f(z, \theta), \sigma^2)
$$
Then, we take our input image, and we calculate the likelihood, based on this Gaussian, by simply plugging in the numbers:
$$
p(X \mid z) = \mathcal{N}(\text{input image}; \text{neural net output}, \text{hyperparameter})
$$
one way to sample is to apply argmax(softmax($\alpha_j$))
That is hardly "sampling", given that you deterministically pick the largest $\alpha_j$ every time. (also, you said that $\alpha$ is the unnormalized probability but that doesn't make sense seeing as log probabilities go into the softmax). The correct way to sample would be sample(softmax($x$)), where $x$ are the logits. Indeed, the goal of gumbel-softmax is not to replace the softmax operation as you've written it, but the sampling operation:
We can replace sample($p$) where $p$ are a vector of probabilities with argmax($\log p + g$) where $g$ is the gumbel noise. Of course, this is equivalent to argmax($x + g$) where $x$ are again the logits. To conclude, sample(softmax($x$)) and argmax($x+g)$ are equivalent procedures.
Then, if the goal was to have the full distribution over possible outcomes for $z_j$, we can use softmax transformation on top of the perturbation with Gumbel noise.
In fact you already have a distribution over all possible outcomes.
However, argmax($x+g$) is not differentiable wrt $x$, therefore to backpropagate we replace its gradient with the gradient of softmax($(x+g)\tau^{-1}$). When $\tau \rightarrow 0$, the expression approaches argmax.
Picking a reasonable, small values of $\tau$ will ensure a good estimate of the gradient while ensuring that the gradients are numerically well behaved.
and $\tau=1$ just makes the two equations identical
In fact, there is no special significance to $\tau = 1$. Rather, $\tau \rightarrow 0$ makes the gradient estimate unbiased but high in variance, where as larger values of $\tau$ add more bias to the gradient estimate but lower the variance.
Best Answer
The answer is "yes" in one sense and "no" in another sense.
Suppose $X \sim \operatorname{Gamma}(\alpha,\beta)$. Let $F_{\alpha,\beta}$ denote the cdf of a Gamma distribution. Then define $\epsilon = \Phi^{-1}[F_{\alpha,\beta}(X)]$. Then if you simulate $\epsilon \sim N(0,1)$ you can get the relevant gamma distribution by setting $X = F^{-1}_{\alpha,\beta}[\Phi(\epsilon)]$. This is a consequence of the probability integral transform. Additionally, the transform $T(\epsilon ; \alpha, \beta) = F^{-1}_{\alpha,\beta}[\Phi(\epsilon)]$ is differentiable. So you could use this idea with the reparametrization trick, at least in principle, to improve your stochastic variational inference. This implies that, in a liberal sense, the answer is "yes, there is a reparameterization trick", and in fact there is one for essentially any family of continuous distributions. If this seems sort of ad-hoc, notice that if you apply this trick with the Gaussian family in place of the gamma, you get back exactly the usual reparameterization trick.
In a more restrictive sense, I would say the answer is "no". The function $F^{-1}$ above is not available in closed form, so things are not so convenient to the point where we might disqualify this approach. Alternatively, there is no reason to restrict ourselves to $\epsilon \sim N(0,1)$, and we might just ask for $\epsilon \sim Q$ for some standard distribution $Q$ that is easy to sample from, such that $T(\epsilon; \alpha, \beta) \sim \text{Gamma}(\alpha,\beta)$ where $T$ is also easy to compute.
If you find such a transformation $T$ and standard distribution $Q$, let me know because I would be interested in it. The main problem is the shape parameter $\alpha$. If I know $\alpha$, then I can take $T(\epsilon; \alpha, \beta) = \epsilon \beta$ and set $\epsilon \sim \mbox{Gamma}(\alpha,1)$, because the gamma family with know $\alpha$ is a scale family. The shape parameter does not have any nice algebraic properties as far as I know, aside from the fact that $X_1 + X_2 \sim \mbox{Gamma}(\alpha_1 + \alpha_2, 1)$ provided that $X_i \sim \mbox{Gamma}(\alpha_i, 1)$ and they are independent. It's not clear how to take advantage of this. A negative result for us is that, if such a convenient $T$ existed, then R would probably use that to sample from the a generic gamma distribution, but instead it uses rejection sampling.