Skip to content

← JAX Stochasticity step 25 of 25

Hard Primitives

Gaussian Reconstruction and KL

Jump to the task

Why this matters

A Gaussian variational model balances two terms: how well a latent variable explains the observation, and how far its distribution is from the prior. The evidence lower bound (ELBO) combines an expected log likelihood with a penalty KL(q || prior).

This exercise practices the reconstruction and KL calculations using a simplified score. It evaluates reconstruction at the mean mu instead of averaging over the latent distribution. That score is not the exact ELBO and is not guaranteed to be a lower bound on the log evidence.

Use a diagonal Gaussian q with mean vector mu and standard deviations exp(log_sigma), a standard normal prior, and a likelihood x | z that is Gaussian with mean z and identity covariance. There is no learned decoder in this exercise.

Worked mini-example

import jax.numpy as jnp

mu = jnp.array([1.0, -1.0])
log_sigma = jnp.array([0.0, 0.0])  # Unit standard deviations.
x = jnp.array([2.0, -1.0])

reconstruction = -0.5 * jnp.sum((x - mu) ** 2)  # -0.5
variance = jnp.exp(2 * log_sigma)
kl = 0.5 * jnp.sum(variance + mu ** 2 - 1 - 2 * log_sigma)  # 1.0
score = reconstruction - kl  # -1.5

The reconstruction penalizes squared error at the mean. The KL term penalizes the difference between q and the standard normal prior. Subtracting the KL makes a larger mismatch reduce the score.

Problem

Implement gaussian_elbo(mu, log_sigma, x). The entrypoint retains its original name; return the following simplified score:

reconstruction = -0.5 * sum((x - mu)^2)
kl = 0.5 * sum(exp(2 * log_sigma) + mu^2 - 1 - 2 * log_sigma)
score = reconstruction - kl

All inputs are finite float32 vectors with the same nonzero shape (d,). log_sigma contains log standard deviations, not log variances. Assume variances are positive and all intermediate calculations remain finite in float32. Return a float32 scalar. Do not draw random samples or add the corrections described below; the fixtures check the displayed score.

How this differs from the ELBO

For this identity-covariance likelihood, the exact expected squared error is sum((x - mu)^2 + exp(2 * log_sigma)). The true ELBO is therefore:

score - 0.5 * sum(exp(2 * log_sigma)) - 0.5 * d * log(2 * pi)

The first correction depends on the variational parameters. Omitting it changes the optimization problem. The second is a constant when the observation dimension and likelihood covariance are fixed.

Setting noise to zero is evaluation at the mean, not a random Monte Carlo draw. Sampling with a fixed random key can be reproducible without replacing the noise with zero. A genuine one-sample ELBO estimate is noisy and also need not be a lower bound on each individual draw; the expectation has the bound property.

Common pitfalls

  • Convert log standard deviation to variance with exp(2 * log_sigma).
  • Subtract KL. Do not negate both terms or return a loss to minimize.
  • Sum over dimensions, rather than averaging. The KL is for the whole diagonal Gaussian vector.
  • If mu = x = 0 and log_sigma = 0, this simplified score is zero. The true ELBO includes the variance and normalization corrections above.