How does the reparameterization trick for variational autoencoders (VAE) work? Is there an intuitive and easy explanation without simplifying the underlying math? And why do we need the 'trick'?

  • 1
    One part of the answer is to notice that all Normal distributions are just scaled and translated versions of Normal(1, 0). To draw from Normal(mu, sigma) you can draw from Normal(1, 0), multiply by sigma (scale), and add mu (translate). – monk Jan 24 at 23:53

Assume we have a normal distribution $q$ that is parameterized by $\theta$, specifically $q_{\theta}(x) = N(\theta,1)$. We want to solve the below problem $$ \text{min}_{\theta} \quad E_q[x^2] $$ This is of course a rather silly problem and the optimal $\theta$ is obvious. However, here we just want to understand how the reparameterization trick helps in calculating the gradient of this objective $E_q[x^2]$.

One way to calculate $\nabla_{\theta} E_q[x^2]$ is as follows $$ \nabla_{\theta} E_q[x^2] = \nabla_{\theta} \int q_{\theta}(x) x^2 dx = \int x^2 \nabla_{\theta} q_{\theta}(x) \frac{q_{\theta}(x)}{q_{\theta}(x)} dx = \int q_{\theta}(x) \nabla_{\theta} \log q_{\theta}(x) x^2 dx = E_q[x^2 \nabla_{\theta} \log q_{\theta}(x)] $$

For our example where $q_{\theta}(x) = N(\theta,1)$, this method gives $$ \nabla_{\theta} E_q[x^2] = E_q[x^2 (x-\theta)] $$

Reparameterization trick is a way to rewrite the expectation so that the distribution with respect to which we take the gradient is independent of parameter $\theta$. To achieve this, we need to make the stochastic element in $q$ independent of $\theta$. Hence, we write $x$ as $$ x = \theta + \epsilon, \quad \epsilon \sim N(0,1) $$ Then, we can write $$ E_q[x^2] = E_p[(\theta+\epsilon)^2] $$ where $p$ is the distribution of $\epsilon$, i.e., $N(0,1)$. Now we can write the derivative of $E_q[x^2]$ as follows $$ \nabla_{\theta} E_q[x^2] = \nabla_{\theta} E_p[(\theta+\epsilon)^2] = E_p[2(\theta+\epsilon)] $$

Here is an IPython notebook I have written that looks at the variance of these two ways of calculating gradients. http://nbviewer.jupyter.org/github/gokererdogan/Notebooks/blob/master/Reparameterization%20Trick.ipynb

  • 1
    What is the "obvious" theta for the first equation? – gwg Apr 29 at 17:12
  • it's 0. one way to see that is to note that E[x^2] = E[x]^2 + Var(x), which is theta^2 + 1 in this case. So theta=0 minimizes this objective. – goker May 7 at 18:38
  • So, it depends completely on the problem? For say min_\theta E_q [ |x|^(1/4) ] it might be completely different? – Anne van Rossum May 12 at 10:49
  • What depends on the problem? The optimal theta? If so, yes it certainly depends on the problem. – goker May 13 at 17:32

After reading through Kingma's NIPS 2015 workshop slides, I realized that we need the reparameterization trick in order to backpropagate through a random node.

Intuitively, in its original form, VAEs sample from a random node $z$ which is approximated by the parametric model $q(z \mid \phi, x)$ of the true posterior. Backprop cannot flow through a random node.

Introducing a new parameter $\epsilon$ allows us to reparameterize $z$ in a way that allows backprop to flow through the deterministic nodes.

original and reparameterised form

  • 1
    Why is $z$ deterministic now on the right? – bringingdownthegauss Aug 4 at 22:42
  • It's not, but it's not a "source of randomness" - this role has been taken over by $\epsilon$. – quant_dev Sep 17 at 20:47
  • Note that this method has been proposed multiple times before 2014: blog.shakirm.com/2015/10/… – quant_dev Sep 17 at 20:49

A reasonable example of the mathematics of the "reparameterization trick" is given in goker's answer, but some motivation could be helpful. (I don't have permissions to comment on that answer; thus here is a separate answer.)

In short, we want to compute some value $G_\theta$ of the form, $$G_\theta = \nabla_{\theta}E_{x\sim q_\theta}[\ldots]$$

