Worked examples — Reparameterization trick

This is the "roll up your sleeves" companion to the parent note. There we built the transform $z=\mu +\sigma ϵ$ and saw why it lets gradients flow. Here we hit every kind of input the trick can face — positive and negative noise, the mean at zero, the variance collapsing to zero, many dimensions, and even the case where the trick cannot be used at all.

---

The scenario matrix

Every situation the trick can be handed falls into one of these cells. The examples below are labelled with the cell they cover, so by the end no cell is left dark.

#	Cell class	What is special	Example
C1	$ϵ>0$, ordinary $\mu ,\sigma$	baseline sign of gradient	Ex 1
C2	$ϵ<0$	sign of $\partial f/\partial \sigma$ flips	Ex 2
C3	$\mu =0$	shift term vanishes	Ex 3
C4	$\sigma \to 0$ (degenerate)	distribution collapses to a point	Ex 4
C5	Multivariate, diagonal $\Sigma$	element-wise trick, Jacobian	Ex 5
C6	Monte-Carlo, $S$ samples	averaging, variance of estimate	Ex 6
C7	Real-world word problem	end-to-end VAE gradient step	Ex 7
C8	Exam twist — discrete latent	trick fails, needs alternative	Ex 8

Recall Why these are the only cells

$z=\mu +\sigma ϵ$ has exactly two learnable knobs ($\mu ,\sigma$) and one random driver ($ϵ$). The behaviour changes only with the sign of $ϵ$ (C1/C2), the boundary values of the knobs ($\mu =0$ in C3, $\sigma =0$ in C4), the dimension (C5), how we estimate the expectation (C6), and whether the latent is even continuous at all (C8). That exhausts the space.

We reuse one loss throughout so the arithmetic stays comparable: $f(z)=(z-3{)}^2,\frac{\partial f}{\partial z}=2(z-3).$ This loss just says "we want $z$ near the target $3$."

---

Example 1 — Cell C1: positive noise, ordinary parameters

Forecast: $z$ lands a bit above $2.0$ (the nudge is positive), still below target $3$, so both gradients should be negative (we want to push $\mu ,\sigma$ up to raise $z$ toward 3).

1. Compute $z$. $z=\mu +\sigma ϵ=2.0+0.5\times 1.3=2.65$. Why this step? The deterministic transform turns the abstract "sample" into a concrete number we can differentiate.
2. Loss. $f=(2.65-3{)}^2=0.1225$. Why this step? We need the current error to know how hard to push.
3. Upstream gradient. $\partial f/\partial z=2(2.65-3)=-0.7$. Why this step? The chain rule needs "how loss changes with $z$" before we split into $\mu$ and $\sigma$.
4. Local Jacobian. $\partial z/\partial \mu =1$, $\partial z/\partial \sigma =ϵ=1.3$. Why this step? These are the only derivatives the trick introduces — differentiating $z=\mu +\sigma ϵ$, a shift contributes $1$ and a stretch contributes $ϵ$; without them the chain rule has no link from $z$ back to the parameters.
5. Chain. $\partial f/\partial \mu =-0.7\times 1=-0.7$; $\partial f/\partial \sigma =-0.7\times 1.3=-0.91$.

Verify: Both negative, matching the forecast. Sanity: nudging $\mu$ up by a tiny $h$ gives $z=2.65+h$, loss $\approx (2.65+h-3{)}^2$, whose slope at $h=0$ is $2(-0.35)=-0.7$. ✓

---

Example 2 — Cell C2: negative noise flips the $\sigma$ gradient

Forecast: Now $z$ dips below $\mu$, further from $3$, so $\partial f/\partial \mu$ stays negative — but since $\partial z/\partial \sigma =ϵ<0$, raising $\sigma$ now lowers $z$, so $\partial f/\partial \sigma$ should flip positive.

