D7.3 · LESSON 3 OF 6

Reparameterization

Explain the different training objectives of autoencoders, VAEs, GANs and diffusion models.

THE IDEA

Sampling that gradients can pass through

The VAE must sample a code z from the encoder's distribution, and sampling is random: you cannot differentiate through a random draw. The fix is to move the randomness outside:

z = μ + σ · ε, ε ~ N(0, 1)

The noise ε is drawn separately and treated as a constant input. Now z is an ordinary function of μ and σ, with ∂z/∂μ = 1 and ∂z/∂σ = ε. The sample is still random (a new ε each time), but the loss can send gradients to the encoder.

mean μ2scale σ0.5noise ε-1

z = μ + σ·ε = 2 + 0.5·-1 = 1.5 ∂z/∂μ = 1 ∂z/∂σ = ε = -1 (randomness lives only in ε, so gradients reach μ and σ)

Blue dot: μ. Orange dot: the sampled z. Moving ε moves z; the encoder's μ and σ stay differentiable.

Move μ, σ and ε. z is the orange dot.

Predict first. Mean 2, scale 0.5, noise −1. What latent is produced?

NumPy

eps = rng.standard_normal(n)     # randomness, no parameters in it
z = mu + sigma * eps             # differentiable in mu and sigma

Chapter context

Generative modeling tries to represent or sample from a data distribution. Different model families construct that distribution differently; a plausible-looking sample alone does not tell us how the model was trained or how well it covers the data.

TRY IT YOURSELF

Part 1. Do it hands-on

Reparameterise a sample, differentiate through it by hand, then code the sampler and its gradients.

mean μ2scale σ0.5noise ε-1

z = μ + σ·ε = 2 + 0.5·-1 = 1.5 ∂z/∂μ = 1 ∂z/∂σ = ε = -1 (randomness lives only in ε, so gradients reach μ and σ)

Blue dot: μ. Orange dot: the sampled z. Moving ε moves z; the encoder's μ and σ stay differentiable.

Fixed numbers: μ = 2, σ = 0.5, ε = −1, loss L = (z − 3)².

z dL/dμ dL/dσ mean of z for ε ∈ [−1,0,1,2]

To do: Mean 2, scale 0.5, noise ε = −1. What latent z = μ + σ·ε is produced? Type it
To do: Loss L = (z − 3)² with the same numbers. Type dL/dμ (the noise ε stays fixed)
To do: Type dL/dσ for the same loss
To do: Use the four noise draws ε = [−1, 0, 1, 2] with μ = 2, σ = 0.5. What is the mean of the four latents? Type it

Part 2. Do it in code

reparameterize(mu, sigma, eps): z = mu + sigma * eps (works for numbers and arrays).
reparam_grads(mu, sigma, eps, target): for the loss L = (z − target)² return (z, dL/dmu, dL/dsigma) treating eps as a fixed constant.
sample_latents(mu, sigma, n, seed): draw n latents with np.random.default_rng(seed).standard_normal(n) as the noise and reparameterize. Return an array of length n.

import numpy as np

def reparameterize(mu, sigma, eps):
    pass

def reparam_grads(mu, sigma, eps, target):
    pass

def sample_latents(mu, sigma, n, seed):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

For mean 2, scale .5 and noise -1, what latent is produced?

A1.5, because z = μ + σ·ε = 2 + 0.5·(−1).B−0.5, because the noise is subtracted from the scale.C2.5, because the noise is added to the mean without being scaled.

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

A VAE uses $z=\mu+\sigma\odot\epsilon$ , with $\epsilon\sim\mathcal N(0,I)$ . A standard diffusion forward marginal is

x_t=\sqrt{\bar\alpha_t}x_0+\sqrt{1-\bar\alpha_t}\epsilon.

Here alpha-bar summarizes retained signal. This is a noising equation, not a complete sampling algorithm.

A VAE optimizes an evidence lower bound:

\operatorname{ELBO}=E_{q_\phi(z\mid x)}[\log p_\theta(x\mid z)]-D_{KL}(q_\phi(z\mid x)\Vert p(z)).

For diagonal Gaussian q and standard normal p, its KL is $\frac12\sum_j(\mu_j^2+\sigma_j^2-1-\log\sigma_j^2)$ . The original GAN minimax objective is

\min_G\max_D E_{x\sim p_{data}}[\log D(x)]+E_{z\sim p(z)}[\log(1-D(G(z)))].

Practical generator losses often use a non-saturating variant. A common diffusion noise-prediction loss is $E\|\epsilon-\epsilon_\theta(x_t,t)\|^2$ under a specified noising schedule.

Worked example

For a scalar diffusion example, take x0=2, alpha-bar=0.64 and epsilon=−1. The signal term is 0.8×2=1.6 and the noise term is 0.6×(−1)=−0.6, giving xt=1.0. The clean value cannot be inferred from xt alone without additional assumptions or a learned model.

For q=N(1,1) against a standard normal prior, the one-dimensional VAE KL is 0.5×(1²+1−1−ln 1)=0.5 nats. This is only the regularization term, not the complete VAE objective.

Reference code

The following example uses the Python standard library and does not call a remote service.

from math import sqrt, isclose
x0, alpha, noise = 2,.64,-1
xt=sqrt(alpha)*x0+sqrt(1-alpha)*noise
assert isclose(xt,1)
assert 2+.5*(-1) == 1.5
from math import log
assert isclose(.5*(1**2+1**2-1-log(1**2)),.5)

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

A VAE's approximate posterior is not identical to its prior. GAN loss is not a universal sample-quality score. A diffusion noising equation is not sufficient to implement the reverse model.