VAE ELBO loss

Background

A variational autoencoder (VAE) is trained to maximize the evidence lower bound (ELBO), equivalently to minimize a loss with two terms: a reconstruction term (how well the decoder rebuilds the input) and a KL term that pulls the encoder's latent distribution $q(z\mid x)=\mathcal N(\mu,\sigma^2)$ toward a standard-normal prior. The KL term is what regularizes the latent space so it can be sampled to generate new data.

Problem statement

Implement vae_elbo_loss(x, x_recon, mu, logvar) returning the mean negative-ELBO loss. Use squared-error reconstruction (summed over features) and the closed-form Gaussian KL:

$$\text{recon}_i = \sum_j (x_{ij} - \hat x_{ij})^2, \qquad \text{KL}_i = -\tfrac12 \sum_j \big(1 + \log\sigma_j^2 - \mu_j^2 - \sigma_j^2\big)$$ $$\mathcal L = \frac{1}{N}\sum_i \big(\text{recon}_i + \text{KL}_i\big), \qquad \sigma_j^2 = e^{\,\text{logvar}_j}$$

Input

• x, x_recon — np.ndarray (N, d), input and its reconstruction.
• mu, logvar — np.ndarray (N, d), the encoder's latent mean and log-variance.

Output

A scalar float: the mean loss over the batch.

Examples

Example 1

Input:  x = [[1, 0]], x_recon = [[1, 0]], mu = [[0, 0]], logvar = [[0, 0]]
Output: 0.0

Explanation: perfect reconstruction makes the SSE term 0; with $\mu=0$ and $\sigma^2=1$ the posterior equals the prior, so KL is $-\tfrac12\sum(1+0-0-1)=0$. The loss is 0.

Constraints

• Convert logvar to variance with exp inside the KL.
• Sum both terms over the feature axis per example, then average over the batch.
• Return reconstruction plus KL (the quantity minimized = negative ELBO).

Notes

• The KL is always $\ge 0$ and is 0 exactly when $\mu=0,\ \sigma^2=1$ — the posterior matches the prior.
• Using logvar (rather than $\sigma$) keeps the variance positive and the optimization numerically stable.