1. $z$. $z=2.0+0.5\times (-1.3)=1.35$. Why this step? Negative $ϵ$ pulls the sample below the mean; we need the concrete $z$ before any derivative.
2. $\partial f/\partial z$. $2(1.35-3)=-3.3$. Why this step? The chain rule always starts from "how the loss reacts to $z$"; here we are much further below target, so the pull is stronger.
3. Jacobian. $\partial z/\partial \mu =1$, $\partial z/\partial \sigma =ϵ=-1.3$. Why this step? Same transform as Ex 1, so the same two derivatives — but the stretch derivative now carries the sign of the noise, which is exactly what will flip the $\sigma$-gradient.
4. Chain. $\partial f/\partial \mu =-3.3\times 1=-3.3$; $\partial f/\partial \sigma =-3.3\times (-1.3)=4.29$. Why this step? Multiplying the upstream derivative by each local derivative is the chain rule that carries the loss's demand all the way back to $\mu$ and $\sigma$.

Verify: $\partial f/\partial \sigma =+4.29>0$ — the sign flipped exactly as forecast, because a bigger stretch drags this particular sample down and away from the target. This is the whole point of C2: the same encoder gets opposite $\sigma$-advice depending on the sign of the noise, and averaging over many $ϵ$ (Ex 6) is what makes the direction reliable. ✓

---

Example 3 — Cell C3: mean sitting at zero

Forecast: With $\mu =0$ the shift term disappears, so $z=\sigma ϵ$ is pure scaled noise. We want $z\to 0$, so we expect gradients that push $\sigma$ down.

1. $z$. $z=0+2.0\times 0.4=0.8$. Why this step? The $\mu$ term is literally $0$; $z$ is entirely the stretched nudge, and we still need its value to differentiate.
2. $\partial f/\partial z$. $f=z^2\Rightarrow 2z=2\times 0.8=1.6$. Why this step? A new loss means a new upstream derivative; everything downstream reuses this number.
3. Jacobian. $\partial z/\partial \mu =1$, $\partial z/\partial \sigma =ϵ=0.4$. Why this step? Differentiating $z=\mu +\sigma ϵ$ gives the same two local slopes as always; note $\mu =0$ does not change them, which is the whole lesson of this cell.
4. Chain. $\partial f/\partial \mu =1.6\times 1=1.6$; $\partial f/\partial \sigma =1.6\times 0.4=0.64$. Why this step? The chain rule multiplies upstream by local; this converts "loss wants smaller $z$" into concrete parameter directions.

Verify: $\partial f/\partial \sigma =0.64>0$ → gradient descent decreases $\sigma$, shrinking the noise so $z$ hugs $0$. Exactly the "collapse toward the target" behaviour we forecast. Note $\mu =0$ did not break anything: the Jacobian $\partial z/\partial \mu =1$ is independent of the value of $\mu$. ✓

---

Example 4 — Cell C4: variance collapsing ($\sigma \to 0$), the degenerate limit

Forecast: As spread vanishes, the "bell" narrows to a spike at $\mu$; the sample becomes deterministic. The $\mu$-gradient should survive. The $\sigma$-gradient does not vanish — it settles at a finite value — but as we will qualify, that gradient can no longer buy much loss reduction per unit of $z$-movement because $z$ has stopped wandering.

1. Sample limit. $z=\mu +\sigma ϵ\to \mu =2.65$ for every $ϵ$. Why this step? Any finite $ϵ$ times $\sigma \to 0$ is $0$, so the shift term $\mu$ is all that remains.
2. Distribution limit. $\mathcal{N}(\mu ,{\sigma }^2)\to$ a Dirac spike at $\mu$ — no randomness left. Why this step? Variance ${\sigma }^2\to 0$ means all probability mass piles onto one point; this is what "degenerate" means.
3. $\mu$-gradient. $\partial f/\partial \mu =2(\mu -3)\cdot 1=2(2.65-3)=-0.7$, unchanged by $\sigma$. Why this step? $\partial z/\partial \mu =1$ regardless of $\sigma$, so the mean channel is always alive — the trick never loses its grip on $\mu$.
4. $\sigma$-gradient. $\partial f/\partial \sigma =2(z-3)ϵ$. At the limit $z\to \mu$, this $\to 2(\mu -3)ϵ=-0.7ϵ$ — a nonzero, finite value for any $ϵ\ne 0$. Why this step? We differentiate the same product $2(z-3)ϵ$ and take the limit of $z$; the derivative itself does not disappear.

