How does backprop work through the random sampling layer in a variational autoencoder?

Question

Implementations of variational autoencoders that I've looked at all include a sampling layer as the last layer of the encoder block. The encoder learns to generate a mean and standard deviation for each input, and samples from it to get the input's representation in latent space. The decoder then attempts to decode this back out to match the inputs.

My question: How does backpropagation handle the random sampling step?

Random sampling is not a deterministic function and doesn't have a derivative. In order to train the encoder, gradient updates must somehow propagate back from the loss, through the sampling layer.

I did my best to hunt for the source code for tensorflow's autodifferentiation of this function, but couldn't find it. Here's an example of a keras implementation of that sampling step, from the keras docs, in which tf.keras.backend.random_normal is used for the sampling.

class Sampling(layers.Layer):
    """Uses (z_mean, z_log_var) to sample z, the vector encoding a digit."""

    def call(self, inputs):
        z_mean, z_log_var = inputs
        batch = tf.shape(z_mean)[0]
        dim = tf.shape(z_mean)[1]
        epsilon = tf.keras.backend.random_normal(shape=(batch, dim))
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon

Answer 1

You do not backpropagate with respect to $\epsilon$, which is the random sample or random variable (depending on how you look at it). You backpropagate with respect to the mean $\mu$ and variance $\sigma$ of the latent Gaussian (the variational distribution). Note that, although $z$ is a random sample (and not just a sample), because it's computed as a function of $\epsilon$ (a random sample, once it's sampled from e.g. $\mathcal{N}(0, 1)$), $\mu$ and $\sigma$ are not: these are learnable parameters and are deterministic.

Having said this, note that we use the reparametrization trick in the VAE, so we compute the random sample as $z = g_\phi(\epsilon, x))$, where $g_\phi$ is a deterministic function (encoder) parametrized by $\phi$ (the weights of the neural network that represents the encoder). In case the random variable $z \sim \mathcal{N}(\mu, \sigma^2)$, then we can express the random variable (so also the random sample) as follows $z=g_\phi(\epsilon, x)) = \mu+\sigma \epsilon$. So, as you can see from the code, we sample $\epsilon$ from some prior $p(\epsilon)$ (e.g. $\mathcal{N}(0, 1)$), then we compute $z$ deterministically.

Why is this reparametrization trick useful? The authors of the VAE paper explain it.

This reparameterization is useful for our case since it can be used to rewrite an expectation w.r.t $q_{\phi}(\mathbf{z} \mid \mathbf{x})$ such that the Monte Carlo estimate of the expectation is differentiable w.r.t. $\phi$. A proof is as follows. Given the deterministic mapping $\mathbf{z}=g_{\phi}(\boldsymbol{\epsilon}, \mathbf{x})$ we know that $q_{\phi}(\mathbf{z} \mid \mathbf{x}) \prod_{i} d z_{i}=$ $p(\boldsymbol{\epsilon}) \prod_{i} d \epsilon_{i}$. Therefore $^{1}, \int q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x}) f(\mathbf{z}) d \mathbf{z}=\int p(\boldsymbol{\epsilon}) f(\mathbf{z}) d \boldsymbol{\epsilon}=\int p(\boldsymbol{\epsilon}) f\left(g_{\boldsymbol{\phi}}(\boldsymbol{\epsilon}, \mathbf{x})\right) d \boldsymbol{\epsilon}$. It follows that a differentiable estimator can be constructed: $\int q_{\phi}(\mathbf{z} \mid \mathbf{x}) f(\mathbf{z}) d \mathbf{z} \simeq \frac{1}{L} \sum_{l=1}^{L} f\left(g_{\phi}\left(\mathbf{x}, \boldsymbol{\epsilon}^{(l)}\right)\right)$ where $\boldsymbol{\epsilon}^{(l)} \sim p(\boldsymbol{\epsilon}).$

Note that $\mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x})} \left[ f(\mathbf{z}) \right] = \int q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x}) f(\mathbf{z}) d \mathbf{z} = \int p(\boldsymbol{\epsilon}) f\left(g_{\boldsymbol{\phi}}(\boldsymbol{\epsilon}, \mathbf{x})\right) d \boldsymbol{\epsilon} = \mathbb{E}_{p(\boldsymbol{\epsilon})} \left[ f\left(g_{\boldsymbol{\phi}}(\boldsymbol{\epsilon}, \mathbf{x})\right) \right]$, which can be estimated with $\frac{1}{L} \sum_{l=1}^{L} f\left(g_{\phi}\left(\mathbf{x}, \boldsymbol{\epsilon}^{(l)}\right)\right)$ where $\boldsymbol{\epsilon}^{(l)} \sim p(\boldsymbol{\epsilon})$. In other words, you can sample $L$ $z$ in the way we did in order to estimate $\mathbb{E}_{\color{blue}{q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x})}} \left[ \color{red}{f(\mathbf{z})} \right]$ (this are Monte Carlo estimates of the expectation).

