Reparameterization trick

The Core Problem

In a VAE, the encoder outputs parameters of a distribution $q_\phi(z|x)$ , typically $\mathcal{N}(\mu_\phi(x), \sigma_\phi^2(x))$ . We need to:

Sample $z \sim q_\phi(z|x)$
Compute loss involving $z$
Backpropagate to update $\phi$ (the encoder parameters)

The Reparameterization Trick: Derivation from Scratch

Step-by-Step Construction

Start with the property of Gaussian distributions:

If $\epsilon \sim \mathcal{N}(0, 1)$ (standard normal), then for any $\mu$ and $\sigma > 0$ :

$z = \mu + \sigma \cdot \epsilon \sim \mathcal{N}(\mu, \sigma^2)$

Why this step? This is the location-scale property of the Gaussian. We're shifting ( $\mu$ ) and stretching ( $\sigma$ ) a standard normal.

Proof:

$\mathbb{E}[z] = \mathbb{E}[\mu + \sigma\epsilon] = \mu + \sigma \cdot 0 = \mu$ ✓
$\text{Var}(z) = \text{Var}(\mu + \sigma\epsilon) = \sigma^2 \cdot \text{Var}(\epsilon) = \sigma^2 \cdot 1 = \sigma^2$ ✓

Now the critical rewrite:

Instead of: $z \sim q_\phi(z|x) = \mathcal{N}(\mu_\phi(x), \sigma_\phi^2(x))$

We write: $\epsilon \sim \mathcal{N}(0, 1)$ $z = \mu_\phi(x) + \sigma_\phi(x) \cdot \epsilon$

Why this step? Now $z$ is a deterministic function of $\phi$ (through $\mu_\phi$ and $\sigma_\phi$ ) and the random noise $\epsilon$ . The randomness lives only in $\epsilon$ , which doesn't depend on our parameters.

Gradient Flow

The loss typically looks like: $\mathcal{L}(\phi) = \mathbb{E}_{q_\phi(z|x)}[f(z, x)]$

Old way (doesn't work): $\nabla_\phi \mathcal{L} = \mathbb{E}_{q_\phi(z|x)}\left[\nabla_\phi f(z, x)\right] \quad \text{← STUCK! Can't compute } \nabla_\phi z$

New way (reparameterized): $\mathcal{L}(\phi) = \mathbb{E}_{p(\epsilon)}[f(g_\phi(\epsilon, x), x)]$

$\nabla_\phi \mathcal{L} = \mathbb{E}_{p(\epsilon)}\left[\nabla_\phi f(g_\phi(\epsilon, x), x)\right]$

Why this step? The expectation is now over $p(\epsilon)$ , which doesn't depend on $\phi$ . We can swap the gradient and expectation:

$= \mathbb{E}_{p(\epsilon)}\left[\frac{\partial f}{\partial z} \cdot \frac{\partial g_\phi}{\partial \phi}\right]$

The chain rule works! $\frac{\partial g_\phi}{\partial \phi}$ involves simple derivatives:

$\frac{\partial}{\partial \mu_\phi}(\mu_\phi + \sigma_\phi \epsilon) = 1$
$\frac{\partial}{\partial \sigma_\phi}(\mu_\phi + \sigma_\phi \epsilon) = \epsilon$

Worked example Worked Example 1: Single Sample Gradient

Setup: Encoder outputs $\mu = 2.0$ , $\sigma = .5$ . We sample $\epsilon = 1.3$ from $\mathcal{N}(0,1)$ . Loss function: $f(z) = (z - 3)^2$ (want $z$ close to 3). Step 1: Compute $z$ via reparameterization. $z = \mu + \sigma \cdot \epsilon = 2.0 + 0.5 \times 1.3 = 2.65$

Why this step? This gives us a concrete sample using our deterministic transform.

Step 2: Compute loss. $f(z) = (2.65 - 3)^2 = 0.1225$

Step 3: Compute $\frac{\partial f}{\partial z}$ (gradient of loss w.r.t. latent). $\frac{\partial f}{\partial z} = 2(z - 3) = 2(2.65 - 3) = -0.7$

