Reparameterization Trick

## The problem: a sampler node has no derivative Training a latent-variable model means optimising something that contains an expectation over a latent variable whose distribution the network itself emits — an objective of the form `E over z ~ q(z|x) of f(z)`. The encoder maps an input to distribution parameters (for a Gaussian latent, a mean vector `mu(x)` and a positive scale `sigma(x)`), a value `z` is drawn from that distribution, and the decoder consumes `z`. Write that literally and the forward graph is `mu, sigma -> [sampler] -> z -> decoder -> loss`. Reverse-mode differentiation walks back from the loss, computes `dL/dz`, arrives at the sampler — and halts. The sampler consumed `mu` and `sigma`, but the number it emitted came from an external random source; inside the graph the output is not a deterministic function of the inputs, so there is no local Jacobian `dz/dmu` or `dz/dsigma` to multiply by. The decoder still trains, because it sits downstream of `z`. The encoder receives nothing from the reconstruction term, which is the half of the model whose whole job is to place inputs in latent space. Worth being precise: the gradient of the expectation with respect to the encoder's parameters exists and is perfectly well defined. What is broken is the *sample path* — the naive graph gives you no way to estimate it. ## The rewrite Draw `eps` from a standard normal — a fixed base distribution with no learnable parameters — and set, elementwise, ``` z = mu + sigma * eps ``` A location-scale shift of a standard normal is a normal with that mean and that scale, so the distribution of `z` is exactly what it was before. What changed is where the randomness enters: `eps` is now an input leaf that no parameter touches, and everything between the parameters and `z` is ordinary differentiable arithmetic. The local derivatives are `dz/dmu = 1` and `dz/dsigma = eps`, so backpropagation gives `dL/dmu = dL/dz` and `dL/dsigma = dL/dz * eps`. The encoder is back on the path. Formally, writing `z = g(theta, eps)`, the trick lets you swap the order of differentiation and expectation: ``` grad_theta E_q[f(z)] = grad_theta E_eps[f(g(theta, eps))] = E_eps[ grad_theta f(g(theta, eps)) ] ``` The swap is legitimate because the density of `eps` no longer depends on `theta` — the parameter has moved from the measure into the integrand. A single `eps` per example therefore gives an unbiased estimate of the gradient, which is why one draw per example per step is standard practice: the batch average and the many steps of training do the averaging for you. ## Which distributions admit the rewrite Any location-scale family works directly: Gaussian, Laplace, logistic, and a uniform on a learned interval via `a + (b - a) * u`. More generally, any distribution you can express as a deterministic, differentiable transform of a parameter-free base variable is reparameterizable — that is the same construction normalizing flows are built from. Gamma and Dirichlet latents are awkward but not hopeless; implicit and accept-reject reparameterizations exist for them. Discrete latents are the real wall. Any map from continuous noise to a discrete value is piecewise constant, so its derivative with respect to the parameters is zero almost everywhere and undefined at the jumps — there is no useful pathwise gradient to be had. That is why a categorical or Bernoulli latent forces a different choice: relax the sample into something continuous, quantize against a codebook and copy the gradient past the lookup, or keep the sample exact and switch to a score-function estimator that only needs the *value* of the objective. ## Practical notes and common bugs - The scale must stay positive, so it is emitted through an unconstrained output and mapped to a positive number; forgetting that produces silent NaNs the first time the network pushes it negative. - Reusing one `eps` for a whole batch is a classic mistake. Each example's gradient stays unbiased, but the draws are perfectly correlated, so the batch average stops averaging the noise down and the latent samples all move together. - Detaching the scale from the graph is another: the model then learns a mean but never learns how uncertain it should be, and the scale drifts wherever the rest of the objective pushes it. - Nothing about the objective changes. The trick is purely a change in how the gradient is estimated, not in what is being optimised. - At generation time you sample `eps` as usual; when you want a deterministic embedding or reconstruction you use `mu` alone and drop the noise.