The Reparameterization Trick: Differentiable Sampling

In variational autoencoders, training requires taking gradients through an expectation over the latent variable. But standard backpropagation cannot flow through a random sampling node.

If we sample z directly from a Gaussian distribution with mean mu and standard deviation sigma, the sampling operation acts as a barrier to reverse mode differentiation.

The reparameterization trick isolates the stochasticity outside the computational graph. We draw an auxiliary random vector epsilon from a standard normal distribution that does not depend on our network parameters.

We then form z through a deterministic affine transformation: mu plus sigma multiplied by epsilon. Each line joins a sample of epsilon below to its z above. Now the path from the parameters to the loss is deterministic, and the orange arrow is the gradient step it gives.

This works for location scale families, such as the Gaussian. For discrete latent variables there is no such map, and the Gumbel softmax relaxation is used instead.

Your turn. Drag the mean, and set the two standard deviations. Follow the gradient, and bring the expected loss below one tenth.

By separating deterministic parameters from external randomness, the reparameterization trick unlocks scalable deep generative modeling.