3
$\begingroup$

I am trying to understand the reparameterization trick (RPT) used in the calculation of stochastic backpropagation. There are already some excellent answers here and here.

Under usual notation, we can represent the RPT $$\nabla_\theta \mathbb{E}_{p(x;\theta)}[f(x)] = \mathbb{E}_{p(\epsilon)}[\nabla_\theta f(g(\epsilon, \theta))]$$

The overarching objective of this equation is to show that both the gradients ( one with reparamterization and the other with out reparameterization) give rise to the same estimates.

My question:

So, I am trying to get my hands-dirty by calculating the gradients using the reparameterization trick.

Say, we observe a set of univariate random variables $x_1,\dots,x_n$. Some oracle has told us that these random variables come from a Normal distribution $N(\mu,\sigma^2)$.

This is a silly problem obviously, and we can get the estimate of $\mu$ and $\sigma$ from many techniques like MLE, MAP, so on. But, I just want to see how we go about using stochastic gradient descent and reparameterization trick.

Now, let's say we have a node that spits out random variables $x_1, \dots, x_n$. We cannot use the simple backpropagation as the output of the node - $y$ - is a stochastic one. Assume, our loss function $f(y)$ is squared error i.e. $(y-x_i)^2$. Using the LHS of RPT equation,

\begin{align}\nabla_\theta \mathbb{E}_{p(x;\theta)}[f(x)] &= \nabla_\theta \int \sum_{i=1}^{n}(y-x_i)^2 p(y) dy \\ &= \nabla_\theta \int \sum_{i=1}^{n}(y^2+x_i^{2} - 2yx_i) p(y) dy \\ &= \nabla_\theta \sum_{i=1}^{n}(\mu^2+\sigma^2 +x_i^2 -2*(\mu)*(x_i)) \\ \end{align}

Now, if we take the specific gradients with respect to $\mu$ and $\sigma^2$ and set to zero, we can see the updates as $\mu = \frac{\sum_{i=1}^{n}x_i}{n}$ and I am unable to get the estimate of $\sigma^2$. Theoretically, I think the estimate of $\sigma^2 = (\frac{x_i^2}{n} - \mu^2)$

Now I am trying to get the same using the RHS of RPT equation,

\begin{align}\mathbb{E}_{p(\epsilon)}[\nabla_\theta f(g(\epsilon, \theta))] &= \int \nabla_\theta \sum_{i=1}^{n}(\mu+\sigma \epsilon-x_i)^2 p(\epsilon) d\epsilon \\ &= \int \sum_{i=1}^{n} (2(\mu+\sigma \epsilon-x_i) \nabla_\theta (\mu+\sigma \epsilon-x_i)) p(\epsilon) d\epsilon \\ \end{align}

Again, if we take the gradient with respect to $\mu$ specifically, we will have \begin{align} \mathbb{E}_{p(\epsilon)}[\nabla_\mu f(g(\epsilon, \mu,\sigma))] = \int \sum_{i=1}^{n} (2(\mu+\sigma \epsilon-x_i) p(\epsilon) d\epsilon \\ \end{align}

Setting the above equation to zero, we can see that $$\mu = \frac{\sum_{i=1}^{n}x_i}{n}$$

This is the same as obtained from the LHS of RPT equation.

However, I try to show the same for the $\sigma^2$ estimate that both LHS and RHS of RPT equation give the same estimate but it isn't working out.

Can you please help me show the estimate for $\sigma^2$ using the Reparameterization trick for the simple Normal distribution case.

Edit:

Also, can we show from this simple case, that variance of the estimate obtained using RPT is lower than the one obtained without RPT?

$\endgroup$
1
$\begingroup$

To my knowledge, reparameterization trick is used to make the Monte-Carlo estimated ELBO have lower variance. In your formulation you seem to use it in other cases. Also, you shouldn't consider doing calculation by hand given very few machine learning problems can be solved analytically. Have you checked the VAE's code? $\mu$ and $\sigma$ are learnt in a nerual network.

To know more details about VAE, you can check my blog post.

$\endgroup$
  • 1
    $\begingroup$ This is not true. The reparameterisation trick is not used to make the loss function differentiable with respect to a parameter. It is used to decrease the variance. $\endgroup$ – nbro Jun 10 at 15:19
  • $\begingroup$ @nbro Indeed.. Withour reparameterization, the estimator could be \begin{align} \nabla_{\phi} E_{q_\phi(z)} \left [ f_{\theta, \phi}(z) \right] &=\color{red } { E_{q_\phi(z|x)} [ (\log p_\theta(x,z) - \log q_\phi(z|x) ) \nabla_{\phi}\log{q_\phi(z|x) } ] } \\ &\approx \color{blue}{\frac{1}{L} \sum_{l=1}^L f_{\theta, \phi}(z) \nabla_{\phi} \log q_\phi(z^{(i)} |x)} \end{align}, which the VAE paper claims have high variance. $\endgroup$ – JP Zhang Jun 10 at 15:30
  • $\begingroup$ The VAE paper is quite badly written or confusing, IMHO, so I understand that people do not want to read the original research papers, which are usually a little bit painful to follow. In any case, there is a lot of misinformation around, and the primary source of information should be the original work (especially, if it has been accepted by serious journals). $\endgroup$ – nbro Jun 10 at 15:34
  • $\begingroup$ @JPZhang: Great blog. Keep up the good work $\endgroup$ – kasa Jun 11 at 9:47

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

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