VAE: Reparameterisation and the ELBO

A variational autoencoder's encoder outputs a Gaussian $q(z|x) = \mathcal{N}(\mu, \sigma^2)$ for each input. Training maximises the ELBO, which means minimising reconstruction error plus the KL divergence from the prior $\mathcal{N}(0, I)$.

Implement:

• reparameterize(mu, logvar, eps): $z = \mu + e^{\frac{1}{2}\log\sigma^2} \cdot \epsilon$. This lets gradients flow through the sampling.
• kl_divergence(mu, logvar): the closed form $-\frac{1}{2}\sum_j(1 + \log\sigma_j^2 - \mu_j^2 - \sigma_j^2)$, summed over latent dimensions and averaged over the batch
• kl_grads(mu, logvar): gradients of that (batch-averaged) KL with respect to mu and logvar
• vae_loss(x, x_hat, mu, logvar, beta=1.0): the mean over the batch of (summed squared reconstruction error + β · KL per sample)