Verify: The trick stays perfectly differentiable in the degenerate limit — no division by $\sigma$ anywhere (contrast the density of a Variational Autoencoder (VAE), which does divide by $\sigma$). Careful qualification: the $\sigma$-gradient is not zero, so an optimizer step would still move $\sigma$. What is really happening at $\sigma \approx 0$ is that different noise draws $ϵ$ produce almost identical samples $z\approx \mu$, so the per-draw gradients $-0.7ϵ$ scatter symmetrically about $0$; their expectation $\mathbb{E}[-0.7ϵ]=0$ is what fades, not any individual gradient. Thus $\sigma$ receives conflicting single-sample pushes that average out, while $\mu$ keeps a clean, consistent signal. This is one practical reason we optimize $\log \sigma$ (keeping $\sigma >0$) and never land exactly on the spike. At $\mu =2.65,ϵ=1.3$ the single-draw $\sigma$-gradient limit is $-0.7\times 1.3=-0.91$, matching Ex 1. ✓

---

Example 5 — Cell C5: multivariate diagonal Gaussian

Forecast: Because the covariance is diagonal, every dimension is its own copy of the 1-D trick — no cross-terms. The Jacobian is diagonal, so $\partial z_i/\partial {\mu }_j=0$ unless $i=j$.

1. Element-wise sample. $z_i={\mu }_i+{\sigma }_i{ϵ}_i$: $z=(0+0.5\cdot 1,1+2\cdot (-0.5),-2+1\cdot 2)=(0.5,0,0)$. Why this step? The Hadamard form $z=\mu +\sigma \odot ϵ$ (using the $\odot$ just defined) applies the scalar trick coordinate by coordinate, because a diagonal covariance means the coordinates are independent.
2. Upstream. $\partial f/\partial z_i=2z_i=(1,0,0)$. Why this step? $f$ is a sum of squares, so its derivative w.r.t. each $z_i$ touches only that coordinate — the loss separates across dimensions.
3. $\mu$-gradient. $\partial f/\partial {\mu }_i=2z_i\cdot 1=(1,0,0)$. Why this step? Each coordinate's shift derivative is $1$ and, crucially, $\partial z_i/\partial {\mu }_j=0$ for $i\ne j$, so the chain rule keeps every dimension in its own lane.
4. $\sigma$-gradient. $\partial f/\partial {\sigma }_i=2z_i\cdot {ϵ}_i=(1\cdot 1,0\cdot (-0.5),0\cdot 2)=(1,0,0)$. Why this step? The stretch derivative for coordinate $i$ is its own noise ${ϵ}_i$; the diagonal Jacobian means no coordinate's noise leaks into another's gradient, which is why backprop here is cheap.

Verify: Dimensions $2$ and $3$ landed exactly on the target $0$, so their gradients are $0$ — the optimizer leaves them alone. Only dimension $1$ (off target at $0.5$) produces a push. The diagonal Jacobian means each coordinate's gradient depends only on its own ${ϵ}_i$. ✓

---

Example 6 — Cell C6: Monte-Carlo gradient over $S$ samples

Forecast: Each sample gives $\partial f/\partial \mu =2(z-3)$. Positive/negative noises spread $z$ around $\mu =2.0$; averaging should give a stable "$\mu$ is below $3$, push it up" signal near $2(\mu -3)=-2.0$.

1. Samples. $z^{(s)}=2+0.5{ϵ}^{(s)}=(2.65,1.35,2.2,1.8)$. Why this step? One transform per drawn noise gives the $S$ concrete latents we will differentiate at.
2. Per-sample $\mu$-gradients. $2(z^{(s)}-3)=(-0.7,-3.3,-1.6,-2.4)$. Why this step? $\partial z/\partial \mu =1$ so the $\mu$-gradient equals the upstream derivative at each sample.
3. Average. ${\hat{g}}_{\mu }=\frac{1}{4}(-0.7-3.3-1.6-2.4)=\frac{1}{4}(-8.0)=-2.0$. Why this step? Monte-Carlo integration: the mean of samples is an unbiased estimate of the expectation we cannot compute in closed form for real networks.

