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

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

• It is likely this will be done using the reparameterisation trick Commented Dec 19, 2021 at 1:37

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.

• Nice. Worth adding LHS to eq1. Commented Mar 17, 2022 at 16:58

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.

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.

• You don't have to leave the "notice" in your new answer. Moderators can convert answer-like comments to comments or even delete them. This is not weird here because we expect users to provide answers with details and some users might have written a comment as an answer because they don't know how the site works or maybe they don't have the reputation to comment. Having said that, you can post another answer, as you did, once you have written a more detailed one.
– nbro
Commented Mar 18, 2022 at 9:19
• I just wanted to say, that I was not happy, and the moderator is now removing my notice. My emotion was not acknowledged -- it was REMOVED. I don't feel my voice respected anymore. People are not machines. Commented Mar 18, 2022 at 10:11
• I removed the noticed, and I am currently not a mod. So it was not a moderator that removed the notice. People that have more than 2k rep here can edit posts without confirmation, like I did. Our site's model discourages people to leave unnecessary details in the questions or answers, with respect to the question being asked. Things like "Thanks!" or "Have a nice day" are discouraged in questions or answers. It's not because we're nice, but because we want to create an archive of questions and answers where people can easily find their answer without losing time reading unnecessary details.
– nbro
Commented Mar 18, 2022 at 11:17
• I understand your feelings, but you have to understand that this site is not meant for people to share their opinions, but to give answers to questions. Anyway, thanks for trying to contributing to site. If you haven't done it yet, you might want to have a look at our site on-topic page.
– nbro
Commented Mar 18, 2022 at 11:20
• Ok thanks. I am happy to see we are actually trying to communicate. I agree with the general intention that to keep the site well-structured, but editing is a strong privilege and the mechanism could definitely be improved to balance between motivating contribution while keep the quality high. Commented Mar 18, 2022 at 11:34