The Reparameterization Trick

When working with variational autoencoders (VAEs), you encounter a core challenge: the model's encoder outputs parameters of a probability distribution (typically the mean $μ$ and standard deviation $σ$ of a Gaussian). To generate a latent variable $z$, you must sample from this distribution. However, sampling is a non-differentiable operation, which means that gradients cannot flow backward through the sampling step. This blocks the gradient-based optimization needed to train VAEs using standard techniques like backpropagation.

The reparameterization trick is a clever solution that allows you to sidestep the non-differentiability of sampling. Instead of sampling $z$ directly from a distribution parameterized by $μ$ and $σ$, you rewrite the sampling process as a deterministic function of the distribution parameters and some auxiliary random noise. Specifically, you sample $ε$ from a standard normal distribution ($N(0, 1)$) and then compute the latent variable as:

Here, the randomness is isolated in $ε$, which is independent of the parameters and can be sampled in a way that does not interfere with gradient flow. The computation of $z$ is now a differentiable function of $μ$ and $σ$, so gradients can propagate through the encoder network during training. This enables you to optimize the VAE end-to-end using gradient descent.