Medium Primitives

VAE ELBO Loss

A variational autoencoder (VAE) encodes each input as a distribution over latent vectors, then decodes a sampled vector to reconstruct the input. This exercise combines two training terms:

Reconstruction error measures how far the decoder output is from the input.
KL divergence penalizes a difference between the encoder's distribution and a standard normal prior (zero mean and unit variance).

The encoder distribution is a diagonal Gaussian. mu gives its means, and log_var gives the natural logarithms of its variances. Coordinates are independent under this distribution; their variances are exp(log_var).

Your task

Implement vae_elbo_loss(reconstructed, original, mu, log_var):

reconstructed and original have shape (N, d).
mu and log_var have shape (N, d_z).
All inputs are floating-point tensors. N, d, and d_z are positive; the input dimension d need not equal the latent dimension d_z.

For each example, sum (reconstructed - original) ** 2 over its d input features. Average those sums over the batch to obtain the reconstruction loss.

For each example, compute

KL = 0.5 * sum(mu ** 2 + exp(log_var) - 1 - log_var)

over its d_z latent coordinates. Average those KL values over the batch. Return the reconstruction loss plus the KL loss as a Python float. Do not average over features or latent coordinates, and do not add a separate weight to either term.

For one example with original = [1, 2], reconstructed = [2, 2], mu = [1, 0], and log_var = [0, 0], reconstruction contributes 1 and KL contributes 0.5. The total is 1.5. A perfect reconstruction with zero means and zero log variances gives zero loss.

How this relates to the ELBO

The evidence lower bound (ELBO) is an objective that VAEs maximize. Its negative combines expected negative log likelihood with KL divergence. The sum of squared errors used here corresponds to a Gaussian decoder with fixed variance 0.5 per input coordinate, after dropping an additive constant. A reconstruction from a sampled latent vector supplies an estimate of the expectation. Other likelihoods or variances give different reconstruction losses or scales; squared error is not a universal VAE rule.

You receive the reconstructed values here, so no sampling is required. Use inputs for which the exponentials and total are representable. This exercise requests a float for comparison; a training loop should keep the loss as a tensor for differentiation.

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 ELBO Loss

Your task

How this relates to the ELBO

Encode a distribution over latent values

Balance reconstruction with a prior

Your code, saved here

Stuck?

PyTorch reference solution