Reparameterization trick for VAE backpropagation

What's actually being asked

Why can't we just backprop through a random sample z ~ q(z|x), how the reparameterization trick fixes that, the derivation of z = μ + σ⊙ε, and why this matters for training.

Answer

The problem — sampling has no gradient

In a VAE, the encoder (parameters φ) looks at input x and outputs the parameters of a distribution, μ(x) and σ(x), defining q(z|x) = N(μ, σ²). To train the decoder we need an actual latent vector z, drawn from this distribution.

The naive way is z ~ N(μ(x), σ(x)²) — sample directly. Yaha dikkat yeh hai: sampling is a stochastic operation sitting between the encoder and the decoder in the computation graph. Backpropagation needs every node to be a deterministic, differentiable function of its inputs so the chain rule can compute ∂L/∂μ and ∂L/∂σ. A random draw isn't a function of μ and σ in that sense — perturbing μ slightly doesn't move the sampled z in any smooth, well-defined direction, so gradients can't flow back through the sampling node into the encoder weights. Without a gradient here, the encoder can never learn.

The fix — reparameterization

The trick is to move all the randomness outside the parameters we need gradients for. Instead of sampling z directly from N(μ,σ²), sample a fixed, parameter-free noise variable and apply a deterministic transform:

z = μ + σ ⊙ ε, where ε ~ N(0, I)

(⊙ is elementwise multiplication.) This works because of the location–scale property of the Gaussian: if ε ~ N(0,1), then for any constants μ, σ, the affine transform μ + σε is itself Gaussian with E[μ + σε] = μ + σ·E[ε] = μ Var(μ + σε) = σ²·Var(ε) = σ²·1 = σ²

So μ + σε ~ N(μ, σ²) — exactly the distribution we wanted to sample from. Nothing about what gets sampled has changed; only how it's computed has.

Now look at the computation graph: ε is sampled once, treated as a constant input (no gradient needed through it), and z is a plain differentiable function of μ and σ — addition and elementwise product. So: ∂z/∂μ = 1, ∂z/∂σ = ε

These are ordinary, well-defined gradients, and the chain rule carries ∂L/∂z straight back into ∂L/∂μ and ∂L/∂σ, and from there into every encoder weight.

Significance in VAE training

• It makes the whole encoder → latent → decoder pipeline end-to-end differentiable, so both networks train with a single backward pass under ordinary SGD/Adam — no special-case gradient estimator needed.
• It keeps the actual stochasticity (needed because a VAE is a probabilistic model that must sample, not just encode deterministically) while removing stochasticity from the parts of the graph that need gradients.
• Without it, we'd have to fall back on high-variance score-function estimators (like REINFORCE) for continuous latents — technically possible, but far noisier and much slower to converge in practice.

Worked example

Source numbers aren't given, so here's the method on illustrative values. Say encoder outputs μ = [0.2, -0.1], σ = [0.5, 0.3] for a 2-D latent. Sample ε = [1.0, -2.0] from N(0,I). Then: z = [0.2 + 0.5×1.0, -0.1 + 0.3×(-2.0)] = [0.7, -0.7] This z is fed to the decoder to produce x̂. If the reconstruction loss gives ∂L/∂z = [0.4, -0.2], then ∂L/∂μ = [0.4, -0.2] directly, and ∂L/∂σ = [0.4×1.0, -0.2×(-2.0)] = [0.4, 0.4] — both flow straight into the encoder.

Exam tip

Don't say "the trick removes randomness" — it doesn't; ε is still random. What it removes is randomness from the path we differentiate through. Examiners specifically like to check whether you know ∂z/∂σ = ε, not just the sampling formula.