Visual walkthrough — Reparameterization trick

---

Step 1 — What a "sample from a distribution" actually looks like

WHAT. Our story stars a machine (called the encoder) that, given a picture $x$, spits out two numbers: a center $\mu$ and a spread $\sigma$. Then a dice-roll picks a point $z$ near that center. We write this dice-roll as

$z\sim \mathcal{N}(\mu ,{\sigma }^2)$

• $z$ ::: the random point we draw (a "latent code")
• $\sim$ ::: read "is drawn from" — the wavy line means randomness happens here
• $\mathcal{N}$ ::: the Gaussian (bell curve) distribution
• $\mu$ ::: the center of the bell — where samples cluster
• $\sigma$ ::: the spread — how wide the bell is ($\sigma >0$ always)

WHY. In a VAE the whole point is that $z$ is a guess — a randomly chosen code that could have made $x$. Randomness is a feature, not a bug. But randomness is exactly what will fight us in Step 3.

PICTURE. Look at the bell below. The peak sits at $\mu$; the wider the curve, the bigger $\sigma$. A sampled $z$ is like dropping a dart that lands somewhere under the bell — more often near the middle.

---

Step 2 — What "gradient" means, and why we want one

WHAT. Training means: nudge the encoder's knobs to make a loss $\mathcal{L}$ smaller. The tool that tells us which way to nudge is the gradient, written ${\nabla }_{\varphi }\mathcal{L}$.

• $\mathcal{L}$ ::: the loss — one number saying "how bad are we" (small = good)
• $\varphi$ ::: all the encoder's internal dials (its weights). $\mu$ and $\sigma$ come out of these.
• ${\nabla }_{\varphi }\mathcal{L}$ ::: the arrow of fastest increase of $\mathcal{L}$ as we wiggle $\varphi$; we step against it to go downhill

WHY. Gradients are the currency of backpropagation. To move the encoder at all, we need $\frac{\partial \mathcal{L}}{\partial \mu }$ and $\frac{\partial \mathcal{L}}{\partial \sigma }$ — how the loss reacts when we slide the center or squeeze the spread.

PICTURE. Below, the loss is a valley. The gradient is the steepest-downhill arrow at your feet. No arrow → you're standing blind → no learning.

---

Step 3 — The blockage: gradients cannot cross a dice-roll

WHAT. Try to compute $\frac{\partial z}{\partial \mu }$ directly on $z\sim \mathcal{N}(\mu ,{\sigma }^2)$. It has no value. Nudge $\mu$ a hair; the random $z$ jumps to a fresh random place with no fixed relationship to the nudge.

WHY. A derivative asks "if input moves by a tiny amount, how much does output move?" That question needs the output to be a deterministic function of the input. A dice-roll is not a function of $\mu$ — the same $\mu$ gives different $z$ every time. So the chain of gradients from $\mathcal{L}$ back to $\varphi$ snaps at the sampling node.

PICTURE. The computation graph: a solid green arrow means "gradients flow." The stochastic node (the dice) is a red wall — the arrow dies there.

---

Step 4 — The key fact: a shift and a stretch build any Gaussian

WHAT. Take one fixed dice — the standard normal $ϵ\sim \mathcal{N}(0,1)$ (center $0$, spread $1$). Then

$z=\mu +\sigma \cdot ϵ$

is distributed exactly as $\mathcal{N}(\mu ,{\sigma }^2)$.

• $ϵ$ ::: noise from a fixed bell that never changes, no matter what $\varphi$ is
• $\mu +$ ::: shift — slide the whole bell so its center lands on $\mu$
• $\sigma \cdot$ ::: stretch — widen the bell by a factor $\sigma$

WHY. This is the location–scale property of Gaussians. It lets us manufacture the exact same randomness as Step 1, but sourced from a dice that doesn't know about $\varphi$. Check the two knobs land right:

• center: $\mathbb{E}[\mu +\sigma ϵ]=\mu +\sigma \cdot 0=\mu$ ✓
• spread: $\text{Var}(\mu +\sigma ϵ)={\sigma }^2\cdot 1={\sigma }^2$ ✓

PICTURE. Watch the standard bell (left) get slid by $\mu$ and widened by $\sigma$ into our target bell (right). Same shape, relocated and rescaled.

---

Step 5 — Reroute the randomness: move the dice off to the side

WHAT. Rebuild the computation graph. The dice-roll now happens on $ϵ$, before $\mu$ and $\sigma$ ever touch it. The path that carries $\varphi$ into $z$ is the plain arithmetic $z=\mu +\sigma ϵ$ — no dice on it.

WHY. Randomness still exists (we still draw $ϵ$), but it lives on a branch that has no $\varphi$ in it. The main trunk from $\varphi$ to $z$ is now a smooth, differentiable function $g_{\varphi }(ϵ,x)={\mu }_{\varphi }(x)+{\sigma }_{\varphi }(x)ϵ$. The red wall is gone from the trunk.

PICTURE. Compare with Step 3. The dice moved to a side-input; the trunk $\varphi \to \mu ,\sigma \to z$ is now solid green all the way. Gradients cross freely.

---

Step 6 — Send the gradient through the smooth path (worked numbers)