In the case of the VAE, we want to optimize the ELBO, which is the following objective function

$$\mathcal{L}\left(\boldsymbol{\theta}, \boldsymbol{\phi} ; \mathbf{x}^{(i)}\right)=\underbrace{-D_{K L}\left(q_{\boldsymbol{\phi}}\left(\mathbf{z} \mid \mathbf{x}^{(i)}\right) \| p_{\boldsymbol{\theta}}(\mathbf{z})\right)}_{\text{KL divergence}}+ \underbrace{\mathbb{E}_{\color{blue}{q_{\phi}\left(\mathbf{z} \mid \mathbf{x}^{(i)}\right)}}\left[\color{red}{\log p_{\boldsymbol{\theta}}\left(\mathbf{x}^{(i)} \mid \mathbf{z}\right)}\right]}_{\text{likelihood}}$$

which we can estimate with Monte Carlo estimates of the second term (the likelihood term)

$$ \widetilde{\mathcal{L}}^{B}\left(\boldsymbol{\theta}, \boldsymbol{\phi} ; \mathbf{x}^{(i)}\right)=-D_{K L}\left(q_{\boldsymbol{\phi}}\left(\mathbf{z} \mid \mathbf{x}^{(i)}\right) \| p_{\boldsymbol{\theta}}(\mathbf{z})\right)+ \underbrace{\frac{1}{L} \sum_{l=1}^{L}\left(\log p_{\boldsymbol{\theta}}\left(\mathbf{x}^{(i)} \mid \mathbf{z}^{(i, l)}\right)\right)}_{\text{likelihood}} $$ where $$ \text { where } \quad \mathbf{z}^{(i, l)}=g_{\phi}\left(\boldsymbol{\epsilon}^{(i, l)}, \mathbf{x}^{(i)}\right) \text { and } \boldsymbol{\epsilon}^{(l)} \sim p(\boldsymbol{\epsilon}) $$

Here, the likelihood is computed with the neural network, for example, in practice, you use the cross-entropy of the output of your decoder, which gets as input the input to the decoder, hence $z$.

You can ignore the KL divergence now because, in the case of Gaussians, it can be computed analytically.

Now, what if we didn't use the reparametrization trick? Could we still backpropagate with respect to $\phi$? The authors of the VAE write

The usual (naïve) Monte Carlo gradient estimator for this type of problem is: $\nabla_{\phi} \mathbb{E}_{q_{\phi}(\mathbf{z})}[f(\mathbf{z})]=\mathbb{E}_{q_{\phi}(\mathbf{z})}\left[f(\mathbf{z}) \nabla_{q_{\phi}(\mathbf{z})} \log q_{\phi}(\mathbf{z})\right] \simeq \frac{1}{L} \sum_{l=1}^{L} f(\mathbf{z}) \nabla_{q_{\phi}\left(\mathbf{z}^{(l)}\right)} \log q_{\phi}\left(\mathbf{z}^{(l)}\right)$ where $\mathbf{z}^{(l)} \sim q_{\phi}\left(\mathbf{z} \mid \mathbf{x}^{(i)}\right) .$ This gradient estimator exhibits exhibits very high variance (see e.g. [BJP12])

So, the answer is yes, as opposed to what many people might say around, but with another estimator (which may remind you of some equations you have seen in reinforcement learning related to REINFORCE, if you are familiar with this), i.e. $$\frac{1}{L} \sum_{l=1}^{L} f(\mathbf{z}) \nabla_{q_{\phi}\left(\mathbf{z}^{(l)}\right)} \log q_{\phi}\left(\mathbf{z}^{(l)}\right) \tag{1}\label{1},$$ which has high variance.

So, in the end, the reparametrization trick can be viewed as a variance reduction technique. There are others, like control variates or Flipout (used e.g. in the context of Bayesian neural networks).

