Explain the reparameterization trick

Derive a gradient estimator for an expected loss whose sampled latent variable depends on trainable distribution parameters.

Problem

Context. A Gaussian latent variable has trainable mean and positive standard deviation, and a differentiable loss depends on its sample. Training requires the derivative of the expected loss, not merely a derivative of the density.

Task. Express sampling through parameter-independent noise, derive the pathwise gradient, and explain how Monte Carlo estimation uses it. Compare what can change for discrete samples or nondifferentiable transformations.

Constraints. State conditions permitting differentiation under the expectation. Hold base noise fixed during a sample’s derivative; distinguish that sample estimate from the exact expectation gradient. Do not assume every distribution or objective supports the same transformation.

The problem

Let $F(\varphi )=E_{q_{\varphi }(z\mid x)}[f(z)]$ for fixed x. A sampler returning a detached draw supplies no ordinary local autodiff path to its distribution parameters. This is a property of the representation and estimator, not proof that every stochastic expectation is nondifferentiable.

The trick

Reparameterize the sample:

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

For positive standard deviation sigma, the Gaussian sample has the desired distribution while its transform is differentiable. When differentiation and expectation can be interchanged,

If f also depends explicitly on phi, include that additional derivative. For $\mu =1$, $\sigma =2$, $ϵ=0.5$, and $f(z)=(z-3{)}^2$, the sample is two and loss one. The local derivative is $2(z-3)=-2$, so $df/d\mu =-2$ and $df/d\sigma =-1$ for this fixed draw.

The expected loss instead is $F=(\mu -3{)}^2+{\sigma }^2=8$, with derivatives $(-4,4)$. Noise values minus one and plus one give sample gradients $(-8,8)$ and $(0,0)$, averaging to $(-4,4)$. That two-point check matches this quadratic’s moments; it is not Gaussian sampling or a general integration rule.

Lower-variance gradients

Pathwise estimates can reduce variance relative to score-function estimates, but no distribution-free variance ordering or sample-count guarantee holds. Score-function methods, baselines, analytic terms, and other estimators remain useful under their respective conditions.

What an L5 answer sounds like

A complete explanation states the noise law, transform, held-fixed sample, distributional identity, and interchange assumptions before applying the chain rule.

What an L6 answer adds

Change the Encoder’s Scale Output

New task: The encoder emits log variance r instead of standard deviation sigma.

Picture and change: Use $\sigma =\exp (r/2)$ and retain the same epsilon, chaining through $d\sigma /dr=\sigma /2$.

Check: At $r=\log 4$, sigma is two; the reference sample’s $df/dr=(-1)(2/2)=-1$. Using variance four as the standard deviation changes the distribution and derivative.

Use when: An unconstrained encoder output needs conversion to positive scale.

Change the Latent Support

New task: Replace the Gaussian latent with Bernoulli probability p.

Picture and change: Thresholding uniform noise gives $z=1[u<p]$, whose sample derivative is zero almost everywhere.

Check: Nevertheless, $E[z]=p$ has derivative one; averaging those zero derivatives fails to recover it.

Use when: Discrete support breaks the smooth sample-path argument. Use a suitable score-function estimator or a justified relaxation, not a Gaussian substitution.

For framework use, keep mu/log-variance/noise axes aligned, commonly [B,D], and preserve the draw’s differentiable transform. PyTorch Gaussian rsample() supplies that path; sample() ordinarily does not. K samples require $O(KBD)$ noise/transform work plus downstream computation. Detaching z blocks the path; fresh noise invalidates the finite-difference comparison.