Without the "trick", we can often rewrite this, per goker's answer, as $E_{x\sim q_\theta}[G^{est}_\theta(x)]$, where, $$G^{est}_\theta(x) = \ldots\frac{1}{q_\theta(x)}\nabla_{\theta}q_\theta(x) = \ldots\nabla_{\theta} \log(q_\theta(x))$$

If we draw an $x$ from $q_\theta$, then $G^{est}_\theta$ is an unbiased estimate of $G_\theta$. This is an example of "importance sampling" for Monte Carlo integration. If the $\theta$ represented some outputs of a computational network, we could use this in back-propagatation (apply the chain rule) to find derivatives with respect to network parameters.

The key point is that $G^{est}_\theta$ is often a very bad (high variance) estimate. Even if you average over a large number of samples, you may find that its average seems to systematically undershoot (or overshoot) $G_\theta$.

A fundamental problem is that essential contributions to $G_\theta$ may come from values of $x$ which are very rare (i.e., $x$ values for which $q_\theta(x)$ is small). The factor of $\frac{1}{q_\theta(x)}$ is scaling up your estimate to account for this, but if you never see such a value of $x$, that scaling won't help. The goodness or badness of $q_\theta$ (i.e.,the quality of the estimate, $G^{est}_\theta$, for $x$ drawn from $q_\theta$) may depend on $\theta$, which may be far from optimum (e.g., an arbitrarily chosen initial value). It is a little like the story of the drunk person who looks for his keys near the streetlight (because that's where he can see/sample) rather than near where he dropped them.

The "reparameterization trick" sometimes address this problem. Using goker's notation, the trick is to rewrite $x$ as a function of a random variable, $\epsilon$, with a distribution, $p$, that does not depend on $\theta$, and then rewrite the expectation in $G_\theta$ as an expectation over $p$,

$$G_\theta = \nabla_\theta E_{\epsilon\sim p}[J(\theta,\epsilon)] = E_{\epsilon\sim p}[ \nabla_\theta J(\theta,\epsilon)]$$ for some $J(\theta,\epsilon)$.

The reparameterization trick is especially useful when the new estimator, $\nabla_\theta J(\theta,\epsilon)$, no longer has the problems mentioned above (i.e., when we are able to choose $p$ so that getting a good estimate does not depend on drawing rare values of $\epsilon$). This can be facilitated (but is not guaranteed) by the fact that $p$ does not depend on $\theta$ and that we can choose $p$ to be a simple unimodal distribution.

However, the reparamerization trick may even "work" when $\nabla_\theta J(\theta,\epsilon)$ is not a good estimator of $G_\theta$. In particular, in the context of optimization (e.g., where we are using the gradient for stochastic gradient descent), during a short interval of optimization, $t$, where $\theta$ doesn't change much, we sample $\epsilon$ a few times, $\{\epsilon_{t+1},\epsilon_{t+2},...,\epsilon_{t+N}\}$. During that interval, we are effectively optimizing a function, $$J^{eff}(\theta,t) = \frac{1}{N}\sum_{i=1}^N J(\theta, \epsilon_{t+i})$$

Since the values of $\epsilon$ that we draw at the start of optimization are going to be similar throughout optimization, we might expect that the effective objective function, $J^{eff}(\theta,t)$, would not depend too much on $t$. This effective objective function may be optimized with the usual approaches to optimization (e.g., stochastic gradient descent) even if it is different from the original objective function (corresponding to $N \rightarrow \infty$).

Even if there are "important" values of $\epsilon$ which are very rare ("important" with respect to computing $G_\theta$ accurately, or finding a global optimum), we don't see them during optimization but we also don't see them when we use our model (if our model is a generative model) -- that is, we optimize the average payoff (objective function value) that we are likely to get in actual use, which is what we care about, not the expected value we are unlikely to realize (along the lines of the St. Petersburg Paradox). Further, even if $J(\theta,\epsilon)$ has the potential (i.e., for some $\theta$) to depend in a fine-grained way on $\epsilon$ -- so that there would be a lot of noise in $J^{eff}(\theta,t)$ for those values of $\theta$ -- the average gradient will be correspondingly weakened in those directions (in $\theta$-space) -- so we might expect the optimization to move away from those values of $\theta$, towards parts of $\theta$-space that are smoother in $\epsilon$. In short, there is reason to hope that, even in complex cases, the reparameterization trick may effectively constrain the expressive power of our model, but allow our optimization process to converge to a "good-enough" model.

