VAE Loss (Negative ELBO)

Compute the loss of a variational autoencoder: the negative evidence lower bound (ELBO), split into its two parts.

A VAE is trained to do two things at once:

1. Reconstruct the input. The decoder outputs x_recon, a probability for each pixel. The reconstruction loss is binary cross entropy, summed over the features of each example:

   recon = -sum_j (x_j * log(x_recon_j) + (1 - x_j) * log(1 - x_recon_j))
   

2. Keep the latent codes close to the prior N(0, I). The KL divergence from the encoder's Gaussian N(mu, exp(log_var)) to the standard normal has a closed form, summed over latent dimensions:

   kl = -0.5 * sum_k (1 + log_var_k - mu_k^2 - exp(log_var_k))
   

The loss of one example is recon + kl. Average each of the three quantities over the batch and return (loss, recon, kl).

Function signature

def vae_loss(x, x_recon, mu, log_var):

Example

Input

x       = np.array([[1.0, 0.0]])
x_recon = np.array([[0.8, 0.1]])
mu      = np.array([[0.5, -0.5]])
log_var = np.array([[0.0, 0.0]])
vae_loss(x, x_recon, mu, log_var)

Output

(0.57850407, 0.32850407, 0.25)

Reasoning

The reconstruction loss is -(ln 0.8 + ln 0.9) = 0.2231 + 0.1054 = 0.3285. The latent variance is already 1, so the KL term only pays for the mean being off zero: -0.5 * ((1 + 0 - 0.25 - 1) * 2) = 0.25. The total is 0.5785.

Constraints and notes

• x and x_recon have shape (batch, d). mu and log_var have shape (batch, k).
• Clip x_recon into [1e-7, 1 - 1e-7] before taking logs.
• Sum over features and latent dimensions, then average over the batch.