Skip to content
mlmentorship

Explain the reparameterization trick

How VAEs propagate gradients through a sampling step. The senior answer explains the why (you can't differentiate through a sample) and the how (move the randomness outside the parameters).

Published · 6 min read ·Core ·Advanced

Math / derivation · active recall

Practice before you read

15 minutes. State assumptions, derive cleanly, check dimensions or limiting cases, then interpret the result.

Back to my workbook How practice works

Math / derivation · closed-book attempt

Explain the reparameterization trick

State assumptions, derive cleanly, check dimensions or limiting cases, then interpret the result.

15:00recommended time

Closing or reloading clears the scratchpad. Only score, weak rubric dimensions, attempt count, and retry date can be stored locally.

30-second answer map

Visual first · depth when needed

See how replacing a parameterized random sample with fixed-distribution noise plus a deterministic transform creates an ordinary backpropagation path from the loss to the encoder parameters.

Preparing the visual…

Asked in: ML breadth and generative-model interviews.

A standard generative-model question. The L4 candidate states the formula. The L6 candidate explains why naive sampling breaks gradients and how the trick fixes it.

The problem

Suppose you want to train a model that samples a latent variable z from a distribution q_phi(z | x) and uses z to reconstruct x. Loss L(theta, phi) depends on z, which is a sample.

To train, you need dL / dphi. But z = sample(q_phi) is a stochastic operation; the gradient through a sample isn’t well-defined in general.

The trick

Reparameterize the sample:

z = mu_phi(x) + sigma_phi(x) * epsilon,  where epsilon ~ N(0, I)

The randomness now comes from epsilon, which doesn’t depend on phi. The sample z is now a deterministic function of (phi, epsilon). The gradient dL / dphi is computed straightforwardly via chain rule.

In short: instead of “sample z from q_phi,” do “sample epsilon from a fixed distribution, then transform deterministically.”

Learning objective

Move randomness outside the parameterized path so ordinary backpropagation can reach the encoder.

Naive sampling compared with a reparameterized Gaussian sample Two stacked computation graphs use solid arrows for the forward pass and dashed arrows for gradients. In the naive graph, encoder parameters phi define q phi of z given x, which produces a random draw z and then a loss. The reverse gradient reaches z but an X marks that ordinary backpropagation stops at the stochastic draw. In the reparameterized graph, the encoder produces mu and sigma while epsilon is sampled independently from a standard normal. Both feed the deterministic transform z equals mu plus sigma times epsilon, then the loss. Dashed arrows form an unbroken reverse path from the loss through z, mu and sigma, and back to phi. No gradient is needed through epsilon. 1 · NAIVE SAMPLE: STOCHASTIC NODE BREAKS THE PATH φencoder params qφ(z|x)distribution zrandom draw L(z)loss ∂L/∂z X · ordinary backprop stops at the draw 2 · REPARAMETERIZED: SAME DISTRIBUTION, OPEN PATH encoder φinput x μφ, σφdifferentiable ε ∼ N(0,I)independent of φ z = μ + σεdeterministic L(z)loss chain rule: ∂L/∂z → ∂L/∂(μ,σ) → ∂L/∂φ epsilon supplies a value; no derivative through its draw is needed
Read it this way: follow solid arrows to sample in either panel. Then follow the dashed gradient backward. The naive graph meets a stochastic draw with no ordinary local derivative; the reparameterized graph holds the sampled epsilon fixed and differentiates through z = mu + sigma * epsilon all the way to phi. The distribution of z is unchanged. Original construction checked against Kingma and Welling's VAE formulation.

Lower-variance gradients

This is the foundational trick for variational autoencoders (VAE). Without it, you’d have to use REINFORCE / score-function gradients, which have much higher variance and need many samples to be useful.

What an L5 answer sounds like

“When you sample z = sample(q_phi(z | x)) and use z in a downstream loss, you can’t backprop through the sampling because the operation is stochastic.

The trick: rewrite the sample as a deterministic function of (parameters, noise), where the noise comes from a fixed distribution.

For a Gaussian: z = mu(x) + sigma(x) * epsilon, epsilon ~ N(0, I). Now z is differentiable w.r.t. mu and sigma, which are computed by the encoder network. Gradients flow through normally.

Used in VAEs, in some RL algorithms (Gumbel-softmax for discrete actions is the discrete analog), in normalizing flows, in continuous-control policy gradient methods.”

This is L5. Mechanism explained, examples given.

What an L6 answer adds

“…some additional points:

It only works for distributions you can reparameterize. Gaussian, Laplace, exponential, uniform: easy. Discrete distributions: hard (requires Gumbel-softmax with continuous relaxation). General distributions: requires implicit differentiation tricks.

Variance is much lower than score-function gradients. For a 1D Gaussian, reparameterization gradient variance scales with the Jacobian of the network; score-function variance scales with 1 / sigma^2 of the sample, which can be huge. Empirically, reparameterization needs 1-10 samples to estimate a useful gradient; REINFORCE often needs 1000+.

For discrete latents (e.g., latent-variable models with categorical z), Gumbel-softmax / concrete distribution is the standard relaxation. Use a continuous relaxation that’s differentiable; anneal the temperature toward zero to make samples nearly discrete. Trade-off: differentiable but biased.

In LLM-RL (RLHF, DPO), reparameterization isn’t used because text generation is discrete and Gumbel-softmax over a vocabulary of 50K tokens is impractical. RLHF uses score-function gradients (PPO), accepting the variance cost. DPO sidesteps this entirely with an analytical objective that doesn’t need sampling at all.”

Tells that get you a strong-hire vote

  • You explain why sampling breaks gradients before stating the trick.
  • You give the Gaussian formula explicitly.
  • You bring up Gumbel-softmax for discrete latents.
  • You compare variance to score-function gradients.
  • You mention DPO sidesteps this for LLM RL.

Tells that get you down-leveled

  • Stating the formula without explanation.
  • Confusion about why naive sampling doesn’t work.
  • No knowledge of discrete-relaxation alternatives.
  • Treating reparameterization as a VAE-only trick (it’s broader).

Common follow-up

“How does this apply to RL?”

The L6 answer:

“Two cases. For continuous-action policies (e.g., robotic control with a Gaussian policy), reparameterize the action sample and backprop through the value function (this is the basic trick behind DDPG and SAC). For discrete-action policies (e.g., RL on discrete decisions), reparameterization doesn’t directly apply; you use score-function gradients (REINFORCE, A2C, PPO) and accept the variance cost, sometimes mitigated by control variates and baselines. The ‘continuous vs discrete’ choice often dominates the algorithm choice in modern RL.”

Related: RLHF and DPO, cross-entropy and softmax, and Bayesian versus frequentist.

Found an error or stale claim? Report a correction or send a private note.