The reparameterization trick

deep-learninggenerative-modelingreparameterization-trickgradients

The formulation

Goal: $\nabla_\phi\,\mathbb{E}_{z\sim q_\phi(z\mid x)}\big[f(z)\big]$, a gradient of an expectation over a $\phi$-dependent distribution. Write $z$ instead as a deterministic function of $\phi$ and a separate noise variable $\epsilon$ drawn from a fixed distribution with no $\phi$ in it:

$$ z = g_\phi(\epsilon,x), \qquad \epsilon\sim p(\epsilon) \quad\text{(fixed, independent of $\phi$)}. $$

The expectation over $q_\phi$ becomes an expectation over the fixed $p(\epsilon)$:

$$ \mathbb{E}_{z\sim q_\phi(z\mid x)}\big[f(z)\big] = \mathbb{E}_{\epsilon\sim p(\epsilon)}\big[f(g_\phi(\epsilon,x))\big]. $$

$\phi$ now sits only inside the function being averaged, not in the distribution being averaged over, so $\nabla_\phi$ can move inside:

$$ \nabla_\phi\,\mathbb{E}_{\epsilon\sim p(\epsilon)}\big[f(g_\phi(\epsilon,x))\big] = \mathbb{E}_{\epsilon\sim p(\epsilon)}\big[\nabla_\phi f(g_\phi(\epsilon,x))\big] \ \approx\ \frac1S\sum_{s=1}^S\nabla_\phi f\big(g_\phi(\epsilon_s,x)\big), \quad \epsilon_s\sim p(\epsilon). $$

Intuition: the randomness did not disappear, it moved to $\epsilon$, a part of the computation $\phi$ has no say over; everything downstream of $\epsilon$ is an ordinary differentiable function of $\phi$ again.