The first thing to note about \ref{1} is that we do not need to take the derivative with respect to $f$. The second thing is that $\nabla_{q_{\phi}\left(\mathbf{z}^{(l)}\right)} \log q_{\phi}\left(\mathbf{z}^{(l)}\right)$ is the score function.

Now, don't ask me how to calculate this gradient $\nabla_{q_{\phi}\left(\mathbf{z}^{(l)}\right)} \log q_{\phi}\left(\mathbf{z}^{(l)}\right)$ (because I am bad at math). However, note that $ \mathbf{z}^{(l)} \sim q_{\phi}\left(\mathbf{z} \mid \mathbf{x}^{(i)}\right) $, so we treat $\mathbf{z}^{(l)}$ as a fixed random sample. By the way, I am not sure whether they used $q_{\phi}\left(\mathbf{z} \mid \mathbf{x}^{(i)}\right)$ differently from $q_{\phi}\left(\mathbf{z} \right)$. I don't think so. I think they are the same and I think we can still back-propagate with respect to $\phi$, even if we use it to sample $\mathbf{z}^{(i)}$, because, once this latter is sampled, it can be treated as fixed (random) sample.

I also note that I think they meant $\frac{1}{L} \sum_{l=1}^{L} f(\mathbf{z}^{(l)}) \nabla_{q_{\phi}\left(\mathbf{z}^{(l)}\right)} \log q_{\phi}\left(\mathbf{z}^{(l)}\right)$, i.e. they forget to use $\mathbf{z}^{(l)}$ rather than just $\mathbf{z}$ as input for $f$ in the Monte Carto estimate. You can see in section 3 of [BJP12] and section 4.2. of [1] they do it like this, and it makes sense. So, the VAE paper has sloppy stuff in there.

My intuition of why this has high variance is because you sample something according to the variational distribution, but then you update this variational distribution, and you continuously do this. However, I am not sure this is the right intuition. In [BJP12], they say (in section 4) this MC estimate has high variance and that you need a lot of samples $\mathbf{x}$. I don't know exactly why this is the case because I didn't fully read this paper yet.

Answer 2

As already mentioned in the comment, the reason, why the does the backpropagation still work is the Reparametrization Trick.

For variational autoencoder (VAE) neural networks to be learned predict parameters of the random distribution - the mean $\mu_{\theta} (x)$ and the variance $\sigma_{\phi} (x)$ for the case on normal distribution. Here $\theta$ and $\phi$ are parameters, that can be learned during the backprogation.

At each sampling step, one generates random number $\varepsilon$ from the normal distribution with zero mean and unit variance $\mathcal{N}(0, 1)$ and then multiplies this number by $\sigma_{\phi} (x)$ and adds $\mu_{\theta} (x)$: $$ \mu_{\theta} (x) + \varepsilon \sigma_{\phi} (x) $$ The output of this operation has a normal distribution $\mathcal{N}(\mu_{\theta} (x), \sigma_{\phi} (x))$. You do not have to backpropagate through the epsilon, and this number is just multiplied by $\sigma_{\phi} (x) $.

The whole procedure is nicely differentiable with respect to the $\theta$ and $\phi$.

However, note , that the possibility to apply this trick is allowed due to the nice property of normal distribution: $$ \mathcal{N}(a, b) \sim a + b \mathcal{N}(0, 1) $$ Arbitrary distribution doesn't satisfy this property.

In more general case, you would need to apply some other trick to propagate the gradients.

Simple strategy is the Straight through estimator, when one copies the gradient from the output layer to the input through non-differentialbe operation. Or more accurate Gumbel-Softmax softmax.

Answer 3

Sampling per-se is non-differentiable in some point of view, but you can use the reparametrisation in some special scenario.

Differentiation, mathematically, is the act of taking derivatives of functions wrt variables. Sampling is the act of generating a set of elements given some parameters, and differentiating a function that outputs a set is, to my knowledge, not well discussed in the setting of machine learning. More often, sampling is followed by an expectation operator, where you takes average of a function across your samples. And this expectation operator, is differentiable.

$$E_\pi(X(s))={1\over N} \sum_{i=1,\{s_i\}\sim\pi}^N X(s_i)=\sum_s X(s)P_\pi(s)$$

In fact, backpropagating through a sampling operation has a well known name -- it's called "policy gradient" in reinforcement learning, where you study how to take samples (known as policy) so as to minimise error (known as maximising the expected reward).

Not sure why this connection has not been well studied yet... maybe NN people just love a fully deterministic function approximator.