The reparameterization trick

The Problem with Sampling

Recall the VAE training loop from lesson 14-3:

1. Encoder outputs $\mu_\phi(x)$ and $\sigma_\phi(x)$
2. Sample $z \sim \mathcal{N}(\mu_\phi(x), \sigma_\phi(x)^2 \cdot I)$
3. Decoder produces $\hat{x} = g_\theta(z)$
4. Compute loss, backpropagate

The reparameterization trick is the single idea that makes VAE training work at all. Without it, gradients cannot flow through the sampling step, and the entire architecture is untrainable. It also appears in many other models with stochastic components — diffusion models, normalizing flows, and stochastic policies in RL all use versions of this trick.

Step 2 is the problem. Backpropagation needs to compute $\partial L / \partial \mu_\phi$ and $\partial L / \partial \sigma_\phi$ — the gradients with respect to the encoder parameters. To do that it needs to propagate gradients through the sampling step, which requires $\partial z / \partial \mu$ and $\partial z / \partial \sigma$.

But $z$ is a random draw. It is not a deterministic function of $\mu$ and $\sigma$. Every time step 2 runs, a different $z$ comes out. There is no fixed derivative to place in the computational graph. The gradient is undefined.

The Trick: Move the Randomness Outside

Instead of sampling $z$ directly, rewrite:

Check: is $z$ still distributed as $\mathcal{N}(\mu, \sigma^2 I)$? Yes — adding a constant to a Gaussian shifts its mean, scaling by $\sigma$ scales its standard deviation. The distribution is identical to before.

What has changed: the randomness now lives entirely in $\varepsilon$, which is independent of the encoder parameters. The path from $\mu$ and $\sigma$ to $z$ is now a deterministic function:

This is differentiable. The local gradients are:

Backprop can now flow gradients from the decoder loss all the way back through $z$, through $\mu$ and $\sigma$, and into the encoder network weights. The encoder is trainable.

Visualizing the Computational Graph

Without reparameterization (broken):

The arrow labeled "sample" has no backward pass. Gradients stop at $z$ and never reach $\mu$ or $\sigma$.

With reparameterization (working):

The path $\mu \to z \to \hat{x} \to L$ and $\sigma \to z \to \hat{x} \to L$ are fully differentiable. The noise $\varepsilon$ is a leaf node with no parameters — its gradient is never needed.

A Complete Numerical Walk-Through

Suppose encoder outputs $\mu = 0.6$, $\sigma = 0.4$ for a single latent dimension. We draw $\varepsilon = -0.7$.

Forward pass:

The decoder receives $z = 0.32$ and produces some output. Suppose after computing the full loss, $\partial L / \partial z = 1.5$.

Backward pass through the reparameterization:

These gradients flow backward into the encoder network. The KL gradient (from lesson 14-4) adds on top:

The encoder weights are updated using the sum of both gradient contributions — reconstruction signal (via reparameterization) and regularization signal (from the KL term directly).

Generalizing Beyond Gaussians

The reparameterization trick applies to any distribution with a tractable inverse CDF (quantile function). To sample from distribution $p$:

1. Draw $u \sim \text{Uniform}(0, 1)$
2. Return $z = F^{-1}(u)$ where $F^{-1}$ is the inverse CDF

Since $F^{-1}$ is differentiable (when it exists), the path $\theta \to z$ is differentiable. This works for Exponential, Logistic, Laplace, Beta (approximately), and others. For distributions without easy inverse CDFs, alternatives like implicit reparameterization (Figurnov et al., 2018) generalize the idea further.