FFVB: Reparameterization Trick

This section describes a control variate technique for variance reduction.

See: Variational Bayes Introduction, Fixed Form Variational Bayes

---

Reparameterization Trick

$$ The reparameterization trick is an attractive alternative to the control variate method. Suppose that for $\theta \sim q_{\lambda }(\cdot )$, there exists a deterministic function $g(\lambda ,\epsilon )$ such that $\theta =g(\lambda ,\epsilon )\sim q_{\lambda }(\cdot )$ where $\epsilon \sim p_{\epsilon }(\cdot )$.

We emphasize that $p_{\epsilon }(\cdot )$ must not depend on $\lambda$. For example, if $q_{\lambda }(\theta )=\mathcal{N}\mathcal{(}\theta \mathcal{;}\mu \mathcal{,}{\sigma }^{\mathcal{2}}\mathcal{)}$ then $\theta =\mu +\sigma \epsilon$ with $\epsilon \sim \mathcal{N}\mathcal{(}\mathcal{0}\mathcal{,}\mathcal{1}\mathcal{)}$. Writing $\text{LB}(\lambda )$ as an expectation with respect to $p_{\epsilon }(\cdot )$

$$\begin{array}{rl}\text{LB}(\lambda )& ={\mathbb{E}}_{\epsilon \sim p_{\epsilon }}(h_{\lambda }(g(\epsilon ,\lambda ))),\end{array}$$

where ${\mathbb{E}}_{\epsilon \sim p_{\epsilon }}(\cdot )$ denotes expectation with respect to $p_{\epsilon }(\cdot )$, and differentiating under the integral sign gives

$$\begin{array}{rl}{\mathrm{\nabla }}_{\lambda }\text{LB}(\lambda )& ={\mathbb{E}}_{\epsilon \sim p_{\epsilon }}({\mathrm{\nabla }}_{\lambda }g(\lambda ,\epsilon {)}^{\mathrm{\top }}{\mathrm{\nabla }}_{\theta }{h}_{\lambda }(\theta ))+{\mathbb{E}}_{\epsilon \sim p_{\epsilon }}({\mathrm{\nabla }}_{\lambda }{h}_{\lambda }(\theta ))\end{array}$$

where the $\theta$ within $h_{\lambda }(\theta )$ is understood as $\theta =g(\epsilon ,\lambda )$ with $\lambda$ fixed.

In particular, the gradient ${\mathrm{\nabla }}_{\lambda }{h}_{\lambda }(\theta )$ is taken when $\theta$ is not considered as a function of $\lambda$. Here, with some abuse of notation, ${\mathrm{\nabla }}_{\lambda }g(\lambda ,\epsilon )$ denotes the Jacobian matrix of size $d_{\theta }\times d_{\lambda }$ of the vector-valued function $\theta =g(\lambda ,\epsilon )$. Note that

$$\begin{array}{rl}{\mathbb{E}}_{\epsilon \sim p_{\epsilon }}({\mathrm{\nabla }}_{\lambda }{h}_{\lambda }(\theta ))={\mathbb{E}}_{\epsilon \sim q_{\epsilon }}({\mathrm{\nabla }}_{\lambda }{h}_{\lambda }(\theta =g(\epsilon ,\lambda )))& =-{\mathbb{E}}_{\epsilon \sim q_{\epsilon }}({\mathrm{\nabla }}_{\lambda }\log q_{\lambda }(\theta =g(\epsilon ,\lambda )))\\ & =-{\mathbb{E}}_{\theta \sim q_{\lambda }}({\mathrm{\nabla }}_{\lambda }\log q_{\lambda }(\theta ))=0,\end{array}$$

hence

$$\begin{array}{}\text{(24)}& {\mathrm{\nabla }}_{\lambda }\text{LB}(\lambda )={\mathbb{E}}_{\epsilon \sim q_{\epsilon }}({\mathrm{\nabla }}_{\lambda }g(\lambda ,\epsilon {)}^{\mathrm{\top }}{\mathrm{\nabla }}_{\theta }{h}_{\lambda }(\theta )).\end{array}$$

The gradient $\text{(24)}$ can be estimated unbiasedly using i.i.d samples ${\epsilon }_s\sim p_{\epsilon }(\cdot )$, $s=1,\dots ,S$, as

$$\begin{array}{}\text{(25)}& \widehat{{\mathrm{\nabla }}_{\lambda }\text{LB}}(\lambda )=\frac{1}{S}\sum _{s=1}^S{\mathrm{\nabla }}_{\lambda }g(\lambda ,{\epsilon }_s{)}^{\mathrm{\top }}{\mathrm{\nabla }}_{\theta }\{h_{\lambda }(g(\lambda ,{\epsilon }_s))\}.\end{array}$$

The reparametrization gradient estimator $\text{(25)}$ is often more efficient than alternative approaches to estimating the lower bound gradient, partly because it takes into account the information from the gradient ${\mathrm{\nabla }}_{\theta }{h}_{\lambda }(\theta )$.

In typical VB applications, the number of Monte Carlo samples $S$ used in estimating the lower bound gradient can be as small as 5 if the reparameterization trick is used, while the control variates method requires an $S$ of about hundreds or more. However, there is a dilemma about choosing $S$ that we must be careful of.

With the reparameterization trick, a small $S$ might be enough for estimating the lower bound gradient, however, we still need a moderate $S$ in order to obtain a good estimate of the lower bound if lower bound is used in the stopping criterion. Also, compared to score-function gradient, FFVB approaches that use reparameterization gradient require not only the model-specific function $h(\theta )$ but also its gradient ${\mathrm{\nabla }}_{\theta }h(\theta )$.

Algorithm 6 provides a detailed implementation of the FFVB approach that uses the reparameterization trick and adaptive learning. A small modification of Algorithm 6 (not presented) gives the implementation of the FFVB approach that uses the reparameterization trick and natural gradient.

---

Next: GVB with Cholesky decomposed covariance

---