Variational Autoencoder

The problem, from first principles

A Variational Autoencoder (VAE) learns a probabilistic latent space representation. The encoder maps input x to mean/logvar of a Gaussian distribution. The decoder samples from this distribution to reconstruct x. The loss combines reconstruction (MSE/BCE) and KL divergence regularization.

Where q is the encoder, p is the decoder, and p(z) = N(0, I).

From theory to code

Implement:

vae_loss(x_recon, x_true, mu, logvar)

Computes the evidence lower bound (ELBO) loss for VAE training.

Constraints

• x_recon, x_true shape: (N, D)
• mu, logvar shape: (N, latent_dim)
• Return scalar loss.
• Reconstruction loss is MSE.