The Reparameterization Trick and Gradient Flow Through Sampling

A computation graph is differentiable only where every node is a deterministic function of its inputs. A node that draws a random number is not: its output changes with the seed, not with the parameters. This is the entire problem of training a model with a stochastic hidden layer, and the two families of solutions to it differ in a way that matters for variance, not just for convenience.

Consider the standard encoder. A network maps an input x to a mean and a log variance, and the latent code is a sample from the resulting Gaussian. The sample feeds a decoder, the decoder produces a reconstruction, and the reconstruction gives a loss. The gradient of that loss with respect to the encoder weights must pass through the sampling node. It cannot, in the literal sense, because the node has no derivative with respect to the parameters: perturbing the mean by a tiny amount does not continuously perturb the realised sample, it merely changes the distribution the sample came from.

Two ways to see a random draw

There are two distinct ways to write the statement “z is drawn from a normal distribution with mean μ and standard deviation σ”. The first treats the draw as an opaque primitive: z ← sample(μ, σ). The second writes the draw as a transformation of a fixed, parameter-free random variable:

z = μ + σ · ε,    ε ~ N(0, 1)

The two statements describe the same distribution. The difference is that the second exposes μ and σ as ordinary multiplicands in an arithmetic expression. Once the noise ε is drawn, the rest of the expression is deterministic and differentiable, and the gradient of the loss with respect to μ is the gradient with respect to z, which is exactly what the decoder already computes. Nothing needs to be known about the normal distribution’s density to differentiate the sample; the density only determines how ε is generated, and that step involves no parameters.

This rewriting is the reparameterisation trick. It works for any distribution that can be expressed as a deterministic function of a parameter-free base distribution, which includes the Gaussian, the uniform, the Laplace, the exponential, and the logistic. It fails for discrete distributions, because a categorical draw cannot be written as a differentiable function of a fixed uniform variate; the usual workaround, the Gumbel-max construction with a temperature, is a relaxation rather than an identity.

The score-function estimator and its variance

The general-purpose alternative requires no such structure. For any distribution q(z | φ),

∇_φ E_q[ f(z) ] = E_q[ f(z) · ∇_φ log q(z | φ) ]

which is an identity obtained by moving the gradient inside the integral and multiplying by q/q. It is unbiased for any differentiable density with support that does not depend on φ. Its cost is variance. The estimator multiplies a scalar-valued f by a vector-valued score, and for a d-dimensional latent the score has d components whose magnitudes grow with d. When f is roughly constant across samples, the expectation of the product is zero but individual realisations are not, so the estimator’s signal-to-noise ratio collapses as the dimension grows.

The Gaussian score is −(z − μ)/σ², which is linear in the deviation of the sample from the mean. Samples far from the mean therefore contribute large terms, and in high dimensions samples are far from the mean in at least some coordinates almost always. The variance also depends on the scale of f, so a reconstruction loss in the thousands of nats produces a proportionally noisier estimator than one in the tens.

The practical consequence is that the score-function estimator needs many samples per gradient step to be usable, and its variance reduction techniques (baselines, control variates, local expectations) are all attempts to subtract the part of f that does not depend on z. Reparameterisation avoids the problem entirely by not multiplying by a score.

Deriving the reparameterized gradient

Write the bound as an expectation over the noise distribution rather than over the latent:

L = E_{ε ~ N(0,1)}[ log p(x | μ + σ·ε) ] − KL( q(z | x) || p(z) )

The expectation is now over a distribution that carries no parameters. The gradient with respect to the encoder parameters passes through the sample path: ∂L/∂μ = E[ ∂log p(x | z)/∂z ], and ∂L/∂σ = E[ ∂log p(x | z)/∂z · ε ]. The second expression shows why the estimator carries information about the scale: ε has unit variance, so the product has variance equal to the variance of the decoder gradient, and the estimate of the scale gradient is as well conditioned as the estimate of the location gradient.

The KL term is usually computed analytically rather than estimated, which removes its variance from the total. It is also differentiable in closed form, so there is no reason to sample it. The one-sample estimator of the reconstruction term is unbiased for the gradient of the expectation, but it is not unbiased for the gradient of the log of the expectation; the distinction matters when the decoder likelihood is very peaked, because then a single sample can produce a gradient far from the mean gradient.

import numpy as np

def encode(x, W_mu, W_lv):
    return x @ W_mu, x @ W_lv          # mean and log variance, not sigma

def sample(mu, log_var, rng):
    eps = rng.standard_normal(mu.shape)   # parameter-free noise
    return mu + np.exp(0.5 * log_var) * eps, eps