Why this step? We need to know how $f$ changes with $z$ to apply the chain rule.

Step 4: Compute $\frac{\partial z}{\partial \mu}$ (how latent changes with mean). $\frac{\partial z}{\partial \mu} = \frac{\partial}{\partial \mu}(\mu + \sigma \epsilon) = 1$ Why this step? The reparameterization makes $z$ linearly dependent on $\mu$ .

Step 5: Chain rule for $\mu$ . $\frac{\partial f}{\partial \mu} = \frac{\partial f}{\partial z} \cdot \frac{\partial z}{\partial \mu} = -0.7 \times 1 = -0.7$

Interpretation: Increase $\mu$ to increase $z$ and reduce loss. Gradient descent will update: $\mu \leftarrow 2.0 - \text{lr} \times (-0.7) = 2.0 + 0.7 \times \text{lr}$ . ✓

Step 6: Compute $\frac{\partial z}{\partial \sigma}$ . $\frac{\partial z}{\partial \sigma} = \frac{\partial}{\partial \sigma}(\mu + \sigma \epsilon) = \epsilon = 1.3$ Why this step? The noise $\epsilon$ acts as the sensitivity factor for $\sigma$ .

Step 7: Chain rule for $\sigma$ . $\frac{\partial f}{\partial \sigma} = \frac{\partial f}{\partial z} \cdot \frac{\partial z}{\partial \sigma} = -0.7 \times 1.3 = -0.91$

Interpretation: Increasing $\sigma$ (with positive $\epsilon$ ) increases $z$ , reducing loss. ✓

Worked example Worked Example 2: Monte Carlo Estimation

Setup: We don't have analytic expectation, so we approximate with $S$ samples.

$\nabla_\phi \mathcal{L} = \mathbb{E}_{p(\epsilon)}\left[\nabla_\phi f(g_\phi(\epsilon, x), x)\right] \approx \frac{1}{S}\sum_{s=1}^S \nabla_\phi f(g_\phi(\epsilon^{(s)}, x), x)$ Step 1: Draw $S$ independent samples $\epsilon^{(1)}, \ldots, \epsilon^{(S)} \sim \mathcal{N}(0, 1)$ .

Why this step? Monte Carlo integration—we're approximating the expectation by sampling.

Step 2: For each sample $s$ , compute: $z^{(s)} = \mu_\phi(x) + \sigma_\phi(x) \cdot \epsilon^{(s)}$

Step 3: Compute individual gradients: $g^{(s)} = \nabla_\phi f(z^{(s)}, x)$

Step 4: Average: $\hat{g} = \frac{1}{S}\sum_{s=1}^S g^{(s)}$

Why this step? The law of large numbers: as $S \to \infty$ , $\hat{g} \to \mathbb{E}[g]$ . Even small $S$ (like 1-10) works well in practice due to the smoothness of neural network losses.

Practical VAE: Often $S=1$ (single sample) per training step! The stochasticity across mini-batches provides enough gradient signal.

Worked example Worked Example 3: Multivariate Gaussian

Setup: Latent is $d$ -dimensional with diagonal covariance.

$q_\phi(z|x) = \mathcal{N}(\mu_\phi(x), \text{diag}(\sigma_\phi^2(x)))$

Reparameterization: $\epsilon \sim \mathcal{N}(0, I_d) \quad \text{(d independent standard normals)}$ $z_i = \mu_{\phi,i}(x) + \sigma_{\phi,i}(x) \cdot \epsilon_i, \quad i=1,\ldots,d$

Why this step? With diagonal covariance, each dimension is independent. We just apply the trick element-wise.

In vector form: $z = \mu_\phi(x) + \sigma_\phi(x) \odot \epsilon$ where $\odot$ is element-wise multiplication (Hadamard product). Gradient: Each $\frac{\partial z_i}{\partial \mu_{\phi,j}} = \delta_{ij}$ and $\frac{\partial z_i}{\partial \sigma_{\phi,j}} = \epsilon_i \delta_{ij}$ . The Jacobian is diagonal, so backprop is efficient!

