Variational Autoencoder (VAE) Loss (ELBO)

Sequence Models & Generative Models DS practice problem on Onlearn.

Difficulty: medium.

Topics: Variational Autoencoder (VAE) Loss (ELBO), Kullback-Leibler Divergence, Reparameterization Trick, Evidence Lower Bound, Reconstruction Loss, Posterior Approximation, Probabilistic Graphical Models, Information Theory, Bayesian Inference, Deep Learning Optimization, Manifold Learning, Latent Variable Models, Variational Inference, Monte Carlo Methods, Stochastic Gradient Estimation, Regularization Techniques.

Implement a function that computes the Variational Autoencoder (VAE) loss, also known as the negative Evidence Lower Bound (ELBO). The VAE loss consists of two components: 1. Reconstruction Loss : Measures how well the decoder reconstructs the original input from the latent representation, computed as the mean (over the batch) of the sum of squared differences between original and reconstructed inputs. 2. KL Divergence : Measures how much the learned latent distribution diverges from a standard normal prior. The encoder outputs the mean (mu) and log variance (log var) of a Gaussian distribution in latent space. Your function should accept: x: original input data, shape (batch size, features) x reconstructed: reconstructed data from the decoder, shape (batch size, features) mu: mean vector of the latent Gaussian, shape (batch size, latent dim) log var: log variance vector of the latent Gaussian, shape (batch size, latent dim) Return a tuple of three floats: (total loss, reconstruction loss, kl divergence), where total loss = reconstruction loss + kl divergence. Use only NumPy.