def backward(dloss_dz, log_var, eps):
    # z = mu + exp(0.5 * log_var) * eps, so d z / d mu = 1
    dloss_dmu = dloss_dz
    # d z / d log_var = 0.5 * exp(0.5 * log_var) * eps
    dloss_dlogvar = 0.5 * np.exp(0.5 * log_var) * eps * dloss_dz
    return dloss_dmu, dloss_dlogvar

The backward function is the piece that is easy to get wrong. Differentiating with respect to the log variance rather than the standard deviation introduces a factor of one half, and forgetting it produces a model that still trains but whose learned variances are systematically mis-scaled. The factor arises because the chain rule is applied to exp(0.5 · log_var), not to log_var directly.

Why the estimator is not simply “lower variance”

Reparameterisation is usually described as a variance reduction technique, and that description is incomplete. The two estimators are not estimating the same quantity in a comparable way: the score-function estimator estimates the gradient of an expectation over a parameterised measure, while the reparameterized estimator estimates the gradient of an expectation over a fixed measure composed with a parameterised map. When the map is linear in the parameters and the base distribution has finite variance, the reparameterized estimator’s variance is bounded by the variance of the integrand times the squared scale, independent of the latent dimension. The score-function estimator’s variance grows with the dimension because the score itself grows.

There are cases where the ordering reverses. If f is discontinuous in z, the reparameterized path cannot differentiate through the discontinuity, while the score-function estimator still produces a valid estimate because it never differentiates f. If the base distribution has heavy tails, the reparameterized estimator inherits them multiplied by the integrand’s scale. And if the parameterisation is misspecified, for instance a scale parameter that can go negative so that the map is not a bijection onto the support, the estimator is biased in a way that no amount of sampling fixes.

Where the trick breaks and what replaces it

The construction requires a base distribution whose dimension matches the latent and a map that is differentiable almost everywhere. Three situations break it in practice.

Discrete latents have no such map. The Gumbel-softmax relaxation replaces the categorical sample with a continuous approximation controlled by a temperature, which recovers the discrete distribution in the zero-temperature limit but has gradients that vanish or explode at the extremes. The straight-through estimator simply substitutes the identity for the sampling derivative, which is biased and works because the bias is often small relative to the gradient magnitude.

Distributions with constrained support need a bijection. A gamma or Dirichlet latent can be reparameterised through a composition of inverse-CDF approximations and log transforms, but each transform adds a Jacobian term to the density, and the composition’s numerical stability at the tails determines whether the gradient survives. The implicit reparameterisation approach differentiates the inverse CDF numerically, which extends coverage but requires the CDF to be differentiable and cheap.

Deterministic decoders with non-Gaussian observation models break the assumption that the reconstruction term is smooth in z. A decoder that quantises its output, or one that produces a categorical distribution over pixels, has a piecewise-constant or saturating dependence on the latent, and the reparameterized gradient through it is zero or uninformative in exactly the regions where the model is confident.

The construction is used wherever a continuous latent needs a gradient, which includes the variational autoencoder described in the bound it optimises and the stochastic encoders of sequential latent models. Diffusion models avoid the problem from the other direction: their training objective is a regression on noise, with the noise drawn independently of the parameters, so no sampling node sits between the parameters and the loss. That structural choice is why denoising score matching can be trained with plain squared error on a quantity that is defined through a stochastic process.

The bookkeeping that determines whether it trains

Three implementation details decide whether the estimator is correct, independent of the mathematics.

The noise must be drawn fresh per example per step, from a generator that is independent of the model’s parameters. Reusing a fixed noise tensor across steps turns the stochastic objective into a deterministic one with a fixed set of samples, which is not the same objective and can produce a model that fits those samples and nothing else.

The variance parameterisation must be constrained. Predicting log variance rather than variance keeps the scale positive without clipping, and predicting it in a range that excludes extremely negative values prevents the collapse of the sample to its mean. The clamp is a modelling decision, not a numerical afterthought: a variance of zero makes the latent deterministic and removes the stochasticity the model was supposed to have.

The KL term’s analytic form assumes the prior is the standard normal and the approximate posterior is diagonal. Changing the prior to a mixture or a learned distribution invalidates the closed form and forces either a Monte Carlo estimate of the KL or a different divergence. The Monte Carlo estimate reintroduces variance into a term that had none, and its variance is not reduced by reparameterisation because the KL is already written as an expectation over the same samples.

For all three, the failure is silent: the loss decreases, the samples look plausible, and the latent structure is wrong in a way that only shows when the model is asked to interpolate or to generate from a region of the prior that training did not cover.