The Reparameterisation Trick

You want to train a network whose forward pass includes drawing a sample. The encoder outputs a mean and a variance, you sample a latent from that Gaussian, the decoder reconstructs from it, and you compute a loss. Then you call backward, and there is nothing to differentiate: torch.normal(mu, sigma) is a stochastic node, and the gradient of a sample with respect to the parameters of the distribution it came from is not defined in the usual sense. The trick that resolves this is an algebraic identity so simple it is easy to underrate, and it is what made variational autoencoders trainable at all.

Moving the randomness out of the path

Instead of sampling \(z \sim \mathcal{N}(\mu_\theta, \sigma_\theta^2)\), sample \(\epsilon \sim \mathcal{N}(0, 1)\) from a distribution with no parameters, and compute

The distribution of \(z\) is identical. But now \(z\) is a deterministic, differentiable function of \(\mu_\theta\) and \(\sigma_\theta\), with \(\epsilon\) entering as an input from outside the computational graph, exactly like a data batch. Autograd handles it without knowing anything special happened: \(\partial z / \partial \mu = 1\) and \(\partial z / \partial \sigma = \epsilon\).

The move is general to any distribution admitting a differentiable transform of parameter-free noise. Gaussians are the easy case. Location-scale families in general work. The Gumbel-max reparameterisation extends the idea to categorical variables by relaxing the argmax to a softmax, giving the Gumbel-softmax or concrete distribution, which trades exactness for differentiability with a temperature parameter controlling how close the relaxation sits to a true one-hot sample.

The variance argument

An alternative exists and predates it. The score-function estimator, known as REINFORCE in reinforcement learning, uses the identity

This requires no differentiable path through the sampler at all and works for discrete variables. It is also, in practice, close to unusable for high-dimensional continuous latents without heavy variance reduction, because the estimator only sees \(f(z)\) as a scalar multiplier and gets no information about which direction in \(z\) would have increased it. The reparameterised estimator differentiates through \(f\) itself, so it uses the gradient \(\nabla_z f\), which carries far more information per sample.

The practical gap is large. Reparameterised gradients typically train a VAE with a single Monte Carlo sample per data point; score-function gradients on the same model need baselines, control variates, and many more samples to reach comparable variance. This is why almost every continuous latent variable model uses reparameterisation, and why discrete latent variable models remain harder.

Where it shows up beyond VAEs

Bayesian neural networks with weight uncertainty use it to sample weights differentiably. Variational inference over any continuous parameter uses it. Stochastic policies in continuous-control reinforcement learning use it, which is exactly the difference between the reparameterised soft actor-critic update and a policy-gradient update. Any time a model must "sample and then be trained on what it sampled", this is the mechanism that makes the training signal flow backwards through the sample.

When it breaks

Discrete latents have no exact version. The Gumbel-softmax relaxation is biased: you train on a soft sample and deploy with a hard one, and the gap between them grows as temperature falls. Annealing temperature during training trades bias against variance, and there is no setting that removes both. Straight-through estimators make the opposite choice, using a hard sample forward and a soft gradient backward, accepting a gradient that does not correspond to the forward computation.

Low variance is not the same as low bias in the objective. Reparameterisation gives an unbiased gradient of the ELBO, but the ELBO is itself a bound, so a low-variance gradient of a loose bound still optimises the wrong thing. Tighter bounds, such as the importance-weighted objective, reduce that gap at the cost of more samples per step.

The variance parameter needs care at the boundary. Networks predict \(\log \sigma\) rather than \(\sigma\) so the value stays positive and the scale stays numerically reasonable. Predicting \(\sigma\) directly and clamping it produces zero gradients whenever the clamp is active, which shows up as a latent dimension that stops learning and never recovers.

Very small sigma degenerates. If the encoder drives \(\sigma \to 0\), the sample becomes deterministic and the model becomes an ordinary autoencoder, with the KL term pushed to whatever the mean alone incurs. That is not a numerical bug; it is posterior collapse arriving through the variance channel, and it is invisible if you only monitor total loss.