WHAT. With $z=\mu +\sigma ϵ$ the chain rule now runs. Using the parent's Example 1: encoder gives $\mu =2.0$, $\sigma =0.5$; we draw $ϵ=1.3$; the loss is $f(z)=(z-3{)}^2$.

$z=\underset{\mu }{\underbrace{2.0}}+\underset{\sigma }{\underbrace{0.5}}\cdot \underset{ϵ}{\underbrace{1.3}}=2.65$

$\frac{\partial f}{\partial z}=2(z-3)=2(2.65-3)=-0.7$

• $\frac{\partial f}{\partial z}$ ::: how the loss reacts to a bump in the code $z$ (negative → raising $z$ lowers loss)

Now the two local derivatives of the smooth path:

$\frac{\partial z}{\partial \mu }=1,\frac{\partial z}{\partial \sigma }=ϵ=1.3$

• $\partial z/\partial \mu =1$ ::: shift by $\mu$ moves $z$ one-for-one
• $\partial z/\partial \sigma =ϵ$ ::: the drawn noise is the sensitivity of $z$ to the stretch

Chain them:

$\frac{\partial f}{\partial \mu }=\frac{\partial f}{\partial z}\cdot \frac{\partial z}{\partial \mu }=-0.7\times 1=-0.7$ $\frac{\partial f}{\partial \sigma }=\frac{\partial f}{\partial z}\cdot \frac{\partial z}{\partial \sigma }=-0.7\times 1.3=-0.91$

WHY. Both gradients are finite numbers — gradient descent can now push $\mu$ up (toward the target $3$) and adjust $\sigma$. This is the payoff: the wall of Step 3 is fully crossed.

PICTURE. The graph carries real numbers backward: $-0.7$ flows into $\mu$, $-0.91$ into $\sigma$, each factor labelled on its edge.

---

Step 7 — Edge and degenerate cases (nothing left hidden)

WHAT & WHY, case by case:

(a) $ϵ<0$ flips the $\sigma$-gradient. Redo Step 6 with the same $\mu ,\sigma$ but $ϵ=-1.3$: $z=2.0+0.5(-1.3)=1.35,\frac{\partial f}{\partial z}=2(1.35-3)=-3.3$ $\frac{\partial f}{\partial \sigma }=-3.3\times (-1.3)=+4.29$ Now the $\sigma$-gradient is positive → "shrink $\sigma$." Makes sense: a negative $ϵ$ with a big $\sigma$ dragged $z$ below the target, so we want less spread. The sign of $\partial f/\partial \sigma$ correctly tracks the drawn $ϵ$.

(b) $\sigma \to 0$ (collapse to a point). The stretch vanishes: $z=\mu$ deterministically. The $\mu$-gradient still equals $\partial f/\partial z$ (the $\times 1$ path survives). The $\sigma$-gradient is $ϵ\cdot \partial f/\partial z$ and stays finite. So the trick degrades gracefully into an ordinary deterministic layer — no division, no blow-up.

(c) Multivariate, $d$ dimensions with diagonal covariance. Draw $ϵ\sim \mathcal{N}(0,I_d)$ ($d$ independent standard normals) and apply element-wise: $z={\mu }_{\varphi }(x)+{\sigma }_{\varphi }(x)\odot ϵ$ where $\odot$ is entrywise multiply. Each coordinate is its own copy of Step 4, so the Jacobian is diagonal and backprop stays cheap.

(d) Non-reparameterizable distributions. If the sample is discrete (a category), you cannot write it as smooth-shift-plus-fixed-noise — there's no continuous stretch. Then this exact trick fails; you reach for the Gumbel-Softmax relaxation, or the score-function estimator used in policy gradients. The parent's ELBO objective, whose KL term and reconstruction term both depend on $z$, only backprops cleanly because the Gaussian latent is reparameterizable.

PICTURE. Four mini-panels: $ϵ<0$ (sign flip), $\sigma \to 0$ (collapse), the $d$-dim Hadamard product, and the discrete "no smooth path" wall.

---

The one-picture summary

Everything above in a single frame: the fixed dice sits off on a side branch (blue), $\varphi$ flows through center-and-stretch (green) into $z$, and the gradient walks all the way back. The dashed red is the old dead path we abandoned.

Recall Feynman retelling — say it like you'd explain to a friend

We wanted to train a machine that makes a random guess $z$, and to train it we need gradients — little arrows saying "nudge this dial that way." But you can't send an arrow through a dice-roll: change the dial, the dice still lands wherever it wants, so there's no "how much did the output move." Dead end.

The fix is sneaky and simple. A bell curve with center $\mu$ and width $\sigma$ is just the standard bell (center 0, width 1) picked up, slid over to $\mu$, and stretched by $\sigma$. So instead of rolling a $\mu ,\sigma$-dice, we roll a plain standard-dice $ϵ$ once, off to the side, then compute $z=\mu +\sigma ϵ$ with ordinary arithmetic. The randomness never touches the dials anymore — it's a side ingredient. The path from the dials to $z$ is now plain multiply-and-add, and gradients stroll right through it: $\partial z/\partial \mu =1$, $\partial z/\partial \sigma =ϵ$. Same random behavior as before, but now trainable. If the width shrinks to zero it just becomes a normal deterministic layer; if the thing you're sampling is a discrete category there's no smooth stretch to exploit, and you need a different trick entirely.