VAE Reparameterization Trick - Generative Models

VAE Reparameterization Trick

What you'll build

The sampling step of a VAE: draw z ~ N(μ, σ²) in a way that lets gradients flow back into the encoder outputs μ and log σ².

Background

A VAE encoder outputs a Gaussian per input, and the decoder needs a sample from it. "Sample from N(μ, σ²)" is not a differentiable operation in μ and σ, so backprop would stop there. The trick moves the randomness into an input that does not depend on the parameters:

• draw ε ~ N(0, I) (independent of μ, σ)
• set z = μ + σ ⊙ ε, a deterministic, differentiable function of μ and σ

Then ∂z/∂μ = 1 and ∂z/∂σ = ε, so ordinary backprop trains the encoder.

Formula

σ = exp(0.5 · log_var)          # log_var = log σ², so this is sqrt(σ²)
ε ~ N(0, I)                     # shape (batch, latent_dim)
z = μ + σ ⊙ ε                   # elementwise

Encoders predict log_var rather than σ because it is unconstrained (any real number) while σ must be positive.

Inputs

• mu: float array (batch, latent_dim)
• log_var: float array (batch, latent_dim), any real values (very negative means σ ≈ 0)

Expected return format

• z: float array (batch, latent_dim)

Conventions (the tests rely on these)

• The starter code calls np.random.seed(42) first. Keep it, then draw the noise once with np.random.randn(*mu.shape), so the tests can reproduce your ε.
• With log_var very negative, σ underflows to 0 and z = mu exactly.

Constraints

NumPy only, vectorised (no Python loops).

Follow-up

• Why does the KL term in the VAE loss stop σ from collapsing to 0?
• How would you reparameterize a categorical latent? (Gumbel-softmax.)

Input: mu = np.array([[0.0, 1.0]]), log_var = np.array([[0.0, 0.0]])
Output: array([[0.4967, 0.8617]])

Input: mu = np.array([[0.0, 1.0]]), log_var = np.log(np.array([[4.0, 4.0]]))
Output: array([[0.9934, 0.7235]])

import numpy as np

def reparameterize(mu, log_var):
    """
    Sample z ~ N(mu, exp(log_var)) with the reparameterization trick.

    Args:
        mu: Mean of the latent distribution, shape (batch, latent_dim)
        log_var: Log variance of the latent distribution, shape (batch, latent_dim)

    Returns:
        z: Sampled latent vectors, shape (batch, latent_dim)
    """
    np.random.seed(42)  # keep: draw eps once with np.random.randn(*mu.shape)
    # Your code here
    pass