Easy Primitives

VAE Reparameterization Trick

A VAE encoder predicts a distribution over latent vectors rather than a single vector. For a diagonal Gaussian, it supplies a mean mu and a log variance log_var for each coordinate.

To draw a sample, first draw independent standard normal noise eps, then shift and scale it:

z = mu + exp(0.5 * log_var) * eps

The factor exp(0.5 * log_var) is the standard deviation. exp(log_var) is the variance, so using it as the scale would give the wrong distribution. For a coordinate with mean 3, variance 4, and noise draw -0.5, the result is 3 + 2 * (-0.5) = 2.

This is the reparameterization trick. For a fixed noise draw, the sample depends differentiably on the encoder's outputs. Its derivative with respect to the mean is 1, and with respect to log variance is 0.5 * exp(0.5 * log_var) * eps. The noise is independent of the encoder parameters and is held fixed during that forward/backward computation.

Your task

Implement reparameterize(mu, log_var, seed). Both tensors have shape (N, d), with positive dimensions, and contain float32 values. Return one sample tensor of the same shape using the formula above.

The seeded noise draw is part of the task, not just its distribution:

In PyTorch

Seed a fresh CPU torch.Generator with int(seed), then call torch.randn(mu.shape, generator=gen) once. The exercise uses CPU inputs and float32 noise.

In JAX

Create key = jax.random.PRNGKey(int(seed)), then call jax.random.normal(key, mu.shape) once. Use Threefry with the runtime-pinned setting jax_threefry_partitionable = True.

Do not split the JAX key before this draw or consume extra random samples. The frameworks produce different samples from the same seed, so each mode has its own expected values. Results are compared numerically; you do not need to reproduce the other framework's random sequence.

For zero means and zero log variances, your output is exactly the noise draw. For very negative log variances, the scale becomes tiny and the result approaches the mean. No KL or reconstruction loss is needed here.

Reference: Auto-Encoding Variational Bayes.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Encode a distribution over latent values

A variational autoencoder (VAE) usually predicts a Gaussian distribution for each input. For a diagonal Gaussian, the encoder outputs a mean mu and log variance ell for each latent coordinate. The standard deviation is exp(ell / 2). Using exp(ell) would give the variance instead.

To sample a latent vector, draw independent standard normal noise epsilon and compute

$$z=\mu+\sigma\epsilon,\qquad\sigma=\exp(\ell/2).$$

This is the reparameterization trick. For a fixed noise draw, gradients can flow through the mean and scale: dz/dmu = 1 and dz/dell = 0.5 * sigma * epsilon. The random draw itself does not depend on the encoder parameters.

Balance reconstruction with a prior

The decoder tries to reconstruct the input from the sample. A second term keeps the encoder's distribution close to a chosen prior, commonly a standard normal distribution. Training usually minimizes a reconstruction loss plus this KL-divergence penalty, equivalent to minimizing the negative evidence lower bound when the likelihood and weighting match.

For a diagonal Gaussian and standard normal prior, the KL penalty is

$$\frac12\sum_j\left(\mu_j^2+e^{\ell_j}-1-\ell_j\right).$$

At mu = 0 and ell = 0, the distributions match and the penalty is zero. Squared-error reconstruction corresponds to a Gaussian likelihood with an appropriate scale. Follow the task's specified loss and reduction, since these set the balance between reconstruction and the KL term.

Sum latent coordinates and average examples as requested. Keep the same noise sample for a forward/backward pair; resampling would change the computation being differentiated. Very large log variances can overflow when exponentiated.

Detecting runtime…

VAE Reparameterization Trick

Your task

In PyTorch

In JAX

Encode a distribution over latent values

Balance reconstruction with a prior

Your code, saved here

Stuck?

PyTorch reference solution