Verify: The exact expectation is ${\mathbb{E}}_{ϵ}[2(\mu +\sigma ϵ-3)]=2(\mu -3)=2(2.0-3)=-2.0$, because $\mathbb{E}[ϵ]=0$. Our 4-sample estimate hit $-2.0$ on the nose here (the symmetric $\pm$ noises cancelled perfectly). In real training $S=1$ often suffices because the estimate is unbiased and mini-batches supply extra averaging. ✓

---

Example 7 — Cell C7: real-world word problem (one VAE step)

Forecast: The two loss terms fight: reconstruction pulls up, prior pulls down, so $z$ should settle between $0$ and $0.9$. We expect $\mu$ to move toward that balance point ($0.45$) and $\log \sigma$ to shrink (prior discourages large spread).

1. Recover $\sigma$. $\sigma =e^{-0.7}\approx 0.4966$. Why this step? The network stores $\log \sigma$ precisely so $\sigma =e^{\log \sigma }>0$ automatically; we must convert to a real spread before we can sample.
2. Sample. $z=0.3+0.4966\times 0.5\approx 0.5483$. Why this step? The reparameterized transform turns the stored parameters plus fixed noise into a differentiable latent.
3. Upstream. $\partial f/\partial z=2(z-0.9)+2z=4z-1.8=4(0.5483)-1.8\approx 0.3932$. Why this step? $f$ is a sum of two squared terms; differentiating each and adding gives how the total loss reacts to $z$ — the start of every chain.
4. $\mu$-gradient. $\partial f/\partial \mu =0.3932\times 1=0.3932$. Why this step? The shift derivative $\partial z/\partial \mu =1$ passes the upstream signal straight through to $\mu$.
5. $\log \sigma$-gradient. Chain twice: $\frac{\partial f}{\partial \log \sigma }=\frac{\partial f}{\partial z}\cdot \underset{ϵ}{\underbrace{\frac{\partial z}{\partial \sigma }}}\cdot \underset{\sigma }{\underbrace{\frac{\partial \sigma }{\partial \log \sigma }}}=0.3932\times 0.5\times 0.4966\approx 0.09763$. Why this step? $z$ depends on $\sigma$ (through the stretch $ϵ$), and $\sigma =e^{\log \sigma }$ depends on $\log \sigma$; because we optimize the log, we need this extra link $\partial \sigma /\partial \log \sigma =\sigma$, which is exactly why the log-parameterization stays differentiable.
6. Update. $\mu \leftarrow 0.3-0.1\times 0.3932=0.2607$; $\log \sigma \leftarrow -0.7-0.1\times 0.09763=-0.70976$. Why this step? Gradient descent subtracts learning-rate times gradient; updating $\log \sigma$ (not $\sigma$) keeps the spread positive no matter how many steps we take.

Verify: Both gradients positive ⇒ $\mu$ decreased (from sample $0.5483$, above the balance $0.45$, so pull $z$ down) and $\log \sigma$ decreased (spread shrinks under prior pressure), matching the forecast. Every step flowed through $z=\mu +\sigma ϵ$ using ordinary backprop — no non-differentiable sampling anywhere. ✓

---

Example 8 — Cell C8: exam twist, when the trick cannot be used

Forecast: The location–scale trick needs a continuous variable you can smoothly shift and stretch. A category index is discrete — nudging a parameter can't move $z$ "half a category." Expect the trick to fail and a relaxation to step in.

First, name the encoder's outputs. The encoder produces class probabilities ${\pi }_1,{\pi }_2,{\pi }_3$ — one non-negative number per category, summing to $1$ — the softmax of the encoder's raw scores. So ${\pi }_k$ is simply "the model's believed probability that the latent equals category $k$." (In our worked check we use $\pi =(0.2,0.3,0.5)$.)

