Variational Autoencoders & ELBO Derivation

Standard deterministic autoencoders compress high-dimensional inputs $x$ into discrete latent codes $z = f(x)$. However, because deterministic autoencoders do not enforce continuity across latent space, the resulting manifold contains arbitrary unmapped regions (holes) where decoding novel vectors generates degraded artifacts.

Variational Autoencoders (VAEs) (Kingma and Welling, 2013) reformulate representation learning into probabilistic Bayesian inference, mapping each sample to a continuous probability distribution and optimizing the Evidence Lower Bound (ELBO).

1. Mathematical Derivation of the Evidence Lower Bound (ELBO)#

Let $x \in \mathbb{R}^D$ be observed data and $z \in \mathbb{R}^d$ be unobserved latent factors with prior distribution $p(z) = \mathcal{N}(0, I)$.

The marginal likelihood of observation $x$ integrates over all possible latent states:

Because true posterior $p(z | x) = \frac{p(x|z)p(z)}{p(x)}$ is computationally intractable to evaluate in high dimensions, we introduce a parameterized variational approximation $q_\phi(z | x) \approx p(z | x)$.

2. Closed-Form Analytical Derivation of Gaussian KL Divergence#

Assume the variational posterior is a diagonal multivariate Gaussian $q_\phi(z | x) = \mathcal{N}(\mu, \text{diag}(\sigma^2))$ and the prior is a standard normal distribution $p(z) = \mathcal{N}(0, I)$ in $\mathbb{R}^d$.

We compute the integral:

Evaluating the one-dimensional case for coordinate $j$:

1. Entropy of $q(z_j)$:

1. Cross-entropy under prior $p(z_j) = \mathcal{N}(0, 1)$:

Since $\mathbb{E}_{q}[z_j^2] = \mu_j^2 + \sigma_j^2$:

Subtracting the two terms:

Summing across all $d$ latent dimensions yields the exact analytical formula:

3. The Reparameterization Trick#

To compute gradients of the reconstruction loss with respect to encoder parameters $\phi$:

If we sample $z$ directly from stochastic distribution $\mathcal{N}(\mu(x), \Sigma(x))$, the sampling step creates a non-differentiable node in the computation graph. Gradients cannot backpropagate through random number generators.

The Reparameterization Trick reformulates random variable $z$ as a deterministic function of inputs, parameters, and an external auxiliary noise variable $\epsilon$:

Now, the expectation is taken over the parameter-free noise distribution $p(\epsilon)$:

Gradients flow seamlessly into $\mu$ and $\sigma$ via standard backpropagation:

4. From-Scratch Implementation: Zero Dependencies#

The following implementations compute the complete VAE Loss Function (Reconstruction Mean Squared Error + Closed-Form Gaussian KL Divergence) and reparameterization sampling from scratch in pure Python and pure Rust.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
import math
import random
from typing import List, Tuple

class VAEPrimitives:
    @staticmethod
    def reparameterize(mu: List[float], log_var: List[float]) -> List[float]:
        # z = mu + exp(0.5 * log_var) * epsilon, where epsilon ~ N(0, 1)
        z = []
        for m, lv in zip(mu, log_var):
            std = math.exp(0.5 * lv)
            eps = random.gauss(0.0, 1.0)
            z.append(m + std * eps)
        return z

    @staticmethod
    def kl_divergence_loss(mu: List[float], log_var: List[float]) -> float:
        # D_KL = -0.5 * sum(1 + log_var - mu^2 - exp(log_var))
        d_kl = 0.0
        for m, lv in zip(mu, log_var):
            d_kl += -0.5 * (1.0 + lv - (m ** 2) - math.exp(lv))
        return d_kl

    @staticmethod
    def reconstruction_loss(x_true: List[float], x_recon: List[float]) -> float:
        # Mean Squared Error as negative Gaussian log-likelihood surrogate
        return sum((t - r) ** 2 for t, r in zip(x_true, x_recon)) / len(x_true)

    @classmethod
    def total_vae_loss(
        cls,
        x_true: List[float],
        x_recon: List[float],
        mu: List[float],
        log_var: List[float],
        beta_kl: float = 1.0
    ) -> Tuple[float, float, float]:
        recon = cls.reconstruction_loss(x_true, x_recon)
        kl = cls.kl_divergence_loss(mu, log_var)
        total = recon + beta_kl * kl
        return total, recon, kl

if __name__ == "__main__":
    random.seed(42)
    # Latent vector dimension 3
    mu = [0.1, -0.2, 0.05]
    log_var = [-0.1, 0.0, 0.2]

    # Sample latent code z
    z_sample = VAEPrimitives.reparameterize(mu, log_var)
    print("Reparameterized Latent Vector z:")
    print(" ", [round(v, 4) for v in z_sample])

    # Compute loss on synthetic reconstruction
    x_target = [1.0, 0.0, 0.5, 0.2]
    x_pred = [0.95, 0.05, 0.48, 0.22]
    total_loss, recon_l, kl_l = VAEPrimitives.total_vae_loss(x_target, x_pred, mu, log_var)

    print("\nLoss Components:")
    print(f"  Reconstruction Loss (MSE): {recon_l:.6f}")
    print(f"  KL Divergence Penalty:     {kl_l:.6f}")
    print(f"  Total VAE Loss:            {total_loss:.6f}")

5. Active Knowledge Checks#