Reparameterization trick

---

This note is from a discussion with Claude Opus 4.7 when reading VAE tutorial.

---

The trick that makes VAE training work end-to-end. It lets gradients flow through a sampling step.

The problem

We want to maximize an objective of the form:

$$\mathcal{L}(\varphi )={\mathbb{E}}_{q_{\varphi }(z)}[f(z)]$$

Gradient descent needs ${\nabla }_{\varphi }\mathcal{L}$. But $\varphi$ controls the distribution we sample from, not the function inside, so:

$${\nabla }_{\varphi }{\mathbb{E}}_{q_{\varphi }(z)}[f(z)]\ne {\mathbb{E}}_{q_{\varphi }(z)}[{\nabla }_{\varphi }f(z)]$$

The expectation itself depends on $\varphi$ through the sampling density. The gradient doesn’t commute with the expectation.

The trick

If $z$ can be written as a deterministic function of $\varphi$ and a $\varphi$-free noise variable:

$$z=g_{\varphi }(ϵ),ϵ\sim p(ϵ)$$

then the expectation rewrites with the fixed base distribution outside:

$${\mathbb{E}}_{q_{\varphi }(z)}[f(z)]={\mathbb{E}}_{p(ϵ)}[f(g_{\varphi }(ϵ))]$$

Now the gradient passes through cleanly:

$${\nabla }_{\varphi }{\mathbb{E}}_{p(ϵ)}[f(g_{\varphi }(ϵ))]={\mathbb{E}}_{p(ϵ)}[{\nabla }_{\varphi }f(g_{\varphi }(ϵ))]$$

A Monte Carlo estimate: sample $ϵ$, compute ${\nabla }_{\varphi }f(g_{\varphi }(ϵ))$ via autograd, done.

Gaussian case (the VAE one)

For $q_{\varphi }(z|x)=\mathcal{N}(z;{\mu }_{\varphi }(x),{\sigma }_{\varphi }(x{)}^{2}I)$:

$$z={\mu }_{\varphi }(x)+{\sigma }_{\varphi }(x)\odot ϵ,ϵ\sim \mathcal{N}(0,I)$$

The encoder outputs $\mu$ and $\log {\sigma }^2$ (logvar, for numerical stability — exponentiating keeps ${\sigma }^2$ positive without constraints). The sample $z$ is a deterministic function of $\mu$, $\sigma$, $ϵ$; backprop flows through $\mu$ and $\sigma$ into the encoder.

This is why VAE encoder heads output two things rather than a sample.

Versus the Score function estimator

The score-function (REINFORCE) estimator works for any $q_{\varphi }$ without requiring a reparameterization:

$${\nabla }_{\varphi }{\mathbb{E}}_{q_{\varphi }}[f(z)]={\mathbb{E}}_{q_{\varphi }}[f(z){\nabla }_{\varphi }\log q_{\varphi }(z)]$$

It’s universal but high-variance — $f(z)$ multiplies the score, so noise in $f$ amplifies into the gradient estimate. Variance reduction (baselines, control variates) helps but rarely closes the gap.

Reparameterization is lower-variance because it uses gradient information from $f$ directly (pathwise derivative carries shape information about $f$, not just scalar values). When applicable, prefer it.

Other applications

The same pathwise-gradient construction appears wherever a network needs to inject learnable-scale noise and still get gradients through the noise scale:

• Noisy top-k gating in Mixture of Experts (Shazeer 2017): $H(x{)}_i=(x{W}_g{)}_i+ϵ\cdot \mathrm{softplus}((x{W}_n{)}_i)$ with $ϵ\sim \mathcal{N}(0,1)$. Structurally identical to $z=\mu +\sigma \odot ϵ$ but for gating logits rather than posterior samples. Purpose is exploration and load balancing across experts, not posterior approximation, but the mechanism is the same pathwise gradient through $W_n$.
• Stochastic policies in continuous-control RL (SAC): action $a={\mu }_{\varphi }(s)+{\sigma }_{\varphi }(s)\odot ϵ$, gradients flow into the policy network through the action. Replaces high-variance Score function policy gradients with low-variance pathwise gradients of the Q-value — one of the reasons SAC works well at scale.

Limitations

Reparameterization requires $z$ to be a differentiable function of $\varphi$ given $ϵ$. This rules out:

• Discrete latents: can’t differentiate through a categorical sample. Workarounds: Gumbel-Softmax (continuous relaxation, biased but low-variance), straight-through estimator (zero-bias for the forward pass, biased gradient — the choice in VQ-VAE), or fall back to Score function gradients. See also DPO for discrete sampling.
• Distributions without nice reparameterizations: gamma, Dirichlet, etc. Generalized reparameterization (Ruiz et al. 2016) and implicit reparameterization (Figurnov et al. 2018) extend the trick using inverse CDFs or rejection sampling.

The discrete case is exactly what motivates VQ-VAE’s codebook + straight-through gradient design.