1. Test the requirement. The identity $z=\mu +\sigma ϵ$ produces a real number; a class label is an integer with no meaningful in-between. Sampling a category from $\pi$ is a step function of $\pi$ — flat, then a sudden jump to a different label — so $\partial z/\partial {\pi }_k$ is $0$ almost everywhere and undefined at the jumps. Why this step? This is precisely the non-differentiability the parent note warned about: a gradient of $0$ everywhere carries no direction to improve $\pi$.
2. Replacement — Gumbel-Softmax. Introduce a temperature $\tau >0$ (a knob that controls how "soft" the output is; small $\tau$ = crisp, large $\tau$ = blurry). Draw one Gumbel noise per class, $g_k=-\log (-\log u_k)$ with $u_k\sim \text{Uniform}(0,1)$, then form a soft one-hot vector $y=(y_1,y_2,y_3)$ whose $k$-th entry is $y_k=\frac{\exp ((\log {\pi }_k+g_k)/\tau )}{\sum _{j=1}^3\exp ((\log {\pi }_j+g_j)/\tau )}.$ Why this step? The Gumbel noise $g_k$ plays the role of $ϵ$ (parameter-free randomness), and the softmax with temperature $\tau$ makes $y$ a continuous, differentiable stand-in for the discrete pick. Crucially ${\sum }_k{y}_k=1$ by construction (a normalized softmax), so $y$ is a valid probability vector we can backprop through — see 5.2.03-Gumbel-Softmax.
3. Hard vs soft, and the $\tau$ limits. As $\tau \to 0^{+}$, one entry of $y$ rushes to $1$ and the rest to $0$: $y$ becomes a hard one-hot vector, recovering the true categorical draw. As $\tau \to \infty$, all logits are crushed together and $y\to (\frac{1}{3},\frac{1}{3},\frac{1}{3})$, a uniform blur. Why this step? Covering both limits shows why we train at a moderate $\tau$: small enough that $y$ resembles a real category, large enough that gradients still flow. (A common trick — "straight-through" — uses the hard one-hot in the forward pass but the soft $y$'s gradient in the backward pass.)
4. The other classic alternative. The score-function / REINFORCE estimator (see 6.1.02-Policy-Gradients) works for any distribution — discrete included — but has much higher variance than reparameterization. That trade-off is exactly why reparameterization is preferred whenever the latent is continuous, and Gumbel-Softmax is used to make a discrete latent continuous-ish.

Verify: With $\pi =(0.2,0.3,0.5)$, any Gumbel draw and any $\tau >0$, the soft one-hot $y$ sums to $1$ (a genuine probability distribution), so the relaxed sample is well-defined and differentiable — confirming the replacement is valid where the plain trick was not. This example marks the boundary of the matrix: outside it, we switch tools rather than force $z=\mu +\sigma ϵ$. ✓

---

Recall Self-check

In Ex 2, why did $\partial f/\partial \sigma$ become positive while it was negative in Ex 1? ::: Because $\partial z/\partial \sigma =ϵ$, and $ϵ$ flipped sign ($+1.3\to -1.3$); a larger $\sigma$ now pulls the sample away from the target. In the $\sigma \to 0$ limit (Ex 4), does any single-sample $\sigma$-gradient vanish? ::: No — each stays a finite $-0.7ϵ$; it is the expectation over $ϵ$ that goes to $0$, so $\sigma$ gets conflicting pushes that cancel while $\mu$ keeps a clean signal. Why must $\sigma >0$, and how do we guarantee it in code? ::: A spread cannot be zero or negative; we output $\log \sigma$ and set $\sigma =e^{\log \sigma }$, which is always positive. Why does the $S=4$ Monte-Carlo estimate in Ex 6 equal the exact value? ::: The symmetric $\pm ϵ$ pairs made $\mathbb{E}[ϵ]$ cancel exactly; in general it is only an unbiased estimate. Which single property of a distribution decides whether the trick applies at all? ::: Whether it is continuous and reparameterizable as a smooth function of parameter-free noise (Ex 8).

See also: 4.5.02-ELBO-Derivation · 4.5.04 Reparameterization trick (Hinglish)