VAE Loss Function - Generative Models

VAE Loss (Negative ELBO)

What you'll build

The training loss of a Gaussian VAE: a reconstruction term plus the KL divergence between the encoder's distribution and the standard-normal prior.

Background

Maximising the evidence lower bound (ELBO) is the same as minimising

L = E_q[ -log p(x | z) ] + KL( q(z | x) || N(0, I) )

The first term makes the decoder reproduce its input; the second keeps every posterior close to the prior, so random z ~ N(0, I) decode to realistic samples. With a Gaussian decoder the first term becomes a squared error; with a diagonal-Gaussian encoder the KL has a closed form.

Formula

For example i with latent dimension d:

KL_i = -0.5 · Σ_{j=1..d} ( 1 + log_var_ij − mu_ij² − exp(log_var_ij) )

This course uses the standard convention (Kingma & Welling; the official PyTorch VAE example), averaged over the batch:

• recon_loss = mean((x − x_reconstructed)²) over all elements (plain MSE)
• kl_loss = mean_i KL_i: sum over latent dims, mean over the batch
• total_loss = recon_loss + kl_loss

The capstone E2E: Variational Autoencoder uses exactly the same convention.

Inputs

• x, x_reconstructed: float arrays (batch, features)
• mu, log_var: float arrays (batch, latent_dim)

Expected return format

A tuple of three Python floats, each rounded to 4 decimals, in this order: (total_loss, recon_loss, kl_loss)

Edge cases

• mu = 0, log_var = 0 (posterior equals the prior) gives kl_loss = 0.0.
• Each latent dim's KL is ≥ 0, so kl_loss ≥ 0.

Constraints

NumPy only, vectorised.

Follow-up

• Why weight the KL with β (β-VAE), and what is "posterior collapse"?
• How does the loss change for binary data (BCE reconstruction)?

Input: x = [[1.0, 2.0]], x_reconstructed = [[0.0, 0.0]], mu = [[0.5, -0.5]], log_var = [[0.0, 1.0]]
Output: (3.1091, 2.5, 0.6091)

Input: x = x_reconstructed = zeros((2, 3)), mu = [[1.0, 1.0], [0.0, 0.0]], log_var = zeros((2, 2))
Output: (0.5, 0.0, 0.5)

import numpy as np

def vae_loss(x, x_reconstructed, mu, log_var):
    """
    Compute the VAE loss (negative ELBO).

    Args:
        x: Original input, shape (batch, features)
        x_reconstructed: Decoder output, shape (batch, features)
        mu: Latent mean, shape (batch, latent_dim)
        log_var: Latent log variance, shape (batch, latent_dim)

    Returns:
        (total_loss, recon_loss, kl_loss): Python floats rounded to 4 decimals.
        recon_loss is the MSE over all elements; kl_loss sums over latent dims
        and averages over the batch; total_loss = recon_loss + kl_loss.
    """
    # Your code here
    pass