I hope that helps.

  • "The factor of 1/qθ(x) is scaling up your estimate to account for this, but if you never see such a value of x, that scaling won't help." Can you explain more? – czxttkl Nov 10 at 19:52
  • @czxttkl In practice we estimate expected values with a finite number of samples. If $q_\theta$ is very small for some $x$, then we may be very unlikely to sample such an $x$. So even though $G_{\theta}^{est}(x)$ includes a large factor of $1/q_\theta$, and may make a meaningful contribution to the true expected value, it may be excluded from our estimate of the expected value for any reasonable number of samples. – Seth Bruder Nov 19 at 17:59

Let me explain first, why do we need Reparameterization trick in VAE.

VAE has encoder and decoder. Decoder randomly samples from true posterior Z~ q(z∣ϕ,x). To implement encoder and decoder as a neural network, you need to backpropogate through random sampling and that is the problem because backpropogation cannot flow through random node; to overcome this obstacle, we use reparameterization trick .

Now lets come to trick. Since our posterior is normally distributed, we can approximate it with another normal distribution . We approximate Z with normally distributed ε.

enter image description here

But how this is relevant ?

Now instead of saying that Z is sampled from q(z∣ϕ,x) , we can say Z is a function that takes parameter (ε,( µ, L)) and these µ, L comes from upper neural network (encoder). Therefore while backpropogation all we need is partial derivatives w.r.t. µ, L and ε is irrelevant for taking derivatives.

enter image description here

  • Best video to understand this concept. I would recommend to watch complete video for better understanding but if you want to understand only reparameterization trick then watch from 8 minute. youtube.com/channel/UCNIkB2IeJ-6AmZv7bQ1oBYg – Sherlock Feb 27 at 23:42

I thought the explanation found in Stanford CS228 course on probabilistic graphical models was very good. It can be found here: https://ermongroup.github.io/cs228-notes/extras/vae/

I've summarized/copied the important parts here for convenience/my own understanding (although I strongly recommend just checking out the original link).

So, our problem is that we have this gradient we want to calculate: $$\nabla_\phi \mathbb{E}_{z\sim q(z|x)}[f(x,z)]$$

If you're familiar with score function estimators (I believe REINFORCE is just a special case of this), you'll notice that is pretty much the problem they solve. However, the score function estimator has a high variance, leading to difficulties in learning models much of the time.

So, under certain conditions, we can express the distribution $q_\phi (z|x)$ as a 2-step process.

First we sample a noise variable $\epsilon$ from a simple distribution $p(\epsilon)$ like the standard Normal. Next, we apply a deterministic transformation $g_\phi(\epsilon, x)$ that maps the random noise onto this more complex distribution. This second part is not always possible, but it is true for many interesting classes of $q_\phi$.

As an example, let's use a very simple q from which we sample.

$$z \sim q_{\mu, \sigma} = \mathcal{N}(\mu, \sigma)$$ Now, instead of sampling from $q$, we can rewrite this as $$ z = g_{\mu, \sigma}(\epsilon) = \mu + \epsilon\cdot\sigma$$ where $\epsilon \sim \mathcal{N}(0, 1)$.

Now, instead of needing to get the gradient of an expectation of q(z), we can rewrite it as the gradient of an expectation with respect to the simpler function $p(\epsilon)$.

$$\nabla_\phi \mathbb{E}_{z\sim q(z|x)}[f(x,z)] = \mathbb{E}_{\epsilon \sim p(\epsilon)}[\nabla_\phi f(x,g(\epsilon, x))]$$

This has lower variance, for imo, non-trivial reasons. Check part D of the appendix here for an explanation: https://arxiv.org/pdf/1401.4082.pdf

We have our probablistic model. And want to recover parameters of the model. We reduce our task to optimizing variational lower bound (VLB). To do this we should be able make two things:

  • calculate VLB
  • get gradient of VLB

Authors suggest using Monte Carlo Estimator for both. And actually they introduce this trick to get more precise Monte Carlo Gradient Estimator of VLB.

It's just improvement of numerical method.

Your Answer

By clicking "Post Your Answer", you acknowledge that you have read our updated terms of service, privacy policy and cookie policy, and that your continued use of the website is subject to these policies.

Not the answer you're looking for? Browse other questions tagged or ask your own question.