Why It Works: The Math

Proof sketch: Change of variables. If $z = \mu + \sigma \epsilon$ , then $dz = \sigma \, d\epsilon$ , and: $p_z(z) dz = p_\epsilon\left(\frac{z-\mu}{\sigma}\right) \frac{dz}{\sigma}$

The distributions match by construction. The expectations are equal.

The gradient trick works because:

The expectation over $\epsilon$ doesn't depend on $\phi$
Leibniz integral rule: $\nabla_\phi \int p(\epsilon) f(g_\phi(\epsilon)) d\epsilon = \int p(\epsilon) \nabla_\phi f(g_\phi(\epsilon)) d\epsilon$
We can push the gradient inside the expectation

Practical Implementation

# PyTorch pseudocode
def reparameterize(mu, log_var):
    """
    mu: [batch, latent_dim] - mean from encoder
    log_var: [batch, latent_dim] - log variance (for numerical stability)
    """
    std = torch.exp(0.5 * log_var)  # sigma = exp(log_var / 2)
    eps = torch.randn_like(std)      # sample epsilon ~ N(0, 1)
    z = mu + std * eps               # reparameterization
    return z

Common Mistakes

Common mistake Mistake 1: Forgetting to Detach

$\epsilon$ Wrong code:

eps = torch.randn_like(std)
eps.requires_grad = True  # DON'T DO THIS
z = mu + std * eps

Why it feels right: "Everything in a neural network needs gradients, right?"

Why it's wrong: $\epsilon$ is sampled from a fixed distribution. Gradients w.r.t. $\epsilon$ are meaningless—we're not optimizing the random seed! Setting requires_grad=True wastes computation and can cause bugs.

The fix: Keep $\epsilon$ as-is. torch.randn_like() creates a tensor with requires_grad=False by default. The gradient will automatically flow through mu and std via the addition and multiplication operations.

Common mistake Mistake 2: Using

$\sigma^2$ Instead of $\sigma$ Wrong:

var = torch.exp(log_var)
z = mu + var * eps  # multiplying by variance, not std!

Why it feels right: The encoder outputs variance, so just use it directly.

Why it's wrong: The reparameterization formula is $z = \mu + \sigma \epsilon$ , not $z = \mu + \sigma^2 \epsilon$ . Using variance instead of standard deviation scales the noise incorrectly. If $\sigma = 0.5$ , then $\sigma^2 = 0.25$ —you'd be underestimating the spread by half!

The fix: Take the square root. Since we typically parameterize with log variance (for numerical stability): $\sigma = \exp(\frac{1}{2}\log \sigma^2) = \exp(0.5 \times \text{log\_var})$

Generalizations

The trick extends beyond Gaussians:

Distribution	Reparameterization
Gaussian	$z = \mu + \sigma \epsilon$ , $\epsilon \sim \mathcal{N}(0,1)$
Exponential	$z = -\frac{1}{\lambda}\log u$ , $u \sim \text{Uniform}(0,1)$
Logistic	$z = \mu + s \log\frac{u}{1-u}$ , $u \sim \text{Uniform}(0,1)$
Laplace	$z = \mu - b \cdot \text{sgn}(u-0.5) \log(1-2

Key requirement: The distribution must have a location-scale form or an invertible CDF (inverse transform sampling).

Recall Explain to a 12-Year-Old

Imagine you're trying to teach a robot to draw. The robot has a "creativity knob" (that's $\sigma$ ) and a "center point" knob (that's $\mu$ ). Every time it draws, it picks a random direction and distance to move from the center.

The problem: if the drawing is too wobbly (random), the robot can't learn which knob settings make good art. It's like trying to improve your basketball shot while someone randomly pushes your arm—you can't tell if it was your aim or the push!

The reparameterization trick is like this: instead of the robot randomly picking both direction AND distance (too random to learn from), we give it a standard random step (like always rolling a regular die), and then the robot scales that step with its knobs. Now if the drawing is bad, the robot knows it was the knob settings, not the randomness.

The knobs control a transform of a fixed random step. That transform is smooth and learnable!

Connections

4.5.01-Variational-Autoencoders - the primary application
4.5.02-ELBO-Derivation - where this trick enables gradient-based optimization
4.3.01-KL-Divergence - the KL term in ELBO has a closed form Gaussians
5.2.03-Gumbel-Softmax - the discrete analog
2.4.01-Backpropagation - why gradients need differentiable paths
6.1.02-Policy-Gradients - contrast with REINFORCE (score function method)

Flashcards

#flashcards/ai-ml

What is the core problem that the reparameterization trick solves? :: In VAEs, we need to backpropagate through a stochastic sampling node $z \sim q_\phi(z|x)$ , but sampling is non-differentiable w.r.t. the distribution parameters $\phi$ . Gradients cannot flow through randomness.

What is the reparameterization trick for a Gaussian distribution?

Instead of sampling

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

directly, sample

\epsilon \sim \mathcal{N}(0, 1)

from a fixed distribution, then compute

z = \mu + \sigma \cdot \epsilon

. This makes

z

a differentiable function of

\mu

and

\sigma

Why can we move the gradient inside the expectation after reparameterization?

After reparameterization, the expectation is over

p(\epsilon)

which doesn't depend on parameters

\phi

\nabla_\phi \mathbb{E}_{p(\epsilon)}[f(g_\phi(\epsilon)] = \mathbb{E}_{p(\epsilon)}[\nabla_\phi f(g_\phi(\epsilon))]

. The Leibniz integral rule allows swapping gradient and integral when the integration domain is parameter-free.

What is

\frac{\partial z}{\partial \mu}

when

z = \mu + \sigma \epsilon

\frac{\partial z}{\partial \mu} = 1

(a constant). This means changes in

\mu

directly shift

z

by the same amount.

What is

\frac{\partial z}{\partial \sigma}

when

z = \mu + \sigma \epsilon

\frac{\partial z}{\partial \sigma} = \epsilon

. The gradient depends on the sampled noise value. When

\epsilon

is large,

\sigma

has more influence on

z

In practice, how many samples

S

do VAEs use for the reparameterization gradient estimator?

Often just

S=1

(single sample per training step). The mini-batch stochasticity across many training steps provides sufficient gradient signal despite the variance of a single-sample estimator.

Why do we typically parameterize with log-variance instead of variance?

Numerical stability. Variance must be positive, but log-variance can be any real number. Computing

\sigma = \exp(0.5 \cdot \text{log\_var})

ensures

\sigma > 0

and avoids optimization constraints.

What is the correct formula to get

\sigma

from log-variance?

\sigma = \exp(0.5 \cdot \log\sigma^2) = \exp(\frac{1}{2} \text{log\_var})

, because

\log \sigma = \frac{1}{2}\log\sigma^2

Can the reparameterization trick be applied to discrete distributions like Categorical?

Not directly. Discrete sampling has no gradient (argmax is non-differentiable). Alternatives include Gumbel-Softmax (continuous relaxation) or REINFORCE (score function estimator).

For a multivariate Gaussian with diagonal covariance, how does reparameterization work?

Apply element-wise:

\epsilon \sim \mathcal{N}(0, I_d)

and

z = \mu + \sigma \odot \epsilon

where

\odot

is element-wise multiplication. Each dimension is independent.

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Dekho, VAE jaise generative models train karte waqt ek bada problem ata hai: humein encoder se nikle hue random sample $z$ ke through backpropagation karna padta hai

The Core Problem

The Reparameterization Trick: Derivation from Scratch

Step-by-Step Construction

Gradient Flow

Why It Works: The Math

Practical Implementation

Common Mistakes

Generalizations

Connections

Flashcards

Concept Map

Hinglish (regional understanding)

Go deeper — visual, from zero

Test yourself — Generative Models

Connections