Reparameterization trick

Motivation

Wikipedia

 The reparameterization trick allows for the efficient computation of gradients through random variables, enabling the optimization of parametric probability models using stochastic gradient descent

Consider a function of the form:

$$L(\theta )={\mathbb{E}}_{X\sim P_{\theta }}[F(X)]$$

Info

Our goal is to answer the question: How can we compute ${\nabla }_{\theta }L(\theta )$?

Before calculating ${\nabla }_{\theta }L(\theta )$, let’s discuss how to calculate $L(\theta )$. We can use Monte Carlo methods to estimate the expectation by repeatedly sampling $x\sim P_{\theta }$, calculating $Y=F(X)$, and then averaging all $Y$:

$$\begin{array}{rl}L(\theta )& \approx \frac{1}{n}\sum _{i=1}^n{Y}_i\\ \text{where}{Y}_i& =F(X_i),X_i\sim P_{\theta }\end{array}$$

The following is the computational graph that describes how to generate a single sample $Y$ from the parameters $\theta$.

graph TD;  
θ{θ}-->P(P);
P-->X{X};
X-->F;
F-->Y{Y};
P(P)-.->θ{θ};
X{X}-.->P;
F-.->X;
Y{Y}-.->F;

The shape and style of nodes/edges define their semantics:

• Diamond nodes represent data ($\theta ,X,Y$), and act as the input/output of functions.
• Rounded boxes represent stochastic functions ($P_{\theta }$)
• Rectangular boxes represent deterministic functions ($F$):
• Full arrows show the direction data flows for generating a sample $Y$ based on parameters $\theta$.
• The dashed arrows show the direction that gradients flow when differentiating the output $Y$ with respect to the parameters $\theta$.

The computational graph makes it clear how to calculate ${\nabla }_{\theta }L$: We simple accumulate the gradients along the dashed arrows. However, an issue arises since we need to differentiate through a stochastic function $P_{\theta }$, which is not well-defined.

We discuss two options for handling this: the REINFORCE estimator, and the reparameterization trick.

REINFORCE Estimator

The REINFORCE estimator uses the log-derivative trick to rewrite ${\nabla }_{\theta }L(\theta )$ as an expectation that we can calculate directly:

$$\begin{array}{rl}{\nabla }_{\theta }L(\theta )& ={\nabla }_{\theta }{\mathbb{E}}_{x\sim P_{\theta }}[f(x)]\\ & ={\nabla }_{\theta }\sum _x{P}_{\theta }(x)f(x)\text{(Def. of Expectation)}\\ & =\sum _x{\nabla }_{\theta }(P_{\theta }(x)f(x))\text{(Linearity of Expectation)}\\ & =\sum _{x}f(x){\nabla }_{\theta }{P}_{\theta }(x)+P_{\theta }(x){\nabla }_{\theta }f(x)\text{(Product rule)}\\ & =\sum _{x}f(x){\nabla }_{\theta }{P}_{\theta }(x)(f\text{ does not depend on }\theta )\\ & =\sum _{x}f(x)P_{\theta }(x){\nabla }_{\theta }\log P_{\theta }(x)(\text{Log-derivative trick})\\ & ={\mathbb{E}}_{x\sim P_{\theta }}[f(x){\nabla }_{\theta }\log P_{\theta }(x)]\end{array}$$

This final expression is called the REINFORCE estimator for ${\nabla }_{\theta }L(\theta )$. Similar to how we approximated $L(\theta )$, we can approximate this expectation by sampling $X\sim P_{\theta }$ and averaging $f(x){\nabla }_{\theta }\log P_{\theta }(x)$. One issue with this estimator is that it has high variance since it involves multiplying $f(x)$ and ${\nabla }_{\theta }\log P_{\theta }(x)$.

Instead of using the high-variance REINFORCE estimator, we can use the reparameterization trick.

Reparameterization trick

References

• REINFORCE vs Reparameterization Trick